[Mlir-commits] [mlir] [mlir][nvgpu] Don't swizzle shared memory rows that aren't a power of 2 (PR #237179)

[Aditya Dandwate]

https://github.com/adityadandwate updated https://github.com/llvm/llvm-project/pull/237179

>From b62575ef71ee3bece1bf3f7ea259dfec62f59b45 Mon Sep 17 00:00:00 2001
From: aditya dandwate <adandwate at polymagelabs.com>
Date: Fri, 30 Oct 2026 16:50:34 +0530
Subject: [PATCH] [mlir][nvgpu] Don't swizzle shared memory rows that aren't a
 power of 2

`nvgpu-optimize-shared-memory` miscompiles memrefs whose rows aren't a power
of 2 long. The swizzle XORs the index of each 16-byte vector in a row with low
bits of the row index. When a row holds a power of 2 number of vectors, the
result is always an index in the same row, so the swizzle only reorders each
row. Otherwise, the result can be past the end of the row: the access then
lands in the next row and overwrites one of its vectors, or past the end of
the buffer for the last rows.

Skip such memrefs and leave their accesses unswizzled.

Assisted-by: Claude
---
 .../NVGPU/Transforms/OptimizeSharedMemory.cpp |  6 ++++
 .../Dialect/NVGPU/optimize-shared-memory.mlir | 30 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index e06b15febf145e..ac3102d4b1815b 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -182,6 +182,12 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   if (rowsPerLine >= threadGroupSize)
     return failure();
 
+  // The XOR swizzle stays within a row only if the row size is a power of 2.
+  // Otherwise it moves some accesses into the next row, where they overwrite
+  // other elements.
+  if (!llvm::isPowerOf2_64(rowSize))
+    return failure();
+
   // Get sets of operations within the function that read/write to shared
   // memory.
   SmallVector<Operation *, 16> shmReadOps;
diff --git a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
index a7d0e5d30e3dd6..127090d2eb2548 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -241,6 +241,36 @@ func.func @abort_if_subview(%arg0: memref<128x128xf16>,
 
 // -----
 
+// A row of 96 f16 elements is 192 bytes, i.e. 12 vectors of 16 bytes, which is
+// not a power of 2. The XOR swizzle would move accesses into the next row;
+// don't swizzle.
+
+// CHECK: @no_swizzle_non_power_of_2_row_f16([[arg0:%.+]]: memref<{{.*}}>, [[ldRow:%.+]]: index, [[ldCol:%.+]]: index, [[stRow:%.+]]: index, [[stCol:%.+]]: index, [[fragRow:%.+]]: index, [[fragCol:%.+]]: index)
+func.func @no_swizzle_non_power_of_2_row_f16(%arg0: memref<128x128xf16>,
+                               %ldRow: index, %ldCol: index,
+                               %stRow: index, %stCol: index,
+                               %fragRow: index, %fragCol :index)
+                                -> vector<4x2xf16> {
+  // CHECK: [[shm:%.+]] = memref.alloc
+  %shm = memref.alloc() : memref<64x96xf16, 3>
+
+  // CHECK-NOT: arith.xori
+  // CHECK: nvgpu.device_async_copy [[arg0]][[[ldRow]], [[ldCol]]], [[shm]][[[stRow]], [[stCol]]]
+  %0 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shm[%stRow, %stCol], 8
+      : memref<128x128xf16> to memref<64x96xf16, 3>
+  %1 = nvgpu.device_async_create_group %0
+  nvgpu.device_async_wait %1 numGroups = 1
+
+  // CHECK-NOT: arith.xori
+  // CHECK: nvgpu.ldmatrix [[shm]][[[fragRow]], [[fragCol]]]
+  %mat = nvgpu.ldmatrix %shm[%fragRow, %fragCol] numTiles = 4 transpose = false
+      : memref<64x96xf16, 3> -> vector<4x2xf16>
+
+  return %mat: vector<4x2xf16>
+}
+
+// -----
+
 // Ensure this case not crash
 
 // CHECK-LABEL: func @test_0_d