[CINN] Add argmin/argmax op in frontend (#71956)

lshpku · web-flow · commit d9ca4a870fb7 · 2025-04-01T10:57:35.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/ir/ops.yaml b/paddle/cinn/hlir/dialect/operator/ir/ops.yaml
@@ -1,3 +1,19 @@
+- op : argmax
+  args : (Tensor x, int64_t[] axis, bool keepdim=false, DataType dtype=DataType::INT64)
+  output : Tensor(out)
+  infer_meta :
+    func : ReduceSumInferMeta
+    param : [x, axis, keepdim, dtype]
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
+- op : argmin
+  args : (Tensor x, int64_t[] axis, bool keepdim=false, DataType dtype=DataType::INT64)
+  output : Tensor(out)
+  infer_meta :
+    func : ReduceSumInferMeta
+    param : [x, axis, keepdim, dtype]
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+
 - op : broadcast
   args : (Tensor x, int64_t[] broadcast_axes,  int64_t[] out_shape)
   output : Tensor(out)
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc b/paddle/cinn/hlir/dialect/operator/transforms/cinn_to_pd_util.cc
@@ -123,6 +123,37 @@ const auto& handler_reduce_prod_op =
   return pd_op;
 };
 
+template <typename TARGET_OP>
+::pir::Operation* ConvertArgMinMaxOp(
+    ::pir::Operation* op,
+    ::pir::IrMapping& ir_mapping,        // NOLINT
+    ::pir::PatternRewriter& rewriter) {  // NOLINT
+  VLOG(6) << "transform " << op->name() << " from cinn_op to pd_op";
+
+  auto attrs = op->attributes();
+  auto* ctx = ::pir::IrContext::Instance();
+
+  ::pir::ArrayAttribute attr_axis =
+      attrs.at("axis").dyn_cast<::pir::ArrayAttribute>();
+  if (attr_axis.empty()) {
+    attrs["axis"] = ::pir::Int64Attribute::get(ctx, 0);
+    attrs["flatten"] = ::pir::BoolAttribute::get(ctx, true);
+  } else {
+    attrs["axis"] = attr_axis.at(0);
+    attrs["flatten"] = ::pir::BoolAttribute::get(ctx, false);
+  }
+
+  attrs["keepdims"] = attrs.at("keepdim");
+  attrs.erase("keepdim");
+
+  auto pd_op = rewriter.Build<TARGET_OP>(
+      ir_mapping.Lookup(op->operand_source(0)), attrs);
+  for (uint32_t i = 0; i < op->num_results(); ++i) {
+    ir_mapping.Add(op->result(i), pd_op->result(i));
+  }
+  return pd_op;
+}
+
 ::pir::Operation* ConvertSliceOp(::pir::Operation* op,
                                  ::pir::IrMapping& ir_mapping,        // NOLINT
                                  ::pir::PatternRewriter& rewriter) {  // NOLINT
@@ -404,6 +435,16 @@ REGISTER_TRANSFORM_RULES(reduce_prod_op,
                          cinn::dialect::ReduceProdOp::name(),
                          cinn::dialect::details::handler_reduce_prod_op);
 
+REGISTER_TRANSFORM_RULES(
+    argmin_op,
+    cinn::dialect::ArgminOp::name(),
+    cinn::dialect::details::ConvertArgMinMaxOp<paddle::dialect::ArgminOp>);
+
+REGISTER_TRANSFORM_RULES(
+    argmax_op,
+    cinn::dialect::ArgmaxOp::name(),
+    cinn::dialect::details::ConvertArgMinMaxOp<paddle::dialect::ArgmaxOp>);
+
 REGISTER_TRANSFORM_RULES(slice_op,
                          cinn::dialect::SliceOp::name(),
                          cinn::dialect::details::ConvertSliceOp);
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/pd_to_cinn_pass.cc
@@ -170,6 +170,53 @@ class ReduceMinMaxOpPattern : public pir::OpRewritePattern<SOURCE_OP> {
   }
 };
 
