[Mlir-commits] [flang] [mlir] [flang] Emit llvm.assume for array bounds constraints (PR #178811)

[Sebastian Pop]

https://github.com/sebpop created https://github.com/llvm/llvm-project/pull/178811

Emit llvm.assume intrinsics for array bounds checking during the FIR to LLVM conversion. This implements array bounds constraints per Fortran 2018 standard section 9.5.3.3.2: "The value of each subscript expression shall be within the bounds for its dimension unless the array section has size zero."

For non-boxed arrays with known shape, the conversion emits two assumes per dimension:
- assume(index >= lower_bound)
- assume(index <= lower_bound + extent - 1)

This enables LLVM loop nest optimizations to take advantage of the guaranteed bounds information.

>From d6059e5a7aa9f45e0f1517500d394645a4dcd750 Mon Sep 17 00:00:00 2001
From: Sebastian Pop <spop at nvidia.com>
Date: Wed, 28 Jan 2026 10:03:01 -0600
Subject: [PATCH] [flang] Emit llvm.assume for array bounds constraints

Emit llvm.assume intrinsics for array bounds checking during the FIR to
LLVM conversion. This implements array bounds constraints per Fortran
2018 standard section 9.5.3.3.2: "The value of each subscript expression
shall be within the bounds for its dimension unless the array section
has size zero."

For non-boxed arrays with known shape, the conversion emits two assumes
per dimension:
- assume(index >= lower_bound)
- assume(index <= lower_bound + extent - 1)

This enables LLVM loop nest optimizations to take advantage of the
guaranteed bounds information.
---
 flang/lib/Optimizer/CodeGen/CodeGen.cpp       | 28 +++++++++
 flang/test/Fir/array-bounds-assume.fir        | 58 +++++++++++++++++++
 .../LLVMIR/LLVMToLLVMIRTranslation.cpp        | 11 ++++
 3 files changed, 97 insertions(+)
 create mode 100644 flang/test/Fir/array-bounds-assume.fir

diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 625701725003f..6668ccbfbbd5f 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2699,6 +2699,34 @@ struct XArrayCoorOpConversion
       mlir::Value lb =
           isShifted ? integerCast(loc, rewriter, idxTy, operands[shiftOffset])
                     : one;
+
+      // Emit array bounds assumes per Fortran 2018 standard section 9.5.3.3.2:
+      // "The value of each subscript expression shall be within the bounds for
+      // its dimension unless the array section has size zero."
+      // For non-boxed arrays with known shape, emit: assume(index >= lb) and
+      // assume(index <= lb + extent - 1).
+      // The llvm.array_bounds attribute marks these assumes for removal before
+      // vectorization to prevent IR bloat and avoid impacting cost models.
+      if (!baseIsBoxed && !coor.getShape().empty()) {
+        mlir::Value extent =
+            integerCast(loc, rewriter, idxTy, operands[shapeOffset]);
+        // Compute upper bound: ub = lb + extent - 1.
+        mlir::Value lbPlusExtent =
+            mlir::LLVM::AddOp::create(rewriter, loc, idxTy, lb, extent, nsw);
+        mlir::Value ub = mlir::LLVM::SubOp::create(rewriter, loc, idxTy,
+                                                   lbPlusExtent, one, nsw);
+        // Create bounds check conditions.
+        mlir::Value lbCheck = mlir::LLVM::ICmpOp::create(
+            rewriter, loc, mlir::LLVM::ICmpPredicate::sge, index, lb);
+        mlir::Value ubCheck = mlir::LLVM::ICmpOp::create(
+            rewriter, loc, mlir::LLVM::ICmpPredicate::sle, index, ub);
+        // Emit assumes for lower and upper bounds with array_bounds marker.
+        auto lbAssume = mlir::LLVM::AssumeOp::create(rewriter, loc, lbCheck);
+        lbAssume->setAttr("llvm.array_bounds", rewriter.getUnitAttr());
+        auto ubAssume = mlir::LLVM::AssumeOp::create(rewriter, loc, ubCheck);
+        ubAssume->setAttr("llvm.array_bounds", rewriter.getUnitAttr());
+      }
+
       mlir::Value step = one;
       bool normalSlice = isSliced;
       // Compute zero based index in dimension i of the element, applying
diff --git a/flang/test/Fir/array-bounds-assume.fir b/flang/test/Fir/array-bounds-assume.fir
new file mode 100644
index 0000000000000..1fcde26585e8e
--- /dev/null
+++ b/flang/test/Fir/array-bounds-assume.fir
@@ -0,0 +1,58 @@
+// Test that array bounds assumes are generated for fir.array_coor operations.
+// This tests that Flang emits llvm.assume intrinsics for array bounds constraints.
+
+// RUN: tco %s | FileCheck %s
+
+// Test 1D array with constant shape.
+func.func @array_coor_1d(%addr : !fir.ref<!fir.array<10xi32>>, %idx : index) -> i32 {
+  %c10 = arith.constant 10 : index
+  %shape = fir.shape %c10 : (index) -> !fir.shape<1>
+  %ref = fir.array_coor %addr(%shape) %idx : (!fir.ref<!fir.array<10xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+  %val = fir.load %ref : !fir.ref<i32>
+  return %val : i32
+}
+
+// CHECK-LABEL: define i32 @array_coor_1d
+// CHECK: %[[LB_CHECK:.*]] = icmp sge i64 %{{.*}}, 1
+// CHECK: %[[UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 10
+// CHECK: call void @llvm.assume(i1 %[[LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: call void @llvm.assume(i1 %[[UB_CHECK]]){{.*}}!llvm.array.bounds
+
+// Test 2D array with constant shape.
+func.func @array_coor_2d(%addr : !fir.ref<!fir.array<10x20xi32>>, %i : index, %j : index) -> i32 {
+  %c10 = arith.constant 10 : index
+  %c20 = arith.constant 20 : index
+  %shape = fir.shape %c10, %c20 : (index, index) -> !fir.shape<2>
+  %ref = fir.array_coor %addr(%shape) %i, %j : (!fir.ref<!fir.array<10x20xi32>>, !fir.shape<2>, index, index) -> !fir.ref<i32>
+  %val = fir.load %ref : !fir.ref<i32>
+  return %val : i32
+}
+
+// CHECK-LABEL: define i32 @array_coor_2d
+// First dimension (i) bounds check: 1 <= i <= 10.
+// CHECK: %[[I_LB_CHECK:.*]] = icmp sge i64 %{{.*}}, 1
+// CHECK: %[[I_UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 10
+// CHECK: call void @llvm.assume(i1 %[[I_LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: call void @llvm.assume(i1 %[[I_UB_CHECK]]){{.*}}!llvm.array.bounds
+// Second dimension (j) bounds check: 1 <= j <= 20.
+// CHECK: %[[J_LB_CHECK:.*]] = icmp sge i64 %{{.*}}, 1
+// CHECK: %[[J_UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 20
+// CHECK: call void @llvm.assume(i1 %[[J_LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: call void @llvm.assume(i1 %[[J_UB_CHECK]]){{.*}}!llvm.array.bounds
+
+// Test array with shifted lower bounds.
+func.func @array_coor_shifted(%addr : !fir.ref<!fir.array<10xi32>>, %idx : index) -> i32 {
+  %c5 = arith.constant 5 : index
+  %c10 = arith.constant 10 : index
+  %shape = fir.shape_shift %c5, %c10 : (index, index) -> !fir.shapeshift<1>
+  %ref = fir.array_coor %addr(%shape) %idx : (!fir.ref<!fir.array<10xi32>>, !fir.shapeshift<1>, index) -> !fir.ref<i32>
+  %val = fir.load %ref : !fir.ref<i32>
+  return %val : i32
+}
+
+// CHECK-LABEL: define i32 @array_coor_shifted
+// Bounds check with lower bound 5: 5 <= idx <= 14.
+// CHECK: %[[LB_CHECK:.*]] = icmp sge i64 %{{.*}}, 5
+// CHECK: %[[UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 14
+// CHECK: call void @llvm.assume(i1 %[[LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: call void @llvm.assume(i1 %[[UB_CHECK]]){{.*}}!llvm.array.bounds
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 31636abfff27f..3a907fdf3b2ae 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -773,6 +773,17 @@ amendOperationImpl(Operation &op, ArrayRef<llvm::Instruction *> instructions,
       inst->setMetadata(llvm::LLVMContext::MD_mmra, mmraMd);
     return success();
   }
+
+  // Handle llvm.array_bounds attribute - marks assumes as array bounds checks
+  // that should be dropped before vectorization.
+  if (name == "llvm.array_bounds") {
+    llvm::LLVMContext &ctx = moduleTranslation.getLLVMContext();
+    llvm::MDNode *md = llvm::MDNode::get(ctx, {});
+    for (llvm::Instruction *inst : instructions)
+      inst->setMetadata("llvm.array.bounds", md);
+    return success();
+  }
+
   return success();
 }