【Infer Symbolic Shape No.125】【BUAA】Add BroadcastTensor, changed 3 files (#68180)

Whsjrczr · web-flow · commit 0056291512c4 · 2024-09-29T10:51:44.000+08:00
* add broadcast tensor

* rerun

* size_t -&gt; int

* add shape

* delete size_t

* size_t
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
@@ -652,12 +652,41 @@ bool BilinearOpInferSymbolicShape(
 //   return true;
 // }
 
-// bool BroadcastTensorsOpInferSymbolicShape(pir::Operation *op,
-//                                           pir::InferSymbolicShapeContext
-//                                           *infer_context) {
-//   // pass
-//   return true;
-// }
+bool BroadcastTensorsOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &input_shape_or_data_list =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0))
+          .dyn_cast<symbol::TensorListShapeOrDataDimExprs>();
+  // 1. Find Output rank = max(Inputs rank)
+  int target_rank = 0;
+  for (const auto &input_shape_or_data : input_shape_or_data_list) {
+    int tmp_rank = input_shape_or_data.shape().size();
+    target_rank = std::max(target_rank, tmp_rank);
+  }
+  // 2. Output dim(axis=x) = max(Inputs dim(axis=x))
+  std::vector<symbol::DimExpr> out_shape;
+  symbol::DimExprBuilder builder;
+  for (int i = 0; i < target_rank; i++) {
+    auto tmp_dim = symbol::DimExpr{1};
+    for (const auto &input_shape_or_data : input_shape_or_data_list) {
+      int axis = input_shape_or_data.shape().size();
+      axis = i - target_rank + axis;
+      if (axis >= 0) {
+        infer_context->AddBroadcastableCstr(input_shape_or_data.shape()[axis],
+                                            tmp_dim);
+        tmp_dim = builder.Broadcast(input_shape_or_data.shape()[axis], tmp_dim);
+      }
+    }
+    out_shape.emplace_back(tmp_dim);
+  }
+  symbol::TensorListShapeOrDataDimExprs out_shapes;
+  for (size_t i = 0; i < input_shape_or_data_list.size(); i++) {
+    out_shapes.emplace_back(out_shape);
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0), symbol::ShapeOrDataDimExprs{out_shapes});
+  return true;
+}
 
 bool BilinearInterpOpInferSymbolicShape(
     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
@@ -24,7 +24,7 @@ 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(BroadcastTensors)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(BroadcastTensors)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchFc)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(BatchNorm_)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -684,6 +684,7 @@
     func: broadcast_tensors
     data_type : input
   backward: broadcast_tensors_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : c_allgather
   args : (Tensor x, int ring_id, int nranks, bool use_calc_stream)
