[Dynamic Shape] Provide cinn::reshape InferSymbolic (#61407)

jiahy0825 · web-flow · commit 2676124a71b5 · 2024-02-02T15:13:05.000+08:00
* [Dynamic Shape] Provide cinn::reshape InferSymbolic

* Fix grab int attribute error
diff --git a/paddle/cinn/hlir/dialect/operator/ir/ops.yaml b/paddle/cinn/hlir/dialect/operator/ir/ops.yaml
@@ -71,6 +71,7 @@
     func : ReshapeInferMeta
   kernel :
     func : reshape
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : scale
   args : (Tensor x, float scale=1.0, float bias=0.0, bool bias_after_scale=true)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc
@@ -20,9 +20,23 @@
 #include "paddle/pir/core/builtin_type_interfaces.h"
 #include "paddle/pir/dialect/shape/ir/shape_attribute.h"
 
+template <typename T>
+struct AttributeTrait;
+
+template <>
+struct AttributeTrait<std::int64_t> {
+  using value_type = ::pir::Int64Attribute;
+};
+
+template <>
+struct AttributeTrait<int> {
+  using value_type = ::pir::Int32Attribute;
+};
+
 template <typename T = int64_t>
 std::vector<T> GetVectorAttr(const ::pir::Operation *op,
                              const std::string &name) {
+  using value_type = typename AttributeTrait<T>::value_type;
   const auto &attr_map = op->attributes();
   PADDLE_ENFORCE(
       attr_map.count(name),
@@ -36,12 +50,12 @@ std::vector<T> GetVectorAttr(const ::pir::Operation *op,
   auto array_list = val.dyn_cast<::pir::ArrayAttribute>().AsVector();
   std::vector<T> vec_res;
   if (array_list.size() > 0) {
-    PADDLE_ENFORCE_EQ(array_list[0].isa<::pir::Int64Attribute>(),
+    PADDLE_ENFORCE_EQ(array_list[0].isa<value_type>(),
                       true,
                       phi::errors::Unimplemented(
                           "the 0th elementwise MUST be ir::Int64Attribute"));
     for (size_t i = 0; i < array_list.size(); ++i) {
-      vec_res.push_back(array_list[i].dyn_cast<::pir::Int64Attribute>().data());
+      vec_res.push_back(array_list[i].dyn_cast<value_type>().data());
     }
   }
   return vec_res;
@@ -1221,6 +1235,23 @@ bool ReduceSumOpInferSymbolicShape(
   return ReduceInferSymbolicShape(op, shape_analysis);
 }
 
+bool ReshapeOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  std::vector<int> shape = GetVectorAttr<int>(op, "shape");
+
+  std::vector<symbol::DimExpr> out_dims;
+  for (int dim : shape) {
+    out_dims.emplace_back(static_cast<std::int64_t>(dim));
+  }
+  symbol::ShapeOrDataDimExprs shape_data{
+      symbol::TensorShapeOrDataDimExprs(out_dims)};
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+  op->set_attribute(
+      "symbolic_shape",
+      pir::shape::SymbolAttribute::get(pir::IrContext::Instance(), shape_data));
+  return true;
+}
+
 }  // namespace cinn::dialect
 
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::InferSymbolicShapeInterface)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h
@@ -236,6 +236,8 @@ bool ReduceProdOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
 bool ReduceSumOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+bool ReshapeOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
 
 }  // namespace cinn::dialect
 
