[Prim][PIR] Support group_norm rank greater than 5 (#68774)

cubehan3 · web-flow · commit e7c761513862 · 2024-10-21T19:44:15.000+08:00
* support group_norm rank greater than 5

* add test cases
diff --git a/paddle/fluid/primitive/composite/composite.h b/paddle/fluid/primitive/composite/composite.h
@@ -1117,9 +1117,11 @@ std::tuple<Tensor, Tensor, Tensor> group_norm_decomp(
         common::errors::Unimplemented("Only support NCHW and NHWC format."));
   }
   size_t rank = x.shape().size();
-  if (rank < 3 || rank > 5) {
+  if (rank < 3) {
     PADDLE_THROW(common::errors::Unimplemented(
-        "Only support NCHW and NHWC format in rank {3, 4, 5}."));
+        "Only support NCHW and NHWC format in rank higher or equal to 3. "
+        "Current rank: %zu",
+        rank));
   }
 
   auto org_dtype = x.dtype();
diff --git a/paddle/fluid/primitive/rule/vjp/details.h b/paddle/fluid/primitive/rule/vjp/details.h
@@ -2260,9 +2260,11 @@ void group_norm_grad(const Tensor& x,
   DataLayout data_layout_ = common::StringToDataLayout(data_layout);
   std::vector<int64_t> x_dims = x.shape();
   int rank = x_dims.size();
-  if (rank < 3 || rank > 5) {
+  if (rank < 3) {
     PADDLE_THROW(common::errors::Unimplemented(
-        "Only support NCHW and NHWC format in rank {3, 4, 5}."));
+        "Only support NCHW and NHWC format in rank higher or equal to 3. "
+        "Current rank: %zu",
+        rank));
   }
   int N = x_dims[0];
   int C;
diff --git a/test/ir/pir/cinn/test_cinn_group_norm.py b/test/ir/pir/cinn/test_cinn_group_norm.py
@@ -152,5 +152,85 @@ def test_prim(self):
         np.testing.assert_allclose(cinn_bias_grad, dy_bias_grad, rtol=1e-5)
 
 
+class TestGroupNormSubGraphRank6NHWC(TestGroupNormSubGraph):
+    def setUp(self):
+        paddle.seed(2024)
+        self.shape = [80, 4, 4, 4, 4, 128]
+        self.dtype = "float32"
+        self.data_format = "NHWC"
+        self.prepare_data()
+
+    def test_prim(self):
+        cinn_out, cinn_x_grad, cinn_weight_grad, cinn_bias_grad = self.eval(
+            use_cinn=False, use_prim=True
+        )
+        dy_out, dy_x_grad, dy_weight_grad, dy_bias_grad = self.eval(
+            use_cinn=False
+        )
+        np.testing.assert_allclose(cinn_x_grad, dy_x_grad, atol=1e-5, rtol=1e-5)
+        np.testing.assert_allclose(cinn_weight_grad, dy_weight_grad, atol=5e-4)
+        np.testing.assert_allclose(cinn_bias_grad, dy_bias_grad, rtol=1e-5)
+
+
+class TestGroupNormSubGraphRank6NCHW(TestGroupNormSubGraph):
+    def setUp(self):
+        paddle.seed(2024)
+        self.shape = [80, 128, 8, 4, 4, 4]
+        self.dtype = "float32"
+        self.data_format = "NCHW"
+        self.prepare_data()
+
+    def test_prim(self):
+        cinn_out, cinn_x_grad, cinn_weight_grad, cinn_bias_grad = self.eval(
+            use_cinn=False, use_prim=True
+        )
+        dy_out, dy_x_grad, dy_weight_grad, dy_bias_grad = self.eval(
+            use_cinn=False
+        )
+        np.testing.assert_allclose(cinn_x_grad, dy_x_grad, atol=1e-5, rtol=1e-5)
+        np.testing.assert_allclose(cinn_weight_grad, dy_weight_grad, atol=5e-4)
+        np.testing.assert_allclose(cinn_bias_grad, dy_bias_grad, rtol=1e-5)
+
+
+class TestGroupNormSubGraphRank7NHWC(TestGroupNormSubGraph):
+    def setUp(self):
+        paddle.seed(2024)
+        self.shape = [80, 4, 4, 4, 4, 4, 128]
+        self.dtype = "float32"
+        self.data_format = "NHWC"
+        self.prepare_data()
+
+    def test_prim(self):
+        cinn_out, cinn_x_grad, cinn_weight_grad, cinn_bias_grad = self.eval(
+            use_cinn=False, use_prim=True
+        )
+        dy_out, dy_x_grad, dy_weight_grad, dy_bias_grad = self.eval(
+            use_cinn=False
+        )
+        np.testing.assert_allclose(cinn_x_grad, dy_x_grad, atol=1e-5, rtol=1e-5)
+        np.testing.assert_allclose(cinn_weight_grad, dy_weight_grad, atol=5e-4)
+        np.testing.assert_allclose(cinn_bias_grad, dy_bias_grad, rtol=1e-5)
+
+
+class TestGroupNormSubGraphRank7NCHW(TestGroupNormSubGraph):
+    def setUp(self):
+        paddle.seed(2024)
+        self.shape = [80, 128, 8, 4, 4, 4, 4]
+        self.dtype = "float32"
+        self.data_format = "NCHW"
+        self.prepare_data()
+
+    def test_prim(self):
+        cinn_out, cinn_x_grad, cinn_weight_grad, cinn_bias_grad = self.eval(
+            use_cinn=False, use_prim=True
+        )
+        dy_out, dy_x_grad, dy_weight_grad, dy_bias_grad = self.eval(
+            use_cinn=False
+        )
+        np.testing.assert_allclose(cinn_x_grad, dy_x_grad, atol=1e-5, rtol=1e-5)
+        np.testing.assert_allclose(cinn_weight_grad, dy_weight_grad, atol=5e-4)
+        np.testing.assert_allclose(cinn_bias_grad, dy_bias_grad, rtol=1e-5)
+
+
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1117,9 +1117,11 @@ std::tuple<Tensor, Tensor, Tensor> group_norm_decomp(`
`1117`	`1117`	`common::errors::Unimplemented("Only support NCHW and NHWC format."));`
`1118`	`1118`	`}`
`1119`	`1119`	`size_t rank = x.shape().size();`
`1120`		`- if (rank < 3 \|\| rank > 5) {`
	`1120`	`+ if (rank < 3) {`
`1121`	`1121`	`PADDLE_THROW(common::errors::Unimplemented(`
`1122`		`- "Only support NCHW and NHWC format in rank {3, 4, 5}."));`
	`1122`	`+ "Only support NCHW and NHWC format in rank higher or equal to 3. "`
	`1123`	`+ "Current rank: %zu",`
	`1124`	`+ rank));`
`1123`	`1125`	`}`
`1124`	`1126`
`1125`	`1127`	`auto org_dtype = x.dtype();`