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

Sebastian Pop via flang-commits flang-commits at lists.llvm.org
Wed Jul 1 21:27:42 PDT 2026


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

>From 5eabf504e20c205dc91376d1c9e7ded04ccabc03 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.
---
 .../include/flang/Optimizer/CodeGen/CodeGen.h |  6 ++
 .../flang/Optimizer/Passes/CommandLineOpts.h  |  5 ++
 flang/lib/Optimizer/CodeGen/CodeGen.cpp       | 44 +++++++++++
 .../lib/Optimizer/Passes/CommandLineOpts.cpp  |  6 ++
 flang/lib/Optimizer/Passes/Pipelines.cpp      |  5 ++
 flang/test/Fir/array-bounds-assume.fir        | 76 +++++++++++++++++++
 .../mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td   |  8 ++
 .../LLVMIR/LLVMToLLVMIRTranslation.cpp        | 12 +++
 .../test/Target/LLVMIR/llvmir-intrinsics.mlir | 12 +++
 9 files changed, 174 insertions(+)
 create mode 100644 flang/test/Fir/array-bounds-assume.fir

diff --git a/flang/include/flang/Optimizer/CodeGen/CodeGen.h b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
index b7a9397edfe6d..dd4a4b4838f17 100644
--- a/flang/include/flang/Optimizer/CodeGen/CodeGen.h
+++ b/flang/include/flang/Optimizer/CodeGen/CodeGen.h
@@ -65,6 +65,12 @@ struct FIRToLLVMPassOptions {
   // Conversion pass of the MLIR complex dialect.
   Fortran::frontend::CodeGenOptions::ComplexRangeKind ComplexRange =
       Fortran::frontend::CodeGenOptions::ComplexRangeKind::CX_Full;
+
+  // Emit llvm.assume intrinsics describing array bounds constraints for
+  // non-boxed array accesses. These provide extra information to loop
+  // optimizations but also increase IR size, so they are only emitted at
+  // higher optimization levels (see getFIRToLLVMPassOptions).
+  bool emitArrayBoundsAssumes = false;
 };
 
 /// Convert FIR to the LLVM IR dialect with default options.
diff --git a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
index 882f02032a3b8..6ee2e41936840 100644
--- a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
+++ b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
@@ -45,6 +45,11 @@ extern llvm::cl::opt<bool> ignoreMissingTypeDescriptors;
 /// differs from most compilers).
 extern llvm::cl::opt<bool> skipExternalRttiDefinition;
 
+/// Shared option in tools to control the emission of llvm.assume intrinsics
+/// describing array bounds constraints for non-boxed array accesses. Even when
+/// enabled, these assumes are only emitted at higher optimization levels.
+extern llvm::cl::opt<bool> enableArrayBoundsAssumes;
+
 /// Default optimization level used to create Flang pass pipeline is O0.
 extern llvm::OptimizationLevel defaultOptLevel;
 
diff --git a/flang/lib/Optimizer/CodeGen/CodeGen.cpp b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
index 74ad23a050fae..6b6c2612b9b52 100644
--- a/flang/lib/Optimizer/CodeGen/CodeGen.cpp
+++ b/flang/lib/Optimizer/CodeGen/CodeGen.cpp
@@ -2883,6 +2883,13 @@ struct XArrayCoorOpConversion
     mlir::LLVM::GEPNoWrapFlags gepFlags =
         mlir::LLVM::GEPNoWrapFlags::nusw | mlir::LLVM::GEPNoWrapFlags::nuw;
 
