【Infer Symbolic Shape BUAA No.70】Add lerp op (#66837)

Marcusryz · web-flow · commit 5914bd9e8dea · 2024-08-05T10:48:45.000+08:00
* Add lerp CINN.

* update
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
@@ -846,6 +846,57 @@ bool GroupNormOpInferSymbolicShape(
   return true;
 }
 
+bool LerpOpInferSymbolicShape(pir::Operation *op,
+                              pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  const auto &y_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(1));
+  const auto &w_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(2));
+  const auto &x_shape = x_shape_or_data.shape();
+  const auto &y_shape = y_shape_or_data.shape();
+  const auto &w_shape = w_shape_or_data.shape();
+  size_t x_ndims = x_shape.size();
+  size_t y_ndims = y_shape.size();
+  size_t w_ndims = w_shape.size();
+  std::vector<symbol::DimExpr> out1_shape;
+  std::vector<symbol::DimExpr> out2_shape;
+  if (x_ndims > y_ndims) {
+    out1_shape.assign(x_shape.begin(), x_shape.end());
+  } else if (x_ndims < y_ndims) {
+    out1_shape.assign(y_shape.begin(), y_shape.end());
+  } else {
+    symbol::DimExprBuilder builder;
+    for (size_t i = 0; i < x_ndims; ++i) {
+      out1_shape.emplace_back(builder.Broadcast(x_shape[i], y_shape[i]));
+      infer_context->AddBroadcastableCstr(x_shape[i], y_shape[i]);
+    }
+  }
+  size_t out1_ndims = out1_shape.size();
+  if (w_ndims > out1_ndims) {
+    out2_shape.assign(w_shape.begin(), w_shape.end());
+  } else if (w_ndims < out1_ndims) {
+    out2_shape.assign(out1_shape.begin(), out1_shape.end());
+  } else {
+    symbol::DimExprBuilder builder;
+    for (size_t i = 0; i < w_ndims; ++i) {
+      out2_shape.emplace_back(builder.Broadcast(w_shape[i], out1_shape[i]));
+      infer_context->AddBroadcastableCstr(w_shape[i], out1_shape[i]);
+    }
+  }
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out2_shape)});
+  return true;
+}
+
+bool Lerp_OpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  return LerpOpInferSymbolicShape(op, infer_context);
+}
+
 bool LayerNormOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   // Get the shapes of input tensors
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
@@ -34,6 +34,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(CrossEntropyWithSoftmax_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullWithTensor)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FlashAttn)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GroupNorm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lerp)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Lerp_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LayerNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Linspace)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(LinearInterp)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2678,6 +2678,7 @@
     func : lerp
   inplace : (x -> out)
   backward : lerp_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : lgamma
   args : (Tensor x)
