[PIR][DynamicShape] Add InferSymbolicShapeInterface for prod (#61342)

Add InferSymbolicShapeInterface for prod

Author: lanxianghit

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.cc

@@ -389,6 +389,30 @@ bool SumOpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool ProdOpInferSymbolicShape(pir::Operation *op,
+                              pir::ShapeConstraintIRAnalysis *shape_analysis) {
+  auto attributes = op->attributes();
+  bool keepdim = attributes["keep_dim"].dyn_cast<pir::BoolAttribute>().data();
+
+  bool reduce_all =
+      attributes["reduce_all"].dyn_cast<pir::BoolAttribute>().data();
+
+  auto axis_gen_op = op->operand_source(1).defining_op();
+  if (axis_gen_op->isa<paddle::dialect::FullIntArrayOp>()) {
+    std::vector<int64_t> axis = GetVectorAttr(
+        axis_gen_op->dyn_cast<paddle::dialect::FullIntArrayOp>(), "value");
+    return ReduceInferDim(op, shape_analysis, axis, keepdim, reduce_all);
+  } else {
+    // TODO(lanxianghit): deal with other source: pir::VectorType,
+    // paddle::dialect::DenseTensorType
+    PADDLE_THROW(
+        phi::errors::Unimplemented("ProdOpInferSymbolicShape: 'axis' only "
+                                   "support FullIntArrayOp's result now."));
+  }
+
+  return true;
+}
+
 bool ReshapeOpInferSymbolicShape(
     pir::Operation *op, pir::ShapeConstraintIRAnalysis *shape_analysis) {
   pir::Value operand_source_shape = op->operand_source(1);

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape.h

@@ -182,6 +182,10 @@ bool ReluOpInferSymbolicShape(pir::Operation *op,
                               pir::ShapeConstraintIRAnalysis *shape_analysis);
 bool Relu_OpInferSymbolicShape(pir::Operation *op,
                                pir::ShapeConstraintIRAnalysis *shape_analysis);
+
+bool ProdOpInferSymbolicShape(pir::Operation *op,
+                              pir::ShapeConstraintIRAnalysis *shape_analysis);
+
 }  // namespace paddle::dialect
 
 namespace cinn::dialect {

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -1041,6 +1041,7 @@
   kernel :
     func : prod
   backward : prod_grad
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : push_sparse_v2
   args : (Tensor[] ids, Tensor[] w, Tensor[] out_grad_in, int embeddingdim = 11, int tableid = 0, str accessorclass = "", str ctrlabelname = "", int paddingid = 0, bool scalesparsegrad = true, str[] inputnames = {}, bool is_distributed = true)