+    // An assumed-size array's last dimension has an unknown extent that is
+    // encoded as -1, so no meaningful upper-bound assume can be emitted for it.
+    const bool lastDimIsAssumedSize =
+        !coor.getShape().empty() &&
+        llvm::isa_and_nonnull<fir::AssumedSizeExtentOp>(
+            coor.getShape().back().getDefiningOp());
+
     // For each dimension of the array, generate the offset calculation.
     for (unsigned i = 0; i < rank; ++i, ++indexOffset, ++shapeOffset,
                   ++shiftOffset, sliceOffset += 3, sliceOps += 3) {
@@ -2891,6 +2898,43 @@ 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. This
+      // is only emitted at higher optimization levels (see FIRToLLVMPassOptions
+      // and getFIRToLLVMPassOptions).
+      if (options.emitArrayBoundsAssumes && !baseIsBoxed &&
+          !coor.getShape().empty()) {
+        // Lower bound assume: index >= lb.
+        mlir::Value lbCheck = mlir::LLVM::ICmpOp::create(
+            rewriter, loc, mlir::LLVM::ICmpPredicate::sge, index, lb);
+        auto lbAssume = mlir::LLVM::AssumeOp::create(rewriter, loc, lbCheck);
+        lbAssume->setAttr(mlir::LLVM::AssumeOp::getArrayBoundsAttrName(),
+                          rewriter.getUnitAttr());
+
+        // Upper bound assume: index <= lb + extent - 1. The extent of an
+        // assumed-size array's last dimension is unknown, so skip the
+        // upper-bound assume in that case to avoid a bogus assumption.
+        if (!(i == rank - 1 && lastDimIsAssumedSize)) {
+          mlir::Value extent =
+              integerCast(loc, rewriter, idxTy, operands[shapeOffset]);
+          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);
+          mlir::Value ubCheck = mlir::LLVM::ICmpOp::create(
+              rewriter, loc, mlir::LLVM::ICmpPredicate::sle, index, ub);
+          auto ubAssume = mlir::LLVM::AssumeOp::create(rewriter, loc, ubCheck);
+          ubAssume->setAttr(mlir::LLVM::AssumeOp::getArrayBoundsAttrName(),
+                            rewriter.getUnitAttr());
+        }
+      }
+
       mlir::Value step = one;
       bool normalSlice = isSliced;
       // Compute zero based index in dimension i of the element, applying
diff --git a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
index d461c1b9757b5..946c66435118f 100644
--- a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
+++ b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
@@ -45,6 +45,12 @@ cl::opt<bool> skipExternalRttiDefinition(
                    "other compilation units"),
     cl::Hidden);
 
+cl::opt<bool> enableArrayBoundsAssumes(
+    "enable-array-bounds-assumes",
+    cl::desc("emit llvm.assume intrinsics for array bounds constraints of "
+             "non-boxed array accesses (only at -O2 and above)"),
+    cl::init(true), cl::Hidden);
+
 OptimizationLevel defaultOptLevel{OptimizationLevel::O0};
 
 codegenoptions::DebugInfoKind noDebugInfo{codegenoptions::NoDebugInfo};
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index a4e9d0c227817..40d5f97f2893e 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -130,6 +130,11 @@ getFIRToLLVMPassOptions(const MLIRToLLVMPassPipelineConfig &config) {
   options.typeDescriptorsRenamedForAssembly =
       !disableCompilerGeneratedNamesConversion;
   options.ComplexRange = config.ComplexRange;
+  // Only emit array bounds assumes at -O2 and above, as they increase IR size
+  // for a benefit that mainly helps loop optimizations at higher optimization
+  // levels.
+  options.emitArrayBoundsAssumes =
+      enableArrayBoundsAssumes && config.OptLevel.getSpeedupLevel() >= 2;
   return options;
 }
 
