【BUAA】【Infer Symbolic Shape】Add TruncatedGaussianRandom, memcpy (#67490)

uanu2002 · web-flow · commit d83e25d22dcc · 2024-08-21T15:05:18.000+08:00
* add cinn TruncatedGaussianRandom

* add cinn memcpy

* fix
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
@@ -486,11 +486,22 @@ bool TriuIndicesOpInferSymbolicShape(
   return true;
 }
 
-// bool TruncatedGaussianRandomOpInferSymbolicShape(
-//     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
-//   // pass
-//   return true;
-// }
+bool TruncatedGaussianRandomOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const std::vector<int> shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "shape");
+  std::vector<symbol::DimExpr> out_shape;
+  out_shape.reserve(shape.size());
+  for (const int &dim : shape) {
+    out_shape.emplace_back(symbol::DimExpr(static_cast<int64_t>(dim)));
+  }
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{
+          symbol::TensorShapeOrDataDimExprs(out_shape)});
+  return true;
+}
 
 bool UniformOpInferSymbolicShape(
     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
@@ -33,6 +33,6 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Randperm)
 // OP_DECLARE_INFER_SYMBOLIC_SHAPE(RecvV2)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilIndices)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriuIndices)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(TruncatedGaussianRandom)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(TruncatedGaussianRandom)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Uniform)
 }  // namespace paddle::dialect
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.cc
@@ -116,6 +116,7 @@ OP_SAME_OPERANDS_AND_RESULT(Logit)
 OP_SAME_OPERANDS_AND_RESULT(Logit_)
 OP_SAME_OPERANDS_AND_RESULT(Logsigmoid)
 OP_SAME_OPERANDS_AND_RESULT(Logsigmoid_)
+OP_SAME_OPERANDS_AND_RESULT(Memcpy)
 OP_SAME_OPERANDS_AND_RESULT(Mish)
 OP_SAME_OPERANDS_AND_RESULT(Pow)
 OP_SAME_OPERANDS_AND_RESULT(Poisson)
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/same_operands_result.h
@@ -108,6 +108,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logit)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logit_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Logsigmoid_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Memcpy)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Mish)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Poisson)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Pow)
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -561,7 +561,7 @@
   kernel:
     func : memcpy
     param: [x, dst_place_type]
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : min
   args : (Tensor x, IntArray axis={}, bool keepdim=false)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -4875,7 +4875,7 @@
     backend : place
     data_type : dtype
   traits : pir::SideEffectTrait
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : unbind
   args : (Tensor input, int axis = 0)
