[CINN] Fix bug of infer_symbol_shape for slice (#66745)

zyfncg · web-flow · commit c04601ba5d6d · 2024-07-30T17:03:23.000+08:00
* fix bug of infer_symbol_shape for slice

* fix bug

* fix slice
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
@@ -199,9 +199,10 @@ class BlockDimExprsAsserter {
     auto opt_shape_tensor_from_dim_exprs =
         BuildShapeTensorFromDataDimExprs(inputs, output, OpDimExprs4Value);
     if (!opt_shape_tensor_from_dim_exprs.has_value()) return;
+    const auto& output_dims =
+        output.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+    if (::common::contain_unknown_dim(output_dims)) return;
     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>(
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
@@ -218,6 +218,7 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
       op->result(0),
       paddle::dialect::slice_utils::SliceRawInferSymbolicShape(
           op->operand_source(0),
+          op->result(0),
           starts,
           ends,
           axes_raw,
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_slice_utils.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_sym_slice_utils.h
@@ -64,12 +64,15 @@ inline void CheckAndUpdateSliceAttrs(
     int64_t end_i = 0;
     if (ends.at(i).isa<int64_t>()) {
       if (in_dims.at(axis).isa<int64_t>()) {
-        ends.at(i) =
-            (ends.at(i).Get<int64_t>() > in_dims.at(axis).Get<int64_t>())
-                ? in_dims.at(axis)
-                : ends.at(i);
+        ends[i] = std::min(ends.at(i).Get<int64_t>(),
+                           in_dims.at(axis).Get<int64_t>());
+      }
+      if (ends.at(i).Get<int64_t>() < 0) {
+        ends[i] = ends.at(i) + in_dims.at(axis);
+      }
+      if (ends.at(i).isa<int64_t>()) {
+        end_i = ends.at(i).Get<int64_t>();
       }
-      end_i = ends.at(i).Get<int64_t>();
     }
 
     // For both start and end can be negative or positive, we need to handle the
@@ -144,6 +147,7 @@ inline std::vector<int64_t> FormatSliceAxes(
 
 inline ShapeOrData SliceRawInferSymbolicShape(
     const pir::Value x,
+    const pir::Value out,
     const ExprVec &starts_expr,
     const ExprVec &ends_expr,
     const std::vector<int64_t> &axes_raw,
@@ -210,10 +214,16 @@ inline ShapeOrData SliceRawInferSymbolicShape(
     const int64_t start =
         starts_int[0] < 0 ? starts_int[0] + in_shapeordata.data().value().size()
                           : starts_int[0];
-    const int64_t end =
-        static_cast<int64_t>(std::numeric_limits<int>::max()) == ends_int[0]
-            ? in_shapeordata.data().value().size()
-            : ends_int[0];
+    const int64_t end = [&]() -> int64_t {
+      if (ends_int[0] < 0) {
+        return ends_int[0] + in_shapeordata.data().value().size();
+      }
+      if (ends_int[0] ==
+          static_cast<int64_t>(std::numeric_limits<int>::max())) {
+        return in_shapeordata.data().value().size();
+      }
+      return ends_int[0];
+    }();
 
     for (int64_t i = start; i < end; i++) {
       out_data.push_back(in_shapeordata.data().value().at(i));
@@ -225,7 +235,18 @@ inline ShapeOrData SliceRawInferSymbolicShape(
         symbol::TensorShapeOrDataDimExprs(shape, out_data)};
   };
 
-  return in_shapeordata.data().has_value() ? GetDataDimExprs()
-                                           : GetShapeDimExprs();
+  const auto &out_shape = in_shapeordata.data().has_value()
+                              ? GetDataDimExprs()
+                              : GetShapeDimExprs();
+  if (out_shape.data().has_value() && out_shape.shape().empty()) {  // 0D tensor
+    const auto &out_ddim =
+        out.type().dyn_cast<paddle::dialect::DenseTensorType>().dims();
+    if (out_ddim.size() == 1 && out_ddim[0] == 1) {  // value is 1D
+      return symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+          std::vector<symbol::DimExpr>{1}, out_shape.data().value())};
+    }
+  }
+
+  return out_shape;
 }
 }  // namespace paddle::dialect::slice_utils
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
@@ -946,6 +946,7 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   infer_context->SetShapeOrDataForValue(
       res,
       slice_utils::SliceRawInferSymbolicShape(operand_source,
+                                              res,
                                               starts,
                                               ends,
                                               axes_vec,
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -504,7 +504,7 @@ ShapeConstraintIRAnalysis::GetShapeOrDataForValue(Value val) {
       SetSymbolForValueByStaticShape(val);
     } else {
       VLOG(3) << "InferShapeOrDataForValue,  defining_op: "
-              << val.defining_op()->name();
+              << val.defining_op()->name() << " id:" << val.defining_op()->id();
       InferShapeOrDataForValue(val);
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -504,7 +504,7 @@ ShapeConstraintIRAnalysis::GetShapeOrDataForValue(Value val) {`
`504`	`504`	`SetSymbolForValueByStaticShape(val);`
`505`	`505`	`} else {`
`506`	`506`	`VLOG(3) << "InferShapeOrDataForValue, defining_op: "`
`507`		`- << val.defining_op()->name();`
	`507`	`+ << val.defining_op()->name() << " id:" << val.defining_op()->id();`
`508`	`508`	`InferShapeOrDataForValue(val);`
`509`	`509`	`}`
`510`	`510`	`}`