Added abstraction to MemoryLayoutAnalysis policy and removed unnecessary

tests.

include/ttmlir/Dialect/TTNN/Analysis/DFShardingPolicy.h

@@ -7,32 +7,25 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "ttmlir/Dialect/TTNN/Analysis/L1ChainConfig.h"
+#include "ttmlir/Dialect/TTNN/Analysis/MemoryLayoutAnalysisPolicy.h"
 
 namespace mlir::tt::ttnn {
 
 // Process ops in DFS schedulable order and build shard chain configs.
 // Schedule is also produced as a side effect of sharding.
 //
-class DFShardingPolicy {
-private:
-  Operation *rootOp;
-  std::vector<L1ChainConfig> *l1ChainConfigs;
-  llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>> legalLayouts;
-  llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> *schedule;
-  unsigned usableL1CacheSize = 0;
-
+class DFShardingPolicy : public MemoryLayoutAnalysisPolicy {
 public:
   DFShardingPolicy(
       Operation *rootOp, std::vector<L1ChainConfig> &l1ChainConfigs,
       const llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>>
           &legalLayouts,
       llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> &schedule,
       unsigned usableL1CacheSize)
-      : rootOp(rootOp), l1ChainConfigs(&l1ChainConfigs),
-        legalLayouts(legalLayouts), schedule(&schedule),
-        usableL1CacheSize(usableL1CacheSize) {}
+      : MemoryLayoutAnalysisPolicy(rootOp, l1ChainConfigs, legalLayouts,
+                                   schedule, usableL1CacheSize) {}
 
-  void run(const std::unordered_set<Edge> &overrideReshardEdges);
+  void run(const std::unordered_set<Edge> &overrideReshardEdges) final;
 };
 
 } // namespace mlir::tt::ttnn

include/ttmlir/Dialect/TTNN/Analysis/L1InterleavedPolicy.h

@@ -7,29 +7,22 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "ttmlir/Dialect/TTNN/Analysis/L1ChainConfig.h"
+#include "ttmlir/Dialect/TTNN/Analysis/MemoryLayoutAnalysisPolicy.h"
 
 namespace mlir::tt::ttnn {
 
-class L1InterleavedPolicy {
-private:
-  Operation *rootOp;
-  std::vector<L1ChainConfig> *l1ChainConfigs;
-  llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>> legalLayouts;
-  llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> *schedule;
-  unsigned usableL1CacheSize = 0;
-
+class L1InterleavedPolicy : public MemoryLayoutAnalysisPolicy {
 public:
   L1InterleavedPolicy(
       Operation *rootOp, std::vector<L1ChainConfig> &l1ChainConfigs,
       const llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>>
           &legalLayouts,
       llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> &schedule,
       unsigned usableL1CacheSize)
-      : rootOp(rootOp), l1ChainConfigs(&l1ChainConfigs),
-        legalLayouts(legalLayouts), schedule(&schedule),
-        usableL1CacheSize(usableL1CacheSize) {}
+      : MemoryLayoutAnalysisPolicy(rootOp, l1ChainConfigs, legalLayouts,
+                                   schedule, usableL1CacheSize) {}
 
-  void run();
+  void run(const std::unordered_set<Edge> &overrideReshardEdges) final;
 };
 
 } // namespace mlir::tt::ttnn

include/ttmlir/Dialect/TTNN/Analysis/MemoryLayoutAnalysisPolicy.h

@@ -0,0 +1,39 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_DIALECT_TTNN_ANALYSIS_MEMORYLAYOUTANALYSISPOLICY_H
+#define TTMLIR_DIALECT_TTNN_ANALYSIS_MEMORYLAYOUTANALYSISPOLICY_H
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "ttmlir/Dialect/TTNN/Analysis/L1ChainConfig.h"
+
+namespace mlir::tt::ttnn {
+
+class MemoryLayoutAnalysisPolicy {
+protected:
+  Operation *rootOp;
+  std::vector<L1ChainConfig> *l1ChainConfigs;
+  llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>> legalLayouts;
+  llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> *schedule;
+  unsigned usableL1CacheSize = 0;
+
+public:
+  virtual ~MemoryLayoutAnalysisPolicy() {};
+
+  MemoryLayoutAnalysisPolicy(
+      Operation *rootOp, std::vector<L1ChainConfig> &l1ChainConfigs,
+      const llvm::DenseMap<Operation *, std::vector<tt::LayoutAttr>>
+          &legalLayouts,
+      llvm::DenseMap<func::FuncOp, llvm::SmallVector<Operation *>> &schedule,
+      unsigned usableL1CacheSize)
+      : rootOp(rootOp), l1ChainConfigs(&l1ChainConfigs),
+        legalLayouts(legalLayouts), schedule(&schedule),
+        usableL1CacheSize(usableL1CacheSize) {}
+
+  virtual void run(const std::unordered_set<Edge> &overrideReshardEdges) = 0;
+};
+
+} // namespace mlir::tt::ttnn
+
+#endif // TTMLIR_DIALECT_TTNN_ANALYSIS_MEMORYLAYOUTANALYSISPOLICY_H

lib/Dialect/TTNN/Analysis/L1InterleavedPolicy.cpp

@@ -5,11 +5,11 @@
 #include "ttmlir/Dialect/TTNN/Analysis/L1InterleavedPolicy.h"
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Scheduler/Scheduler.h"
-#include <llvm/Support/raw_ostream.h>
 
 namespace mlir::tt::ttnn {
 
-void L1InterleavedPolicy::run() {
+void L1InterleavedPolicy::run(
+    const std::unordered_set<Edge> &overrideReshardEdges) {
   rootOp->walk([&](func::FuncOp func) {
     mlir::tt::scheduler::Scheduler scheduler(&func);
     llvm::SmallVector<mlir::Operation *> scheduleableOps;

lib/Dialect/TTNN/Analysis/MemoryLayoutAnalysis.cpp

@@ -68,7 +68,7 @@ void MemoryLayoutAnalysis::analysisImplementation() {
     L1InterleavedPolicy l1InterleavedPolicy(
         op, l1ChainConfigs, filterL1InterleavedOnly(analysisInput.legalLayouts),
         analysisResult.schedule, analysisInput.usableL1CacheSize);
-    l1InterleavedPolicy.run();
+    l1InterleavedPolicy.run(analysisInput.overrideReshardEdges);
     break;
   }
   }