[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 created https://github.com/llvm/llvm-project/pull/237179

`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

>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] [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