Merge pull request PaddlePaddle#7 from Fridge003/cinn_tmp

fix
Fridge003 · Mar 20, 2024 · 6c64295 · 6c64295
2 parents 4a767ec + a3fb6ff
commit 6c64295
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/paddle/cinn/frontend/cluster_ops/clustering_engine.cc b/paddle/cinn/frontend/cluster_ops/clustering_engine.cc
@@ -95,6 +95,24 @@ common::BfsWalker<const StmtPattern*> ClusteringEngine::MakeAcyclicSameClusterBf
   return common::BfsWalker<const StmtPattern*>(VisitAcyclicClusterNext);
 }
 
+ShardableAxes4ValueT ClusteringEngine::MakeInferedShardableAxes4Value(
+    const std::vector<const StmtPattern*>& stmt_ptrs) {
+  const OpSetPtr ops = [&] {
+    auto ops = std::make_shared<OpSet>();
+    for (const auto* stmt_ptr : stmt_ptrs) {
+      VisitStmtOp(*stmt_ptr, [&](const auto* op) { ops->insert(op); });
+    }
+    return ops;
+  }();
+  auto value2shardable_axes = shardable_axes_inferer_.InferShardableAxes(ops);
+  return [map = std::move(value2shardable_axes)](
+              pir::Value value) -> std::optional<const ShardableAxes*> {
+    const auto& iter = map.find(value);
+    if (iter == map.end()) return std::nullopt;
+    return &iter->second;
+  };
+}
+
 IsAcyclicConnectedT ClusteringEngine::MakePredicatorIsAcyclicConnected(
     const common::TopoWalker<const StmtPattern*>& walker,
     const std::vector<StmtPattern>& stmt_patterns,

diff --git a/paddle/cinn/frontend/cluster_ops/clustering_engine.h b/paddle/cinn/frontend/cluster_ops/clustering_engine.h
@@ -52,14 +52,19 @@ class ClusteringEngine {
     }
     }
 
-  common::BfsWalker<const StmtPattern*> MakeAcyclicSameClusterBfsWalker(
-      const std::vector<StmtPattern>& stmt_patterns);
-
+  using ShardableAxes4ValueT =
+      std::function<std::optional<const ShardableAxes*>(pir::Value)>;
   using IsAcyclicConnectedT =
       std::function<bool(const StmtPattern* src, const StmtPattern* dst)>;
   using ClusterRoot4StmtT =
       std::function<const StmtPattern*(const StmtPattern*)>;
 
+  ShardableAxes4ValueT MakeInferedShardableAxes4Value(
+      const std::vector<const StmtPattern*>& stmt_ptrs);
+
+  common::BfsWalker<const StmtPattern*> MakeAcyclicSameClusterBfsWalker(
+      const std::vector<StmtPattern>& stmt_patterns);
+
   IsAcyclicConnectedT MakePredicatorIsAcyclicConnected(
       const common::TopoWalker<const StmtPattern*>& walker,
       const std::vector<StmtPattern>& stmt_patterns,