[CINN]add ShapeOrDataDimExpr represent for TensorArray type (#65956)

Hongqing-work · web-flow · commit 8de8736f5794 · 2024-07-12T17:11:28.000+08:00
* [CINN]add ShapeOrDataDimExpr represent for TensorArray type

* change name

* fix some print info
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_util.cc
@@ -188,6 +188,9 @@ struct ShapeSignatureGenerator {
             }
           },
           [&](const symbol::TensorListShapeOrDataDimExprs& impl) { return; },
+          [&](const symbol::RankedTensorArrayShapeOrDataDimExprs& impl) {
+            return;
+          },
           [&](const symbol::NullShapeOrDataDimExpr& impl) { return; });
     };
 
@@ -249,6 +252,8 @@ struct ShapeSignatureGenerator {
         [&](const symbol::TensorListShapeOrDataDimExprs& impl) -> ResType {
           return std::make_pair(std::nullopt, std::nullopt);
         },
+        [&](const symbol::RankedTensorArrayShapeOrDataDimExprs& impl)
+            -> ResType { return std::make_pair(std::nullopt, std::nullopt); },
         [&](const symbol::NullShapeOrDataDimExpr& impl) -> ResType {
           return std::make_pair(std::nullopt, std::nullopt);
         });
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/fuse_shape_ops_into_generate_shape_op_pass.cc
@@ -70,6 +70,9 @@ std::vector<pir::Value> FindSourceDenseTensorOfDimTensor(
       [](const symbol::TensorListShapeOrDataDimExprs& dim_expr) {
         return true;
       },
+      [](const symbol::RankedTensorArrayShapeOrDataDimExprs& dim_expr) {
+        return false;
+      },
       [](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return false;
       }};
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/simplify_dim_expr_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/simplify_dim_expr_pass.cc
@@ -51,25 +51,24 @@ void VisitEachValue(const pir::Operation& op, const DoEachT& DoEach) {
   }
 }
 
+std::vector<symbol::DimExpr> SimplifyDimExprVector(
+    const std::vector<symbol::DimExpr>& original_dim_exprs) {
+  std::vector<symbol::DimExpr> simplified_dim_exprs{};
+  for (const symbol::DimExpr& dim_expr : original_dim_exprs) {
+    simplified_dim_exprs.push_back(symbol::SimplifyDimExpr(dim_expr));
+  }
+  return simplified_dim_exprs;
+}
+
 symbol::TensorShapeOrDataDimExprs SimplifyTensorShapeOrData(
     const symbol::TensorShapeOrDataDimExprs& shape_or_data) {
-  const auto& SimplifyDimExpr =
-      [](const std::vector<symbol::DimExpr>& original_dim_expr)
-      -> std::vector<symbol::DimExpr> {
-    std::vector<symbol::DimExpr> simplified_dim_expr{};
-    for (const symbol::DimExpr& dim_expr : original_dim_expr) {
-      simplified_dim_expr.push_back(symbol::SimplifyDimExpr(dim_expr));
-    }
-    return simplified_dim_expr;
-  };
-
   std::vector<symbol::DimExpr> simplified_shape =
-      SimplifyDimExpr(shape_or_data.shape());
+      SimplifyDimExprVector(shape_or_data.shape());
   if (!shape_or_data.data().has_value()) {
     return symbol::ShapeOrData<symbol::DimExpr>(simplified_shape);
   }
   std::vector<symbol::DimExpr> simplified_data =
-      SimplifyDimExpr(shape_or_data.data().value());
+      SimplifyDimExprVector(shape_or_data.data().value());
   return symbol::ShapeOrData<symbol::DimExpr>(simplified_shape,
                                               simplified_data);
 }
@@ -90,7 +89,12 @@ symbol::ShapeOrDataDimExprs SimplifyShapeOrData(
         }
         return symbol::ShapeOrDataDimExprs(simplified_tensor_list);
       },
