[PIR][DynamicShape] Add InferSymbolicShape interfaces for llama ops (#61439)

lanxianghit · web-flow · commit 519f89c36962 · 2024-02-02T09:29:48.000+08:00
* Add InferSymbolicShape interfaces for llama ops
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc
@@ -984,6 +984,105 @@ bool Relu_OpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis) {
   return SameOperandsAndResultShape(op, shape_analysis);
 }
+
+bool ArangeOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool EmbeddingOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  const auto x_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(0));
+  const auto weight_shape_or_data =
+      shape_analysis->GetShapeOrDataForValue(op->operand_source(1));
+  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> &weight_dims = [&] {
+    std::vector<symbol::DimExpr> dims;
+    if (weight_shape_or_data.data().has_value()) {
+      dims = weight_shape_or_data.data().value();
+    } else {
+      dims = weight_shape_or_data.shape();
+    }
+    return dims;
+  }();
+
+  const symbol::ShapeOrDataDimExprs &shape_data = [&] {
+    std::vector<symbol::DimExpr> out_dims = x_dims;
+    // no need to check validation of weight_dims index, since all checks have
+    // been done at corresponding InferMeta
+    out_dims.emplace_back(weight_dims[1]);
+    return symbol::ShapeOrDataDimExprs{
+        symbol::TensorShapeOrDataDimExprs(out_dims)};
+  }();
+
+  shape_analysis->SetShapeOrDataForValue(op->result(0), shape_data);
+
+  return true;
+}
+
+bool SparseWeightEmbeddingOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool ExpandOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool MatmulOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool MaxOpInferSymbolicShape(pir::Operation *op,
+                             pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool TrilOpInferSymbolicShape(pir::Operation *op,
+                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool Tril_OpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return TrilOpInferSymbolicShape(op, shape_analysis);
+}
+
+bool WhereOpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  PADDLE_THROW(phi::errors::Unimplemented(
+      op->name() + " DOES NOT have InferSymbolicShapeInterface!"));
+  return true;
+}
+
+bool Where_OpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  return WhereOpInferSymbolicShape(op, shape_analysis);
+}
 }  // namespace paddle::dialect
 namespace cinn::dialect {
 
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h
@@ -186,6 +186,39 @@ bool Relu_OpInferSymbolicShape(pir::Operation *op,
 bool ProdOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis);
 
+bool ArangeOpInferSymbolicShape(pir::Operation *op,
+                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool EmbeddingOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool SparseWeightEmbeddingOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool ExpandOpInferSymbolicShape(pir::Operation *op,
+                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool MatmulOpInferSymbolicShape(pir::Operation *op,
+                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool MaxOpInferSymbolicShape(pir::Operation *op,
+                             pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool TransposeOpInferSymbolicShape(
+    pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool TrilOpInferSymbolicShape(pir::Operation *op,
+                              pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool Tril_OpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool WhereOpInferSymbolicShape(pir::Operation *op,
+                               pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool Where_OpInferSymbolicShape(pir::Operation *op,
+                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 }  // namespace paddle::dialect
 
 namespace cinn::dialect {
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -457,6 +457,7 @@
     param : [x, weight, padding_idx]
     data_type : weight
   backward : embedding_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : empty
   args : (IntArray shape, DataType dtype=DataType::FLOAT32, Place place=CPUPlace())
@@ -840,6 +841,7 @@
   data_transform :
     support_trans_dtype : x, y
   backward : matmul_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : matrix_rank
   args : (Tensor x, float tol, bool use_default_tol=true, bool hermitian=false)
@@ -867,6 +869,7 @@
   kernel :
     func : max
   backward : max_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : maximum
   args : (Tensor x, Tensor y)
@@ -1426,6 +1429,7 @@
     func : tril
   inplace: (x -> out)
   backward : tril_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : tril_indices
   args : (int rows, int cols, int offset, DataType dtype, Place place={})
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -77,6 +77,7 @@
     backend : place
   data_transform :
     support_trans_dtype : start, end, step
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : assign
   args : (Tensor x)
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -909,6 +909,7 @@
     func : expand
     data_type : x
   backward : expand_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : expand_as
   args : (Tensor x, Tensor y, int[] target_shape = {})
@@ -3015,6 +3016,7 @@
     func : where
   inplace: (x -> out)
   backward : where_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : yolo_box
   args : (Tensor x, Tensor img_size, int[] anchors={}, int class_num = 1, float conf_thresh = 0.01, int downsample_ratio = 32, bool clip_bbox = true, float scale_x_y=1.0, bool iou_aware=false, float iou_aware_factor=0.5)
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
@@ -0,0 +1,106 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import paddle
+from paddle.static import InputSpec
+
+
+def get_sym_shape_str_for_op(net, input_spec, op_name):
+    forward_program = net.forward.get_concrete_program(*input_spec)[
+        1
+    ].infer_program.forward_program
+    all_sym_shape_str = []
+    for op in forward_program.global_block().ops:
+        if op.name() == op_name:
+            all_sym_shape_str.append(op.attrs()['sym_shape_str'])
+
+    return all_sym_shape_str
+
+
+def apply_to_static(net, use_cinn, input_spec=None):
+    build_strategy = paddle.static.BuildStrategy()
+    build_strategy.build_cinn_pass = use_cinn
+    return paddle.jit.to_static(
+        net,
+        input_spec=input_spec,
+        build_strategy=build_strategy,
+        full_graph=True,
+    )
+
+
+class TestBase(unittest.TestCase):
+    def setUp(self):
+        paddle.seed(2022)
+        self.prepare_data()
+
+    def prepare_data(self):
+        pass
+
+    def test_eval_symbolic(self):
+        pass
+
+
+class EmbeddingNet(paddle.nn.Layer):
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.embedding = paddle.nn.Embedding(
+            num_embeddings,
+            embedding_dim,
+            weight_attr=paddle.ParamAttr(
+                initializer=paddle.nn.initializer.XavierNormal()
+            ),
+        )
+
+    def forward(self, x):
+        out = self.embedding(x)
+
+        return out
+
+
+class TestEmbeddingOpInferSymbolicShape(TestBase):
+    def prepare_data(self):
+        self.x_shape = [1, 2048]
+        self.num_embeddings = 32000
+        self.embedding_dim = 768
+        self.x = paddle.randint(low=0, high=768, shape=self.x_shape)
+        self.expected_sym_shape = 'shape[S0, S1, 768], data[NULL]'
+
+    def test_eval_symbolic(self):
+        net = EmbeddingNet(self.num_embeddings, self.embedding_dim)
+        input_spec = [
+            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, 'builtin.shadow_output'
+        )
+        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.x)
+        return out
+
+
+if __name__ == '__main__':
+    unittest.main()
