PaddlePaddle
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.cc
Copy file name to clipboard
+10-2Lines changed: 10 additions & 2 deletions b/‎paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.cc
Copy file name to clipboard
+10-2Lines changed: 10 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions b/‎paddle/fluid/pir/dialect/kernel/ir/kernel_type.cc
Copy file name to clipboard
+13Lines changed: 13 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/kernel/ir/kernel_type.h
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions b/‎paddle/fluid/pir/dialect/kernel/ir/kernel_type.h
Copy file name to clipboard
+33Lines changed: 33 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/kernel/ir/type_storage.h
Copy file name to clipboard
+62Lines changed: 62 additions & 0 deletions b/‎paddle/fluid/pir/dialect/kernel/ir/type_storage.h
Copy file name to clipboard
+62Lines changed: 62 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.cc
Copy file name to clipboard
+29Lines changed: 29 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/fluid/pir/dialect/operator/ir/manual_api.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
@@ -300,6 +300,9 @@ void BuildValue(pir::Value value,
     var->GetMutable<phi::DenseTensor>();
   } else if (value.type().isa<paddle::dialect::AllocatedSelectedRowsType>()) {
     var->GetMutable<phi::SelectedRows>();
+  } else if (value.type()
+                 .isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
+    var->GetMutable<phi::TensorArray>();
   } else if (value.type().isa<pir::VectorType>()) {
     auto tensor_array = var->GetMutable<VariableRefArray>();
     tensor_array->clear();
 
@@ -187,6 +187,9 @@ void BuildPhiContext(pir::Operation* op,
     if (var->IsType<phi::DenseTensor>()) {
       const phi::TensorBase* tensor_in = &(var->Get<phi::DenseTensor>());
       ctx->EmplaceBackInput(InType(tensor_in));
+    } else if (var->IsType<phi::TensorArray>()) {
+      const phi::TensorBase* tensor_in = &(var->Get<phi::TensorArray>());
+      ctx->EmplaceBackInput(InType(tensor_in));
     } else if (var->IsType<VariableRefArray>()) {
       InListType inputs;
       auto& variable_array = var->Get<VariableRefArray>();
@@ -411,6 +414,13 @@ void BuildPhiContext(pir::Operation* op,
                 ->Get<phi::SelectedRows>()))));
       VLOG(8) << "ctx->EmplaceBackOutput SelectedRows: "
               << value_exec_info.GetVarName(out_ptr);
+    } else if (out_ptr.type()
+                   .isa<paddle::dialect::AllocatedDenseTensorArrayType>()) {
+      ctx->EmplaceBackOutput(OutType(const_cast<phi::TensorArray*>(
+          &(inner_scope->FindVar(value_exec_info.GetVarName(out_ptr))
+                ->Get<phi::TensorArray>()))));
+      VLOG(8) << "ctx->EmplaceBackOutput TensorArray: "
+              << value_exec_info.GetVarName(out_ptr);
     } else if (out_ptr.type().isa<pir::VectorType>()) {
       OutListType outputs;
       auto& variable_array =
 
@@ -32,9 +32,9 @@ KernelDialect::KernelDialect(pir::IrContext *context)
 }
 
 void KernelDialect::initialize() {
-  RegisterTypes<paddle::dialect::AllocatedDenseTensorType>();
   RegisterTypes<paddle::dialect::AllocatedDenseTensorType,
-                paddle::dialect::AllocatedSelectedRowsType>();
+                paddle::dialect::AllocatedSelectedRowsType,
+                paddle::dialect::AllocatedDenseTensorArrayType>();
   RegisterOps<dialect::PhiKernelOp, dialect::LegacyKernelOp>();
   RegisterAttributes<paddle::dialect::KernelAttribute>();
 }
@@ -64,6 +64,14 @@ void KernelDialect::PrintType(pir::Type type, std::ostream &os) const {
     }
     tensor_type.dtype().Print(os);
     os << ">";
+  } else if (type.isa<AllocatedDenseTensorArrayType>()) {
+    AllocatedDenseTensorArrayType tensor_array_type =
+        type.dyn_cast<AllocatedDenseTensorArrayType>();
+
+    os << phi::AllocationTypeStr(tensor_array_type.place().GetType()) << "_";
+    os << "tensor_array<";
+    tensor_array_type.dtype().Print(os);
+    os << ">";
   }
 }
 
 