-      [&](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
+      [](const symbol::RankedTensorArrayShapeOrDataDimExprs& tensor_array) {
+        return symbol::ShapeOrDataDimExprs(
+            symbol::RankedTensorArrayShapeOrDataDimExprs(
+                SimplifyDimExprVector(tensor_array.GetShapeHint())));
+      },
+      [](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return symbol::ShapeOrDataDimExprs(null_shape_or_data);
       }};
   return std::visit(lambdas, shape_or_data.variant());
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/local_infer_symbolic_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/local_infer_symbolic_util.cc
@@ -75,13 +75,21 @@ void InitLocalShapeAnalysis(const pir::Operation& op,
         }
         return symbol::ShapeOrDataDimExprs(ret);
       };
+  auto NewSymbolReplacedTensorArray =
+      [&](const symbol::RankedTensorArrayShapeOrDataDimExprs&
+              tensor_array_shape) {
+        return symbol::ShapeOrDataDimExprs(
+            symbol::RankedTensorArrayShapeOrDataDimExprs(
+                NewSymbolReplacedDimExprs(tensor_array_shape.GetShapeHint())));
+      };
   auto NewSymbolReplacedNull =
       [&](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return symbol::ShapeOrDataDimExprs(null_shape_or_data);
       };
   auto GetNewSymbolReplaced = [&](const auto& value_dim_exprs) {
     auto patterns = common::Overloaded{NewSymbolReplacedTensor,
                                        NewSymbolReplacedTensorList,
+                                       NewSymbolReplacedTensorArray,
                                        NewSymbolReplacedNull};
     return std::visit(patterns, value_dim_exprs.variant());
   };
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc b/paddle/cinn/hlir/dialect/operator/transforms/lowering_pass/collect_sym_expr.cc
@@ -73,6 +73,14 @@ void VisitEachDimExpr(const symbol::ShapeOrDataDimExprs& shape_or_data,
           VisitEachDimExprFromTensorShapeOrData(tensor_shape_or_data, DoEach);
         }
       },
+      [&](const symbol::RankedTensorArrayShapeOrDataDimExprs& tensor_array) {
+        PADDLE_THROW(phi::errors::Fatal(
+            "Dead code, TensorArray should not be handled in backend."));
+        for (const symbol::DimExpr& dim_expr : tensor_array.GetShapeHint()) {
+          DoEach(dim_expr);
+        }
+        return;
+      },
       [&](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return;
       }};
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pir_to_py_code_converter.cc b/paddle/cinn/hlir/dialect/operator/transforms/pir_to_py_code_converter.cc
@@ -671,6 +671,10 @@ struct PirToPyCodeConverterHelper {
         [](const symbol::TensorListShapeOrDataDimExprs& impl) {
           return ConvertTensorListShapeOrData(impl);
         },
+        [](const symbol::RankedTensorArrayShapeOrDataDimExprs& impl) {
+          // TODO(Hongqing-work): support tensor_array to py
+          return std::string("self.s_tensor_array()");
+        },
         [](const symbol::NullShapeOrDataDimExpr& impl) {
           return std::string("self.s_null()");
         });
diff --git a/paddle/fluid/pir/dialect/operator/utils/shape_analysis_utils.cc b/paddle/fluid/pir/dialect/operator/utils/shape_analysis_utils.cc
@@ -45,6 +45,11 @@ symbol::ShapeOrDataDimExprs ClearDataInfo(
         }
         return symbol::ShapeOrDataDimExprs{new_shape_exprs};
       },
+      [](const symbol::RankedTensorArrayShapeOrDataDimExprs& shape_exprs) {
+        return symbol::ShapeOrDataDimExprs{
+            symbol::RankedTensorArrayShapeOrDataDimExprs{
+                shape_exprs.GetShapeHint()}};
+      },
       [](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return symbol::ShapeOrDataDimExprs{null_shape_or_data};
       }};
diff --git a/paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h b/paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h
@@ -114,12 +114,41 @@ class ShapeOrData {
   std::optional<std::vector<T>> data_;
 };
 
