[PIR][DynamicShape] Fix SliceOp InferSymblilcShape && fix yield Op (#62068)

zhangboSJTU · web-flow · commit e99bcaa81c6d · 2024-02-27T10:23:13.000+08:00
* fix slice

* fix yield Op
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/paddle_op_infer_sym.cc
@@ -325,7 +325,16 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
   const auto &GetDataDimExprs = [&]() -> symbol::ShapeOrDataDimExprs {
     const std::vector<symbol::DimExpr> out_data = [&] {
       std::vector<symbol::DimExpr> out_data;
-      for (int64_t i = starts[0]; i < ends[0]; i++) {
+      const int64_t start =
+          starts[0] < 0
+              ? starts[0] + operand_shape_or_data.data().value().size()
+              : starts[0];
+      const int64_t end =
+          static_cast<int64_t>(std::numeric_limits<int>::max()) == ends[0]
+              ? operand_shape_or_data.data().value().size()
+              : ends[0];
+
+      for (int64_t i = start; i < end; i++) {
         out_data.push_back(operand_shape_or_data.data().value()[i]);
       }
       return out_data;
@@ -352,7 +361,7 @@ bool SliceOpInferSymbolicShape(pir::Operation *op,
                  static_cast<int64_t>(std::numeric_limits<int>::max());
     };
     for (size_t i = 0; i < axes.size(); ++i) {
-      int64_t axis = axes[i];
+      const int64_t axis = axes[i];
       auto end =
           IsMaxInt(dim_expr_ends[i]) ? out_shape[axis] : dim_expr_ends[i];
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -723,11 +723,6 @@ std::vector<std::vector<pir::Value>> WhileOp::Vjp(
 
 bool WhileOp::InferSymbolicShape(
     pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  VLOG(3) << "############ WhileOp::InferSymbolicShape start...";
-  pir::Program *body_program = body().parent_program();
-  VLOG(3) << "##### WhileOp::InferSymbolicShape: sub_program id = "
-          << body_program->module_op().operation()->id();
-
   for (auto &value : block_args()) {
     std::vector<symbol::DimExpr> sym_dims;
     const std::vector<int64_t> &dims =
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -128,10 +128,23 @@ struct ShadowOutputOpInferSymbolicShapeInterfaceModel
       : InferSymbolicShapeInterface::Concept(InferSymbolicShape) {}
 };
 
+struct YieldOpInferSymbolicShapeInterfaceModel
+    : public InferSymbolicShapeInterface::Concept {
+  static inline bool InferSymbolicShape(
+      pir::Operation* op, pir::ShapeConstraintIRAnalysis* shape_analysis) {
+    // Since YieldOp has no output, just return true
+    return true;
+  }
+
+  YieldOpInferSymbolicShapeInterfaceModel()
+      : InferSymbolicShapeInterface::Concept(InferSymbolicShape) {}
+};
+
 OperatorDialect::OperatorDialect(pir::IrContext* ctx)
     : pir::Dialect(name(), ctx, pir::TypeId::get<OperatorDialect>()) {
   initialize();
   ctx->GetOrRegisterDialect<::pir::ControlFlowDialect>();
+
   auto info = ctx->GetRegisteredOpInfo(pir::TuplePushOp::name());
   info.AttachInterface(std::move(
       pir::InterfaceValue::Get<VjpInterface, TuplePushOpVjpInterfaceModel>()));
@@ -151,6 +164,11 @@ OperatorDialect::OperatorDialect(pir::IrContext* ctx)
       std::move(pir::InterfaceValue::Get<
                 InferSymbolicShapeInterface,
                 ShadowOutputOpInferSymbolicShapeInterfaceModel>()));
+
+  info = ctx->GetRegisteredOpInfo(pir::YieldOp::name());
+  info.AttachInterface(std::move(
+      pir::InterfaceValue::Get<InferSymbolicShapeInterface,
+                               YieldOpInferSymbolicShapeInterfaceModel>()));
 }
 
 void PrintTypeImpl(pir::Type type, std::ostream& os) {
diff --git a/paddle/pir/include/dialect/control_flow/ir/cf_op.h b/paddle/pir/include/dialect/control_flow/ir/cf_op.h
@@ -14,16 +14,13 @@
 
 #pragma once
 #include <functional>
-#include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h"
 #include "paddle/pir/include/core/builder.h"
 #include "paddle/pir/include/core/op_base.h"
 #include "paddle/pir/include/core/op_trait.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_interface.h"
 
 namespace pir {
-class IR_API YieldOp : public Op<YieldOp,
-                                 SideEffectTrait,
-                                 paddle::dialect::InferSymbolicShapeInterface> {
+class IR_API YieldOp : public Op<YieldOp, SideEffectTrait> {
  public:
   using Op::Op;
   static const char *name() { return "cf.yield"; }
@@ -34,8 +31,6 @@ class IR_API YieldOp : public Op<YieldOp,
                     OperationArgument &argument,  // NOLINT
                     const std::vector<Value> &Value);
   void VerifySig() {}
-
-  bool InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis);
 };
 
 ///
diff --git a/paddle/pir/src/dialect/control_flow/ir/cf_op.cc b/paddle/pir/src/dialect/control_flow/ir/cf_op.cc
@@ -25,13 +25,6 @@ void YieldOp::Build(Builder &builder,
   argument.AddInputs(inputs);
 }
 
-bool YieldOp::InferSymbolicShape(
-    pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  VLOG(3) << "############ YieldOp::InferSymbolicShape start...";
-  // YieldOp has no output, just return true
-  return true;
-}
-
 void TuplePushOp::Build(Builder &builder,             // NOLINT
                         OperationArgument &argument,  // NOLINT
                         Value inlet,
