[Mlir-commits] [mlir] [mlir][AMDGPU] Do not request nuw for negative-stride memrefs (PR #221419)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Samarth Narang (snarang181)

<details>
<summary>Changes</summary>

AMDGPUToROCDL asked getStridedElementPtr for inbounds|nuw unconditionally when lowering amdgpu.global_transpose_load and non-speculative amdgpu.global_prefetch. nuw promises that each index*stride term does not wrap as an unsigned value, which a negative stride cannot satisfy: it is a large unsigned bit pattern, so idx * (-1 as u64) wraps for any positive index and LLVM may treat the resulting address as poison. 

Gate nuw on all strides being statically non-negative, matching what MemRefToLLVM already does for memref.load/store (#<!-- -->204309). A dynamic stride also fails the check, since non-negativity cannot be established for one. We do not touch the speculative prefetch path as it already requests no flags.

Fixes #<!-- -->221412.

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


3 Files Affected:

- (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+17-3) 
- (modified) mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir (+16) 
- (modified) mlir/test/Conversion/AMDGPUToROCDL/global_transpose_load.mlir (+19) 


``````````diff
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 2a846339b562e..42d0999ac9fc4 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -43,6 +43,21 @@ namespace mlir {
 using namespace mlir;
 using namespace mlir::amdgpu;
 
+/// Returns the GEP no-wrap flags valid for address arithmetic on `memRefType`.
+/// `inbounds` always holds, since these ops require in-bounds indices. `nuw`
+/// additionally promises that no `index * stride` term wraps as an unsigned
+/// value, which only holds when every stride is statically non-negative. This
+/// matches what MemRefToLLVM does for `memref.load`/`store`.
+static LLVM::GEPNoWrapFlags getNoWrapFlags(MemRefType memRefType) {
+  auto [strides, offset] = memRefType.getStridesAndOffset();
+  LLVM::GEPNoWrapFlags flags = LLVM::GEPNoWrapFlags::inbounds;
+  if (llvm::all_of(strides, [](int64_t stride) {
+        return ShapedType::isStatic(stride) && stride >= 0;
+      }))
+    flags = flags | LLVM::GEPNoWrapFlags::nuw;
+  return flags;
+}
+
 // Define commonly used chipsets versions for convenience.
 constexpr Chipset kGfx908 = Chipset(9, 0, 8);
 constexpr Chipset kGfx90a = Chipset(9, 0, 0xa);
@@ -2356,7 +2371,7 @@ struct GlobalTransposeLoadOpLowering
 
     Value srcPtr = getStridedElementPtr(
         rewriter, loc, srcMemRefType, adaptor.getSrc(), adaptor.getSrcIndices(),
-        LLVM::GEPNoWrapFlags::inbounds | LLVM::GEPNoWrapFlags::nuw);
+        getNoWrapFlags(srcMemRefType));
 
     size_t numElements = resultType.getNumElements();
     size_t elementTypeSize =
@@ -4505,8 +4520,7 @@ struct GlobalPrefetchOpLowering
     MemRefType memRefType = op.getSrc().getType();
     Location loc = op->getLoc();
     auto inboundsFlags = isSpeculative ? LLVM::GEPNoWrapFlags::none
-                                       : LLVM::GEPNoWrapFlags::inbounds |
-                                             LLVM::GEPNoWrapFlags::nuw;
+                                       : getNoWrapFlags(memRefType);
     Value prefetchPtr = getStridedElementPtr(
         rewriter, loc, memRefType, descriptor, indices, inboundsFlags);
 
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir b/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
index f71de64cd071f..02a56c5c4b386 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
@@ -31,3 +31,19 @@ func.func @glb_prefetch3(%src : memref<64x64xf16, #gpu.address_space<global>>, %
   amdgpu.global_prefetch %src[%i, %j] NT_RT SE speculative : memref<64x64xf16, #gpu.address_space<global>>
   func.return
 }
+
+// -----
+
+// As above: a negative stride cannot carry the `nuw` promise.
+
+// CHECK-LABEL: @glb_prefetch_negative_stride
+func.func @glb_prefetch_negative_stride(
+    %src : memref<64x64xf16, strided<[64, -1], offset: 63>, #gpu.address_space<global>>,
+    %i : i64, %j : i64) {
+  // CHECK: llvm.getelementptr inbounds %{{.*}}[%{{.*}}]
+  // CHECK-NOT: inbounds|nuw
+  // CHECK: rocdl.global.prefetch
+  amdgpu.global_prefetch %src[%i, %j] HT WGP
+    : memref<64x64xf16, strided<[64, -1], offset: 63>, #gpu.address_space<global>>
+  func.return
+}
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/global_transpose_load.mlir b/mlir/test/Conversion/AMDGPUToROCDL/global_transpose_load.mlir
index f378d7232d7b3..f22c6abc5f382 100644
--- a/mlir/test/Conversion/AMDGPUToROCDL/global_transpose_load.mlir
+++ b/mlir/test/Conversion/AMDGPUToROCDL/global_transpose_load.mlir
@@ -55,3 +55,22 @@ func.func @global_transpose_load_16xi6(%i : index, %j : index,
          : memref<128x32xi8, #gpu.address_space<global>> -> vector<16xi6>
   return %0 : vector<16xi6>
 }
+
+// -----
+
+// A negative stride makes the `index * stride` term wrap when computed as an
+// unsigned value, so `nuw` must not be requested even though the signed
+// address is correct and in bounds.
+
+// CHECK-LABEL: func @global_transpose_load_negative_stride
+func.func @global_transpose_load_negative_stride(%i : index, %j : index,
+    %src : memref<128x256xf16, strided<[256, -1], offset: 255>, #gpu.address_space<global>>)
+    -> vector<8xf16> {
+  // CHECK: llvm.mul %{{.*}}, %{{.*}} overflow<nsw> : i64
+  // CHECK-NOT: overflow<nsw, nuw>
+  // CHECK: llvm.getelementptr inbounds %{{.*}}[%{{.*}}]
+  // CHECK-NOT: inbounds|nuw
+  %0 = amdgpu.global_transpose_load %src[%i, %j]
+         : memref<128x256xf16, strided<[256, -1], offset: 255>, #gpu.address_space<global>> -> vector<8xf16>
+  return %0 : vector<8xf16>
+}

``````````

</details>


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