[Optimizer] Add ability to pass in custom constraint check function t…

…o ShardSolver. (#1573)
tenstorrent · Dec 13, 2024 · b0c0c2b · b0c0c2b
1 parent ed84a56
commit b0c0c2b
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 4 deletions.
diff --git a/include/ttmlir/Dialect/TTNN/Analysis/ShardSolver.h b/include/ttmlir/Dialect/TTNN/Analysis/ShardSolver.h
@@ -287,7 +287,10 @@ class ShardSolver {
               const std::vector<OpL1MemSpec> &shardSpecs,
               const llvm::DenseSet<Operation *> &shardedOps,
               const unsigned usableL1CacheSize,
-              const std::unordered_set<Edge> &overrideReshardEdges);
+              const std::unordered_set<Edge> &overrideReshardEdges,
+              std::function<bool(mlir::Operation *, TTNNLayoutAttr const &,
+                                 mlir::Operation *, TTNNLayoutAttr const &)>
+                  customCheckShardCompatible = nullptr);
   RemainingLayoutAttrs at(Operation *operation) const;
   void set(Operation *operation, TTNNLayoutAttr const &layout);
   static bool supportsInterleavedInputShardedOutput(Operation *op);
@@ -310,6 +313,9 @@ class ShardSolver {
 
   llvm::DenseMap<Operation *, TTNNLayoutAttr> selectedOpLayout;
   std::unordered_set<Edge> memReconfigEdges;
+  std::function<bool(mlir::Operation *, TTNNLayoutAttr const &,
+                     mlir::Operation *, TTNNLayoutAttr const &)>
+      customCheckShardCompatible;
 };
 
 } // namespace mlir::tt::ttnn

diff --git a/lib/Dialect/TTNN/Analysis/ShardSolver.cpp b/lib/Dialect/TTNN/Analysis/ShardSolver.cpp
@@ -19,10 +19,14 @@ ShardSolver::ShardSolver(
     const std::vector<OpL1MemSpec> &shardSpecs,
     const llvm::DenseSet<Operation *> &shardedOps,
     const unsigned usableL1CacheSize,
-    const std::unordered_set<Edge> &overrideReshardEdges)
+    const std::unordered_set<Edge> &overrideReshardEdges,
+    std::function<bool(Operation *, TTNNLayoutAttr const &, Operation *,
+                       TTNNLayoutAttr const &)>
+        customCheckShardCompatible)
     : legalLayouts(&legalLayouts), shardSpecs(&shardSpecs),
       shardedOps(&shardedOps), usableL1CacheSize(usableL1CacheSize),
-      memReconfigEdges(overrideReshardEdges) {
+      memReconfigEdges(overrideReshardEdges),
+      customCheckShardCompatible(customCheckShardCompatible) {
   pathSets.reserve(shardSpecs.size());
   pathSetIds.reserve(shardSpecs.size());
   bitsets.reserve(shardedOps.size());
@@ -505,6 +509,13 @@ bool ShardSolver::checkShardCompatible(
     Operation *producerOp, TTNNLayoutAttr const &producerLayout,
     Operation *consumerOp, TTNNLayoutAttr const &consumerLayout) const {
 
+  // Custom(test) hook for shard compatibility check.
+  //
+  if (customCheckShardCompatible) {
+    return customCheckShardCompatible(producerOp, producerLayout, consumerOp,
+                                      consumerLayout);
+  }
+
   // TEMP : Dummy mock implementation, will be replaced.
   //
 

diff --git a/test/unittests/Optimizer/TestShardSolver.cpp b/test/unittests/Optimizer/TestShardSolver.cpp
@@ -217,8 +217,26 @@ TEST_F(ShardSolverBase, VerifyProduceMaxCoreUsage) {
   addLayoutForOp(op, legalLayouts, BufferType::L1,
                  TensorMemoryLayout::BlockSharded, 1, 1);
 
+  // Create custom checkShardCompatible function.
+  //
+  std::function<bool(mlir::Operation *, TTNNLayoutAttr const &,
+                     mlir::Operation *, TTNNLayoutAttr const &)>
+      checkShardCompatible = [](mlir::Operation *producerOp,
+                                TTNNLayoutAttr const &producerLayout,
+                                mlir::Operation *consumerOp,
+                                TTNNLayoutAttr const &consumerLayout) {
+        // Simple shard compat assumption. Try to keep same shard layout.
+        //
+        if (producerLayout.getMemLayout() != consumerLayout.getMemLayout()) {
+          return false;
+        }
+
+        return true;
+      };
+
   ShardSolver shardSolver(legalLayouts, opL1MemSpecs, l1ChainedOps,
-                          usableL1CacheSize, overrideReshardEdges);
+                          usableL1CacheSize, overrideReshardEdges,
+                          checkShardCompatible);
 
   llvm::DenseMap<mlir::Operation *, llvm::SmallVector<float, 64>>
       accMaxCoreUsage = shardSolver.produceMaxCoreUsage();