[CINN] Register InferSymbolicShape for cinn_op.argmin/argmax

Author: lshpku

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

@@ -260,6 +260,16 @@ bool ReduceSumOpInferSymbolicShape(
   return ReduceInferSymbolicShape(op, infer_context);
 }
 
+bool ArgminOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  return ReduceInferSymbolicShape(op, infer_context);
+}
+
+bool ArgmaxOpInferSymbolicShape(pir::Operation *op,
+                                pir::InferSymbolicShapeContext *infer_context) {
+  return ReduceInferSymbolicShape(op, infer_context);
+}
+
 bool ReshapeOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   std::vector<int> shape =

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

@@ -23,6 +23,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceMax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceMin)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceProd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReduceSum)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Argmin)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(Argmax)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Reshape)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Slice)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Split)