[PIR] [DynamicShape] Add infer sym op for pd.conv3d pd.randint pd.assign_value pd.triu builtin.set_parameter && pd.split_with_num (#62955)

zhangboSJTU · web-flow · commit e504f06dae2f · 2024-03-25T18:49:18.000+08:00
* add conv3d &amp;&amp; randint

* add assign op

* add triu

* add split_with_num

* add built.set_parameter
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
@@ -150,6 +150,11 @@ bool Conv2dOpInferSymbolicShape(
   return true;
 }
 
+bool Conv3dOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return Conv2dOpInferSymbolicShape(op, shape_analysis);
+}
+
 bool EmbeddingOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto x_shape_or_data =
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h
@@ -19,6 +19,7 @@
 namespace paddle::dialect {
 
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv2d)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Conv3d)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Embedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(SparseWeightEmbedding)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ExpandAs)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc
@@ -72,11 +72,25 @@ bool ArangeOpInferSymbolicShape(
 
 bool AssignValueOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
+  const std::vector<int> shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "shape");
+  std::vector<symbol::DimExpr> sym_dims;
+  sym_dims.reserve(shape.size());
+  for (const int &dim : shape) {
+    sym_dims.emplace_back(symbol::DimExpr(static_cast<int64_t>(dim)));
+  }
+
+  symbol::ShapeOrDataDimExprs shape_data{
+      symbol::TensorShapeOrDataDimExprs(sym_dims)};
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
   return true;
 }
 
+bool AssignValue_OpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return AssignValueOpInferSymbolicShape(op, shape_analysis);
+}
+
 bool DataOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
   const auto &attributes = op->attributes();
@@ -248,17 +262,36 @@ bool GaussianOpInferSymbolicShape(
 
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
-        op->name() +
-        " 's InferSymbolicShape interface is NOT implemented now."));
+        "Currently shape must comes from FullIntArrayOp in GaussianOp's "
+        "InferSymbolicShape."));
     return true;
   }
 }
 
 bool RandintOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
+  const auto &shape_gen_op = op->operand_source(0).defining_op();
+
+  if (shape_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
+    std::vector<int64_t> shape = details::GetVectorAttr(
+        shape_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
+    std::vector<symbol::DimExpr> sym_dims;
+    sym_dims.reserve(shape.size());
+    for (const int64_t &dim : shape) {
+      sym_dims.emplace_back(symbol::DimExpr(dim));
+    }
+
+    symbol::ShapeOrDataDimExprs shape_data{
+        symbol::TensorShapeOrDataDimExprs(sym_dims)};
+    shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+    return true;
+
+  } else {
+    PADDLE_THROW(phi::errors::Unimplemented(
+        "Currently shape must comes from FullIntArrayOp in RandintOp's "
+        "InferSymbolicShape."));
+    return true;
+  }
 }
 
 bool TrilIndicesOpInferSymbolicShape(
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h
@@ -19,6 +19,7 @@
 namespace paddle::dialect {
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Arange)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AssignValue)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(AssignValue_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Data)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Empty)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Feed)
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
@@ -123,6 +123,8 @@ OP_SAME_OPERANDS_AND_RESULT(Tanh)
 OP_SAME_OPERANDS_AND_RESULT(Tanh_)
 OP_SAME_OPERANDS_AND_RESULT(Tril)
 OP_SAME_OPERANDS_AND_RESULT(Tril_)
