[Prim][PIR]group_norm supports bwd nhwc (#63171)

cyber-pioneer · web-flow · commit ed50854a2807 · 2024-04-03T17:21:34.000+08:00
* gn decomp support nhwc

* add decomp test case

* fix gn decomp rule

* fix atol

* polish code

* fix bug

* gix test case shape

* fix test case atol

* gn back supports nhwc

* gn support nhwc

* debug

* fix test case

* polish code

* polish rule

* fix acc
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -1513,11 +1513,28 @@ void group_norm_grad(const Tensor& x,
   // cal d_bias:
   // d_bias = sum(dy, axes=(0,2,3))
   DataLayout data_layout_ = common::StringToDataLayout(data_layout);
-  if (data_layout_ != DataLayout::kNCHW) {
-    // TODO(chengyanfu): Subsequent support NHWC
+  std::vector<int64_t> x_dims = x.shape();
+  int rank = x_dims.size();
+  int N = x_dims[0];
+  int C;
+  int hw;
+  std::vector<int64_t> reduce_axis;
+
+  if (data_layout_ == DataLayout::kNCHW) {
+    C = x_dims[1];
+    hw = x_dims[2] * x_dims[3];
+    reduce_axis = {2, 3};
+  } else if (data_layout_ == DataLayout::kNHWC) {
+    C = x_dims[rank - 1];
+    hw = x_dims[1] * x_dims[2];
+    reduce_axis = {1, 2};
+  } else {
     PADDLE_THROW(phi::errors::InvalidArgument("Unsupported storage order: %s",
                                               data_layout));
   }
+
+  int g_num = C / groups;
+
   Tensor x_data = x;
   Tensor out_grad_data = out_grad;
 
@@ -1531,37 +1548,38 @@ void group_norm_grad(const Tensor& x,
     out_grad_data = cast<T>(out_grad, phi::DataType::FLOAT32);
   }
 
-  std::vector<int64_t> x_dims = common::vectorize<int64_t>(x.dims());
-  auto add_axis = std::vector<int64_t>({-1});
-  const int N = x_dims[0];
-  const int C = x_dims[1];
+  auto shape_group = std::vector<int64_t>({N, groups, g_num});
 
-  const int hw = x_dims[2] * x_dims[3];
-  const int g_num = C / groups;
+  std::vector<int64_t> whole_group_shape;
+  if (data_layout_ == DataLayout::kNCHW) {
+    whole_group_shape = std::vector<int64_t>({N, groups, g_num, -1});
+  } else {
+    whole_group_shape = std::vector<int64_t>({N, -1, groups, g_num});
+  }
+  auto var_eps = variance + epsilon;
 
-  auto reduce_axis = IntArray(std::vector<int64_t>({2, 3}));
-  auto shape_group = IntArray(std::vector<int64_t>({N, groups, g_num}));
-  auto whole_group_shape =
-      IntArray(std::vector<int64_t>({N, groups, g_num, hw}));
+  auto inv_std = rsqrt<T>(var_eps);
 
-  auto scale_ptr = scale.get_ptr();
-  auto bias_ptr = bias.get_ptr();
-  auto sqrt_element = 1.0 / (variance + epsilon);
-  auto inv_std = elementwise_pow<T>(
-      sqrt_element,
-      full<T>(
-          common::vectorize(sqrt_element.dims()), 0.5, sqrt_element.dtype()));
   auto inv_std_mul_s = inv_std / hw / g_num;
   auto dtype = x_data.dtype();
   auto sum_y_grad_mul_x =
       sum<T>(out_grad_data * x_data, reduce_axis, dtype, false);
   auto sum_y_grad = sum<T>(out_grad_data, reduce_axis, dtype, false);
+
+  Tensor scale_data;
+  if (scale) {
+    scale_data = scale.get();
+  }
+  Tensor bias_data;
+  if (bias) {
+    bias_data = bias.get();
+  }
+
   if (x_grad) {
     Tensor d1;
     Tensor d2;
     Tensor p1;
-    if (scale_ptr) {
-      auto scale_data = scale.get();
+    if (scale) {
       if (scale_data.dtype() == phi::DataType::FLOAT16 ||
           scale_data.dtype() == phi::DataType::BFLOAT16) {
         scale_data = cast<T>(scale_data, phi::DataType::FLOAT32);
@@ -1573,18 +1591,23 @@ void group_norm_grad(const Tensor& x,
       p1 = reshape<T>(inv_std, std::vector<int64_t>({N, groups, 1})) *
            reshape<T>(scale_data, std::vector<int64_t>({1, groups, g_num}));
     } else {
-      d1 = (reshape<T>(sum_y_grad_mul_x, shape_group))
-               .sum(std::vector<int64_t>({2}), dtype, false);
-      d2 = (reshape<T>(sum_y_grad, shape_group))
-               .sum(std::vector<int64_t>({2}), dtype, false);
-      p1 = (reshape<T>(inv_std, std::vector<int64_t>({N, groups, 1})))
-               .expand(IntArray(shape_group));
+      d1 = (reshape<T>(sum_y_grad_mul_x, shape_group)).sum({2}, dtype, false);
+      d2 = (reshape<T>(sum_y_grad, shape_group)).sum({2}, dtype, false);
+      p1 = (reshape<T>(inv_std, {N, groups, 1})).expand(shape_group);
     }
 
     auto p2 = (d2 * mean - d1) * (inv_std_mul_s * inv_std * inv_std);
     auto p3 = -p2 * mean - d2 * inv_std_mul_s;
-    auto first_shape = get_unsqueeze_dims(p1, std::vector<int64_t>({3}));
-    auto second_shape = get_unsqueeze_dims(p2, std::vector<int64_t>({2, 3}));
+    std::vector<int64_t> first_shape;
+    std::vector<int64_t> second_shape;
+    if (data_layout_ == DataLayout::kNCHW) {
+      first_shape = get_unsqueeze_dims(p1, {3});
+      second_shape = get_unsqueeze_dims(p2, {2, 3});
+    } else {
+      first_shape = get_unsqueeze_dims(p1, {1});
+      second_shape = get_unsqueeze_dims(p2, {1, 3});
+    }
+
     p1 = reshape<T>(p1, first_shape);
     p2 = reshape<T>(p2, second_shape);
     p3 = reshape<T>(p3, second_shape);
@@ -1599,29 +1622,24 @@ void group_norm_grad(const Tensor& x,
 
     set_output<T>(x_grad_data, x_grad);
   }
+
   if (scale_grad) {
-    if (scale_ptr) {
-      auto third_shape = get_unsqueeze_dims(mean, std::vector<int64_t>({2}));
+    if (scale) {
+      auto third_shape = get_unsqueeze_dims(mean, {2});
       auto tmp1 = (reshape<T>(sum_y_grad_mul_x, shape_group) -
                    reshape<T>(sum_y_grad, shape_group) *
                        reshape<T>(mean, third_shape)) *
                   reshape<T>(inv_std, third_shape);
-      auto scale_grad_tmp = reshape<T>(
-          tmp1.sum(std::vector<int64_t>({0}), scale_ptr->dtype(), false),
-          IntArray(std::vector<int64_t>({C})));
+      auto scale_grad_tmp =
+          reshape<T>(tmp1.sum({0}, scale->dtype(), false), {C});
       set_output<T>(scale_grad_tmp, scale_grad);
-    } else {
-      scale_grad = nullptr;
     }
   }
 
   if (bias_grad) {
-    if (bias_ptr) {
-      auto bias_grad_tmp =
-          sum_y_grad.sum(std::vector<int64_t>({0}), bias_ptr->dtype(), false);
+    if (bias) {
+      auto bias_grad_tmp = sum_y_grad.sum({0}, bias->dtype(), false);
       set_output<T>(bias_grad_tmp, bias_grad);
-    } else {
-      bias_grad = nullptr;
     }
   }
 }
diff --git a/test/legacy_test/op_test.py b/test/legacy_test/op_test.py
@@ -1636,7 +1636,7 @@ def _compare_expect_and_actual_outputs(
                 np.testing.assert_allclose(
                     expect_out,
                     actual_out,
-                    rtol=1e-05,
+                    rtol=1e-03 if self.dtype == np.uint16 else 1e-5,
                     atol=inplace_atol,
                     err_msg='Operator ('
                     + self.op_type
diff --git a/test/legacy_test/test_group_norm_op.py b/test/legacy_test/test_group_norm_op.py
@@ -180,7 +180,7 @@ def test_check_grad(self):
             self.do_compare_between_place()
             return
 
-        check_prim_grad = True if self.data_format == "NCHW" else False
+        check_prim_grad = True
 
         self.rev_comp_atol = 1e-12
         self.rev_comp_rtol = 1e-12
@@ -234,7 +234,7 @@ def test_check_grad(self):
         if self.compare_between_place:
             return
 
-        check_prim_grad = True if self.data_format == "NCHW" else False
+        check_prim_grad = True
         self.rev_comp_atol = 1e-2
         self.rev_comp_rtol = 1e-2
         place = core.CUDAPlace(0)
@@ -265,7 +265,7 @@ def setUp(self):
         self.data_format = "NCHW"
         self.dtype = np.uint16
         self.shape = (2, 100, 3, 5)
-        self.attrs = {'epsilon': 1e-5, 'groups': 2, 'data_layout': "NCHW"}
+        self.attrs = {'epsilon': 1e-5, 'groups': 10, 'data_layout': "NCHW"}
         self.compare_between_place = False
         self.init_test_case()
 
@@ -313,12 +313,12 @@ def test_check_grad(self):
         if self.compare_between_place:
             return
 
-        check_prim_grad = True if self.data_format == "NCHW" else False
+        check_prim_grad = True
 
         self.rev_comp_atol = 1e-2
         self.rev_comp_rtol = 1e-2
         # prim bf16 has diff in windows
-        if sys.platform == "win32":
+        if sys.platform == "win32" or self.data_format == "NHWC":
             self.rev_comp_atol = 5e-2
             self.rev_comp_rtol = 5e-2
         place = core.CUDAPlace(0)
@@ -363,7 +363,7 @@ def init_test_case(self):
 
 class TestGroupNormBF16Op2(TestGroupNormBF16Op):
     def init_test_case(self):
-        self.attrs['groups'] = 4
+        self.attrs['groups'] = 10
 
 
 class TestGroupNormOpBigEps1(TestGroupNormOp):
@@ -398,7 +398,7 @@ def init_test_case(self):
 
 class TestGroupNormOp1_With_NHWC(TestGroupNormOp):
     def init_test_case(self):
-        self.attrs['groups'] = 1
+        self.attrs['groups'] = 2
         self.data_format = "NHWC"
 
 
@@ -411,7 +411,7 @@ def init_test_case(self):
 class TestGroupNormFP16Op_With_NHWC(TestGroupNormFP16OP):
     def init_test_case(self):
         self.no_need_check_inplace = True
-        self.attrs['groups'] = 1
+        self.attrs['groups'] = 10
         self.data_format = "NHWC"
         self.attrs['epsilon'] = 0.5
         self.shape = (1, 100, 4, 4)
@@ -435,13 +435,16 @@ class TestGroupNormBF16Op_With_NHWC(TestGroupNormBF16Op):
     def setUp(self):
         self.op_type = "group_norm"
         self.python_api = group_norm_wrapper
+        self.public_python_api = group_norm_wrapper
         self.python_out_sig = ["Y"]
         self.data_format = "NHWC"
+        self.prim_op_type = "comp"
+
         self.dtype = np.uint16
-        self.shape = (1, 3, 5, 100)
+        self.shape = (1, 3, 5, 512)
         self.attrs = {
             'epsilon': 5e-2,
-            'groups': 2,
+            'groups': 32,
             'data_layout': self.data_format,
         }
         self.compare_between_place = False
@@ -458,7 +461,7 @@ def setUp(self):
             .reshape(self.shape)
             .astype(np.float32)
         )
-        scale = np.sin(np.arange(self.shape[3])).astype(np.float32)
+        scale = np.ones(self.shape[3]).astype(np.float32)
         bias = np.sin(np.arange(self.shape[3])).astype(np.float32)
         output, mean, var = group_norm_naive(
             input,
@@ -477,9 +480,14 @@ def setUp(self):
         self.outputs = {'Y': output, 'Mean': mean, 'Variance': var}
 
     def test_check_output(self):
-        rtol = 2e-2
         place = core.CUDAPlace(0)
-        self.check_output_with_place(place, rtol=rtol, check_pir=True)
+        self.check_output_with_place(
+            place,
+            rtol=2e-2,
+            inplace_atol=1e-3,
+            check_pir=True,
+            check_prim_pir=True,
+        )
 
 
 class TestGroupNormOpBigEps1_With_NHWC(TestGroupNormOp):
