[PIR]Support set_dynamic_shape and refine DoValueCheck (#59722)

0x45f · web-flow · commit 04423a72aee9 · 2023-12-12T15:00:10.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -261,48 +261,38 @@ std::set<std::string> GetRegisterDataType(const std::string& op_name) {
   return data_type;
 }
 
-void DoValueCheck(const pir::Value& value,
-                  const std::string& input_name,
-                  const std::set<std::string>& expected_dtype,
-                  const std::string& op_name) {
+std::string GetValueDataType(const pir::Value& value) {
   if (value.type().isa<pir::DenseTensorType>()) {
-    std::string value_type = phi::DataTypeToString(dialect::TransToPhiDataType(
+    return phi::DataTypeToString(dialect::TransToPhiDataType(
         value.type().dyn_cast<pir::DenseTensorType>().dtype()));
-    if (expected_dtype.find(value_type) == expected_dtype.end()) {
-      std::ostringstream joined;
-      std::copy(expected_dtype.begin(),
-                expected_dtype.end(),
-                std::ostream_iterator<std::string>(joined, ","));
-      PADDLE_THROW(phi::errors::InvalidArgument(
-          "Check data type error for op: %s, input: %s, %s.dtype is %s, and "
-          "expected_dtype is %s",
-          op_name,
-          input_name,
-          input_name,
-          value_type,
-          joined.str()));
-    }
   } else if (value.type().isa<paddle::dialect::SelectedRowsType>()) {
-    std::string value_type = phi::DataTypeToString(dialect::TransToPhiDataType(
+    return phi::DataTypeToString(dialect::TransToPhiDataType(
         value.type().dyn_cast<paddle::dialect::SelectedRowsType>().dtype()));
-    if (expected_dtype.find(value_type) == expected_dtype.end()) {
-      std::ostringstream joined;
-      std::copy(expected_dtype.begin(),
-                expected_dtype.end(),
-                std::ostream_iterator<std::string>(joined, ","));
-      PADDLE_THROW(phi::errors::InvalidArgument(
-          "Check data type error for op: %s, input: %s, %s.dtype is %s, and "
-          "expected_dtype is %s",
-          op_name,
-          input_name,
-          input_name,
-          value_type,
-          joined.str()));
-    }
   } else {
+    PADDLE_THROW(
+        phi::errors::InvalidArgument("Currently, we can only get dtype for "
+                                     "DenseTensorType and SelectedRowsType."));
+  }
+}
+
+void DoValueCheck(const pir::Value& value,
+                  const std::string& input_name,
+                  const std::set<std::string>& expected_dtype,
+                  const std::string& op_name) {
+  std::string value_type = GetValueDataType(value);
+  if (expected_dtype.find(value_type) == expected_dtype.end()) {
+    std::ostringstream joined;
+    std::copy(expected_dtype.begin(),
+              expected_dtype.end(),
+              std::ostream_iterator<std::string>(joined, ", "));
     PADDLE_THROW(phi::errors::InvalidArgument(
-        "Currently, we can only get dtype for dense "
-        "tensor."));
+        "Check data type error for op: %s, input: %s, %s.dtype: %s, and "
+        "expected_dtype: %s",
+        op_name,
+        input_name,
+        input_name,
+        value_type,
+        joined.str()));
   }
 }
 
@@ -336,10 +326,10 @@ void CheckDataType(const phi::DataType& dtype,
     std::ostringstream joined;
     std::copy(expected_dtype.begin(),
               expected_dtype.end(),
-              std::ostream_iterator<std::string>(joined, ","));
+              std::ostream_iterator<std::string>(joined, ", "));
     PADDLE_THROW(phi::errors::InvalidArgument(
-        "Check data type error for op: %s, dtype is %s, and "
-        "expected_dtype is %s",
+        "Check data type error for op: %s, dtype: %s, and "
+        "expected_dtype: %s",
         op_name,
         str_dtype,
         joined.str()));
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -810,6 +810,34 @@ void BindType(py::module *m) {
         print_stream << self;
         return print_stream.str();
       });
+
+  m->def("create_shaped_type",
+         [](Type &type, const std::vector<int> &shape) -> Type {
+           if (type.isa<DenseTensorType>()) {
+             DenseTensorType src_type = type.dyn_cast<DenseTensorType>();
+             DenseTensorType dst_type =
+                 DenseTensorType::get(pir::IrContext::Instance(),
+                                      src_type.dtype(),
+                                      phi::make_ddim(shape),
+                                      src_type.data_layout(),
+                                      src_type.lod(),
+                                      src_type.offset());
+             return dst_type;
+           } else if (type.isa<SelectedRowsType>()) {
+             SelectedRowsType src_type = type.dyn_cast<SelectedRowsType>();
+             SelectedRowsType dst_type =
+                 SelectedRowsType::get(pir::IrContext::Instance(),
+                                       src_type.dtype(),
+                                       phi::make_ddim(shape),
+                                       src_type.data_layout(),
+                                       src_type.lod(),
+                                       src_type.offset());
+             return dst_type;
+           } else {
+             PADDLE_THROW(phi::errors::InvalidArgument(
+                 "Currently, we can only set shape for dense tensor"));
+           }
+         });
 }
 
 void BindAttribute(py::module *m) {
diff --git a/python/paddle/jit/dy2static/utils_helper.py b/python/paddle/jit/dy2static/utils_helper.py
@@ -187,10 +187,14 @@ def type_from_annotation(annotation):
 
 def set_dynamic_shape(variable, shape_list):
     if paddle.base.dygraph.base.in_to_static_mode():
-        assert isinstance(
-            variable, paddle.base.framework.Variable
-        ), "In to_static mode, variable must be a Variable."
-        variable.desc.set_shape(shape_list)
+        if isinstance(variable, paddle.base.framework.Variable):
+            variable.desc.set_shape(shape_list)
+        elif isinstance(variable, paddle.pir.Value):
+            variable.set_shape(shape_list)
+        else:
+            raise TypeError(
+                "In to_static mode, variable must be a Variable or Value"
+            )
     else:
         # in dygraph mode, dynamic shape is not needed, just do nothing.
         return
diff --git a/python/paddle/pir/__init__.py b/python/paddle/pir/__init__.py
@@ -32,6 +32,7 @@
     set_insertion_point,
     translate_to_pir,
     translate_to_pir_with_param_map,
+    create_shaped_type,
 )
 
 from . import core  # noqa: F401
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -480,6 +480,19 @@ def append(self, var):
 
         array_write(x=var, i=array_length(self), array=self)
 
+    def set_shape(self, shape):
+        assert (
+            paddle.base.dygraph.base.in_to_static_mode()
+        ), "We only support call 'set_shape' in to_static mode."
+
+        if self.is_dense_tensor_type() or self.is_selected_row_type():
+            type = paddle.pir.create_shaped_type(self.type(), shape)
+            self.set_type(type)
+        else:
+            raise ValueError(
+                "Currently, we can only set shape for dense and selected_row tensor"
+            )
+
     import paddle
 
     value_methods = [
@@ -495,6 +508,7 @@ def append(self, var):
         ('clone', clone),
         ('clear_gradient', clear_gradient),
         ('append', append),
+        ('set_shape', set_shape),
         (
             '__add__',
             _binary_creator_('__add__', paddle.tensor.add, False, _scalar_add_),
diff --git a/test/dygraph_to_static/test_set_dynamic_shape.py b/test/dygraph_to_static/test_set_dynamic_shape.py
@@ -14,7 +14,11 @@
 
 import unittest
 
-from dygraph_to_static_utils import Dy2StTestBase, test_ast_only
+from dygraph_to_static_utils import (
+    Dy2StTestBase,
+    test_ast_only,
+    test_legacy_and_pt_and_pir,
+)
 
 import paddle
 
@@ -37,6 +41,18 @@ def dygraph_func(loop_number):
         actual_shape = paddle.jit.to_static(dygraph_func)(loop_num).shape
         self.assertEqual(expected_shape, actual_shape)
 
+    @test_ast_only
+    @test_legacy_and_pt_and_pir
+    def test_pir_ast(self):
+        def dygraph_func():
+            mask = paddle.randn([2, 2])
+            paddle.jit.dy2static.utils_helper.set_dynamic_shape(mask, [-1, 2])
+            return mask.shape[0]
+
+        out = paddle.jit.to_static(dygraph_func)()
+        self.assertTrue(isinstance(out, paddle.Tensor))
+        self.assertEqual(out.numpy(), 2)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_pow.py b/test/legacy_test/test_pow.py
@@ -212,6 +212,18 @@ def test_errors(self):
         y = int(np.random.rand() * 10)
         self.assertRaises(TypeError, paddle.pow, x, str(y))
 
+    def test_pir_error(self):
+        with paddle.pir_utils.IrGuard():
+
+            def x_dtype_error():
+                with paddle.static.program_guard(
+                    paddle.static.Program(), paddle.static.Program()
+                ):
+                    x = paddle.static.data('x', [2, 2], dtype='int8')
+                    out = paddle.pow(x, 2)
+
+            self.assertRaises(ValueError, x_dtype_error)
+
 
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@`
`32`	`32`	`set_insertion_point,`
`33`	`33`	`translate_to_pir,`
`34`	`34`	`translate_to_pir_with_param_map,`
	`35`	`+ create_shaped_type,`
`35`	`36`	`)`
`36`	`37`
`37`	`38`	`from . import core # noqa: F401`