[PIR] Skip more checks at compile time (#67234)

SigureMo · web-flow · commit 00a40a10191e · 2024-08-10T12:07:14.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -71,6 +71,7 @@
     'AccuracyInferMeta',
     'BoxCoderInferMeta',
     'InstanceNormInferMeta',
+    'GroupNormInferMeta',
     'LayerNormInferMeta',
     'MatchMatrixTensorInferMeta',
     'MultiClassNMSInferMeta',
diff --git a/paddle/phi/infermeta/ternary.cc b/paddle/phi/infermeta/ternary.cc
@@ -937,7 +937,8 @@ void GroupNormInferMeta(const MetaTensor& x,
                         const std::string& data_layout_str,
                         MetaTensor* y,
                         MetaTensor* mean,
-                        MetaTensor* variance) {
+                        MetaTensor* variance,
+                        MetaConfig config) {
   PADDLE_ENFORCE_NE(y,
                     nullptr,
                     common::errors::InvalidArgument(
@@ -966,34 +967,37 @@ void GroupNormInferMeta(const MetaTensor& x,
   const int64_t channel_num =
       (data_layout == DataLayout::kNCHW ? x_dim[1] : x_dim[x_dim.size() - 1]);
   auto batch_size = x_dim[0];
-  PADDLE_ENFORCE_LE(
-      groups,
-      channel_num,
-      common::errors::InvalidArgument(
-          "The Attr(groups) of Op(group_norm) must be less than or "
-          "equal to the number of channels. But received: groups "
-          "is [%s], channels is [%s], the Attr(data_layout) "
-          "is [%s]. The error may come from wrong data_layout setting.",
-          groups,
-          channel_num,
-          data_layout_str));
-  PADDLE_ENFORCE_GE(
-      groups,
-      1,
-      common::errors::InvalidArgument(
-          "The Attr(groups) of Op(group_norm) must be "
-          "greater than or equal to 1. But received: groups is [%s].",
-          groups));
-  PADDLE_ENFORCE_EQ(
-      channel_num % groups,
-      0,
-      common::errors::InvalidArgument(
-          "Expected number of channels in input to be divisible by "
-          "num_groups, but got input channel is %d and num_groups is %d",
-          channel_num,
-          groups));
+  bool need_check = channel_num != -1 || config.is_runtime;
+  if (need_check) {
+    PADDLE_ENFORCE_LE(
+        groups,
+        channel_num,
+        common::errors::InvalidArgument(
+            "The Attr(groups) of Op(group_norm) must be less than or "
+            "equal to the number of channels. But received: groups "
+            "is [%s], channels is [%s], the Attr(data_layout) "
+            "is [%s]. The error may come from wrong data_layout setting.",
+            groups,
+            channel_num,
+            data_layout_str));
+    PADDLE_ENFORCE_GE(
+        groups,
+        1,
+        common::errors::InvalidArgument(
+            "The Attr(groups) of Op(group_norm) must be "
+            "greater than or equal to 1. But received: groups is [%s].",
+            groups));
+    PADDLE_ENFORCE_EQ(
+        channel_num % groups,
+        0,
+        common::errors::InvalidArgument(
+            "Expected number of channels in input to be divisible by "
+            "num_groups, but got input channel is %d and num_groups is %d",
+            channel_num,
+            groups));
+  }
 
-  if (scale) {
+  if (scale && need_check) {
     PADDLE_ENFORCE_EQ(
         scale.dims().size(),
         1UL,
@@ -1015,7 +1019,7 @@ void GroupNormInferMeta(const MetaTensor& x,
             channel_num,
             data_layout_str));
   }
-  if (bias) {
+  if (bias && need_check) {
     PADDLE_ENFORCE_EQ(
         bias.dims().size(),
         1UL,
diff --git a/paddle/phi/infermeta/ternary.h b/paddle/phi/infermeta/ternary.h
@@ -187,7 +187,8 @@ void GroupNormInferMeta(const MetaTensor& x,
                         const std::string& data_layout,
                         MetaTensor* y,
                         MetaTensor* mean,
-                        MetaTensor* variance);
+                        MetaTensor* variance,
+                        MetaConfig config = MetaConfig());
 
 void LayerNormInferMeta(const MetaTensor& x,
                         const MetaTensor& scale,
diff --git a/python/paddle/nn/functional/conv.py b/python/paddle/nn/functional/conv.py
@@ -423,24 +423,22 @@ def conv1d(
         )
     num_channels = x.shape[channel_dim]
     num_filters = weight.shape[0]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) "
-            f"should be defined. Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv1d should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv1d should not be zero")
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
-            "the channel of input must be divisible by groups,"
+            "The channel of input must be divisible by groups,"
             f"received: the channel of input is {num_channels}, the shape of input is {x.shape}"
             f", the groups is {groups}"
         )
-    if num_filters % groups != 0:
+    if num_filters % groups != 0 and (
+        in_dynamic_mode() or (num_filters != -1 and groups != -1)
+    ):
         raise ValueError(
-            "the number of filters must be divisible by groups,"
+            "The number of filters must be divisible by groups,"
             f"received: the number of filters is {num_filters}, the shape of weight is {weight.shape}"
             f", the groups is {groups}"
         )
@@ -666,24 +664,23 @@ def conv2d(
         )
     num_channels = x.shape[channel_dim]
     num_filters = weight.shape[0]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) "
-            f"should be defined. Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv2d should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv2d should be not be zero.")
+
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
-            "the channel of input must be divisible by groups,"
+            "The channel of input must be divisible by groups,"
             f"received: the channel of input is {num_channels}, the shape of input is {x.shape}"
             f", the groups is {groups}"
         )
-    if num_filters % groups != 0:
+    if num_filters % groups != 0 and (
+        in_dynamic_mode() or (num_filters != -1 and groups != -1)
+    ):
         raise ValueError(
-            "the number of filters must be divisible by groups,"
+            "The number of filters must be divisible by groups,"
             f"received: the number of filters is {num_filters}, the shape of weight is {weight.shape}"
             f", the groups is {groups}"
         )
@@ -922,16 +919,12 @@ def conv1d_transpose(
         )
 
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) "
-            f"should be defined. Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv1d_transpose should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv1d_transpose should not be zero.")
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
             "the channel of input must be divisible by groups,"
             f"received: the channel of input is {num_channels}, the shape of input is {x.shape}"
@@ -1200,16 +1193,12 @@ def conv2d_transpose(
             f"Input weight should be 4D tensor, but received weight with the shape of {weight.shape}"
         )
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) "
-            f"should be defined. Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv2d_transpose should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv2d_transpose should not be zero.")
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
             "the channel of input must be divisible by groups,"
             f"received: the channel of input is {num_channels}, the shape of input is {x.shape}"
