Add verifier for TTNNLayoutAttr (#2244)

This change adds verifier for `TTNNLayoutAttr`.
I've also removed optional `TenorMemoryLayout` parameter from some
functions in TTNNLayout.cpp and replaced it map from buffer types to
default tensor memory layouts, since there were situations where for
example under certain conditions we constructed `TTNNLayoutAttr` with
`SystemMemory` and `Interleaved`.

include/ttmlir/Dialect/TTNN/IR/TTNNOpsAttrs.td

@@ -172,6 +172,8 @@ def TTNN_TTNNLayoutAttr: TTNN_Attr<"TTNNLayout", "ttnn_layout"> {
     llvm::SmallVector<int64_t> getTiledShape(ArrayRef<int64_t> logicalTensorShape) const;
     AffineMap replaceMemoryMapSymbolsWithShardShape(AffineMap physicalMemoryMap) const;
   }];
+
+  let genVerifyDecl = 1;
 }
 
 #endif  // TTMLIR_TTMLIR_DIALECT_TTNN_TTNNOPSATTRS_TD

lib/Dialect/TTNN/IR/TTNNOpsAttrs.cpp

@@ -65,6 +65,34 @@ bool TTNNLayoutAttr::hasInterleavedDRAMTensorMemoryLayout() const {
          (getMemLayout().getValue() == TensorMemoryLayout::Interleaved);
 }
 
+// Checks:
+// 1. If memory layout is present then:
+//   - System memory buffer type is not allowed
+//   - DRAM buffer type must have Interleaved memory layout
+// 2. If memory layout is not present then:
+//   - Buffer type must be SystemMemory
+llvm::LogicalResult verifyBufferAndMemoryLayout(
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
+    BufferType bufferType, TensorMemoryLayoutAttr memLayoutAttr) {
+  if (memLayoutAttr) {
+    if (bufferType == BufferType::SystemMemory) {
+      return emitError()
+             << "Memory layout is not allowed for SystemMemory buffer type.";
+    }
+
+    if (bufferType == BufferType::DRAM &&
+        memLayoutAttr.getValue() != TensorMemoryLayout::Interleaved) {
+      return emitError()
+             << "DRAM buffer type must have Interleaved memory layout.";
+    }
+  } else if (bufferType != BufferType::SystemMemory) {
+    return emitError()
+           << "Memory layout is required for non-SystemMemory buffer type.";
+  }
+
+  return ::llvm::success();
+}
+
 // Calculate the logical shape of the shard.
 //
 // Shard is defined as a piece of the tensor that is mapped to a single grid