diff --git a/flang/test/Fir/array-bounds-assume.fir b/flang/test/Fir/array-bounds-assume.fir
new file mode 100644
index 0000000000000..ac0b87b9ffd11
--- /dev/null
+++ b/flang/test/Fir/array-bounds-assume.fir
@@ -0,0 +1,76 @@
+// 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: call void @llvm.assume(i1 %[[LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: %[[UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 10
+// 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: call void @llvm.assume(i1 %[[I_LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: %[[I_UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 10
+// 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: call void @llvm.assume(i1 %[[J_LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: %[[J_UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 20
+// 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: call void @llvm.assume(i1 %[[LB_CHECK]]){{.*}}!llvm.array.bounds
+// CHECK: %[[UB_CHECK:.*]] = icmp sle i64 %{{.*}}, 14
+// CHECK: call void @llvm.assume(i1 %[[UB_CHECK]]){{.*}}!llvm.array.bounds
+
+// Test assumed-size array: the extent of the last dimension is unknown
+// (encoded as -1), so only the lower-bound assume is emitted for it; emitting
+// an upper-bound assume would be a bogus assumption.
+func.func @array_coor_assumed_size(%addr : !fir.ref<!fir.array<?xi32>>, %idx : index) -> i32 {
+  %ext = fir.assumed_size_extent : index
+  %shape = fir.shape %ext : (index) -> !fir.shape<1>
+  %ref = fir.array_coor %addr(%shape) %idx : (!fir.ref<!fir.array<?xi32>>, !fir.shape<1>, index) -> !fir.ref<i32>
+  %val = fir.load %ref : !fir.ref<i32>
+  return %val : i32
+}
+
+// CHECK-LABEL: define i32 @array_coor_assumed_size
+// CHECK: %[[LB_CHECK:.*]] = icmp sge i64 %{{.*}}, 1
+// CHECK: call void @llvm.assume(i1 %[[LB_CHECK]]){{.*}}!llvm.array.bounds
+// No upper-bound assume for the assumed-size last dimension.
+// CHECK-NOT: call void @llvm.assume
+// CHECK: ret i32
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index 73629b95a06e5..1d1e86c4bd077 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -655,6 +655,14 @@ def LLVM_AssumeOp
                    "Value":$ptr1, "Value":$ptr2)>
   ];
 
+  let extraClassDeclaration = [{
+    /// Name of the discardable attribute marking an `llvm.assume` as an array
+    /// bounds assumption. It is translated to the `llvm.array.bounds` metadata
+    /// on the corresponding LLVM IR intrinsic call, which passes such as
+    /// DropUnnecessaryAssumes use to selectively drop these assumptions.
+    static StringRef getArrayBoundsAttrName() { return "llvm.array_bounds"; }
+  }];
+
   let hasVerifier = 1;
 }
 
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 7cd8c3c77c15f..fd89fcdec33f7 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -903,6 +903,18 @@ amendOperationImpl(Operation &op, ArrayRef<llvm::Instruction *> instructions,
       inst->setMetadata(llvm::LLVMContext::MD_mmra, mmraMd);
     return success();
   }
+
+  // Handle the array bounds attribute on llvm.assume: it marks assumes as
+  // array bounds checks that should be dropped before vectorization. Translate
+  // it to the "llvm.array.bounds" metadata on the LLVM IR intrinsic call.
+  if (name == LLVM::AssumeOp::getArrayBoundsAttrName()) {
+    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();
 }
 
diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index ea393dd445eda..df356d9ef45c4 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
@@ -1415,6 +1415,15 @@ llvm.func @vector_scmp(%a: vector<4 x i32>, %b: vector<4 x i32>) -> vector<4 x i
   llvm.return %0 : vector<4 x i32>
 }
 
+// The `llvm.array_bounds` unit attribute on `llvm.intr.assume` is translated to
+// the `llvm.array.bounds` metadata (an empty node) on the LLVM IR call.
+// CHECK-LABEL: @assume_array_bounds
+llvm.func @assume_array_bounds(%cond: i1) {
+  // CHECK: call void @llvm.assume(i1 %{{.+}}){{.*}}!llvm.array.bounds ![[BOUNDS:[0-9]+]]
+  llvm.intr.assume %cond : i1 {llvm.array_bounds}
+  llvm.return
+}
+
 // Check that intrinsics are declared with appropriate types.
 // CHECK-DAG: declare float @llvm.fma.f32(float, float, float)
 // CHECK-DAG: declare <8 x float> @llvm.fma.v8f32(<8 x float>, <8 x float>, <8 x float>) #0
@@ -1622,3 +1631,6 @@ llvm.func @vector_scmp(%a: vector<4 x i32>, %b: vector<4 x i32>) -> vector<4 x i
 // CHECK-DAG: declare range(i2 -1, -2) i2 @llvm.scmp.i2.i32(i32, i32)
 // CHECK-DAG: declare range(i32 -1, 2) <4 x i32> @llvm.scmp.v4i32.v4i32(<4 x i32>, <4 x i32>)
 // CHECK-DAG: declare void @llvm.fake.use(...)
+
+// The array bounds metadata node is empty.
+// CHECK-DAG: ![[BOUNDS]] = !{}



More information about the flang-commits mailing list