Add mechanism for overriding with corresponding unittests.

tenstorrent · Nov 29, 2024 · 47e7e40 · 47e7e40
1 parent 40f1737
commit 47e7e40
Show file tree

Hide file tree

Showing 6 changed files with 326 additions and 444 deletions.
diff --git a/include/ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h b/include/ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h
@@ -27,18 +27,23 @@ struct MemoryLayoutAnalysisPolicyTypeParser
     return false;
   }
 
-  static void print(llvm::raw_ostream &os,
-                    const MemoryLayoutAnalysisPolicyType &value) {
-    llvm::StringRef policy;
+  static std::string toString(const MemoryLayoutAnalysisPolicyType &value) {
+    std::string res;
     switch (value) {
     case MemoryLayoutAnalysisPolicyType::DFSharding:
-      policy = "DFSharding";
+      res += "DFSharding";
       break;
     case MemoryLayoutAnalysisPolicyType::L1Interleaved:
-      policy = "L1Interleaved";
+      res += "L1Interleaved";
       break;
     }
-    os << "memory-layout-analysis-policy=" << policy << "\n";
+    return res;
+  }
+
+  static void print(llvm::raw_ostream &os,
+                    const MemoryLayoutAnalysisPolicyType &value) {
+    os << "memory-layout-analysis-policy="
+       << MemoryLayoutAnalysisPolicyTypeParser::toString(value) << "\n";
   }
 };
 

diff --git a/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h b/include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h
@@ -2,11 +2,13 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#pragma once
+
 #ifndef TTMLIR_DIALECT_TTNN_PIPELINES_TTNNPIPELINES_H
 #define TTMLIR_DIALECT_TTNN_PIPELINES_TTNNPIPELINES_H
 
 #include "ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h"
-#include "ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h"
+#include "ttmlir/Dialect/TTNN/Utils/Utils.h"
 
 #include "mlir/Pass/PassOptions.h"
 

diff --git a/include/ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h b/include/ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h
@@ -2,84 +2,20 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#pragma once
+
 #ifndef TTMLIR_DIALECT_TTNN_UTILS_OPTIMIZEROVERRIDES_H
 #define TTMLIR_DIALECT_TTNN_UTILS_OPTIMIZEROVERRIDES_H
 
