Skip to content

Commit cce9117

Browse files
cferry-AMDcferry-AMD
authored
Merge pull request #482 from Xilinx/corentin.affine_expr_bounds
[FXML-5704] Compute affine expression bounds
2 parents 2b2e860 + a8c0978 commit cce9117

File tree

7 files changed

+678
-0
lines changed

7 files changed

+678
-0
lines changed

‎mlir/include/mlir/Analysis/AffineExprBounds.h

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
//===- AffineExprBounds.h - Compute bounds of affine expressions *- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This header file defines an analysis of affine expressions to compute their
10+
// ranges (lower/upper bounds) in a given context.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
#ifndef MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H
14+
#define MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H
15+
16+
#include "mlir/IR/AffineExprVisitor.h"
17+
#include "mlir/IR/Attributes.h"
18+
#include "mlir/IR/BuiltinAttributes.h"
19+
#include "mlir/Interfaces/InferIntRangeInterface.h"
20+
21+
#include "mlir/IR/AffineExpr.h"
22+
#include "mlir/IR/AffineMap.h"
23+
#include "mlir/Support/LogicalResult.h"
24+
25+
using namespace mlir;
26+
27+
/// This visitor computes the bounds of affine expressions, using as context the
28+
/// bounds of the dimensions of the expression.
29+
///
30+
/// Example:
31+
/// Given bounds 0 <= d0 <= 99 and 0 <= d1 <= 199, we can compute the bounds
32+
/// of the following expression:
33+
/// lb(2 * d0 + 3 * d1) = 0
34+
/// ub(2 * d0 + 3 * d1) = 795
35+
///
36+
/// * The bounds given in the context are inclusive, and the bounds returned
37+
/// are also inclusive.
38+
/// * If bounds are not available for a dimension, std::nullopt can be used
39+
/// instead. The bounds of an expression that involves it will be std::nullopt.
40+
/// * Limitations:
41+
/// - Parametric expressions (using symbols) are not supported.
42+
/// - Unsigned FloorDiv is currently not supported.
43+
class AffineExprBoundsVisitor
44+
: public AffineExprVisitor<AffineExprBoundsVisitor, LogicalResult> {
45+
public:
46+
/// Initialize the context (bounds) with APInt. All bounds must have the same
47+
/// signedness and bit width.
48+
AffineExprBoundsVisitor(ArrayRef<std::optional<APInt>> constLowerBounds,
49+
ArrayRef<std::optional<APInt>> constUpperBounds,
50+
bool boundsSigned, uint64_t bitWidth,
51+
MLIRContext *context);
52+
53+
/// Initialize the context (bounds) with 64-bit signed integers. This allows
54+
/// to directly map index-type values such as Linalg op bounds, which are
55+
/// represented as int64_t.
56+
AffineExprBoundsVisitor(ArrayRef<std::optional<int64_t>> constLowerBounds,
57+
ArrayRef<std::optional<int64_t>> constUpperBounds,
58+
MLIRContext *context);
59+
60+
/// Get the upper bound of \p expr using the context bounds.
61+
std::optional<APInt> getUpperBound(AffineExpr expr);
62+
std::optional<int64_t> getIndexUpperBound(AffineExpr expr);
63+
64+
/// Get the lower bound of \p expr using the context bounds.
65+
std::optional<APInt> getLowerBound(AffineExpr expr);
66+
std::optional<int64_t> getIndexLowerBound(AffineExpr expr);
67+
68+
// These methods are directly called by the AffineExprVisitor base class.
69+
LogicalResult visitMulExpr(AffineBinaryOpExpr expr);
70+
LogicalResult visitAddExpr(AffineBinaryOpExpr expr);
71+
LogicalResult visitDimExpr(AffineDimExpr expr);
72+
LogicalResult visitSymbolExpr(AffineSymbolExpr expr);
73+
LogicalResult visitConstantExpr(AffineConstantExpr expr);
74+
LogicalResult visitCeilDivExpr(AffineBinaryOpExpr expr);
75+
LogicalResult visitFloorDivExpr(AffineBinaryOpExpr expr);
76+
LogicalResult visitModExpr(AffineBinaryOpExpr expr);
77+
78+
private:
79+
bool boundsSigned;
80+
uint64_t bitWidth;
81+
void inferBinOpRange(
82+
AffineBinaryOpExpr expr,
83+
const std::function<ConstantIntRanges(ArrayRef<ConstantIntRanges>)>
84+
&opInference);
85+
86+
/// Bounds that have been computed for subexpressions are memoized and reused.
87+
llvm::DenseMap<AffineExpr, APInt> lb;
88+
llvm::DenseMap<AffineExpr, APInt> ub;
89+
};
90+
91+
#endif // MLIR_ANALYSIS_AFFINEEXPRBOUNDS_H

