[CINN] Fix bug of infer symbolic shape for some op (#63789)

zyfncg · web-flow · commit 060e5c564e13 · 2024-04-25T11:37:18.000+08:00
* fix bug of infer symbolic shape for some op

* fix bug of check symbol infer pass

* fix bug
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/check_infer_symbolic_pass.cc
@@ -25,6 +25,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/pattern_rewrite/pattern_rewrite_driver.h"
@@ -260,7 +261,8 @@ class BlockDimExprsAsserter {
                 input_tensors, output_dim_expr_attrs, symbol_bindings)
             .out();
     return builder_
-        .Build<paddle::dialect::CastOp>(out_shape_value, phi::DataType::INT32)
+        .Build<cinn::dialect::GenerateShapeOp>(
+            input_tensors, output_dim_expr_attrs, symbol_bindings)
         .out();
   }
 
@@ -275,7 +277,18 @@ class BlockDimExprsAsserter {
     auto opt_shape_tensor_from_dim_exprs =
         BuildShapeTensorFromDataDimExprs(inputs, output, OpDimExprs4Value);
     if (!opt_shape_tensor_from_dim_exprs.has_value()) return;
-    AddAssertEqual(op, opt_shape_tensor_from_dim_exprs.value(), output);
+    pir::Value flatten_output = [&] {
+      const auto& output_dims =
+          output.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+      if (output_dims.size() > 1) {
+        return builder_
+            .Build<paddle::dialect::FlattenOp>(
+                output, 0, output_dims.size() - 1)
+            .out();
+      }
+      return output;
+    }();
+    AddAssertEqual(op, opt_shape_tensor_from_dim_exprs.value(), flatten_output);
   }
 
   size_t GetNumel(pir::Value value) {
@@ -302,12 +315,28 @@ class BlockDimExprsAsserter {
                           "received lhs's numel is [%d], rhs's numel is [%d]",
                           lhs_numel,
                           rhs_numel));
+
+    pir::Value rhs_value = [&] {
+      const auto& lhs_dtype =
+          lhs.type().dyn_cast<paddle::dialect::DenseTensorType>().dtype();
+      const auto& rhs_dtype =
+          rhs.type().dyn_cast<paddle::dialect::DenseTensorType>().dtype();
+      if (lhs_dtype != rhs_dtype) {
+        return builder_
+            .Build<paddle::dialect::CastOp>(
+                rhs, paddle::dialect::TransToPhiDataType(lhs_dtype))
+            .out();
+      }
+      return rhs;
+    }();
+
     pir::Value lhs_eq_rhs =
-        builder_.Build<paddle::dialect::EqualOp>(lhs, rhs).out();
+        builder_.Build<paddle::dialect::EqualOp>(lhs, rhs_value).out();
     pir::Value all_eq =
         builder_.Build<paddle::dialect::AllOp>(lhs_eq_rhs).out();
     pir::Value assert_data =
-        builder_.Build<pir::CombineOp>(std::vector<pir::Value>{lhs, rhs}).out();
+        builder_.Build<pir::CombineOp>(std::vector<pir::Value>{lhs, rhs_value})
+            .out();
     auto assert_op = builder_.Build<paddle::dialect::AssertOp>(
         all_eq, assert_data, lhs_numel);
     const std::string error_msg = "Check [" + op->name() + "_" +
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/split_generate_shape_into_shape_ops_pass.cc
@@ -127,15 +127,29 @@ struct CachedDimExprToValueConverter {
           ->Build<paddle::dialect::CastOp>(value, phi::DataType::INT64)
           .out();
     };
