[PIR+CINN]Unify Name Mapper between Paddle/CINN Dialect and AST Backend (#59203)

Aurelius84 · web-flow · commit 8bc2a36a6860 · 2023-11-23T19:08:32.000+08:00
refine IsSupportCinn

fix UT
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass.cc
@@ -1395,7 +1395,8 @@ class GeneralFusionMergePassHelper {
       }
       // master node
       for (auto& node : consumer->master_ops) {
-        if (GetOpKind(node->name()) == OpPatternKind::kReduction) {
+        if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==
+            OpPatternKind::kReduction) {
           fused_group->master_ops.insert(node);
         }
       }
@@ -1474,7 +1475,8 @@ class GeneralFusionMergePassHelper {
            ++consumer) {
         ::pir::Operation* master_node = nullptr;
         for (auto& node : (*consumer)->master_ops) {
-          if (GetOpKind(node->name()) != OpPatternKind::kReduction) {
+          if (hlir::framework::pir::CompatibleInfo::OpKind(*node) !=
+              OpPatternKind::kReduction) {
             master_node = node;
             break;
           }
@@ -1609,7 +1611,8 @@ class GeneralFusionMergePassHelper {
       }
       // master nodes
       for (auto& node : producer->master_ops) {
-        if (GetOpKind(node->name()) == OpPatternKind::kReduction) {
+        if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==
+            OpPatternKind::kReduction) {
           fused_group->master_ops.insert(node);
         }
       }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass_utils.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_pass_utils.h
@@ -185,7 +185,8 @@ static bool ReduceFuseReduce1(const OpGroupPtr& first,
   // }
   std::unique_ptr<cinn::dialect::ir::OpNode> reducer_0 = nullptr;
   for (auto op : first.GetGroup()->CollectOps()) {
-    if (GetOpKind(op->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*op) ==
+        OpPatternKind::kReduction) {
       reducer_0.reset(new cinn::dialect::ir::OpNode(op));
       break;
     }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_util.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/group_with_group_merge_util.h
@@ -223,7 +223,8 @@ inline bool elementwise_fuse_reduce(const std::shared_ptr<ir::Group>& first,
   // if reduce using block_reduce, can't fuse producer.
   ::pir::Operation* reducer = nullptr;
   for (auto& node : second->master_ops) {
-    if (GetOpKind(node->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==
+        OpPatternKind::kReduction) {
       reducer = node;
       break;
     }
@@ -291,7 +292,8 @@ inline bool broadcast_fuse_reduce(const std::shared_ptr<ir::Group>& first,
   }
   ::pir::Operation* reducer = nullptr;
   for (auto& node : second->master_ops) {
-    if (GetOpKind(node->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==
+        OpPatternKind::kReduction) {
       reducer = node;
       break;
     }
@@ -339,7 +341,7 @@ inline bool horizontal_relation(const std::shared_ptr<ir::Group>& first,
                             OpPatternKind kind) {
     std::unordered_set<::pir::Operation*> selected;
     for (auto node : nodes) {
-      if (GetOpKind(node->name()) == kind) {
+      if (hlir::framework::pir::CompatibleInfo::OpKind(*node) == kind) {
         selected.insert(node);
       }
     }
@@ -425,7 +427,8 @@ inline bool reduce_fuse_broadcast(const std::shared_ptr<ir::Group>& first,
   // required that each consumer of type Broadcast meet the same shape after
   // broadcast as before reduce.
   for (auto& node_in_master : first->master_ops) {
-    if (GetOpKind(node_in_master->name()) != OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*node_in_master) !=
+        OpPatternKind::kReduction) {
       continue;
     }
     ::pir::Operation* reducer = node_in_master;
@@ -488,7 +491,8 @@ inline bool reduce_fuse_broadcast(const std::shared_ptr<ir::Group>& first,
             visited_set.insert(consumer);
             candidates.push(consumer);
           }
-          if (GetOpKind(consumer->name()) == OpPatternKind::kBroadcast &&
+          if (hlir::framework::pir::CompatibleInfo::OpKind(*consumer) ==
+                  OpPatternKind::kBroadcast &&
               second->OpSet().find(consumer) != second->OpSet().end()) {
             broadcasters.insert(consumer);
           }
@@ -552,7 +556,8 @@ inline bool reduce_fuse_reduce(const std::shared_ptr<ir::Group>& first,
   }
   ::pir::Operation* reducer_0 = nullptr;
   for (auto& reducer : first->master_ops) {
-    if (GetOpKind(reducer->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*reducer) ==
+        OpPatternKind::kReduction) {
       reducer_0 = reducer;
       break;
     }
@@ -561,7 +566,8 @@ inline bool reduce_fuse_reduce(const std::shared_ptr<ir::Group>& first,
 
   ::pir::Operation* reducer_1 = nullptr;
   for (auto& reducer : second->master_ops) {
-    if (GetOpKind(reducer->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*reducer) ==
+        OpPatternKind::kReduction) {
       reducer_1 = reducer;
       break;
     }
@@ -598,7 +604,8 @@ inline bool reduce_fuse_reduce(const std::shared_ptr<ir::Group>& first,
     auto shared_size = 0;
     for (auto& fusion_group : {first, second}) {
       for (auto* master : fusion_group->master_ops) {
-        if (GetOpKind(master->name()) == OpPatternKind::kReduction) {
+        if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
+            OpPatternKind::kReduction) {
           shared_size += GetSharedSize(master);
         }
       }
@@ -619,7 +626,8 @@ inline bool reduce_fuse_reduce(const std::shared_ptr<ir::Group>& first,
     auto shared_size = 0;
     for (auto& fusion_group : {first, second}) {
       for (auto* master : fusion_group->master_ops) {
-        if (GetOpKind(master->name()) == OpPatternKind::kReduction) {
+        if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
+            OpPatternKind::kReduction) {
           shared_size += GetSharedSize(master);
         }
       }
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_node.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_node.h
@@ -30,7 +30,7 @@ class OpNode {
       : node_(node), input_tensors_(node), output_tensors_(node) {}
 
   OpPatternKind kind() const {
-    auto kind = GetOpKind(node_->name());
+    auto kind = hlir::framework::pir::CompatibleInfo::OpKind(*node_);
     if (kind == OpPatternKind::kBroadcast) {
       // As binary op was defined as broadcast, actually it should be
       // element-wise.
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_pass.cc
@@ -33,45 +33,6 @@ namespace cinn {
 namespace dialect {
 namespace ir {
 
-std::unordered_map<std::string, OpPatternKind> OpKindMap = {
-    {"pd_op.add", OpPatternKind::kElementWise},
-    {"pd_op.subtract", OpPatternKind::kElementWise},
-    {"pd_op.multiply", OpPatternKind::kElementWise},
-    {"pd_op.divide", OpPatternKind::kElementWise},
-    {"pd_op.sqrt", OpPatternKind::kElementWise},
-    {"pd_op.rsqrt", OpPatternKind::kElementWise},
-    {"pd_op.full", OpPatternKind::kElementWise},
-    {"pd_op.relu", OpPatternKind::kElementWise},
-    {"pd_op.exp", OpPatternKind::kElementWise},
-    {"pd_op.sin", OpPatternKind::kElementWise},
-    {"pd_op.cos", OpPatternKind::kElementWise},
-    {"pd_op.pow", OpPatternKind::kElementWise},
-    {"pd_op.elementwise_pow", OpPatternKind::kElementWise},
-    {"pd_op.sum", OpPatternKind::kReduction},
-    {"cinn_op.reshape", OpPatternKind::kElementWise},
-    {"pd_op.cast", OpPatternKind::kElementWise},
-    {"pd_op.greater_than", OpPatternKind::kElementWise},
-    {"pd_op.greater_equal", OpPatternKind::kElementWise},
-    {"pd_op.transpose", OpPatternKind::kInjective},
-    {"pd_op.gather_nd", OpPatternKind::kInjective},
-    {"cinn_op.scale", OpPatternKind::kElementWise},
-    {"cinn_op.concat", OpPatternKind::kInjective},
-    {"cinn_op.slice", OpPatternKind::kInjective},
-    {"cinn_op.reduce_sum", OpPatternKind::kReduction},
-    {"cinn_op.reduce_max", OpPatternKind::kReduction},
-    {"cinn_op.broadcast", OpPatternKind::kBroadcast},
-    {"cinn_op.uniform_random", OpPatternKind::kElementWise}};
-
-OpPatternKind GetOpKind(const std::string& op_name) {
-  auto found_it = OpKindMap.find(op_name);
-  if (found_it == OpKindMap.end()) {
-    PADDLE_THROW(phi::errors::Unavailable(
-        "not support [%s] op yet in op kind map", op_name));
-  }
-
-  return found_it->second;
-}
-
 std::vector<pir::Operation*> GetProducerOpsReverseSort(
     pir::Operation* op,
     const std::unordered_map<pir::Operation*, size_t>& op2id) {
@@ -323,7 +284,8 @@ class OpFusionPassHelper {
         }
 
         // group type
-        group->op_pattern_kind = GetOpKind(op->name());
+        group->op_pattern_kind =
+            hlir::framework::pir::CompatibleInfo::OpKind(*op);
         // use current op as master op for schedule
         group->master_ops.insert(op);
 
@@ -389,7 +351,8 @@ class OpFusionPassHelper {
  private:
   void DoOpFusion() {
     for (auto consumer : ops_) {
-      auto consumer_kind = GetOpKind(consumer->name());
+      auto consumer_kind =
+          hlir::framework::pir::CompatibleInfo::OpKind(*consumer);
       // kNonFusible op can't fuse any other op.
       if (consumer_kind == OpPatternKind::kNonFusible) {
         continue;
@@ -418,7 +381,8 @@ class OpFusionPassHelper {
           continue;
         }
         // kNonFusible op can't fuse any other op.
-        auto producer_kind = GetOpKind(producer->name());
+        auto producer_kind =
+            hlir::framework::pir::CompatibleInfo::OpKind(*producer);
         if (producer_kind == OpPatternKind::kNonFusible) {
           continue;
         }
@@ -625,13 +589,17 @@ class OpFusionPassHelper {
   }
 
   bool CanFuse(::pir::Operation* producer, const ::pir::Operation* consumer) {
-    auto& relation = fusion_relation_map_[GetOpKind(producer->name())];
+    auto& relation =
+        fusion_relation_map_[hlir::framework::pir::CompatibleInfo::OpKind(
+            *producer)];
     // first step: check producer can be fused into consumer
-    if (relation.op_kind.count(GetOpKind(consumer->name()))) {
+    if (relation.op_kind.count(
+            hlir::framework::pir::CompatibleInfo::OpKind(*consumer))) {
       auto& consumer_group = fusion_groups_[consumer];
       // second step: check producer can be fused into consumer group
       VLOG(3) << "Call ConditionFunction, Producer Op Pattern : "
-              << GetOpKind(producer->name()) << " , Consumer Group Pattern : "
+              << hlir::framework::pir::CompatibleInfo::OpKind(*producer)
+              << " , Consumer Group Pattern : "
               << consumer_group->op_pattern_kind;
 
       return relation.fusion_op_kind[consumer_group->op_pattern_kind](
diff --git a/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h b/paddle/cinn/hlir/dialect/operator/transforms/group_merge/op_with_group_merge_util.h
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "paddle/cinn/hlir/framework/pir/group.h"
+#include "paddle/cinn/hlir/framework/pir/utils.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/pir/core/operation.h"
@@ -61,8 +62,6 @@ std::vector<T> GetVectorAttr(const ::pir::Operation* op,
   return vec_res;
 }
 
-OpPatternKind GetOpKind(const std::string& op_name);
-
 phi::DDim GetFirstInputShape(const ::pir::Operation* op);
 
 phi::DDim GetValueShape(const ::pir::Value& value);
@@ -114,7 +113,8 @@ inline bool reduce_fuse_reduce(::pir::Operation* producer,
                                const std::shared_ptr<Group>& consumer) {
   ::pir::Operation* reducer = NULL;
   for (auto* master : consumer->master_ops) {
-    if (GetOpKind(master->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
+        OpPatternKind::kReduction) {
       reducer = master;
       break;
     }
@@ -157,7 +157,8 @@ inline bool reduce_fuse_reduce(::pir::Operation* producer,
     if (input_shape_same || without_last_dim) {
       auto shared_size = GetSharedSize(producer);
       for (auto* master : consumer->master_ops) {
-        if (GetOpKind(master->name()) == OpPatternKind::kReduction) {
+        if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
+            OpPatternKind::kReduction) {
           shared_size += GetSharedSize(master);
         }
       }
@@ -207,7 +208,8 @@ inline bool is_horizontal_relation(::pir::Operation* producer,
   };
 
   for (auto op : consumer->ops_set) {
-    if (GetOpKind(op->name()) != consumer->op_pattern_kind) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*op) !=
+        consumer->op_pattern_kind) {
       continue;
     }
     if (check_depency(op)) {
@@ -228,7 +230,8 @@ inline bool horizontal_or_vertical_reduce_relation(
   // reducer op in fusion op.
   ::pir::Operation* reducer = NULL;
   for (auto* master : consumer->master_ops) {
-    if (GetOpKind(master->name()) == OpPatternKind::kReduction) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*master) ==
+        OpPatternKind::kReduction) {
       reducer = master;
       break;
     }
@@ -385,7 +388,8 @@ inline bool reduce_fuse_broadcast(::pir::Operation* producer,
       };
 
   for (auto op : consumer->ops_set) {
-    if (GetOpKind(op->name()) != OpPatternKind::kBroadcast) {
+    if (hlir::framework::pir::CompatibleInfo::OpKind(*op) !=
+        OpPatternKind::kBroadcast) {
       continue;
     }
 
diff --git a/paddle/cinn/hlir/framework/pir/utils.cc b/paddle/cinn/hlir/framework/pir/utils.cc
diff --git a/paddle/fluid/pir/transforms/build_cinn_pass.cc b/paddle/fluid/pir/transforms/build_cinn_pass.cc

Original file line number	Diff line number	Diff line change
`@@ -1395,7 +1395,8 @@ class GeneralFusionMergePassHelper {`
`1395`	`1395`	`}`
`1396`	`1396`	`// master node`
`1397`	`1397`	`for (auto& node : consumer->master_ops) {`
`1398`		`- if (GetOpKind(node->name()) == OpPatternKind::kReduction) {`
	`1398`	`+ if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==`
	`1399`	`+ OpPatternKind::kReduction) {`
`1399`	`1400`	`fused_group->master_ops.insert(node);`
`1400`	`1401`	`}`
`1401`	`1402`	`}`
`@@ -1474,7 +1475,8 @@ class GeneralFusionMergePassHelper {`
`1474`	`1475`	`++consumer) {`
`1475`	`1476`	`::pir::Operation* master_node = nullptr;`
`1476`	`1477`	`for (auto& node : (*consumer)->master_ops) {`
`1477`		`- if (GetOpKind(node->name()) != OpPatternKind::kReduction) {`
	`1478`	`+ if (hlir::framework::pir::CompatibleInfo::OpKind(*node) !=`
	`1479`	`+ OpPatternKind::kReduction) {`
`1478`	`1480`	`master_node = node;`
`1479`	`1481`	`break;`
`1480`	`1482`	`}`
`@@ -1609,7 +1611,8 @@ class GeneralFusionMergePassHelper {`
`1609`	`1611`	`}`
`1610`	`1612`	`// master nodes`
`1611`	`1613`	`for (auto& node : producer->master_ops) {`
`1612`		`- if (GetOpKind(node->name()) == OpPatternKind::kReduction) {`
	`1614`	`+ if (hlir::framework::pir::CompatibleInfo::OpKind(*node) ==`
	`1615`	`+ OpPatternKind::kReduction) {`
`1613`	`1616`	`fused_group->master_ops.insert(node);`
`1614`	`1617`	`}`
`1615`	`1618`	`}`