+using NullShapeOrDataDimExpr = std::monostate;
 using TensorShapeOrDataDimExprs = ShapeOrData<DimExpr>;
 using TensorListShapeOrDataDimExprs = std::vector<TensorShapeOrDataDimExprs>;
-using NullShapeOrDataDimExpr = std::monostate;
-using ShapeOrDataDimExprsBase = std::variant<NullShapeOrDataDimExpr,
-                                             TensorShapeOrDataDimExprs,
-                                             TensorListShapeOrDataDimExprs>;
+
+/* TensorArray can append tensors dynamically. In a static graph, we only
+ * store the shape of one element as a hint, because we assume that all elements
+ * in the TensorArray have the same rank, and with equal constraints on specific
+ * dimensions. */
+class RankedTensorArrayShapeOrDataDimExprs {
+ public:
+  RankedTensorArrayShapeOrDataDimExprs() = default;
+  explicit RankedTensorArrayShapeOrDataDimExprs(
+      const std::vector<DimExpr>& shape)
+      : shape_hint_{shape} {}
+  const std::vector<DimExpr>& GetShapeHint() const { return shape_hint_; }
+  bool operator==(const RankedTensorArrayShapeOrDataDimExprs& other) const {
+    if (shape_hint_.size() != other.shape_hint_.size()) return false;
+    for (size_t i = 0; i < shape_hint_.size(); ++i) {
+      DimExpr dim0 = symbol::SimplifyDimExpr(shape_hint_[i]);
+      DimExpr dim1 = symbol::SimplifyDimExpr(other.shape_hint_[i]);
+      if (dim0 != dim1) return false;
+    }
+
+    return true;
+  }
+
+ private:
+  std::vector<DimExpr> shape_hint_;
+};
+
+using ShapeOrDataDimExprsBase =
+    std::variant<NullShapeOrDataDimExpr,
+                 TensorShapeOrDataDimExprs,
+                 TensorListShapeOrDataDimExprs,
+                 RankedTensorArrayShapeOrDataDimExprs>;
 
 class ShapeOrDataDimExprs : public ShapeOrDataDimExprsBase {
  public:
@@ -131,6 +160,10 @@ class ShapeOrDataDimExprs : public ShapeOrDataDimExprsBase {
       const TensorListShapeOrDataDimExprs& tensor_list_dim_exprs)
       : ShapeOrDataDimExprsBase(tensor_list_dim_exprs) {}
 
+  ShapeOrDataDimExprs(const RankedTensorArrayShapeOrDataDimExprs&
+                          tensor_array_dim_exprs)  // NOLINT
+      : ShapeOrDataDimExprsBase(tensor_array_dim_exprs) {}
+
   ShapeOrDataDimExprs(const NullShapeOrDataDimExpr& null_dim_expr)  // NOLINT
       : ShapeOrDataDimExprsBase(null_dim_expr) {}
 
@@ -227,6 +260,14 @@ struct hash<symbol::TensorListShapeOrDataDimExprs> {
   }
 };
 
+template <>
+struct hash<symbol::RankedTensorArrayShapeOrDataDimExprs> {
+  std::size_t operator()(
+      const symbol::RankedTensorArrayShapeOrDataDimExprs& obj) const {
+    return std::hash<std::vector<symbol::DimExpr>>()(obj.GetShapeHint());
+  }
+};
+
 template <>
 struct hash<symbol::ShapeOrDataDimExprs> {
   std::size_t operator()(const symbol::ShapeOrDataDimExprs& obj) const {
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -273,6 +273,11 @@ InferSymbolicShapeContext::SimplifyBroadcastForShapeOrData(
         }
         return symbol::ShapeOrDataDimExprs(simplified_tensor_list);
       },
+      [&](const symbol::RankedTensorArrayShapeOrDataDimExprs& tensor_array) {
+        symbol::RankedTensorArrayShapeOrDataDimExprs simplified_tensor_array(
+            DimExprsVisitor(tensor_array.GetShapeHint()));
+        return symbol::ShapeOrDataDimExprs(simplified_tensor_array);
+      },
       [&](const symbol::NullShapeOrDataDimExpr& null_shape_or_data) {
         return symbol::ShapeOrDataDimExprs(null_shape_or_data);
       });
