[Mlir-commits] [mlir] 23b794f - [mlir][Affine][NFC] Define AffineForOp operands in ODS (#67694)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Sep 29 01:47:31 PDT 2023
Author: Matthias Springer
Date: 2023-09-29T10:47:28+02:00
New Revision: 23b794f720066d41497d9836b23e80a5a0876cbd
URL: https://github.com/llvm/llvm-project/commit/23b794f720066d41497d9836b23e80a5a0876cbd
DIFF: https://github.com/llvm/llvm-project/commit/23b794f720066d41497d9836b23e80a5a0876cbd.diff
LOG: [mlir][Affine][NFC] Define AffineForOp operands in ODS (#67694)
Modernize affine dialect ops: Define LB, UB, step and inits as operands
in TableGen.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp
mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/Affine/Transforms/PipelineDataTransfer.cpp
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/test/Dialect/Affine/ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
index 56b4a609e62c001..f070d0488619063 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.h
@@ -499,27 +499,28 @@ class AffineBound {
AffineForOp getAffineForOp() { return op; }
AffineMap getMap() { return map; }
- unsigned getNumOperands() { return opEnd - opStart; }
- Value getOperand(unsigned idx) { return op.getOperand(opStart + idx); }
+ unsigned getNumOperands() { return operands.size(); }
+ Value getOperand(unsigned idx) {
+ return op.getOperand(operands.getBeginOperandIndex() + idx);
+ }
using operand_iterator = AffineForOp::operand_iterator;
using operand_range = AffineForOp::operand_range;
- operand_iterator operandBegin() { return op.operand_begin() + opStart; }
- operand_iterator operandEnd() { return op.operand_begin() + opEnd; }
+ operand_iterator operandBegin() { return operands.begin(); }
+ operand_iterator operandEnd() { return operands.end(); }
operand_range getOperands() { return {operandBegin(), operandEnd()}; }
private:
// 'affine.for' operation that contains this bound.
AffineForOp op;
- // Start and end positions of this affine bound operands in the list of
- // the containing 'affine.for' operation operands.
- unsigned opStart, opEnd;
+ // Operands of the affine map.
+ OperandRange operands;
// Affine map for this bound.
AffineMap map;
- AffineBound(AffineForOp op, unsigned opStart, unsigned opEnd, AffineMap map)
- : op(op), opStart(opStart), opEnd(opEnd), map(map) {}
+ AffineBound(AffineForOp op, OperandRange operands, AffineMap map)
+ : op(op), operands(operands), map(map) {}
friend class AffineForOp;
};
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index d8ef0506d0822d7..f90a7b23ec12e18 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,12 @@ 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,
+ AffineMapAttr:$lowerBoundMap,
+ AffineMapAttr:$upperBoundMap,
+ IndexAttr:$step);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -251,26 +257,10 @@ def AffineForOp : Affine_Op<"for",
using BodyBuilderFn =
function_ref<void(OpBuilder &, Location, Value, ValueRange)>;
- static StringRef getStepAttrStrName() { return "step"; }
- static StringRef getLowerBoundAttrStrName() { return "lower_bound"; }
- static StringRef getUpperBoundAttrStrName() { return "upper_bound"; }
-
BlockArgument getInductionVar() { return getBody()->getArgument(0); }
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.
@@ -283,20 +273,7 @@ def AffineForOp : Affine_Op<"for",
AffineBound getUpperBound();
/// Returns loop step.
- int64_t getStep() {
- return ::llvm::cast<IntegerAttr>(*(*this)->getInherentAttr(getStepAttrStrName())).getInt();
- }
-
- /// Returns affine map for the lower bound.
- AffineMap getLowerBoundMap() { return getLowerBoundMapAttr().getValue(); }
- AffineMapAttr getLowerBoundMapAttr() {
- return ::llvm::cast<AffineMapAttr>(*(*this)->getInherentAttr(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()));
- }
+ int64_t getStepAsInt() { return getStep().getSExtValue(); }
/// Set lower bound. The new bound must have the same number of operands as
/// the current bound map. Otherwise, 'replaceForLowerBound' should be used.
@@ -305,18 +282,10 @@ def AffineForOp : Affine_Op<"for",
/// current bound map. Otherwise, 'replaceForUpperBound' should be used.
void setUpperBound(ValueRange operands, AffineMap map);
- /// Set the lower bound map without changing operands.
- void setLowerBoundMap(AffineMap map);
-
- /// Set the upper bound map without changing operands.
- void setUpperBoundMap(AffineMap map);
-
/// 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.
@@ -325,7 +294,9 @@ def AffineForOp : Affine_Op<"for",
}
/// Number of operands controlling the loop: lb and ub.
- unsigned getNumControlOperands() { return getOperation()->getNumOperands() - getNumIterOperands(); }
+ unsigned getNumControlOperands() {
+ return getOperation()->getNumOperands() - getNumIterOperands();
+ }
/// Get the number of loop-carried values.
unsigned getNumIterOperands();
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..c61bd566c7676f1 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,19 +1897,28 @@ 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());
// Add an attribute for the step.
- result.addAttribute(getStepAttrStrName(),
+ result.addAttribute(getStepAttrName(result.name),
builder.getIntegerAttr(builder.getIndexType(), step));
// Add the lower bound.
- result.addAttribute(getLowerBoundAttrStrName(), AffineMapAttr::get(lbMap));
+ result.addAttribute(getLowerBoundMapAttrName(result.name),
+ AffineMapAttr::get(lbMap));
result.addOperands(lbOperands);
// Add the upper bound.
- result.addAttribute(getUpperBoundAttrStrName(), AffineMapAttr::get(ubMap));
+ result.addAttribute(getUpperBoundMapAttrName(result.name),
+ AffineMapAttr::get(ubMap));
result.addOperands(ubOperands);
result.addOperands(iterArgs);
@@ -1991,8 +2000,9 @@ static ParseResult parseBound(bool isLower, OperationState &result,
failed(p.parseOptionalKeyword(isLower ? "max" : "min"));
auto &builder = p.getBuilder();
- auto boundAttrStrName = isLower ? AffineForOp::getLowerBoundAttrStrName()
- : AffineForOp::getUpperBoundAttrStrName();
+ auto boundAttrStrName =
+ isLower ? AffineForOp::getLowerBoundMapAttrName(result.name)
+ : AffineForOp::getUpperBoundMapAttrName(result.name);
// Parse ssa-id as identity map.
SmallVector<OpAsmParser::UnresolvedOperand, 1> boundOpInfos;
@@ -2083,21 +2093,27 @@ 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")) {
result.addAttribute(
- AffineForOp::getStepAttrStrName(),
+ getStepAttrName(result.name),
builder.getIntegerAttr(builder.getIndexType(), /*value=*/1));
} else {
SMLoc stepLoc = parser.getCurrentLocation();
IntegerAttr stepAttr;
if (parser.parseAttribute(stepAttr, builder.getIndexType(),
- AffineForOp::getStepAttrStrName().data(),
+ getStepAttrName(result.name).data(),
result.attributes))
return failure();
@@ -2130,6 +2146,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 +2224,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 +2243,12 @@ 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=*/{getLowerBoundMapAttrName(getOperation()->getName()),
+ getUpperBoundMapAttrName(getOperation()->getName()),
+ getStepAttrName(getOperation()->getName()),
+ getOperandSegmentSizeAttr()});
}
/// Fold the constant bounds of a loop.
@@ -2307,7 +2331,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();
@@ -2438,61 +2462,25 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
}
AffineBound AffineForOp::getLowerBound() {
- auto lbMap = getLowerBoundMap();
- return AffineBound(AffineForOp(*this), 0, lbMap.getNumInputs(), lbMap);
+ return AffineBound(*this, getLowerBoundOperands(), getLowerBoundMap());
}
AffineBound AffineForOp::getUpperBound() {
- auto lbMap = getLowerBoundMap();
- auto ubMap = getUpperBoundMap();
- return AffineBound(AffineForOp(*this), lbMap.getNumInputs(),
- lbMap.getNumInputs() + ubMap.getNumInputs(), ubMap);
+ return AffineBound(*this, getUpperBoundOperands(), getUpperBoundMap());
}
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);
-
- (*this)->setAttr(getLowerBoundAttrStrName(), AffineMapAttr::get(map));
+ getLowerBoundOperandsMutable().assign(lbOperands);
+ setLowerBoundMap(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);
-
- (*this)->setAttr(getUpperBoundAttrStrName(), AffineMapAttr::get(map));
-}
-
-void AffineForOp::setLowerBoundMap(AffineMap map) {
- auto lbMap = getLowerBoundMap();
- assert(lbMap.getNumDims() == map.getNumDims() &&
- lbMap.getNumSymbols() == map.getNumSymbols());
- assert(map.getNumResults() >= 1 && "bound map has at least one result");
- (void)lbMap;
- (*this)->setAttr(getLowerBoundAttrStrName(), AffineMapAttr::get(map));
-}
-
-void AffineForOp::setUpperBoundMap(AffineMap map) {
- auto ubMap = getUpperBoundMap();
- assert(ubMap.getNumDims() == map.getNumDims() &&
- ubMap.getNumSymbols() == map.getNumSymbols());
- assert(map.getNumResults() >= 1 && "bound map has at least one result");
- (void)ubMap;
- (*this)->setAttr(getUpperBoundAttrStrName(), AffineMapAttr::get(map));
+ getUpperBoundOperandsMutable().assign(ubOperands);
+ setUpperBoundMap(map);
}
bool AffineForOp::hasConstantLowerBound() {
@@ -2519,19 +2507,9 @@ 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()};
+ return {operand_begin(), operand_begin() + getLowerBoundOperands().size() +
+ getUpperBoundOperands().size()};
}
bool AffineForOp::matchingBoundOperandList() {
@@ -2565,7 +2543,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() {
@@ -2586,7 +2564,7 @@ FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
inits.append(newInitOperands.begin(), newInitOperands.end());
AffineForOp newLoop = rewriter.create<AffineForOp>(
getLoc(), getLowerBoundOperands(), getLowerBoundMap(),
- getUpperBoundOperands(), getUpperBoundMap(), getStep(), inits);
+ getUpperBoundOperands(), getUpperBoundMap(), getStepAsInt(), inits);
// Generate the new yield values and append them to the scf.yield operation.
auto yieldOp = cast<AffineYieldOp>(getBody()->getTerminator());
@@ -2633,8 +2611,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..19ae1584842aea8 100644
--- a/mlir/test/Dialect/Affine/ops.mlir
+++ b/mlir/test/Dialect/Affine/ops.mlir
@@ -1,5 +1,5 @@
// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s
+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file %s -mlir-print-op-generic | FileCheck -check-prefix=GENERIC %s
// Check that the attributes for the affine operations are round-tripped.
// Check that `affine.yield` is visible in the generic form.
@@ -42,17 +42,23 @@ func.func @empty() {
return
}
+// -----
+
+// GENERIC: #[[$map:.*]] = affine_map<() -> (0)>
+// GENERIC: #[[$map1:.*]] = affine_map<() -> (10)>
+
// Check that an explicit affine.yield is not printed in custom format.
// Check that no extra terminator is introduced.
// CHECK-LABEL: @affine.yield
+// CHECK-GENERIC-LABEL: @affine.yield
func.func @affine.yield() {
// CHECK: affine.for
// CHECK-NEXT: }
//
- // GENERIC: "affine.for"() ({
+ // GENERIC: "affine.for"() <{lowerBoundMap = #[[$map]], operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index, upperBoundMap = #[[$map1]]}> ({
// GENERIC-NEXT: ^bb0(%{{.*}}: index):
// GENERIC-NEXT: "affine.yield"() : () -> ()
- // GENERIC-NEXT: }) {lower_bound = #map, step = 1 : index, upper_bound = #map1} : () -> ()
+ // GENERIC-NEXT: }) : () -> ()
affine.for %i = 0 to 10 {
"affine.yield"() : () -> ()
}
More information about the Mlir-commits
mailing list