support printing name of cinn kernel (#71076)

zyfncg · web-flow · commit 430d69255c4d · 2025-02-11T10:41:25.000+08:00
diff --git a/paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc b/paddle/cinn/hlir/dialect/operator/ir/op_dialect.cc
@@ -87,7 +87,7 @@ void OperatorDialect::PrintAttribute(pir::Attribute attr,
   } else if (attr.isa<CINNKernelInfoAttribute>()) {
     auto cinn_kernel_info = attr.dyn_cast<CINNKernelInfoAttribute>();
 
-    os << "(" << cinn_kernel_info.data().fn_ptr;
+    os << "(" << cinn_kernel_info.data().fn_name;
     os << ')';
   } else if (attr.isa<FusionTrackerPtrAttribute>()) {
     auto tracker = attr.dyn_cast<FusionTrackerPtrAttribute>();
diff --git a/paddle/cinn/hlir/framework/pir/utils.cc b/paddle/cinn/hlir/framework/pir/utils.cc
@@ -657,8 +657,29 @@ std::string CompatibleInfo::OpName(const ::pir::Operation& op) {
   return OpNameAfterStripDialect(op);
 }
 
+std::string ShortenOpName(const std::string& name) {
+  static const std::unordered_map<std::string, std::string> OP_SHORT_NAMES = {
+      {"fill_constant", "full"},
+      {"reduce_sum", "sum"},
+      {"reduce_max", "r_max"},
+      {"reduce_min", "r_min"},
+      {"reduce_prod", "prod"},
+      {"elementwise_add", "add"},
+      {"elementwise_mul", "mul"},
+      {"subtract", "sub"},
+      {"divide", "div"},
+      {"broadcast_to", "bc"},
+      {"generate_shape", "gs"},
+      {"yield_store", "yield"},
+  };
+  if (OP_SHORT_NAMES.count(name)) {
+    return OP_SHORT_NAMES.at(name);
+  }
+  return name;
+}
+
 std::string CompatibleInfo::OpFuncName(const ::pir::Operation& op) {
-  std::string op_name = OpName(op);
+  std::string op_name = ShortenOpName(OpName(op));
   std::string func_name =
       cinn::common::Context::Global().NewName("fn_" + op_name);
   return func_name;
@@ -668,7 +689,7 @@ std::string CompatibleInfo::GroupOpsName(
     const std::vector<::pir::Operation*>& ops) {
   std::string name = "fn_";
   for (auto* op : ops) {
-    name += OpName(*op);
+    name += ShortenOpName(OpName(*op));
     name += "_";
   }
   return cinn::common::Context::Global().NewName(name);