@@ -1483,21 +1472,19 @@ def conv3d(
         )
     num_channels = x.shape[channel_dim]
     num_filters = weight.shape[0]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv3d should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv3d should not be 0.")
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
             "The number of input channels must be divisible by Attr(groups). "
             f"Received: number of channels({num_channels}), groups({groups})."
         )
-    if num_filters % groups != 0:
+    if num_filters % groups != 0 and (
+        in_dynamic_mode() or (num_filters != -1 and groups != -1)
+    ):
         raise ValueError(
             "The number of filters must be divisible by Attr(groups). "
             f"Received: number of filters({num_filters}), groups({groups})."
@@ -1691,17 +1678,12 @@ def conv3d_transpose(
             f"Input weight should be 5D tensor, but received weight with the shape of {weight.shape}"
         )
     num_channels = x.shape[channel_dim]
-    num_filters = weight.shape[1]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
-    if groups <= 0:
-        raise ValueError(
-            f"The groups of conv3d_transpose should be greater than 0. Received groups: {groups}"
-        )
-    if num_channels % groups != 0:
+
+    if groups == 0:
+        raise ValueError("The groups of conv3d_transpose should not be zero.")
+    if num_channels % groups != 0 and (
+        in_dynamic_mode() or (num_channels != -1 and groups != -1)
+    ):
         raise ValueError(
             "The number of input channels must be divisible by Attr(groups). "
             f"Received: number of channels({num_channels}), groups({groups})."
diff --git a/python/paddle/sparse/nn/functional/conv.py b/python/paddle/sparse/nn/functional/conv.py
@@ -72,11 +72,6 @@ def _conv3d(
     channel_last = data_format == "NDHWC"
     channel_dim = -1 if channel_last else 1
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
 
     padding, padding_algorithm = _update_padding_nd(padding, channel_last, dims)
     stride = convert_to_list(stride, dims, 'stride')
@@ -161,11 +156,6 @@ def _conv2d(
             f"Input x should be 4D tensor, but received x with the shape of {x.shape}"
         )
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
 
     padding, padding_algorithm = _update_padding_nd(padding, channel_last, dims)
     stride = convert_to_list(stride, dims, 'stride')
@@ -247,11 +237,6 @@ def _conv3d_igemm(
             f"Input x should be 5D tensor, but received x with the shape of {x.shape}"
         )
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
 
     padding, padding_algorithm = _update_padding_nd(padding, channel_last, dims)
     stride = convert_to_list(stride, dims, 'stride')
@@ -331,11 +316,6 @@ def _conv2d_igemm(
             f"Input x should be 4D tensor, but received x with the shape of {x.shape}"
         )
     num_channels = x.shape[channel_dim]
-    if num_channels < 0:
-        raise ValueError(
-            f"The channel dimension of the input({x.shape}) should be defined. "
-            f"Received: {num_channels}."
-        )
 
     padding, padding_algorithm = _update_padding_nd(padding, channel_last, dims)
     stride = convert_to_list(stride, dims, 'stride')
diff --git a/test/dygraph_to_static/test_dynamic_shape_infermeta.py b/test/dygraph_to_static/test_dynamic_shape_infermeta.py
@@ -72,6 +72,27 @@ def test_depthwise_conv2d(self):
             [InputSpec(shape=[None, None, None, None], dtype='float32')],
         )
 
+    @test_pir_only
+    @test_ast_only
+    def test_group_norm(self):
+        self.check_dynamic_shape(
+            paddle.nn.GroupNorm(3, 3),
+            [paddle.randn([1, 3, 32, 32])],
+            [InputSpec(shape=[None, None, None, None], dtype='float32')],
+        )
+
+    @test_pir_only
+    @test_ast_only
+    def test_functional_conv(self):
+        self.check_dynamic_shape(
+            paddle.nn.functional.conv2d,
+            [paddle.randn([1, 3, 32, 32]), paddle.randn([3, 3, 3, 3])],
+            [
+                InputSpec(shape=[None, None, None, None], dtype='float32'),
+                InputSpec(shape=[None, None, None, None], dtype='float32'),
+            ],
+        )
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_functional_conv3d_transpose.py b/test/legacy_test/test_functional_conv3d_transpose.py
@@ -266,4 +266,5 @@ def setUp(self):
 
 
 if __name__ == "__main__":
+    paddle.enable_static()
     unittest.main()
diff --git a/test/sot/test_trace_list_arg.py b/test/sot/test_trace_list_arg.py
@@ -22,6 +22,7 @@
 )
 
 import paddle
+from paddle.jit.sot.utils.envs import with_allow_dynamic_shape_guard
 
 
 def foo(x: list[paddle.Tensor], y: list[paddle.Tensor]):
@@ -46,7 +47,8 @@ def test_foo(self):
             self.assert_results(foo, [a], [c])  # Cache miss
             self.assertEqual(cache.translate_count, 2)
 
-    def test_bar(self):
+    @with_allow_dynamic_shape_guard(False)
+    def test_bar_static_shape(self):
         a = [paddle.to_tensor(1), paddle.to_tensor(2), paddle.to_tensor(3)]
         b = [paddle.to_tensor([2, 3]), paddle.to_tensor(4), paddle.to_tensor(5)]
 
@@ -58,6 +60,19 @@ def test_bar(self):
             self.assert_results(bar, b, 1, 1)  # Cache hit
             self.assertEqual(cache.translate_count, 2)
 
+    @with_allow_dynamic_shape_guard(True)
+    def test_bar_dynamic_shape(self):
+        # TODO: Fix this after implement symbolic fallback mechanism
+        a = [paddle.to_tensor(1), paddle.to_tensor(2), paddle.to_tensor(3)]
+        b = [paddle.to_tensor([2, 3]), paddle.to_tensor(4), paddle.to_tensor(5)]
+        with test_instruction_translator_cache_context() as cache:
+            self.assert_results(bar, a, 1, 1)
+            self.assertEqual(cache.translate_count, 1)
+            self.assert_results(bar, a, 2, 0)  # Cache hit, but break graph
+            self.assertEqual(cache.translate_count, 3)
+            self.assert_results(bar, b, 1, 1)  # Cache hit
+            self.assertEqual(cache.translate_count, 3)
+
 
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -266,4 +266,5 @@ def setUp(self):`
`266`	`266`
`267`	`267`
`268`	`268`	`if __name__ == "__main__":`
	`269`	`+ paddle.enable_static()`
`269`	`270`	`unittest.main()`