[PIR] Add enforce for ArrayWrite & add dims for DenseTensorArrayType & open test (#61450)

chen2016013 · web-flow · commit e06048dd6d58 · 2024-02-07T12:25:49.000+08:00
* Add enforce for ArrayWrite &amp; open test

* fix concat

* fix test

* add dims for densetensorarraytype

* update 'is_dense_tensor_array_type'

* update body_block.arg's type to outer value
diff --git a/paddle/fluid/ir_adaptor/translator/type_translator.cc b/paddle/fluid/ir_adaptor/translator/type_translator.cc
@@ -102,10 +102,12 @@ TypeTranslator::TypeTranslator() {
                   << "[" << var_desc.Name() << "] from LOD_TENSOR_ARRAY";
          pir::Type dtype =
              this->operator[](var_desc.GetDataType())(ctx, var_desc);
+         phi::DDim dims = common::make_ddim(var_desc.GetShape());
          DenseTensorTypeStorage::DataLayout layout =
              DenseTensorTypeStorage::DataLayout::UNDEFINED;
 
-         return paddle::dialect::DenseTensorArrayType::get(ctx, dtype, layout);
+         return paddle::dialect::DenseTensorArrayType::get(
+             ctx, dtype, dims, layout);
        }},
       {VarType::SELECTED_ROWS,
        [&](pir::IrContext* ctx, const VarDesc& var_desc) -> pir::Type {
diff --git a/paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc b/paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc
@@ -73,6 +73,10 @@ const pir::Type& AllocatedDenseTensorArrayType::dtype() const {
   return storage()->dense_tensor_array_type_.dtype();
 }
 
+const pir::DDim& AllocatedDenseTensorArrayType::dims() const {
+  return storage()->dense_tensor_array_type_.dims();
+}
+
 const phi::DataLayout& AllocatedDenseTensorArrayType::data_layout() const {
   return storage()->dense_tensor_array_type_.data_layout();
 }
diff --git a/paddle/fluid/pir/dialect/kernel/ir/kernel_type.h b/paddle/fluid/pir/dialect/kernel/ir/kernel_type.h
@@ -120,9 +120,10 @@ class AllocatedDenseTensorArrayType
   static AllocatedDenseTensorArrayType get(pir::IrContext *ctx,
                                            const phi::Place &place,
                                            const pir::Type &dtype,
+                                           const phi::DDim &dims,
                                            const phi::DataLayout &layout) {
     dialect::DenseTensorArrayType type =
-        dialect::DenseTensorArrayType::get(ctx, dtype, layout);
+        dialect::DenseTensorArrayType::get(ctx, dtype, dims, layout);
 
     return pir::TypeManager::template get<AllocatedDenseTensorArrayType>(
         ctx, place, type);
@@ -132,6 +133,8 @@ class AllocatedDenseTensorArrayType
 
   const pir::Type &dtype() const;
 
+  const pir::DDim &dims() const;
+
   const phi::DataLayout &data_layout() const;
 };
 
diff --git a/paddle/fluid/pir/dialect/kernel/ir/type_storage.h b/paddle/fluid/pir/dialect/kernel/ir/type_storage.h
@@ -191,6 +191,7 @@ struct AllocatedDenseTensorArrayTypeStorage : public pir::TypeStorage {
                           dialect::DenseTensorArrayTypeStorage::HashValue(
                               dialect::DenseTensorArrayTypeStorage::ParamKey(
                                   dense_tensor_array_type.dtype(),
+                                  dense_tensor_array_type.dims(),
                                   dense_tensor_array_type.data_layout())));
     return hash_value;
   }
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.cc b/paddle/fluid/pir/dialect/operator/ir/op_type.cc
@@ -31,6 +31,7 @@ const size_t& SelectedRowsType::offset() const { return storage()->offset_; }
 const pir::Type& DenseTensorArrayType::dtype() const {
   return storage()->dtype_;
 }
+const phi::DDim& DenseTensorArrayType::dims() const { return storage()->dims_; }
 
 const phi::DataLayout& DenseTensorArrayType::data_layout() const {
   return storage()->layout_;
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_type.h b/paddle/fluid/pir/dialect/operator/ir/op_type.h
@@ -53,6 +53,8 @@ class DenseTensorArrayType
 
   const pir::Type &dtype() const;
 
+  const phi::DDim &dims() const;
+
   const phi::DataLayout &data_layout() const;
 };
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/type_storage.h b/paddle/fluid/pir/dialect/operator/ir/type_storage.h
@@ -109,21 +109,24 @@ struct SelectedRowsTypeStorage : public pir::TypeStorage {
 
 struct DenseTensorArrayTypeStorage : public pir::TypeStorage {
   using DataLayout = phi::DataLayout;
+  using DDim = phi::DDim;
   ///
   /// \brief Declare ParamKey according to parameter type.
   ///
-  using ParamKey = std::tuple<pir::Type, phi::DataLayout>;
+  using ParamKey = std::tuple<pir::Type, phi::DDim, phi::DataLayout>;
 
   DenseTensorArrayTypeStorage(const pir::Type& dtype,
+                              const phi::DDim& dims,
                               const phi::DataLayout& layout)
-      : dtype_(dtype), layout_(layout) {}
+      : dtype_(dtype), dims_(dims), layout_(layout) {}
 
   ///
   /// \brief Each derived TypeStorage must define a Construct method, which
   /// StorageManager uses to construct a derived TypeStorage.
   ///
   static DenseTensorArrayTypeStorage* Construct(const ParamKey& key) {
-    return new DenseTensorArrayTypeStorage(std::get<0>(key), std::get<1>(key));
+    return new DenseTensorArrayTypeStorage(
+        std::get<0>(key), std::get<1>(key), std::get<2>(key));
   }
 
   ///
@@ -134,28 +137,32 @@ struct DenseTensorArrayTypeStorage : public pir::TypeStorage {
     // hash dtype
     hash_value =
         pir::hash_combine(hash_value, std::hash<pir::Type>()(std::get<0>(key)));
+    // hash dims
+    hash_value =
+        pir::hash_combine(hash_value, std::hash<phi::DDim>()(std::get<1>(key)));
     // hash layout
     hash_value = pir::hash_combine(
         hash_value,
         std::hash<std::underlying_type<phi::DataLayout>::type>()(
             static_cast<std::underlying_type<phi::DataLayout>::type>(
-                std::get<1>(key))));
+                std::get<2>(key))));
     return hash_value;
   }
 
   ///
   /// \brief Each derived TypeStorage needs to overload operator==.
   ///
   bool operator==(const ParamKey& key) const {
-    return ParamKey(dtype_, layout_) == key;
+    return ParamKey(dtype_, dims_, layout_) == key;
   }
 
