[PIR][Dy2St] Set TensorArray dtype after first write if its type is undefined (#67981)

SigureMo · web-flow · commit b99a3cd07397 · 2024-09-16T13:22:23.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -2099,33 +2099,33 @@ std::vector<pir::Type> ArrayWrite_Op::InferMeta(
       x_type.offset());
   paddle::dialect::IrMetaTensor meta_x(&dense_x);
 
-  paddle::dialect::IrTensor dense_out;
-  paddle::dialect::IrMetaTensor meta_out(&dense_out);
+  paddle::dialect::IrTensor dense_array_out;
+  paddle::dialect::IrMetaTensor meta_out(&dense_array_out);
 
   phi::ArrayWriteInferMeta(
       meta_array, meta_x, &meta_out, phi::MetaConfig(false, false));
 
   std::vector<pir::Type> argument_outputs;
   pir::Type out_type = paddle::dialect::DenseTensorArrayType::get(
       pir::IrContext::Instance(),
-      paddle::dialect::TransToIrDataType(dense_array.dtype()),
+      paddle::dialect::TransToIrDataType(dense_array_out.dtype()),
       x_type.dims(),
-      dense_array.layout());
+      dense_array_out.layout());
   // update array's dims as x's dims.
   // TOOD(chenxi67) Do not change if dim is set by custom
   if (array_.type().isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
     array_.set_type(paddle::dialect::AllocatedDenseTensorArrayType::get(
         pir::IrContext::Instance(),
         place,
-        array_type.dtype(),
+        paddle::dialect::TransToIrDataType(dense_array_out.dtype()),
         x_type.dims(),
-        array_type.data_layout()));
+        dense_array_out.layout()));
   } else if (array_.type().isa<paddle::dialect::DenseTensorArrayType>()) {
-    array_.set_type(
-        paddle::dialect::DenseTensorArrayType::get(pir::IrContext::Instance(),
-                                                   array_type.dtype(),
-                                                   x_type.dims(),
-                                                   array_type.data_layout()));
+    array_.set_type(paddle::dialect::DenseTensorArrayType::get(
+        pir::IrContext::Instance(),
+        paddle::dialect::TransToIrDataType(dense_array_out.dtype()),
+        x_type.dims(),
+        dense_array_out.layout()));
   }
 
   argument_outputs.push_back(out_type);
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.h b/paddle/fluid/pir/dialect/operator/utils/utils.h
@@ -31,7 +31,9 @@ using VariantType = phi::Attribute;
 // TODO(zhangbo): The builtin type needs to cover all data types of
 // phi::DataType.
 static inline phi::DataType TransToPhiDataType(pir::Type dtype) {
-  if (dtype.isa<pir::BFloat16Type>()) {
+  if (dtype.isa<pir::UndefinedType>()) {
+    return phi::DataType::UNDEFINED;
+  } else if (dtype.isa<pir::BFloat16Type>()) {
     return phi::DataType::BFLOAT16;
   } else if (dtype.isa<pir::Float16Type>()) {
     return phi::DataType::FLOAT16;
@@ -76,6 +78,8 @@ static inline pir::Type TransToIrDataType(phi::DataType dtype,
     ctx = pir::IrContext::Instance();
   }
   switch (dtype) {
+    case phi::DataType::UNDEFINED:
+      return pir::UndefinedType::get(ctx);
     case phi::DataType::BFLOAT16:
       return pir::BFloat16Type::get(ctx);
     case phi::DataType::FLOAT16:
diff --git a/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h b/paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
@@ -613,6 +613,10 @@ Json AttrTypeWriter::WriteBuiltInType(const pir::Type& type) {
     VLOG(8) << "Write DenseTensorType ... ";
     return pir::serializeTypeToJsonIncludeWriteType<pir::DenseTensorType>(
         type.dyn_cast<pir::DenseTensorType>());
+  } else if (type.isa<pir::UndefinedType>()) {
+    PADDLE_THROW(common::errors::PreconditionNotMet(
+        "Unexpected type pir::UndefinedType, "
+        "it should be replace with a concrete type when ArrayWrite."));
   } else {
     PADDLE_ENFORCE(false,
                    common::errors::InvalidArgument(
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -1180,6 +1180,22 @@ phi::KernelKey GetKernelKey(
     return {backend, phi::DataLayout::ANY, dtype};
   }
 
+  if (op->isa<CreateArrayOp>()) {
+    VLOG(6) << "CreateArrayOp doesn't need a kernel";
+    auto backend = paddle::experimental::ParseBackend(place);
+    auto dtype = op->result(0)
+                     .type()
+                     .dyn_cast<paddle::dialect::DenseTensorArrayType>()
+                     .dtype();
+
+    phi::KernelKey res(
+        backend, phi::DataLayout::ANY, TransToPhiDataType(dtype));
+    if (NeedFallBackCpu(op, kernel_fn_str, res)) {
+      res.set_backend(phi::Backend::CPU);
+    }
+    return res;
+  }
+
   phi::Backend kernel_backend = phi::Backend::UNDEFINED;
   phi::DataLayout kernel_layout = phi::DataLayout::UNDEFINED;
   phi::DataType kernel_dtype = phi::DataType::UNDEFINED;
diff --git a/paddle/fluid/pybind/control_flow_api.cc b/paddle/fluid/pybind/control_flow_api.cc
@@ -358,8 +358,8 @@ std::vector<Value> PyWhileOp::OptimizeUpdate() {
 
   for (size_t operand_index = 1u, arg_index = 0u; operand_index < operand_num;
        ++operand_index) {
+    operand_source(operand_index).set_type(body_block.arg(arg_index).type());
     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/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -150,17 +150,21 @@ 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(),
-                      common::errors::InvalidArgument(
-                          "The dtype (%s) of input x shall be same as "
-                          "dtype (%d) of array.",
-                          x.dtype(),
-                          array.dtype()));
+  phi::DataType out_dtype = array.dtype();
+  if (x.dtype() != phi::DataType::UNDEFINED) {
+    if (array.dtype() == phi::DataType::UNDEFINED) {
+      out_dtype = x.dtype();
+    } else {
+      PADDLE_ENFORCE_EQ(array.dtype(),
+                        x.dtype(),
+                        common::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_dtype(out_dtype);
   out->set_layout(array.layout());
 }
 
diff --git a/paddle/pir/include/core/builder.h b/paddle/pir/include/core/builder.h
@@ -23,6 +23,7 @@
 
 namespace pir {
 class Type;
+class UndefinedType;
 class UInt8Type;
 class Int8Type;
 class Int16Type;
@@ -136,6 +137,7 @@ class Builder {
   template <typename OpTy, typename... Args>
   OpTy Build(Args &&...args);
 
+  IR_API UndefinedType undefined_type();
   IR_API BoolType bool_type();
   IR_API UInt8Type uint8_type();
   IR_API Int8Type int8_type();
diff --git a/paddle/pir/include/core/builtin_type.h b/paddle/pir/include/core/builtin_type.h
@@ -95,6 +95,7 @@ class IR_API DenseTensorType : public Type::TypeBase<DenseTensorType,
   };
 
 #define FOREACH_BUILTIN_TYPE(__macro)      \
+  __macro(UndefinedType, "t_undefined");   \
   __macro(BFloat16Type, "t_bf16");         \
   __macro(Float16Type, "t_f16");           \
   __macro(Float32Type, "t_f32");           \
@@ -117,6 +118,7 @@ FOREACH_BUILTIN_TYPE(DECLARE_BUILTIN_TYPE)
 
 }  // namespace pir
 
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::UndefinedType)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::UInt8Type)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::Int8Type)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::VectorType)
diff --git a/paddle/pir/src/core/builtin_dialect.cc b/paddle/pir/src/core/builtin_dialect.cc
@@ -29,7 +29,8 @@ BuiltinDialect::BuiltinDialect(IrContext* context)
 
 void BuiltinDialect::initialize() {
   // Register all built-in types defined in builtin_type.h.
-  RegisterTypes<BFloat16Type,
+  RegisterTypes<UndefinedType,
+                BFloat16Type,
                 Float16Type,
                 Float32Type,
                 Float64Type,
diff --git a/paddle/pir/src/core/builtin_type.cc b/paddle/pir/src/core/builtin_type.cc
@@ -53,6 +53,7 @@ DenseTensorType DenseTensorType::dyn_cast_impl(Type type) {
 
 }  // namespace pir
 
+IR_DEFINE_EXPLICIT_TYPE_ID(pir::UndefinedType)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::UInt8Type)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::Int8Type)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::VectorType)
diff --git a/paddle/pir/src/core/ir_context.cc b/paddle/pir/src/core/ir_context.cc
@@ -153,6 +153,7 @@ class IrContextImpl {
   // TypeStorage uniquer and cache instances.
   StorageManager registed_type_storage_manager_;
   // Cache some built-in type objects.
+  UndefinedType undefined_type;
   BFloat16Type bfp16_type;
   Float16Type fp16_type;
   Float32Type fp32_type;
@@ -198,6 +199,7 @@ IrContext::IrContext() : impl_(new IrContextImpl()) {
   GetOrRegisterDialect<BuiltinDialect>();
   VLOG(10) << "==============================================";
 
+  impl_->undefined_type = TypeManager::get<UndefinedType>(this);
   impl_->bfp16_type = TypeManager::get<BFloat16Type>(this);
   impl_->fp16_type = TypeManager::get<Float16Type>(this);
   impl_->fp32_type = TypeManager::get<Float32Type>(this);
@@ -341,6 +343,10 @@ const AbstractAttribute &AbstractAttribute::lookup(TypeId type_id,
   return *abstract_attribute;
 }
 
+UndefinedType UndefinedType::get(IrContext *ctx) {
+  return ctx->impl().undefined_type;
+}
+
 BFloat16Type BFloat16Type::get(IrContext *ctx) {
   return ctx->impl().bfp16_type;
 }
diff --git a/paddle/pir/src/core/ir_printer.cc b/paddle/pir/src/core/ir_printer.cc
@@ -44,7 +44,9 @@ void BasicIrPrinter::PrintType(Type type) {
     return;
   }
 
-  if (type.isa<BFloat16Type>()) {
+  if (type.isa<UndefinedType>()) {
+    os << "undefined";
+  } else if (type.isa<BFloat16Type>()) {
     os << "bf16";
   } else if (type.isa<Float16Type>()) {
     os << "f16";
diff --git a/python/paddle/jit/dy2static/convert_operators.py b/python/paddle/jit/dy2static/convert_operators.py
@@ -196,7 +196,12 @@ def _convert_tensor_arrray_if_necessary(setterhelper, push_pop_names):
 
     def maybe_to_tensor_array(v):
         if isinstance(v, list):
-            return paddle.tensor.create_array("float32", initialized_list=v)
+            dtype = (
+                paddle.base.libpaddle.DataType.UNDEFINED
+                if use_pir_api()
+                else "float32"
+            )
+            return paddle.tensor.create_array(dtype, initialized_list=v)
         else:
             return v
 
diff --git a/python/paddle/tensor/array.py b/python/paddle/tensor/array.py
@@ -276,7 +276,8 @@ def array_write(
         if array is None:
             array = paddle._pir_ops.create_array(x.dtype)
 
-        x = paddle.cast(x, array.dtype)
+        if array.dtype != paddle.base.libpaddle.DataType.UNDEFINED:
+            x = paddle.cast(x, array.dtype)
         paddle._pir_ops.array_write_(array, x, i)
         return array
     else:
@@ -360,7 +361,8 @@ def create_array(
             dtype = paddle.base.framework.convert_np_dtype_to_dtype_(dtype)
         out = paddle._pir_ops.create_array(dtype)
         for val in array:
-            val = paddle.cast(val, dtype)
+            if dtype != paddle.base.libpaddle.DataType.UNDEFINED:
+                val = paddle.cast(val, dtype)
             paddle._pir_ops.array_write_(out, val, array_length(out))
         return out
     else:
diff --git a/test/dygraph_to_static/test_for_enumerate.py b/test/dygraph_to_static/test_for_enumerate.py
@@ -173,7 +173,7 @@ def for_iter_var_list(x):
     # 1. prepare data, ref test_list.py
     x = paddle.to_tensor(x)
     iter_num = paddle.tensor.fill_constant(shape=[1], value=5, dtype="int32")
-    a = paddle.tensor.create_array("int32")
+    a = []
     for i in range(iter_num):
         a.append(x + i)
     # 2. iter list[var]
@@ -188,7 +188,7 @@ def for_enumerate_var_list(x):
     # 1. prepare data, ref test_list.py
     x = paddle.to_tensor(x)
     iter_num = paddle.tensor.fill_constant(shape=[1], value=5, dtype="int32")
-    a = paddle.tensor.create_array("int32")
+    a = []
     for i in range(iter_num):
         a.append(x + i)
     # 2. iter list[var]
diff --git a/test/dygraph_to_static/test_list.py b/test/dygraph_to_static/test_list.py
@@ -21,6 +21,7 @@
     IrMode,
     ToStaticMode,
     disable_test_case,
+    test_ast_only,
     test_legacy_and_pt_and_pir,
 )
 
@@ -373,5 +374,22 @@ def test_to_static(self):
         self.assertEqual(res, 48.0)
 
 
+def tensor_array_dtype():
+    l = []
+    for i in range(paddle.to_tensor(3)):
+        l.append(i)
+    return l[0]
+
+
+class TestTensorArrayDtype(Dy2StTestBase):
+    @test_ast_only
+    @test_legacy_and_pt_and_pir
+    def test_tensor_array_dtype(self):
+        fn = tensor_array_dtype
+        static_fn = paddle.jit.to_static(fn)
+        st_out = static_fn()
+        self.assertEqual(st_out.dtype, paddle.int64)
+
+
 if __name__ == '__main__':
     unittest.main()
