add batch_norm_GB

Author: Yin Hongyun

diopi_test/python/configs/diopi_configs.py

@@ -152,6 +152,46 @@
         ),
     ),
 
+    "batch_norm_GB": dict(
+        name=["batch_norm_GB"],
+        interface=['CustomizedTest'],
+        dtype=[np.float32, np.float16, np.float64],
+        atol=1e-3,
+        rtol=1e-4,
+        atol_half=1e-1,
+        rtol_half=1e-2,
+        para=dict(
+            training=[True, True, True],
+            momentum=[0.01, 0.01, 0.01],
+            axis=[0, 1, 2],
+            eps=[1e-4, 1e-4, 1e-4],
+        ),
+        tensor_para=dict(
+            args=[
+                {
+                    "ins": ["input"],
+                    "shape": ((2, 64, 32, 32),(2, 64, 32, 32),(2, 64, 32, 32)),
+                    "gen_fn": "Genfunc.randn",
+                },
+                {
+                    "ins": ["running_mean"],
+                    "shape": ((2,), (64,), (32,)),
+                    "gen_fn": "Genfunc.zeros",
+                },
+                {
+                    "ins": ["running_var"],
+                    "shape": ((2,), (64,), (32,)),
+                    "gen_fn": "Genfunc.ones",
+                },
+                {
+                    "ins": ["weight", "bias"],
+                    "shape": ((2,), (64,), (32,)),
+                    "gen_fn": "Genfunc.randn",
+                },
+            ]
+        ),
+    ),
+    
     # FIXME batch_norm输入0size的张量报错
     'batch_norm': dict(
         name=["batch_norm"],

diopi_test/python/conformance/customized_test.py

@@ -891,3 +891,21 @@ def pool3d(input, kernel_size, stride, padding, dilation, ceil_mode, count_inclu
     def layer_normGB(input, weight, bias, eps, normalized_shape):
         return torch.nn.functional.layer_norm(input=input, weight=weight, bias=bias, eps=eps, normalized_shape=normalized_shape)
 
+    def batch_norm_GB(input, running_mean, running_var, weight, bias, training=False, momentum=0.1, eps=1e-05, axis=1):
+        dim = input.dim()
+        dims = list(range(dim))
+        dims.remove(axis)
+        dims.insert(1, axis)
+        permuted_input = input.permute(dims)
+        out = torch.nn.functional.batch_norm(
+            permuted_input,
+            running_mean,
+            running_var,
+            weight=weight,
+            bias=bias,
+            training=training,
+            momentum=momentum,
+            eps=eps,
+        )
+        out = out.permute(dims)
+        return out

diopi_test/python/conformance/diopi_functions.py

@@ -2823,6 +2823,52 @@ def batch_norm(
     return out
 
 
+def batch_norm_GB(
+    input,
+    running_mean,
+    running_var,
+    weight,
+    bias,
+    training=False,
+    momentum=0.1,
+    eps=1e-05,
+    axis=1
+) -> Tensor:
+    dim = input.size().len
+    dim = [i for i in range(dim) if i!= axis]
+    dtype = Dtype.float32 if input.get_dtype() == Dtype.float16 else None
+    _, save_mean = reduce_op_process(input, dim, dtype=dtype)
+    save_invstd = raw_like(save_mean)
+
+    if not training:
+        assert (
+            running_mean is not None and running_var is not None
+        ), "if not trainging, running_mean and running_var must be defined"
+
+    out = raw_like(input)
+    func = check_function("diopiBatchNormGB")
+    ret = func(
+        input.context(),
+        out,
+        save_mean,
+        save_invstd,
+        input,
+        weight,
+        bias,
+        running_mean,
+        running_var,
+        training,
+        momentum,
+        eps,
+        axis
+    )
+
+    check_returncode(ret)
+    GLOBAL_STATE["batch_norm_save_mean"] = save_mean
+    GLOBAL_STATE["batch_norm_save_invstd"] = save_invstd
+    return out
+
+
 def batch_norm_stats(input, eps):
     func = check_function("diopiBatchNormStats")
     # cuda accumulate dtype mapping

diopi_test/python/conformance/global_op_list.py

@@ -11,6 +11,7 @@
     "conv2d": ["2d", "input", "weight"],
     "conv3d": ["3d", "input", "weight"],
     "batch_norm": ["input"],
+    "batch_norm_GB": ["input", "running_mean", "running_var"],
     "adaptive_avg_pool2d": ["2d", "input"],
     "adaptive_max_pool2d": ["2d", "input"],
     "adaptive_avg_pool3d": ["3d", "input"],
@@ -64,6 +65,7 @@
 
 ops_with_states = {
     "batch_norm": {"running_mean", "running_var"},
+    "batch_norm_GB": {"running_mean", "running_var"},
     "sgd": {"buf", "param"},
     "fill_": {"input"},
     "zero_": {"input"},

impl/torch/functions/functions.cpp

@@ -2557,6 +2557,29 @@ diopiError_t diopiBatchNorm(diopiContextHandle_t ctx, diopiTensorHandle_t out, d
     return diopiSuccess;
 }
 
+diopiError_t diopiBatchNormGB(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiTensorHandle_t save_mean, diopiTensorHandle_t save_invstd,
+                            diopiConstTensorHandle_t input, diopiConstTensorHandle_t weight, diopiConstTensorHandle_t bias, diopiTensorHandle_t running_mean,
+                            diopiTensorHandle_t running_var, bool training, double momentum, double eps, int64_t axis) {
+    impl::aten::setCurStream(ctx);
+    auto atInput = impl::aten::buildATen(input);
+    auto atWeight = impl::aten::buildATen(weight);
+    auto atBias = impl::aten::buildATen(bias);
+    auto atRunningMean = impl::aten::buildATen(running_mean);
+    auto atRunningVar = impl::aten::buildATen(running_var);
+    auto atOut = impl::aten::buildATen(out);
+    auto atSaveMean = impl::aten::buildATen(save_mean);
+    auto atSaveInvstd = impl::aten::buildATen(save_invstd);
+
+    std::vector<int64_t> dims(atInput.dim());
+    std::iota(dims.begin(), dims.end(), 0); 
+    std::swap(dims[1], dims[axis]);
+    auto permutedInput = atInput.permute(dims);
+    CALL_ATEN_CUDA_FUNC(
+        native_batch_norm_out, atOut, atSaveMean, atSaveInvstd, permutedInput, atWeight, atBias, atRunningMean, atRunningVar, training, momentum, eps);
+    atOut = atOut.permute(dims);
+    return diopiSuccess;
+}
+
 diopiError_t diopiSlice(diopiContextHandle_t ctx, diopiTensorHandle_t null_out, diopiConstTensorHandle_t input, int64_t dim, int64_t start, int64_t end,
                         int64_t step) {
     impl::aten::setCurStream(ctx);

proto/include/diopi/functions.h

@@ -120,6 +120,13 @@ DIOPI_API diopiError_t diopiBatchNorm(diopiContextHandle_t ctx, diopiTensorHandl
                                       diopiConstTensorHandle_t input, diopiConstTensorHandle_t weight, diopiConstTensorHandle_t bias,
                                       diopiTensorHandle_t running_mean, diopiTensorHandle_t running_var, bool training, double momentum, double eps);
 
+/**
+ * @brief Applies Batch Normalization.
+ */
+DIOPI_API diopiError_t diopiBatchNormGB(diopiContextHandle_t ctx, diopiTensorHandle_t out, diopiTensorHandle_t save_mean, diopiTensorHandle_t save_invstd,
+                                      diopiConstTensorHandle_t input, diopiConstTensorHandle_t weight, diopiConstTensorHandle_t bias,
+                                      diopiTensorHandle_t running_mean, diopiTensorHandle_t running_var, bool training, double momentum, double eps, int64_t axis);
+
 /**
  * @brief Computes the mean and inverse standard deviation across a batch of data for Synchronized Batch Normalization (SyncBN).
  * @param[in] ctx Context environment.