[CINN] 【Infer Symbolic Shape BUAA 】Add Frobenius_norm op (#68098)

crazyxiaoxi · Jeff114514 · web-flow · commit c52c3f188014 · 2024-09-14T15:30:26.000+08:00
* first try

* fix prodop

---------

Co-authored-by: Jeff114514 &lt;928430299@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.cc
@@ -1347,12 +1347,15 @@ bool Flatten_OpInferSymbolicShape(
   return FlattenOpInferSymbolicShape(op, infer_context);
 }
 
-// bool FrobeniusNormOpInferSymbolicShape(pir::Operation *op,
-//                                        pir::InferSymbolicShapeContext
-//                                        *infer_context) {
-//   // pass
-//   return true;
-// }
+bool FrobeniusNormOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  bool keepdim = op->attribute<pir::BoolAttribute>("keep_dim").data();
+  bool reduce_all = op->attribute<pir::BoolAttribute>("reduce_all").data();
+
+  const auto &axis = details::GetVectorAttr(op, "axis");
+
+  return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
+}
 
 bool FoldOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
@@ -2623,21 +2626,9 @@ bool ProdOpInferSymbolicShape(pir::Operation *op,
   bool keepdim = GetBoolAttr(op, "keepdim");
   bool reduce_all = GetBoolAttr(op, "reduce_all");
 
-  auto axis_gen_op = op->operand_source(1).defining_op();
-  if (axis_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
-    std::vector<int64_t> axis = details::GetVectorAttr(
-        axis_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
-    return details::ReduceInferDim(
-        op, infer_context, axis, keepdim, reduce_all);
-  } else {
-    // TODO(lanxianghit): deal with other source: pir::VectorType,
-    // paddle::dialect::DenseTensorType
-    PADDLE_THROW(
-        common::errors::Unimplemented("ProdOpInferSymbolicShape: 'axis' only "
-                                      "support FullIntArrayOp's result now."));
-  }
+  const auto &axis = details::GetVectorAttr(op, "axis");
 
-  return true;
+  return details::ReduceInferDim(op, infer_context, axis, keepdim, reduce_all);
 }
 
 bool QrOpInferSymbolicShape(pir::Operation *op,
diff --git a/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h b/paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/unary_infer_sym.h
@@ -73,7 +73,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Flatten_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Fold)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FakeQuantizeDequantizeAbsMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(FullBatchSizeLike)
-// OP_DECLARE_INFER_SYMBOLIC_SHAPE(FrobeniusNorm)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(FrobeniusNorm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Inverse)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(GumbelSoftmax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IdentityLoss)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -2027,7 +2027,7 @@
   kernel :
     func : frobenius_norm
   backward : frobenius_norm_grad
-  # interfaces : paddle::dialect::InferSymbolicShapeInterface
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : ftrl
   args: (Tensor param, Tensor squared_accumulator, Tensor linear_accumulator, Tensor grad, Tensor learning_rate, float l1=0.0f, float l2=0.0f, float lr_power=-0.5f)
