[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 87ec900f8864aa23cb817974b6b2edad0a01d941 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 1/3] [mlir][nvgpu] Don't swizzle shared memory rows that
 aren't a power of 2

`optimizeSharedMemoryReadsAndWrites` XORs the index of each 128-bit vector
in a row with the low `log2(vectors per row)` bits of the row index, rounded
down. When the number of vectors per row isn't a power of 2, the result can
leave the row. In a `memref<64x96xf16, 3>`, a row holds 12 vectors and
vectors 8-11 are moved to 12-15, which belong to the next row, so stores
overwrite each other and loads return the wrong data.

Skip such memrefs, as is already done for rows too short to need swizzling.

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

diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index e06b15febf145e..0f12231d5c3c6f 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -182,6 +182,10 @@ 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.
+  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..4421e0f41fcf83 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -241,6 +241,35 @@ func.func @abort_if_subview(%arg0: memref<128x128xf16>,
 
 // -----
 
+// A row of 96 f16 elements is 12 vectors of 128 bits, 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.token
+  nvgpu.device_async_wait[%1] { numGroups = 1 : i32}
+
+  // CHECK-NOT: arith.xori
+  // CHECK: nvgpu.ldmatrix [[shm]][[[fragRow]], [[fragCol]]]
+  %mat = nvgpu.ldmatrix %shm[%fragRow, %fragCol] {numTiles = 4 : i32, transpose = false}
+      : memref<64x96xf16, 3> -> vector<4x2xf16>
+
+  return %mat: vector<4x2xf16>
+}
+
+// -----
+
 // Ensure this case not crash
 
 // CHECK-LABEL: func @test_0_d

>From 848e706afabfed0bf0528f9358380363140371d7 Mon Sep 17 00:00:00 2001
From: aditya dandwate <adandwate at polymagelabs.com>
Date: Fri, 30 Oct 2026 17:03:16 +0530
Subject: [PATCH 2/3] Use the current assembly syntax of the NVGPU ops in the
 new test

---
 mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
index 4421e0f41fcf83..2f39753702b2b9 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -257,12 +257,12 @@ func.func @no_swizzle_non_power_of_2_row_f16(%arg0: memref<128x128xf16>,
   // 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.token
-  nvgpu.device_async_wait[%1] { numGroups = 1 : i32}
+  %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 : i32, transpose = false}
+  %mat = nvgpu.ldmatrix %shm[%fragRow, %fragCol] numTiles = 4 transpose = false
       : memref<64x96xf16, 3> -> vector<4x2xf16>
 
   return %mat: vector<4x2xf16>

>From 32567525fc601de143e687ea6fa7064777924da9 Mon Sep 17 00:00:00 2001
From: aditya dandwate <adandwate at polymagelabs.com>
Date: Fri, 30 Oct 2026 17:08:03 +0530
Subject: [PATCH 3/3] Say what the swizzle does to rows that aren't a power of
 2 long

---
 mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index 0f12231d5c3c6f..ac3102d4b1815b 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -183,6 +183,8 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
     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();