Merge pull request PaddlePaddle#55 from tc20042008/xk-cinn-trivalop-fuse

refine signature of group_pattern_util.ClusterIntoGroupPatternsFromOpList
feifei-111 · Mar 11, 2024 · c0dd054 · c0dd054
2 parents 5ff4943 + d3d6926
commit c0dd054
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/paddle/cinn/frontend/group_pattern_util.cc b/paddle/cinn/frontend/group_pattern_util.cc
@@ -852,12 +852,38 @@ GroupPattern FuseToGroupPattern(const std::vector<pir::Operation*>& ops) {
   return stmt_patterns;
 }
 
+class ClusteringHelper {
+ public:
+  ClusteringHelper(
+      const pir::ShapeConstraintIRAnalysis* shape_analysis,
+      const std::vector<pir::Operation*>& ops,
+      const OpsClusteringSpec& clustering_spec)
+    : shape_analysis_(shape_analysis), ops_(ops), clustering_spec_(clustering_spec) {
+    this->IsInThisOpList = MakePredicatorIsInThisFusionOp(ops);
+    this->IsInjectiveSource =
+        MakePredicatorIsInjectiveSource(ops_, this->IsInThisOpList);
+  }
+
+  std::vector<ConditionalGroupPattern> ClusterIntoGroupPatterns() {
+    LOG(FATAL) << "TODO(tianchao)";
+  }
+
+ private:
+  const pir::ShapeConstraintIRAnalysis* shape_analysis_;
+  const std::vector<pir::Operation*> ops_;
+  const OpsClusteringSpec clustering_spec_;
+  std::function<bool(const pir::Operation*)> IsInThisOpList;
+  std::function<bool(const pir::Operation*)> IsInjectiveSource;
+};
+
 }  // namespace
 
 std::vector<ConditionalGroupPattern> ClusterIntoGroupPatternsFromOpList(
+    const pir::ShapeConstraintIRAnalysis* shape_analysis,
     const std::vector<pir::Operation*>& ops,
-    const OpsClusteringSpec& clusteringSpec) {
-  // TODO();
+    const OpsClusteringSpec& clustering_spec) {
+  ClusteringHelper helper(shape_analysis, ops, clustering_spec);
+  return helper.ClusterIntoGroupPatterns();
 }
 
 GroupPattern GenerateGroupPatternFromOpList(

diff --git a/paddle/cinn/frontend/group_pattern_util.h b/paddle/cinn/frontend/group_pattern_util.h
@@ -28,8 +28,9 @@ struct OpsClusteringSpec {
 };
 
 std::vector<ConditionalGroupPattern> ClusterIntoGroupPatternsFromOpList(
+    const pir::ShapeConstraintIRAnalysis* shape_analysis,
     const std::vector<pir::Operation*>& ops,
-    const OpsClusteringSpec& clusteringSpec);
+    const OpsClusteringSpec& clustering_spec);
 
 GroupPattern GenerateGroupPatternFromOpList(
     const std::vector<pir::Operation*>& ops);