[Mlir-commits] [mlir] [mlir][Interfaces] Allow integer types for `ValueBoundsOpInterface` (PR #196082)
Matthias Springer
llvmlistbot at llvm.org
Wed May 6 07:50:50 PDT 2026
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/196082
>From b4920caef81657e262557ba7f1eb23ac28dee525 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Wed, 6 May 2026 13:33:27 +0000
Subject: [PATCH] [mlir][Interfaces] Allow integer types for
`ValueBoundsOpInterface`
---
.../Dialect/Affine/Transforms/Transforms.h | 12 +--
.../Dialect/Arith/Transforms/Transforms.h | 12 +--
.../IR/ScalableValueBoundsConstraintSet.h | 7 +-
.../mlir/Interfaces/ValueBoundsOpInterface.h | 49 +++++++++----
.../Affine/Transforms/ReifyValueBounds.cpp | 23 ++++--
.../Arith/Transforms/ReifyValueBounds.cpp | 46 ++++++++----
.../Linalg/Transforms/HoistPadding.cpp | 2 +-
.../Dialect/Linalg/Transforms/Hoisting.cpp | 3 +-
.../lib/Dialect/Linalg/Transforms/Padding.cpp | 2 +-
.../Dialect/Linalg/Transforms/Promotion.cpp | 2 +-
.../Linalg/Transforms/TilingInterfaceImpl.cpp | 5 +-
.../Transforms/IndependenceTransforms.cpp | 2 +-
.../Transforms/IndependenceTransforms.cpp | 3 +-
.../IR/ScalableValueBoundsConstraintSet.cpp | 8 +-
.../lib/Interfaces/ValueBoundsOpInterface.cpp | 73 +++++++++++--------
.../Affine/invalid-reify-bound-dim.mlir | 2 +-
.../Arith/value-bounds-op-interface-impl.mlir | 39 ++++++++++
.../Dialect/Affine/TestReifyValueBounds.cpp | 21 ++++--
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 10 ++-
mlir/test/lib/Dialect/Test/TestOps.td | 1 +
20 files changed, 217 insertions(+), 105 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
index 84adb8e6a1e6d..f8f6d4f5928b0 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h
@@ -80,19 +80,19 @@ FailureOr<AffineApplyOp> decompose(RewriterBase &rewriter, AffineApplyOp op);
/// `stopCondition` is met.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
FailureOr<OpFoldResult>
reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
const ValueBoundsConstraintSet::Variable &var,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Reify a bound for the given index-typed value in terms of SSA values for
/// which `stopCondition` is met. If no stop condition is specified, reify in
/// terms of the operands of the owner op.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
///
/// Example:
/// %0 = arith.addi %a, %b : index
@@ -107,19 +107,19 @@ reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
FailureOr<OpFoldResult> reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Reify a bound for the specified dimension of the given shaped value in terms
/// of SSA values for which `stopCondition` is met. If no stop condition is
/// specified, reify in terms of the operands of the owner op.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
FailureOr<OpFoldResult> reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim,
ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Materialize an already computed bound with Affine dialect ops.
///
diff --git a/mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h
index ffd367ef11abc..c079bb9aea6e1 100644
--- a/mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h
@@ -27,19 +27,19 @@ namespace arith {
/// `stopCondition` is met.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
FailureOr<OpFoldResult>
reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
const ValueBoundsConstraintSet::Variable &var,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Reify a bound for the given index-typed value in terms of SSA values for
/// which `stopCondition` is met. If no stop condition is specified, reify in
/// terms of the operands of the owner op.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
///
/// Example:
/// %0 = arith.addi %a, %b : index
@@ -54,19 +54,19 @@ reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
FailureOr<OpFoldResult> reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
const ValueBoundsConstraintSet::StopConditionFn &stopCondition = nullptr,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Reify a bound for the specified dimension of the given shaped value in terms
/// of SSA values for which `stopCondition` is met. If no stop condition is
/// specified, reify in terms of the operands of the owner op.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
+/// `options.closedUB` is set to "true", upper bounds are also closed.
FailureOr<OpFoldResult> reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim,
const ValueBoundsConstraintSet::StopConditionFn &stopCondition = nullptr,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
} // namespace arith
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
index 216b94fc445c7..b32f4a777b90a 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.h
@@ -32,8 +32,8 @@ struct ScalableValueBoundsConstraintSet
ScalableValueBoundsConstraintSet(
MLIRContext *context,
ValueBoundsConstraintSet::StopConditionFn stopCondition,
- unsigned vscaleMin, unsigned vscaleMax)
- : RTTIExtends(context, stopCondition,
+ unsigned vscaleMin, unsigned vscaleMax, ValueBoundsOptions options = {})
+ : RTTIExtends(context, stopCondition, options,
/*addConservativeSemiAffineBounds=*/true),
vscaleMin(vscaleMin), vscaleMax(vscaleMax) {};
@@ -71,7 +71,8 @@ struct ScalableValueBoundsConstraintSet
static FailureOr<ConstantOrScalableBound>
computeScalableBound(Value value, std::optional<int64_t> dim,
unsigned vscaleMin, unsigned vscaleMax,
- presburger::BoundType boundType, bool closedUB = true,
+ presburger::BoundType boundType,
+ ValueBoundsOptions options = {/*closedUB=*/true},
const StopConditionFn &stopCondition = nullptr);
/// Get the value of vscale. Returns `nullptr` vscale as not been encountered.
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index 0590cadca8f9e..434f786c161d8 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -52,9 +52,22 @@ class HyperrectangularSlice {
// Profiled: 488K calls, avg=1.5+-0.5. N=2 covers >90% of cases inline.
using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>, 2>;
+/// Options that control value bound computation.
+struct ValueBoundsOptions {
+ /// By default, lower/equal bounds are closed and upper bounds are open. If
+ /// `closedUB` is set to "true", upper bounds are also closed.
+ bool closedUB = false;
+
+ /// If set to "true", integer-typed SSA values are treated like index-typed
+ /// SSA values. The value bounds infrastructure assumes that such integer
+ /// computations do not overflow. If set to "false", integer-typed SSA values
+ /// are rejected.
+ bool allowIntegerType = false;
+};
+
/// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
-/// constraint system and mapping of constrained variables to index-typed
-/// values or dimension sizes of shaped values.
+/// constraint system and mapping of constrained variables to index-typed values
+/// or dimension sizes of shaped values.
///
/// Interface implementations of `ValueBoundsOpInterface` use `addBounds` to
/// insert constraints about their results and/or region block arguments into
@@ -176,19 +189,21 @@ class ValueBoundsConstraintSet
/// `ValueBoundsOpInterface` for each visited value.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
- /// `closedUB` is set to "true", upper bounds are also closed.
+ /// `options.closedUB` is set to "true", upper bounds are also closed.
static LogicalResult
computeBound(AffineMap &resultMap, ValueDimList &mapOperands,
presburger::BoundType type, const Variable &var,
- StopConditionFn stopCondition, bool closedUB = false);
+ StopConditionFn stopCondition, ValueBoundsOptions options = {});
/// Compute a bound in terms of the values/dimensions in `dependencies`. The
/// computed bound consists of only constant terms and dependent values (or
/// dimension sizes thereof).
- static LogicalResult
- computeDependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
- presburger::BoundType type, const Variable &var,
- ValueDimList dependencies, bool closedUB = false);
+ static LogicalResult computeDependentBound(AffineMap &resultMap,
+ ValueDimList &mapOperands,
+ presburger::BoundType type,
+ const Variable &var,
+ ValueDimList dependencies,
+ ValueBoundsOptions options = {});
/// Compute a bound in that is independent of all values in `independencies`.
///
@@ -198,10 +213,12 @@ class ValueBoundsConstraintSet
/// must be made independent of loop induction variables (in the case of "for"
/// loops). Loop induction variables are the independencies; they may not
/// appear in the computed bound.
- static LogicalResult
- computeIndependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
- presburger::BoundType type, const Variable &var,
- ValueRange independencies, bool closedUB = false);
+ static LogicalResult computeIndependentBound(AffineMap &resultMap,
+ ValueDimList &mapOperands,
+ presburger::BoundType type,
+ const Variable &var,
+ ValueRange independencies,
+ ValueBoundsOptions options = {});
/// Compute a constant bound for the given variable.
///
@@ -216,11 +233,11 @@ class ValueBoundsConstraintSet
/// computed.
///
/// By default, lower/equal bounds are closed and upper bounds are open. If
- /// `closedUB` is set to "true", upper bounds are also closed.
+ /// `options.closedUB` is set to "true", upper bounds are also closed.
static FailureOr<int64_t>
computeConstantBound(presburger::BoundType type, const Variable &var,
const StopConditionFn &stopCondition = nullptr,
- bool closedUB = false);
+ ValueBoundsOptions options = {});
/// Compute a constant delta between the given two values. Return "failure"
/// if a constant delta could not be determined.
@@ -330,6 +347,7 @@ class ValueBoundsConstraintSet
ValueBoundsConstraintSet(MLIRContext *ctx,
const StopConditionFn &stopCondition,
+ ValueBoundsOptions options = {},
bool addConservativeSemiAffineBounds = false);
/// Return "true" if, based on the current state of the constraint system,
@@ -435,6 +453,9 @@ class ValueBoundsConstraintSet
/// The current stop condition function.
StopConditionFn stopCondition = nullptr;
+ /// Options that control value bound computation.
+ ValueBoundsOptions options;
+
/// Should conservative bounds be added for semi-affine expressions.
bool addConservativeSemiAffineBounds = false;
};
diff --git a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
index 9537d3e75c26a..1bd068d6151c6 100644
--- a/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp
@@ -9,6 +9,7 @@
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
@@ -19,12 +20,13 @@ using namespace mlir::affine;
FailureOr<OpFoldResult> mlir::affine::reifyValueBound(
OpBuilder &b, Location loc, presburger::BoundType type,
const ValueBoundsConstraintSet::Variable &var,
- ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
+ ValueBoundsConstraintSet::StopConditionFn stopCondition,
+ ValueBoundsOptions options) {
// Compute bound.
AffineMap boundMap;
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeBound(
- boundMap, mapOperands, type, var, stopCondition, closedUB)))
+ boundMap, mapOperands, type, var, stopCondition, options)))
return failure();
// Reify bound.
@@ -41,8 +43,12 @@ OpFoldResult affine::materializeComputedBound(
std::optional<int64_t> dim = valueDim.second;
if (!dim.has_value()) {
- // This is an index-typed value.
- assert(value.getType().isIndex() && "expected index type");
+ // This is an index-typed/integer-typed value.
+ assert(
+ (value.getType().isIndex() || value.getType().isSignlessInteger()) &&
+ "expected index or signless integer type");
+ if (value.getType().isSignlessInteger())
+ value = arith::IndexCastOp::create(b, loc, b.getIndexType(), value);
operands.push_back(value);
continue;
}
@@ -82,7 +88,7 @@ OpFoldResult affine::materializeComputedBound(
FailureOr<OpFoldResult> mlir::affine::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,
- bool closedUB) {
+ ValueBoundsOptions options) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
// We are trying to reify a bound for `value` in terms of the owning op's
@@ -94,17 +100,18 @@ FailureOr<OpFoldResult> mlir::affine::reifyShapedValueDimBound(
};
return reifyValueBound(b, loc, type, {value, dim},
stopCondition ? stopCondition : reifyToOperands,
- closedUB);
+ options);
}
FailureOr<OpFoldResult> mlir::affine::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
- ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
+ ValueBoundsConstraintSet::StopConditionFn stopCondition,
+ ValueBoundsOptions options) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
return v != value;
};
return reifyValueBound(b, loc, type, value,
stopCondition ? stopCondition : reifyToOperands,
- closedUB);
+ options);
}
diff --git a/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp b/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
index 127563c8f4967..6f02575767e8e 100644
--- a/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
+++ b/mlir/lib/Dialect/Arith/Transforms/ReifyValueBounds.cpp
@@ -18,6 +18,18 @@
using namespace mlir;
using namespace mlir::arith;
+static bool isIndexLikeType(Type type, ValueBoundsOptions options) {
+ return type.isIndex() || (options.allowIntegerType && type.isInteger());
+}
+
+static Value castToIndexValue(OpBuilder &b, Location loc, Value value) {
+ if (value.getType().isIndex())
+ return value;
+ assert(value.getType().isSignlessInteger() &&
+ "expected index or signless integer type");
+ return IndexCastOp::create(b, loc, b.getIndexType(), value);
+}
+
/// Build Arith IR for the given affine map and its operands.
static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
ValueRange operands) {
@@ -28,10 +40,12 @@ static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
return ConstantIndexOp::create(b, loc,
cast<AffineConstantExpr>(e).getValue());
case AffineExprKind::DimId:
- return operands[cast<AffineDimExpr>(e).getPosition()];
+ return castToIndexValue(b, loc,
+ operands[cast<AffineDimExpr>(e).getPosition()]);
case AffineExprKind::SymbolId:
- return operands[cast<AffineSymbolExpr>(e).getPosition() +
- map.getNumDims()];
+ return castToIndexValue(
+ b, loc,
+ operands[cast<AffineSymbolExpr>(e).getPosition() + map.getNumDims()]);
case AffineExprKind::Add: {
auto binaryExpr = cast<AffineBinaryOpExpr>(e);
return AddIOp::create(b, loc, buildExpr(binaryExpr.getLHS()),
@@ -66,13 +80,13 @@ static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
OpBuilder &b, Location loc, presburger::BoundType type,
const ValueBoundsConstraintSet::Variable &var,
- ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
+ ValueBoundsConstraintSet::StopConditionFn stopCondition,
+ ValueBoundsOptions options) {
// Compute bound.
AffineMap boundMap;
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeBound(
- boundMap, mapOperands, type, var, std::move(stopCondition),
- closedUB)))
+ boundMap, mapOperands, type, var, std::move(stopCondition), options)))
return failure();
// Materialize tensor.dim/memref.dim ops.
@@ -82,8 +96,9 @@ FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
std::optional<int64_t> dim = valueDim.second;
if (!dim.has_value()) {
- // This is an index-typed value.
- assert(value.getType().isIndex() && "expected index type");
+ // This is an index-typed/integer-typed value.
+ assert(isIndexLikeType(value.getType(), options) &&
+ "expected index or integer type");
operands.push_back(value);
continue;
}
@@ -109,10 +124,11 @@ FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
}
// No arith ops are needed if the bound is a single SSA value.
if (auto expr = dyn_cast<AffineDimExpr>(boundMap.getResult(0)))
- return static_cast<OpFoldResult>(operands[expr.getPosition()]);
- if (auto expr = dyn_cast<AffineSymbolExpr>(boundMap.getResult(0)))
return static_cast<OpFoldResult>(
- operands[expr.getPosition() + boundMap.getNumDims()]);
+ castToIndexValue(b, loc, operands[expr.getPosition()]));
+ if (auto expr = dyn_cast<AffineSymbolExpr>(boundMap.getResult(0)))
+ return static_cast<OpFoldResult>(castToIndexValue(
+ b, loc, operands[expr.getPosition() + boundMap.getNumDims()]));
// General case: build Arith ops.
return static_cast<OpFoldResult>(buildArithValue(b, loc, boundMap, operands));
}
@@ -120,7 +136,7 @@ FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
FailureOr<OpFoldResult> mlir::arith::reifyShapedValueDimBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
int64_t dim, const ValueBoundsConstraintSet::StopConditionFn &stopCondition,
- bool closedUB) {
+ ValueBoundsOptions options) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
// We are trying to reify a bound for `value` in terms of the owning op's
@@ -132,18 +148,18 @@ FailureOr<OpFoldResult> mlir::arith::reifyShapedValueDimBound(
};
return reifyValueBound(b, loc, type, {value, dim},
stopCondition ? stopCondition : reifyToOperands,
- closedUB);
+ options);
}
FailureOr<OpFoldResult> mlir::arith::reifyIndexValueBound(
OpBuilder &b, Location loc, presburger::BoundType type, Value value,
const ValueBoundsConstraintSet::StopConditionFn &stopCondition,
- bool closedUB) {
+ ValueBoundsOptions options) {
auto reifyToOperands = [&](Value v, std::optional<int64_t> d,
ValueBoundsConstraintSet &cstr) {
return v != value;
};
return reifyValueBound(b, loc, type, value,
stopCondition ? stopCondition : reifyToOperands,
- closedUB);
+ options);
}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index 867171583b010..85ec1c19384de 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -483,7 +483,7 @@ HoistPaddingAnalysis::getHoistedPackedTensorSizes(RewriterBase &rewriter,
return !isa<affine::AffineMinOp, affine::AffineMaxOp,
affine::AffineApplyOp>(op);
},
- /*closedUB=*/true);
+ ValueBoundsOptions{/*closedUB=*/true});
assert(succeeded(loopUb) && "could not get upper bound");
Value ubVal = getValueOrCreateConstantIndexOp(rewriter, loc, *loopUb);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index e1dc40d6d37d9..a573b4a54dbba 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -227,7 +227,8 @@ void mlir::linalg::hoistRedundantVectorTransfers(Operation *root,
FailureOr<int64_t> maxLb =
ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, lb,
- /*stopCondition=*/nullptr, /*closedUB=*/true);
+ /*stopCondition=*/nullptr,
+ ValueBoundsOptions{/*closedUB=*/true});
if (failed(maxLb))
return;
FailureOr<int64_t> minUb =
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp b/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
index dd8437985197d..9b8ec2f186a1b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Padding.cpp
@@ -116,7 +116,7 @@ FailureOr<bool> PaddedShape::initialize(linalg::LinalgOp opToPad,
presburger::BoundType::UB,
{opOperand->get(),
/*dim=*/i},
- /*stopCondition=*/nullptr, /*closedUB=*/true);
+ /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
if (failed(upperBound)) {
LLVM_DEBUG(
DBGS() << "----could not compute a bounding box for padding\n");
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
index 6519c4f64dd05..cd71db7ec783a 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Promotion.cpp
@@ -261,7 +261,7 @@ FailureOr<PromotionInfo> mlir::linalg::promoteSubviewAsNewBuffer(
FailureOr<int64_t> upperBound =
ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, rangeValue.size,
- /*stopCondition=*/nullptr, /*closedUB=*/true);
+ /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
size = failed(upperBound)
? getValueOrCreateConstantIndexOp(b, loc, rangeValue.size)
: arith::ConstantIndexOp::create(b, loc, *upperBound);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index 7ed07e1ec9a01..4eaa7bf0233c6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -1157,7 +1157,8 @@ struct PackOpTiling
FailureOr<int64_t> cstTileSize =
ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, sizes[dim],
- /*stopCondition=*/nullptr, /*closedUB=*/true);
+ /*stopCondition=*/nullptr,
+ ValueBoundsOptions{/*closedUB=*/true});
std::optional<int64_t> cstInnerSize =
getConstantIntValue(dimAndTileMapping[dim]);
@@ -1326,7 +1327,7 @@ static UnpackTileDimInfo getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp,
info.isAlignedToInnerTileSize = false;
FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, tileSize,
- /*stopCondition=*/nullptr, /*closedUB=*/true);
+ /*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
std::optional<int64_t> cstInnerSize = getConstantIntValue(innerTileSize);
if (!failed(cstSize) && cstInnerSize) {
if (*cstSize % *cstInnerSize == 0)
diff --git a/mlir/lib/Dialect/MemRef/Transforms/IndependenceTransforms.cpp b/mlir/lib/Dialect/MemRef/Transforms/IndependenceTransforms.cpp
index d5e2b97e501e6..e4ee029d93c44 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/IndependenceTransforms.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/IndependenceTransforms.cpp
@@ -25,7 +25,7 @@ static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc,
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeIndependentBound(
boundMap, mapOperands, presburger::BoundType::UB, ofr, independencies,
- /*closedUB=*/true)))
+ ValueBoundsOptions{/*closedUB=*/true})))
return failure();
return affine::materializeComputedBound(b, loc, boundMap, mapOperands);
}
diff --git a/mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp b/mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp
index 9fd27d328694e..c80c7fb334f34 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/IndependenceTransforms.cpp
@@ -27,8 +27,7 @@ static FailureOr<OpFoldResult> makeIndependent(OpBuilder &b, Location loc,
ValueDimList mapOperands;
if (failed(ValueBoundsConstraintSet::computeIndependentBound(
boundMap, mapOperands, presburger::BoundType::UB, value,
- independencies,
- /*closedUB=*/true)))
+ independencies, ValueBoundsOptions{/*closedUB=*/true})))
return failure();
return mlir::affine::materializeComputedBound(b, loc, boundMap, mapOperands);
}
diff --git a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
index 2986f4c2d607d..ed41d19d627a3 100644
--- a/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
+++ b/mlir/lib/Dialect/Vector/IR/ScalableValueBoundsConstraintSet.cpp
@@ -42,8 +42,8 @@ char ScalableValueBoundsConstraintSet::ID = 0;
FailureOr<ConstantOrScalableBound>
ScalableValueBoundsConstraintSet::computeScalableBound(
Value value, std::optional<int64_t> dim, unsigned vscaleMin,
- unsigned vscaleMax, presburger::BoundType boundType, bool closedUB,
- const StopConditionFn &stopCondition) {
+ unsigned vscaleMax, presburger::BoundType boundType,
+ ValueBoundsOptions options, const StopConditionFn &stopCondition) {
using namespace presburger;
assert(vscaleMin <= vscaleMax);
@@ -56,7 +56,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
ScalableValueBoundsConstraintSet scalableCstr(
value.getContext(), stopCondition ? stopCondition : defaultStopCondition,
- vscaleMin, vscaleMax);
+ vscaleMin, vscaleMax, options);
int64_t pos = scalableCstr.insert(value, dim, /*isSymbol=*/false);
scalableCstr.processWorklist();
@@ -98,7 +98,7 @@ ScalableValueBoundsConstraintSet::computeScalableBound(
SmallVector<AffineMap, 1> lowerBound(1), upperBound(1);
scalableCstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lowerBound,
- &upperBound, closedUB);
+ &upperBound, options.closedUB);
auto invalidBound = [](auto &bound) {
return !bound[0] || bound[0].getNumResults() != 1;
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index c40a0e0114752..02579673d2293 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -77,6 +77,14 @@ static std::optional<int64_t> getConstantIntValue(OpFoldResult ofr) {
return std::nullopt;
}
+static bool isIndexOrIntegerType(Type type) {
+ return type.isIndex() || type.isInteger();
+}
+
+static bool isIndexLikeType(Type type, ValueBoundsOptions options) {
+ return type.isIndex() || (options.allowIntegerType && type.isInteger());
+}
+
ValueBoundsConstraintSet::Variable::Variable(OpFoldResult ofr)
: Variable(ofr, std::nullopt) {}
@@ -90,7 +98,7 @@ ValueBoundsConstraintSet::Variable::Variable(OpFoldResult ofr,
std::optional<int64_t> dim) {
Builder b(ofr.getContext());
if (auto constInt = ::getConstantIntValue(ofr)) {
- assert(!dim && "expected no dim for index-typed values");
+ assert(!dim && "expected no dim for index/integer-typed values");
map = AffineMap::get(/*dimCount=*/0, /*symbolCount=*/0,
b.getAffineConstantExpr(*constInt));
return;
@@ -100,7 +108,8 @@ ValueBoundsConstraintSet::Variable::Variable(OpFoldResult ofr,
if (dim) {
assert(isa<ShapedType>(value.getType()) && "expected shaped type");
} else {
- assert(value.getType().isIndex() && "expected index type");
+ assert(isIndexOrIntegerType(value.getType()) &&
+ "expected index or integer type");
}
#endif // NDEBUG
map = AffineMap::get(/*dimCount=*/0, /*symbolCount=*/1,
@@ -158,8 +167,8 @@ ValueBoundsConstraintSet::Variable::Variable(AffineMap map,
ValueBoundsConstraintSet::ValueBoundsConstraintSet(
MLIRContext *ctx, const StopConditionFn &stopCondition,
- bool addConservativeSemiAffineBounds)
- : builder(ctx), stopCondition(stopCondition),
+ ValueBoundsOptions options, bool addConservativeSemiAffineBounds)
+ : builder(ctx), stopCondition(stopCondition), options(options),
addConservativeSemiAffineBounds(addConservativeSemiAffineBounds) {
assert(stopCondition && "expected non-null stop condition");
}
@@ -167,8 +176,9 @@ ValueBoundsConstraintSet::ValueBoundsConstraintSet(
char ValueBoundsConstraintSet::ID = 0;
#ifndef NDEBUG
-static void assertValidValueDim(Value value, std::optional<int64_t> dim) {
- if (value.getType().isIndex()) {
+static void assertValidValueDim(Value value, std::optional<int64_t> dim,
+ ValueBoundsOptions options) {
+ if (isIndexLikeType(value.getType(), options)) {
assert(!dim.has_value() && "invalid dim value");
} else if (auto shapedType = dyn_cast<ShapedType>(value.getType())) {
assert(*dim >= 0 && "invalid dim value");
@@ -206,7 +216,7 @@ void ValueBoundsConstraintSet::addBound(BoundType type, int64_t pos,
AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
std::optional<int64_t> dim) {
#ifndef NDEBUG
- assertValidValueDim(value, dim);
+ assertValidValueDim(value, dim, options);
#endif // NDEBUG
// Check if the value/dim is statically known. In that case, an affine
@@ -268,7 +278,7 @@ int64_t ValueBoundsConstraintSet::insert(Value value,
std::optional<int64_t> dim,
bool isSymbol, bool addToWorklist) {
#ifndef NDEBUG
- assertValidValueDim(value, dim);
+ assertValidValueDim(value, dim, options);
#endif // NDEBUG
ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
@@ -344,7 +354,7 @@ int64_t ValueBoundsConstraintSet::insert(const Variable &var, bool isSymbol) {
int64_t ValueBoundsConstraintSet::getPos(Value value,
std::optional<int64_t> dim) const {
#ifndef NDEBUG
- assertValidValueDim(value, dim);
+ assertValidValueDim(value, dim, options);
#endif // NDEBUG
LDBG() << "Getting pos for: " << value
<< " (dim: " << dim.value_or(kIndexValue)
@@ -471,15 +481,16 @@ void ValueBoundsConstraintSet::projectOutAnonymous(
LogicalResult ValueBoundsConstraintSet::computeBound(
AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
- const Variable &var, StopConditionFn stopCondition, bool closedUB) {
+ const Variable &var, StopConditionFn stopCondition,
+ ValueBoundsOptions options) {
MLIRContext *ctx = var.getContext();
- int64_t ubAdjustment = closedUB ? 0 : 1;
+ int64_t ubAdjustment = options.closedUB ? 0 : 1;
Builder b(ctx);
mapOperands.clear();
// Process the backward slice of `value` (i.e., reverse use-def chain) until
// `stopCondition` is met.
- ValueBoundsConstraintSet cstr(ctx, stopCondition);
+ ValueBoundsConstraintSet cstr(ctx, stopCondition, options);
int64_t pos = cstr.insert(var, /*isSymbol=*/false);
assert(pos == 0 && "expected first column");
cstr.processWorklist();
@@ -573,9 +584,10 @@ LogicalResult ValueBoundsConstraintSet::computeBound(
Value value = valueDim.first;
int64_t dim = valueDim.second;
if (dim == ValueBoundsConstraintSet::kIndexValue) {
- // An index-type value is used: can be used directly in the affine.apply
- // op.
- assert(value.getType().isIndex() && "expected index type");
+ // An index-typed/integer-typed value is used: it can be used directly in
+ // the computed bound.
+ assert(isIndexLikeType(value.getType(), options) &&
+ "expected index or integer type");
mapOperands.push_back(std::make_pair(value, std::nullopt));
continue;
}
@@ -592,18 +604,20 @@ LogicalResult ValueBoundsConstraintSet::computeBound(
LogicalResult ValueBoundsConstraintSet::computeDependentBound(
AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
- const Variable &var, ValueDimList dependencies, bool closedUB) {
+ const Variable &var, ValueDimList dependencies,
+ ValueBoundsOptions options) {
return computeBound(
resultMap, mapOperands, type, var,
[&](Value v, std::optional<int64_t> d, ValueBoundsConstraintSet &cstr) {
return llvm::is_contained(dependencies, std::make_pair(v, d));
},
- closedUB);
+ options);
}
LogicalResult ValueBoundsConstraintSet::computeIndependentBound(
AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
- const Variable &var, ValueRange independencies, bool closedUB) {
+ const Variable &var, ValueRange independencies,
+ ValueBoundsOptions options) {
// Return "true" if the given value is independent of all values in
// `independencies`. I.e., neither the value itself nor any value in the
// backward slice (reverse use-def chain) is contained in `independencies`.
@@ -632,12 +646,12 @@ LogicalResult ValueBoundsConstraintSet::computeIndependentBound(
[&](Value v, std::optional<int64_t> d, ValueBoundsConstraintSet &cstr) {
return isIndependent(v);
},
- closedUB);
+ options);
}
FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType type, const Variable &var,
- const StopConditionFn &stopCondition, bool closedUB) {
+ const StopConditionFn &stopCondition, ValueBoundsOptions options) {
// Default stop condition if none was specified: Keep adding constraints until
// a bound could be computed.
int64_t pos = 0;
@@ -647,12 +661,13 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
};
ValueBoundsConstraintSet cstr(
- var.getContext(), stopCondition ? stopCondition : defaultStopCondition);
+ var.getContext(), stopCondition ? stopCondition : defaultStopCondition,
+ options);
pos = cstr.populateConstraints(var.map, var.mapOperands);
assert(pos == 0 && "expected `map` is the first column");
// Compute constant bound for `valueDim`.
- int64_t ubAdjustment = closedUB ? 0 : 1;
+ int64_t ubAdjustment = options.closedUB ? 0 : 1;
if (auto bound = cstr.cstr.getConstantBound64(type, pos))
return type == BoundType::UB ? *bound + ubAdjustment : *bound;
return failure();
@@ -661,7 +676,7 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
void ValueBoundsConstraintSet::populateConstraints(Value value,
std::optional<int64_t> dim) {
#ifndef NDEBUG
- assertValidValueDim(value, dim);
+ assertValidValueDim(value, dim, options);
#endif // NDEBUG
// `getExpr` pushes the value/dim onto the worklist (unless it was already
@@ -686,8 +701,8 @@ ValueBoundsConstraintSet::computeConstantDelta(Value value1, Value value2,
std::optional<int64_t> dim1,
std::optional<int64_t> dim2) {
#ifndef NDEBUG
- assertValidValueDim(value1, dim1);
- assertValidValueDim(value2, dim2);
+ assertValidValueDim(value1, dim1, /*options=*/{});
+ assertValidValueDim(value2, dim2, /*options=*/{});
#endif // NDEBUG
Builder b(value1.getContext());
@@ -973,14 +988,14 @@ ValueBoundsConstraintSet::BoundBuilder::operator[](int64_t dim) {
assert(!this->dim.has_value() && "dim was already set");
this->dim = dim;
#ifndef NDEBUG
- assertValidValueDim(value, this->dim);
+ assertValidValueDim(value, this->dim, cstr.options);
#endif // NDEBUG
return *this;
}
void ValueBoundsConstraintSet::BoundBuilder::operator<(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim);
+ assertValidValueDim(value, this->dim, cstr.options);
#endif // NDEBUG
cstr.addBound(BoundType::UB, cstr.getPos(value, this->dim), expr);
}
@@ -995,14 +1010,14 @@ void ValueBoundsConstraintSet::BoundBuilder::operator>(AffineExpr expr) {
void ValueBoundsConstraintSet::BoundBuilder::operator>=(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim);
+ assertValidValueDim(value, this->dim, cstr.options);
#endif // NDEBUG
cstr.addBound(BoundType::LB, cstr.getPos(value, this->dim), expr);
}
void ValueBoundsConstraintSet::BoundBuilder::operator==(AffineExpr expr) {
#ifndef NDEBUG
- assertValidValueDim(value, this->dim);
+ assertValidValueDim(value, this->dim, cstr.options);
#endif // NDEBUG
cstr.addBound(BoundType::EQ, cstr.getPos(value, this->dim), expr);
}
diff --git a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
index e72457c776d69..c2fe61b0b92f9 100644
--- a/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
+++ b/mlir/test/Dialect/Affine/invalid-reify-bound-dim.mlir
@@ -28,7 +28,7 @@ func.func @test_invalid_reify_int_value(%size: index) -> (index) {
%zero = arith.constant 0 : index
%int_val = arith.constant 1 : index
- // expected-error at +1 {{'test.reify_bound' op unexpected 'dim' attribute for index variable}}
+ // expected-error at +1 {{'test.reify_bound' op unexpected 'dim' attribute for index/integer variable}}
%dim = "test.reify_bound"(%int_val) {dim = 1 : i64} : (index) -> index
return %dim: index
diff --git a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
index 66de08bf4e219..78802a282f096 100644
--- a/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir
@@ -55,6 +55,29 @@ func.func @arith_muli(%a: index) -> index {
// -----
+// CHECK: #[[$map_muli_i32:.*]] = affine_map<()[s0] -> (s0 * 7)>
+// CHECK-LABEL: func @arith_muli_integer(
+// CHECK-SAME: %[[a:.*]]: i32
+// CHECK: %[[cast:.*]] = arith.index_cast %[[a]] : i32 to index
+// CHECK: %[[apply:.*]] = affine.apply #[[$map_muli_i32]]()[%[[cast]]]
+// CHECK: return %[[apply]]
+// CHECK-ARITH-LABEL: func @arith_muli_integer(
+// CHECK-ARITH-SAME: %[[a:.*]]: i32
+// CHECK-ARITH: %[[c7:.*]] = arith.constant 7 : i32
+// CHECK-ARITH: arith.muli %[[a]], %[[c7]] : i32
+// CHECK-ARITH-DAG: %[[cast:.*]] = arith.index_cast %[[a]] : i32 to index
+// CHECK-ARITH-DAG: %[[c7_reified:.*]] = arith.constant 7 : index
+// CHECK-ARITH: %[[mul:.*]] = arith.muli %[[cast]], %[[c7_reified]] : index
+// CHECK-ARITH: return %[[mul]]
+func.func @arith_muli_integer(%a: i32) -> index {
+ %c7 = arith.constant 7 : i32
+ %product = arith.muli %a, %c7 : i32
+ %0 = "test.reify_bound"(%product) {allow_integer_type} : (i32) -> (index)
+ return %0 : index
+}
+
+// -----
+
func.func @arith_muli_non_pure(%a: index, %b: index) -> index {
%0 = arith.muli %a, %b : index
// Semi-affine expressions (such as "symbol * symbol") are not supported.
@@ -101,6 +124,22 @@ func.func @arith_const() -> index {
// -----
+// CHECK-LABEL: func @arith_addi_integer_constant()
+// CHECK: %[[c12:.*]] = arith.constant 12 : index
+// CHECK: return %[[c12]]
+// CHECK-ARITH-LABEL: func @arith_addi_integer_constant()
+// CHECK-ARITH: %[[c12:.*]] = arith.constant 12 : index
+// CHECK-ARITH: return %[[c12]]
+func.func @arith_addi_integer_constant() -> index {
+ %c5 = arith.constant 5 : i32
+ %c7 = arith.constant 7 : i32
+ %sum = arith.addi %c5, %c7 : i32
+ %0 = "test.reify_bound"(%sum) {allow_integer_type, constant} : (i32) -> (index)
+ return %0 : index
+}
+
+// -----
+
// CHECK-LABEL: func @arith_select(
func.func @arith_select(%c: i1) -> (index, index) {
// CHECK: arith.constant 5 : index
diff --git a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
index 3569a7399ad0a..4aa1b134a1078 100644
--- a/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -106,6 +107,8 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
bool constant = op.getConstant();
bool scalable = op.getScalable();
+ ValueBoundsOptions options;
+ options.allowIntegerType = op.getAllowIntegerType();
// Prepare stop condition. By default, reify in terms of the op's
// operands. No stop condition is used when a constant was requested.
@@ -132,14 +135,16 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
FailureOr<OpFoldResult> reified = failure();
if (constant) {
auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
- boundType, {value, dim}, /*stopCondition=*/nullptr);
+ boundType, {value, dim}, /*stopCondition=*/nullptr, options);
if (succeeded(reifiedConst))
reified = FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
} else if (scalable) {
auto loc = op->getLoc();
+ options.closedUB = true;
auto reifiedScalable =
vector::ScalableValueBoundsConstraintSet::computeScalableBound(
- value, dim, *op.getVscaleMin(), *op.getVscaleMax(), boundType);
+ value, dim, *op.getVscaleMin(), *op.getVscaleMax(), boundType,
+ options);
if (succeeded(reifiedScalable)) {
SmallVector<std::pair<Value, std::optional<int64_t>>, 1> vscaleOperand;
if (reifiedScalable->map.getNumInputs() == 1) {
@@ -152,11 +157,12 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
}
} else {
if (useArithOps) {
- reified = arith::reifyValueBound(rewriter, op->getLoc(), boundType,
- op.getVariable(), stopCondition);
+ reified =
+ arith::reifyValueBound(rewriter, op->getLoc(), boundType,
+ op.getVariable(), stopCondition, options);
} else {
reified = reifyValueBound(rewriter, op->getLoc(), boundType,
- op.getVariable(), stopCondition);
+ op.getVariable(), stopCondition, options);
}
}
if (failed(reified)) {
@@ -169,9 +175,10 @@ static LogicalResult testReifyValueBounds(FunctionOpInterface funcOp,
rewriter.replaceOp(op, val);
return WalkResult::skip();
}
- Value constOp = arith::ConstantIndexOp::create(
+ auto attr = cast<IntegerAttr>(cast<Attribute>(*reified));
+ Value constOp = arith::ConstantOp::create(
rewriter, op->getLoc(),
- cast<IntegerAttr>(cast<Attribute>(*reified)).getInt());
+ rewriter.getIntegerAttr(op.getResult().getType(), attr.getInt()));
rewriter.replaceOp(op, constOp);
return WalkResult::skip();
});
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 340b44b14dd96..ec79b26c650a8 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1062,11 +1062,15 @@ LogicalResult ReifyBoundOp::verify() {
if (isa<ShapedType>(getVar().getType())) {
if (!getDim().has_value())
return emitOpError("expected 'dim' attribute for shaped type variable");
- } else if (getVar().getType().isIndex()) {
+ } else if (getVar().getType().isIndex() || getVar().getType().isInteger()) {
+ if (getVar().getType().isInteger() && !getAllowIntegerType())
+ return emitOpError("integer variable requires 'allow_integer_type'");
if (getDim().has_value())
- return emitOpError("unexpected 'dim' attribute for index variable");
+ return emitOpError(
+ "unexpected 'dim' attribute for index/integer variable");
} else {
- return emitOpError("expected index-typed variable or shape type variable");
+ return emitOpError(
+ "expected index-typed/integer-typed variable or shape type variable");
}
if (getConstant() && getScalable())
return emitOpError("'scalable' and 'constant' are mutually exlusive");
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 348ff5d7f4ea0..2df06579984a7 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2630,6 +2630,7 @@ def ReifyBoundOp : TEST_Op<"reify_bound", [Pure]> {
OptionalAttr<I64Attr>:$dim,
DefaultValuedAttr<StrAttr, "\"EQ\"">:$type,
UnitAttr:$constant,
+ UnitAttr:$allow_integer_type,
UnitAttr:$scalable,
OptionalAttr<I64Attr>:$vscale_min,
OptionalAttr<I64Attr>:$vscale_max);
More information about the Mlir-commits
mailing list