@@ -439,12 +467,28 @@ TTNNLayoutAttr TTNNLayoutAttr::withBufferType(::mlir::MLIRContext *context,
   TensorMemoryLayoutAttr memLayoutAttr = getMemLayout();
   tt::GridAttr grid = getGrid();
 
-  // For SystemMemory we need to clear memory layout attribute
+  // For SystemMemory we need to clear memory layout and set grid to 1x1.
   if (memorySpace == BufferType::SystemMemory) {
     memLayoutAttr = TensorMemoryLayoutAttr{};
     grid = tt::GridAttr::get(context, grid.getShape().size());
   }
 
+  // For DRAM we need to set memory layout to interleaved and set grid to 1x1.
+  if (memorySpace == BufferType::DRAM) {
+    memLayoutAttr =
+        TensorMemoryLayoutAttr::get(context, TensorMemoryLayout::Interleaved);
+    grid = tt::GridAttr::get(context, grid.getShape().size());
+  }
+
+  // For L1 we will inherit the memory layout if its set.
+  // Otherwise we will set it to interleaved.
+  if (memorySpace == BufferType::L1) {
+    memLayoutAttr = getMemLayout()
+                        ? getMemLayout()
+                        : TensorMemoryLayoutAttr::get(
+                              context, TensorMemoryLayout::Interleaved);
+  }
+
   return TTNNLayoutAttr::get(
       context, getLinear(), grid,
       buildMemRef<BufferType, BufferTypeAttr>(context, getScalarShardShape(),
@@ -539,12 +583,6 @@ TTNNLayoutAttr TTNNLayoutAttr::get(
     Type elementType, BufferType bufferType, GridAttr grid,
     TensorMemoryLayoutAttr memLayoutAttr,
     ArrayRef<std::pair<std::int64_t, std::int64_t>> collapseIntervals) {
-
-  // TensorMemoryLayout for SystemMemory is always null
-  if (bufferType == BufferType::SystemMemory) {
-    memLayoutAttr = nullptr;
-  }
-
   // Construct a new affine map which will be used to map from logical
   // space to physical space.
   AffineMap linear = collapsedLinearAffineMap(
@@ -567,6 +605,14 @@ TTNNLayoutAttr TTNNLayoutAttr::get(
   return get(context, linear, grid, memRefType, memLayoutAttr);
 }
 
+::llvm::LogicalResult TTNNLayoutAttr::verify(
+    ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, AffineMap,
+    GridAttr, MemRefType memref, TensorMemoryLayoutAttr memLayout) {
+  BufferType bufferType =
+      mlir::cast<BufferTypeAttr>(memref.getMemorySpace()).getValue();
+  return verifyBufferAndMemoryLayout(emitError, bufferType, memLayout);
+}
+
 // Construct a new MemoryConfig
 //
 // This function creates a deep copy of the current MemoryConfigAttr and
@@ -602,15 +648,9 @@ ::llvm::LogicalResult MemoryConfigAttr::verify(
     ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
     BufferTypeAttr bufferType, ShardSpecAttr shardSpec,
     TensorMemoryLayoutAttr tensorMemoryLayout) {
-  // Verify that we don't have tensorMemoryLayout for BufferType::SystemMemory
-  if (bufferType.getValue() == BufferType::SystemMemory && tensorMemoryLayout) {
-    emitError() << "MemoryConfig with SystemMemory buffer type cannot have "
-                   "tensor memory layout.";
-    return ::llvm::failure();
-  }
+  return verifyBufferAndMemoryLayout(emitError, bufferType.getValue(),
+                                     tensorMemoryLayout);
 
   // TODO(#2140): Once we complete #1628, we should add a verifier for
   // ShardSpecAttr. ShardSpecAttr is only valid if the buffer type is L1.
-
-  return ::llvm::success();
 }

lib/Dialect/TTNN/Transforms/TTNNLayout.cpp

@@ -21,12 +21,15 @@ namespace mlir::tt::ttnn {
 static const std::array<std::pair<int64_t, int64_t>, 1> g_defaultCollapseDims =
     {{{0, -1}}};
 
-// Default memory space for tensors on device
-static const BufferType g_defaultMemorySpaceDevice = BufferType::DRAM;
+static const llvm::SmallDenseMap<BufferType, std::optional<TensorMemoryLayout>,
+                                 4>
+    g_bufferLayoutMap = {
+        {BufferType::DRAM, TensorMemoryLayout::Interleaved},
+        {BufferType::L1, TensorMemoryLayout::Interleaved},
+        {BufferType::SystemMemory, std::nullopt},
+};
 
-// Default memory layout for tensors on device
-static const TensorMemoryLayout g_defaultMemoryLayout =
-    TensorMemoryLayout::Interleaved;
+static const BufferType g_defaultMemorySpaceDevice = BufferType::DRAM;
 
 //===----------------------------------------------------------------------===//
 // Helper methods
@@ -45,16 +48,21 @@ inline Location appendInputSuffix(Location loc, int64_t operandIndex) {
   return loc;
 }
 
-static TTNNLayoutAttr createLayoutAttr(
-    MLIRContext *ctx, GridAttr deviceGrid, RankedTensorType type,
-    std::optional<BufferType> bufferTypeOpt = std::nullopt,
-    std::optional<TensorMemoryLayout> memoryLayoutOpt = std::nullopt,
-    std::optional<bool> isTiledOpt = std::nullopt) {
+static TensorMemoryLayoutAttr getMemoryLayoutAttr(MLIRContext *ctx,
+                                                  BufferType bufferType) {
+  std::optional<TensorMemoryLayout> layout = g_bufferLayoutMap.at(bufferType);
+  if (layout) {
+    return TensorMemoryLayoutAttr::get(ctx, layout.value());
+  }
+
+  return TensorMemoryLayoutAttr{};
+}
+
+static TTNNLayoutAttr
+createLayoutAttr(MLIRContext *ctx, GridAttr deviceGrid, RankedTensorType type,
+                 BufferType bufferType = g_defaultMemorySpaceDevice,
+                 bool isTiled = true) {
 
-  BufferType bufferType = bufferTypeOpt.value_or(g_defaultMemorySpaceDevice);
-  TensorMemoryLayout memoryLayout =
-      memoryLayoutOpt.value_or(g_defaultMemoryLayout);
-  bool isTiled = isTiledOpt.value_or(true);
   std::int64_t deviceGridRank = deviceGrid.getShape().size();
   // Default to single core grid
   GridAttr tensorGrid = GridAttr::get(ctx, deviceGridRank);
@@ -65,9 +73,11 @@ static TTNNLayoutAttr createLayoutAttr(
   // Force TileType for tensors
   auto elementType = isTiled ? TileType::get(ctx, type.getElementType())
                              : type.getElementType();
-  return TTNNLayoutAttr::get(
-      ctx, type.getShape(), elementType, bufferType, tensorGrid,
-      TensorMemoryLayoutAttr::get(ctx, memoryLayout), collapseDimsRef);
+
+  TensorMemoryLayoutAttr memoryLayoutAttr =
+      getMemoryLayoutAttr(ctx, bufferType);
+  return TTNNLayoutAttr::get(ctx, type.getShape(), elementType, bufferType,
+                             tensorGrid, memoryLayoutAttr, collapseDimsRef);
 }
 
 //===----------------------------------------------------------------------===//
@@ -117,7 +127,7 @@ class TTNNLayoutTensorTypeSystemMemoryConverter : public TypeConverter {
       }
 
       TTNNLayoutAttr newLayout = createLayoutAttr(
-          ctx, deviceGrid, type, BufferType::SystemMemory, std::nullopt, false);
+          ctx, deviceGrid, type, BufferType::SystemMemory, false);
       // Convert mlir data types to tt data types
       Type elementType = mlir::tt::dataTypeToElementType(
           ctx, elementTypeToDataType(type.getElementType()));
@@ -210,10 +220,12 @@ class TTNNLayoutTensorTypeRewriter : public RewritePattern {
 // Given desired buffer type, memory layout and type checks if the input tensor
 // needs to be converted to the desired layout. If it does, creates a new
 // EmptyOp/ConstantOp/EmptyOp + ToLayoutOp depending on the input.
-static std::optional<Value>
-createToLayoutOp(PatternRewriter &rewriter, Location loc, Value input,
-                 BufferType desiredBufferType,
-                 TensorMemoryLayoutAttr desiredMemLayoutAttr, bool tiled) {
+static std::optional<Value> createToLayoutOp(PatternRewriter &rewriter,
+                                             Location loc, Value input,
+                                             BufferType desiredBufferType,
+                                             bool tiled) {
+  TensorMemoryLayoutAttr desiredMemLayoutAttr =
+      getMemoryLayoutAttr(rewriter.getContext(), desiredBufferType);
 
   // Get type
   RankedTensorType ty = mlir::cast<RankedTensorType>(input.getType());
@@ -309,9 +321,9 @@ static bool changeLayoutToHost(DestinationStyleOpInterface &op,
                                OpOperand &operand, PatternRewriter &rewriter,
                                bool isDPSResult) {
   Location newLoc = appendInputSuffix(op.getLoc(), operand.getOperandNumber());
-  std::optional<Value> layout = createToLayoutOp(
-      rewriter, newLoc, operand.get(), BufferType::SystemMemory,
-      nullptr /* desiredMemLayoutAttr */, false /* tiled */);
+  std::optional<Value> layout =
+      createToLayoutOp(rewriter, newLoc, operand.get(),
+                       BufferType::SystemMemory, false /* tiled */);
   if (layout.has_value()) {
     rewriter.modifyOpInPlace(op, [&]() {
       op->setOperand(operand.getOperandNumber(), *layout);
@@ -385,10 +397,7 @@ class TTNNLayoutDPSOperandsRewriter
 
       // Given the operand constraint, create the desired layout for the operand
       std::optional<Value> desiredLayout = createToLayoutOp(
-          rewriter, newLoc, operand.get(), g_defaultMemorySpaceDevice,
-          TensorMemoryLayoutAttr::get(rewriter.getContext(),
-                                      g_defaultMemoryLayout),
-          isTiled);
+          rewriter, newLoc, operand.get(), g_defaultMemorySpaceDevice, isTiled);
 
       // If layout changed update the operand
       if (desiredLayout) {
@@ -451,9 +460,9 @@ class TTNNLayoutHoistedFuncCallRewriter
     size_t locIdx = 0;
     for (auto operand : callOp.getOperands()) {
       Location newLoc = appendInputSuffix(callOp.getLoc(), locIdx++);
-      std::optional<Value> optionalLayoutOp = createToLayoutOp(
-          rewriter, newLoc, operand, BufferType::SystemMemory,
-          nullptr /* desiredMemLayoutAttr */, false /* tiled */);
+      std::optional<Value> optionalLayoutOp =
+          createToLayoutOp(rewriter, newLoc, operand, BufferType::SystemMemory,
+                           false /* tiled */);
       fromDeviceOperands.push_back(
           optionalLayoutOp.has_value() ? optionalLayoutOp.value() : operand);
     }
@@ -576,7 +585,7 @@ class TTNNLayoutFuncInputOutputTypeRewriter
   RankedTensorType toSystemMemoryType(MLIRContext *ctx,
                                       RankedTensorType ty) const {
     TTNNLayoutAttr newLayout = createLayoutAttr(
-        ctx, deviceGrid, ty, BufferType::SystemMemory, std::nullopt, false);
+        ctx, deviceGrid, ty, BufferType::SystemMemory, false /* isTiledOpt */);
     auto newType =
         RankedTensorType::get(ty.getShape(), ty.getElementType(), newLayout);
     return newType;
@@ -632,18 +641,12 @@ class TTNNLayoutFuncReturnRewriter
       BufferType desiredBufferType =
           forceHost ? BufferType::SystemMemory : g_defaultMemorySpaceDevice;
 
-      TensorMemoryLayoutAttr desiredMemLayoutAttr =
-          forceHost ? nullptr
-                    : TensorMemoryLayoutAttr::get(rewriter.getContext(),
-                                                  g_defaultMemoryLayout);
-
       bool isTiled = !forceHost;
 
       Location newLoc =
           appendInputSuffix(op.getLoc(), operand.getOperandNumber());
-      std::optional<Value> updatedLayout =
-          createToLayoutOp(rewriter, newLoc, operand.get(), desiredBufferType,
-                           desiredMemLayoutAttr, isTiled);
+      std::optional<Value> updatedLayout = createToLayoutOp(
+          rewriter, newLoc, operand.get(), desiredBufferType, isTiled);
       if (updatedLayout.has_value()) {
         rewriter.modifyOpInPlace(op, [&]() {
           op.setOperand(operand.getOperandNumber(), *updatedLayout);

test/ttmlir/Dialect/TTNN/data_movement/to_memory_config/to_memory_config_tests_negative.mlir

@@ -22,7 +22,7 @@ module {
 #device_tile_layout2 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x3x!tt.tile<32x32, f32>, #dram>, <interleaved>>
 module{
   func.func @forward(%arg0: tensor<32x32xf32, #device_tile_layout1>) -> tensor<32x96xf32, #device_tile_layout2> {
-    // CHECK: error: 'ttnn.to_memory_config' op Output tensor layout memory space must match memory config memory space.
+    // CHECK: error: DRAM buffer type must have Interleaved memory layout.
     %1 = "ttnn.to_memory_config"(%arg0) <{memory_config = #ttnn.memory_config<#dram, <<1x3>>, <single_bank>>}> : (tensor<32x32xf32, #device_tile_layout1>) -> tensor<32x96xf32, #device_tile_layout2>
     return %1 : tensor<32x96xf32, #device_tile_layout2>
   }
@@ -49,10 +49,10 @@ module {
 #dram = #ttnn.buffer_type<dram>
 #system_memory = #ttnn.buffer_type<system_memory>
 #device_tile_layout1 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #dram>, <interleaved>>
-#device_tile_layout2 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x3x!tt.tile<32x32, f32>, #system_memory>, <interleaved>>
+#device_tile_layout2 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x3x!tt.tile<32x32, f32>, #system_memory>>
 module {
   func.func @forward(%arg0: tensor<32x32xf32, #device_tile_layout1>) -> tensor<32x96xf32, #device_tile_layout2> {
-    // CHECK: error: MemoryConfig with SystemMemory buffer type cannot have tensor memory layout.
+    // CHECK: error: Memory layout is not allowed for SystemMemory buffer type.
     %1 = "ttnn.to_memory_config"(%arg0) <{memory_config = #ttnn.memory_config<#system_memory, <<1x4>>, <interleaved>>}> : (tensor<32x32xf32, #device_tile_layout1>) -> tensor<32x96xf32, #device_tile_layout2>
     return %1 : tensor<32x96xf32, #device_tile_layout2>
   }

test/ttmlir/Dialect/TTNN/layout_encoding.mlir

@@ -0,0 +1,42 @@
+// RUN: not ttmlir-opt --split-input-file %s 2>&1 | FileCheck %s
+// Negative tests for TTNNLayoutAttr
+
+// Verify that verification fails if tensor memory layout is set for SystemMemory buffer type.
+#system_memory = #ttnn.buffer_type<system_memory>
+#ttnn_layout = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<32x32xf32, #system_memory>, <interleaved>>
+#ttnn_layout1 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #system_memory>, <interleaved>>
+module {
+  func.func @forward(%arg0: tensor<32x32xf32, #ttnn_layout>) -> tensor<32x96xf32, #ttnn_layout1> {
+    // CHECK: error: Memory layout is not allowed for SystemMemory buffer type.
+    %1 = "ttnn.to_layout"(%arg0) <{layout = #ttnn.layout<tile>}>  : (tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout1>
+    return %1 : tensor<32x32xf32, #ttnn_layout1>
+  }
+}
+
+// -----
+
+// Verify that verification fails if tensor memory layout is set to anything other than interleaved for DRAM buffer type.
+#dram = #ttnn.buffer_type<dram>
+#ttnn_layout = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<32x32xf32, #dram>, <block_sharded>>
+#ttnn_layout1 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #dram>, <block_sharded>>
+module {
+  func.func @forward(%arg0: tensor<32x32xf32, #ttnn_layout>) -> tensor<32x96xf32, #ttnn_layout1> {
+    // CHECK: error: DRAM buffer type must have Interleaved memory layout.
+    %1 = "ttnn.to_layout"(%arg0) <{layout = #ttnn.layout<tile>}>  : (tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout1>
+    return %1 : tensor<32x32xf32, #ttnn_layout1>
+  }
+}
+
+// -----
+
+// Verify that verification fails if tensor memory layout is not set for DRAM/L1 buffer type.
+#l1 = #ttnn.buffer_type<l1>
+#ttnn_layout = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<32x32xf32, #l1>>
+#ttnn_layout1 = #ttnn.ttnn_layout<(d0, d1) -> (d0, d1), <1x1>, memref<1x1x!tt.tile<32x32, f32>, #l1>>
+module {
+  func.func @forward(%arg0: tensor<32x32xf32, #ttnn_layout>) -> tensor<32x96xf32, #ttnn_layout1> {
+    // CHECK: error: Memory layout is required for non-SystemMemory buffer type.
+    %1 = "ttnn.to_layout"(%arg0) <{layout = #ttnn.layout<tile>}>  : (tensor<32x32xf32, #ttnn_layout>) -> tensor<32x32xf32, #ttnn_layout1>
+    return %1 : tensor<32x32xf32, #ttnn_layout1>
+  }
+}

test/unittests/OpModel/TTNN/Conversion/TestConversion.cpp

@@ -449,14 +449,20 @@ TEST_P(MlirToTtnnConversionMemoryConfig, MemoryConfig) {
 }
 
 INSTANTIATE_TEST_SUITE_P(
-    ToMemoryConfig, MlirToTtnnConversionMemoryConfig,
-    ::testing::Combine(
-        ::testing::Values(mlir::tt::ttnn::BufferType::DRAM,
-                          mlir::tt::ttnn::BufferType::L1),
-        ::testing::Values(mlir::tt::ttnn::TensorMemoryLayout::Interleaved,
-                          mlir::tt::ttnn::TensorMemoryLayout::HeightSharded,
-                          mlir::tt::ttnn::TensorMemoryLayout::WidthSharded,
-                          mlir::tt::ttnn::TensorMemoryLayout::BlockSharded)));
+    ToMemoryConfig, MlirToTtnnConversionMemoryConfig, ::testing::ValuesIn([] {
+      using mlir::tt::ttnn::BufferType;
+      using mlir::tt::ttnn::TensorMemoryLayout;
+
+      std::vector<std::tuple<BufferType, TensorMemoryLayout>>
+          validCombinations = {
+              {BufferType::DRAM,
+               TensorMemoryLayout::Interleaved}, // DRAM only with Interleaved
+              {BufferType::L1, TensorMemoryLayout::Interleaved},
+              {BufferType::L1, TensorMemoryLayout::HeightSharded},
+              {BufferType::L1, TensorMemoryLayout::WidthSharded},
+              {BufferType::L1, TensorMemoryLayout::BlockSharded}};
+      return validCombinations;
+    }()));
 
 //================================================================================
 // getTensorLayout

test/unittests/OpModel/TTNN/OpModelFixture.h

@@ -80,13 +80,17 @@ class OpModelFixture : public ::testing::Test {
             ? virtualGrid.value()
             : GetVirtualGridShape(tensorShape, tensorMemoryLayout);
 
+    auto memLayoutAttr = bufferType == mlir::tt::ttnn::BufferType::SystemMemory
+                             ? mlir::tt::ttnn::TensorMemoryLayoutAttr{}
+                             : mlir::tt::ttnn::TensorMemoryLayoutAttr::get(
+                                   &context, tensorMemoryLayout);
+
     return mlir::tt::ttnn::TTNNLayoutAttr::get(
         &context, tensorShape,
         mlir::tt::TileType::get(&context, builder.getBF16Type()), bufferType,
         CreateGrid(&context, tensorMemoryLayout, virtualGridSelected,
                    physicalGrid),
-        mlir::tt::ttnn::TensorMemoryLayoutAttr::get(&context,
-                                                    tensorMemoryLayout));
+        memLayoutAttr);
   }
 
   mlir::tt::ttnn::TTNNLayoutAttr CreateRowMajorLayout(