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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-affine

Author: purnima shrivastava (purnima-nlp)

<details>
<summary>Changes</summary>

`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


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


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+4-1) 
- (modified) mlir/test/Dialect/Affine/constant-fold.mlir (+12) 


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

``````````

</details>


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