[Infer Symbolic Shape No.126] Add symbol_infer_interface for assign_ops (#68947)

ooooo-create · web-flow · commit aa857e5db763 · 2024-11-12T14:17:16.000+08:00
* Add symbol_infer_interface for assign_ops

* Revert open sym_infer check

* Open pir optest test &amp;&amp; refine interface

* Fix
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.cc
@@ -645,12 +645,24 @@ bool BilinearOpInferSymbolicShape(
   return true;
 }
 
-// bool AssignPosOpInferSymbolicShape(pir::Operation *op,
-//                                    pir::InferSymbolicShapeContext
-//                                    *infer_context) {
-//   // pass
-//   return true;
-// }
+bool AssignPosOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &eff_num_len_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  if (eff_num_len_shape_or_data.data()
+          .has_value()) {  // accoding to the kernel code
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+            {eff_num_len_shape_or_data.data()->at(0)})});
+  } else {
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(
+            {infer_context->GetNextSymName()})});
+  }
+  return true;
+}
 
 bool BroadcastTensorsOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/multiary_infer_sym.h
@@ -23,7 +23,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Addmm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Addmm_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(AddN)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Auc)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(AssignPos)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(AssignPos)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BroadcastTensors)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchFc)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchNorm)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -402,6 +402,7 @@
   kernel :
     func : assign_pos
   traits : paddle::dialect::ForwardOnlyTrait
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : assign_value_
   args : (Tensor output, int[] shape, DataType dtype, Scalar[] values, Place place = {})
diff --git a/test/legacy_test/test_assign_pos_op.py b/test/legacy_test/test_assign_pos_op.py
@@ -78,6 +78,7 @@ def setUp(self):
         y = count(x, 16)
         cum_count = np.cumsum(y).astype(x.dtype)
         self.op_type = "assign_pos"
+        self.python_api = utils._assign_pos
         self.inputs = {
             'X': x,
             "cum_count": cum_count,
@@ -89,7 +90,12 @@ def setUp(self):
     def test_forward(self):
         paddle.enable_static()
         np.testing.assert_allclose = get_redefined_allclose(self.cum_count)
-        self.check_output_with_place(paddle.CUDAPlace(0), check_dygraph=False)
+        self.check_output_with_place(
+            paddle.CUDAPlace(0),
+            check_dygraph=False,
+            check_pir=True,
+            check_symbol_infer=False,
+        )
 
 
 @unittest.skipIf(
