[Mlir-commits] [mlir] [mlir][affine] Avoid folding out-of-bounds constant loads (PR #224076)

[purnima shrivastava]

https://github.com/purnima-nlp created https://github.com/llvm/llvm-project/pull/224076

`AffineLoadOp::fold` could crash when folding a constant load with out-of-bounds indices, including loads from zero-sized constant memrefs.

Check `ElementsAttr::isValidIndex` before indexing the constant attribute and skip folding when the access is invalid.

Adds a regression test for `memref<0xi8>`.

Fixes #223949


>From cf64fa66b8929314d194479557c12298bdd65bfd Mon Sep 17 00:00:00 2001
From: Purnima Shrivastava <purnimashrivastava05 at .com>
Date: Wed, 16 Sep 2026 22:17:31 +0530
Subject: [PATCH] [mlir][affine] Avoid folding out-of-bounds constant loads

---
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp    |  5 ++++-
 mlir/test/Dialect/Affine/constant-fold.mlir | 12 ++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 83a9e06374457..0b14c7de07058 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -3520,7 +3520,10 @@ OpFoldResult AffineLoadOp::fold(FoldAdaptor adaptor) {
   auto indices =
       llvm::map_to_vector<4>(getAffineMap().getConstantResults(),
                              [](int64_t v) -> uint64_t { return v; });
-  return cstAttr.getValues<Attribute>()[indices];
+  ElementsAttr elementsAttr = cstAttr;
+  if (!elementsAttr.isValidIndex(indices))
+    return {};
+  return elementsAttr.getValues<Attribute>()[indices];
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Affine/constant-fold.mlir b/mlir/test/Dialect/Affine/constant-fold.mlir
index 8bddacc024751..ace83e49974c7 100644
--- a/mlir/test/Dialect/Affine/constant-fold.mlir
+++ b/mlir/test/Dialect/Affine/constant-fold.mlir
@@ -90,3 +90,15 @@ func.func @affine_apply_poison_division_zero_2() -> index {
   %0 = affine.apply affine_map<(d0)[s0] -> (d0 mod (s0 - s0))>(%c16)[%c16]
   return %0 : index
 }
+
+// -----
+
+memref.global "private" constant @empty : memref<0xi8> = dense<>
+
+// CHECK-LABEL: func.func @load_from_empty_global
+func.func @load_from_empty_global() -> i8 {
+  %0 = memref.get_global @empty : memref<0xi8>
+  // CHECK: affine.load %{{.*}}[0] : memref<0xi8>
+  %1 = affine.load %0[0] : memref<0xi8>
+  return %1 : i8
+}