-  ParamKey GetAsKey() const { return ParamKey(dtype_, layout_); }
+  ParamKey GetAsKey() const { return ParamKey(dtype_, dims_, layout_); }
 
   ///
   /// \brief DenseTensorTypeStorage include five parameters: dtype, layout
   ///
   pir::Type dtype_;
+  phi::DDim dims_;
   phi::DataLayout layout_;
 };
 
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -524,8 +524,11 @@ static pir::Type BuildOutputType(pir::Type type,
         type, place, out_dtype, ctx);
   } else if (type.isa<DenseTensorArrayType>()) {
     auto array_type = type.dyn_cast<DenseTensorArrayType>();
-    return AllocatedDenseTensorArrayType::get(
-        ctx, place, array_type.dtype(), array_type.data_layout());
+    return AllocatedDenseTensorArrayType::get(ctx,
+                                              place,
+                                              array_type.dtype(),
+                                              array_type.dims(),
+                                              array_type.data_layout());
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
         "BuildOutputType only support DenseTensorType and SelectedRowsType"));
@@ -564,7 +567,7 @@ static pir::Type BuildOutputType(pir::Type type,
   } else if (type.isa<DenseTensorArrayType>()) {
     auto array_type = type.dyn_cast<DenseTensorArrayType>();
     return AllocatedDenseTensorArrayType::get(
-        ctx, place, array_type.dtype(), layout);
+        ctx, place, array_type.dtype(), array_type.dims(), layout);
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
         "BuildOutputType only support DenseTensorType and SelectedRowsType"));
