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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-nvgpu

@llvm/pr-subscribers-mlir

Author: Aditya Dandwate (adityadandwate)

<details>
<summary>Changes</summary>

`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

---
Full diff: https://github.com/llvm/llvm-project/pull/237179.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp (+6) 
- (modified) mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir (+29) 


``````````diff
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

``````````

</details>


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