[Mlir-commits] [mlir] [mlir][Affine][NFC] Define AffineForOp operands in ODS (PR #67694)
Matthias Springer
llvmlistbot at llvm.org
Thu Sep 28 08:18:50 PDT 2023
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/67694
Modernize affine dialect ops: Define LB, UB, step and inits as operands in TableGen.
Note: There are still some operands that have not been updated to TableGen yet.
>From 6411564107407f7c52a5eb7d4530f27ad01d888b Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Thu, 28 Sep 2023 17:16:11 +0200
Subject: [PATCH] [mlir][Affine][NFC] Define AffineForOp operands in ODS
Define LB, UB, step and inits as operands in TableGen.
Note: There are still some operands that have not been updated to TableGen yet.
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 31 +++-----
.../AffineToStandard/AffineToStandard.cpp | 3 +-
mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp | 2 +-
.../Affine/Analysis/AffineStructures.cpp | 2 +-
.../Dialect/Affine/Analysis/LoopAnalysis.cpp | 2 +-
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 79 +++++++++----------
.../Transforms/PipelineDataTransfer.cpp | 2 +-
.../Affine/Transforms/SuperVectorize.cpp | 6 +-
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 43 +++++-----
mlir/lib/Dialect/Affine/Utils/Utils.cpp | 4 +-
mlir/test/Dialect/Affine/ops.mlir | 4 +-
11 files changed, 83 insertions(+), 95 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index d8ef0506d0822d7..43ad0a82ff58f07 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -117,7 +117,8 @@ def AffineApplyOp : Affine_Op<"apply", [Pure]> {
}
def AffineForOp : Affine_Op<"for",
- [AutomaticAllocationScope, ImplicitAffineTerminator, ConditionallySpeculatable,
+ [AttrSizedOperandSegments, AutomaticAllocationScope,
+ ImplicitAffineTerminator, ConditionallySpeculatable,
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
"getSingleUpperBound", "replaceWithAdditionalYields"]>,
@@ -147,7 +148,7 @@ def AffineForOp : Affine_Op<"for",
from a lower bound to an upper bound by a stride. The stride, represented by
`step`, is a positive constant integer which defaults to "1" if not present.
The lower and upper bounds specify a half-open range: the range includes the
- lower bound but does not include the upper bound.
+ lower bound but does not include the upper bound.
The lower and upper bounds of a `affine.for` operation are represented as an
application of an affine mapping to a list of SSA values passed to the map.
@@ -229,7 +230,10 @@ def AffineForOp : Affine_Op<"for",
explicitly present. The number and types of the "affine.for" results must
match the initial values in the `iter_args` binding and the yield operands.
}];
- let arguments = (ins Variadic<AnyType>);
+ let arguments = (ins Variadic<Index>:$lowerBoundOperands,
+ Variadic<Index>:$upperBoundOperands,
+ Variadic<AnyType>:$inits,
+ IndexAttr:$step);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -259,19 +263,10 @@ def AffineForOp : Affine_Op<"for",
Block::BlockArgListType getRegionIterArgs() {
return getBody()->getArguments().drop_front();
}
- Operation::operand_range getInits() {
- return getOperands().drop_front(getNumControlOperands());
- }
// TODO: provide iterators for the lower and upper bound operands
// if the current access via getLowerBound(), getUpperBound() is too slow.
- /// Returns operands for the lower bound map.
- operand_range getLowerBoundOperands();
-
- /// Returns operands for the upper bound map.
- operand_range getUpperBoundOperands();
-
/// Returns operands for the lower and upper bound maps with the operands
/// for the lower bound map in front of those for the upper bound map.
operand_range getControlOperands();
@@ -283,19 +278,17 @@ def AffineForOp : Affine_Op<"for",
AffineBound getUpperBound();
/// Returns loop step.
- int64_t getStep() {
- return ::llvm::cast<IntegerAttr>(*(*this)->getInherentAttr(getStepAttrStrName())).getInt();
- }
+ int64_t getStepAsInt() { return getStep().getSExtValue(); }
/// Returns affine map for the lower bound.
AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); }
AffineMapAttr getLowerBoundMapAttr() {
- return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(getLowerBoundAttrStrName()));
+ return ::llvm::cast<AffineMapAttr>(getOperation()->getAttr(getLowerBoundAttrStrName()));
}
/// Returns affine map for the upper bound. The upper bound is exclusive.
AffineMap getUpperBoundMap() { return getUpperBoundMapAttr().getValue(); }
AffineMapAttr getUpperBoundMapAttr() {
- return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(getUpperBoundAttrStrName()));
+ return ::llvm::cast<AffineMapAttr>(getOperation()->getAttr(getUpperBoundAttrStrName()));
}
/// Set lower bound. The new bound must have the same number of operands as
@@ -314,9 +307,7 @@ def AffineForOp : Affine_Op<"for",
/// Set loop step.
void setStep(int64_t step) {
assert(step > 0 && "step has to be a positive integer constant");
- auto *context = getLowerBoundMap().getContext();
- (*this)->setAttr(StringAttr::get(context, getStepAttrStrName()),
- IntegerAttr::get(IndexType::get(context), step));
+ setStep(APInt(/*numBits=*/64, step, /*isSigned=*/true));
}
/// Returns number of region arguments for loop-carried values.
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 456fd487352fb5d..7dbbf015182f39f 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -156,7 +156,8 @@ class AffineForLowering : public OpRewritePattern<AffineForOp> {
Location loc = op.getLoc();
Value lowerBound = lowerAffineLowerBound(op, rewriter);
Value upperBound = lowerAffineUpperBound(op, rewriter);
- Value step = rewriter.create<arith::ConstantIndexOp>(loc, op.getStep());
+ Value step =
+ rewriter.create<arith::ConstantIndexOp>(loc, op.getStepAsInt());
auto scfForOp = rewriter.create<scf::ForOp>(loc, lowerBound, upperBound,
step, op.getInits());
rewriter.eraseBlock(scfForOp.getBody());
diff --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index aff620ad0e70608..94d875d678df293 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
@@ -86,7 +86,7 @@ static Operation::operand_range getUpperBoundOperands(AffineForOp forOp) {
// materialize a corresponding constant using builder.
static Value getOrCreateStep(AffineForOp forOp, OpBuilder &builder) {
return builder.create<arith::ConstantIndexOp>(forOp.getLoc(),
- forOp.getStep());
+ forOp.getStepAsInt());
}
// Get a Value for the loop lower bound. If the value requires computation,
diff --git a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
index 5f32505690263fc..6ed3ba14fe15229 100644
--- a/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
@@ -74,7 +74,7 @@ FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
return failure();
}
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
if (step != 1) {
if (!forOp.hasConstantLowerBound())
LLVM_DEBUG(forOp.emitWarning("domain conservatively approximated"));
diff --git a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
index c97e99c0a0c19da..d56db64eac08261 100644
--- a/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
+++ b/mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
@@ -39,7 +39,7 @@ void mlir::affine::getTripCountMapAndOperands(
AffineForOp forOp, AffineMap *tripCountMap,
SmallVectorImpl<Value> *tripCountOperands) {
MLIRContext *context = forOp.getContext();
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
int64_t loopSpan;
if (forOp.hasConstantBounds()) {
int64_t lb = forOp.getConstantLowerBound();
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 6c060c90e24af82..c2dd84e6e67ca7d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -613,11 +613,11 @@ static int64_t getLargestKnownDivisor(AffineExpr e, ArrayRef<Value> operands) {
// LoopLikeOpInterface.
if (AffineForOp forOp = getForInductionVarOwner(operand)) {
if (forOp.hasConstantLowerBound() && forOp.getConstantLowerBound() == 0) {
- operandDivisor = forOp.getStep();
+ operandDivisor = forOp.getStepAsInt();
} else {
uint64_t lbLargestKnownDivisor =
forOp.getLowerBoundMap().getLargestKnownDivisorOfMapExprs();
- operandDivisor = std::gcd(lbLargestKnownDivisor, forOp.getStep());
+ operandDivisor = std::gcd(lbLargestKnownDivisor, forOp.getStepAsInt());
}
}
return operandDivisor;
@@ -695,7 +695,7 @@ static std::optional<int64_t> getUpperBound(Value iv) {
if (forOp.hasConstantLowerBound()) {
return forOp.getConstantUpperBound() - 1 -
(forOp.getConstantUpperBound() - forOp.getConstantLowerBound() - 1) %
- forOp.getStep();
+ forOp.getStepAsInt();
}
return forOp.getConstantUpperBound() - 1;
}
@@ -1897,6 +1897,13 @@ void AffineForOp::build(OpBuilder &builder, OperationState &result,
"upper bound operand count does not match the affine map");
assert(step > 0 && "step has to be a positive integer constant");
+ // Set variadic segment sizes.
+ result.addAttribute(
+ getOperandSegmentSizeAttr(),
+ builder.getDenseI32ArrayAttr({static_cast<int32_t>(lbOperands.size()),
+ static_cast<int32_t>(ubOperands.size()),
+ static_cast<int32_t>(iterArgs.size())}));
+
for (Value val : iterArgs)
result.addTypes(val.getType());
@@ -2083,10 +2090,16 @@ ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
return failure();
// Parse loop bounds.
- if (parseBound(/*isLower=*/true, result, parser) ||
- parser.parseKeyword("to", " between bounds") ||
- parseBound(/*isLower=*/false, result, parser))
+ int64_t numOperands = result.operands.size();
+ if (parseBound(/*isLower=*/true, result, parser))
+ return failure();
+ int64_t numLbOperands = result.operands.size() - numOperands;
+ if (parser.parseKeyword("to", " between bounds"))
return failure();
+ numOperands = result.operands.size();
+ if (parseBound(/*isLower=*/false, result, parser))
+ return failure();
+ int64_t numUbOperands = result.operands.size() - numOperands;
// Parse the optional loop step, we default to 1 if one is not present.
if (parser.parseOptionalKeyword("step")) {
@@ -2130,6 +2143,12 @@ ParseResult AffineForOp::parse(OpAsmParser &parser, OperationState &result) {
}
}
+ result.addAttribute(
+ getOperandSegmentSizeAttr(),
+ builder.getDenseI32ArrayAttr({static_cast<int32_t>(numLbOperands),
+ static_cast<int32_t>(numUbOperands),
+ static_cast<int32_t>(operands.size())}));
+
// Parse the body region.
Region *body = result.addRegion();
if (regionArgs.size() != result.types.size() + 1)
@@ -2202,8 +2221,8 @@ void AffineForOp::print(OpAsmPrinter &p) {
p << " to ";
printBound(getUpperBoundMapAttr(), getUpperBoundOperands(), "min", p);
- if (getStep() != 1)
- p << " step " << getStep();
+ if (getStepAsInt() != 1)
+ p << " step " << getStepAsInt();
bool printBlockTerminators = false;
if (getNumIterOperands() > 0) {
@@ -2221,10 +2240,10 @@ void AffineForOp::print(OpAsmPrinter &p) {
p << ' ';
p.printRegion(getRegion(), /*printEntryBlockArgs=*/false,
printBlockTerminators);
- p.printOptionalAttrDict((*this)->getAttrs(),
- /*elidedAttrs=*/{getLowerBoundAttrStrName(),
- getUpperBoundAttrStrName(),
- getStepAttrStrName()});
+ p.printOptionalAttrDict(
+ (*this)->getAttrs(),
+ /*elidedAttrs=*/{getLowerBoundAttrStrName(), getUpperBoundAttrStrName(),
+ getStepAttrStrName(), getOperandSegmentSizeAttr()});
}
/// Fold the constant bounds of a loop.
@@ -2307,7 +2326,7 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
namespace {
/// Returns constant trip count in trivial cases.
static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
if (!forOp.hasConstantBounds() || step <= 0)
return std::nullopt;
int64_t lb = forOp.getConstantLowerBound();
@@ -2452,28 +2471,14 @@ AffineBound AffineForOp::getUpperBound() {
void AffineForOp::setLowerBound(ValueRange lbOperands, AffineMap map) {
assert(lbOperands.size() == map.getNumInputs());
assert(map.getNumResults() >= 1 && "bound map has at least one result");
-
- SmallVector<Value, 4> newOperands(lbOperands.begin(), lbOperands.end());
-
- auto ubOperands = getUpperBoundOperands();
- newOperands.append(ubOperands.begin(), ubOperands.end());
- auto iterOperands = getInits();
- newOperands.append(iterOperands.begin(), iterOperands.end());
- (*this)->setOperands(newOperands);
-
+ getLowerBoundOperandsMutable().assign(lbOperands);
(*this)->setAttr(getLowerBoundAttrStrName(), AffineMapAttr::get(map));
}
void AffineForOp::setUpperBound(ValueRange ubOperands, AffineMap map) {
assert(ubOperands.size() == map.getNumInputs());
assert(map.getNumResults() >= 1 && "bound map has at least one result");
-
- SmallVector<Value, 4> newOperands(getLowerBoundOperands());
- newOperands.append(ubOperands.begin(), ubOperands.end());
- auto iterOperands = getInits();
- newOperands.append(iterOperands.begin(), iterOperands.end());
- (*this)->setOperands(newOperands);
-
+ getUpperBoundOperandsMutable().assign(ubOperands);
(*this)->setAttr(getUpperBoundAttrStrName(), AffineMapAttr::get(map));
}
@@ -2519,16 +2524,6 @@ void AffineForOp::setConstantUpperBound(int64_t value) {
setUpperBound({}, AffineMap::getConstantMap(value, getContext()));
}
-AffineForOp::operand_range AffineForOp::getLowerBoundOperands() {
- return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs()};
-}
-
-AffineForOp::operand_range AffineForOp::getUpperBoundOperands() {
- return {operand_begin() + getLowerBoundMap().getNumInputs(),
- operand_begin() + getLowerBoundMap().getNumInputs() +
- getUpperBoundMap().getNumInputs()};
-}
-
AffineForOp::operand_range AffineForOp::getControlOperands() {
return {operand_begin(), operand_begin() + getLowerBoundMap().getNumInputs() +
getUpperBoundMap().getNumInputs()};
@@ -2565,7 +2560,7 @@ std::optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
std::optional<OpFoldResult> AffineForOp::getSingleStep() {
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getStep()));
+ return OpFoldResult(b.getI64IntegerAttr(getStepAsInt()));
}
std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
@@ -2633,8 +2628,8 @@ Speculation::Speculatability AffineForOp::getSpeculatability() {
//
// For Step != 1, the loop may not terminate. We can add more smarts here if
// needed.
- return getStep() == 1 ? Speculation::RecursivelySpeculatable
- : Speculation::NotSpeculatable;
+ return getStepAsInt() == 1 ? Speculation::RecursivelySpeculatable
+ : Speculation::NotSpeculatable;
}
/// Returns true if the provided value is the induction variable of a
diff --git a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
index 7029251a3720cb7..d169aca2c971f24 100644
--- a/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
@@ -107,7 +107,7 @@ static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
// Create 'iv mod 2' value to index the leading dimension.
auto d0 = bInner.getAffineDimExpr(0);
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
auto modTwoMap =
AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, d0.floorDiv(step) % 2);
auto ivModTwoOp = bInner.create<AffineApplyOp>(forOp.getLoc(), modTwoMap,
diff --git a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
index 85c2602aa266d67..f07ed626a916384 100644
--- a/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
@@ -1021,7 +1021,7 @@ static Value createMask(AffineForOp vecForOp, VectorizationState &state) {
if (vecForOp.hasConstantBounds()) {
int64_t originalTripCount =
vecForOp.getConstantUpperBound() - vecForOp.getConstantLowerBound();
- if (originalTripCount % vecForOp.getStep() == 0)
+ if (originalTripCount % vecForOp.getStepAsInt() == 0)
return nullptr;
}
@@ -1296,9 +1296,9 @@ static Operation *vectorizeAffineForOp(AffineForOp forOp,
unsigned vectorDim = loopToVecDimIt->second;
assert(vectorDim < strategy.vectorSizes.size() && "vector dim overflow");
int64_t forOpVecFactor = strategy.vectorSizes[vectorDim];
- newStep = forOp.getStep() * forOpVecFactor;
+ newStep = forOp.getStepAsInt() * forOpVecFactor;
} else {
- newStep = forOp.getStep();
+ newStep = forOp.getStepAsInt();
}
// Get information about reduction kinds.
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 9d8ed9b4ac93387..fb8a0a7c330cf22 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -78,7 +78,7 @@ getCleanupLoopLowerBound(AffineForOp forOp, unsigned unrollFactor,
// these affine.apply's make up the cleanup loop lower bound.
SmallVector<AffineExpr, 4> bumpExprs(tripCountMap.getNumResults());
SmallVector<Value, 4> bumpValues(tripCountMap.getNumResults());
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
for (unsigned i = 0, e = tripCountMap.getNumResults(); i < e; i++) {
auto tripCountExpr = tripCountMap.getResult(i);
bumpExprs[i] = (tripCountExpr - tripCountExpr % unrollFactor) * step;
@@ -189,8 +189,9 @@ static AffineForOp generateShiftedLoop(
assert(lbMap.getNumInputs() == lbOperands.size());
assert(ubMap.getNumInputs() == ubOperands.size());
- auto loopChunk = b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap,
- ubOperands, ubMap, srcForOp.getStep());
+ auto loopChunk =
+ b.create<AffineForOp>(srcForOp.getLoc(), lbOperands, lbMap, ubOperands,
+ ubMap, srcForOp.getStepAsInt());
auto loopChunkIV = loopChunk.getInductionVar();
auto srcIV = srcForOp.getInductionVar();
@@ -208,7 +209,7 @@ static AffineForOp generateShiftedLoop(
auto ivRemap = bodyBuilder.create<AffineApplyOp>(
srcForOp.getLoc(),
bodyBuilder.getSingleDimShiftAffineMap(
- -static_cast<int64_t>(srcForOp.getStep() * shift)),
+ -static_cast<int64_t>(srcForOp.getStepAsInt() * shift)),
loopChunkIV);
operandMap.map(srcIV, ivRemap);
} else {
@@ -255,7 +256,7 @@ LogicalResult mlir::affine::affineForOpBodySkew(AffineForOp forOp,
assert(isOpwiseShiftValid(forOp, shifts) &&
"shifts will lead to an invalid transformation\n");
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
unsigned numChildOps = shifts.size();
@@ -535,7 +536,7 @@ static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
origLowerBoundExpr);
// Get the origLoopStep as an affine expression.
- AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStep());
+ AffineExpr origLoopStep = b.getAffineConstantExpr(origLoop.getStepAsInt());
// Insert ub as inter-tile ((loop IV - origlb) * tilingParameter) +
// (tilingParameter * origLoopStep) + origlb.
@@ -557,7 +558,7 @@ static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
newIntraTileLoop.setUpperBound(ubOperands, ubMap);
// Original loop step must be preserved.
- newIntraTileLoop.setStep(origLoop.getStep());
+ newIntraTileLoop.setStep(origLoop.getStepAsInt());
}
/// Set lower and upper bounds of inter-tile loops for parametric tiling.
@@ -654,7 +655,7 @@ static void setInterTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
newLoop.setUpperBound(ubOperands, ubMap);
// Original loop step must be preserved.
- newLoop.setStep(origLoop.getStep());
+ newLoop.setStep(origLoop.getStepAsInt());
}
/// Constructs and sets new loop bounds after tiling for the case of
@@ -708,7 +709,7 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap());
// If the step size of original loop is x and tileSize is y then after
// tiling the tile space loops' step size becomes x*y.
- newLoops[i].setStep(tileSizes[i] * origLoops[i].getStep());
+ newLoops[i].setStep(tileSizes[i] * origLoops[i].getStepAsInt());
}
// Bounds for intra-tile loops.
for (unsigned i = 0; i < width; i++) {
@@ -720,14 +721,14 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
newLoops[width + i].setLowerBound(
/*operands=*/newLoops[i].getInductionVar(), lbMap);
// The step sizes of intra-tile loops is just the original loops' step size.
- newLoops[width + i].setStep(origLoops[i].getStep());
+ newLoops[width + i].setStep(origLoops[i].getStepAsInt());
// Set the upper bound.
if (mayBeConstantCount && *mayBeConstantCount < tileSizes[i]) {
// Trip count is less than the tile size: upper bound is lower bound +
// trip count * stepSize.
- AffineMap ubMap = b.getSingleDimShiftAffineMap(*mayBeConstantCount *
- origLoops[i].getStep());
+ AffineMap ubMap = b.getSingleDimShiftAffineMap(
+ *mayBeConstantCount * origLoops[i].getStepAsInt());
newLoops[width + i].setUpperBound(
/*operands=*/newLoops[i].getInductionVar(), ubMap);
} else if (largestDiv % tileSizes[i] != 0) {
@@ -757,7 +758,7 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
AffineExpr dim = b.getAffineDimExpr(origUbMap.getNumDims());
// The new upper bound map is the original one with an additional
// expression i + tileSize * stepSize (of original loop) appended.
- boundExprs.push_back(dim + tileSizes[i] * origLoops[i].getStep());
+ boundExprs.push_back(dim + tileSizes[i] * origLoops[i].getStepAsInt());
boundExprs.append(origUbMap.getResults().begin(),
origUbMap.getResults().end());
AffineMap ubMap =
@@ -767,8 +768,8 @@ constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops,
} else {
// No need of the min expression.
AffineExpr dim = b.getAffineDimExpr(0);
- AffineMap ubMap =
- AffineMap::get(1, 0, dim + tileSizes[i] * origLoops[i].getStep());
+ AffineMap ubMap = AffineMap::get(
+ 1, 0, dim + tileSizes[i] * origLoops[i].getStepAsInt());
newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap);
}
}
@@ -1061,7 +1062,7 @@ LogicalResult mlir::affine::loopUnrollByFactor(
auto yieldedValues = forOp.getBody()->getTerminator()->getOperands();
// Scale the step of loop being unrolled by unroll factor.
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
forOp.setStep(step * unrollFactor);
generateUnrolledLoop(
forOp.getBody(), forOp.getInductionVar(), unrollFactor,
@@ -1244,7 +1245,7 @@ LogicalResult mlir::affine::loopUnrollJamByFactor(AffineForOp forOp,
}
// Scale the step of loop being unroll-jammed by the unroll-jam factor.
- int64_t step = forOp.getStep();
+ int64_t step = forOp.getStepAsInt();
forOp.setStep(step * unrollJamFactor);
auto forOpIV = forOp.getInductionVar();
@@ -1571,7 +1572,7 @@ static void augmentMapAndBounds(OpBuilder &b, Value iv, AffineMap *map,
static SmallVector<AffineForOp, 8>
stripmineSink(AffineForOp forOp, uint64_t factor,
ArrayRef<AffineForOp> targets) {
- auto originalStep = forOp.getStep();
+ auto originalStep = forOp.getStepAsInt();
auto scaledStep = originalStep * factor;
forOp.setStep(scaledStep);
@@ -1659,7 +1660,7 @@ LogicalResult mlir::affine::coalesceLoops(MutableArrayRef<AffineForOp> loops) {
OpBuilder builder(outermost);
for (AffineForOp loop : loops) {
// We only work on normalized loops.
- if (loop.getStep() != 1 || !loop.hasConstantLowerBound() ||
+ if (loop.getStepAsInt() != 1 || !loop.hasConstantLowerBound() ||
loop.getConstantLowerBound() != 0)
return failure();
}
@@ -2606,7 +2607,7 @@ static AffineIfOp createSeparationCondition(MutableArrayRef<AffineForOp> loops,
for (auto loop : loops) {
(void)loop;
// TODO: Non-unit stride is not an issue to generalize to.
- assert(loop.getStep() == 1 && "point loop step expected to be one");
+ assert(loop.getStepAsInt() == 1 && "point loop step expected to be one");
// Mark everything symbols for the purpose of finding a constant diff pair.
cst.setDimSymbolSeparation(/*newSymbolCount=*/cst.getNumDimAndSymbolVars() -
1);
@@ -2673,7 +2674,7 @@ createFullTiles(MutableArrayRef<AffineForOp> inputNest,
FlatAffineValueConstraints cst;
for (auto loop : inputNest) {
// TODO: straightforward to generalize to a non-unit stride.
- if (loop.getStep() != 1) {
+ if (loop.getStepAsInt() != 1) {
LLVM_DEBUG(llvm::dbgs()
<< "[tile separation] non-unit stride not implemented\n");
return failure();
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index cef5080daad29bc..346320cf5bbaa82 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -367,7 +367,7 @@ mlir::affine::affineParallelize(AffineForOp forOp,
loc, ValueRange(reducedValues).getTypes(), reductionKinds,
llvm::ArrayRef(lowerBoundMap), lowerBoundOperands,
llvm::ArrayRef(upperBoundMap), upperBoundOperands,
- llvm::ArrayRef(forOp.getStep()));
+ llvm::ArrayRef(forOp.getStepAsInt()));
// Steal the body of the old affine for op.
newPloop.getRegion().takeBody(forOp.getRegion());
Operation *yieldOp = &newPloop.getBody()->back();
@@ -570,7 +570,7 @@ LogicalResult mlir::affine::normalizeAffineFor(AffineForOp op,
Location loc = op.getLoc();
OpBuilder opBuilder(op);
- int64_t origLoopStep = op.getStep();
+ int64_t origLoopStep = op.getStepAsInt();
// Construct the new upper bound value map.
AffineMap oldLbMap = op.getLowerBoundMap();
diff --git a/mlir/test/Dialect/Affine/ops.mlir b/mlir/test/Dialect/Affine/ops.mlir
index 1063f2a7ecba489..1fb1503c3b0c4b4 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -49,10 +49,10 @@ func.func @affine.yield() {
// CHECK: affine.for
// CHECK-NEXT: }
//
- // GENERIC: "affine.for"() ({
+ // GENERIC: "affine.for"() <{operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index}> ({
// GENERIC-NEXT: ^bb0(%{{.*}}: index):
// GENERIC-NEXT: "affine.yield"() : () -> ()
- // GENERIC-NEXT: }) {lower_bound = #map, step = 1 : index, upper_bound = #map1} : () -> ()
+ // GENERIC-NEXT: }) {lower_bound = #map, upper_bound = #map1} : () -> ()
affine.for %i = 0 to 10 {
"affine.yield"() : () -> ()
}
More information about the Mlir-commits
mailing list