[Mlir-commits] [mlir] [mlir][nvgpu] Expose which memrefs optimize-shared-memory can swizzle (NFC) (PR #237180)

[Aditya Dandwate]

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

>From 13bb80d9f21893f12780b637e43f82231d00bdb8 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/2] [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 moves some accesses out of their row and into the next
one, so data is overwritten.

For example, in `memref<64x96xf16, 3>` each row holds 12 128-bit vectors, and
vector `c` of row `r` is moved to `c ^ (r & 7)`. Vector 8 of row 4 and vector 5
of row 5 then both land at vector 60 of the buffer, so one store overwrites the
other and a later load of either returns the wrong value. Vector 8 of row 63
lands at vector 771, past the end of the 768-vector buffer.

Skip such memrefs and leave their accesses unswizzled.

Assisted-by: Claude
---
 .../NVGPU/Transforms/OptimizeSharedMemory.cpp |  6 ++++
 .../Dialect/NVGPU/optimize-shared-memory.mlir | 29 +++++++++++++++++++
 2 files changed, 35 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..2f39753702b2b9 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_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

>From a9d0e740c4d9f573bf86166d576632ca903adcf2 Mon Sep 17 00:00:00 2001
From: aditya dandwate <adandwate at polymagelabs.com>
Date: Fri, 30 Oct 2026 16:51:11 +0530
Subject: [PATCH 2/2] [mlir][nvgpu] Expose which memrefs optimize-shared-memory
 can swizzle (NFC)

Move the checks on a memref's type out of
`optimizeSharedMemoryReadsAndWrites` into a new public function,
`canSwizzleSharedMemoryAccesses`. A pipeline that pads the shared memory
buffers that won't be swizzled can then ask for that decision instead of
duplicating it.

Assisted-by: Claude
---
 .../Dialect/NVGPU/Transforms/Transforms.h     | 10 +++++
 .../NVGPU/Transforms/OptimizeSharedMemory.cpp | 42 ++++++++++---------
 2 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/mlir/include/mlir/Dialect/NVGPU/Transforms/Transforms.h b/mlir/include/mlir/Dialect/NVGPU/Transforms/Transforms.h
index 2ff9ac43851d84..c2beab4a1fed1b 100644
--- a/mlir/include/mlir/Dialect/NVGPU/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/NVGPU/Transforms/Transforms.h
@@ -16,6 +16,7 @@
 #include "mlir/IR/Operation.h"
 
 namespace mlir {
+class MemRefType;
 class RewriterBase;
 
 namespace nvgpu {
@@ -24,6 +25,15 @@ namespace nvgpu {
 /// Passes
 ///
 
+/// Returns true if `optimizeSharedMemoryReadsAndWrites` can swizzle accesses to
+/// shared memory buffers of type `memRefType`. Their rows, i.e., the innermost
+/// dimension, must be of static size, longer than 16 bytes, and span a power of
+/// 2 number of bits. The swizzle XORs the index of each 128-bit vector in a row
+/// with bits of the row index; it keeps the result within the row only if a row
+/// is a power of 2 number of vectors. Memrefs with vector element types are not
+/// supported.
+bool canSwizzleSharedMemoryAccesses(MemRefType memRefType);
+
 /// Optimizes vectorized accesses to a shared memory buffer specified by
 /// memrefValue. This transformation assumes the following:
 /// 1) All relevant accesses to `memrefValue` are contained with `parentOp`.
diff --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index ac3102d4b1815b..f82db90c38efff 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -145,34 +145,21 @@ getShmReadAndWriteOps(Operation *parentOp, Value shmMemRef,
   return success();
 }
 
-llvm::LogicalResult
-mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
-                                                Value memrefValue) {
-  auto memRefType = dyn_cast<MemRefType>(memrefValue.getType());
-  if (!memRefType || !NVGPUDialect::hasSharedMemoryAddressSpace(memRefType))
-    return failure();
-
+bool mlir::nvgpu::canSwizzleSharedMemoryAccesses(MemRefType memRefType) {
   // Not support 0D MemRefs.
   if (memRefType.getRank() == 0)
-    return failure();
+    return false;
 
   // Only support memrefs with scalar element types (i.e., int or float).
   // Memrefs with vector element types are not supported.
   if (!memRefType.getElementType().isIntOrFloat())
-    return failure();
-
-  // Abort if the given value has any sub-views; we do not do any alias
-  // analysis.
-  bool hasSubView = false;
-  parentOp->walk([&](memref::SubViewOp subView) { hasSubView = true; });
-  if (hasSubView)
-    return failure();
+    return false;
 
   // Check if this is necessary given the assumption of 128b accesses:
   // If dim[rank-1] is small enough to fit 8 rows in a 128B line.
   const int64_t rowSize = memRefType.getDimSize(memRefType.getRank() - 1);
   if (ShapedType::isDynamic(rowSize) || rowSize == 0)
-    return failure();
+    return false;
 
   const int64_t rowsPerLine =
       (8 * kSharedMemoryLineSizeBytes / memRefType.getElementTypeBitWidth()) /
@@ -180,12 +167,29 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   const int64_t threadGroupSize =
       1LL << (7 - llvm::Log2_64(kDefaultVectorSizeBits / 8));
   if (rowsPerLine >= threadGroupSize)
-    return failure();
+    return false;
 
   // 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 llvm::isPowerOf2_64(rowSize);
+}
+
+llvm::LogicalResult
+mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
+                                                Value memrefValue) {
+  auto memRefType = dyn_cast<MemRefType>(memrefValue.getType());
+  if (!memRefType || !NVGPUDialect::hasSharedMemoryAddressSpace(memRefType))
+    return failure();
+
+  if (!canSwizzleSharedMemoryAccesses(memRefType))
+    return failure();
+
+  // Abort if the given value has any sub-views; we do not do any alias
+  // analysis.
+  bool hasSubView = false;
+  parentOp->walk([&](memref::SubViewOp subView) { hasSubView = true; });
+  if (hasSubView)
     return failure();
 
   // Get sets of operations within the function that read/write to shared