‎mlir/lib/Analysis/AffineExprBounds.cpp

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,186 @@
1+
//===- AffineExprBounds.h - Compute bounds of affine expressions *- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// This file implements an analysis of affine expressions to compute their
10+
// ranges (lower/upper bounds) in a given context.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
#include "mlir/Analysis/AffineExprBounds.h"
14+
15+
#include "mlir/IR/AffineExprVisitor.h"
16+
#include "mlir/IR/AffineMap.h"
17+
#include "mlir/IR/BuiltinAttributes.h"
18+
#include "mlir/Interfaces/InferIntRangeInterface.h"
19+
#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
20+
#include "llvm/ADT/APInt.h"
21+
22+
#include <cstdint>
23+
24+
using namespace mlir;
25+
26+
AffineExprBoundsVisitor::AffineExprBoundsVisitor(
27+
ArrayRef<std::optional<APInt>> constLowerBounds,
28+
ArrayRef<std::optional<APInt>> constUpperBounds, bool boundsSigned,
29+
uint64_t bitWidth, MLIRContext *context)
30+
: boundsSigned(boundsSigned), bitWidth(bitWidth) {
31+
assert(constLowerBounds.size() == constUpperBounds.size());
32+
for (unsigned i = 0; i < constLowerBounds.size(); i++) {
33+
if (constLowerBounds[i].has_value()) {
34+
lb[getAffineDimExpr(i, context)] = constLowerBounds[i].value();
35+
}
36+
if (constUpperBounds[i].has_value()) {
37+
ub[getAffineDimExpr(i, context)] = constUpperBounds[i].value();
38+
}
39+
}
40+
}
41+
42+
AffineExprBoundsVisitor::AffineExprBoundsVisitor(
43+
ArrayRef<std::optional<int64_t>> constLowerBounds,
44+
ArrayRef<std::optional<int64_t>> constUpperBounds, MLIRContext *context)
45+
: boundsSigned(true), bitWidth(64) {
46+
assert(constLowerBounds.size() == constUpperBounds.size());
47+
// Convert int64_ts to APInts.
48+
for (unsigned i = 0; i < constLowerBounds.size(); i++) {
49+
if (constLowerBounds[i].has_value()) {
50+
lb[getAffineDimExpr(i, context)] =
51+
APInt(64, constLowerBounds[i].value(), /*isSigned=*/true);
52+
}
53+
if (constUpperBounds[i].has_value()) {
54+
ub[getAffineDimExpr(i, context)] =
55+
APInt(64, constUpperBounds[i].value(), /*isSigned=*/true);
56+
}
57+
}
58+
}
59+
60+
std::optional<APInt> AffineExprBoundsVisitor::getUpperBound(AffineExpr expr) {
61+
// Use memoized bound if available.
62+
auto i = ub.find(expr);
63+
if (i != ub.end()) {
64+
return i->second;
65+
}
66+
// Compute the bound otherwise.
67+
if (failed(walkPostOrder(expr))) {
68+
return std::nullopt;
69+
}
70+
return ub[expr];
71+
}
72+
73+
std::optional<APInt> AffineExprBoundsVisitor::getLowerBound(AffineExpr expr) {
74+
// Use memoized bound if available.
75+
auto i = lb.find(expr);
76+
if (i != lb.end()) {
77+
return i->second;
78+
}
79+
// Compute the bound otherwise.
80+
if (failed(walkPostOrder(expr))) {
81+
return std::nullopt;
82+
}
83+
return lb[expr];
84+
}
85+
86+
std::optional<int64_t>
87+
AffineExprBoundsVisitor::getIndexUpperBound(AffineExpr expr) {
88+
std::optional<APInt> apIntResult = getUpperBound(expr);
89+
if (!apIntResult)
90+
return std::nullopt;
91+
92+
return apIntResult->getSExtValue();
93+
}
94+
95+
std::optional<int64_t>
96+
AffineExprBoundsVisitor::getIndexLowerBound(AffineExpr expr) {
97+
std::optional<APInt> apIntResult = getLowerBound(expr);
98+
if (!apIntResult)
99+
return std::nullopt;
100+
101+
return apIntResult->getSExtValue();
102+
}
103+
104+
ConstantIntRanges getRange(APInt lb, APInt ub, bool boundsSigned) {
105+
return ConstantIntRanges::range(lb, ub, boundsSigned);
106+
}
107+
108+
/// Wrapper around the intrange::infer* functions that infers the range of
109+
/// binary operations on two ranges.
110+
void AffineExprBoundsVisitor::inferBinOpRange(
111+
AffineBinaryOpExpr expr,
112+
const std::function<ConstantIntRanges(ArrayRef<ConstantIntRanges>)>
113+
&opInference) {
114+
ConstantIntRanges lhsRange =
115+
getRange(lb[expr.getLHS()], ub[expr.getLHS()], boundsSigned);
116+
ConstantIntRanges rhsRange =
117+
getRange(lb[expr.getRHS()], ub[expr.getRHS()], boundsSigned);
118+
ConstantIntRanges result = opInference({lhsRange, rhsRange});
119+
120+
lb[expr] = (boundsSigned) ? result.smin() : result.umin();
121+
ub[expr] = (boundsSigned) ? result.smax() : result.umax();
122+
}
123+
124+
// Visitor method overrides.
125+
LogicalResult AffineExprBoundsVisitor::visitMulExpr(AffineBinaryOpExpr expr) {
126+
inferBinOpRange(expr, [](ArrayRef<ConstantIntRanges> ranges) {
127+
return intrange::inferMul(ranges);
128+
});
129+
return success();
130+
}
131+
LogicalResult AffineExprBoundsVisitor::visitAddExpr(AffineBinaryOpExpr expr) {
132+
inferBinOpRange(expr, [](ArrayRef<ConstantIntRanges> ranges) {
133+
return intrange::inferAdd(ranges);
134+
});
135+
return success();
136+
}
137+
LogicalResult
138+
AffineExprBoundsVisitor::visitCeilDivExpr(AffineBinaryOpExpr expr) {
139+
inferBinOpRange(
140+
expr, [boundsSigned = boundsSigned](ArrayRef<ConstantIntRanges> ranges) {
141+
if (boundsSigned) {
142+
return intrange::inferCeilDivS(ranges);
143+
}
144+
return intrange::inferCeilDivU(ranges);
145+
});
146+
return success();
147+
}
148+
LogicalResult
149+
AffineExprBoundsVisitor::visitFloorDivExpr(AffineBinaryOpExpr expr) {
150+
// There is no inferFloorDivU in the intrange library. We only offer
151+
// computation of bounds for signed floordiv operations.
152+
if (boundsSigned) {
153+
inferBinOpRange(expr, [](ArrayRef<ConstantIntRanges> ranges) {
154+
return intrange::inferFloorDivS(ranges);
155+
});
156+
return success();
157+
}
158+
return failure();
159+
}
160+
LogicalResult AffineExprBoundsVisitor::visitModExpr(AffineBinaryOpExpr expr) {
161+
inferBinOpRange(
162+
expr, [boundsSigned = boundsSigned](ArrayRef<ConstantIntRanges> ranges) {
163+
if (boundsSigned) {
164+
return intrange::inferRemS(ranges);
165+
}
166+
return intrange::inferRemU(ranges);
167+
});
168+
return success();
169+
}
170+
LogicalResult AffineExprBoundsVisitor::visitDimExpr(AffineDimExpr expr) {
171+
if (lb.find(expr) == lb.end() || ub.find(expr) == ub.end()) {
172+
return failure();
173+
}
174+
return success();
175+
}
176+
LogicalResult AffineExprBoundsVisitor::visitSymbolExpr(AffineSymbolExpr expr) {
177+
return failure();
178+
}
179+
LogicalResult
180+
AffineExprBoundsVisitor::visitConstantExpr(AffineConstantExpr expr) {
181+
APInt apIntVal =
182+
APInt(bitWidth, static_cast<uint64_t>(expr.getValue()), boundsSigned);
183+
lb[expr] = apIntVal;
184+
ub[expr] = apIntVal;
185+
return success();
186+
}

‎mlir/lib/Analysis/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ set(LLVM_OPTIONAL_SOURCES
2121
add_subdirectory(Presburger)
2222

2323
add_mlir_library(MLIRAnalysis
24+
AffineExprBounds.cpp
2425
AliasAnalysis.cpp
2526
CallGraph.cpp
2627
DataFlowFramework.cpp

0 commit comments

Comments
 (0)