diff --git a/paddle/pir/src/dialect/shape/utils/shape_or_data_expr.cc b/paddle/pir/src/dialect/shape/utils/shape_or_data_expr.cc
@@ -15,28 +15,27 @@
 #include "paddle/pir/include/dialect/shape/utils/shape_or_data_expr.h"
 
 namespace symbol {
+std::vector<DimExpr> SubstituteDimExprVector(
+    const std::vector<DimExpr>& original_dim_expr,
+    const std::unordered_map<DimExpr, DimExpr>& substitution_pattern) {
+  std::vector<DimExpr> substituted_dim_expr{};
+  for (const DimExpr& dim_expr : original_dim_expr) {
+    const auto& tmp_dim_expr =
+        SubstituteDimExpr(dim_expr, substitution_pattern);
+    substituted_dim_expr.push_back(SimplifyDimExpr(tmp_dim_expr));
+  }
+  return substituted_dim_expr;
+}
+
 TensorShapeOrDataDimExprs SubstituteTensorShapeOrData(
     const TensorShapeOrDataDimExprs& shape_or_data,
     const std::unordered_map<DimExpr, DimExpr>& substitution_pattern) {
-  auto SubstituteOneDimExpr =
-      [](const std::vector<DimExpr>& original_dim_expr,
-         const std::unordered_map<DimExpr, DimExpr>& substitution_pattern)
-      -> std::vector<DimExpr> {
-    std::vector<DimExpr> substituted_dim_expr{};
-    for (const DimExpr& dim_expr : original_dim_expr) {
-      const auto& tmp_dim_expr =
-          SubstituteDimExpr(dim_expr, substitution_pattern);
-      substituted_dim_expr.push_back(SimplifyDimExpr(tmp_dim_expr));
-    }
-    return substituted_dim_expr;
-  };
-
   std::vector<DimExpr> substituted_shape =
-      SubstituteOneDimExpr(shape_or_data.shape(), substitution_pattern);
+      SubstituteDimExprVector(shape_or_data.shape(), substitution_pattern);
   if (!shape_or_data.data().has_value()) {
     return ShapeOrData<DimExpr>(substituted_shape);
   } else {
-    std::vector<DimExpr> substituted_data = SubstituteOneDimExpr(
+    std::vector<DimExpr> substituted_data = SubstituteDimExprVector(
         shape_or_data.data().value(), substitution_pattern);
     return ShapeOrData<DimExpr>(substituted_shape, substituted_data);
   }
@@ -59,6 +58,12 @@ ShapeOrDataDimExprs SubstituteShapeOrData(
         }
         return ShapeOrDataDimExprs(substituted_tensor_list);
       },
+      [&](const RankedTensorArrayShapeOrDataDimExprs& tensor_array) {
+        RankedTensorArrayShapeOrDataDimExprs substituted_tensor_array(
+            SubstituteDimExprVector(tensor_array.GetShapeHint(),
+                                    substitution_pattern));
+        return ShapeOrDataDimExprs(substituted_tensor_array);
+      },
       [&](const NullShapeOrDataDimExpr& null_shape_or_data) {
         return ShapeOrDataDimExprs(null_shape_or_data);
       }};
@@ -89,6 +94,10 @@ std::ostream& operator<<(std::ostream& stream,
           }
         }
       },
+      [&](const RankedTensorArrayShapeOrDataDimExprs& tensor_array_shape_data) {
+        stream << "TensorArray with shape hint: "
+               << tensor_array_shape_data.GetShapeHint();
+      },
       [&](const NullShapeOrDataDimExpr& null_shape_data) {
         stream << "shape[NULL], data[NULL]";
       }};