-#include <llvm/Support/CommandLine.h>
-
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h"
+
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
+#include "ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h"
+#include "ttmlir/Dialect/TTNN/Utils/Utils.h"
 
 namespace mlir::tt::ttnn {
 
-struct OutputLayoutOverrideParams {
-
-  SmallVector<int64_t, 2> grid;
-  BufferType bufferType;
-  TensorMemoryLayout tensorMemoryLayout; // INTERLEAVED / SHARDED etc...
-  Layout memoryLayout;                   // ROW_MAJOR / TILE
-  mlir::tt::DataType dataType;
-
-  bool operator==(const OutputLayoutOverrideParams rhs) const {
-    return grid[0] == rhs.grid[0] && grid[1] == rhs.grid[1] &&
-           bufferType == rhs.bufferType &&
-           tensorMemoryLayout == rhs.tensorMemoryLayout &&
-           memoryLayout == rhs.memoryLayout && dataType == rhs.dataType;
-  }
-
-  bool operator!=(const OutputLayoutOverrideParams &rhs) const {
-    return !(*this == rhs);
-  }
-};
-
-struct InputLayoutOverrideParams {
-
-  SmallVector<int64_t> operandIdxes;
-
-  bool operator==(const InputLayoutOverrideParams &rhs) const {
-    if (operandIdxes.size() != rhs.operandIdxes.size()) {
-      return false;
-    }
-    for (std::size_t i = 0; i < operandIdxes.size(); i++) {
-      if (operandIdxes[i] != rhs.operandIdxes[i]) {
-        return false;
-      }
-    }
-    return true;
-  }
-
-  bool operator!=(const InputLayoutOverrideParams &rhs) const {
-    return !(*this == rhs);
-  }
-};
-
-struct OutputLayoutOverrideParser
-    : public llvm::cl::parser<llvm::StringMap<OutputLayoutOverrideParams>> {
-public:
-  OutputLayoutOverrideParser(llvm::cl::Option &opt)
-      : llvm::cl::parser<llvm::StringMap<OutputLayoutOverrideParams>>(opt) {}
-
-  bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
-             llvm::StringMap<OutputLayoutOverrideParams> &value);
-
-  static void print(llvm::raw_ostream &os,
-                    const llvm::StringMap<OutputLayoutOverrideParams> &value);
-};
-
-struct InputLayoutOverrideParser
-    : public llvm::cl::parser<llvm::StringMap<InputLayoutOverrideParams>> {
-public:
-  InputLayoutOverrideParser(llvm::cl::Option &opt)
-      : llvm::cl::parser<llvm::StringMap<InputLayoutOverrideParams>>(opt) {}
-
-  bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
-             llvm::StringMap<InputLayoutOverrideParams> &value);
-
-  static void print(llvm::raw_ostream &os,
-                    const llvm::StringMap<InputLayoutOverrideParams> &value);
-};
-
 class OptimizerOverridesHandler {
 public:
   OptimizerOverridesHandler() {};
@@ -135,6 +71,10 @@ class OptimizerOverridesHandler {
                                tt::DataType);
 
 private:
+  // Options for the TTIR to TTNN backend pipeline,
+  // we use them to extract the names and the deafulat values.
+  TTIRToTTNNBackendPipelineOptions pipelineOptions;
+
   // Flags for enabling/disabling the optimizer passes
   bool enableOptimizer = true;
 

diff --git a/include/ttmlir/Dialect/TTNN/Utils/Utils.h b/include/ttmlir/Dialect/TTNN/Utils/Utils.h
@@ -5,10 +5,89 @@
 #ifndef TTMLIR_DIALECT_TTNN_UTILS_UTILS_H
 #define TTMLIR_DIALECT_TTNN_UTILS_UTILS_H
 
+#include <llvm/Support/CommandLine.h>
+
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOpsTypes.h"
 
+namespace mlir::tt::ttnn {
+
+struct OutputLayoutOverrideParams {
+
+  SmallVector<int64_t, 2> grid;
+  BufferType bufferType;
+  TensorMemoryLayout tensorMemoryLayout; // INTERLEAVED / SHARDED etc...
+  Layout memoryLayout;                   // ROW_MAJOR / TILE
+  mlir::tt::DataType dataType;
+
+  bool operator==(const OutputLayoutOverrideParams rhs) const {
+    return grid[0] == rhs.grid[0] && grid[1] == rhs.grid[1] &&
+           bufferType == rhs.bufferType &&
+           tensorMemoryLayout == rhs.tensorMemoryLayout &&
+           memoryLayout == rhs.memoryLayout && dataType == rhs.dataType;
+  }
+
+  bool operator!=(const OutputLayoutOverrideParams &rhs) const {
+    return !(*this == rhs);
+  }
+};
+
+struct InputLayoutOverrideParams {
+
+  SmallVector<int64_t> operandIdxes;
+
+  bool operator==(const InputLayoutOverrideParams &rhs) const {
+    if (operandIdxes.size() != rhs.operandIdxes.size()) {
+      return false;
+    }
+    for (std::size_t i = 0; i < operandIdxes.size(); i++) {
+      if (operandIdxes[i] != rhs.operandIdxes[i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  bool operator!=(const InputLayoutOverrideParams &rhs) const {
+    return !(*this == rhs);
+  }
+};
+
+struct OutputLayoutOverrideParser
+    : public llvm::cl::parser<llvm::StringMap<OutputLayoutOverrideParams>> {
+public:
+  OutputLayoutOverrideParser(llvm::cl::Option &opt)
+      : llvm::cl::parser<llvm::StringMap<OutputLayoutOverrideParams>>(opt) {}
+
+  bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
+             llvm::StringMap<OutputLayoutOverrideParams> &value);
+
+  static std::string
+  toString(const llvm::StringMap<OutputLayoutOverrideParams> &);
+
+  static void print(llvm::raw_ostream &os,
+                    const llvm::StringMap<OutputLayoutOverrideParams> &value);
+};
+
+struct InputLayoutOverrideParser
+    : public llvm::cl::parser<llvm::StringMap<InputLayoutOverrideParams>> {
+public:
+  InputLayoutOverrideParser(llvm::cl::Option &opt)
+      : llvm::cl::parser<llvm::StringMap<InputLayoutOverrideParams>>(opt) {}
+
+  bool parse(llvm::cl::Option &opt, StringRef argName, StringRef arg,
+             llvm::StringMap<InputLayoutOverrideParams> &value);
+
+  static std::string
+  toString(const llvm::StringMap<InputLayoutOverrideParams> &);
+
+  static void print(llvm::raw_ostream &os,
+                    const llvm::StringMap<InputLayoutOverrideParams> &value);
+};
+
+} // namespace mlir::tt::ttnn
+
 namespace mlir::tt::ttnn::utils {
 
 // Map tt::MemorySpace to ttnn::BufferType