[Prim][Pir] Add support for instanceNorm1D NC data layout (#69460)

cubehan3 · web-flow · commit 78cca74a0abe · 2024-11-25T14:03:10.000+08:00
* add support for NC data layout

* change instancenorm static shape implementation

* polish testcase
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h b/paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h
@@ -750,105 +750,101 @@ std::tuple<Tensor, Tensor, Tensor> instance_norm_decomp(
     const paddle::optional<Tensor>& scale,
     const paddle::optional<Tensor>& bias,
     float epsilon) {
-  if (has_dynamic_shape(x.shape())) {
-    auto org_dtype = x.dtype();
-    Tensor x_cast = ConverToMT<T>(x);
+  auto org_dtype = x.dtype();
+  Tensor x_cast = ConverToMT<T>(x);
+  const std::vector<int64_t> x_dims = x.shape();
 
+  if (has_dynamic_shape(x_dims)) {
     std::vector<int64_t> axis;
-    auto x_dim = x.shape();
-    for (size_t i = 2; i < x_dim.size(); i++) {
+    for (size_t i = 2; i < x_dims.size(); i++) {
       axis.push_back(static_cast<int64_t>(i));
     }
+    bool reduce_axes_empty = axis.empty();
 
     // out = (x - mean(x)) / sqrt(var + epsilon))
     // var = mean((x-mean(x))^2)
-    auto mean_ = mean_decomp<T>(x_cast, axis, true);
+    auto mean_ =
+        reduce_axes_empty ? x_cast : mean_decomp<T>(x_cast, axis, true);
     auto difference = x_cast - mean_;
     auto var_tmp1 = difference * difference;
-    auto variance = mean_decomp<T>(var_tmp1, axis, true);
+    auto variance =
+        reduce_axes_empty ? var_tmp1 : mean_decomp<T>(var_tmp1, axis, true);
     auto var_shape = shape<T>(variance);
     auto var_tmp3 =
         variance + full_scalar<T>(epsilon, variance.dtype(), variance.place());
     auto rsqrt_var = rsqrt<T>(var_tmp3);
     auto out = difference * rsqrt_var;
 
-    int dim_size = x_dim.size();
+    int dim_size = x_dims.size();
     auto x_shape_tensor = shape<T>(x);
     std::vector<Tensor> slice_shape_concat;
 
-    auto shape_1 = full<T>({1}, 1, x_shape_tensor.dtype());
-    auto shape_2 =
-        cast<T>(get_slice<T>(x_shape_tensor, 1), x_shape_tensor.dtype());
-    auto shape_3 = full<T>({dim_size - 2}, 1, x_shape_tensor.dtype());
-
-    slice_shape_concat.push_back(shape_1);
-    slice_shape_concat.push_back(shape_2);
-    slice_shape_concat.push_back(shape_3);
+    slice_shape_concat.push_back(full<T>({1}, 1, x_shape_tensor.dtype()));
+    slice_shape_concat.push_back(
+        cast<T>(get_slice<T>(x_shape_tensor, 1), x_shape_tensor.dtype()));
+    if (dim_size > 2) {
+      slice_shape_concat.push_back(
+          full<T>({dim_size - 2}, 1, x_shape_tensor.dtype()));
+    }
     auto slice_shape_tensor = concat<T>(slice_shape_concat, 0);
 
-    Tensor scale_cast;
     if (scale) {
-      scale_cast = backend::reshape<T>(scale.get(), slice_shape_tensor);
+      auto scale_cast = backend::reshape<T>(scale.get(), slice_shape_tensor);
       scale_cast = ConverToMT<T>(scale_cast);
       out = out * scale_cast;
     }
-    Tensor bias_cast;
+
     if (bias) {
-      bias_cast = backend::reshape<T>(bias.get(), slice_shape_tensor);
+      auto bias_cast = backend::reshape<T>(bias.get(), slice_shape_tensor);
       bias_cast = ConverToMT<T>(bias_cast);
       out = out + bias_cast;
     }
 
     std::vector<int64_t> res_shape(1, -1);
     auto mean_out = reshape<T>(mean_, res_shape);
     auto variance_out = reshape<T>(rsqrt_var, res_shape);
-
-    Tensor res;
-    res = ConverToOrig<T>(out, org_dtype);
+    auto res = ConverToOrig<T>(out, org_dtype);
 
     return std::make_tuple(res, mean_out, variance_out);
   }
-  auto org_dtype = x.dtype();
-  Tensor x_cast = ConverToMT<T>(x);
 
-  std::vector<int64_t> axis;
-  auto x_dim = x.shape();
-  for (size_t i = 2; i < x_dim.size(); i++) {
-    axis.push_back(static_cast<int64_t>(i));
-  }
+  // static shape
+  const int64_t N = x_dims[0];
+  const int64_t C = x_dims[1];
+  const int64_t NxC = N * C;
+  const int64_t sample_size = x.numel() / N / C;
+  const std::vector<int64_t> shape{NxC, sample_size};
+  const std::vector<int64_t> rdims{1};
 
-  // out = (x - mean(x)) / sqrt(var + epsilon))
-  // var = mean((x-mean(x))^2)
-  auto mean_ = mean_decomp<T>(x_cast, axis, true);
-  auto difference = x_cast - mean_;
-  auto var_tmp1 = difference * difference;
-  auto variance = mean_decomp<T>(var_tmp1, axis, true);
+  auto x_arr = reshape<T>(x_cast, shape);
+  auto mean_ = mean_decomp<T>(x_arr, rdims, true);
+  auto difference = x_arr - mean_;
+  auto variance = mean_decomp<T>(difference * difference, rdims, true);
   auto var_tmp3 = variance + epsilon;
   auto rsqrt_var = rsqrt<T>(var_tmp3);
   auto out = difference * rsqrt_var;
 
-  std::vector<int64_t> slice_shape(x_dim.size(), 1);
-  slice_shape[1] = x_dim[1];
+  std::vector<int64_t> slice_shape(x_dims.size(), 1);
+  slice_shape[1] = x_dims[1];
 
-  Tensor scale_cast;
+  out = reshape<T>(out, x_dims);
   if (scale) {
-    scale_cast = reshape<T>(scale.get(), slice_shape);
+    auto scale_cast = reshape<T>(scale.get(), slice_shape);
     scale_cast = ConverToMT<T>(scale_cast);
     out = out * scale_cast;
   }
-  Tensor bias_cast;
+
   if (bias) {
-    bias_cast = reshape<T>(bias.get(), slice_shape);
+    auto bias_cast = reshape<T>(bias.get(), slice_shape);
     bias_cast = ConverToMT<T>(bias_cast);
     out = out + bias_cast;
   }
 
   std::vector<int64_t> res_shape(1, -1);
   auto mean_out = reshape<T>(mean_, res_shape);
   auto variance_out = reshape<T>(rsqrt_var, res_shape);
+  auto res = ConverToOrig<T>(out, org_dtype);
 
-  Tensor res;
-  res = ConverToOrig<T>(out, org_dtype);
   return std::make_tuple(res, mean_out, variance_out);
 }
 
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -1412,19 +1412,23 @@ void instance_norm_grad(const Tensor& x,
   // x_grad = scale * inv_var * (y_grad - y_grad.mean(2,3) - x_hat * (y_grad *
   // x_hat).mean((h,w)))
   if (x_grad) {
+    bool is_reduce_empty = reduce_axes.empty();
     Tensor scale_data_tensor =
         scale.get_ptr() ? scale.get()
                         : full<T>(IntArray({c}), 1., x.dtype(), x.place());
     auto unsqueeze_shape = get_unsqueeze_dims(scale_data_tensor, n_reduce_axes);
     auto scale_data = reshape<T>(scale_data_tensor, unsqueeze_shape);
     auto promoted_scale = ConverToMT<T>(scale_data);
-    auto result =
-        (promoted_scale * std_inv) *
-        (promoted_y_grad -
-         promoted_y_grad.sum(reduce_axes, promoted_y_grad.dtype(), true) / hw -
-         (x_hat * ((promoted_y_grad * x_hat)
-                       .sum(reduce_axes, promoted_y_grad.dtype(), true) /
-                   hw)));
+    auto tmp1 =
+        is_reduce_empty
+            ? promoted_y_grad
+            : promoted_y_grad.sum(reduce_axes, promoted_y_grad.dtype(), true);
+    auto tmp2 = is_reduce_empty
+                    ? (promoted_y_grad * x_hat)
+                    : (promoted_y_grad * x_hat)
+                          .sum(reduce_axes, promoted_y_grad.dtype(), true);
+    auto result = (promoted_scale * std_inv) *
+                  (promoted_y_grad - tmp1 / hw - (x_hat * tmp2 / hw));
     set_output<T>(ConverToOrig<T>(result, x.dtype()), x_grad);
   }
   // scale_grad = x_hat * y_grad.sum(n, h, w)
diff --git a/paddle/phi/kernels/gpu/instance_norm_grad_kernel.cu b/paddle/phi/kernels/gpu/instance_norm_grad_kernel.cu
@@ -385,14 +385,6 @@ void InstanceNormGradKernel(const Context &dev_ctx,
   dims = {1, NxC, H, W, D};
   strides = {NxC * H * W * D, H * W * D, W * D, D, 1};
 
-  if ((H * W * D) == 1) {
-    phi::Copy(dev_ctx, d_y, dev_ctx.GetPlace(), false, d_x);
-    phi::funcs::SetConstant<GPUContext, BatchNormParamType<T>> functor;
-    functor(dev_ctx, d_scale, static_cast<BatchNormParamType<T>>(0));
-    functor(dev_ctx, d_bias, static_cast<BatchNormParamType<T>>(0));
-    return;
-  }
-
 #ifdef PADDLE_WITH_HIP
   miopenTensorDescriptor_t data_desc_;
   miopenTensorDescriptor_t in_param_desc_;
diff --git a/test/legacy_test/test_instance_norm_op.py b/test/legacy_test/test_instance_norm_op.py
@@ -25,8 +25,8 @@
 
 def _reference_instance_norm_naive(x, scale, bias, epsilon, mean, var):
     x_shape = x.shape
-    if len(x_shape) == 2:
-        x = np.reshape(x, (x.shape[0], x.shape[1], 1, 1))
+    if len(x_shape) < 4:
+        x = np.reshape(x, (x.shape[0], x.shape[1], -1, 1))
     n, c, h, w = x.shape
 
     mean_tile = np.reshape(mean, (n, c, 1, 1))
@@ -40,7 +40,7 @@ def _reference_instance_norm_naive(x, scale, bias, epsilon, mean, var):
     bias_tile = np.reshape(bias, (1, c, 1, 1))
     bias_tile = np.tile(bias_tile, (n, 1, h, w))
     y = scale_tile * x_norm + bias_tile
-    if len(x_shape) == 2:
+    if len(x_shape) < 4:
         y = np.reshape(y, x_shape)
     return y, mean, var
 
@@ -80,6 +80,8 @@ def _reference_instance_norm_grad(x, d_y, scale, mean, var, epsilon):
 
 
 def _cal_mean_variance(x, epsilon, mean_shape):
+    if len(x.shape) < 4:
+        x = np.reshape(x, (x.shape[0], x.shape[1], -1, 1))
     mean = np.reshape(np.mean(x, axis=(2, 3)), mean_shape)
     var = np.reshape(np.var(x, axis=(2, 3)), mean_shape)
     return mean, var
@@ -197,6 +199,64 @@ def init_test_case(self):
         )
 
 
+class TestInstanceNormCaseNCL(TestInstanceNormOp):
+    def init_test_case(self):
+        x_shape = [2, 100, 4]
+        n, c = x_shape[0], x_shape[1]
+        self.epsilon = 1e-05
+        dtype = np.float32
+        scale_shape = [c]
+        mean_shape = [n * c]
+        np.random.seed()
+        self.x_np = np.random.random_sample(x_shape).astype(dtype)
+        self.scale_np = np.ones(scale_shape).astype(dtype)
+        self.bias_np = np.zeros(scale_shape).astype(dtype)
+        self.mean_np, self.var_np = _cal_mean_variance(
+            self.x_np, self.epsilon, mean_shape
+        )
+        self.fw_comp_atol = 1e-5
+
+    def test_check_output(self):
+        self.check_output(check_pir=True, check_prim_pir=True)
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X', 'Scale', 'Bias'],
+            'Y',
+            check_pir=True,
+            check_prim_pir=True,
+        )
+
+
+class TestInstanceNormCaseNC(TestInstanceNormOp):
+    def init_test_case(self):
+        x_shape = [2, 100]
+        n, c = x_shape[0], x_shape[1]
+        self.epsilon = 1e-05
+        dtype = np.float32
+        scale_shape = [c]
+        mean_shape = [n * c]
+        np.random.seed()
+        self.x_np = np.random.random_sample(x_shape).astype(dtype)
+        self.scale_np = np.ones(scale_shape).astype(dtype)
+        self.bias_np = np.zeros(scale_shape).astype(dtype)
+        self.mean_np, self.var_np = _cal_mean_variance(
+            self.x_np, self.epsilon, mean_shape
+        )
+        self.fw_comp_atol = 2e-5
+
+    def test_check_output(self):
+        self.check_output(atol=2e-5, check_pir=True, check_prim_pir=True)
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X', 'Scale', 'Bias'],
+            'Y',
+            check_pir=True,
+            check_prim_pir=True,
+        )
+
+
 class TestElasticNormOp(unittest.TestCase):
     def init_test_case(self):
         self.epsilon = 1e-5
diff --git a/test/legacy_test/test_instance_norm_op_v2.py b/test/legacy_test/test_instance_norm_op_v2.py
@@ -37,13 +37,12 @@ def instance_norm_wrapper(
 def _reference_instance_norm(x, scale, bias, epsilon):
     prev_x_shape = x.shape
     if len(x.shape) < 4:
-        N, C = x.shape[0], x.shape[1]
-        x = np.reshape(x, (N, C, -1, 1))
+        x = np.reshape(x, (x.shape[0], x.shape[1], -1, 1))
 
     N, C, H, W = x.shape
     mean = np.mean(x, axis=(2, 3), keepdims=True)
     variance = np.var(x, axis=(2, 3), keepdims=True)
-    std = np.sqrt(variance) + epsilon
+    std = np.sqrt(variance + epsilon)
     x_norm = (x - mean) / std
     scale = scale.reshape([1, C, 1, 1])
     bias = bias.reshape([1, C, 1, 1])
@@ -319,6 +318,34 @@ def test_check_grad(self):
         )
 
 
+class TestInstanceNormWithNC(TestInstanceNormFP32OP):
+    def init_shape(self):
+        self.shape = [4, 100]
+
+    def set_err_thre(self):
+        super().set_err_thre()
+        self.fw_comp_atol = 3e-5
+
+    def test_check_output(self):
+        self.check_output(
+            atol=self.atol,
+            check_pir=True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
+        )
+
+    def test_check_grad(self):
+        self.check_grad(
+            ['X', 'Scale', 'Bias'],
+            'Y',
+            check_pir=True,
+            check_prim_pir=(
+                False if os.getenv("FLAGS_enable_pir_in_executor") else True
+            ),
+        )
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda()
     or not core.is_float16_supported(core.CUDAPlace(0)),
diff --git a/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py b/test/prim/pir_prim/test_prim_sub_graph_dynamic_shape.py
@@ -701,6 +701,19 @@ def setUp(self):
         self.tol = 5e-6
 
 
+class TestPrimInstancenormNC(TestPrimBase):
+    def setUp(self):
+        np.random.seed(2023)
+        self.shape_x = [2, 128]
+        self.dtype_x = "float32"
+        self.init_x_shape = [None, None]
+        self.x = np.random.random(self.shape_x).astype(self.dtype_x)
+        self.net = instance_norm_net
+        self.necessary_ops = "pd_op.instance_norm"
+        self.enable_cinn = False
+        self.tol = 2e-5
+
+
 class TestPrimGroupNorm1(TestPrimBase):
     def setUp(self):
         np.random.seed(2023)
