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

[Samarth Narang]

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

>From 8a4e3b3ca894a4895c8d5e959bb1d997a8f8fed4 Mon Sep 17 00:00:00 2001
From: sanarang <sanarang at nvidia.com>
Date: Sat, 5 Sep 2026 07:48:36 +0000
Subject: [PATCH] [mlir][AMDGPU] Do not request nuw for negative-stride memrefs

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. The signed address arithmetic
is correct and the access is in bounds, so the false promise is the whole
defect.

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. The
speculative prefetch path is unchanged; it already requests no flags.

Fixes #221412.
---
 .../AMDGPUToROCDL/AMDGPUToROCDL.cpp           | 22 +++++++++++++++----
 .../AMDGPUToROCDL/global-prefetch.mlir        | 16 ++++++++++++++
 .../AMDGPUToROCDL/global_transpose_load.mlir  | 19 ++++++++++++++++
 3 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 6916c2a0eff46..26579c7002bad 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 =
@@ -4504,9 +4519,8 @@ struct GlobalPrefetchOpLowering
     MemRefDescriptor descriptor(memRef);
     MemRefType memRefType = op.getSrc().getType();
     Location loc = op->getLoc();
-    auto inboundsFlags = isSpeculative ? LLVM::GEPNoWrapFlags::none
-                                       : LLVM::GEPNoWrapFlags::inbounds |
-                                             LLVM::GEPNoWrapFlags::nuw;
+    auto inboundsFlags =
+        isSpeculative ? LLVM::GEPNoWrapFlags::none : 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>
+}