+OP_SAME_OPERANDS_AND_RESULT(Triu)
+OP_SAME_OPERANDS_AND_RESULT(Triu_)
 OP_SAME_OPERANDS_AND_RESULT(Trunc)
 OP_SAME_OPERANDS_AND_RESULT(Trunc_)
 
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
@@ -115,6 +115,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tanh)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tanh_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tril)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tril_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Triu)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Triu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Trunc)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Trunc_)
 
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
@@ -634,8 +634,36 @@ bool SplitOpInferSymbolicShape(pir::Operation *op,
 
 bool SplitWithNumOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
+  int64_t axis = op->operand_source(1)
+                     .defining_op<paddle::dialect::FullOp>()
+                     .attributes()
+                     .at("value")
+                     .dyn_cast<paddle::dialect::ScalarAttribute>()
+                     .data()
+                     .to<int64_t>();
+  const auto &attributes = op->attributes();
+  int num = attributes.at("num").dyn_cast<pir::Int32Attribute>().data();
+  const auto &x_s_or_d =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
+  int rank = x_s_or_d.shape().size();
+  axis = axis < 0 ? axis + rank : axis;
+
+  symbol::DimExpr input_axis_dim = x_s_or_d.shape().at(axis);
+  symbol::DimExpr axis_shape = input_axis_dim / symbol::DimExpr{num};
+
+  const auto &out_s_d = [&] {
+    std::vector<symbol::DimExpr> out_s_d;
+    for (size_t i = 0; i < x_s_or_d.shape().size(); ++i) {
+      const auto &sym_dim =
+          axis == static_cast<int64_t>(i) ? axis_shape : x_s_or_d.shape()[i];
+      out_s_d.push_back(sym_dim);
+    }
+    return symbol::TensorShapeOrDataDimExprs(out_s_d);
+  }();
+
+  symbol::TensorListShapeOrDataDimExprs outs_s_d(num, out_s_d);
+  shape_analysis->SetShapeOrDataForValue(op->result(0),
+                                         symbol::ShapeOrDataDimExprs{outs_s_d});
   return true;
 }
 
@@ -783,18 +811,6 @@ bool Transpose_OpInferSymbolicShape(
   return TransposeOpInferSymbolicShape(op, shape_analysis);
 }
 
-bool TriuOpInferSymbolicShape(pir::Operation *op,
-                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " 's InferSymbolicShape interface is NOT implemented now."));
-  return true;
-}
-
-bool Triu_OpInferSymbolicShape(pir::Operation *op,
-                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  return TriuOpInferSymbolicShape(op, shape_analysis);
-}
-
 bool SqueezeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   IR_ENFORCE(op->num_operands() == 2,
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
@@ -53,8 +53,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Tile)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Topk)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Transpose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Transpose_)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Triu)
-OP_DECLARE_INFER_SYMBOLIC_SHAPE(Triu_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unbind)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Unique)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(UniqueConsecutive)
diff --git a/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc b/paddle/fluid/pir/dialect/operator/ir/op_dialect.cc
@@ -131,6 +131,17 @@ struct ParameterOpInferSymbolicShapeInterfaceModel
       : InferSymbolicShapeInterface::Concept(InferSymbolicShape) {}
 };
 
