[Mlir-commits] [mlir] [mlir][nvgpu] Expose which memrefs optimize-shared-memory can swizzle (NFC) (PR #237180)
Aditya Dandwate
llvmlistbot at llvm.org
Fri Oct 30 04:31:22 PDT 2026
https://github.com/adityadandwate created https://github.com/llvm/llvm-project/pull/237180
Depends on #237179; only the last commit is new.
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
>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/2] [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 1bc819002ec8b5a7a75072024cdf0abff0a33f13 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 0f12231d5c3c6f..6fb87dc70a3a19 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,10 +167,27 @@ 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.
- 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
More information about the Mlir-commits
mailing list