@@ -2151,6 +2154,7 @@ std::vector<pir::Value> BuildInputs(
               ctx,
               out_place,
               new_in_alloc_type.dtype(),
+              new_in_alloc_type.dims(),
               new_in_alloc_type.data_layout());
           new_in = AddPlaceTransferOp(
               new_in, out_type, in_place, out_place, kernel_key, block);
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
@@ -241,6 +241,7 @@ std::vector<Value> PyWhileOp::OptimizeUpdate() {
   for (size_t operand_index = 1u, arg_index = 0u; operand_index < operand_num;
        ++operand_index) {
     if (yield_op.operand_source(operand_index) == body_block.arg(arg_index)) {
+      operand_source(operand_index).set_type(body_block.arg(arg_index).type());
       body_block.arg(arg_index).ReplaceAllUsesWith(
           operand_source(operand_index));
       body_block.EraseArg(arg_index);
diff --git a/paddle/phi/core/meta_tensor.cc b/paddle/phi/core/meta_tensor.cc
@@ -47,6 +47,17 @@ DDim MetaTensor::dims() const {
   }
 }
 
+size_t MetaTensor::size() const {
+  ValidCheck(*this);
+  PADDLE_ENFORCE_EQ(
+      phi::TensorArray::classof(tensor_),
+      true,
+      phi::errors::InvalidArgument(
+          "The current MetaTensor is not initialized by TensorArray."));
+  phi::TensorArray* tensor_array = static_cast<phi::TensorArray*>(tensor_);
+  return tensor_array->size();
+}
+
 DDim MetaTensor::dims(int64_t index) const {
   ValidCheck(*this);
   PADDLE_ENFORCE_EQ(
diff --git a/paddle/phi/core/meta_tensor.h b/paddle/phi/core/meta_tensor.h
@@ -64,6 +64,7 @@ class TEST_API MetaTensor {
 
   virtual int64_t numel() const;
   virtual DDim dims() const;
+  size_t size() const;  // Returns the number of tensors in TensorArray.
   DDim dims(int64_t index) const;
   virtual DataType dtype() const;
   virtual DataLayout layout() const;
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -144,6 +144,16 @@ void ArrayWriteInferMeta(const MetaTensor& array,
                          const MetaTensor& x,
                          MetaTensor* out,
                          MetaConfig config) {
+  if (array.dtype() != phi::DataType::UNDEFINED &&
+      x.dtype() != phi::DataType::UNDEFINED) {
+    PADDLE_ENFORCE_EQ(array.dtype(),
+                      x.dtype(),
+                      phi::errors::InvalidArgument(
+                          "The dtype (%s) of input x shall be same as "
+                          "dtype (%d) of array.",
+                          x.dtype(),
+                          array.dtype()));
+  }
   out->set_dtype(array.dtype());
   out->set_layout(array.layout());
 }
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -154,7 +154,39 @@ void ArrayToTensorInferMeta(const MetaTensor& x,
                             MetaTensor* out,
                             MetaTensor* out_index,
                             MetaConfig config) {
-  if (config.is_runtime) return;
+  if (config.is_runtime) {
+    PADDLE_ENFORCE_EQ(
+        (x.is_tensor_array()),
+        true,
+        phi::errors::InvalidArgument(
+            "The dtype of 1st input in ArrayToTensor must be TensorArray, but "
+            "received [%s]",
+            x.dtype()));
+    size_t x_size = x.size();
+    std::vector<phi::DDim> vec_dims;
+    for (size_t i = 0; i < x_size; ++i) {
+      phi::DDim dims = x.dims(i);
+      if (!vec_dims.empty()) {
+        PADDLE_ENFORCE_EQ(
+            (vec_dims.back() == dims),
+            true,
+            phi::errors::InvalidArgument(
+                "The dims of input in ArrayToTensor must be equal, but "
+                "received"));
+      }
+      vec_dims.push_back(dims);
+    }
+    if (vec_dims.empty()) return;
+    auto dim_vec = common::vectorize<int>(vec_dims[0]);
+    if (use_stack) {
+      dim_vec.insert(dim_vec.begin() + axis, dim_vec.size() * dim_vec[axis]);
+    } else {
+      dim_vec[axis] = dim_vec.size() * dim_vec[axis];
+    }
+    phi::DDim dims = common::make_ddim(dim_vec);
+    out->set_dims(dims);
+    return;
+  }
   auto dims = x.dims();
   // if the shape is empty
   if (dims == common::make_ddim({0UL})) return;
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -1292,12 +1292,23 @@ def concat(x, axis=0, name=None):
              [14, 15, 16]])
     """
     input = x
-    if in_dynamic_or_pir_mode():
+    if in_dynamic_mode():
         if isinstance(axis, Variable):
             axis = axis.item(0)
         if not isinstance(input, (Variable, paddle.pir.Value)):
             input = [t for t in input if t.shape.count(0) == 0]
         return _C_ops.concat(input, axis)
+    elif in_pir_mode():
+        if isinstance(input, paddle.pir.Value):
+            assert input.is_dense_tensor_array_type(), (
+                "If the element of concat op is Value, "
+                "dtype of the element must be Tensorarray"
+            )
+            out, _ = _C_ops.array_to_tensor(input, axis, False)
+            return out
+        if not isinstance(input, paddle.pir.Value):
+            input = [t for t in input if t.shape.count(0) == 0]
+        return _C_ops.concat(input, axis)
     else:
         check_type(input, 'input', (list, tuple, Variable), 'concat')
         if not isinstance(input, Variable):
diff --git a/test/dygraph_to_static/test_list.py b/test/dygraph_to_static/test_list.py
@@ -48,7 +48,7 @@ def test_list_append_in_if(x):
     if x.numpy()[0] > 0:
         a.append(x)
     else:
-        a.append(paddle.full(shape=[1, 2], fill_value=9, dtype="int64"))
+        a.append(paddle.full(shape=[1, 2], fill_value=9, dtype="float32"))
     # TODO(Aurelius84): Currently, run_program_op doesn't support output LoDTensorArray.
     return a[0]
 
@@ -128,7 +128,7 @@ def test_list_append_in_while_loop_with_stack(x, iter_num):
 def test_tensor_array_slice(x, iter_num):
     a = []
     for i in range(paddle.to_tensor(3)):
-        a.append(paddle.to_tensor(i))
+        a.append(paddle.to_tensor(float(i)))
     t = a[1:3]
     return a[2]
 
@@ -218,7 +218,7 @@ def setUp(self):
         self.init_dygraph_func()
 
     def init_data(self):
-        self.input = np.random.random(3).astype('int32')
+        self.input = np.random.random(3).astype('float32')
 
     def init_dygraph_func(self):
         self.all_dygraph_funcs = [
@@ -277,7 +277,7 @@ def init_dygraph_func(self):
 
 class TestListInWhileLoop(TestListWithoutControlFlowConfig):
     def init_data(self):
-        self.input = np.random.random(3).astype('int32')
+        self.input = np.random.random(3).astype('float32')
         self.iter_num = 3
 
     def init_dygraph_func(self):
@@ -298,6 +298,7 @@ def train(self, to_static=False):
             return self.result_to_numpy(res)
 
     @disable_test_case((ToStaticMode.AST, IrMode.PT))
+    @test_legacy_and_pt_and_pir
     def test_transformed_static_result(self):
         self.compare_transformed_static_result()
 
@@ -320,7 +321,7 @@ def init_dygraph_func(self):
         ]
 
 
-class TestListInForLoopWithConcat(TestListInWhileLoopWithStack):
+class TestListInForLoopWithConcat(TestListInWhileLoop):
     def init_dygraph_func(self):
         self.all_dygraph_funcs = [
             test_list_append_in_for_loop_with_concat,
@@ -337,9 +338,6 @@ def init_dygraph_func(self):
     def init_data(self):
         self.input = np.random.random((3, 4)).astype('float32')
 
-    def test_transformed_static_result(self):
-        self.compare_transformed_static_result()
-
 
 class ListWithCondNet(paddle.nn.Layer):
     def __init__(self):
@@ -366,6 +364,7 @@ def forward(self, x, index, *args):
 
 
 class TestListWithCondGradInferVarType(Dy2StTestBase):
+    @test_legacy_and_pt_and_pir
     def test_to_static(self):
         net = ListWithCondNet()
         x = paddle.to_tensor([2, 3, 4], dtype='float32')
diff --git a/test/dygraph_to_static/test_slice.py b/test/dygraph_to_static/test_slice.py
@@ -47,7 +47,7 @@ def test_slice_in_if(x):
     if x.numpy()[0] > 0:
         a.append(x)
     else:
-        a.append(paddle.full(shape=[1, 2], fill_value=9, dtype="int32"))
+        a.append(paddle.full(shape=[1, 2], fill_value=9, dtype="float32"))
 
     if x.numpy()[0] > 0:
         a[0] = x
@@ -116,7 +116,7 @@ def setUp(self):
         self.dygraph_func = None
 
     def init_input(self):
-        self.input = np.random.random(3).astype('int32')
+        self.input = np.random.random(3).astype('float32')
 
     def init_dygraph_func(self):
         raise NotImplementedError(
diff --git a/test/legacy_test/test_slice_op.py b/test/legacy_test/test_slice_op.py
@@ -1186,8 +1186,8 @@ class TestSliceTensorArray(unittest.TestCase):
     def test_slice_range(self):
         with paddle.pir_utils.IrGuard():
             arr = paddle.tensor.create_array("int32")
-            x = paddle.static.data("x", shape=[2, 2], dtype="float32")
-            y = paddle.static.data("y", shape=[1, 2], dtype="float32")
+            x = paddle.static.data("x", shape=[2, 2], dtype="int32")
+            y = paddle.static.data("y", shape=[1, 2], dtype="int32")
 
             zero = paddle.tensor.creation.fill_constant([], 'int64', 0)
             paddle.tensor.array_write(x, zero, array=arr)
@@ -1200,8 +1200,8 @@ def test_slice_range(self):
     def test_slice_item(self):
         with paddle.pir_utils.IrGuard():
             arr = paddle.tensor.create_array("int32")
-            x = paddle.static.data("x", shape=[2, 2], dtype="float32")
-            y = paddle.static.data("y", shape=[1, 2], dtype="float32")
+            x = paddle.static.data("x", shape=[2, 2], dtype="int32")
+            y = paddle.static.data("y", shape=[1, 2], dtype="int32")
 
             zero = paddle.tensor.creation.fill_constant([], 'int64', 0)
             paddle.tensor.array_write(x, zero, array=arr)

Original file line number	Diff line number	Diff line change
`@@ -73,6 +73,10 @@ const pir::Type& AllocatedDenseTensorArrayType::dtype() const {`
`73`	`73`	`return storage()->dense_tensor_array_type_.dtype();`
`74`	`74`	`}`
`75`	`75`
	`76`	`+const pir::DDim& AllocatedDenseTensorArrayType::dims() const {`
	`77`	`+ return storage()->dense_tensor_array_type_.dims();`
	`78`	`+}`
	`79`	`+`
`76`	`80`	`const phi::DataLayout& AllocatedDenseTensorArrayType::data_layout() const {`
`77`	`81`	`return storage()->dense_tensor_array_type_.data_layout();`
`78`	`82`	`}`
Original file line number	Diff line number	Diff line change
`@@ -191,6 +191,7 @@ struct AllocatedDenseTensorArrayTypeStorage : public pir::TypeStorage {`
`191`	`191`	`dialect::DenseTensorArrayTypeStorage::HashValue(`
`192`	`192`	`dialect::DenseTensorArrayTypeStorage::ParamKey(`
`193`	`193`	`dense_tensor_array_type.dtype(),`
	`194`	`+ dense_tensor_array_type.dims(),`
`194`	`195`	`dense_tensor_array_type.data_layout())));`
`195`	`196`	`return hash_value;`
`196`	`197`	`}`
Original file line number	Diff line number	Diff line change
`@@ -31,6 +31,7 @@ const size_t& SelectedRowsType::offset() const { return storage()->offset_; }`
`31`	`31`	`const pir::Type& DenseTensorArrayType::dtype() const {`
`32`	`32`	`return storage()->dtype_;`
`33`	`33`	`}`
	`34`	`+const phi::DDim& DenseTensorArrayType::dims() const { return storage()->dims_; }`
`34`	`35`
`35`	`36`	`const phi::DataLayout& DenseTensorArrayType::data_layout() const {`
`36`	`37`	`return storage()->layout_;`