[CINN]Add InferSymbolicShape for StackCreateOp (#68056)

Hongqing-work · web-flow · commit 0f5d5a6e9aab · 2024-09-11T18:45:46.000+08:00
* [CINN]Add InferSymbolicShape for StackCreateOp

* fix

* change pir include

* change to pir::InferSymbolicShapeInterface

* export infer_symbolic_shape for inference lib
diff --git a/cmake/inference_lib.cmake b/cmake/inference_lib.cmake
@@ -399,6 +399,11 @@ copy(
   SRCS ${PADDLE_SOURCE_DIR}/paddle/pir/include/dialect/shape/utils/*.h
   DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/pir/dialect/shape/utils/
 )
+copy(
+  inference_lib_dist
+  SRCS ${PADDLE_SOURCE_DIR}/paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/*.h
+  DSTS ${PADDLE_INFERENCE_INSTALL_DIR}/paddle/include/paddle/pir/dialect/shape/interface/infer_symbolic_shape/
+)
 copy(
   inference_lib_dist
   SRCS ${PADDLE_SOURCE_DIR}/paddle/pir/include/pass/*.h
diff --git a/paddle/pir/include/dialect/control_flow/ir/cf_op.h b/paddle/pir/include/dialect/control_flow/ir/cf_op.h
@@ -20,6 +20,7 @@
 #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"
+#include "paddle/pir/include/dialect/shape/interface/infer_symbolic_shape/infer_symbolic_shape.h"
 
 namespace pir {
 class IR_API YieldOp : public Op<YieldOp, SideEffectTrait> {
@@ -100,7 +101,9 @@ class IR_API TuplePopOp : public Op<TuplePopOp, SideEffectTrait> {
   TuplePushOp tuple_push_op() { return container_interface().tuple_push_op(); }
 };
 
-class IR_API StackCreateOp : public Op<StackCreateOp, ContainerOpInterface> {
+class IR_API StackCreateOp : public Op<StackCreateOp,
+                                       ContainerOpInterface,
+                                       pir::InferSymbolicShapeInterface> {
  public:
   using Op::Op;
   static const char *name() { return "cf.stack_create"; }
@@ -110,6 +113,8 @@ class IR_API StackCreateOp : public Op<StackCreateOp, ContainerOpInterface> {
                     OperationArgument &argument);  // NOLINT
   void VerifySig();
 
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
+
   Value container() { return result(0); }
   Value stack() { return result(0); }
   Value inlet() { return result(1); }
diff --git a/paddle/pir/src/dialect/control_flow/ir/cf_op.cc b/paddle/pir/src/dialect/control_flow/ir/cf_op.cc
@@ -200,6 +200,16 @@ void StackCreateOp::VerifySig() {
   VLOG(4) << "End Verifying for StackCreateOp.";
 }
 
+bool StackCreateOp::InferSymbolicShape(
+    pir::InferSymbolicShapeContext *infer_context) {
+  const auto &null_shape_or_data =
+      symbol::ShapeOrDataDimExprs(symbol::NullShapeOrDataDimExpr());
+  infer_context->SetShapeOrDataForValue(result(0), null_shape_or_data);
+  infer_context->SetShapeOrDataForValue(result(1), null_shape_or_data);
+  infer_context->SetShapeOrDataForValue(result(2), null_shape_or_data);
+  return true;
+}
+
 size_t StackCreateOp::tuple_size() { return tuple_push_op().tuple_size(); }
 
 Value StackCreateOp::inlet_element(size_t index) {