+    auto FlattenValueIfNeed = [&](pir::Value value) {
+      const auto& dims =
+          value.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+      if (dims.size() <= 1) {
+        return value;
+      }
+      return rewriter
+          ->Build<paddle::dialect::FlattenOp>(value, 0, dims.size() - 1)
+          .out();
+    };
     if (tensor_dim.value.type()
             .dyn_cast<paddle::dialect::DenseTensorType>()
             .dims()
             .size() == 0) {
-      return CastToInt64IfNeed(tensor_dim.value);
+      return CastToInt64IfNeed(
+          rewriter
+              ->Build<paddle::dialect::ReshapeOp>(tensor_dim.value,
+                                                  std::vector<int64_t>{1})
+              .out());
     }
     return CastToInt64IfNeed(rewriter
                                  ->Build<paddle::dialect::SliceOp>(
-                                     tensor_dim.value,
+                                     FlattenValueIfNeed(tensor_dim.value),
                                      std::vector<int64_t>{0LL},
                                      std::vector<int64_t>{tensor_dim.axis},
                                      std::vector<int64_t>{tensor_dim.axis + 1},
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc
@@ -294,14 +294,9 @@ bool GatherNdOpInferSymbolicShape(
   const auto &index_shape_or_data =
       shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
 
-  const std::vector<symbol::DimExpr> &x_sym_shape =
-      x_shape_or_data.data().has_value() ? x_shape_or_data.data().value()
-                                         : x_shape_or_data.shape();
-
+  const std::vector<symbol::DimExpr> &x_sym_shape = x_shape_or_data.shape();
   const std::vector<symbol::DimExpr> &index_sym_shape =
-      index_shape_or_data.data().has_value()
-          ? index_shape_or_data.data().value()
-          : index_shape_or_data.shape();
+      index_shape_or_data.shape();
 
   int x_dims_size = x_sym_shape.size();
   int index_dims_size = index_sym_shape.size();
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/cinn_op_infer_sym.cc
@@ -126,18 +126,6 @@ bool ReshapeOpInferSymbolicShape(
   std::vector<int> shape =
       paddle::dialect::details::GetVectorAttr<int>(op, "shape");
 
-  const symbol::ShapeOrDataDimExprs &x_dim_expr =
-      shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
-  if (x_dim_expr.data().has_value()) {
-    if (shape.size() == 1 && shape.front() == 1) {
-      shape_analysis->SetShapeOrDataForValue(
-          op->result(0),
-          symbol::TensorShapeOrDataDimExprs(std::vector<symbol::DimExpr>{1},
-                                            x_dim_expr.data().value()));
-      return true;
-    }
-  }
-
   const auto &GetProduct = [&](const auto &dim_exprs, const auto &Filter) {
     symbol::DimExpr product{1};
     for (const auto &dim_expr : dim_exprs) {
@@ -170,8 +158,10 @@ bool ReshapeOpInferSymbolicShape(
     return target_shape;
   }();
 
-  const auto &original_shape =
-      shape_analysis->GetShapeOrDataForValue(op->operand_source(0)).shape();
+  const symbol::ShapeOrDataDimExprs &x_dim_expr =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
+
+  const auto &original_shape = x_dim_expr.shape();
 
   const auto &out_dims = [&] {
     const auto &numel =
@@ -193,8 +183,14 @@ bool ReshapeOpInferSymbolicShape(
     return out_dims;
   }();
 
-  symbol::ShapeOrDataDimExprs shape_data{
-      symbol::TensorShapeOrDataDimExprs(out_dims)};
+  symbol::ShapeOrDataDimExprs shape_data = [&] {
+    if (x_dim_expr.data().has_value()) {
+      return symbol::TensorShapeOrDataDimExprs(out_dims,
+                                               x_dim_expr.data().value());
+    }
+    return symbol::TensorShapeOrDataDimExprs(out_dims);
+  }();
+
   shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
 
   return true;
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -476,23 +476,6 @@ bool ReshapeOpInferSymbolicShape(
       shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
   const symbol::ShapeOrDataDimExprs &shape_dim_expr =
       shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
-  if (x_dim_expr.data().has_value()) {
-    const auto &shape_data = details::GetExprVecFromData(shape_dim_expr);
-    auto IsOne = [](const symbol::DimExpr &expr) {
-      return expr.isa<int64_t>() && expr.dyn_cast<int64_t>() == 1;
-    };
-    if (shape_data.size() == 1 && IsOne(shape_data.at(0))) {
-      shape_analysis->SetShapeOrDataForValue(
-          op->result(0),
-          symbol::TensorShapeOrDataDimExprs(shape_data,
-                                            x_dim_expr.data().value()));
-      shape_analysis->SetShapeOrDataForValue(
-          op->result(1),
-          CreateShapeOrDataForXShape(
-              shape_analysis->GetShapeOrDataForValue(op->operand_source(0))));
-      return true;
-    }
-  }
 
   const auto &GetProduct = [&](const auto &dim_exprs, const auto &Filter) {
     symbol::DimExpr product{1};
@@ -544,20 +527,15 @@ bool ReshapeOpInferSymbolicShape(
     return out_dims;
   }();
 
-  symbol::ShapeOrDataDimExprs shape_data{
-      symbol::TensorShapeOrDataDimExprs(out_dims)};
-
-  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
-
-  const auto UNUSED &x_shape = [&] {
-    std::vector<symbol::DimExpr> x_shape{symbol::DimExpr(0)};
-    const auto &original_shape =
-        shape_analysis->GetShapeOrDataForValue(op->operand_source(0)).shape();
-    for (const auto &dim : original_shape) {
-      x_shape.push_back(dim);
+  symbol::ShapeOrDataDimExprs shape_data = [&] {
+    if (x_dim_expr.data().has_value()) {
+      return symbol::TensorShapeOrDataDimExprs(out_dims,
+                                               x_dim_expr.data().value());
     }
-    return x_shape;
+    return symbol::TensorShapeOrDataDimExprs(out_dims);
   }();
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   shape_analysis->SetShapeOrDataForValue(
       op->result(1),
       CreateShapeOrDataForXShape(
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
@@ -48,7 +48,23 @@ class ShapeOrData {
               shape[0].template Get<std::int64_t>(),
               data.size()));
     } else {
-      IR_THROW("Size of shape should be 0 or 1, but got %d", shape.size());
+      int64_t numel = 1;
+      for (const auto& expr : shape) {
+        PADDLE_ENFORCE_EQ(expr.template isa<int64_t>(),
+                          true,
+                          ::common::errors::InvalidArgument(
+                              "When data has value, the expr of shape should "
+                              "be int, but got %s.",
+                              ToString(expr)));
+        numel *= expr.template Get<int64_t>();
+      }
+      PADDLE_ENFORCE_EQ(numel,
+                        data.size(),
+                        ::common::errors::InvalidArgument(
+                            "Size of data should be the same as "
+                            "product of value[%d] of shape, but got [%d].",
+                            numel,
+                            data.size()));
     }
   }
 