+struct SetParameterOpInferSymbolicShapeInterfaceModel
+    : public InferSymbolicShapeInterface::Concept {
+  static inline bool InferSymbolicShape(
+      pir::Operation* op, pir::ShapeConstraintIRAnalysis* shape_analysis) {
+    return true;
+  }
+
+  SetParameterOpInferSymbolicShapeInterfaceModel()
+      : InferSymbolicShapeInterface::Concept(InferSymbolicShape) {}
+};
+
 struct ShadowOutputOpInferSymbolicShapeInterfaceModel
     : public InferSymbolicShapeInterface::Concept {
   static inline bool InferSymbolicShape(
@@ -240,6 +251,16 @@ OperatorDialect::OperatorDialect(pir::IrContext* ctx)
   info.AttachInterface(
       pir::InterfaceValue::Get<InferSymbolicShapeInterface,
                                YieldOpInferSymbolicShapeInterfaceModel>());
+
+  info = ctx->GetRegisteredOpInfo(pir::SetParameterOp::name());
+  info.AttachInterface(pir::InterfaceValue::Get<
+                       InferSymbolicShapeInterface,
+                       SetParameterOpInferSymbolicShapeInterfaceModel>());
+
+  info = ctx->GetRegisteredOpInfo(pir::SliceOp::name());
+  info.AttachInterface(
+      pir::InterfaceValue::Get<InferSymbolicShapeInterface,
+                               SliceOpInferSymbolicShapeInterfaceModel>());
 }
 
 void PrintTypeImpl(pir::Type type, std::ostream& os) {
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -132,6 +132,7 @@
     param : [shape, dtype, values]
     data_type : dtype
     backend : place > output
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : barrier
   args : (Tensor x, int ring_id=0)
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -613,6 +613,7 @@
     func : conv3d
     data_type : input
   backward : conv3d_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : conv3d_transpose
   args : (Tensor x, Tensor filter, int[] strides={1, 1, 1}, int[] paddings={0, 0, 0}, int[] output_padding={}, int[] output_size={}, str padding_algorithm="EXPLICIT", int groups=1, int[] dilations={1, 1, 1}, str data_format="NCHW")
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_binary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_binary_op.py
@@ -200,5 +200,33 @@ def test_eval_symbolic(self):
         return True
 
 
+class Conv3dNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+        self.conv = paddle.nn.Conv3D(4, 6, (3, 3, 3))
+
+    def forward(self, x):
+        z = paddle.empty(shape=[2, 4, 8, 8, 8])
+        out = self.conv(z)
+        return out
+
+
+class Conv3dOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.expected = ['shape[2, 6, 6, 6, 6], data[NULL]']
+
+    def test_eval_symbolic(self):
+        net = Conv3dNet()
+
+        x_spec = InputSpec(shape=[None, None, None], dtype='float32')
+
+        input_spec = [x_spec]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+        check_infer_results(net, input_spec, 'pd_op.conv3d', self.expected)
+
+        return True
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_nullary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_nullary_op.py
@@ -14,6 +14,7 @@
 
 import unittest
 
+import numpy as np
 from test_infer_sym_shape_utils import (
     TestBase,
     apply_to_static,
@@ -62,6 +63,33 @@ def test_eval_symbolic(self):
         return out
 
 
+class AssignNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        data = paddle.empty(shape=[3, 3])
+        array = np.array([[1, 1], [3, 4], [1, 3]]).astype(np.int64)
+        out = paddle.assign(array, data)
+        return out
+
+
+class AssignOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.expected = ['shape[3, 2], data[NULL]']
+
+    def test_eval_symbolic(self):
+        net = AssignNet()
+        x_spec = InputSpec(shape=[None, None, 2], dtype='float32')
+        input_spec = [x_spec]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+        check_infer_results(
+            net, input_spec, 'pd_op.assign_value_', self.expected
+        )
+        return True
+
+
 class EmptyNet(paddle.nn.Layer):
     def __init__(self):
         super().__init__()
@@ -113,5 +141,28 @@ def test_eval_symbolic(self):
         return True
 
 
+class RandintNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        out = paddle.randint(low=-5, high=5, shape=[12, 32])
+        return out
+
+
+class RandintOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.expected = ['shape[12, 32], data[NULL]']
+
+    def test_eval_symbolic(self):
+        net = RandintNet()
+        x_spec = InputSpec(shape=[None, None, 2], dtype='float32')
+        input_spec = [x_spec]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+        check_infer_results(net, input_spec, 'pd_op.randint', self.expected)
+        return True
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/ir/pir/cinn/symbolic/test_infer_sym_shape_unary_op.py b/test/ir/pir/cinn/symbolic/test_infer_sym_shape_unary_op.py
@@ -481,5 +481,44 @@ def test_eval_symbolic(self):
         return True
 
 
+class SplitWithNumNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        data = paddle.empty(shape=[4, 6, 5])
+        out0, out1, out2 = paddle.split(data, num_or_sections=3, axis=1)
+        out0, out1, out2 = paddle.split(x, num_or_sections=3, axis=1)
+        return out0, out1, out2
+
+
+class SplitWithNumOpInferSymbolicShapeTest(TestBase):
+    def prepare_data(self):
+        self.cases = [np.random.rand(4, 6, 5)]
+        self.expected = [
+            "shape[4, 2, 5], data[NULL], shape[4, 2, 5], data[NULL], shape[4, 2, 5], data[NULL]",
+            "shape[S0, Mul(S1, 1 / (3)), S2], data[NULL], shape[S0, Mul(S1, 1 / (3)), S2], data[NULL], shape[S0, Mul(S1, 1 / (3)), S2], data[NULL]",
+        ]
+
+    def test_eval_symbolic(self):
+        net = SplitWithNumNet()
+
+        for i in range(len(self.cases)):
+            x = self.cases[i]
+            x_spec = InputSpec(
+                shape=[None for index in range(len(x.shape))], dtype='float32'
+            )
+            input_spec = [x_spec]
+            net = apply_to_static(net, False, input_spec)
+            net.eval()
+
+            # check the infer result
+            check_infer_results(
+                net, input_spec, 'pd_op.split_with_num', self.expected
+            )
+
+        return True
+
+
 if __name__ == '__main__':
     unittest.main()
