[Mlir-commits] [llvm] [mlir] [mlir][vector] Fold `in_bounds` for transfers with loop-derived indices (PR #215340)
Dhairyashil R G
llvmlistbot at llvm.org
Sat Aug 29 06:00:43 PDT 2026
https://github.com/dhairyashilRG updated https://github.com/llvm/llvm-project/pull/215340
>From 3e385b1e518fd615a7725e974603b69f10ca0328 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Fri, 28 Aug 2026 22:02:46 +0530
Subject: [PATCH] [mlir][vector] Add a pass to infer the in_bounds attribute
Adds `-vector-infer-in-bounds`, an opt-in pass that marks dimensions of
`vector.transfer_read` and `vector.transfer_write` in-bounds when
value-bounds analysis can prove the transfer, including its starting point,
stays within the source.
The canonicalizer's existing fold only handles constant indices. After
tiling and vectorization the index of a transfer is normally a loop
induction variable or an affine expression of one, so the fold almost never
fires on real code, and every such transfer keeps `in_bounds = false`. That
forces a masked or guarded lowering: on AArch64 NEON, where masked memory
operations are not legal, `ScalarizeMaskedMemIntrin` expands them into
per-lane test-and-branch sequences.
Running `ValueBoundsConstraintSet` from the folder itself is not an option.
`TransferReadOp::fold(FoldAdaptor)` carries no options and no pass context,
so the analysis could not be gated, and every canonicalization of every
function would pay for it. An opt-in pass is what reviewers of the earlier
folder-based version asked for, and it keeps the cost where it is wanted.
The pass uses a `walk` rather than a greedy pattern set: the attribute of one
transfer never affects that of another, so there is no fixpoint to reach and
each operation is visited exactly once. It only ever adds information -- a
dimension already marked in-bounds is left alone, and no dimension is ever
marked out-of-bounds -- so it cannot invalidate an assumption the input
already made.
Bounds are only claimed where they can be proved:
- a dynamic source dimension has no static size to compare against;
- a scalable vector dimension holds `vscale * N` elements, so its static
size is only a lower bound;
- the lower bound of the index must be non-negative, because `in_bounds`
covers the starting point of the transfer;
- the largest index the enclosing loops can produce must still leave room
for a full vector. `closedUB` is required here: the bound wanted is the
largest attainable index, not one past it.
Adds 19 tests covering induction variables, `affine.apply` indices, step
alignment, `scf.for`, tensor sources, permuted and broadcast dimensions, and
negative cases for each of the four conditions above.
---
.../mlir/Dialect/Vector/Transforms/Passes.td | 18 +
.../Dialect/Vector/Transforms/CMakeLists.txt | 2 +
.../Vector/Transforms/VectorInferInBounds.cpp | 154 ++++++++
.../Vector/vector-infer-in-bounds.mlir | 345 ++++++++++++++++++
.../llvm-project-overlay/mlir/BUILD.bazel | 1 +
5 files changed, 520 insertions(+)
create mode 100644 mlir/lib/Dialect/Vector/Transforms/VectorInferInBounds.cpp
create mode 100644 mlir/test/Dialect/Vector/vector-infer-in-bounds.mlir
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/Passes.td b/mlir/include/mlir/Dialect/Vector/Transforms/Passes.td
index 9431a4d8e240f..fc5aa3d0ba52a 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/Passes.td
@@ -39,4 +39,22 @@ def LowerVectorToFromElementsToShuffleTree
let summary = "Lower `vector.to_elements` and `vector.from_elements` to a tree of `vector.shuffle` operations";
}
+def VectorInferInBounds : Pass<"vector-infer-in-bounds", "func::FuncOp"> {
+ let summary = "Infer the `in_bounds` attribute of vector transfer operations";
+ let description = [{
+ Marks dimensions of `vector.transfer_read` and `vector.transfer_write`
+ operations as in-bounds when value-bounds analysis can prove that the
+ transfer, including its starting point, stays within the source.
+
+ The canonicalizer already handles transfers whose indices are constants.
+ This pass additionally handles indices that are loop induction variables,
+ or affine expressions of them, which is the common shape after tiling and
+ vectorization. That analysis is too expensive to run from a folder, so it
+ is opt-in here.
+
+ The pass only ever adds information: a dimension already marked in-bounds
+ is left untouched, and no dimension is ever marked out-of-bounds.
+ }];
+}
+
#endif // MLIR_DIALECT_VECTOR_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
index dfe873f1a1b8d..99f79026668b6 100644
--- a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
@@ -21,6 +21,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
VectorDropLeadUnitDim.cpp
VectorEmulateMaskedLoadStore.cpp
VectorEmulateNarrowType.cpp
+ VectorInferInBounds.cpp
VectorInsertExtractStridedSliceRewritePatterns.cpp
VectorLinearize.cpp
VectorTransferOpTransforms.cpp
@@ -54,6 +55,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
MLIRSubsetOpInterface
MLIRTensorDialect
MLIRTransforms
+ MLIRValueBoundsOpInterface
MLIRVectorDialect
MLIRVectorInterfaces
MLIRVectorUtils
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorInferInBounds.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorInferInBounds.cpp
new file mode 100644
index 0000000000000..2acf3f493fdb0
--- /dev/null
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorInferInBounds.cpp
@@ -0,0 +1,154 @@
+//===- VectorInferInBounds.cpp - Infer in_bounds for transfer ops ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass that infers the `in_bounds` attribute of
+// `vector.transfer_read` and `vector.transfer_write` operations whose indices
+// are not constants, using value-bounds analysis.
+//
+// The canonicalizer already marks a transfer in-bounds when its indices are
+// constants. That covers very little real code: after tiling and vectorization
+// the index of a transfer is typically a loop induction variable, or an affine
+// expression of one, and the canonicalizer gives up. Value-bounds analysis can
+// still bound such an index, but running it from a folder would impose the cost
+// on every canonicalization of every function, so it lives here, in an opt-in
+// pass, instead.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace mlir {
+namespace vector {
+#define GEN_PASS_DEF_VECTORINFERINBOUNDS
+#include "mlir/Dialect/Vector/Transforms/Passes.h.inc"
+} // namespace vector
+} // namespace mlir
+
+#define DEBUG_TYPE "vector-infer-in-bounds"
+
+using namespace mlir;
+using namespace mlir::vector;
+
+namespace {
+
+/// Returns "true" if the transfer of `op` along vector dimension `resultIdx`,
+/// which indexes the source at `indicesIdx`, is provably within the bounds of
+/// the source.
+template <typename TransferOp>
+static bool isProvablyInBounds(TransferOp op, int64_t resultIdx,
+ int64_t indicesIdx) {
+ // A dynamic source dimension has no static size to compare against.
+ if (op.getShapedType().isDynamicDim(indicesIdx))
+ return false;
+ // A scalable vector dimension holds `vscale * N` elements, so its static size
+ // is only a lower bound and cannot prove that the transfer fits.
+ if (op.getVectorType().getScalableDims()[resultIdx])
+ return false;
+
+ int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
+ int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
+ // Largest index at which a full vector still fits. Computed as a subtraction
+ // rather than adding to the index, which could overflow.
+ int64_t maxStart = sourceSize - vectorSize;
+ if (maxStart < 0)
+ return false;
+
+ Value index = op.getIndices()[indicesIdx];
+
+ // The transfer is in bounds if even the largest index the enclosing loops can
+ // produce still leaves room for a full vector. `closedUB` is required because
+ // the bound wanted here is the largest attainable index, not one past it.
+ FailureOr<int64_t> maxIndex = ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::UB, index, /*stopCondition=*/nullptr,
+ ValueBoundsOptions{/*closedUB=*/true});
+ if (failed(maxIndex) || *maxIndex > maxStart)
+ return false;
+
+ // `in_bounds` promises that the transfer stays within the source *including
+ // its starting point*, so the smallest attainable index must be non-negative
+ // as well. Queried only once the upper bound holds, so that an index that
+ // fails that pays for one query rather than two.
+ FailureOr<int64_t> minIndex = ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType::LB, index);
+ return succeeded(minIndex) && *minIndex >= 0;
+}
+
+/// Recomputes the `in_bounds` attribute of `op`, marking a dimension in-bounds
+/// when `isProvablyInBounds` can prove it. Dimensions already marked in-bounds
+/// are left alone: this only ever adds information.
+template <typename TransferOp>
+static void inferInBounds(TransferOp op) {
+ // TODO: Support the 0-d corner case, which has no vector dimension to mark.
+ if (op.getTransferRank() == 0)
+ return;
+
+ AffineMap permutationMap = op.getPermutationMap();
+ bool changed = false;
+ SmallVector<bool, 4> newInBounds;
+ newInBounds.reserve(op.getTransferRank());
+ // Indices of the non-broadcast dims, needed when handling broadcast dims.
+ SmallVector<unsigned> nonBcastDims;
+
+ // 1. Process the non-broadcast dims.
+ for (unsigned i = 0; i < op.getTransferRank(); ++i) {
+ // 1.1. Already in-bounds, nothing to prove.
+ if (op.isDimInBounds(i)) {
+ newInBounds.push_back(true);
+ continue;
+ }
+ // 1.2. Marked out-of-bounds; try to prove otherwise.
+ bool inBounds = false;
+ if (auto dimExpr = dyn_cast<AffineDimExpr>(permutationMap.getResult(i))) {
+ inBounds = isProvablyInBounds(op, /*resultIdx=*/i,
+ /*indicesIdx=*/dimExpr.getPosition());
+ nonBcastDims.push_back(i);
+ }
+ newInBounds.push_back(inBounds);
+ changed |= inBounds;
+ }
+
+ // 2. Handle the broadcast dims. A broadcast dim reads the same element for
+ // every lane, so it is in-bounds exactly when every non-broadcast dim is.
+ bool allNonBcastDimsInBounds = llvm::all_of(
+ nonBcastDims, [&newInBounds](unsigned idx) { return newInBounds[idx]; });
+ if (allNonBcastDimsInBounds) {
+ for (size_t idx : permutationMap.getBroadcastDims()) {
+ changed |= !newInBounds[idx];
+ newInBounds[idx] = true;
+ }
+ }
+
+ if (!changed)
+ return;
+
+ // OpBuilder is only used as a helper to build a BoolArrayAttr.
+ OpBuilder b(op.getContext());
+ op.setInBoundsAttr(b.getBoolArrayAttr(newInBounds));
+}
+
+struct VectorInferInBoundsPass
+ : public vector::impl::VectorInferInBoundsBase<VectorInferInBoundsPass> {
+
+ void runOnOperation() override {
+ // A `walk` rather than a greedy pattern set: the attribute of one transfer
+ // never affects that of another, so there is no fixpoint to reach and each
+ // op needs to be visited exactly once. Value-bounds queries are expensive
+ // enough that re-running them to a fixpoint would be wasteful.
+ getOperation().walk([](Operation *op) {
+ if (auto readOp = dyn_cast<vector::TransferReadOp>(op))
+ inferInBounds(readOp);
+ else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(op))
+ inferInBounds(writeOp);
+ });
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/Vector/vector-infer-in-bounds.mlir b/mlir/test/Dialect/Vector/vector-infer-in-bounds.mlir
new file mode 100644
index 0000000000000..ef24a43aef141
--- /dev/null
+++ b/mlir/test/Dialect/Vector/vector-infer-in-bounds.mlir
@@ -0,0 +1,345 @@
+// RUN: mlir-opt %s -vector-infer-in-bounds -split-input-file | FileCheck %s
+
+// The index is an affine.for induction variable rather than a constant. The
+// largest value it takes is 384, and 384 + 128 == 512 == the memref dim, so the
+// transfer is exactly in bounds.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_from_loop_iv
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true]} : memref<512xf32>, vector<128xf32>
+func.func @fold_transfer_in_bounds_from_loop_iv(%m: memref<512xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to 385 step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<512xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// Same loop, but the memref is one element shorter: 384 + 128 == 512 > 511.
+// The attribute must not be added -- doing so would be a silent out-of-bounds
+// read.
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_off_by_one
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<511xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_off_by_one(%m: memref<511xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to 385 step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<511xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// The induction variable is always a multiple of the step away from the lower
+// bound, so for `0 to 300 step 128` it only ever takes {0, 128, 256} and
+// 256 + 128 == 384 <= 390. A bound of `ub - 1` = 299 would give 427 > 390 and
+// miss this.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_step_alignment
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true]} : memref<390xf32>, vector<128xf32>
+func.func @fold_transfer_in_bounds_step_alignment(%m: memref<390xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to 300 step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<390xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// The index is an affine.apply over an induction variable rather than the
+// induction variable itself: 384 + 64 + 64 == 512.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_affine_apply_index
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true]} : memref<512xf32>, vector<64xf32>
+func.func @fold_transfer_in_bounds_affine_apply_index(%m: memref<512xf32>, %p: f32) -> vector<64xf32> {
+ %acc = arith.constant dense<0.0> : vector<64xf32>
+ %r = affine.for %i = 0 to 385 step 128 iter_args(%a = %acc) -> (vector<64xf32>) {
+ %idx = affine.apply affine_map<(d0) -> (d0 + 64)>(%i)
+ %v = vector.transfer_read %m[%idx], %p : memref<512xf32>, vector<64xf32>
+ %s = arith.addf %a, %v : vector<64xf32>
+ affine.yield %s : vector<64xf32>
+ }
+ return %r : vector<64xf32>
+}
+
+// -----
+
+// `in_bounds` promises that the starting point is in bounds too, so an index
+// that is provably below the memref must not be folded even though it leaves
+// room for a full vector at the top end (-52 + 128 <= 512).
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_negative_constant_index
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<512xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_negative_constant_index(%m: memref<512xf32>, %p: f32) -> vector<128xf32> {
+ %c = arith.constant -50 : index
+ %v = vector.transfer_read %m[%c], %p : memref<512xf32>, vector<128xf32>
+ return %v : vector<128xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_negative_loop_iv
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<512xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_negative_loop_iv(%m: memref<512xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = -100 to -50 step 16 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<512xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// Dynamic loop bound: no constant bound can be derived for the induction
+// variable, so the transfer must stay potentially out-of-bounds.
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_dynamic_loop_bound
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<512xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_dynamic_loop_bound(%m: memref<512xf32>, %p: f32, %n: index) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to %n step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<512xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// The index is bounded, but the source dimension is dynamic, so there is no
+// static size to compare against and nothing can be proved.
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_dynamic_source_dim
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<?xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_dynamic_source_dim(%m: memref<?xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to 385 step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<?xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// A scalable vector dimension must not be folded even when a constant bound is
+// available for the index: `vector<[4]xf32>` reads `4 * vscale` elements.
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_scalable_loop_iv
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<4xf32>, vector<[4]xf32>
+func.func @no_fold_transfer_in_bounds_scalable_loop_iv(%m: memref<4xf32>, %p: f32) -> vector<[4]xf32> {
+ %acc = arith.constant dense<0.0> : vector<[4]xf32>
+ %r = affine.for %i = 0 to 1 iter_args(%a = %acc) -> (vector<[4]xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<4xf32>, vector<[4]xf32>
+ %s = arith.addf %a, %v : vector<[4]xf32>
+ affine.yield %s : vector<[4]xf32>
+ }
+ return %r : vector<[4]xf32>
+}
+
+// -----
+
+// The write path uses the same bound computation, where an unsound fold is an
+// out-of-bounds store.
+
+// CHECK-LABEL: func @fold_transfer_write_in_bounds_from_loop_iv
+// CHECK: vector.transfer_write %{{.*}} {in_bounds = [true]} : vector<128xf32>, memref<512xf32>
+func.func @fold_transfer_write_in_bounds_from_loop_iv(%m: memref<512xf32>, %v: vector<128xf32>) {
+ affine.for %i = 0 to 385 step 128 {
+ vector.transfer_write %v, %m[%i] : vector<128xf32>, memref<512xf32>
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func @no_fold_transfer_write_in_bounds_off_by_one
+// CHECK: vector.transfer_write
+// CHECK-NOT: in_bounds
+// CHECK: : vector<128xf32>, memref<511xf32>
+func.func @no_fold_transfer_write_in_bounds_off_by_one(%m: memref<511xf32>, %v: vector<128xf32>) {
+ affine.for %i = 0 to 385 step 128 {
+ vector.transfer_write %v, %m[%i] : vector<128xf32>, memref<511xf32>
+ }
+ return
+}
+
+// -----
+
+// Both dims driven by induction variables of a 2-D nest.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_2d_nest
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, true]} : memref<64x512xf32>, vector<4x128xf32>
+func.func @fold_transfer_in_bounds_2d_nest(%m: memref<64x512xf32>, %p: f32) -> vector<4x128xf32> {
+ %acc = arith.constant dense<0.0> : vector<4x128xf32>
+ %r = affine.for %i = 0 to 61 step 4 iter_args(%a = %acc) -> (vector<4x128xf32>) {
+ %r2 = affine.for %j = 0 to 385 step 128 iter_args(%b = %a) -> (vector<4x128xf32>) {
+ %v = vector.transfer_read %m[%i, %j], %p : memref<64x512xf32>, vector<4x128xf32>
+ %s = arith.addf %b, %v : vector<4x128xf32>
+ affine.yield %s : vector<4x128xf32>
+ }
+ affine.yield %r2 : vector<4x128xf32>
+ }
+ return %r : vector<4x128xf32>
+}
+
+// -----
+
+// A tensor source goes through the same path as a memref.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_tensor_source
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true]} : tensor<512xf32>, vector<128xf32>
+func.func @fold_transfer_in_bounds_tensor_source(%t: tensor<512xf32>, %p: f32) -> vector<128xf32> {
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = affine.for %i = 0 to 385 step 128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %t[%i], %p : tensor<512xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ affine.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// scf.for already implements ValueBoundsOpInterface, so it benefits from the
+// bound query without any affine-specific support.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_scf_for
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true]} : memref<512xf32>, vector<128xf32>
+func.func @fold_transfer_in_bounds_scf_for(%m: memref<512xf32>, %p: f32) -> vector<128xf32> {
+ %c0 = arith.constant 0 : index
+ %c385 = arith.constant 385 : index
+ %c128 = arith.constant 128 : index
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = scf.for %i = %c0 to %c385 step %c128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<512xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ scf.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_scf_for_off_by_one
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<511xf32>, vector<128xf32>
+func.func @no_fold_transfer_in_bounds_scf_for_off_by_one(%m: memref<511xf32>, %p: f32) -> vector<128xf32> {
+ %c0 = arith.constant 0 : index
+ %c385 = arith.constant 385 : index
+ %c128 = arith.constant 128 : index
+ %acc = arith.constant dense<0.0> : vector<128xf32>
+ %r = scf.for %i = %c0 to %c385 step %c128 iter_args(%a = %acc) -> (vector<128xf32>) {
+ %v = vector.transfer_read %m[%i], %p : memref<511xf32>, vector<128xf32>
+ %s = arith.addf %a, %v : vector<128xf32>
+ scf.yield %s : vector<128xf32>
+ }
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// A transposing permutation map: vector dim 0 comes from memref dim 1 and vice
+// versa. Each vector dim must be checked against the memref dim it actually
+// maps to. Here iv max is 60: dim 1 (512) has room for 4, dim 0 (64) does not
+// have room for 8 (60 + 8 = 68 > 64), so exactly one dim is in bounds.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_transposed_partial
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, false]
+func.func @fold_transfer_in_bounds_transposed_partial(%m: memref<64x512xf32>, %p: f32) -> vector<4x8xf32> {
+ %acc = arith.constant dense<0.0> : vector<4x8xf32>
+ %r = affine.for %i = 0 to 61 step 4 iter_args(%a = %acc) -> (vector<4x8xf32>) {
+ %v = vector.transfer_read %m[%i, %i], %p {permutation_map = affine_map<(d0, d1) -> (d1, d0)>}
+ : memref<64x512xf32>, vector<4x8xf32>
+ %s = arith.addf %a, %v : vector<4x8xf32>
+ affine.yield %s : vector<4x8xf32>
+ }
+ return %r : vector<4x8xf32>
+}
+
+// -----
+
+// Same map, iv max 56: 56 + 4 = 60 <= 512 and 56 + 8 = 64 <= 64, so both.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_transposed_full
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, true]
+func.func @fold_transfer_in_bounds_transposed_full(%m: memref<64x512xf32>, %p: f32) -> vector<4x8xf32> {
+ %acc = arith.constant dense<0.0> : vector<4x8xf32>
+ %r = affine.for %i = 0 to 57 step 8 iter_args(%a = %acc) -> (vector<4x8xf32>) {
+ %v = vector.transfer_read %m[%i, %i], %p {permutation_map = affine_map<(d0, d1) -> (d1, d0)>}
+ : memref<64x512xf32>, vector<4x8xf32>
+ %s = arith.addf %a, %v : vector<4x8xf32>
+ affine.yield %s : vector<4x8xf32>
+ }
+ return %r : vector<4x8xf32>
+}
+
+// -----
+
+// A broadcast vector dim is only in bounds once every non-broadcast dim is.
+// Here the real dim (128 wide, iv max 384, memref dim 512) fits exactly.
+
+// CHECK-LABEL: func @fold_transfer_in_bounds_broadcast_dim
+// CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, true]
+func.func @fold_transfer_in_bounds_broadcast_dim(%m: memref<64x512xf32>, %p: f32) -> vector<4x128xf32> {
+ %acc = arith.constant dense<0.0> : vector<4x128xf32>
+ %r = affine.for %i = 0 to 64 iter_args(%o = %acc) -> (vector<4x128xf32>) {
+ %r2 = affine.for %j = 0 to 385 step 128 iter_args(%a = %o) -> (vector<4x128xf32>) {
+ %v = vector.transfer_read %m[%i, %j], %p {permutation_map = affine_map<(d0, d1) -> (0, d1)>}
+ : memref<64x512xf32>, vector<4x128xf32>
+ %s = arith.addf %a, %v : vector<4x128xf32>
+ affine.yield %s : vector<4x128xf32>
+ }
+ affine.yield %r2 : vector<4x128xf32>
+ }
+ return %r : vector<4x128xf32>
+}
+
+// -----
+
+// One element shorter: the real dim no longer fits, so the broadcast dim must
+// not be claimed either.
+
+// CHECK-LABEL: func @no_fold_transfer_in_bounds_broadcast_dim_oob
+// CHECK: vector.transfer_read
+// CHECK-NOT: in_bounds
+// CHECK: : memref<64x511xf32>, vector<4x128xf32>
+func.func @no_fold_transfer_in_bounds_broadcast_dim_oob(%m: memref<64x511xf32>, %p: f32) -> vector<4x128xf32> {
+ %acc = arith.constant dense<0.0> : vector<4x128xf32>
+ %r = affine.for %i = 0 to 64 iter_args(%o = %acc) -> (vector<4x128xf32>) {
+ %r2 = affine.for %j = 0 to 385 step 128 iter_args(%a = %o) -> (vector<4x128xf32>) {
+ %v = vector.transfer_read %m[%i, %j], %p {permutation_map = affine_map<(d0, d1) -> (0, d1)>}
+ : memref<64x511xf32>, vector<4x128xf32>
+ %s = arith.addf %a, %v : vector<4x128xf32>
+ affine.yield %s : vector<4x128xf32>
+ }
+ affine.yield %r2 : vector<4x128xf32>
+ }
+ return %r : vector<4x128xf32>
+}
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index c64250f0e0324..a4cd8707fd4b8 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -5234,6 +5234,7 @@ cc_library(
":TensorDialect",
":TransformUtils",
":UBDialect",
+ ":ValueBoundsOpInterface",
":VectorDialect",
":VectorEnumsIncGen",
":VectorInterfaces",
More information about the Mlir-commits
mailing list