[PIR][DynamicShape] Add InferSymbolicShape for arange, expand (#61502)

lanxianghit · web-flow · commit df6f93cd3e21 · 2024-02-04T16:00:06.000+08:00
* Add InferSymbolicShape for arange, expand
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc
@@ -1001,8 +1001,54 @@ bool Relu_OpInferSymbolicShape(pir::Operation *op,
 
 bool ArangeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  PADDLE_THROW(phi::errors::Unimplemented(
-      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  const auto &start_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &end_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &step_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(2));
+
+  const auto start = [&] {
+    symbol::DimExpr expr;
+    if (start_shape_or_data.data().has_value()) {
+      expr = start_shape_or_data.data().value()[0];
+    } else {
+      expr = start_shape_or_data.shape()[0];
+    }
+    return expr;
+  }();
+
+  const auto end = [&] {
+    symbol::DimExpr expr;
+    if (end_shape_or_data.data().has_value()) {
+      expr = end_shape_or_data.data().value()[0];
+    } else {
+      expr = end_shape_or_data.shape()[0];
+    }
+    return expr;
+  }();
+
+  const auto step = [&] {
+    symbol::DimExpr expr;
+    if (step_shape_or_data.data().has_value()) {
+      expr = step_shape_or_data.data().value()[0];
+    } else {
+      expr = step_shape_or_data.shape()[0];
+    }
+    return expr;
+  }();
+
+  const symbol::ShapeOrDataDimExprs &shape_data = [&] {
+    std::vector<symbol::DimExpr> out_dims;
+    // TODO(lanxianghit, jiahy0825): here should be ceil((end - start) / step),
+    // but DimExpr doesn't support ceil and float now
+    out_dims.emplace_back((end - start) / step);
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(out_dims)};
+  }();
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+
   return true;
 }
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -3460,15 +3460,53 @@ void ExpandOp::Build(pir::Builder &builder,
 
 bool ExpandOp::InferSymbolicShape(
     pir::ShapeConstraintIRAnalysis *shape_analysis) {
-  const auto &data_shape = shape_analysis->GetShapeOrDataForValue(shape());
-  IR_ENFORCE(data_shape.data().has_value(),
-             "Value shape comes from ShapeOp, it must have data");
-  const auto &data = data_shape.data().value();
+  const auto x_shape_or_data = shape_analysis->GetShapeOrDataForValue(x());
+  const auto expand_shape_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(shape());
+
+  const std::vector<symbol::DimExpr> &x_dims = [&] {
+    std::vector<symbol::DimExpr> dims;
+    if (x_shape_or_data.data().has_value()) {
+      dims = x_shape_or_data.data().value();
+    } else {
+      dims = x_shape_or_data.shape();
+    }
+    return dims;
+  }();
+
+  const std::vector<symbol::DimExpr> &expand_shape = [&] {
+    std::vector<symbol::DimExpr> dims;
+    if (expand_shape_shape_or_data.data().has_value()) {
+      dims = expand_shape_shape_or_data.data().value();
+    } else {
+      dims = expand_shape_shape_or_data.shape();
+    }
+
+    if (dims.empty()) {
+      dims = std::vector<symbol::DimExpr>(x_dims.size(), -1);
+    }
+    return dims;
+  }();
+
+  std::vector<symbol::DimExpr> out_shape = expand_shape;
+  for (size_t i = 0; i < expand_shape.size(); i++) {
+    if (expand_shape[i] == -1) {  // copy the dim from x
+      // the shape is right aligned
+      int index = i - (expand_shape.size() - x_dims.size());
+      IR_ENFORCE(index >= 0,
+                 "in ExpandOpInferSymbolicShape, the dim to copy must >= 0, "
+                 "but got %d",
+                 index);
+
+      out_shape[i] = x_dims[index];
+    }
+  }
 
-  symbol::ShapeOrDataDimExprs shape_data{
-      symbol::TensorShapeOrDataDimExprs(data)};
+  shape_analysis->SetShapeOrDataForValue(
+      out(),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
 
-  shape_analysis->SetShapeOrDataForValue(out(), shape_data);
   return true;
 }
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/update_ops.yaml b/paddle/fluid/pir/dialect/operator/ir/update_ops.yaml
@@ -15,6 +15,7 @@
     data_type : dtype
     backend : place
   support_tensor : [start, end, step]
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : sequence_mask
   args: (Tensor x, Scalar(int) max_len, DataType out_dtype)
diff --git a/test/ir/pir/cinn/symbolic/test_op_infer_sym_shape.py b/test/ir/pir/cinn/symbolic/test_op_infer_sym_shape.py
@@ -20,7 +20,7 @@
 from paddle.static import InputSpec
 
 
-def get_sym_shape_str_for_op(net, input_spec, op_name):
+def get_sym_shape_str_for_op(net, input_spec, op_name='builtin.shadow_output'):
     forward_program = net.forward.get_concrete_program(*input_spec)[
         1
     ].infer_program.forward_program
@@ -102,5 +102,105 @@ def test_eval_symbolic(self):
         return out
 
 
+class ArangeNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, in_0, in_1, in_2):
+        if in_1 is None:
+            end = in_0
+            out = paddle.arange(end)
+        else:
+            start, end, step = in_0, in_1, in_2
+            out = paddle.arange(start, end, step)
+
+        return out
+
+
+class TestArangeOpInferSymbolicShape(TestBase):
+    def prepare_data(self):
+        self.start = paddle.full([1], 0)
+        self.end = paddle.full([1], 5)
+        self.step = paddle.full([1], 1)
+
+        self.expected_sym_shape = (
+            'shape[Mul(Add(S1, -S0), 1 / (S2))], data[NULL]'
+        )
+
+    def test_eval_symbolic(self):
+        net = ArangeNet()
+
+        input_spec = [
+            InputSpec(shape=[None], dtype='float32'),
+            InputSpec(shape=[None], dtype='float32'),
+            InputSpec(shape=[None], dtype='float32'),
+        ]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+
+        # check the infer result
+        sym_shape_str_list = get_sym_shape_str_for_op(net, input_spec)
+        np.testing.assert_equal(len(sym_shape_str_list), 1)
+        np.testing.assert_equal(
+            sym_shape_str_list[0].find(self.expected_sym_shape),
+            0,
+            'output shape is not expected!',
+        )
+        out = net(self.start, self.end, self.step)
+        return out
+
+
+class ExpandNet(paddle.nn.Layer):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        out = paddle.expand(x, [paddle.shape(y)[1], paddle.shape(y)[0]])
+        out = paddle.expand(x, [7, 5, paddle.shape(y)[0]])
+        out = paddle.expand(x, [7, -1, paddle.shape(y)[0]])
+        out = paddle.expand(x, [7, paddle.shape(y)[1], -1])
+
+        return out
+
+
+class TestExpandOpInferSymbolicShape(TestBase):
+    def prepare_data(self):
+        self.x = paddle.rand([1, 3], 'float32')
+        self.y = paddle.rand([3, 2], 'float32')
+
+        self.expected_sym_shapes = [
+            'shape[S3, S2], data[NULL]',
+            'shape[7, 5, S2], data[NULL]',
+            'shape[7, S0, S2], data[NULL]',
+            'shape[7, S3, S1], data[NULL]',
+        ]
+
+    def test_eval_symbolic(self):
+        net = ExpandNet()
+
+        input_spec = [
+            InputSpec(shape=[None, None], dtype='float32'),
+            InputSpec(shape=[None, None], dtype='float32'),
+        ]
+        net = apply_to_static(net, False, input_spec)
+        net.eval()
+
+        # check the infer result
+        sym_shape_str_list = get_sym_shape_str_for_op(
+            net, input_spec, 'pd_op.expand'
+        )
+        np.testing.assert_equal(
+            len(sym_shape_str_list), len(self.expected_sym_shapes)
+        )
+        for i in range(len(self.expected_sym_shapes)):
+            np.testing.assert_equal(
+                sym_shape_str_list[i].find(self.expected_sym_shapes[i]),
+                0,
+                'output shape is not expected!',
+            )
+        out = net(self.x, self.y)
+        return out
+
+
 if __name__ == '__main__':
     unittest.main()
