[Mlir-commits] [mlir] [MLIR][MemRef] Elide load through offset-shift `reinterpret_cast` (PR #197079)
Alan Li
llvmlistbot at llvm.org
Thu May 14 04:59:32 PDT 2026
https://github.com/lialan updated https://github.com/llvm/llvm-project/pull/197079
>From a49169b8777dab2dc04e177b59e749f17dc8d95c Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Thu, 30 Apr 2026 12:13:19 -0700
Subject: [PATCH 01/13] [MLIR][MemRef] Elide load through offset-shift
reinterpret_cast
Generalize ElideReinterpretCast to also fold loads through
reinterpret_casts that differ from their source only by offset (same
rank, same element type, same memory space, same strides, innermost
stride == 1). The cast offset is folded into the consumer load index:
load %rc[%idx] -> load %src[%idx + (rc_offset - src_offset)].
This complements the existing rank-only pattern
(RewriteLoadFromReinterpretCast) and is needed by SPIR-V codegen, where
narrow-type emulation produces flat-memref reinterpret_casts with
runtime offsets that the SPIR-V conversion cannot lower because the
source rtarray pointer type and the destination fixed-size pointer type
differ.
Restricted to rank-1 sources for now; multi-rank can be added later by
linearizing the offset shift across dimensions.
---
.../mlir/Dialect/MemRef/Transforms/Passes.td | 5 +
.../Transforms/ElideReinterpretCast.cpp | 119 +++++++++++++++++-
2 files changed, 121 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
index 3fb0588df395a..54eaf2cf7be3a 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td
@@ -19,6 +19,11 @@ def ElideReinterpretCastPass : Pass<"memref-elide-reinterpret-cast"> {
operations to obtain compatible shapes with equivalent ops that operate on
compatible shapes directly. This simplifies conversion to EmitC.
}];
+ let dependentDialects = [
+ "affine::AffineDialect",
+ "arith::ArithDialect",
+ "memref::MemRefDialect",
+ ];
}
def ExpandOpsPass : Pass<"memref-expand"> {
diff --git a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
index 41dad1384da75..c1fb797d2047d 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ElideReinterpretCast.cpp
@@ -6,11 +6,13 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Transforms/DialectConversion.h"
@@ -453,6 +455,114 @@ struct RewriteLoadFromReinterpretCast
}
};
+/// Returns true when `rc` is a pure offset-shift reinterpret_cast: source and
+/// result have the same rank, the same element type, the same memory space,
+/// and identical per-rank strides; only the offset (and possibly sizes) differ.
+/// In that form, `load %rc[%idx]` is equivalent to a load on the source at an
+/// adjusted index that absorbs the offset difference.
+///
+/// Restricted to rank-1 sources for now to keep the index transformation
+/// straightforward (innermost stride must equal one). Multi-rank cases can be
+/// added later by linearizing the offset shift across dimensions.
+static bool isPureOffsetShiftRC(memref::ReinterpretCastOp rc) {
+ auto inputTy = dyn_cast<MemRefType>(rc.getSource().getType());
+ auto outputTy = dyn_cast<MemRefType>(rc.getType());
+ if (!inputTy || !outputTy)
+ return false;
+
+ if (inputTy.getRank() != 1 || outputTy.getRank() != 1)
+ return false;
+ if (inputTy.getElementType() != outputTy.getElementType())
+ return false;
+ if (inputTy.getMemorySpace() != outputTy.getMemorySpace())
+ return false;
+
+ int64_t inputOffset, outputOffset;
+ SmallVector<int64_t> inputStrides, outputStrides;
+ if (failed(inputTy.getStridesAndOffset(inputStrides, inputOffset)))
+ return false;
+ if (failed(outputTy.getStridesAndOffset(outputStrides, outputOffset)))
+ return false;
+ if (inputStrides != outputStrides)
+ return false;
+ // Only innermost stride == 1 is supported; otherwise the offset shift
+ // cannot be absorbed into a single index addition.
+ if (inputStrides.back() != 1)
+ return false;
+
+ return true;
+}
+
+/// Rewrites `memref.load` through an offset-shift `reinterpret_cast` by
+/// folding the offset difference into the load index on the source memref.
+///
+/// Shape restriction gated by isPureOffsetShiftRC(): rank-1 source and result,
+/// matching element type / memory space / strides, innermost stride == 1.
+/// Sizes and offsets may differ.
+///
+/// BEFORE
+/// %view = memref.reinterpret_cast %src to offset: [%off], sizes: [N],
+/// strides: [1] : memref<?xi8> to memref<Nxi8, strided<[1], offset: ?>>
+/// %v = memref.load %view[%i] : memref<Nxi8, strided<[1], offset: ?>>
+///
+/// AFTER
+/// %adj = arith.addi %off, %i : index // (or affine.apply for folding)
+/// %v = memref.load %src[%adj] : memref<?xi8>
+struct RewriteLoadFromOffsetShiftReinterpretCast
+ : public OpRewritePattern<memref::LoadOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(memref::LoadOp op,
+ PatternRewriter &rewriter) const override {
+ auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
+ if (!rc)
+ return rewriter.notifyMatchFailure(
+ op, "load source is not a memref.reinterpret_cast");
+ if (!isPureOffsetShiftRC(rc))
+ return rewriter.notifyMatchFailure(
+ op, "reinterpret_cast is not a pure offset shift");
+
+ Location loc = op.getLoc();
+ Value src = rc.getSource();
+ auto inputTy = cast<MemRefType>(src.getType());
+
+ // Pull the source memref's offset from its strided layout. If it's
+ // dynamic, materialize it via memref.extract_strided_metadata.
+ int64_t srcStaticOffset;
+ SmallVector<int64_t> srcStaticStrides;
+ [[maybe_unused]] LogicalResult status =
+ inputTy.getStridesAndOffset(srcStaticStrides, srcStaticOffset);
+ assert(succeeded(status) &&
+ "isPureOffsetShiftRC ensured a strided layout");
+
+ OpFoldResult srcOffsetFR;
+ if (ShapedType::isDynamic(srcStaticOffset)) {
+ auto md = memref::ExtractStridedMetadataOp::create(rewriter, loc, src);
+ srcOffsetFR = md.getOffset();
+ } else {
+ srcOffsetFR = rewriter.getIndexAttr(srcStaticOffset);
+ }
+
+ // shift = rcOffset - srcOffset; newIdx = oldIdx + shift. Use
+ // affine.apply for automatic constant folding.
+ OpFoldResult rcOffsetFR = rc.getMixedOffsets().front();
+ OpFoldResult oldIdxFR = getAsOpFoldResult(op.getIndices().front());
+
+ AffineExpr s0, s1, s2;
+ bindSymbols(rewriter.getContext(), s0, s1, s2);
+ OpFoldResult newIdxFR = affine::makeComposedFoldedAffineApply(
+ rewriter, loc, s0 + s1 - s2, {oldIdxFR, rcOffsetFR, srcOffsetFR});
+ Value newIdx = getValueOrCreateConstantIndexOp(rewriter, loc, newIdxFR);
+
+ // If the reinterpret_cast was only used by this load, drop it.
+ if (rc.getResult().hasOneUse())
+ rewriter.eraseOp(rc);
+ rewriter.replaceOpWithNewOp<memref::LoadOp>(op, src, ValueRange{newIdx});
+ return success();
+ }
+};
+
struct ElideReinterpretCastPass
: public memref::impl::ElideReinterpretCastPassBase<
ElideReinterpretCastPass> {
@@ -472,9 +582,11 @@ struct ElideReinterpretCastPass
auto rc = op.getMemRef().getDefiningOp<memref::ReinterpretCastOp>();
if (!rc)
return true;
- return !isPureRankExpansionOrCollapsingRC(rc);
+ return !isPureRankExpansionOrCollapsingRC(rc) &&
+ !isPureOffsetShiftRC(rc);
});
- target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect>();
+ target.addLegalDialect<affine::AffineDialect, arith::ArithDialect,
+ memref::MemRefDialect>();
if (failed(applyPartialConversion(getOperation(), target,
std::move(patterns))))
signalPassFailure();
@@ -485,6 +597,7 @@ struct ElideReinterpretCastPass
void mlir::memref::populateElideReinterpretCastPatterns(
RewritePatternSet &patterns) {
- patterns.add<CopyToScalarLoadAndStore, RewriteLoadFromReinterpretCast>(
+ patterns.add<CopyToScalarLoadAndStore, RewriteLoadFromReinterpretCast,
+ RewriteLoadFromOffsetShiftReinterpretCast>(
patterns.getContext());
}
>From c01dcd86a550cd23a93bd29be69e1073a9bc0220 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Mon, 11 May 2026 18:44:29 -0700
Subject: [PATCH 02/13] [MLIR][MemRef] Add lit tests for offset-shift
reinterpret_cast elision
Extend mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir with positive
and negative tests covering RewriteLoadFromOffsetShiftReinterpretCast.
Positive:
* static src offset + static rc offset + static load index (constant fold).
* dynamic rc offset, static src offset (index forwarded directly).
* static rc offset, dynamic load index (affine.apply with constant shift).
* dynamic src offset materialized via memref.extract_strided_metadata.
* same-offset cast (shift folds to 0, load index passed through).
Negative (pattern must not fire):
* rank > 1 source/result.
* differing strides between source and reinterpret_cast result.
* innermost stride != 1.
---
.../MemRef/elide-reinterpret-cast.mlir | 161 ++++++++++++++++++
1 file changed, 161 insertions(+)
diff --git a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
index 61b6d480ce7a0..45f03576d568a 100644
--- a/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
+++ b/mlir/test/Dialect/MemRef/elide-reinterpret-cast.mlir
@@ -534,3 +534,164 @@ func.func private @negative_diff_non_unit_size(
%0 = memref.load %reinterpret_cast[%c0, %c98] : memref<1x99xf32>
return
}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Positive tests for offset-shift reinterpret_cast
+//
+// `RewriteLoadFromOffsetShiftReinterpretCast` folds a load through a
+// reinterpret_cast that differs from its source only by offset (rank-1, same
+// element type / memory space / strides, innermost stride == 1). The cast
+// offset is absorbed into the consumer load index:
+// load %rc[%idx] -> load %src[%idx + rcOff - srcOff]
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func.func @offset_shift_static_offsets(
+// CHECK-SAME: %[[SRC:.*]]: memref<16xi8>) -> i8
+func.func @offset_shift_static_offsets(%src: memref<16xi8>) -> i8 {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [4], sizes: [8], strides: [1]
+ : memref<16xi8> to memref<8xi8, strided<[1], offset: 4>>
+ %c2 = arith.constant 2 : index
+ // %adj = 2 + 4 - 0 = 6
+ // CHECK: %[[C6:.*]] = arith.constant 6 : index
+ // CHECK: %[[V:.*]] = memref.load %[[SRC]][%[[C6]]] : memref<16xi8>
+ %v = memref.load %rc[%c2] : memref<8xi8, strided<[1], offset: 4>>
+ // CHECK: return %[[V]]
+ return %v : i8
+}
+
+// -----
+
+// CHECK-LABEL: func.func @offset_shift_dynamic_rc_offset(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<?xi8>) -> i8
+func.func @offset_shift_dynamic_rc_offset(%off: index, %src: memref<?xi8>)
+ -> i8 {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [%off], sizes: [8], strides: [1]
+ : memref<?xi8> to memref<8xi8, strided<[1], offset: ?>>
+ %c0 = arith.constant 0 : index
+ // %adj = 0 + %off - 0 = %off
+ // CHECK: %[[V:.*]] = memref.load %[[SRC]][%[[OFF]]] : memref<?xi8>
+ %v = memref.load %rc[%c0] : memref<8xi8, strided<[1], offset: ?>>
+ // CHECK: return %[[V]]
+ return %v : i8
+}
+
+// -----
+
+// CHECK-LABEL: func.func @offset_shift_dynamic_load_index(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<16xi8>) -> i8
+func.func @offset_shift_dynamic_load_index(%i: index, %src: memref<16xi8>)
+ -> i8 {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [3], sizes: [8], strides: [1]
+ : memref<16xi8> to memref<8xi8, strided<[1], offset: 3>>
+ // %adj = %i + 3 - 0 = %i + 3
+ // CHECK: %[[ADJ:.*]] = affine.apply {{.*}}[%[[I]]]
+ // CHECK: %[[V:.*]] = memref.load %[[SRC]][%[[ADJ]]] : memref<16xi8>
+ %v = memref.load %rc[%i] : memref<8xi8, strided<[1], offset: 3>>
+ // CHECK: return %[[V]]
+ return %v : i8
+}
+
+// -----
+
+// Dynamic source offset: must materialize the source offset via
+// memref.extract_strided_metadata before computing the adjusted index.
+//
+// CHECK-LABEL: func.func @offset_shift_dynamic_src_offset(
+// CHECK-SAME: %[[OFF:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<?xi8, strided<[1], offset: ?>>) -> i8
+func.func @offset_shift_dynamic_src_offset(%off: index,
+ %src: memref<?xi8, strided<[1], offset: ?>>) -> i8 {
+ // CHECK: %{{.*}}, %[[SRCOFF:[a-zA-Z0-9_]+]], %{{.*}}, %{{.*}} = memref.extract_strided_metadata %[[SRC]]
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [%off], sizes: [8], strides: [1]
+ : memref<?xi8, strided<[1], offset: ?>>
+ to memref<8xi8, strided<[1], offset: ?>>
+ %c1 = arith.constant 1 : index
+ // %adj = 1 + %off - %srcOff
+ // CHECK: %[[ADJ:.*]] = affine.apply {{.*}}[%[[OFF]], %[[SRCOFF]]]
+ // CHECK: %[[V:.*]] = memref.load %[[SRC]][%[[ADJ]]]
+ %v = memref.load %rc[%c1] : memref<8xi8, strided<[1], offset: ?>>
+ // CHECK: return %[[V]]
+ return %v : i8
+}
+
+// -----
+
+// Same-offset cast (rc offset equals src offset). The shift folds to 0,
+// so the rewritten load index equals the original index.
+//
+// CHECK-LABEL: func.func @offset_shift_same_offset(
+// CHECK-SAME: %[[I:.*]]: index
+// CHECK-SAME: %[[SRC:.*]]: memref<16xi8>) -> i8
+func.func @offset_shift_same_offset(%i: index, %src: memref<16xi8>) -> i8 {
+ // CHECK-NOT: memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [0], sizes: [8], strides: [1]
+ : memref<16xi8> to memref<8xi8, strided<[1]>>
+ // CHECK: %[[V:.*]] = memref.load %[[SRC]][%[[I]]] : memref<16xi8>
+ %v = memref.load %rc[%i] : memref<8xi8, strided<[1]>>
+ // CHECK: return %[[V]]
+ return %v : i8
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Negative tests for offset-shift reinterpret_cast (must NOT rewrite)
+//===----------------------------------------------------------------------===//
+
+// Rank-2 source/result: pattern is restricted to rank-1.
+//
+// CHECK-LABEL: func.func @negative_offset_shift_rank2(
+func.func @negative_offset_shift_rank2(%src: memref<4x4xi8>) -> i8 {
+ // CHECK: %[[RC:.*]] = memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [4], sizes: [2, 2], strides: [4, 1]
+ : memref<4x4xi8> to memref<2x2xi8, strided<[4, 1], offset: 4>>
+ %c0 = arith.constant 0 : index
+ // CHECK: memref.load %[[RC]]
+ %v = memref.load %rc[%c0, %c0] : memref<2x2xi8, strided<[4, 1], offset: 4>>
+ return %v : i8
+}
+
+// -----
+
+// Element type mismatch is invalid IR for reinterpret_cast in general; the
+// allowed case the pattern must reject is a *stride* mismatch.
+//
+// CHECK-LABEL: func.func @negative_offset_shift_diff_stride(
+func.func @negative_offset_shift_diff_stride(
+ %src: memref<16xi8, strided<[1]>>) -> i8 {
+ // CHECK: %[[RC:.*]] = memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [4], sizes: [4], strides: [2]
+ : memref<16xi8, strided<[1]>>
+ to memref<4xi8, strided<[2], offset: 4>>
+ %c0 = arith.constant 0 : index
+ // CHECK: memref.load %[[RC]]
+ %v = memref.load %rc[%c0] : memref<4xi8, strided<[2], offset: 4>>
+ return %v : i8
+}
+
+// -----
+
+// Innermost stride != 1: offset shift cannot be absorbed into a single index
+// addition without scaling.
+//
+// CHECK-LABEL: func.func @negative_offset_shift_inner_stride_ne_one(
+func.func @negative_offset_shift_inner_stride_ne_one(
+ %src: memref<16xi8, strided<[2]>>) -> i8 {
+ // CHECK: %[[RC:.*]] = memref.reinterpret_cast
+ %rc = memref.reinterpret_cast %src to offset: [4], sizes: [4], strides: [2]
+ : memref<16xi8, strided<[2]>>
+ to memref<4xi8, strided<[2], offset: 4>>
+ %c0 = arith.constant 0 : index
+ // CHECK: memref.load %[[RC]]
+ %v = memref.load %rc[%c0] : memref<4xi8, strided<[2], offset: 4>>
+ return %v : i8
+}
+
>From 7f176180caa5f7b7650183f7ba555de1d4f4ab58 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 09:21:50 -0700
Subject: [PATCH 03/13] [MLIR] Add InferIntDivisibilityOpInterface and
divisibility lattice types
Ports the op interface and lattice value types (ConstantIntDivisibility,
IntegerDivisibility, SetIntDivisibilityFn) from IREE's
IREE::Util::IntegerDivisibilityAnalysis into upstream MLIR. No analysis or
op implementations yet; subsequent commits will add those.
Mirrors the structure of InferIntRangeInterface.
---
mlir/include/mlir/Interfaces/CMakeLists.txt | 1 +
.../InferIntDivisibilityOpInterface.h | 120 ++++++++++++++++++
.../InferIntDivisibilityOpInterface.td | 41 ++++++
mlir/lib/Interfaces/CMakeLists.txt | 2 +
.../InferIntDivisibilityOpInterface.cpp | 11 ++
5 files changed, 175 insertions(+)
create mode 100644 mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
create mode 100644 mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.td
create mode 100644 mlir/lib/Interfaces/InferIntDivisibilityOpInterface.cpp
diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt
index 3cbc9df05f3d7..6461c68423c73 100644
--- a/mlir/include/mlir/Interfaces/CMakeLists.txt
+++ b/mlir/include/mlir/Interfaces/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_interface(DerivedAttributeOpInterface)
add_mlir_interface(DestinationStyleOpInterface)
add_mlir_interface(FunctionInterfaces)
add_mlir_interface(IndexingMapOpInterface)
+add_mlir_interface(InferIntDivisibilityOpInterface)
add_mlir_interface(InferIntRangeInterface)
add_mlir_interface(InferStridedMetadataInterface)
add_mlir_interface(InferTypeOpInterface)
diff --git a/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
new file mode 100644
index 0000000000000..6050ccdfbf99c
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
@@ -0,0 +1,120 @@
+//===- InferIntDivisibilityOpInterface.h - Integer Divisibility -*- C++ -*-===//
+//
+// 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 contains definitions of the integer divisibility inference
+// interface defined in `InferIntDivisibilityOpInterface.td`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE_H
+#define MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE_H
+
+#include "mlir/IR/OpDefinition.h"
+#include <numeric>
+#include <optional>
+
+namespace mlir {
+
+/// Statically known divisibility information for an integer SSA value.
+/// Tracks separate divisors for the unsigned and signed interpretations of
+/// the value so that subsequent analyses can use whichever is more precise.
+class ConstantIntDivisibility {
+public:
+ ConstantIntDivisibility() = default;
+ ConstantIntDivisibility(uint64_t udiv, uint64_t sdiv)
+ : udivVal(udiv), sdivVal(sdiv) {}
+
+ bool operator==(const ConstantIntDivisibility &other) const {
+ return udivVal == other.udivVal && sdivVal == other.sdivVal;
+ }
+
+ uint64_t udiv() const { return this->udivVal; }
+ uint64_t sdiv() const { return this->sdivVal; }
+
+ // Returns the union (computed separately for signed and unsigned bounds)
+ // for this range and `other`.
+ ConstantIntDivisibility getUnion(const ConstantIntDivisibility &other) const {
+ return ConstantIntDivisibility(
+ /*udiv=*/std::gcd(udiv(), other.udiv()),
+ /*sdiv=*/std::gcd(sdiv(), other.sdiv()));
+ }
+
+private:
+ uint64_t udivVal;
+ uint64_t sdivVal;
+
+ friend raw_ostream &operator<<(raw_ostream &os,
+ const ConstantIntDivisibility &div);
+};
+
+inline raw_ostream &operator<<(raw_ostream &os,
+ const ConstantIntDivisibility &div) {
+ os << "ConstantIntDivisibility(udiv = " << div.udivVal
+ << ", sdiv = " << div.sdivVal << ")";
+ return os;
+}
+
+/// This lattice value represents the integer divisibility of an SSA value.
+class IntegerDivisibility {
+public:
+ IntegerDivisibility(ConstantIntDivisibility value)
+ : value(std::move(value)) {}
+ IntegerDivisibility(
+ std::optional<ConstantIntDivisibility> value = std::nullopt)
+ : value(std::move(value)) {}
+ // Gets the minimum divisibility of 1 that is used to indicate that the value
+ // cannot be analyzed further.
+ static IntegerDivisibility getMinDivisibility() {
+ return IntegerDivisibility(ConstantIntDivisibility(1, 1));
+ }
+
+ bool isUninitialized() const { return !value.has_value(); }
+ const ConstantIntDivisibility &getValue() const {
+ assert(!isUninitialized());
+ return *value;
+ }
+
+ bool operator==(const IntegerDivisibility &rhs) const {
+ return value == rhs.value;
+ }
+
+ static IntegerDivisibility join(const IntegerDivisibility &lhs,
+ const IntegerDivisibility &rhs) {
+ if (lhs.isUninitialized()) {
+ return rhs;
+ }
+ if (rhs.isUninitialized()) {
+ return lhs;
+ }
+ return IntegerDivisibility(lhs.getValue().getUnion(rhs.getValue()));
+ }
+
+ void print(raw_ostream &os) const { os << value; }
+
+private:
+ std::optional<ConstantIntDivisibility> value;
+};
+
+inline raw_ostream &operator<<(raw_ostream &os,
+ const IntegerDivisibility &div) {
+ div.print(os);
+ return os;
+}
+
+/// The type of the `setResultDivs` callback provided to ops implementing
+/// InferIntDivisibilityOpInterface. It should be called once for each integer
+/// result value and be passed the ConstantIntDivisibility corresponding to
+/// that value.
+using SetIntDivisibilityFn =
+ llvm::function_ref<void(Value, const ConstantIntDivisibility &)>;
+
+} // end namespace mlir
+
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h.inc"
+
+#endif // MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE_H
diff --git a/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.td b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.td
new file mode 100644
index 0000000000000..c665475e0fd7f
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.td
@@ -0,0 +1,41 @@
+//===- InferIntDivisibilityOpInterface.td - Integer Divisibility -*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for divisibility analysis on scalar integers.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE
+#define MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def InferIntDivisibilityOpInterface :
+ OpInterface<"InferIntDivisibilityOpInterface"> {
+ let description = [{
+ Allows operations to participate in integer divisibility analysis.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Infer the divisibility of the results of this op given the
+ divisibility of its arguments. For each result value, the method
+ should call `setResultDivs` with that `Value` as an argument.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"inferResultDivisibility",
+ /*args=*/(ins
+ "::llvm::ArrayRef<::mlir::IntegerDivisibility>":$argDivs,
+ "::mlir::SetIntDivisibilityFn":$setResultDivs)
+ >
+ ];
+}
+
+#endif // MLIR_INTERFACES_INFERINTDIVISIBILITYOPINTERFACE
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index 41e890cb408ba..267c5b945da62 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -9,6 +9,7 @@ set(LLVM_OPTIONAL_SOURCES
FunctionImplementation.cpp
FunctionInterfaces.cpp
IndexingMapOpInterface.cpp
+ InferIntDivisibilityOpInterface.cpp
InferIntRangeInterface.cpp
InferStridedMetadataInterface.cpp
InferTypeOpInterface.cpp
@@ -67,6 +68,7 @@ add_mlir_library(MLIRFunctionInterfaces
add_mlir_interface_library(IndexingMapOpInterface)
add_mlir_interface_library(InferIntRangeInterface)
+add_mlir_interface_library(InferIntDivisibilityOpInterface)
add_mlir_library(MLIRInferStridedMetadataInterface
InferStridedMetadataInterface.cpp
diff --git a/mlir/lib/Interfaces/InferIntDivisibilityOpInterface.cpp b/mlir/lib/Interfaces/InferIntDivisibilityOpInterface.cpp
new file mode 100644
index 0000000000000..acd7cd9530b5c
--- /dev/null
+++ b/mlir/lib/Interfaces/InferIntDivisibilityOpInterface.cpp
@@ -0,0 +1,11 @@
+//===- InferIntDivisibilityOpInterface.cpp - Integer divisibility inference ==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
+
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.cpp.inc"
>From db0de0a0ba6cc499da70e19a723f8e0beb858c31 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 09:34:07 -0700
Subject: [PATCH 04/13] [MLIR] Address review feedback on
InferIntDivisibilityOpInterface
- Alphabetize lib/Interfaces/CMakeLists.txt entry (Div < Range).
- Fix stale "range" wording in ConstantIntDivisibility::getUnion comment.
- Add explicit to IntegerDivisibility's optional-bearing constructor to
match IntegerValueRange.
---
.../include/mlir/Interfaces/InferIntDivisibilityOpInterface.h | 4 ++--
mlir/lib/Interfaces/CMakeLists.txt | 2 +-
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
index 6050ccdfbf99c..374acee05cb10 100644
--- a/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
+++ b/mlir/include/mlir/Interfaces/InferIntDivisibilityOpInterface.h
@@ -37,7 +37,7 @@ class ConstantIntDivisibility {
uint64_t sdiv() const { return this->sdivVal; }
// Returns the union (computed separately for signed and unsigned bounds)
- // for this range and `other`.
+ // for this divisibility and `other`.
ConstantIntDivisibility getUnion(const ConstantIntDivisibility &other) const {
return ConstantIntDivisibility(
/*udiv=*/std::gcd(udiv(), other.udiv()),
@@ -64,7 +64,7 @@ class IntegerDivisibility {
public:
IntegerDivisibility(ConstantIntDivisibility value)
: value(std::move(value)) {}
- IntegerDivisibility(
+ explicit IntegerDivisibility(
std::optional<ConstantIntDivisibility> value = std::nullopt)
: value(std::move(value)) {}
// Gets the minimum divisibility of 1 that is used to indicate that the value
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index 267c5b945da62..d20d290c45c01 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -67,8 +67,8 @@ add_mlir_library(MLIRFunctionInterfaces
)
add_mlir_interface_library(IndexingMapOpInterface)
-add_mlir_interface_library(InferIntRangeInterface)
add_mlir_interface_library(InferIntDivisibilityOpInterface)
+add_mlir_interface_library(InferIntRangeInterface)
add_mlir_library(MLIRInferStridedMetadataInterface
InferStridedMetadataInterface.cpp
>From 25c02c8058b7e21eb326aea9b30f316956a47d50 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 09:38:59 -0700
Subject: [PATCH 05/13] [MLIR] Add IntegerDivisibilityAnalysis
Ports IREE's IntegerDivisibilityAnalysis (sparse forward dataflow analysis
over the InferIntDivisibilityOpInterface) into upstream MLIR. Mirrors the
shape of IntegerRangeAnalysis. The analysis is a silent no-op until per-op
external models are added in a follow-up commit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../DataFlow/IntegerDivisibilityAnalysis.h | 57 ++++++++
mlir/lib/Analysis/CMakeLists.txt | 3 +
.../DataFlow/IntegerDivisibilityAnalysis.cpp | 135 ++++++++++++++++++
3 files changed, 195 insertions(+)
create mode 100644 mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
create mode 100644 mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
new file mode 100644
index 0000000000000..d4f4f6ae5df0b
--- /dev/null
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
@@ -0,0 +1,57 @@
+//===- IntegerDivisibilityAnalysis.h - Integer divisibility -----*- C++ -*-===//
+//
+// 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 declares the dataflow analysis class for integer divisibility
+// inference. Operations participate in the analysis by implementing
+// `InferIntDivisibilityOpInterface`.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_DATAFLOW_INTEGERDIVISIBILITYANALYSIS_H
+#define MLIR_ANALYSIS_DATAFLOW_INTEGERDIVISIBILITYANALYSIS_H
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
+
+#include <optional>
+
+namespace mlir::dataflow {
+
+class IntegerDivisibilityLattice : public Lattice<IntegerDivisibility> {
+public:
+ using Lattice::Lattice;
+};
+
+class IntegerDivisibilityAnalysis
+ : public SparseForwardDataFlowAnalysis<IntegerDivisibilityLattice> {
+public:
+ using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+
+ // At an entry point, set the lattice to the most pessimistic state,
+ // indicating that no further reasoning can be done.
+ void setToEntryState(IntegerDivisibilityLattice *lattice) override;
+
+ // Visit an operation, invoking the transfer function.
+ LogicalResult
+ visitOperation(Operation *op,
+ ArrayRef<const IntegerDivisibilityLattice *> operands,
+ ArrayRef<IntegerDivisibilityLattice *> results) override;
+
+ /// Visit block arguments or operation results of an operation with region
+ /// control-flow for which values are not defined by region control-flow. This
+ /// function tries to infer the divisibility of loop induction variables based
+ /// on known loop bounds and steps.
+ void visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor,
+ ValueRange successorInputs,
+ ArrayRef<IntegerDivisibilityLattice *> argLattices) override;
+};
+
+} // namespace mlir::dataflow
+
+#endif // MLIR_ANALYSIS_DATAFLOW_INTEGERDIVISIBILITYANALYSIS_H
diff --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt
index db10ebcf2c311..596ffaff428b5 100644
--- a/mlir/lib/Analysis/CMakeLists.txt
+++ b/mlir/lib/Analysis/CMakeLists.txt
@@ -13,6 +13,7 @@ set(LLVM_OPTIONAL_SOURCES
DataFlow/ConstantPropagationAnalysis.cpp
DataFlow/DeadCodeAnalysis.cpp
DataFlow/DenseAnalysis.cpp
+ DataFlow/IntegerDivisibilityAnalysis.cpp
DataFlow/IntegerRangeAnalysis.cpp
DataFlow/LivenessAnalysis.cpp
DataFlow/SparseAnalysis.cpp
@@ -37,6 +38,7 @@ add_mlir_library(MLIRAnalysis
DataFlow/ConstantPropagationAnalysis.cpp
DataFlow/DeadCodeAnalysis.cpp
DataFlow/DenseAnalysis.cpp
+ DataFlow/IntegerDivisibilityAnalysis.cpp
DataFlow/IntegerRangeAnalysis.cpp
DataFlow/LivenessAnalysis.cpp
DataFlow/SparseAnalysis.cpp
@@ -53,6 +55,7 @@ add_mlir_library(MLIRAnalysis
MLIRControlFlowInterfaces
MLIRDataLayoutInterfaces
MLIRFunctionInterfaces
+ MLIRInferIntDivisibilityOpInterface
MLIRInferIntRangeInterface
MLIRInferStridedMetadataInterface
MLIRInferTypeOpInterface
diff --git a/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
new file mode 100644
index 0000000000000..84b88c7ad0342
--- /dev/null
+++ b/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
@@ -0,0 +1,135 @@
+//===- IntegerDivisibilityAnalysis.cpp - Integer divisibility ---*- C++ -*-===//
+//
+// 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 defines the dataflow analysis class for integer divisibility
+// inference. Operations participate in the analysis by implementing
+// `InferIntDivisibilityOpInterface`.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+
+#define DEBUG_TYPE "int-divisibility-analysis"
+
+using llvm::dbgs;
+
+namespace mlir::dataflow {
+
+void IntegerDivisibilityAnalysis::setToEntryState(
+ IntegerDivisibilityLattice *lattice) {
+ propagateIfChanged(lattice,
+ lattice->join(IntegerDivisibility::getMinDivisibility()));
+}
+
+LogicalResult IntegerDivisibilityAnalysis::visitOperation(
+ Operation *op, ArrayRef<const IntegerDivisibilityLattice *> operands,
+ ArrayRef<IntegerDivisibilityLattice *> results) {
+ auto inferrable = dyn_cast<InferIntDivisibilityOpInterface>(op);
+ if (!inferrable) {
+ setAllToEntryStates(results);
+ return success();
+ }
+
+ LLVM_DEBUG(dbgs() << "Inferring divisibility for " << *op << "\n");
+ auto argDivs = llvm::map_to_vector(
+ operands, [](const IntegerDivisibilityLattice *lattice) {
+ return lattice->getValue();
+ });
+ auto joinCallback = [&](Value v, const IntegerDivisibility &newDiv) {
+ auto result = dyn_cast<OpResult>(v);
+ if (!result) {
+ return;
+ }
+ assert(llvm::is_contained(op->getResults(), result));
+
+ LLVM_DEBUG(dbgs() << "Inferred divisibility " << newDiv << "\n");
+ IntegerDivisibilityLattice *lattice = results[result.getResultNumber()];
+ IntegerDivisibility oldDiv = lattice->getValue();
+
+ ChangeResult changed = lattice->join(newDiv);
+
+ // Catch loop results with loop variant bounds and conservatively make
+ // them [-inf, inf] so we don't circle around infinitely often (because
+ // the dataflow analysis in MLIR doesn't attempt to work out trip counts
+ // and often can't).
+ bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
+ return op->hasTrait<OpTrait::IsTerminator>();
+ });
+ if (isYieldedResult && !oldDiv.isUninitialized() &&
+ !(lattice->getValue() == oldDiv)) {
+ LLVM_DEBUG(llvm::dbgs() << "Loop variant loop result detected\n");
+ changed |= lattice->join(IntegerDivisibility::getMinDivisibility());
+ }
+ propagateIfChanged(lattice, changed);
+ };
+
+ inferrable.inferResultDivisibility(argDivs, joinCallback);
+ return success();
+}
+
+void IntegerDivisibilityAnalysis::visitNonControlFlowArguments(
+ Operation *op, const RegionSuccessor &successor, ValueRange successorInputs,
+ ArrayRef<IntegerDivisibilityLattice *> argLattices) {
+ // Get the constant divisibility, or query the lattice for Values.
+ auto getDivFromOfr = [&](std::optional<OpFoldResult> ofr, Block *block,
+ bool isUnsigned) -> uint64_t {
+ if (ofr.has_value()) {
+ if (auto constBound = getConstantIntValue(*ofr)) {
+ return constBound.value();
+ }
+ auto value = cast<Value>(ofr.value());
+ const IntegerDivisibilityLattice *lattice =
+ getLatticeElementFor(getProgramPointBefore(block), value);
+ if (lattice != nullptr && !lattice->getValue().isUninitialized()) {
+ return isUnsigned ? lattice->getValue().getValue().udiv()
+ : lattice->getValue().getValue().sdiv();
+ }
+ }
+ return isUnsigned
+ ? IntegerDivisibility::getMinDivisibility().getValue().udiv()
+ : IntegerDivisibility::getMinDivisibility().getValue().sdiv();
+ };
+
+ // Infer bounds for loop arguments that have static bounds
+ if (auto loop = dyn_cast<LoopLikeOpInterface>(op)) {
+ std::optional<SmallVector<Value>> ivs = loop.getLoopInductionVars();
+ std::optional<SmallVector<OpFoldResult>> lbs = loop.getLoopLowerBounds();
+ std::optional<SmallVector<OpFoldResult>> steps = loop.getLoopSteps();
+ if (!ivs || !lbs || !steps) {
+ return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
+ op, successor, successorInputs, argLattices);
+ }
+ for (auto [iv, lb, step] : llvm::zip_equal(*ivs, *lbs, *steps)) {
+ IntegerDivisibilityLattice *ivEntry = getLatticeElement(iv);
+ Block *block = iv.getParentBlock();
+ uint64_t stepUDiv = getDivFromOfr(step, block, /*unsigned=*/true);
+ uint64_t stepSDiv = getDivFromOfr(step, block, /*unsigned=*/false);
+ uint64_t lbUDiv = getDivFromOfr(lb, block, /*unsigned=*/true);
+ uint64_t lbSDiv = getDivFromOfr(lb, block, /*unsigned=*/false);
+ ConstantIntDivisibility lbDiv(lbUDiv, lbSDiv);
+ ConstantIntDivisibility stepDiv(stepUDiv, stepSDiv);
+
+ // Loop induction variables are computed as `lb + i * step`. The
+ // divisibility for `i * step` is just the divisibility of `step`, so
+ // the total divisibility is obtained by unioning the step divisibility
+ // with the lower bound divisibility, which takes the GCD of the two.
+ ConstantIntDivisibility ivDiv = stepDiv.getUnion(lbDiv);
+ propagateIfChanged(ivEntry, ivEntry->join(ivDiv));
+ }
+ return;
+ }
+
+ return SparseForwardDataFlowAnalysis::visitNonControlFlowArguments(
+ op, successor, successorInputs, argLattices);
+}
+
+} // namespace mlir::dataflow
>From 148de26bf348c834bbd614e3bf7cce97105232ca Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 09:44:43 -0700
Subject: [PATCH 06/13] [MLIR] Address review feedback on
IntegerDivisibilityAnalysis
- Replace stale "[-inf, inf]" widening comment with divisibility-accurate
wording.
- Promote two method comments to /// Doxygen style.
- Add class-level /// doc comments for IntegerDivisibilityLattice and
IntegerDivisibilityAnalysis, mirroring IntegerRangeAnalysis.
- Drop redundant <optional> include from the header.
---
.../DataFlow/IntegerDivisibilityAnalysis.h | 17 ++++++++++++-----
.../DataFlow/IntegerDivisibilityAnalysis.cpp | 8 ++++----
2 files changed, 16 insertions(+), 9 deletions(-)
diff --git a/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
index d4f4f6ae5df0b..3a877647490a3 100644
--- a/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
+++ b/mlir/include/mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h
@@ -18,25 +18,32 @@
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
-#include <optional>
-
namespace mlir::dataflow {
+/// This lattice element represents the integer divisibility of an SSA value.
class IntegerDivisibilityLattice : public Lattice<IntegerDivisibility> {
public:
using Lattice::Lattice;
};
+/// Integer divisibility analysis determines, for each integer-typed SSA
+/// value, a divisor that the value is guaranteed to be a multiple of. It
+/// uses operations that implement `InferIntDivisibilityOpInterface` and
+/// also sets the divisibility of induction variables of loops with known
+/// lower bounds and steps.
+///
+/// This analysis depends on DeadCodeAnalysis, and will be a silent no-op
+/// if DeadCodeAnalysis is not loaded in the same solver context.
class IntegerDivisibilityAnalysis
: public SparseForwardDataFlowAnalysis<IntegerDivisibilityLattice> {
public:
using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
- // At an entry point, set the lattice to the most pessimistic state,
- // indicating that no further reasoning can be done.
+ /// At an entry point, set the lattice to the most pessimistic state,
+ /// indicating that no further reasoning can be done.
void setToEntryState(IntegerDivisibilityLattice *lattice) override;
- // Visit an operation, invoking the transfer function.
+ /// Visit an operation, invoking the transfer function.
LogicalResult
visitOperation(Operation *op,
ArrayRef<const IntegerDivisibilityLattice *> operands,
diff --git a/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp b/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
index 84b88c7ad0342..5be40f297aaa4 100644
--- a/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
+++ b/mlir/lib/Analysis/DataFlow/IntegerDivisibilityAnalysis.cpp
@@ -57,10 +57,10 @@ LogicalResult IntegerDivisibilityAnalysis::visitOperation(
ChangeResult changed = lattice->join(newDiv);
- // Catch loop results with loop variant bounds and conservatively make
- // them [-inf, inf] so we don't circle around infinitely often (because
- // the dataflow analysis in MLIR doesn't attempt to work out trip counts
- // and often can't).
+ // Catch loop results with loop-variant divisibility and conservatively
+ // set them to divisibility 1 (no information) so we don't ratchet
+ // indefinitely (the dataflow analysis in MLIR doesn't attempt to work
+ // out trip counts and often can't).
bool isYieldedResult = llvm::any_of(v.getUsers(), [](Operation *op) {
return op->hasTrait<OpTrait::IsTerminator>();
});
>From a2fee168baa4497f99ba113425b1968de309af7a Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 09:50:44 -0700
Subject: [PATCH 07/13] [MLIR] Add InferIntDivisibilityOpInterface external
models for arith and affine
Ports the per-op divisibility models from IREE's UtilExternalModels.cpp.
Each dialect now provides registerInferIntDivisibilityOpInterfaceExternalModels.
Registration into the global DialectRegistry is wired up in a follow-up commit.
Arith ops covered: ConstantOp, AddIOp, SubIOp, MulIOp, DivUIOp,
MinUIOp/MaxUIOp/MinSIOp/MaxSIOp, SelectOp.
Affine ops covered: AffineApplyOp, AffineMinOp, AffineMaxOp,
AffineDelinearizeIndexOp (the last using AffineExprDivisibilityFinder).
---
.../IR/InferIntDivisibilityOpInterfaceImpl.h | 21 +
.../IR/InferIntDivisibilityOpInterfaceImpl.h | 21 +
mlir/lib/Dialect/Affine/IR/CMakeLists.txt | 2 +
.../InferIntDivisibilityOpInterfaceImpl.cpp | 368 ++++++++++++++++++
mlir/lib/Dialect/Arith/IR/CMakeLists.txt | 3 +
.../InferIntDivisibilityOpInterfaceImpl.cpp | 162 ++++++++
6 files changed, 577 insertions(+)
create mode 100644 mlir/include/mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h
create mode 100644 mlir/include/mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h
create mode 100644 mlir/lib/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.cpp
create mode 100644 mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
diff --git a/mlir/include/mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h
new file mode 100644
index 0000000000000..560bef63b56a8
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h
@@ -0,0 +1,21 @@
+//===- InferIntDivisibilityOpInterfaceImpl.h --------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_AFFINE_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
+#define MLIR_DIALECT_AFFINE_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace affine {
+void registerInferIntDivisibilityOpInterfaceExternalModels(
+ DialectRegistry ®istry);
+} // namespace affine
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
diff --git a/mlir/include/mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h
new file mode 100644
index 0000000000000..0909790b8b7d0
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h
@@ -0,0 +1,21 @@
+//===- InferIntDivisibilityOpInterfaceImpl.h --------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_ARITH_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
+#define MLIR_DIALECT_ARITH_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace arith {
+void registerInferIntDivisibilityOpInterfaceExternalModels(
+ DialectRegistry ®istry);
+} // namespace arith
+} // namespace mlir
+
+#endif // MLIR_DIALECT_ARITH_IR_INFERINTDIVISIBILITYOPINTERFACEIMPL_H
diff --git a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
index 566bc060e5d38..1caf2fa396797 100644
--- a/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/IR/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRAffineDialect
AffineMemoryOpInterfaces.cpp
AffineOps.cpp
AffineValueMap.cpp
+ InferIntDivisibilityOpInterfaceImpl.cpp
InferIntRangeInterfaceImpls.cpp
ValueBoundsOpInterfaceImpl.cpp
@@ -16,6 +17,7 @@ add_mlir_dialect_library(MLIRAffineDialect
MLIRArithDialect
MLIRDialectUtils
MLIRIR
+ MLIRInferIntDivisibilityOpInterface
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
MLIRLoopLikeInterface
diff --git a/mlir/lib/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.cpp
new file mode 100644
index 0000000000000..30850cf0a4df0
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.cpp
@@ -0,0 +1,368 @@
+//===- InferIntDivisibilityOpInterfaceImpl.cpp ----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/AffineExprVisitor.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
+
+#include <cstdlib>
+#include <numeric>
+
+using namespace mlir;
+
+namespace {
+
+static ConstantIntDivisibility
+getDivisibilityOfOperand(Value v, IntegerDivisibility divisibility) {
+ if (!divisibility.isUninitialized()) {
+ return divisibility.getValue();
+ }
+ APInt intVal;
+ if (matchPattern(v, m_ConstantInt(&intVal))) {
+ uint64_t udiv = intVal.getZExtValue();
+ uint64_t sdiv = std::abs(intVal.getSExtValue());
+ return ConstantIntDivisibility(udiv, sdiv);
+ }
+ return ConstantIntDivisibility(1, 1);
+}
+
+/// Visits affine expressions and recursively calculates the divisibilities of
+/// each subexpression. The final divisibilities of the expression and its
+/// subexpressions will be stored in the map for which a reference is provided
+/// to the AffineExprDivisibilityFinder (i.e., `divisibilityMap`).
+class AffineExprDivisibilityFinder
+ : public AffineExprVisitor<AffineExprDivisibilityFinder,
+ ConstantIntDivisibility> {
+public:
+ using ExprDivisibilityMap =
+ llvm::DenseMap<AffineExpr, ConstantIntDivisibility>;
+ AffineExprDivisibilityFinder(ExprDivisibilityMap &divisibilityMap)
+ : divisibilityMap(divisibilityMap) {}
+
+ ConstantIntDivisibility visitConstantExpr(AffineConstantExpr expr) {
+ // Constant expressions are trivial, since they are always static.
+ uint64_t constValue = std::abs(expr.getValue());
+ return ConstantIntDivisibility(constValue, constValue);
+ }
+
+ ConstantIntDivisibility visitDimExpr(AffineDimExpr expr) {
+ // Dim expressions cannot be analyzed further, so return the divisibility
+ // in `divisibilityMap` if it has been populated by the caller, or fallback
+ // to the minimum divisibility.
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ return IntegerDivisibility::getMinDivisibility().getValue();
+ }
+
+ ConstantIntDivisibility visitSymbolExpr(AffineSymbolExpr expr) {
+ // Symbol expressions cannot be analyzed further, so return the divisibility
+ // in `divisibilityMap` if it has been populated by the caller, or fallback
+ // to the minimum divisibility.
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ return IntegerDivisibility::getMinDivisibility().getValue();
+ }
+
+ /// Infer the divisibility of an addition or subtraction expression by
+ /// recursively visiting the LHS and RHS, and then unioning the results.
+ ConstantIntDivisibility visitAddExpr(AffineBinaryOpExpr expr) {
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ // The divisibility of an addition is the GCD of its constituents'
+ // divisibilities.
+ ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
+ ConstantIntDivisibility rhsDiv = visit(expr.getRHS());
+ return lhsDiv.getUnion(rhsDiv);
+ }
+
+ /// Infer the divisibility of a multiplication expression by recursively
+ /// visiting the LHS and RHS, and then multiplying the results.
+ ConstantIntDivisibility visitMulExpr(AffineBinaryOpExpr expr) {
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ // The divisibility of a multiplication is the product of its constituents'
+ // divisibilities.
+ ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
+ ConstantIntDivisibility rhsDiv = visit(expr.getRHS());
+ return ConstantIntDivisibility(lhsDiv.udiv() * rhsDiv.udiv(),
+ lhsDiv.sdiv() * rhsDiv.sdiv());
+ }
+
+ ConstantIntDivisibility visitFloorDivExpr(AffineBinaryOpExpr expr) {
+ return visitDivExpr(expr);
+ }
+
+ ConstantIntDivisibility visitCeilDivExpr(AffineBinaryOpExpr expr) {
+ return visitDivExpr(expr);
+ }
+
+ /// Infer the divisibility of a mod expression. If the RHS is a constant,
+ /// the result divisibility is gcd(lhs_divisibility, rhs_constant), since
+ /// (d * k) mod c is always divisible by gcd(d, c). Furthermore, if the
+ /// LHS divisibility is itself divisible by the constant (i.e., d % c == 0),
+ /// then (d * k) mod c is always zero, represented as divisibility 0.
+ ConstantIntDivisibility visitModExpr(AffineBinaryOpExpr expr) {
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ auto constRhs = dyn_cast<AffineConstantExpr>(expr.getRHS());
+ if (!constRhs || constRhs.getValue() == 0) {
+ return ConstantIntDivisibility(1, 1);
+ }
+ auto constValue = static_cast<uint64_t>(std::abs(constRhs.getValue()));
+ ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
+ // If the LHS is always a multiple of constValue, x mod constValue is
+ // always zero. Divisibility 0 is the lattice top ("divides everything").
+ uint64_t modUDiv = (lhsDiv.udiv() % constValue == 0)
+ ? 0
+ : std::gcd(lhsDiv.udiv(), constValue);
+ uint64_t modSDiv = (lhsDiv.sdiv() % constValue == 0)
+ ? 0
+ : std::gcd(lhsDiv.sdiv(), constValue);
+ return ConstantIntDivisibility(modUDiv, modSDiv);
+ }
+
+private:
+ ConstantIntDivisibility visitInvalidExpr(AffineBinaryOpExpr expr) {
+ return IntegerDivisibility::getMinDivisibility().getValue();
+ }
+
+ /// Helper shared by ceildiv and floordiv implementations. Returns the minimum
+ /// divisibility as a fallback if the divisor is not a constant, because the
+ /// divisibility cannot be inferred in this case. If the divisor is a
+ /// constant, then this function recursively visits the dividend, and returns
+ /// the quotient of the dividend's divisibility with the divisor.
+ ConstantIntDivisibility visitDivExpr(AffineBinaryOpExpr expr) {
+ if (divisibilityMap.contains(expr)) {
+ return divisibilityMap[expr];
+ }
+ auto constRhs = dyn_cast<AffineConstantExpr>(expr.getRHS());
+ // Division by zero is undefined, so return the minimum divisibility.
+ if (!constRhs || constRhs.getValue() == 0) {
+ return ConstantIntDivisibility(1, 1);
+ }
+ auto constValue = static_cast<uint64_t>(std::abs(constRhs.getValue()));
+ ConstantIntDivisibility lhsDiv = visit(expr.getLHS());
+ uint64_t divUDiv =
+ lhsDiv.udiv() % constValue == 0 ? lhsDiv.udiv() / constValue : 1;
+ uint64_t divSDiv =
+ lhsDiv.sdiv() % constValue == 0 ? lhsDiv.sdiv() / constValue : 1;
+ return ConstantIntDivisibility(divUDiv, divSDiv);
+ }
+
+ ExprDivisibilityMap &divisibilityMap;
+};
+
+/// Returns the divisibilities of each AffineMap result based on the
+/// divisibilities of its dims and symbols. The `dimAndSymbolDivisibilities`
+/// should contain the divisibilities of the dims, followed by the
+/// divisibilities of the symbols in ascending order by their positions.
+static SmallVector<ConstantIntDivisibility> getResultDivisibilities(
+ AffineMap map,
+ ArrayRef<ConstantIntDivisibility> dimAndSymbolDivisibilities) {
+ // Seed the AffineExprDivisibilityFinder with the dimAndSymbolDivisibilities.
+ llvm::DenseMap<AffineExpr, ConstantIntDivisibility> exprDivisibilityMap;
+ SmallVector<AffineExpr> inputExprs;
+ inputExprs.append(llvm::map_to_vector(
+ llvm::seq<int64_t>(map.getNumDims()),
+ [&](int64_t dim) { return getAffineDimExpr(dim, map.getContext()); }));
+ inputExprs.append(llvm::map_to_vector(
+ llvm::seq<int64_t>(map.getNumSymbols()),
+ [&](int64_t sym) { return getAffineSymbolExpr(sym, map.getContext()); }));
+ for (auto [expr, divisibility] :
+ llvm::zip_equal(inputExprs, dimAndSymbolDivisibilities)) {
+ exprDivisibilityMap[expr] = divisibility;
+ }
+ AffineExprDivisibilityFinder divisibilityFinder(exprDivisibilityMap);
+
+ // Walk each result expression and compute their divisibilities.
+ SmallVector<ConstantIntDivisibility> resultDivisibilities;
+ for (AffineExpr resultExpr : map.getResults()) {
+ resultDivisibilities.push_back(divisibilityFinder.visit(resultExpr));
+ }
+ return resultDivisibilities;
+}
+
+struct AffineApplyInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ AffineApplyInferIntDivisibilityOpInterface, affine::AffineApplyOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto affineApplyOp = cast<affine::AffineApplyOp>(op);
+ SmallVector<ConstantIntDivisibility> operandDivisibilities;
+ for (auto [operand, divisibility] :
+ llvm::zip(affineApplyOp.getOperands(), argDivs)) {
+ operandDivisibilities.push_back(
+ getDivisibilityOfOperand(operand, divisibility));
+ }
+
+ SmallVector<ConstantIntDivisibility> resultDivisibilities =
+ getResultDivisibilities(affineApplyOp.getMap(), operandDivisibilities);
+ for (auto [result, divisibility] :
+ llvm::zip_equal(affineApplyOp->getResults(), resultDivisibilities)) {
+ setResultDivs(result, divisibility);
+ }
+ }
+};
+
+/// Infer the result divisibility of an affine.min or affine.max operation
+/// based on its operand divisibilities. The result divisibility is the GCD
+/// of the divisibilities of each of the affine map results, because the result
+/// of the affine.min/max op could be any of these results.
+template <typename MinOrMaxTy>
+static void
+inferAffineMinOrMaxResultDivisibility(MinOrMaxTy minOrMaxOp,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) {
+ static_assert(
+ llvm::is_one_of<MinOrMaxTy, affine::AffineMinOp,
+ affine::AffineMaxOp>::value,
+ "MinOrMaxTy must be affine::AffineMinOp or affine::AffineMaxOp");
+ SmallVector<ConstantIntDivisibility> operandDivisibilities;
+ for (auto [operand, divisibility] :
+ llvm::zip(minOrMaxOp.getOperands(), argDivs)) {
+ operandDivisibilities.push_back(
+ getDivisibilityOfOperand(operand, divisibility));
+ }
+
+ SmallVector<ConstantIntDivisibility> resultDivisibilities =
+ getResultDivisibilities(minOrMaxOp.getMap(), operandDivisibilities);
+
+ ConstantIntDivisibility resultDivisibility =
+ resultDivisibilities.pop_back_val();
+ for (auto divisibility : resultDivisibilities) {
+ resultDivisibility = resultDivisibility.getUnion(divisibility);
+ }
+ setResultDivs(minOrMaxOp.getResult(), resultDivisibility);
+}
+
+struct AffineMinInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ AffineMinInferIntDivisibilityOpInterface, affine::AffineMinOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto affineMinOp = cast<affine::AffineMinOp>(op);
+ inferAffineMinOrMaxResultDivisibility(affineMinOp, argDivs, setResultDivs);
+ }
+};
+
+struct AffineMaxInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ AffineMaxInferIntDivisibilityOpInterface, affine::AffineMaxOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto affineMaxOp = cast<affine::AffineMaxOp>(op);
+ inferAffineMinOrMaxResultDivisibility(affineMaxOp, argDivs, setResultDivs);
+ }
+};
+
+struct AffineDelinearizeIndexInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ AffineDelinearizeIndexInferIntDivisibilityOpInterface,
+ affine::AffineDelinearizeIndexOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto delinearizeOp = cast<affine::AffineDelinearizeIndexOp>(op);
+ MLIRContext *ctx = op->getContext();
+
+ // Operands are: [linear_index, dynamic_basis_values...]
+ ConstantIntDivisibility linearDiv =
+ getDivisibilityOfOperand(delinearizeOp.getLinearIndex(), argDivs[0]);
+
+ ArrayRef<int64_t> staticBasis = delinearizeOp.getStaticBasis();
+ int64_t numResults = delinearizeOp.getNumResults();
+
+ // Build affine expressions for each result.
+ // Dim 0 = linear index, symbols = dynamic basis values.
+ AffineExpr linearExpr = getAffineDimExpr(0, ctx);
+
+ // Collect operand divisibilities: [linear_index_div, dynamic_basis_divs...]
+ SmallVector<ConstantIntDivisibility> operandDivs;
+ operandDivs.push_back(linearDiv);
+
+ // Map static/dynamic basis values to affine expressions.
+ int64_t dynIdx = 0;
+ SmallVector<AffineExpr> basisExprs;
+ for (int64_t i = 0, e = static_cast<int64_t>(staticBasis.size()); i < e;
+ ++i) {
+ if (ShapedType::isDynamic(staticBasis[i])) {
+ basisExprs.push_back(getAffineSymbolExpr(dynIdx, ctx));
+ operandDivs.push_back(getDivisibilityOfOperand(
+ delinearizeOp.getDynamicBasis()[dynIdx], argDivs[1 + dynIdx]));
+ dynIdx++;
+ } else {
+ basisExprs.push_back(getAffineConstantExpr(staticBasis[i], ctx));
+ }
+ }
+
+ // The computation basis skips the outer bound if present.
+ bool hasOuter = delinearizeOp.hasOuterBound();
+ int64_t basisStart = hasOuter ? 1 : 0;
+
+ // Each result[i] can be expressed as an affine expression of the linear
+ // index using the effective basis (after dropping outer bound if present).
+ // Effective basis B[k] = basisExprs[basisStart + k], for k = 0..N-2.
+ // Stride s[i] = product of B[i..N-2] = product of
+ // basisExprs[basisStart+i .. end].
+ //
+ // result[0] = x floordiv s[0]
+ // result[i>0] = (x floordiv s[i]) mod B[i-1]
+ // For i=N-1, s[N-1]=1, so result[N-1] = x mod B[N-2].
+
+ AffineExpr stride = getAffineConstantExpr(1, ctx);
+ for (int64_t i = numResults - 1; i >= 0; --i) {
+ AffineExpr resultExpr;
+ if (i == 0) {
+ resultExpr = linearExpr.floorDiv(stride);
+ } else {
+ resultExpr =
+ (linearExpr.floorDiv(stride)) % basisExprs[basisStart + i - 1];
+ }
+
+ AffineMap resultMap = AffineMap::get(1, dynIdx, resultExpr, ctx);
+ SmallVector<ConstantIntDivisibility> divs =
+ getResultDivisibilities(resultMap, operandDivs);
+ setResultDivs(delinearizeOp.getResult(i), divs[0]);
+
+ if (i > 0) {
+ stride = basisExprs[basisStart + i - 1] * stride;
+ }
+ }
+ }
+};
+
+} // namespace
+
+void mlir::affine::registerInferIntDivisibilityOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *context, AffineDialect *dialect) {
+ AffineApplyOp::attachInterface<AffineApplyInferIntDivisibilityOpInterface>(
+ *context);
+ AffineMinOp::attachInterface<AffineMinInferIntDivisibilityOpInterface>(
+ *context);
+ AffineMaxOp::attachInterface<AffineMaxInferIntDivisibilityOpInterface>(
+ *context);
+ AffineDelinearizeIndexOp::attachInterface<
+ AffineDelinearizeIndexInferIntDivisibilityOpInterface>(*context);
+ });
+}
diff --git a/mlir/lib/Dialect/Arith/IR/CMakeLists.txt b/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
index 4beb99ccfdfba..3423e11a7d0f0 100644
--- a/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
@@ -1,6 +1,7 @@
set(LLVM_OPTIONAL_SOURCES
ArithOps.cpp
ArithDialect.cpp
+ InferIntDivisibilityOpInterfaceImpl.cpp
InferIntRangeInterfaceImpls.cpp
ValueBoundsOpInterfaceImpl.cpp
)
@@ -12,6 +13,7 @@ add_public_tablegen_target(MLIRArithCanonicalizationIncGen)
add_mlir_dialect_library(MLIRArithDialect
ArithOps.cpp
ArithDialect.cpp
+ InferIntDivisibilityOpInterfaceImpl.cpp
InferIntRangeInterfaceImpls.cpp
ADDITIONAL_HEADER_DIRS
@@ -24,6 +26,7 @@ add_mlir_dialect_library(MLIRArithDialect
LINK_LIBS PUBLIC
MLIRCastInterfaces
MLIRDialect
+ MLIRInferIntDivisibilityOpInterface
MLIRInferIntRangeCommon
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
diff --git a/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp b/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
new file mode 100644
index 0000000000000..21254362e224a
--- /dev/null
+++ b/mlir/lib/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.cpp
@@ -0,0 +1,162 @@
+//===- InferIntDivisibilityOpInterfaceImpl.cpp ----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/InferIntDivisibilityOpInterface.h"
+
+#include <cstdlib>
+
+using namespace mlir;
+
+namespace {
+
+static ConstantIntDivisibility
+getDivisibilityOfOperand(Value v, IntegerDivisibility divisibility) {
+ if (!divisibility.isUninitialized()) {
+ return divisibility.getValue();
+ }
+ APInt intVal;
+ if (matchPattern(v, m_ConstantInt(&intVal))) {
+ uint64_t udiv = intVal.getZExtValue();
+ uint64_t sdiv = std::abs(intVal.getSExtValue());
+ return ConstantIntDivisibility(udiv, sdiv);
+ }
+ return ConstantIntDivisibility(1, 1);
+}
+
+/// Helper for binary arith ops whose result divisibility is the GCD (union) of
+/// their operands' divisibilities. This covers add, sub, min, and max.
+template <typename OpTy>
+struct ArithBinaryGCDInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<OpTy>, OpTy> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto binOp = cast<OpTy>(op);
+ auto lhsDiv = getDivisibilityOfOperand(binOp.getLhs(), argDivs[0]);
+ auto rhsDiv = getDivisibilityOfOperand(binOp.getRhs(), argDivs[1]);
+ setResultDivs(binOp.getResult(), lhsDiv.getUnion(rhsDiv));
+ }
+};
+
+/// For arith.select, the result divisibility is the GCD of the true and false
+/// operands' divisibilities. The condition (operand 0) is i1 and irrelevant.
+struct ArithSelectInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ ArithSelectInferIntDivisibilityOpInterface, arith::SelectOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto selectOp = cast<arith::SelectOp>(op);
+ // argDivs[0] is the condition (i1), argDivs[1] is true, argDivs[2] is
+ // false.
+ auto trueDiv =
+ getDivisibilityOfOperand(selectOp.getTrueValue(), argDivs[1]);
+ auto falseDiv =
+ getDivisibilityOfOperand(selectOp.getFalseValue(), argDivs[2]);
+ setResultDivs(selectOp.getResult(), trueDiv.getUnion(falseDiv));
+ }
+};
+
+struct ArithConstantInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ ArithConstantInferIntDivisibilityOpInterface, arith::ConstantOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto constOp = cast<arith::ConstantOp>(op);
+ auto constAttr = dyn_cast_if_present<IntegerAttr>(constOp.getValue());
+ if (constAttr) {
+ const APInt &value = constAttr.getValue();
+ uint64_t udiv = value.getZExtValue();
+ uint64_t sdiv = std::abs(value.getSExtValue());
+ setResultDivs(constOp.getResult(), ConstantIntDivisibility(udiv, sdiv));
+ }
+ }
+};
+
+struct ArithMulIInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ ArithMulIInferIntDivisibilityOpInterface, arith::MulIOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto mulOp = cast<arith::MulIOp>(op);
+
+ auto lhsDivisibility = getDivisibilityOfOperand(mulOp.getLhs(), argDivs[0]);
+ auto rhsDivisibility = getDivisibilityOfOperand(mulOp.getRhs(), argDivs[1]);
+
+ uint64_t mulUDiv = lhsDivisibility.udiv() * rhsDivisibility.udiv();
+ uint64_t mulSDiv = lhsDivisibility.sdiv() * rhsDivisibility.sdiv();
+
+ setResultDivs(mulOp.getResult(), ConstantIntDivisibility(mulUDiv, mulSDiv));
+ }
+};
+
+struct ArithDivUIInferIntDivisibilityOpInterface
+ : InferIntDivisibilityOpInterface::ExternalModel<
+ ArithDivUIInferIntDivisibilityOpInterface, arith::DivUIOp> {
+
+ void inferResultDivisibility(Operation *op,
+ ArrayRef<IntegerDivisibility> argDivs,
+ SetIntDivisibilityFn setResultDivs) const {
+ auto divOp = cast<arith::DivUIOp>(op);
+
+ APInt intVal;
+ if (!matchPattern(divOp.getRhs(), m_ConstantInt(&intVal))) {
+ return;
+ }
+
+ auto lhsDivisibility = getDivisibilityOfOperand(divOp.getLhs(), argDivs[0]);
+
+ uint64_t divUDiv = lhsDivisibility.udiv() % intVal.getZExtValue() == 0
+ ? lhsDivisibility.udiv() / intVal.getZExtValue()
+ : 1;
+ uint64_t divSDiv =
+ lhsDivisibility.sdiv() % std::abs(intVal.getSExtValue()) == 0
+ ? lhsDivisibility.sdiv() / std::abs(intVal.getSExtValue())
+ : 1;
+
+ setResultDivs(divOp, ConstantIntDivisibility(divUDiv, divSDiv));
+ }
+};
+
+} // namespace
+
+void mlir::arith::registerInferIntDivisibilityOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *context, ArithDialect *dialect) {
+ ConstantOp::attachInterface<ArithConstantInferIntDivisibilityOpInterface>(
+ *context);
+ MulIOp::attachInterface<ArithMulIInferIntDivisibilityOpInterface>(*context);
+ DivUIOp::attachInterface<ArithDivUIInferIntDivisibilityOpInterface>(
+ *context);
+ AddIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<AddIOp>>(*context);
+ SubIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<SubIOp>>(*context);
+ MinUIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<MinUIOp>>(*context);
+ MaxUIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<MaxUIOp>>(*context);
+ MinSIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<MinSIOp>>(*context);
+ MaxSIOp::attachInterface<
+ ArithBinaryGCDInferIntDivisibilityOpInterface<MaxSIOp>>(*context);
+ SelectOp::attachInterface<ArithSelectInferIntDivisibilityOpInterface>(
+ *context);
+ });
+}
>From a1eb0e8702156441da3493db7bca170b4cfca50a Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 10:17:44 -0700
Subject: [PATCH 08/13] [MLIR] Replace assumeAligned with
divisibility-analysis-driven gating
Strips the bool assumeAligned parameter from
populateMemRefNarrowTypeEmulationPatterns and
populateVectorNarrowTypeEmulationPatterns. Replaces it with an optional
DataFlowSolver* that pattern matchers consult to decide whether a dynamic
offset is provably aligned.
Three-state behavior:
- Provably aligned (udiv % elementsPerByte == 0): take the aligned path.
- Provably misaligned (udiv > 1 and udiv % elementsPerByte != 0): reject
with notifyMatchFailure (MemRef) or fall through to the partial-store
path (Vector).
- Opaque / no solver: behave as the old assumeAligned=true.
Updates TestEmulateNarrowTypePass to construct and load the analysis
before populating patterns.
---
.../Dialect/MemRef/Transforms/Transforms.h | 18 +-
.../Vector/Transforms/VectorRewritePatterns.h | 18 +-
.../Dialect/MemRef/Transforms/CMakeLists.txt | 1 +
.../MemRef/Transforms/EmulateNarrowType.cpp | 211 ++++++++++--------
.../Dialect/Vector/Transforms/CMakeLists.txt | 1 +
.../Transforms/VectorEmulateNarrowType.cpp | 102 +++++----
mlir/test/lib/Dialect/MemRef/CMakeLists.txt | 1 +
.../Dialect/MemRef/TestEmulateNarrowType.cpp | 26 ++-
8 files changed, 232 insertions(+), 146 deletions(-)
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
index 4d6c54d74d2a9..f8a37eceb7b5c 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
@@ -18,6 +18,7 @@
#include "llvm/ADT/STLFunctionalExtras.h"
namespace mlir {
+class DataFlowSolver;
class OpBuilder;
class RewritePatternSet;
class RewriterBase;
@@ -89,15 +90,20 @@ void populateMemRefWideIntEmulationConversions(
/// over wider types.
/// When `disableAtomicRMW` is true, the store patterns generate non-atomic
/// read-modify-write sequences instead of atomic operations.
-/// When `assumeAligned` is true, `memref.subview` and
-/// `memref.reinterpret_cast` patterns accept dynamic offsets under the
-/// alignment contract that the caller guarantees those offsets are a multiple
-/// of `dstBits / srcBits`. When false (the default), dynamic offsets are
-/// rejected to preserve soundness for callers that cannot prove divisibility.
+/// When `solver` is non-null, `memref.subview` and `memref.reinterpret_cast`
+/// patterns consult the `IntegerDivisibilityLattice` for each dynamic offset:
+/// - If the analysis proves the offset is a multiple of `dstBits / srcBits`,
+/// the pattern proceeds.
+/// - If the analysis proves the offset is *not* such a multiple, the pattern
+/// rejects the op via `notifyMatchFailure`.
+/// - If the lattice is uninitialized/opaque, the pattern proceeds (the same
+/// behavior as the legacy `assumeAligned=true` path).
+/// When `solver` is null, the patterns behave as if every dynamic offset were
+/// opaque, i.e. they proceed under the alignment contract.
void populateMemRefNarrowTypeEmulationPatterns(
const arith::NarrowTypeEmulationConverter &typeConverter,
RewritePatternSet &patterns, bool disableAtomicRMW = false,
- bool assumeAligned = false);
+ DataFlowSolver *solver = nullptr);
/// Appends type conversions for emulating memref operations over narrow types
/// with ops over wider types.
diff --git a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
index 7d6d565d5a4f4..097239d3de8f8 100644
--- a/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
+++ b/mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
@@ -20,6 +20,7 @@
namespace mlir {
class ConversionTarget;
+class DataFlowSolver;
class RewritePatternSet;
class TypeConverter;
@@ -394,15 +395,20 @@ void populateVectorMaskMaterializationPatterns(RewritePatternSet &patterns,
/// Appends patterns for emulating vector operations over narrow types with ops
/// over wider types. The `disableAtomicRMW` indicates whether to use a normal
/// read-modify-write sequence instead of using `memref.generic_atomic_rmw` to
-/// perform subbyte storing. When `assumeAligned` is true, store offsets are
-/// assumed to be aligned to container element boundaries, so a store whose
-/// source vector fills whole container elements is emitted as a simple
-/// bitcast + store without checking the offset. Stores that are not divisible
-/// in size are rejected.
+/// perform subbyte storing. When `solver` is non-null, store offsets are
+/// queried against the `IntegerDivisibilityLattice` to decide whether to take
+/// the aligned fast path (bitcast + store) or fall through to the partial-
+/// store path:
+/// - If the lattice proves the dynamic offset is a multiple of the number
+/// of emulated elements per container element, take the fast path.
+/// - If the lattice proves it is *not* such a multiple, fall through to the
+/// partial-store path which is sound for any offset.
+/// - If the lattice is opaque or `solver` is null, take the fast path (this
+/// matches the legacy `assumeAligned=true` behavior).
void populateVectorNarrowTypeEmulationPatterns(
const arith::NarrowTypeEmulationConverter &typeConverter,
RewritePatternSet &patterns, bool disableAtomicRMW = false,
- bool assumeAligned = false);
+ DataFlowSolver *solver = nullptr);
/// Populates patterns for both MeMref flattening and Vector narrow type
/// emulation.
diff --git a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
index 1c5e07f89b338..340a8afb02c7d 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
@@ -27,6 +27,7 @@ add_mlir_dialect_library(MLIRMemRefTransforms
LINK_LIBS PUBLIC
MLIRAffineTransforms
MLIRAffineUtils
+ MLIRAnalysis
MLIRArithDialect
MLIRArithTransforms
MLIRDialectUtils
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index a11e14faa5475..9f59a27bb6106 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -7,6 +7,8 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
@@ -33,16 +35,18 @@ using namespace mlir;
/// Converts a memref::ReinterpretCastOp to the converted type. The result
/// memref is linearized to a rank-1 byte view (or rank-0 if the source is
-/// rank-0). When `assumeAligned` is true, dynamic offsets are accepted under
-/// the alignment contract that the caller guarantees the offset is a multiple
-/// of `dstBits / srcBits`; statically-provable misalignment is rejected.
-/// When `assumeAligned` is false, dynamic offsets are rejected outright since
-/// divisibility cannot be proven from the IR alone.
+/// rank-0). When `solver` is non-null, the dynamic offset (if any) is checked
+/// against the `IntegerDivisibilityLattice`: if the lattice proves the offset
+/// is not a multiple of `dstBits / srcBits`, the pattern rejects via
+/// `notifyMatchFailure`; otherwise (proven multiple or opaque) the pattern
+/// proceeds. When `solver` is null the pattern always proceeds and relies on
+/// the caller's alignment contract. Statically-provable misalignment of the
+/// folded offset is still rejected below.
static LogicalResult
convertCastingOp(ConversionPatternRewriter &rewriter,
memref::ReinterpretCastOp::Adaptor adaptor,
memref::ReinterpretCastOp op, MemRefType newTy,
- bool assumeAligned) {
+ DataFlowSolver *solver) {
if (newTy == op.getType()) {
return rewriter.notifyMatchFailure(
op, "result type was not converted by narrow-type emulation");
@@ -74,14 +78,26 @@ convertCastingOp(ConversionPatternRewriter &rewriter,
op, "result memref is not row-major contiguous");
}
- // Reject dynamic offsets unless the caller has opted into the alignment
- // contract via `assumeAligned`. Without it we cannot prove the offset is a
- // multiple of `dstBits / srcBits`.
- if (!assumeAligned &&
- llvm::is_contained(op.getStaticOffsets(), ShapedType::kDynamic)) {
- return rewriter.notifyMatchFailure(
- op, "dynamic offsets require assumeAligned=true to ensure the offset "
- "is a multiple of dstBits / srcBits");
+ // For a dynamic offset, consult the divisibility lattice if a solver was
+ // provided. Three-state behavior:
+ // - lattice proves `udiv % elementsPerByte == 0` (or is opaque): proceed.
+ // - lattice proves `udiv > 1 && udiv % elementsPerByte != 0`: reject.
+ // - solver == nullptr: proceed (legacy `assumeAligned=true` behavior).
+ int64_t elementsPerByte = dstBits / srcBits;
+ ArrayRef<int64_t> staticOffsets = op.getStaticOffsets();
+ if (solver && !staticOffsets.empty() &&
+ staticOffsets[0] == ShapedType::kDynamic) {
+ Value dynOffset = op.getOffsets()[0];
+ const dataflow::IntegerDivisibilityLattice *lattice =
+ solver->lookupState<dataflow::IntegerDivisibilityLattice>(dynOffset);
+ if (lattice && !lattice->getValue().isUninitialized()) {
+ uint64_t udiv = lattice->getValue().getValue().udiv();
+ if (udiv > 1 && udiv % elementsPerByte != 0) {
+ return rewriter.notifyMatchFailure(
+ op, "dynamic offset is provably not a multiple of "
+ "`dstBits / srcBits`");
+ }
+ }
}
Location loc = op.getLoc();
@@ -447,15 +463,15 @@ struct ConvertMemRefMemorySpaceCast final
// ConvertMemRefReinterpretCast
//===----------------------------------------------------------------------===//
-/// Forwards to `convertCastingOp`, which enforces all preconditions.
-/// `assumeAligned` is propagated from the populate entry point and controls
-/// acceptance of dynamic offsets.
+/// Forwards to `convertCastingOp`, which enforces all preconditions. The
+/// optional `solver` is forwarded so that dynamic offsets can be checked
+/// against the divisibility lattice.
struct ConvertMemRefReinterpretCast final
: OpConversionPattern<memref::ReinterpretCastOp> {
ConvertMemRefReinterpretCast(const TypeConverter &typeConverter,
- MLIRContext *context, bool assumeAligned)
+ MLIRContext *context, DataFlowSolver *solver)
: OpConversionPattern<memref::ReinterpretCastOp>(typeConverter, context),
- assumeAligned(assumeAligned) {}
+ solver(solver) {}
LogicalResult
matchAndRewrite(memref::ReinterpretCastOp op, OpAdaptor adaptor,
@@ -468,11 +484,11 @@ struct ConvertMemRefReinterpretCast final
llvm::formatv("failed to convert memref type: {0}", op.getType()));
}
- return convertCastingOp(rewriter, adaptor, op, newTy, assumeAligned);
+ return convertCastingOp(rewriter, adaptor, op, newTy, solver);
}
private:
- bool assumeAligned;
+ DataFlowSolver *solver;
};
//===----------------------------------------------------------------------===//
@@ -574,17 +590,18 @@ struct ConvertMemrefStore final : OpConversionPattern<memref::StoreOp> {
//===----------------------------------------------------------------------===//
/// Emulating narrow ints on subview have limited support, supporting only
-/// static sizes and stride of 1. When `assumeAligned` is true, dynamic
-/// offsets are accepted under the alignment contract that the caller
-/// guarantees the offset is a multiple of `dstBits / srcBits`. Without that
-/// opt-in, dynamic offsets are rejected. Ideally, the subview should be
-/// folded away before running narrow type emulation, and this pattern should
-/// only run for cases that can't be folded.
+/// static sizes and stride of 1. When `solver` is non-null, each dynamic
+/// offset is queried against the `IntegerDivisibilityLattice`; if any dynamic
+/// offset is provably not a multiple of `dstBits / srcBits` the pattern is
+/// rejected. When `solver` is null, dynamic offsets are accepted under the
+/// caller's alignment contract. Ideally, the subview should be folded away
+/// before running narrow type emulation, and this pattern should only run
+/// for cases that can't be folded.
struct ConvertMemRefSubview final : OpConversionPattern<memref::SubViewOp> {
ConvertMemRefSubview(const TypeConverter &typeConverter, MLIRContext *context,
- bool assumeAligned)
+ DataFlowSolver *solver)
: OpConversionPattern<memref::SubViewOp>(typeConverter, context),
- assumeAligned(assumeAligned) {}
+ solver(solver) {}
LogicalResult
matchAndRewrite(memref::SubViewOp subViewOp, OpAdaptor adaptor,
@@ -627,14 +644,32 @@ struct ConvertMemRefSubview final : OpConversionPattern<memref::SubViewOp> {
"dynamic size is not supported");
}
- // Reject dynamic offsets unless the caller has opted into the alignment
- // contract via `assumeAligned`.
- if (!assumeAligned && llvm::is_contained(subViewOp.getStaticOffsets(),
- ShapedType::kDynamic)) {
- return rewriter.notifyMatchFailure(
- subViewOp,
- "dynamic offsets require assumeAligned=true to ensure the offset "
- "is a multiple of dstBits / srcBits");
+ // For each dynamic offset, consult the divisibility lattice (if a solver
+ // was provided). Three-state behavior:
+ // - lattice proves `udiv % elementsPerByte == 0` (or is opaque): proceed.
+ // - lattice proves `udiv > 1 && udiv % elementsPerByte != 0`: reject.
+ // - solver == nullptr: proceed (legacy `assumeAligned=true` behavior).
+ int64_t elementsPerByte = dstBits / srcBits;
+ if (solver) {
+ ArrayRef<int64_t> staticOffsets = subViewOp.getStaticOffsets();
+ ValueRange dynOffsets = subViewOp.getOffsets();
+ unsigned dynIdx = 0;
+ for (int64_t staticOff : staticOffsets) {
+ if (staticOff != ShapedType::kDynamic)
+ continue;
+ Value dynOffset = dynOffsets[dynIdx++];
+ const dataflow::IntegerDivisibilityLattice *lattice =
+ solver->lookupState<dataflow::IntegerDivisibilityLattice>(
+ dynOffset);
+ if (lattice && !lattice->getValue().isUninitialized()) {
+ uint64_t udiv = lattice->getValue().getValue().udiv();
+ if (udiv > 1 && udiv % elementsPerByte != 0) {
+ return rewriter.notifyMatchFailure(
+ subViewOp, "dynamic offset is provably not a multiple of "
+ "`dstBits / srcBits`");
+ }
+ }
+ }
}
// Transform the offsets, sizes and strides according to the emulation.
@@ -666,7 +701,7 @@ struct ConvertMemRefSubview final : OpConversionPattern<memref::SubViewOp> {
}
private:
- bool assumeAligned;
+ DataFlowSolver *solver;
};
//===----------------------------------------------------------------------===//
@@ -726,7 +761,8 @@ struct ConvertMemRefExpandShape final
void memref::populateMemRefNarrowTypeEmulationPatterns(
const arith::NarrowTypeEmulationConverter &typeConverter,
- RewritePatternSet &patterns, bool disableAtomicRMW, bool assumeAligned) {
+ RewritePatternSet &patterns, bool disableAtomicRMW,
+ DataFlowSolver *solver) {
// Populate `memref.*` conversion patterns.
patterns
@@ -737,7 +773,7 @@ void memref::populateMemRefNarrowTypeEmulationPatterns(
ConvertMemRefAssumeAlignment, ConvertMemRefMemorySpaceCast>(
typeConverter, patterns.getContext());
patterns.add<ConvertMemRefSubview, ConvertMemRefReinterpretCast>(
- typeConverter, patterns.getContext(), assumeAligned);
+ typeConverter, patterns.getContext(), solver);
patterns.insert<ConvertMemrefStore>(typeConverter, patterns.getContext(),
disableAtomicRMW);
memref::populateResolveExtractStridedMetadataPatterns(patterns);
@@ -763,53 +799,52 @@ static SmallVector<int64_t> getLinearizedShape(MemRefType ty, int srcBits,
void memref::populateMemRefNarrowTypeEmulationConversions(
arith::NarrowTypeEmulationConverter &typeConverter) {
- typeConverter.addConversion(
- [&typeConverter](MemRefType ty) -> std::optional<Type> {
- Type elementType = ty.getElementType();
- if (!elementType.isIntOrFloat())
- return ty;
-
- unsigned width = elementType.getIntOrFloatBitWidth();
- unsigned loadStoreWidth = typeConverter.getLoadStoreBitwidth();
- if (width >= loadStoreWidth)
- return ty;
-
- // Currently only handle innermost stride being 1, checking
- SmallVector<int64_t> strides;
- int64_t offset;
- if (failed(ty.getStridesAndOffset(strides, offset)))
- return nullptr;
- if (!strides.empty() && strides.back() != 1)
- return nullptr;
-
- auto newElemTy = IntegerType::get(
- ty.getContext(), loadStoreWidth,
- elementType.isInteger()
- ? cast<IntegerType>(elementType).getSignedness()
- : IntegerType::SignednessSemantics::Signless);
- if (!newElemTy)
- return nullptr;
-
- StridedLayoutAttr layoutAttr;
- // If the offset is 0, we do not need a strided layout as the stride is
- // 1, so we only use the strided layout if the offset is not 0.
- if (offset != 0) {
- if (offset == ShapedType::kDynamic) {
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
- } else {
- // Check if the number of bytes are a multiple of the loadStoreWidth
- // and if so, divide it by the loadStoreWidth to get the offset.
- if ((offset * width) % loadStoreWidth != 0)
- return std::nullopt;
- offset = (offset * width) / loadStoreWidth;
-
- layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
- ArrayRef<int64_t>{1});
- }
- }
+ typeConverter.addConversion([&typeConverter](
+ MemRefType ty) -> std::optional<Type> {
+ Type elementType = ty.getElementType();
+ if (!elementType.isIntOrFloat())
+ return ty;
+
+ unsigned width = elementType.getIntOrFloatBitWidth();
+ unsigned loadStoreWidth = typeConverter.getLoadStoreBitwidth();
+ if (width >= loadStoreWidth)
+ return ty;
+
+ // Currently only handle innermost stride being 1, checking
+ SmallVector<int64_t> strides;
+ int64_t offset;
+ if (failed(ty.getStridesAndOffset(strides, offset)))
+ return nullptr;
+ if (!strides.empty() && strides.back() != 1)
+ return nullptr;
+
+ auto newElemTy = IntegerType::get(
+ ty.getContext(), loadStoreWidth,
+ elementType.isInteger() ? cast<IntegerType>(elementType).getSignedness()
+ : IntegerType::SignednessSemantics::Signless);
+ if (!newElemTy)
+ return nullptr;
+
+ StridedLayoutAttr layoutAttr;
+ // If the offset is 0, we do not need a strided layout as the stride is
+ // 1, so we only use the strided layout if the offset is not 0.
+ if (offset != 0) {
+ if (offset == ShapedType::kDynamic) {
+ layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+ ArrayRef<int64_t>{1});
+ } else {
+ // Check if the number of bytes are a multiple of the loadStoreWidth
+ // and if so, divide it by the loadStoreWidth to get the offset.
+ if ((offset * width) % loadStoreWidth != 0)
+ return std::nullopt;
+ offset = (offset * width) / loadStoreWidth;
+
+ layoutAttr = StridedLayoutAttr::get(ty.getContext(), offset,
+ ArrayRef<int64_t>{1});
+ }
+ }
- return MemRefType::get(getLinearizedShape(ty, width, loadStoreWidth),
- newElemTy, layoutAttr, ty.getMemorySpace());
- });
+ return MemRefType::get(getLinearizedShape(ty, width, loadStoreWidth),
+ newElemTy, layoutAttr, ty.getMemorySpace());
+ });
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
index 112a1db6fe93b..da6bd01e4b801 100644
--- a/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt
@@ -38,6 +38,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
MLIRAffineDialect
MLIRAffineAnalysis
MLIRAffineUtils
+ MLIRAnalysis
MLIRArithDialect
MLIRDialectUtils
MLIRGPUDialect
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index 583cda7ac2810..f8d5e26fc6e7a 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -16,6 +16,8 @@
/// TODO: Support for non-powers-of-two.
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
@@ -113,28 +115,27 @@ static FailureOr<Operation *> getCompressedMaskOp(OpBuilder &rewriter,
auto newMaskType = VectorType::get(maskShape, rewriter.getI1Type());
std::optional<Operation *> newMask =
TypeSwitch<Operation *, std::optional<Operation *>>(maskOp)
- .Case(
- [&](vector::CreateMaskOp createMaskOp)
- -> std::optional<Operation *> {
- OperandRange maskOperands = createMaskOp.getOperands();
- // The `vector.create_mask` op creates a mask arrangement
- // without any zeros at the front. Also, because
- // `numFrontPadElems` is strictly smaller than
- // `numSrcElemsPerDest`, the compressed mask generated by
- // padding the original mask by `numFrontPadElems` will not
- // have any zeros at the front as well.
- AffineExpr s0;
- bindSymbols(rewriter.getContext(), s0);
- s0 = (s0 + numFrontPadElems).ceilDiv(numSrcElemsPerDest);
- OpFoldResult origIndex = getAsOpFoldResult(maskOperands.back());
- OpFoldResult maskIndex = affine::makeComposedFoldedAffineApply(
- rewriter, loc, s0, origIndex);
- SmallVector<Value> newMaskOperands(maskOperands.drop_back());
- newMaskOperands.push_back(
- getValueOrCreateConstantIndexOp(rewriter, loc, maskIndex));
- return vector::CreateMaskOp::create(rewriter, loc, newMaskType,
- newMaskOperands);
- })
+ .Case([&](vector::CreateMaskOp createMaskOp)
+ -> std::optional<Operation *> {
+ OperandRange maskOperands = createMaskOp.getOperands();
+ // The `vector.create_mask` op creates a mask arrangement
+ // without any zeros at the front. Also, because
+ // `numFrontPadElems` is strictly smaller than
+ // `numSrcElemsPerDest`, the compressed mask generated by
+ // padding the original mask by `numFrontPadElems` will not
+ // have any zeros at the front as well.
+ AffineExpr s0;
+ bindSymbols(rewriter.getContext(), s0);
+ s0 = (s0 + numFrontPadElems).ceilDiv(numSrcElemsPerDest);
+ OpFoldResult origIndex = getAsOpFoldResult(maskOperands.back());
+ OpFoldResult maskIndex = affine::makeComposedFoldedAffineApply(
+ rewriter, loc, s0, origIndex);
+ SmallVector<Value> newMaskOperands(maskOperands.drop_back());
+ newMaskOperands.push_back(
+ getValueOrCreateConstantIndexOp(rewriter, loc, maskIndex));
+ return vector::CreateMaskOp::create(rewriter, loc, newMaskType,
+ newMaskOperands);
+ })
.Case([&](vector::ConstantMaskOp constantMaskOp)
-> std::optional<Operation *> {
// Take the shape of mask, compress its trailing dimension:
@@ -511,12 +512,14 @@ namespace {
// Emulate `vector.store` using a multi-byte container type.
//
-// When `assumeAligned` is true, store offsets are assumed to be aligned to
-// container element boundaries, so a store whose source vector fills whole
-// container elements (isDivisibleInSize) is emitted as a simple bitcast +
-// store without checking the offset. Stores that are not divisible in size
-// are rejected. This is useful for downstream users that have already
-// ensured alignment.
+// When `solver` is non-null, store offsets are queried against the
+// `IntegerDivisibilityLattice` to decide whether to take the aligned fast
+// path (bitcast + store) or fall through to the partial-store path. If the
+// lattice proves the innermost dynamic index is a multiple of
+// `containerBits / emulatedBits`, or if the lattice is opaque, the fast path
+// is taken (matches the legacy `assumeAligned=true` behavior). If the
+// lattice proves the index is *not* such a multiple, the fast path is
+// skipped and the partial-store path (which is sound for any offset) runs.
//
// The container type is obtained through Op adaptor and would normally be
// generated via `NarrowTypeEmulationConverter`.
@@ -559,9 +562,9 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
using Base::Base;
ConvertVectorStore(MLIRContext *context, bool disableAtomicRMW,
- bool assumeAligned)
+ DataFlowSolver *solver)
: OpConversionPattern<vector::StoreOp>(context),
- disableAtomicRMW(disableAtomicRMW), assumeAligned(assumeAligned) {}
+ disableAtomicRMW(disableAtomicRMW), solver(solver) {}
LogicalResult
matchAndRewrite(vector::StoreOp op, OpAdaptor adaptor,
@@ -605,10 +608,32 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
// Note, per-element-alignment was already verified above.
bool isDivisibleInSize = origElements % emulatedPerContainerElem == 0;
- // In assume-aligned mode, isDivisibleInSize alone is sufficient — the
- // caller guarantees that store offsets are aligned to container element
- // boundaries.
- if (assumeAligned) {
+ // Decide whether to take the aligned fast path (bitcast + store). The
+ // fast path is sound when the store offset is a multiple of
+ // `emulatedPerContainerElem`. Three-state behavior:
+ // - solver == nullptr: take the fast path (legacy `assumeAligned=true`
+ // behavior).
+ // - lattice opaque/uninitialized: take the fast path.
+ // - lattice proves `udiv % emulatedPerContainerElem == 0`: take it.
+ // - lattice proves the offset is *not* such a multiple: fall through
+ // to the partial-store path below, which is sound for any offset.
+ bool fastPathOffsetAligned = true;
+ if (solver && !op.getIndices().empty()) {
+ Value innerOffset = op.getIndices().back();
+ const dataflow::IntegerDivisibilityLattice *lattice =
+ solver->lookupState<dataflow::IntegerDivisibilityLattice>(
+ innerOffset);
+ if (lattice && !lattice->getValue().isUninitialized()) {
+ uint64_t udiv = lattice->getValue().getValue().udiv();
+ if (udiv % static_cast<uint64_t>(emulatedPerContainerElem) != 0)
+ fastPathOffsetAligned = false;
+ }
+ }
+
+ // In the aligned case, isDivisibleInSize alone is sufficient — the
+ // offset is guaranteed (by solver or by contract) to be aligned to
+ // container element boundaries.
+ if (fastPathOffsetAligned) {
if (!isDivisibleInSize)
return rewriter.notifyMatchFailure(
op, "the source vector does not fill whole container elements "
@@ -852,7 +877,7 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
private:
const bool disableAtomicRMW;
- const bool assumeAligned;
+ DataFlowSolver *const solver;
};
//===----------------------------------------------------------------------===//
@@ -2005,7 +2030,7 @@ struct RewriteBitCastOfTruncI : OpRewritePattern<vector::BitCastOp> {
auto shuffledElementType =
cast<IntegerType>(getElementTypeOrSelf(truncValue.getType()));
Value runningResult;
- for (const BitCastRewriter ::Metadata &metadata :
+ for (const BitCastRewriter::Metadata &metadata :
bcr.precomputeMetadata(shuffledElementType)) {
runningResult = bcr.genericRewriteStep(
rewriter, bitCastOp->getLoc(), truncValue, runningResult, metadata);
@@ -2292,7 +2317,8 @@ struct RewriteVectorTranspose : OpRewritePattern<vector::TransposeOp> {
// The emulated type is inferred from the converted memref type.
void vector::populateVectorNarrowTypeEmulationPatterns(
const arith::NarrowTypeEmulationConverter &typeConverter,
- RewritePatternSet &patterns, bool disableAtomicRMW, bool assumeAligned) {
+ RewritePatternSet &patterns, bool disableAtomicRMW,
+ DataFlowSolver *solver) {
// Populate `vector.*` conversion patterns.
// TODO: #119553 support atomicity
patterns.add<ConvertVectorLoad, ConvertVectorMaskedLoad,
@@ -2303,7 +2329,7 @@ void vector::populateVectorNarrowTypeEmulationPatterns(
// to avoid emitting atomic operations and reduce it to read-modify-write
// sequence for stores if it is known there are no thread contentions.
patterns.insert<ConvertVectorStore>(patterns.getContext(), disableAtomicRMW,
- assumeAligned);
+ solver);
}
void vector::populateVectorNarrowTypeRewritePatterns(
diff --git a/mlir/test/lib/Dialect/MemRef/CMakeLists.txt b/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
index 39457ab2d0bf7..06b07d2156e4e 100644
--- a/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_library(MLIRMemRefTestPasses
MLIRTestDialect
)
mlir_target_link_libraries(MLIRMemRefTestPasses PUBLIC
+ MLIRAnalysis
MLIRPass
MLIRMemRefDialect
MLIRMemRefTransforms
diff --git a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
index bec83a8dcbef9..f984f38a6511f 100644
--- a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
+++ b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
@@ -7,6 +7,10 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
@@ -96,13 +100,25 @@ struct TestEmulateNarrowTypePass
arith::ArithDialect, vector::VectorDialect, memref::MemRefDialect,
affine::AffineDialect>(opLegalCallback);
+ // Run the divisibility analysis on the original IR so the narrow-type
+ // emulation patterns can consult it when deciding whether a dynamic
+ // offset is provably aligned.
+ DataFlowSolver solver;
+ solver.load<dataflow::DeadCodeAnalysis>();
+ solver.load<dataflow::SparseConstantPropagation>();
+ solver.load<dataflow::IntegerDivisibilityAnalysis>();
+ if (failed(solver.initializeAndRun(op))) {
+ signalPassFailure();
+ return;
+ }
+
RewritePatternSet patterns(ctx);
arith::populateArithNarrowTypeEmulationPatterns(typeConverter, patterns);
memref::populateMemRefNarrowTypeEmulationPatterns(
- typeConverter, patterns, disableAtomicRMW, assumeAligned);
+ typeConverter, patterns, disableAtomicRMW, &solver);
vector::populateVectorNarrowTypeEmulationPatterns(
- typeConverter, patterns, disableAtomicRMW, assumeAligned);
+ typeConverter, patterns, disableAtomicRMW, &solver);
if (failed(applyPartialConversion(op, target, std::move(patterns))))
signalPassFailure();
@@ -127,12 +143,6 @@ struct TestEmulateNarrowTypePass
llvm::cl::desc("disable atomic read-modify-write and prefer generating "
"normal sequence"),
llvm::cl::init(false)};
-
- Option<bool> assumeAligned{
- *this, "assume-aligned",
- llvm::cl::desc("assume store offsets are aligned to container element "
- "boundaries"),
- llvm::cl::init(false)};
};
struct TestMemRefFlattenAndVectorNarrowTypeEmulationPass
>From 3feb621a55a6006226a9d722cf1fc1916685c05b Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 10:22:02 -0700
Subject: [PATCH 09/13] [MLIR] Match Vector divisibility gate to MemRef
(require udiv > 1 to reject)
The fast bitcast+store path in ConvertVectorStore was incorrectly bypassed
when the divisibility lattice resolved to getMinDivisibility() (udiv == 1).
Add the udiv > 1 guard so that opaque-but-initialized values fall through
to the fast path, matching the MemRef gates and the spec.
---
mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index f8d5e26fc6e7a..cfff3e40fcc71 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -625,7 +625,8 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
innerOffset);
if (lattice && !lattice->getValue().isUninitialized()) {
uint64_t udiv = lattice->getValue().getValue().udiv();
- if (udiv % static_cast<uint64_t>(emulatedPerContainerElem) != 0)
+ if (udiv > 1 &&
+ udiv % static_cast<uint64_t>(emulatedPerContainerElem) != 0)
fastPathOffsetAligned = false;
}
}
>From 7339fb5d87390e1514f0b1b80b986aeb05a94675 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 10:25:23 -0700
Subject: [PATCH 10/13] [MLIR] Drop shadowing inner elementsPerByte declaration
Caught by code review. The inner declaration at the empty-mixedSizes
branch shadowed the outer one introduced when the divisibility-driven
gate moved earlier in convertCastingOp; both held the same value.
---
mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp | 1 -
1 file changed, 1 deletion(-)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
index 9f59a27bb6106..fb9e341792d25 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp
@@ -109,7 +109,6 @@ convertCastingOp(ConversionPatternRewriter &rewriter,
OpFoldResult newOffset;
OpFoldResult intraOffset;
if (mixedSizes.empty()) {
- int64_t elementsPerByte = dstBits / srcBits;
AffineExpr s0;
bindSymbols(rewriter.getContext(), s0);
newOffset = affine::makeComposedFoldedAffineApply(
>From 50af57733fa60785b573614ad15b7ddfa4b193af Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 11:25:11 -0700
Subject: [PATCH 11/13] [MLIR] Add divisibility-analysis test pass and lit
tests
- Add TestIntegerDivisibilityAnalysisPass which annotates 'test.int_divisibility'
ops with the analyzed divisibility.
- Wire the pass into mlir-opt.
- New lit test mlir/test/Analysis/DataFlow/integer-divisibility.mlir exercising
constants, arith.muli, arith.addi, affine.apply (mul/floordiv/mod), and scf.for
induction variables.
- Update emulate-narrow-type.mlir (MemRef and Vector): drop assume-aligned RUN
flag and add cases that exercise the new analysis-driven gating
(provably-aligned dyn offset accepted, provably-misaligned dyn offset rejected).
- Drop the stale emulate-narrow-type-no-assume-aligned.mlir and
vector-emulate-narrow-type-aligned-store-only.mlir, which depended on the
removed assume-aligned flag.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../DataFlow/integer-divisibility.mlir | 152 ++++++++++++++++++
...emulate-narrow-type-no-assume-aligned.mlir | 23 ---
.../Dialect/MemRef/emulate-narrow-type.mlir | 59 ++++++-
...mulate-narrow-type-aligned-store-only.mlir | 55 -------
.../Vector/vector-emulate-narrow-type.mlir | 34 ++++
mlir/test/lib/Analysis/CMakeLists.txt | 2 +
.../TestIntegerDivisibilityAnalysis.cpp | 96 +++++++++++
mlir/test/lib/Dialect/MemRef/CMakeLists.txt | 2 +
.../Dialect/MemRef/TestEmulateNarrowType.cpp | 4 +
mlir/tools/mlir-opt/mlir-opt.cpp | 2 +
10 files changed, 346 insertions(+), 83 deletions(-)
create mode 100644 mlir/test/Analysis/DataFlow/integer-divisibility.mlir
delete mode 100644 mlir/test/Dialect/MemRef/emulate-narrow-type-no-assume-aligned.mlir
delete mode 100644 mlir/test/Dialect/Vector/vector-emulate-narrow-type-aligned-store-only.mlir
create mode 100644 mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
diff --git a/mlir/test/Analysis/DataFlow/integer-divisibility.mlir b/mlir/test/Analysis/DataFlow/integer-divisibility.mlir
new file mode 100644
index 0000000000000..7f9466e949d4c
--- /dev/null
+++ b/mlir/test/Analysis/DataFlow/integer-divisibility.mlir
@@ -0,0 +1,152 @@
+// RUN: mlir-opt --split-input-file --test-int-divisibility-analysis --allow-unregistered-dialect %s | FileCheck %s
+
+// CHECK-LABEL: @constant
+func.func @constant() -> index {
+ %0 = arith.constant 8 : index
+ // CHECK: divisibility = "udiv = 8, sdiv = 8"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @muli_constant
+func.func @muli_constant(%arg0 : index) -> index {
+ %c4 = arith.constant 4 : index
+ %0 = arith.muli %arg0, %c4 : index
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @addi_gcd_of_muli_operands
+func.func @addi_gcd_of_muli_operands(%arg0 : index, %arg1 : index) -> index {
+ %c8 = arith.constant 8 : index
+ %c12 = arith.constant 12 : index
+ %a = arith.muli %arg0, %c8 : index
+ %b = arith.muli %arg1, %c12 : index
+ %0 = arith.addi %a, %b : index
+ // gcd(8, 12) = 4.
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @addi_same_divisibility
+func.func @addi_same_divisibility(%arg0 : index, %arg1 : index) -> index {
+ %c16 = arith.constant 16 : index
+ %a = arith.muli %arg0, %c16 : index
+ %b = arith.muli %arg1, %c16 : index
+ %0 = arith.addi %a, %b : index
+ // CHECK: divisibility = "udiv = 16, sdiv = 16"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @affine_apply_mul
+func.func @affine_apply_mul(%arg0 : index) -> index {
+ %c2 = arith.constant 2 : index
+ %seed = arith.muli %arg0, %c2 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 * 16)>(%seed)
+ // 2 * 16 = 32.
+ // CHECK: divisibility = "udiv = 32, sdiv = 32"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @affine_apply_mul_then_floordiv
+func.func @affine_apply_mul_then_floordiv(%arg0 : index) -> index {
+ %0 = affine.apply affine_map<(d0) -> (d0 * 16)>(%arg0)
+ %1 = affine.apply affine_map<(d0) -> (d0 floordiv 4)>(%0)
+ // 16 floordiv 4 = 4.
+ // CHECK: divisibility = "udiv = 4, sdiv = 4"
+ %2 = "test.int_divisibility"(%1) : (index) -> index
+ return %2 : index
+}
+
+// -----
+
+// CHECK-LABEL: @affine_apply_mod_zero
+func.func @affine_apply_mod_zero(%arg0 : index) -> index {
+ %0 = affine.apply affine_map<(d0) -> (d0 * 16)>(%arg0)
+ %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%0)
+ // 16 % 16 == 0, so x mod 16 is always 0 -> divisibility 0 (lattice top).
+ // CHECK: divisibility = "udiv = 0, sdiv = 0"
+ %2 = "test.int_divisibility"(%1) : (index) -> index
+ return %2 : index
+}
+
+// -----
+
+// CHECK-LABEL: @affine_apply_constant
+func.func @affine_apply_constant() -> index {
+ %0 = affine.apply affine_map<() -> (64)>()
+ // CHECK: divisibility = "udiv = 64, sdiv = 64"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @scf_for_constant_step
+func.func @scf_for_constant_step() {
+ %c0 = arith.constant 0 : index
+ %c64 = arith.constant 64 : index
+ %c8 = arith.constant 8 : index
+ scf.for %iv = %c0 to %c64 step %c8 {
+ // CHECK: divisibility = "udiv = 8, sdiv = 8"
+ %0 = "test.int_divisibility"(%iv) : (index) -> index
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @scf_for_nontrivial_gcd
+func.func @scf_for_nontrivial_gcd() {
+ %c12 = arith.constant 12 : index
+ %c100 = arith.constant 100 : index
+ %c18 = arith.constant 18 : index
+ scf.for %iv = %c12 to %c100 step %c18 {
+ // gcd(12, 18) = 6.
+ // CHECK: divisibility = "udiv = 6, sdiv = 6"
+ %0 = "test.int_divisibility"(%iv) : (index) -> index
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @scf_for_coprime
+func.func @scf_for_coprime() {
+ %c15 = arith.constant 15 : index
+ %c100 = arith.constant 100 : index
+ %c8 = arith.constant 8 : index
+ scf.for %iv = %c15 to %c100 step %c8 {
+ // gcd(15, 8) = 1.
+ // CHECK: divisibility = "udiv = 1, sdiv = 1"
+ %0 = "test.int_divisibility"(%iv) : (index) -> index
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL: @affine_apply_mul_plus_const
+func.func @affine_apply_mul_plus_const(%arg0 : index) -> index {
+ %c4 = arith.constant 4 : index
+ %seed = arith.muli %arg0, %c4 : index
+ %0 = affine.apply affine_map<(d0) -> (d0 * 8 + 16)>(%seed)
+ // seed has udiv = 4, multiplied by 8 -> 32, then +16. gcd(32,16) = 16.
+ // CHECK: divisibility = "udiv = 16, sdiv = 16"
+ %1 = "test.int_divisibility"(%0) : (index) -> index
+ return %1 : index
+}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type-no-assume-aligned.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type-no-assume-aligned.mlir
deleted file mode 100644
index 3625f91cedbd6..0000000000000
--- a/mlir/test/Dialect/MemRef/emulate-narrow-type-no-assume-aligned.mlir
+++ /dev/null
@@ -1,23 +0,0 @@
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8" --cse --verify-diagnostics --split-input-file %s
-
-// Without `assume-aligned=true`, dynamic offsets in `memref.subview` and
-// `memref.reinterpret_cast` cannot be proven to be multiples of
-// `dstBits / srcBits`. The patterns must reject them so partial conversion
-// fails to legalize the op.
-
-func.func @negative_subview_dynamic_inner_offset_i4(%off: index) -> i4 {
- %c0 = arith.constant 0 : index
- %arr = memref.alloc() : memref<128xi4>
- // expected-error @+1 {{failed to legalize operation 'memref.subview' that was explicitly marked illegal}}
- %subview = memref.subview %arr[%off] [32] [1] : memref<128xi4> to memref<32xi4, strided<[1], offset: ?>>
- %ld = memref.load %subview[%c0] : memref<32xi4, strided<[1], offset: ?>>
- return %ld : i4
-}
-
-// -----
-
-func.func @negative_reinterpret_cast_memref_rank3_dynamic_offset_i4(%arg0: memref<2x4x8xi4>, %off: index) -> memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>> {
- // expected-error @+1 {{failed to legalize operation 'memref.reinterpret_cast' that was explicitly marked illegal}}
- %r = memref.reinterpret_cast %arg0 to offset: [%off], sizes: [4, 4, 8], strides: [32, 8, 1] : memref<2x4x8xi4> to memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
- return %r : memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
-}
diff --git a/mlir/test/Dialect/MemRef/emulate-narrow-type.mlir b/mlir/test/Dialect/MemRef/emulate-narrow-type.mlir
index adc8fe3b36096..efb1e3c1005d5 100644
--- a/mlir/test/Dialect/MemRef/emulate-narrow-type.mlir
+++ b/mlir/test/Dialect/MemRef/emulate-narrow-type.mlir
@@ -1,5 +1,5 @@
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 assume-aligned=true" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=32 assume-aligned=true" --cse --verify-diagnostics --split-input-file %s | FileCheck %s --check-prefix=CHECK32
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=32" --cse --verify-diagnostics --split-input-file %s | FileCheck %s --check-prefix=CHECK32
// Expect no conversions.
func.func @memref_i8() -> i8 {
@@ -267,7 +267,7 @@ func.func @memref_subview_dynamic_inner_offset_i4(%off: index) -> i4 {
func.func @memref_subview_aligned_dynamic_inner_offset_i4(%x: index) -> i4 {
%c0 = arith.constant 0 : index
- %off = affine.apply affine_map<()[s0] -> (s0 * 2)>()[%x]
+ %off = affine.apply affine_map<()[s0] -> (s0 * 8)>()[%x]
%arr = memref.alloc() : memref<128xi4>
%subview = memref.subview %arr[%off] [32] [1] : memref<128xi4> to memref<32xi4, strided<[1], offset: ?>>
%ld = memref.load %subview[%c0] : memref<32xi4, strided<[1], offset: ?>>
@@ -277,10 +277,17 @@ func.func @memref_subview_aligned_dynamic_inner_offset_i4(%x: index) -> i4 {
// CHECK-LABEL: func.func @memref_subview_aligned_dynamic_inner_offset_i4(
// CHECK-SAME: %[[X:[a-zA-Z0-9_]+]]: index
// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<64xi8>
-// CHECK-NOT: affine.apply
-// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ALLOC]][%[[X]]] [16] [1] : memref<64xi8> to memref<16xi8, strided<[1], offset: ?>>
+// CHECK: %[[OFF:.+]] = affine.apply
+// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[ALLOC]][%[[OFF]]] [16] [1] : memref<64xi8> to memref<16xi8, strided<[1], offset: ?>>
// CHECK: memref.load %[[SUBVIEW]]
+// CHECK32-LABEL: func.func @memref_subview_aligned_dynamic_inner_offset_i4(
+// CHECK32-SAME: %[[X:[a-zA-Z0-9_]+]]: index
+// CHECK32: %[[ALLOC:.+]] = memref.alloc() : memref<16xi32>
+// CHECK32-NOT: affine.apply
+// CHECK32: %[[SUBVIEW:.+]] = memref.subview %[[ALLOC]][%[[X]]] [4] [1] : memref<16xi32> to memref<4xi32, strided<[1], offset: ?>>
+// CHECK32: memref.load %[[SUBVIEW]]
+
// -----
func.func @negative_memref_subview_non_contiguous(%idx : index) -> i4 {
@@ -712,3 +719,45 @@ func.func @alloc_non_contiguous() {
func.func @argument_non_contiguous(%arg0 : memref<8x8xi4, strided<[1, 8]>>) {
return
}
+
+// -----
+
+// Divisibility-aware acceptance: the dynamic offset is `%i * 8`, which the
+// `IntegerDivisibilityAnalysis` proves is a multiple of both the i4 -> i8
+// ratio (2) and the i4 -> i32 ratio (8). Lowering succeeds and emits the
+// linearized reinterpret_cast.
+
+func.func @reinterpret_cast_dynamic_offset_divisible_i4(%arg0: memref<2x4x8xi4>, %i: index) -> memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>> {
+ %c8 = arith.constant 8 : index
+ %off = arith.muli %i, %c8 : index
+ %r = memref.reinterpret_cast %arg0 to offset: [%off], sizes: [4, 4, 8], strides: [32, 8, 1] : memref<2x4x8xi4> to memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
+ return %r : memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
+}
+
+// CHECK-LABEL: func @reinterpret_cast_dynamic_offset_divisible_i4(
+// CHECK-SAME: %[[ARG0:.+]]: memref<32xi8>,
+// CHECK-SAME: %[[I:.+]]: index
+// CHECK: %[[NEWOFF:.+]] = affine.apply
+// CHECK: %[[R:.+]] = memref.reinterpret_cast %[[ARG0]] to offset: {{\[}}%[[NEWOFF]]{{\]}}, sizes: [64], strides: [1] : memref<32xi8> to memref<64xi8, strided<[1], offset: ?>>
+// CHECK: return %[[R]]
+
+// CHECK32-LABEL: func @reinterpret_cast_dynamic_offset_divisible_i4(
+// CHECK32-SAME: %[[ARG0:.+]]: memref<8xi32>,
+// CHECK32-SAME: %[[I:.+]]: index
+// CHECK32: %[[R:.+]] = memref.reinterpret_cast %[[ARG0]] to offset: {{.*}}, sizes: [16], strides: [1] : memref<8xi32> to memref<16xi32, strided<[1], offset: ?>>
+// CHECK32: return %[[R]]
+
+// -----
+
+// Divisibility-aware rejection: the dynamic offset is `%i * 3`. With i4 -> i8
+// emulation, the divisor 3 is not a multiple of `elementsPerByte == 2`, so
+// the analysis proves the offset is not multi-element aligned. The pattern
+// must reject and partial conversion must fail.
+
+func.func @negative_reinterpret_cast_dynamic_offset_not_divisible_i4(%arg0: memref<2x4x8xi4>, %i: index) -> memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>> {
+ %c3 = arith.constant 3 : index
+ %off = arith.muli %i, %c3 : index
+ // expected-error @+1 {{failed to legalize operation 'memref.reinterpret_cast' that was explicitly marked illegal}}
+ %r = memref.reinterpret_cast %arg0 to offset: [%off], sizes: [4, 4, 8], strides: [32, 8, 1] : memref<2x4x8xi4> to memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
+ return %r : memref<4x4x8xi4, strided<[32, 8, 1], offset: ?>>
+}
diff --git a/mlir/test/Dialect/Vector/vector-emulate-narrow-type-aligned-store-only.mlir b/mlir/test/Dialect/Vector/vector-emulate-narrow-type-aligned-store-only.mlir
deleted file mode 100644
index 2ef568fa6a741..0000000000000
--- a/mlir/test/Dialect/Vector/vector-emulate-narrow-type-aligned-store-only.mlir
+++ /dev/null
@@ -1,55 +0,0 @@
-// RUN: mlir-opt --test-emulate-narrow-int="memref-load-bitwidth=8 assume-aligned=true" --cse --verify-diagnostics --split-input-file %s | FileCheck %s
-
-/// Aligned store, constant index - the source vector fills whole container
-/// elements. Produces a simple bitcast + store.
-func.func @vector_store_i4_aligned_const(%arg0: vector<8xi4>, %arg1: index, %arg2: index) {
- %0 = memref.alloc() : memref<4x8xi4>
- vector.store %arg0, %0[%arg1, %arg2] : memref<4x8xi4>, vector<8xi4>
- return
-}
-// CHECK-DAG: #[[$MAP:.+]] = affine_map<()[s0, s1] -> (s0 * 4 + s1 floordiv 2)>
-// CHECK: func @vector_store_i4_aligned_const
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: vector<8xi4>
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index
-// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<16xi8>
-// CHECK: %[[INDEX:.+]] = affine.apply #[[$MAP]]()[%[[ARG1]], %[[ARG2]]]
-// CHECK: %[[VEC_I8:.+]] = vector.bitcast %[[ARG0]] : vector<8xi4> to vector<4xi8>
-// CHECK: vector.store %[[VEC_I8]], %[[ALLOC]][%[[INDEX]]] : memref<16xi8>, vector<4xi8>
-
-// -----
-
-/// Aligned store, dynamic index. The source vector (8 x i4 = 32 bits) is a
-/// whole multiple of the container element size (i8 = 8 bits), so no partial
-/// stores are needed. This holds regardless of the dynamic offset.
-func.func @vector_store_i4_aligned_dynamic(%arg0: vector<8xi4>, %arg1: index, %arg2: index, %arg3: index, %arg4: index) {
- %0 = memref.alloc(%arg1, %arg2) : memref<?x?xi4>
- vector.store %arg0, %0[%arg3, %arg4] : memref<?x?xi4>, vector<8xi4>
- return
-}
-// CHECK-DAG: #[[$MAP:.+]] = affine_map<()[s0, s1] -> ((s0 * s1) floordiv 2, s0 floordiv 2)>
-// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0, s1, s2] -> ((s2 + s0 * s1) floordiv 2)>
-// CHECK: func @vector_store_i4_aligned_dynamic
-// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: vector<8xi4>
-// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index
-// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index
-// CHECK: %[[SIZE:.+]] = affine.max #[[$MAP]]()[%[[ARG2]], %[[ARG1]]]
-// CHECK: %[[ALLOC:.+]] = memref.alloc(%[[SIZE]]) : memref<?xi8>
-// CHECK: %[[INDEX:.+]] = affine.apply #[[$MAP1]]()[%[[ARG3]], %[[ARG2]], %[[ARG4]]]
-// CHECK: %[[VEC_I8:.+]] = vector.bitcast %[[ARG0]] : vector<8xi4> to vector<4xi8>
-// CHECK: vector.store %[[VEC_I8]], %[[ALLOC]][%[[INDEX]]] : memref<?xi8>, vector<4xi8>
-
-// -----
-
-/// The source vector does not fill whole container elements (3 x i4 != N x i8),
-/// so the aligned pattern rejects it. With aligned-store-only, no unaligned
-/// pattern is available, so legalization fails.
-func.func @vector_store_i4_not_divisible(%arg0: vector<3xi4>) {
- %0 = memref.alloc() : memref<12xi4>
- %c0 = arith.constant 0 : index
- // expected-error @below {{failed to legalize operation 'vector.store' that was explicitly marked illegal}}
- vector.store %arg0, %0[%c0] : memref<12xi4>, vector<3xi4>
- return
-}
diff --git a/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir b/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
index 98b1f07ef5fb0..7b6b2ae76bb8d 100644
--- a/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
+++ b/mlir/test/Dialect/Vector/vector-emulate-narrow-type.mlir
@@ -758,3 +758,37 @@ func.func @vector_maskedstore_i4_arith_constant(%val_to_store: vector<8xi4>) {
// CHECK: %[[SELECT:.+]] = arith.select %[[MASK]], %[[VAL_TO_STORE]], %[[LOAD_UPCAST]]
// CHECK: %[[SELECT_DOWNCAST:.+]] = vector.bitcast %[[SELECT]]
// CHECK: vector.maskedstore %[[ALLOC]][%[[IDX_FLATTENED]]], %[[COMPRESSED_MASK]], %[[SELECT_DOWNCAST]]
+
+// -----
+
+///----------------------------------------------------------------------------------------
+/// vector.store - divisibility-analysis-driven alignment
+///----------------------------------------------------------------------------------------
+
+/// The innermost dynamic index is `arith.muli %arg, %c8`, which the
+/// `IntegerDivisibilityAnalysis` proves is a multiple of both the i4 -> i8
+/// ratio (2) and the i4 -> i32 ratio (8). The aligned fast path is taken:
+/// `vector.bitcast` + `vector.store`.
+
+func.func @vector_store_i4_dyn_index_divisible(%arg0: vector<8xi4>, %i: index) {
+ %0 = memref.alloc() : memref<32xi4>
+ %c8 = arith.constant 8 : index
+ %idx = arith.muli %i, %c8 : index
+ vector.store %arg0, %0[%idx] : memref<32xi4>, vector<8xi4>
+ return
+}
+
+// CHECK-LABEL: func @vector_store_i4_dyn_index_divisible
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: vector<8xi4>
+// CHECK-SAME: %[[I:[a-zA-Z0-9]+]]: index
+// CHECK: %[[ALLOC:.+]] = memref.alloc() : memref<16xi8>
+// CHECK: %[[INDEX:.+]] = affine.apply
+// CHECK: %[[VEC_I8:.+]] = vector.bitcast %[[ARG0]] : vector<8xi4> to vector<4xi8>
+// CHECK: vector.store %[[VEC_I8]], %[[ALLOC]][%[[INDEX]]] : memref<16xi8>, vector<4xi8>
+
+// CHECK32-LABEL: func @vector_store_i4_dyn_index_divisible
+// CHECK32-SAME: %[[ARG0:[a-zA-Z0-9]+]]: vector<8xi4>
+// CHECK32-SAME: %[[I:[a-zA-Z0-9]+]]: index
+// CHECK32: %[[ALLOC:.+]] = memref.alloc() : memref<4xi32>
+// CHECK32: %[[VEC_I32:.+]] = vector.bitcast %[[ARG0]] : vector<8xi4> to vector<1xi32>
+// CHECK32: vector.store %[[VEC_I32]], %[[ALLOC]]
diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt
index c37671ade37b3..d86af5017f24b 100644
--- a/mlir/test/lib/Analysis/CMakeLists.txt
+++ b/mlir/test/lib/Analysis/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_library(MLIRTestAnalysis
DataFlow/TestDeadCodeAnalysis.cpp
DataFlow/TestDenseBackwardDataFlowAnalysis.cpp
DataFlow/TestDenseForwardDataFlowAnalysis.cpp
+ DataFlow/TestIntegerDivisibilityAnalysis.cpp
DataFlow/TestLivenessAnalysis.cpp
DataFlow/TestSparseBackwardDataFlowAnalysis.cpp
DataFlow/TestStridedMetadataRangeAnalysis.cpp
@@ -27,6 +28,7 @@ add_mlir_library(MLIRTestAnalysis
mlir_target_link_libraries(MLIRTestAnalysis PUBLIC
MLIRAffineDialect
MLIRAnalysis
+ MLIRArithDialect
MLIRFunctionInterfaces
MLIRMemRefDialect
MLIRPass
diff --git a/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
new file mode 100644
index 0000000000000..9290a13605e55
--- /dev/null
+++ b/mlir/test/lib/Analysis/DataFlow/TestIntegerDivisibilityAnalysis.cpp
@@ -0,0 +1,96 @@
+//===- TestIntegerDivisibilityAnalysis.cpp - Test int divisibility --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
+#include "mlir/Analysis/DataFlowFramework.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+namespace {
+struct TestIntegerDivisibilityAnalysisPass
+ : public PassWrapper<TestIntegerDivisibilityAnalysisPass, OperationPass<>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(
+ TestIntegerDivisibilityAnalysisPass)
+
+ StringRef getArgument() const override {
+ return "test-int-divisibility-analysis";
+ }
+ StringRef getDescription() const override {
+ return "Test integer divisibility analysis by annotating "
+ "'test.int_divisibility' ops with the divisibility of their operand.";
+ }
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect, affine::AffineDialect>();
+ arith::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
+ affine::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
+ }
+
+ void runOnOperation() override {
+ Operation *rootOp = getOperation();
+ MLIRContext *context = &getContext();
+
+ // The pass is rooted on `test.int_divisibility` ops, which are expected
+ // to have a single operand for which to annotate divisibility information.
+ SmallVector<std::pair<Operation *, Value>> queryOps;
+ rootOp->walk([&](Operation *op) {
+ if (op->getName().getStringRef() == "test.int_divisibility" &&
+ op->getNumOperands() == 1)
+ queryOps.emplace_back(op, op->getOperand(0));
+ });
+
+ DataFlowSolver solver;
+ // DeadCodeAnalysis is the base analysis that allows the solver to traverse
+ // control flow. It is required by IntegerDivisibilityAnalysis.
+ solver.load<DeadCodeAnalysis>();
+ // SparseConstantPropagation allows the solver to call
+ // visitNonControlFlowArguments and analyze arguments like loop induction
+ // variables.
+ solver.load<SparseConstantPropagation>();
+ solver.load<IntegerDivisibilityAnalysis>();
+ if (failed(solver.initializeAndRun(rootOp)))
+ return signalPassFailure();
+
+ for (auto &[op, value] : queryOps) {
+ const auto *lattice =
+ solver.lookupState<IntegerDivisibilityLattice>(value);
+ if (!lattice || lattice->getValue().isUninitialized()) {
+ op->setAttr("divisibility", StringAttr::get(context, "uninitialized"));
+ continue;
+ }
+
+ // Format for the divisibility information is "udiv = X, sdiv = Y".
+ const auto &div = lattice->getValue().getValue();
+ std::string result;
+ llvm::raw_string_ostream os(result);
+ os << "udiv = " << div.udiv() << ", sdiv = " << div.sdiv();
+ op->setAttr("divisibility", StringAttr::get(context, result));
+ }
+ }
+};
+} // end anonymous namespace
+
+namespace mlir::test {
+void registerTestIntegerDivisibilityAnalysisPass() {
+ PassRegistration<TestIntegerDivisibilityAnalysisPass>();
+}
+} // end namespace mlir::test
diff --git a/mlir/test/lib/Dialect/MemRef/CMakeLists.txt b/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
index 06b07d2156e4e..05638772dc3f2 100644
--- a/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/MemRef/CMakeLists.txt
@@ -10,7 +10,9 @@ add_mlir_library(MLIRMemRefTestPasses
MLIRTestDialect
)
mlir_target_link_libraries(MLIRMemRefTestPasses PUBLIC
+ MLIRAffineDialect
MLIRAnalysis
+ MLIRArithDialect
MLIRPass
MLIRMemRefDialect
MLIRMemRefTransforms
diff --git a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
index f984f38a6511f..171ab19a1805b 100644
--- a/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
+++ b/mlir/test/lib/Dialect/MemRef/TestEmulateNarrowType.cpp
@@ -12,7 +12,9 @@
#include "mlir/Analysis/DataFlow/IntegerDivisibilityAnalysis.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -41,6 +43,8 @@ struct TestEmulateNarrowTypePass
registry
.insert<arith::ArithDialect, func::FuncDialect, memref::MemRefDialect,
vector::VectorDialect, affine::AffineDialect>();
+ arith::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
+ affine::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
}
StringRef getArgument() const final { return "test-emulate-narrow-int"; }
StringRef getDescription() const final {
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index c4754b3a08551..13c0934f34656 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -104,6 +104,7 @@ void registerTestComposeSubView();
void registerTestMultiBuffering();
void registerTestIRVisitorsPass();
void registerTestGenericIRVisitorsPass();
+void registerTestIntegerDivisibilityAnalysisPass();
void registerTestInterfaces();
void registerTestIRVisitorsPass();
void registerTestLastModifiedPass();
@@ -253,6 +254,7 @@ static void registerTestPasses() {
mlir::test::registerTestMultiBuffering();
mlir::test::registerTestIRVisitorsPass();
mlir::test::registerTestGenericIRVisitorsPass();
+ mlir::test::registerTestIntegerDivisibilityAnalysisPass();
mlir::test::registerTestInterfaces();
mlir::test::registerTestIrdlTestDialectConversionPass();
mlir::test::registerTestIRVisitorsPass();
>From 76fd7544d06483167e2cbf2e93e435965467dd5c Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 11:35:01 -0700
Subject: [PATCH 12/13] [MLIR] Fall through to partial-store path on size
mismatch in ConvertVectorStore
The Task 4 rewrite eagerly rejected with notifyMatchFailure inside the fast
bitcast+store path when isDivisibleInSize was false. That mirrored old
assumeAligned=true semantics but broke the long-standing in-tree
vector-emulate-narrow-type-unaligned*.mlir tests, which exercised the
partial-store path with a default (assumeAligned=false) solver-less run.
Make the fast path require both offset alignment AND size divisibility;
otherwise fall through to the partial-store path, which is sound for any
combination of offset and size.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
.../Vector/Transforms/VectorEmulateNarrowType.cpp | 14 +++++---------
1 file changed, 5 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
index cfff3e40fcc71..96486da59f9a3 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorEmulateNarrowType.cpp
@@ -631,15 +631,11 @@ struct ConvertVectorStore final : OpConversionPattern<vector::StoreOp> {
}
}
- // In the aligned case, isDivisibleInSize alone is sufficient — the
- // offset is guaranteed (by solver or by contract) to be aligned to
- // container element boundaries.
- if (fastPathOffsetAligned) {
- if (!isDivisibleInSize)
- return rewriter.notifyMatchFailure(
- op, "the source vector does not fill whole container elements "
- "(not divisible in size)");
-
+ // Take the aligned fast path only when both the offset is aligned to
+ // container element boundaries AND the source vector exactly fills
+ // whole container elements. Otherwise fall through to the partial-store
+ // path below, which is sound for any offset and any size.
+ if (fastPathOffsetAligned && isDivisibleInSize) {
auto stridedMetadata =
memref::ExtractStridedMetadataOp::create(rewriter, loc, op.getBase());
OpFoldResult linearizedIndices;
>From 00a5d235e2bea87f2883535e01ca523063b23c18 Mon Sep 17 00:00:00 2001
From: Alan Li <me at alanli.org>
Date: Wed, 13 May 2026 11:46:42 -0700
Subject: [PATCH 13/13] [MLIR] Register divisibility analysis external models
in RegisterAllDialects
Wires up arith::registerInferIntDivisibilityOpInterfaceExternalModels and
affine::registerInferIntDivisibilityOpInterfaceExternalModels in the
global dialect registry, so any caller that uses registerAllDialects
picks up divisibility inference on arith + affine ops automatically.
This is the last commit of the port -- IntegerDivisibilityAnalysis is now
self-contained and consumable by downstream passes (e.g., the MemRef and
Vector narrow-type emulation patterns) by passing a configured
DataFlowSolver* into the populate function.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply at anthropic.com>
---
mlir/lib/RegisterAllDialects.cpp | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/mlir/lib/RegisterAllDialects.cpp b/mlir/lib/RegisterAllDialects.cpp
index 01a7401db4710..d9819a63f289b 100644
--- a/mlir/lib/RegisterAllDialects.cpp
+++ b/mlir/lib/RegisterAllDialects.cpp
@@ -15,8 +15,10 @@
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/IR/InferIntDivisibilityOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/InferIntDivisibilityOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
#include "mlir/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.h"
@@ -159,10 +161,12 @@ void mlir::registerAllDialects(DialectRegistry ®istry) {
// clang-format on
// Register all external models.
+ affine::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
affine::registerValueBoundsOpInterfaceExternalModels(registry);
arith::registerBufferDeallocationOpInterfaceExternalModels(registry);
arith::registerBufferizableOpInterfaceExternalModels(registry);
arith::registerBufferViewFlowOpInterfaceExternalModels(registry);
+ arith::registerInferIntDivisibilityOpInterfaceExternalModels(registry);
arith::registerShardingInterfaceExternalModels(registry);
arith::registerValueBoundsOpInterfaceExternalModels(registry);
bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
More information about the Mlir-commits
mailing list