[CINN] 【Infer Symbolic Shape BUAA 】Add recv_v2 op (#68447)

crazyxiaoxi · Jeff114514 · web-flow · commit 2512240fa152 · 2024-09-27T15:05:20.000+08:00
* recvv2

* fix

---------

Co-authored-by: Jeff114514 &lt;928430299@qq.com&gt;
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
@@ -428,12 +428,57 @@ bool ReadFileOpInferSymbolicShape(
   return true;
 }
 
-// bool RecvV2OpInferSymbolicShape(pir::Operation *op,
-//                                 pir::InferSymbolicShapeContext
-//                                 *infer_context) {
-//   // pass
-//   return true;
-// }
+bool RecvV2OpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  const int ring_id = op->attribute<pir::Int32Attribute>("ring_id").data();
+  const bool dynamic_shape =
+      op->attribute<pir::BoolAttribute>("dynamic_shape").data();
+  const int peer = op->attribute<pir::Int32Attribute>("peer").data();
+
+  PADDLE_ENFORCE_GE(
+      peer,
+      0,
+      common::errors::InvalidArgument(
+          "The peer (%d) for recv_v2 op must be non-negative.", peer));
+
+  PADDLE_ENFORCE_GE(
+      ring_id,
+      0,
+      common::errors::InvalidArgument(
+          "The ring_id (%d) for recv_v2 op must be non-negative.", ring_id));
+
+  const std::vector<int> out_shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "out_shape");
+  if (!dynamic_shape) {
+    PADDLE_ENFORCE_GE(out_shape.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The size of the output shape must be greater than 0 "
+                          "but the value given is %d.",
+                          out_shape.size()));
+
+    std::vector<symbol::DimExpr> output_shape;
+    for (size_t i = 0; i < out_shape.size(); ++i) {
+      PADDLE_ENFORCE_GE(out_shape[i],
+                        1,
+                        common::errors::InvalidArgument(
+                            "The shape attribute for recv_v2 must be set "
+                            "explicitly, but the %dth element is %d which "
+                            "is less than 1. Or dynamic_shape should be set to "
+                            "True for both send_v2 and recv_v2.",
+                            i,
+                            out_shape[i]));
+      output_shape.push_back(symbol::DimExpr(out_shape[i]));
+    }
+
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(output_shape)});
+  }
+
+  return true;
+}
 
 bool SeedOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
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
@@ -31,7 +31,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Randint)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Randperm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReadFile)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Seed)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(RecvV2)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(RecvV2)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilIndices)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriuIndices)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TruncatedGaussianRandom)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -789,7 +789,7 @@
     func : recv_v2
     param : [ring_id, dynamic_shape, peer, out_shape, dtype, use_calc_stream]
     data_type : dtype
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : remainder
   args : (Tensor x, Tensor y)
