[Optimizer] Add df and layout overrides

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

@@ -6,8 +6,8 @@
 #define TTMLIR_DIALECT_TTNN_ANALYSIS_LEGALGRIDANALYSIS_H
 
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
-#include "ttmlir/Dialect/TT/Utils/OverrideParams.h"
 #include "ttmlir/Dialect/TTNN/Analysis/TTNNAnalysis.h"
+#include "ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h"
 #include "llvm/ADT/StringMap.h"
 
 namespace mlir::tt::ttnn {

include/ttmlir/Dialect/TTNN/Pipelines/TTNNPipelines.h

@@ -7,7 +7,7 @@
 
 #include "mlir/Pass/PassOptions.h"
 #include "ttmlir/Dialect/TT/Utils/MemoryLayoutAnalysisParams.h"
-#include "ttmlir/Dialect/TT/Utils/OverrideParams.h"
+#include "ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h"
 
 namespace mlir::tt::ttnn {
 
@@ -46,17 +46,19 @@ struct TTIRToTTNNBackendPipelineOptions
   // The format is a comma separated list of op names equal to the output layout
   // params separated by ":"
   //
-  // op_name=grid_size:memory_space:tensor_memory_layout
+  // op_name=grid_size:memory_space:tensor_memory_layout:memory_layout:data_type
   //
   // * grid_size=2x2
   // * memory_space: system, mmio, dram or l1
   // * tensor_memory_layout: none, interleaved, single_bank, height_sharded,
   //   width_sharded or block_sharded
+  // * memory_layout: row_major or tile
+  // * data_type: f32, f16, bf16, bfp_f8, bfp_bf8, bfp_f4, bfp_bf4, bfp_f2,
+  //   bfp_bf2, u32, u16, u8
   //
-  // Full Example: "op1=2x2:dram:interleaved,op2=4x4:l1:block_sharded"
+  // Full Example:
+  // "op1=2x2:dram:interleaved:tile:fp32,op2=4x4:l1:block_sharded:row_major:fp16"
   //
-  // This will set the output layout for op1 to grid 2x2,dram,interleaved and
-  // op2 4x4,l1,block_sharded.
   //
   // Note: This option is only valid if optimizerPassEnabled is true.
   //

include/ttmlir/Dialect/TTNN/Transforms/Optimizer.h

@@ -103,14 +103,14 @@ class TTNNOptimizerBase : public ::mlir::OperationPass<::mlir::ModuleOp> {
 
 protected:
   ::mlir::Pass::Option<llvm::StringMap<InputLayoutOverrideParams>,
-                       mlir::tt::InputLayoutOverrideParser>
+                       mlir::tt::ttnn::InputLayoutOverrideParser>
       overrideInputLayout{
           *this, "insert-memreconfig",
           ::llvm::cl::desc(
               "Manually insert memory reconfig op for specific op's operand."),
           ::llvm::cl::init(llvm::StringMap<InputLayoutOverrideParams>())};
   ::mlir::Pass::Option<llvm::StringMap<OutputLayoutOverrideParams>,
-                       mlir::tt::OutputLayoutOverrideParser>
+                       mlir::tt::ttnn::OutputLayoutOverrideParser>
       overrideOutputLayout{
           *this, "override-output-layout",
           ::llvm::cl::desc("Override output tensor layout for specific ops."),

include/ttmlir/Dialect/TTNN/Transforms/Passes.h

@@ -8,10 +8,10 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
-#include "ttmlir/Dialect/TT/Utils/OverrideParams.h"
 #include "ttmlir/Dialect/TTNN/Analysis/MemoryLayoutAnalysis.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNN.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOps.h"
+#include "ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h"
 
 namespace mlir::tt::ttnn {
 

include/ttmlir/Dialect/TTNN/Utils/OptimizerOverrides.h

@@ -0,0 +1,55 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_DIALECT_TT_UTILS_OVERRIDEPARAMS_H
+#define TTMLIR_DIALECT_TT_UTILS_OVERRIDEPARAMS_H
+
+#include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
+#include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
+#include <cstdint>
+#include <llvm/Support/CommandLine.h>
+
+namespace mlir::tt::ttnn {
+
+struct OutputLayoutOverrideParams {
+  SmallVector<int64_t, 2> grid;
+  tt::MemorySpace memorySpace;
+  tt::TensorMemoryLayout tensorMemoryLayout; // INTERLEAVED / SHARDED etc...
+  tt::ttnn::Layout memoryLayout;             // ROW_MAJOR / TILE
+  tt::DataType dataType;
+};
+
+struct InputLayoutOverrideParams {
+  SmallVector<int64_t> operandIdxes;
+};
+
+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);
+};
+
+} // namespace mlir::tt::ttnn
+
+#endif

lib/Dialect/TTNN/Analysis/LegalGridAnalysis.cpp

@@ -6,6 +6,8 @@
 #include "ttmlir/Dialect/TT/IR/TTOpsTypes.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNN.h"
 #include "ttmlir/Dialect/TTNN/IR/TTNNOps.h"
+#include "ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.h"
+#include "ttmlir/Dialect/TTNN/Utils/Utils.h"
 
 namespace mlir::tt::ttnn {
 
@@ -84,12 +86,21 @@ bool LegalGridAnalysis::applyOverrides() {
       mlir::cast<RankedTensorType>(op->getResult(0).getType());
   tt::LayoutAttr layout = mlir::cast<tt::LayoutAttr>(tensorType.getEncoding());
 
+  GridAttr grid =
+      GridAttr::get(op->getContext(), ArrayRef<int64_t>(override.grid));
+
+  // Create element type for the new layout.
+  Type elementType =
+      utils::createRowMajorTypeFromDtype(op->getContext(), override.dataType);
+  if (override.memoryLayout == Layout::Tile) {
+    elementType = TileType::get(op->getContext(), elementType);
+  }
+
   analysisResult.push_back(
-      layout.withMemorySpace(op->getContext(), override.memorySpace)
-          .withMemoryLayout(op->getContext(), override.memoryLayout)
-          .withGrid(op->getContext(), tensorType,
-                    GridAttr::get(op->getContext(),
-                                  ArrayRef<int64_t>(override.grid))));
+      layout.withGrid(op->getContext(), tensorType, grid)
+          .withMemorySpace(op->getContext(), override.memorySpace)
+          .withMemoryLayout(op->getContext(), override.tensorMemoryLayout)
+          .withElementType(op->getContext(), elementType));
 
   return true;
 }