+template <typename SOURCE_OP, typename TARGET_OP>
+class ArgMinMaxOpPattern : public pir::OpRewritePattern<SOURCE_OP> {
+ public:
+  using pir::OpRewritePattern<SOURCE_OP>::OpRewritePattern;
+
+  bool Match(SOURCE_OP op) const override {
+    const bool is_denied = CompatibleInfo::IsDeniedForCinn(*op.operation());
+    return !is_denied && IsDefinedBy<FullOp>(op, 1);
+  }
+
+  void Rewrite(SOURCE_OP op, pir::PatternRewriter &rewriter) const override {
+    const FullOp full_op = CastDefinedTo<FullOp>(op, 1);
+
+    const int64_t axis_value =
+        full_op.attribute("value")
+            .template dyn_cast<paddle::dialect::ScalarAttribute>()
+            .data()
+            .to<int64_t>();
+    const bool flatten = op.attribute("flatten")
+                             .template dyn_cast<::pir::BoolAttribute>()
+                             .data();
+    const bool keepdim = op.attribute("keepdims")
+                             .template dyn_cast<::pir::BoolAttribute>()
+                             .data();
+    const auto &dtype =
+        op.attribute("dtype")
+            .template dyn_cast<paddle::dialect::DataTypeAttribute>()
+            .data();
+
+    // The argmin/argmax has exactly one axis and is only effective when the
+    // `flatten` attr is false.
+    std::vector<int64_t> axis;
+    if (!flatten) {
+      axis = {axis_value};
+    }
+
+    auto cinn_op =
+        rewriter.Build<TARGET_OP>(op->operand_source(0), axis, keepdim, dtype);
+
+    rewriter.ReplaceAllUsesWith(op.result(0), cinn_op.result(0));
+    rewriter.EraseOp(op);
+    if (full_op->use_empty()) {
+      rewriter.EraseOp(full_op);
+    }
+  }
+};
+
 class ProdOpPattern : public pir::OpRewritePattern<paddle::dialect::ProdOp> {
  public:
   using pir::OpRewritePattern<paddle::dialect::ProdOp>::OpRewritePattern;
@@ -1328,6 +1375,12 @@ pir::RewritePatternSet PdOpToCinnOpPass::InitializePatterns(
                                cinn::dialect::ReduceMinOp>>(context);
   ps.Add<ReduceMinMaxOpPattern<paddle::dialect::MaxOp,
                                cinn::dialect::ReduceMaxOp>>(context);
+  ps.Add<
+      ArgMinMaxOpPattern<paddle::dialect::ArgminOp, cinn::dialect::ArgminOp>>(
+      context);
+  ps.Add<
+      ArgMinMaxOpPattern<paddle::dialect::ArgmaxOp, cinn::dialect::ArgmaxOp>>(
+      context);
   ps.Add<ProdOpPattern>(context);
   ps.Add<ReshapeOpPattern>(context);
   ps.Add<PowOpPattern>(context);
diff --git a/paddle/cinn/hlir/op/contrib/argmax.cc b/paddle/cinn/hlir/op/contrib/argmax.cc
@@ -107,7 +107,7 @@ std::shared_ptr<framework::OpStrategy> StrategyForArgmax(
     const framework::NodeAttr &attrs,
     const std::vector<Tensor> &inputs,
     const std::vector<Type> &out_type,
-    const std::vector<std::vector<int>> &output_shapes,
+    const std::vector<std::vector<ir::Dim>> &output_shapes,
     const Target &target) {
   int axis;
   bool keep_dims = false;
@@ -174,10 +174,10 @@ CINN_REGISTER_HELPER(argmax_ops) {
       .describe("This operator implements the op argmax.")
       .set_num_inputs(1)
       .set_num_outputs(1)
-      .set_attr<cinn::hlir::framework::StrategyFunction>(
-          "CINNStrategy", cinn::hlir::op::StrategyForArgmax)
+      .set_attr<cinn::hlir::framework::StrategyFunctionSymbolic>(
+          "CINNStrategySymbolic", cinn::hlir::op::StrategyForArgmax)
       .set_attr<cinn::hlir::framework::OpPatternKind>(
-          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
+          "OpPattern", cinn::hlir::framework::OpPatternKind::kReduction)
       .set_support_level(4);
 
   return true;
diff --git a/paddle/cinn/hlir/op/contrib/argmin.cc b/paddle/cinn/hlir/op/contrib/argmin.cc
@@ -103,7 +103,7 @@ std::shared_ptr<framework::OpStrategy> StrategyForArgmin(
     const framework::NodeAttr &attrs,
     const std::vector<Tensor> &inputs,
     const std::vector<Type> &out_type,
-    const std::vector<std::vector<int>> &output_shapes,
+    const std::vector<std::vector<ir::Dim>> &output_shapes,
     const Target &target) {
   int axis;
   bool keep_dims = false;
@@ -170,10 +170,10 @@ CINN_REGISTER_HELPER(argmin_ops) {
       .describe("This operator implements the op argmin.")
       .set_num_inputs(1)
       .set_num_outputs(1)
-      .set_attr<cinn::hlir::framework::StrategyFunction>(
-          "CINNStrategy", cinn::hlir::op::StrategyForArgmin)
+      .set_attr<cinn::hlir::framework::StrategyFunctionSymbolic>(
+          "CINNStrategySymbolic", cinn::hlir::op::StrategyForArgmin)
       .set_attr<cinn::hlir::framework::OpPatternKind>(
-          "OpPattern", cinn::hlir::framework::OpPatternKind::kNonFusible)
+          "OpPattern", cinn::hlir::framework::OpPatternKind::kReduction)
       .set_support_level(4);
 
   return true;