@@ -65,8 +65,21 @@ size_t AllocatedSelectedRowsType::offset() const {
   return storage()->selected_rows_type_.offset();
 }
 
+const phi::Place& AllocatedDenseTensorArrayType::place() const {
+  return storage()->place_;
+}
+
+const pir::Type& AllocatedDenseTensorArrayType::dtype() const {
+  return storage()->dense_tensor_array_type_.dtype();
+}
+
+const phi::DataLayout& AllocatedDenseTensorArrayType::data_layout() const {
+  return storage()->dense_tensor_array_type_.data_layout();
+}
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorType)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedSelectedRowsType)
+IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorArrayType)
@@ -103,8 +103,41 @@ class AllocatedSelectedRowsType
   size_t offset() const;
 };
 
+class AllocatedDenseTensorArrayType
+    : public pir::Type::TypeBase<AllocatedDenseTensorArrayType,
+                                 pir::Type,
+                                 AllocatedDenseTensorArrayTypeStorage> {
+ public:
+  using Base::Base;
+
+  static AllocatedDenseTensorArrayType get(pir::IrContext *ctx,
+                                           const phi::Place &place,
+                                           dialect::DenseTensorArrayType type) {
+    return pir::TypeManager::template get<AllocatedDenseTensorArrayType>(
+        ctx, place, type);
+  }
+
+  static AllocatedDenseTensorArrayType get(pir::IrContext *ctx,
+                                           const phi::Place &place,
+                                           const pir::Type &dtype,
+                                           const phi::DataLayout &layout) {
+    dialect::DenseTensorArrayType type =
+        dialect::DenseTensorArrayType::get(ctx, dtype, layout);
+
+    return pir::TypeManager::template get<AllocatedDenseTensorArrayType>(
+        ctx, place, type);
+  }
+
+  const phi::Place &place() const;
+
+  const pir::Type &dtype() const;
+
+  const phi::DataLayout &data_layout() const;
+};
+
 }  // namespace dialect
 }  // namespace paddle
 
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorType)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedSelectedRowsType)
+IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AllocatedDenseTensorArrayType)
@@ -152,5 +152,67 @@ struct AllocatedSelectedRowsTypeStorage : public pir::TypeStorage {
   dialect::SelectedRowsType selected_rows_type_;
 };
 
+///
+/// \brief Define Parametric TypeStorage for AllocatedSelectedRowsTypeStorage.
+///
+///
+struct AllocatedDenseTensorArrayTypeStorage : public pir::TypeStorage {
+  using Place = phi::Place;
+  ///
+  /// \brief Declare ParamKey according to parameter type.
+  ///
+  using ParamKey = std::tuple<phi::Place, dialect::DenseTensorArrayType>;
+
+  AllocatedDenseTensorArrayTypeStorage(
+      const phi::Place& place, const dialect::DenseTensorArrayType& type)
+      : place_(place), dense_tensor_array_type_(type) {}
+
+  ///
+  /// \brief Each derived TypeStorage must define a Construct method, which
+  /// StorageManager uses to construct a derived TypeStorage.
+  ///
+  static AllocatedDenseTensorArrayTypeStorage* Construct(const ParamKey& key) {
+    return new AllocatedDenseTensorArrayTypeStorage(std::get<0>(key),
+                                                    std::get<1>(key));
+  }
+
+  ///
+  /// \brief Each derived TypeStorage must provide a HashValue method.
+  ///
+  static std::size_t HashValue(const ParamKey& key) {
+    std::size_t hash_value = 791;
+    // hash place
+    hash_value = pir::hash_combine(hash_value, std::get<0>(key).HashValue());
+
+    // hash dtype
+    auto dense_tensor_array_type = std::get<1>(key);
+    hash_value =
+        pir::hash_combine(hash_value,
+                          dialect::DenseTensorArrayTypeStorage::HashValue(
+                              dialect::DenseTensorArrayTypeStorage::ParamKey(
+                                  dense_tensor_array_type.dtype(),
+                                  dense_tensor_array_type.data_layout())));
+    return hash_value;
+  }
+
+  ///
+  /// \brief Each derived TypeStorage needs to overload operator==.
+  ///
+  bool operator==(const ParamKey& key) const {
+    return ParamKey(place_, dense_tensor_array_type_) == key;
+  }
+
+  ParamKey GetAsKey() const {
+    return ParamKey(place_, dense_tensor_array_type_);
+  }
+
+  ///
+  /// \brief AllocatedDenseTensorArrayTypeStorage include five parameters:
+  /// place, SelectedRowsType
+  ///
+  phi::Place place_;
+  dialect::DenseTensorArrayType dense_tensor_array_type_;
+};
+
 }  // namespace dialect
 }  // namespace paddle
