[CINN] Fix bug of infer_symbol_shape of same_operands_result op (#66880)

zyfncg · web-flow · commit 3c73c0341205 · 2024-08-03T10:28:56.000+08:00
* fix bug of infer_symbol_shape of same_operands_result op

* 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
@@ -66,6 +66,9 @@ class BlockDimExprsAsserter {
 
  private:
   void AssertOpRegions(const pir::Operation* op) {
+    // Inserting too many check op in while block will cause model too slow, so
+    // skip the while block checking.
+    if (op->isa<paddle::dialect::WhileOp>()) return;
     for (std::size_t i = 0; i < op->num_regions(); ++i) {
       for (auto& block : op->region(i)) {
         BlockDimExprsAsserter asserter(GraphDimExprs4Value, ir_ctx_, &block);
@@ -86,6 +89,18 @@ class BlockDimExprsAsserter {
       LOG(INFO) << "skip the checking for [ " << op->name() << " ]";
       return;
     }
+    const bool is_same_operand_result_op = [&] {
+      if (op->num_operands() != 1 || op->num_results() != 1) return false;
+      if (!op->operand_source(0).type().isa<paddle::dialect::DenseTensorType>())
+        return false;
+      if (op->result(0).type().isa<paddle::dialect::DenseTensorType>())
+        return false;
+      return GraphDimExprs4Value(op->operand_source(0)) ==
+             GraphDimExprs4Value(op->result(0));
+    }();
+    // skip the ops which operand and result have same shape
+    if (is_same_operand_result_op) return;
+
     auto OpDimExprs4Value = GetOpDimExprs4Value(op);
     const auto& inputs = [&] {
       std::vector<pir::Value> inputs;
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -14,14 +14,16 @@
 
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h"
 
-#define OP_SAME_OPERANDS_AND_RESULT(name)                                  \
-  bool name##OpInferSymbolicShape(                                         \
-      pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) { \
-    const symbol::ShapeOrDataDimExprs &operand_shape_or_data =             \
-        infer_context->GetShapeOrDataForValue(op->operand_source(0));      \
-    infer_context->SetShapeOrDataForValue(op->result(0),                   \
-                                          operand_shape_or_data);          \
-    return true;                                                           \
+#define OP_SAME_OPERANDS_AND_RESULT(name)                                     \
+  bool name##OpInferSymbolicShape(                                            \
+      pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {    \
+    const auto &operand_shape =                                               \
+        infer_context->GetShapeOrDataForValue(op->operand_source(0)).shape(); \
+    infer_context->SetShapeOrDataForValue(                                    \
+        op->result(0),                                                        \
+        symbol::ShapeOrDataDimExprs{                                          \
+            symbol::TensorShapeOrDataDimExprs(operand_shape)});               \
+    return true;                                                              \
   }
 
 namespace paddle::dialect {
@@ -37,17 +39,13 @@ OP_SAME_OPERANDS_AND_RESULT(Asin)
 OP_SAME_OPERANDS_AND_RESULT(Asin_)
 OP_SAME_OPERANDS_AND_RESULT(Asinh)
 OP_SAME_OPERANDS_AND_RESULT(Asinh_)
-OP_SAME_OPERANDS_AND_RESULT(Assign)
-OP_SAME_OPERANDS_AND_RESULT(Assign_)
 OP_SAME_OPERANDS_AND_RESULT(Atan)
 OP_SAME_OPERANDS_AND_RESULT(Atan_)
 OP_SAME_OPERANDS_AND_RESULT(Atanh)
 OP_SAME_OPERANDS_AND_RESULT(Atanh_)
 OP_SAME_OPERANDS_AND_RESULT(Bernoulli)
 OP_SAME_OPERANDS_AND_RESULT(BitwiseNot)
 OP_SAME_OPERANDS_AND_RESULT(BitwiseNot_)
-OP_SAME_OPERANDS_AND_RESULT(Cast)
-OP_SAME_OPERANDS_AND_RESULT(Cast_)
 OP_SAME_OPERANDS_AND_RESULT(Ceil)
 OP_SAME_OPERANDS_AND_RESULT(Ceil_)
 OP_SAME_OPERANDS_AND_RESULT(Celu)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -38,8 +38,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Atanh_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Bernoulli)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BitwiseNot)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BitwiseNot_)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Ceil)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Ceil_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Celu)
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
@@ -348,6 +348,19 @@ bool AsRealOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool AssignOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)));
+  return true;
+}
+
+bool Assign_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return AssignOpInferSymbolicShape(op, infer_context);
+}
+
 bool BipartiteMatchOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &dist_mat_shape_or_data =
@@ -366,6 +379,19 @@ bool BipartiteMatchOpInferSymbolicShape(
   return true;
 }
 
+bool CastOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      infer_context->GetShapeOrDataForValue(op->operand_source(0)));
+  return true;
+}
+
+bool Cast_OpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  return CastOpInferSymbolicShape(op, infer_context);
+}
+
 bool CholeskyOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -25,7 +25,11 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Argmax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Argmin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AsComplex)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AsReal)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Assign)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Assign_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BipartiteMatch)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cast_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cholesky)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cummax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Cummin)
diff --git a/paddle/pir/src/dialect/shape/utils/dim_expr.cc b/paddle/pir/src/dialect/shape/utils/dim_expr.cc
@@ -47,9 +47,7 @@ DimExpr DimExpr::operator/(const DimExpr& other) const {
   if (this->isa<std::int64_t>() && other.isa<std::int64_t>()) {
     std::int64_t num = this->dyn_cast<std::int64_t>();
     std::int64_t dem = other.dyn_cast<std::int64_t>();
-    if (num % dem == 0) {
-      return num / dem;
-    }
+    return num / dem;
   }
   const DimExpr& reciprocal = Reciprocal<DimExpr>(other);
   DimExpr div_expr = Mul<DimExpr>{List<DimExpr>{*this, reciprocal}};