@@ -119,5 +119,34 @@ pir::OpResult zeros(const std::vector<int64_t>& shape,
   return paddle::dialect::full(shape, 0, dtype, place);
 }
 
+pir::OpResult create_array(phi::DataType dtype) {
+  auto create_array_op = ApiBuilder::Instance()
+                             .GetBuilder()
+                             ->Build<paddle::dialect::CreateArrayOp>(dtype);
+  return create_array_op.out();
+}
+
+pir::OpResult array_length(pir::Value x) {
+  auto array_length_op = ApiBuilder::Instance()
+                             .GetBuilder()
+                             ->Build<paddle::dialect::ArrayLengthOp>(x);
+  return array_length_op.out();
+}
+
+pir::OpResult array_read(pir::Value array, pir::Value i) {
+  auto array_read_op =
+      ApiBuilder::Instance().GetBuilder()->Build<paddle::dialect::ArrayReadOp>(
+          array, i);
+  return array_read_op.out();
+}
+
+pir::OpResult array_write_(pir::Value array, pir::Value x, pir::Value i) {
+  auto array_write_op =
+      ApiBuilder::Instance()
+          .GetBuilder()
+          ->Build<paddle::dialect::ArrayWrite_Op>(array, x, i);
+  return array_write_op.out();
+}
+
 }  // namespace dialect
 }  // namespace paddle
@@ -60,5 +60,13 @@ pir::OpResult zeros(const std::vector<int64_t>& shape,
                     phi::DataType dtype = phi::DataType::FLOAT32,
                     const Place& place = phi::CPUPlace());
 
+pir::OpResult create_array(phi::DataType dtype);
+
+pir::OpResult array_length(pir::Value x);
+
+pir::OpResult array_read(pir::Value array, pir::Value i);
+
+pir::OpResult array_write_(pir::Value array, pir::Value x, pir::Value i);
+
 }  // namespace dialect
 }  // namespace paddle
Original file line number	Diff line number	Diff line change
`@@ -32,9 +32,9 @@ KernelDialect::KernelDialect(pir::IrContext *context)`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`void KernelDialect::initialize() {`
`35`		`- RegisterTypes<paddle::dialect::AllocatedDenseTensorType>();`
`36`	`35`	`RegisterTypes<paddle::dialect::AllocatedDenseTensorType,`
`37`		`- paddle::dialect::AllocatedSelectedRowsType>();`
	`36`	`+ paddle::dialect::AllocatedSelectedRowsType,`
	`37`	`+ paddle::dialect::AllocatedDenseTensorArrayType>();`
`38`	`38`	`RegisterOps<dialect::PhiKernelOp, dialect::LegacyKernelOp>();`
`39`	`39`	`RegisterAttributes<paddle::dialect::KernelAttribute>();`
`40`	`40`	`}`
`@@ -64,6 +64,14 @@ void KernelDialect::PrintType(pir::Type type, std::ostream &os) const {`
`64`	`64`	`}`
`65`	`65`	`tensor_type.dtype().Print(os);`
`66`	`66`	`os << ">";`
	`67`	`+ } else if (type.isa<AllocatedDenseTensorArrayType>()) {`
	`68`	`+ AllocatedDenseTensorArrayType tensor_array_type =`
	`69`	`+ type.dyn_cast<AllocatedDenseTensorArrayType>();`
	`70`	`+`
	`71`	`+ os << phi::AllocationTypeStr(tensor_array_type.place().GetType()) << "_";`
	`72`	`+ os << "tensor_array<";`
	`73`	`+ tensor_array_type.dtype().Print(os);`
	`74`	`+ os << ">";`
`67`	`75`	`}`
`68`	`76`	`}`
`69`	`77`