-
Notifications
You must be signed in to change notification settings - Fork 15.4k
[mlir][SCF] scf.for: Add support for unsigned integer comparison
#153379
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[mlir][SCF] scf.for: Add support for unsigned integer comparison
#153379
Conversation
|
@llvm/pr-subscribers-mlir-spirv @llvm/pr-subscribers-mlir-scf Author: Matthias Springer (matthias-springer) ChangesAdd a new unit attribute to allow for unsigned integer comparison. Example: scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
// body
}Discussion: https://discourse.llvm.org/t/scf-should-scf-for-support-unsigned-comparison/84655 Patch is 27.08 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/153379.diff 19 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0c1c15b85f4c9..41be006225730 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -169,9 +169,13 @@ def ForOp : SCF_Op<"for",
region capturing the loop body. The induction variable is represented as an
argument of this region. This SSA value is a signless integer or index.
The step is a value of same type but required to be positive, the lower and
- upper bounds can be also negative or zero. The lower and upper bounds specify
- a half-open range: the iteration is executed iff the signed comparison of induction
- variable value is less than the upper bound and bigger or equal to the lower bound.
+ upper bounds can be also negative or zero. The lower and upper bounds
+ specify a half-open range: the iteration is executed iff the comparison of
+ induction variable value is less than the upper bound and bigger or equal
+ to the lower bound.
+
+ By default, the integer comparison is signed. If the `unsignedInt` unit
+ attribute is specified, the integer comparison is unsigned.
The body region must contain exactly one block that terminates with
`scf.yield`. Calling ForOp::build will create such a region and insert
@@ -184,8 +188,8 @@ def ForOp : SCF_Op<"for",
... // body
}
...
- // Integer case.
- scf.for %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
+ // Unsigned integer case.
+ scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
... // body
}
```
@@ -258,7 +262,8 @@ def ForOp : SCF_Op<"for",
let arguments = (ins AnySignlessIntegerOrIndex:$lowerBound,
AnySignlessIntegerOrIndex:$upperBound,
AnySignlessIntegerOrIndex:$step,
- Variadic<AnyType>:$initArgs);
+ Variadic<AnyType>:$initArgs,
+ UnitAttr:$unsignedInt);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -266,7 +271,8 @@ def ForOp : SCF_Op<"for",
let builders = [OpBuilder<(ins "Value":$lowerBound, "Value":$upperBound,
"Value":$step, CArg<"ValueRange", "{}">:$initArgs,
CArg<"function_ref<void(OpBuilder &, Location, Value, ValueRange)>",
- "nullptr">)>];
+ "nullptr">,
+ CArg<"bool", "false">:$unsignedInt)>];
let extraClassDeclaration = [{
using BodyBuilderFn =
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index ba448e46913ac..c4f9ff8ab4c22 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -382,8 +382,11 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
// With the body block done, we can fill in the condition block.
rewriter.setInsertionPointToEnd(conditionBlock);
- auto comparison = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::slt, iv, upperBound);
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto comparison =
+ arith::CmpIOp::create(rewriter, loc, predicate, iv, upperBound);
cf::CondBranchOp::create(rewriter, loc, comparison, firstBodyBlock,
ArrayRef<Value>(), endBlock, ArrayRef<Value>());
diff --git a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
index 84cbd869c78ef..fd0339e7943e0 100644
--- a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
+++ b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
@@ -154,6 +154,10 @@ ForLowering::matchAndRewrite(ForOp forOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
Location loc = forOp.getLoc();
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Create an emitc::variable op for each result. These variables will be
// assigned to by emitc::assign ops within the loop body.
SmallVector<Value> resultVariables;
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index dc92367fc58cd..03042d87b8151 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -178,8 +178,14 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
// Generate the rest of the loop header.
rewriter.setInsertionPointToEnd(header);
auto *mergeBlock = loopOp.getMergeBlock();
- auto cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
- newIndVar, adaptor.getUpperBound());
+ Value cmpOp;
+ if (forOp.getUnsignedInt()) {
+ cmpOp = spirv::ULessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ } else {
+ cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ }
spirv::BranchConditionalOp::create(rewriter, loc, cmpOp, body,
ArrayRef<Value>(), mergeBlock,
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index fd530f26db526..807012d5438aa 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -594,7 +594,8 @@ static FailureOr<PackingResult> buildPackingLoopNestImpl(
auto clonedForOp = scf::ForOp::create(
rewriter, loc, bvm.lookupOrDefault(forOp.getLowerBound()),
bvm.lookupOrDefault(forOp.getUpperBound()),
- bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor);
+ bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// Map the induction var, region args and results to the `clonedForOp`.
bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 58986a6009995..8c0faebedfc34 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -55,7 +55,8 @@ static scf::ForOp replaceWithDifferentYield(RewriterBase &rewriter,
scf::ForOp newLoop = scf::ForOp::create(
rewriter, loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(),
- loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loop.getUnsignedInt());
// Generate the new yield with the replaced operand.
auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 0262a1b8a3893..ecd85c9101951 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -318,9 +318,12 @@ void ConditionOp::getSuccessorRegions(
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange initArgs,
- BodyBuilderFn bodyBuilder) {
+ BodyBuilderFn bodyBuilder, bool unsignedInt) {
OpBuilder::InsertionGuard guard(builder);
+ if (unsignedInt)
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
result.addOperands({lb, ub, step});
result.addOperands(initArgs);
for (Value v : initArgs)
@@ -450,6 +453,9 @@ static void printInitializationList(OpAsmPrinter &p,
}
void ForOp::print(OpAsmPrinter &p) {
+ if (getUnsignedInt())
+ p << " unsigned";
+
p << " " << getInductionVar() << " = " << getLowerBound() << " to "
<< getUpperBound() << " step " << getStep();
@@ -462,7 +468,8 @@ void ForOp::print(OpAsmPrinter &p) {
p.printRegion(getRegion(),
/*printEntryBlockArgs=*/false,
/*printBlockTerminators=*/!getInitArgs().empty());
- p.printOptionalAttrDict((*this)->getAttrs());
+ p.printOptionalAttrDict((*this)->getAttrs(),
+ /*elidedAttrs=*/getUnsignedIntAttrName().strref());
}
ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
@@ -472,6 +479,10 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
OpAsmParser::Argument inductionVariable;
OpAsmParser::UnresolvedOperand lb, ub, step;
+ if (succeeded(parser.parseOptionalKeyword("unsigned")))
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
+
// Parse the induction variable followed by '='.
if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual() ||
// Parse loop bounds.
@@ -562,7 +573,7 @@ ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,
inits.append(newInitOperands.begin(), newInitOperands.end());
scf::ForOp newLoop = scf::ForOp::create(
rewriter, getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
- [](OpBuilder &, Location, Value, ValueRange) {});
+ [](OpBuilder &, Location, Value, ValueRange) {}, getUnsignedInt());
newLoop->setAttrs(getPrunedAttributeList(getOperation(), {}));
// Generate the new yield values and append them to the scf.yield operation.
@@ -806,7 +817,8 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,
// 2. Create the new forOp shell.
scf::ForOp newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), newIterOperands);
+ forOp.getStep(), newIterOperands, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
SmallVector<Value, 4> newBlockTransferArgs(newBlock.getArguments().begin(),
@@ -931,7 +943,8 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
scf::ForOp newForOp =
scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),
- forOp.getUpperBound(), forOp.getStep(), newIterArgs);
+ forOp.getUpperBound(), forOp.getStep(), newIterArgs,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
@@ -989,12 +1002,12 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
/// Util function that tries to compute a constant diff between u and l.
/// Returns std::nullopt when the difference between two AffineValueMap is
/// dynamic.
-static std::optional<int64_t> computeConstDiff(Value l, Value u) {
+static std::optional<APInt> computeConstDiff(Value l, Value u) {
IntegerAttr clb, cub;
if (matchPattern(l, m_Constant(&clb)) && matchPattern(u, m_Constant(&cub))) {
llvm::APInt lbValue = clb.getValue();
llvm::APInt ubValue = cub.getValue();
- return (ubValue - lbValue).getSExtValue();
+ return ubValue - lbValue;
}
// Else a simple pattern match for x + c or c + x
@@ -1003,7 +1016,7 @@ static std::optional<int64_t> computeConstDiff(Value l, Value u) {
u, m_Op<arith::AddIOp>(matchers::m_Val(l), m_ConstantInt(&diff))) ||
matchPattern(
u, m_Op<arith::AddIOp>(m_ConstantInt(&diff), matchers::m_Val(l))))
- return diff.getSExtValue();
+ return diff;
return std::nullopt;
}
@@ -1022,13 +1035,15 @@ struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
return success();
}
- std::optional<int64_t> diff =
+ std::optional<APInt> diff =
computeConstDiff(op.getLowerBound(), op.getUpperBound());
if (!diff)
return failure();
// If the loop is known to have 0 iterations, remove it.
- if (*diff <= 0) {
+ bool zeroOrLessIterations =
+ diff->isZero() || (!op.getUnsignedInt() && diff->isNegative());
+ if (zeroOrLessIterations) {
rewriter.replaceOp(op, op.getInitArgs());
return success();
}
@@ -3384,9 +3399,8 @@ ParseResult scf::WhileOp::parse(OpAsmParser &parser, OperationState &result) {
if (functionType.getNumInputs() != operands.size()) {
return parser.emitError(typeLoc)
- << "expected as many input types as operands "
- << "(expected " << operands.size() << " got "
- << functionType.getNumInputs() << ")";
+ << "expected as many input types as operands " << "(expected "
+ << operands.size() << " got " << functionType.getNumInputs() << ")";
}
// Resolve input operands.
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index f8799c52e8797..77a40e07765a2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -769,7 +769,8 @@ struct ForOpInterface
// Construct a new scf.for op with memref instead of tensor values.
auto newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), castedInitArgs);
+ forOp.getStep(), castedInitArgs, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block *loopBody = newForOp.getBody();
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
index bee77800e9d44..04bc14874267f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
@@ -58,9 +58,12 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
auto *beforeBlock = rewriter.createBlock(
&whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, lcvLocs);
rewriter.setInsertionPointToStart(whileOp.getBeforeBody());
- auto cmpOp = arith::CmpIOp::create(
- rewriter, whileOp.getLoc(), arith::CmpIPredicate::slt,
- beforeBlock->getArgument(0), forOp.getUpperBound());
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto cmpOp = arith::CmpIOp::create(rewriter, whileOp.getLoc(), predicate,
+ beforeBlock->getArgument(0),
+ forOp.getUpperBound());
scf::ConditionOp::create(rewriter, whileOp.getLoc(), cmpOp.getResult(),
beforeBlock->getArguments());
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index 1130538e51fb7..13dc8991a4735 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -449,9 +449,9 @@ scf::ForOp LoopPipelinerInternal::createKernelLoop(
arith::MulIOp::create(rewriter, loc, step, maxStageValue);
newUb = arith::SubIOp::create(rewriter, loc, ub, maxStageByStep);
}
- auto newForOp =
- scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb,
- forOp.getStep(), newLoopArg);
+ auto newForOp = scf::ForOp::create(
+ rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb, forOp.getStep(),
+ newLoopArg, /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// When there are no iter args, the loop body terminator will be created.
// Since we always create it below, remove the terminator if it was created.
if (!newForOp.getBody()->empty())
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 4752c0837c1c5..99bbaf5ab9efa 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -256,6 +256,10 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
LogicalResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override {
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Do not peel already peeled loops.
if (forOp->hasAttr(kPeeledLoopLabel))
return failure();
diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
index 1b07b77546030..a18cfebf86043 100644
--- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
@@ -116,7 +116,8 @@ class ConvertForOpTypes
llvm::getSingleElement(adaptor.getLowerBound()),
llvm::getSingleElement(adaptor.getUpperBound()),
llvm::getSingleElement(adaptor.getStep()),
- flattenValues(adaptor.getInitArgs()));
+ flattenValues(adaptor.getInitArgs()),
+ /*bodyBuilder=*/nullptr, op.getUnsignedInt());
// Reserve whatever attributes in the original op.
newOp->setAttrs(op->getAttrs());
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index c0e47ee1e74fc..473aacb77ea6e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -797,7 +797,8 @@ FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForOp>(
inits.append(newInitOperands.begin(), newInitOperands.end());
auto newLoop = scf::ForOp::create(
rewriter, loc, loopOp.getLowerBound(), loopOp.getUpperBound(),
- loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loopOp.getUnsignedInt());
// Move the loop body to the new op.
Block *loopBody = loopOp.getBody();
@@ -935,7 +936,8 @@ static LogicalResult addInitOperandsToLoopNest(
auto newLoop = scf::ForOp::create(
rewriter, forLoop.getLoc(), forLoop.getLowerBound(),
forLoop.getUpperBound(), forLoop.getStep(), newInits,
- [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {});
+ [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {},
+ forLoop.getUnsignedInt());
// Merge the body of the new loop with the body of the old loops.
SmallVector<Value> sourceBlockArgs;
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 57317951d609c..51196960fc7bc 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1233,6 +1233,7 @@ static void getPerfectlyNestedLoopsImpl(
static Loops stripmineSink(scf::ForOp forOp, Value factor,
ArrayRef<scf::ForOp> targets) {
+ assert(!forOp.getUnsignedInt() && "unsigned loops are not supported");
auto originalStep = forOp.getStep();
auto iv = forOp.getInductionVar();
@@ -1241,6 +1242,8 @@ static Loops stripmineSink(scf::ForOp forOp, Value factor,
Loops innerLoops;
for (auto t : targets) {
+ assert(!t.getUnsignedInt() && "unsigned loops are not supported");
+
// Save information for splicing ops out of t when done
auto begin = t.getBody()->begin();
auto nOps = t.getBody()->getOperations().size();
@@ -1415,6 +1418,8 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
Rewriter...
[truncated]
|
|
@llvm/pr-subscribers-mlir-linalg Author: Matthias Springer (matthias-springer) ChangesAdd a new unit attribute to allow for unsigned integer comparison. Example: scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
// body
}Discussion: https://discourse.llvm.org/t/scf-should-scf-for-support-unsigned-comparison/84655 Patch is 27.08 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/153379.diff 19 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0c1c15b85f4c9..41be006225730 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -169,9 +169,13 @@ def ForOp : SCF_Op<"for",
region capturing the loop body. The induction variable is represented as an
argument of this region. This SSA value is a signless integer or index.
The step is a value of same type but required to be positive, the lower and
- upper bounds can be also negative or zero. The lower and upper bounds specify
- a half-open range: the iteration is executed iff the signed comparison of induction
- variable value is less than the upper bound and bigger or equal to the lower bound.
+ upper bounds can be also negative or zero. The lower and upper bounds
+ specify a half-open range: the iteration is executed iff the comparison of
+ induction variable value is less than the upper bound and bigger or equal
+ to the lower bound.
+
+ By default, the integer comparison is signed. If the `unsignedInt` unit
+ attribute is specified, the integer comparison is unsigned.
The body region must contain exactly one block that terminates with
`scf.yield`. Calling ForOp::build will create such a region and insert
@@ -184,8 +188,8 @@ def ForOp : SCF_Op<"for",
... // body
}
...
- // Integer case.
- scf.for %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
+ // Unsigned integer case.
+ scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
... // body
}
```
@@ -258,7 +262,8 @@ def ForOp : SCF_Op<"for",
let arguments = (ins AnySignlessIntegerOrIndex:$lowerBound,
AnySignlessIntegerOrIndex:$upperBound,
AnySignlessIntegerOrIndex:$step,
- Variadic<AnyType>:$initArgs);
+ Variadic<AnyType>:$initArgs,
+ UnitAttr:$unsignedInt);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -266,7 +271,8 @@ def ForOp : SCF_Op<"for",
let builders = [OpBuilder<(ins "Value":$lowerBound, "Value":$upperBound,
"Value":$step, CArg<"ValueRange", "{}">:$initArgs,
CArg<"function_ref<void(OpBuilder &, Location, Value, ValueRange)>",
- "nullptr">)>];
+ "nullptr">,
+ CArg<"bool", "false">:$unsignedInt)>];
let extraClassDeclaration = [{
using BodyBuilderFn =
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index ba448e46913ac..c4f9ff8ab4c22 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -382,8 +382,11 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
// With the body block done, we can fill in the condition block.
rewriter.setInsertionPointToEnd(conditionBlock);
- auto comparison = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::slt, iv, upperBound);
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto comparison =
+ arith::CmpIOp::create(rewriter, loc, predicate, iv, upperBound);
cf::CondBranchOp::create(rewriter, loc, comparison, firstBodyBlock,
ArrayRef<Value>(), endBlock, ArrayRef<Value>());
diff --git a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
index 84cbd869c78ef..fd0339e7943e0 100644
--- a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
+++ b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
@@ -154,6 +154,10 @@ ForLowering::matchAndRewrite(ForOp forOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
Location loc = forOp.getLoc();
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Create an emitc::variable op for each result. These variables will be
// assigned to by emitc::assign ops within the loop body.
SmallVector<Value> resultVariables;
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index dc92367fc58cd..03042d87b8151 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -178,8 +178,14 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
// Generate the rest of the loop header.
rewriter.setInsertionPointToEnd(header);
auto *mergeBlock = loopOp.getMergeBlock();
- auto cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
- newIndVar, adaptor.getUpperBound());
+ Value cmpOp;
+ if (forOp.getUnsignedInt()) {
+ cmpOp = spirv::ULessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ } else {
+ cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ }
spirv::BranchConditionalOp::create(rewriter, loc, cmpOp, body,
ArrayRef<Value>(), mergeBlock,
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index fd530f26db526..807012d5438aa 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -594,7 +594,8 @@ static FailureOr<PackingResult> buildPackingLoopNestImpl(
auto clonedForOp = scf::ForOp::create(
rewriter, loc, bvm.lookupOrDefault(forOp.getLowerBound()),
bvm.lookupOrDefault(forOp.getUpperBound()),
- bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor);
+ bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// Map the induction var, region args and results to the `clonedForOp`.
bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 58986a6009995..8c0faebedfc34 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -55,7 +55,8 @@ static scf::ForOp replaceWithDifferentYield(RewriterBase &rewriter,
scf::ForOp newLoop = scf::ForOp::create(
rewriter, loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(),
- loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loop.getUnsignedInt());
// Generate the new yield with the replaced operand.
auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 0262a1b8a3893..ecd85c9101951 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -318,9 +318,12 @@ void ConditionOp::getSuccessorRegions(
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange initArgs,
- BodyBuilderFn bodyBuilder) {
+ BodyBuilderFn bodyBuilder, bool unsignedInt) {
OpBuilder::InsertionGuard guard(builder);
+ if (unsignedInt)
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
result.addOperands({lb, ub, step});
result.addOperands(initArgs);
for (Value v : initArgs)
@@ -450,6 +453,9 @@ static void printInitializationList(OpAsmPrinter &p,
}
void ForOp::print(OpAsmPrinter &p) {
+ if (getUnsignedInt())
+ p << " unsigned";
+
p << " " << getInductionVar() << " = " << getLowerBound() << " to "
<< getUpperBound() << " step " << getStep();
@@ -462,7 +468,8 @@ void ForOp::print(OpAsmPrinter &p) {
p.printRegion(getRegion(),
/*printEntryBlockArgs=*/false,
/*printBlockTerminators=*/!getInitArgs().empty());
- p.printOptionalAttrDict((*this)->getAttrs());
+ p.printOptionalAttrDict((*this)->getAttrs(),
+ /*elidedAttrs=*/getUnsignedIntAttrName().strref());
}
ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
@@ -472,6 +479,10 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
OpAsmParser::Argument inductionVariable;
OpAsmParser::UnresolvedOperand lb, ub, step;
+ if (succeeded(parser.parseOptionalKeyword("unsigned")))
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
+
// Parse the induction variable followed by '='.
if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual() ||
// Parse loop bounds.
@@ -562,7 +573,7 @@ ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,
inits.append(newInitOperands.begin(), newInitOperands.end());
scf::ForOp newLoop = scf::ForOp::create(
rewriter, getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
- [](OpBuilder &, Location, Value, ValueRange) {});
+ [](OpBuilder &, Location, Value, ValueRange) {}, getUnsignedInt());
newLoop->setAttrs(getPrunedAttributeList(getOperation(), {}));
// Generate the new yield values and append them to the scf.yield operation.
@@ -806,7 +817,8 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,
// 2. Create the new forOp shell.
scf::ForOp newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), newIterOperands);
+ forOp.getStep(), newIterOperands, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
SmallVector<Value, 4> newBlockTransferArgs(newBlock.getArguments().begin(),
@@ -931,7 +943,8 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
scf::ForOp newForOp =
scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),
- forOp.getUpperBound(), forOp.getStep(), newIterArgs);
+ forOp.getUpperBound(), forOp.getStep(), newIterArgs,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
@@ -989,12 +1002,12 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
/// Util function that tries to compute a constant diff between u and l.
/// Returns std::nullopt when the difference between two AffineValueMap is
/// dynamic.
-static std::optional<int64_t> computeConstDiff(Value l, Value u) {
+static std::optional<APInt> computeConstDiff(Value l, Value u) {
IntegerAttr clb, cub;
if (matchPattern(l, m_Constant(&clb)) && matchPattern(u, m_Constant(&cub))) {
llvm::APInt lbValue = clb.getValue();
llvm::APInt ubValue = cub.getValue();
- return (ubValue - lbValue).getSExtValue();
+ return ubValue - lbValue;
}
// Else a simple pattern match for x + c or c + x
@@ -1003,7 +1016,7 @@ static std::optional<int64_t> computeConstDiff(Value l, Value u) {
u, m_Op<arith::AddIOp>(matchers::m_Val(l), m_ConstantInt(&diff))) ||
matchPattern(
u, m_Op<arith::AddIOp>(m_ConstantInt(&diff), matchers::m_Val(l))))
- return diff.getSExtValue();
+ return diff;
return std::nullopt;
}
@@ -1022,13 +1035,15 @@ struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
return success();
}
- std::optional<int64_t> diff =
+ std::optional<APInt> diff =
computeConstDiff(op.getLowerBound(), op.getUpperBound());
if (!diff)
return failure();
// If the loop is known to have 0 iterations, remove it.
- if (*diff <= 0) {
+ bool zeroOrLessIterations =
+ diff->isZero() || (!op.getUnsignedInt() && diff->isNegative());
+ if (zeroOrLessIterations) {
rewriter.replaceOp(op, op.getInitArgs());
return success();
}
@@ -3384,9 +3399,8 @@ ParseResult scf::WhileOp::parse(OpAsmParser &parser, OperationState &result) {
if (functionType.getNumInputs() != operands.size()) {
return parser.emitError(typeLoc)
- << "expected as many input types as operands "
- << "(expected " << operands.size() << " got "
- << functionType.getNumInputs() << ")";
+ << "expected as many input types as operands " << "(expected "
+ << operands.size() << " got " << functionType.getNumInputs() << ")";
}
// Resolve input operands.
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index f8799c52e8797..77a40e07765a2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -769,7 +769,8 @@ struct ForOpInterface
// Construct a new scf.for op with memref instead of tensor values.
auto newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), castedInitArgs);
+ forOp.getStep(), castedInitArgs, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block *loopBody = newForOp.getBody();
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
index bee77800e9d44..04bc14874267f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
@@ -58,9 +58,12 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
auto *beforeBlock = rewriter.createBlock(
&whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, lcvLocs);
rewriter.setInsertionPointToStart(whileOp.getBeforeBody());
- auto cmpOp = arith::CmpIOp::create(
- rewriter, whileOp.getLoc(), arith::CmpIPredicate::slt,
- beforeBlock->getArgument(0), forOp.getUpperBound());
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto cmpOp = arith::CmpIOp::create(rewriter, whileOp.getLoc(), predicate,
+ beforeBlock->getArgument(0),
+ forOp.getUpperBound());
scf::ConditionOp::create(rewriter, whileOp.getLoc(), cmpOp.getResult(),
beforeBlock->getArguments());
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index 1130538e51fb7..13dc8991a4735 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -449,9 +449,9 @@ scf::ForOp LoopPipelinerInternal::createKernelLoop(
arith::MulIOp::create(rewriter, loc, step, maxStageValue);
newUb = arith::SubIOp::create(rewriter, loc, ub, maxStageByStep);
}
- auto newForOp =
- scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb,
- forOp.getStep(), newLoopArg);
+ auto newForOp = scf::ForOp::create(
+ rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb, forOp.getStep(),
+ newLoopArg, /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// When there are no iter args, the loop body terminator will be created.
// Since we always create it below, remove the terminator if it was created.
if (!newForOp.getBody()->empty())
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 4752c0837c1c5..99bbaf5ab9efa 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -256,6 +256,10 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
LogicalResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override {
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Do not peel already peeled loops.
if (forOp->hasAttr(kPeeledLoopLabel))
return failure();
diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
index 1b07b77546030..a18cfebf86043 100644
--- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
@@ -116,7 +116,8 @@ class ConvertForOpTypes
llvm::getSingleElement(adaptor.getLowerBound()),
llvm::getSingleElement(adaptor.getUpperBound()),
llvm::getSingleElement(adaptor.getStep()),
- flattenValues(adaptor.getInitArgs()));
+ flattenValues(adaptor.getInitArgs()),
+ /*bodyBuilder=*/nullptr, op.getUnsignedInt());
// Reserve whatever attributes in the original op.
newOp->setAttrs(op->getAttrs());
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index c0e47ee1e74fc..473aacb77ea6e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -797,7 +797,8 @@ FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForOp>(
inits.append(newInitOperands.begin(), newInitOperands.end());
auto newLoop = scf::ForOp::create(
rewriter, loc, loopOp.getLowerBound(), loopOp.getUpperBound(),
- loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loopOp.getUnsignedInt());
// Move the loop body to the new op.
Block *loopBody = loopOp.getBody();
@@ -935,7 +936,8 @@ static LogicalResult addInitOperandsToLoopNest(
auto newLoop = scf::ForOp::create(
rewriter, forLoop.getLoc(), forLoop.getLowerBound(),
forLoop.getUpperBound(), forLoop.getStep(), newInits,
- [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {});
+ [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {},
+ forLoop.getUnsignedInt());
// Merge the body of the new loop with the body of the old loops.
SmallVector<Value> sourceBlockArgs;
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 57317951d609c..51196960fc7bc 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1233,6 +1233,7 @@ static void getPerfectlyNestedLoopsImpl(
static Loops stripmineSink(scf::ForOp forOp, Value factor,
ArrayRef<scf::ForOp> targets) {
+ assert(!forOp.getUnsignedInt() && "unsigned loops are not supported");
auto originalStep = forOp.getStep();
auto iv = forOp.getInductionVar();
@@ -1241,6 +1242,8 @@ static Loops stripmineSink(scf::ForOp forOp, Value factor,
Loops innerLoops;
for (auto t : targets) {
+ assert(!t.getUnsignedInt() && "unsigned loops are not supported");
+
// Save information for splicing ops out of t when done
auto begin = t.getBody()->begin();
auto nOps = t.getBody()->getOperations().size();
@@ -1415,6 +1418,8 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
Rewriter...
[truncated]
|
|
@llvm/pr-subscribers-mlir Author: Matthias Springer (matthias-springer) ChangesAdd a new unit attribute to allow for unsigned integer comparison. Example: scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
// body
}Discussion: https://discourse.llvm.org/t/scf-should-scf-for-support-unsigned-comparison/84655 Patch is 27.08 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/153379.diff 19 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0c1c15b85f4c9..41be006225730 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -169,9 +169,13 @@ def ForOp : SCF_Op<"for",
region capturing the loop body. The induction variable is represented as an
argument of this region. This SSA value is a signless integer or index.
The step is a value of same type but required to be positive, the lower and
- upper bounds can be also negative or zero. The lower and upper bounds specify
- a half-open range: the iteration is executed iff the signed comparison of induction
- variable value is less than the upper bound and bigger or equal to the lower bound.
+ upper bounds can be also negative or zero. The lower and upper bounds
+ specify a half-open range: the iteration is executed iff the comparison of
+ induction variable value is less than the upper bound and bigger or equal
+ to the lower bound.
+
+ By default, the integer comparison is signed. If the `unsignedInt` unit
+ attribute is specified, the integer comparison is unsigned.
The body region must contain exactly one block that terminates with
`scf.yield`. Calling ForOp::build will create such a region and insert
@@ -184,8 +188,8 @@ def ForOp : SCF_Op<"for",
... // body
}
...
- // Integer case.
- scf.for %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
+ // Unsigned integer case.
+ scf.for unsigned %iv_32 = %lb_32 to %ub_32 step %step_32 : i32 {
... // body
}
```
@@ -258,7 +262,8 @@ def ForOp : SCF_Op<"for",
let arguments = (ins AnySignlessIntegerOrIndex:$lowerBound,
AnySignlessIntegerOrIndex:$upperBound,
AnySignlessIntegerOrIndex:$step,
- Variadic<AnyType>:$initArgs);
+ Variadic<AnyType>:$initArgs,
+ UnitAttr:$unsignedInt);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -266,7 +271,8 @@ def ForOp : SCF_Op<"for",
let builders = [OpBuilder<(ins "Value":$lowerBound, "Value":$upperBound,
"Value":$step, CArg<"ValueRange", "{}">:$initArgs,
CArg<"function_ref<void(OpBuilder &, Location, Value, ValueRange)>",
- "nullptr">)>];
+ "nullptr">,
+ CArg<"bool", "false">:$unsignedInt)>];
let extraClassDeclaration = [{
using BodyBuilderFn =
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index ba448e46913ac..c4f9ff8ab4c22 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -382,8 +382,11 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
// With the body block done, we can fill in the condition block.
rewriter.setInsertionPointToEnd(conditionBlock);
- auto comparison = arith::CmpIOp::create(
- rewriter, loc, arith::CmpIPredicate::slt, iv, upperBound);
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto comparison =
+ arith::CmpIOp::create(rewriter, loc, predicate, iv, upperBound);
cf::CondBranchOp::create(rewriter, loc, comparison, firstBodyBlock,
ArrayRef<Value>(), endBlock, ArrayRef<Value>());
diff --git a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
index 84cbd869c78ef..fd0339e7943e0 100644
--- a/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
+++ b/mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
@@ -154,6 +154,10 @@ ForLowering::matchAndRewrite(ForOp forOp, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const {
Location loc = forOp.getLoc();
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Create an emitc::variable op for each result. These variables will be
// assigned to by emitc::assign ops within the loop body.
SmallVector<Value> resultVariables;
diff --git a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
index dc92367fc58cd..03042d87b8151 100644
--- a/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
+++ b/mlir/lib/Conversion/SCFToSPIRV/SCFToSPIRV.cpp
@@ -178,8 +178,14 @@ struct ForOpConversion final : SCFToSPIRVPattern<scf::ForOp> {
// Generate the rest of the loop header.
rewriter.setInsertionPointToEnd(header);
auto *mergeBlock = loopOp.getMergeBlock();
- auto cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
- newIndVar, adaptor.getUpperBound());
+ Value cmpOp;
+ if (forOp.getUnsignedInt()) {
+ cmpOp = spirv::ULessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ } else {
+ cmpOp = spirv::SLessThanOp::create(rewriter, loc, rewriter.getI1Type(),
+ newIndVar, adaptor.getUpperBound());
+ }
spirv::BranchConditionalOp::create(rewriter, loc, cmpOp, body,
ArrayRef<Value>(), mergeBlock,
diff --git a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
index fd530f26db526..807012d5438aa 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp
@@ -594,7 +594,8 @@ static FailureOr<PackingResult> buildPackingLoopNestImpl(
auto clonedForOp = scf::ForOp::create(
rewriter, loc, bvm.lookupOrDefault(forOp.getLowerBound()),
bvm.lookupOrDefault(forOp.getUpperBound()),
- bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor);
+ bvm.lookupOrDefault(forOp.getStep()), hoistedPackedTensor,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// Map the induction var, region args and results to the `clonedForOp`.
bvm.map(forOp.getInductionVar(), clonedForOp.getInductionVar());
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
index 58986a6009995..8c0faebedfc34 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -55,7 +55,8 @@ static scf::ForOp replaceWithDifferentYield(RewriterBase &rewriter,
scf::ForOp newLoop = scf::ForOp::create(
rewriter, loop.getLoc(), loop.getLowerBound(), loop.getUpperBound(),
- loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loop.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loop.getUnsignedInt());
// Generate the new yield with the replaced operand.
auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 0262a1b8a3893..ecd85c9101951 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -318,9 +318,12 @@ void ConditionOp::getSuccessorRegions(
void ForOp::build(OpBuilder &builder, OperationState &result, Value lb,
Value ub, Value step, ValueRange initArgs,
- BodyBuilderFn bodyBuilder) {
+ BodyBuilderFn bodyBuilder, bool unsignedInt) {
OpBuilder::InsertionGuard guard(builder);
+ if (unsignedInt)
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
result.addOperands({lb, ub, step});
result.addOperands(initArgs);
for (Value v : initArgs)
@@ -450,6 +453,9 @@ static void printInitializationList(OpAsmPrinter &p,
}
void ForOp::print(OpAsmPrinter &p) {
+ if (getUnsignedInt())
+ p << " unsigned";
+
p << " " << getInductionVar() << " = " << getLowerBound() << " to "
<< getUpperBound() << " step " << getStep();
@@ -462,7 +468,8 @@ void ForOp::print(OpAsmPrinter &p) {
p.printRegion(getRegion(),
/*printEntryBlockArgs=*/false,
/*printBlockTerminators=*/!getInitArgs().empty());
- p.printOptionalAttrDict((*this)->getAttrs());
+ p.printOptionalAttrDict((*this)->getAttrs(),
+ /*elidedAttrs=*/getUnsignedIntAttrName().strref());
}
ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
@@ -472,6 +479,10 @@ ParseResult ForOp::parse(OpAsmParser &parser, OperationState &result) {
OpAsmParser::Argument inductionVariable;
OpAsmParser::UnresolvedOperand lb, ub, step;
+ if (succeeded(parser.parseOptionalKeyword("unsigned")))
+ result.addAttribute(getUnsignedIntAttrName(result.name),
+ builder.getUnitAttr());
+
// Parse the induction variable followed by '='.
if (parser.parseOperand(inductionVariable.ssaName) || parser.parseEqual() ||
// Parse loop bounds.
@@ -562,7 +573,7 @@ ForOp::replaceWithAdditionalYields(RewriterBase &rewriter,
inits.append(newInitOperands.begin(), newInitOperands.end());
scf::ForOp newLoop = scf::ForOp::create(
rewriter, getLoc(), getLowerBound(), getUpperBound(), getStep(), inits,
- [](OpBuilder &, Location, Value, ValueRange) {});
+ [](OpBuilder &, Location, Value, ValueRange) {}, getUnsignedInt());
newLoop->setAttrs(getPrunedAttributeList(getOperation(), {}));
// Generate the new yield values and append them to the scf.yield operation.
@@ -806,7 +817,8 @@ mlir::scf::replaceAndCastForOpIterArg(RewriterBase &rewriter, scf::ForOp forOp,
// 2. Create the new forOp shell.
scf::ForOp newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), newIterOperands);
+ forOp.getStep(), newIterOperands, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
SmallVector<Value, 4> newBlockTransferArgs(newBlock.getArguments().begin(),
@@ -931,7 +943,8 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
scf::ForOp newForOp =
scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),
- forOp.getUpperBound(), forOp.getStep(), newIterArgs);
+ forOp.getUpperBound(), forOp.getStep(), newIterArgs,
+ /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block &newBlock = newForOp.getRegion().front();
@@ -989,12 +1002,12 @@ struct ForOpIterArgsFolder : public OpRewritePattern<scf::ForOp> {
/// Util function that tries to compute a constant diff between u and l.
/// Returns std::nullopt when the difference between two AffineValueMap is
/// dynamic.
-static std::optional<int64_t> computeConstDiff(Value l, Value u) {
+static std::optional<APInt> computeConstDiff(Value l, Value u) {
IntegerAttr clb, cub;
if (matchPattern(l, m_Constant(&clb)) && matchPattern(u, m_Constant(&cub))) {
llvm::APInt lbValue = clb.getValue();
llvm::APInt ubValue = cub.getValue();
- return (ubValue - lbValue).getSExtValue();
+ return ubValue - lbValue;
}
// Else a simple pattern match for x + c or c + x
@@ -1003,7 +1016,7 @@ static std::optional<int64_t> computeConstDiff(Value l, Value u) {
u, m_Op<arith::AddIOp>(matchers::m_Val(l), m_ConstantInt(&diff))) ||
matchPattern(
u, m_Op<arith::AddIOp>(m_ConstantInt(&diff), matchers::m_Val(l))))
- return diff.getSExtValue();
+ return diff;
return std::nullopt;
}
@@ -1022,13 +1035,15 @@ struct SimplifyTrivialLoops : public OpRewritePattern<ForOp> {
return success();
}
- std::optional<int64_t> diff =
+ std::optional<APInt> diff =
computeConstDiff(op.getLowerBound(), op.getUpperBound());
if (!diff)
return failure();
// If the loop is known to have 0 iterations, remove it.
- if (*diff <= 0) {
+ bool zeroOrLessIterations =
+ diff->isZero() || (!op.getUnsignedInt() && diff->isNegative());
+ if (zeroOrLessIterations) {
rewriter.replaceOp(op, op.getInitArgs());
return success();
}
@@ -3384,9 +3399,8 @@ ParseResult scf::WhileOp::parse(OpAsmParser &parser, OperationState &result) {
if (functionType.getNumInputs() != operands.size()) {
return parser.emitError(typeLoc)
- << "expected as many input types as operands "
- << "(expected " << operands.size() << " got "
- << functionType.getNumInputs() << ")";
+ << "expected as many input types as operands " << "(expected "
+ << operands.size() << " got " << functionType.getNumInputs() << ")";
}
// Resolve input operands.
diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index f8799c52e8797..77a40e07765a2 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -769,7 +769,8 @@ struct ForOpInterface
// Construct a new scf.for op with memref instead of tensor values.
auto newForOp = scf::ForOp::create(
rewriter, forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(),
- forOp.getStep(), castedInitArgs);
+ forOp.getStep(), castedInitArgs, /*bodyBuilder=*/nullptr,
+ forOp.getUnsignedInt());
newForOp->setAttrs(forOp->getAttrs());
Block *loopBody = newForOp.getBody();
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
index bee77800e9d44..04bc14874267f 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForToWhile.cpp
@@ -58,9 +58,12 @@ struct ForLoopLoweringPattern : public OpRewritePattern<ForOp> {
auto *beforeBlock = rewriter.createBlock(
&whileOp.getBefore(), whileOp.getBefore().begin(), lcvTypes, lcvLocs);
rewriter.setInsertionPointToStart(whileOp.getBeforeBody());
- auto cmpOp = arith::CmpIOp::create(
- rewriter, whileOp.getLoc(), arith::CmpIPredicate::slt,
- beforeBlock->getArgument(0), forOp.getUpperBound());
+ arith::CmpIPredicate predicate = forOp.getUnsignedInt()
+ ? arith::CmpIPredicate::ult
+ : arith::CmpIPredicate::slt;
+ auto cmpOp = arith::CmpIOp::create(rewriter, whileOp.getLoc(), predicate,
+ beforeBlock->getArgument(0),
+ forOp.getUpperBound());
scf::ConditionOp::create(rewriter, whileOp.getLoc(), cmpOp.getResult(),
beforeBlock->getArguments());
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index 1130538e51fb7..13dc8991a4735 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -449,9 +449,9 @@ scf::ForOp LoopPipelinerInternal::createKernelLoop(
arith::MulIOp::create(rewriter, loc, step, maxStageValue);
newUb = arith::SubIOp::create(rewriter, loc, ub, maxStageByStep);
}
- auto newForOp =
- scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb,
- forOp.getStep(), newLoopArg);
+ auto newForOp = scf::ForOp::create(
+ rewriter, forOp.getLoc(), forOp.getLowerBound(), newUb, forOp.getStep(),
+ newLoopArg, /*bodyBuilder=*/nullptr, forOp.getUnsignedInt());
// When there are no iter args, the loop body terminator will be created.
// Since we always create it below, remove the terminator if it was created.
if (!newForOp.getBody()->empty())
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index 4752c0837c1c5..99bbaf5ab9efa 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -256,6 +256,10 @@ struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
LogicalResult matchAndRewrite(ForOp forOp,
PatternRewriter &rewriter) const override {
+ if (forOp.getUnsignedInt())
+ return rewriter.notifyMatchFailure(forOp,
+ "unsigned loops are not supported");
+
// Do not peel already peeled loops.
if (forOp->hasAttr(kPeeledLoopLabel))
return failure();
diff --git a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
index 1b07b77546030..a18cfebf86043 100644
--- a/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp
@@ -116,7 +116,8 @@ class ConvertForOpTypes
llvm::getSingleElement(adaptor.getLowerBound()),
llvm::getSingleElement(adaptor.getUpperBound()),
llvm::getSingleElement(adaptor.getStep()),
- flattenValues(adaptor.getInitArgs()));
+ flattenValues(adaptor.getInitArgs()),
+ /*bodyBuilder=*/nullptr, op.getUnsignedInt());
// Reserve whatever attributes in the original op.
newOp->setAttrs(op->getAttrs());
diff --git a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
index c0e47ee1e74fc..473aacb77ea6e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
@@ -797,7 +797,8 @@ FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForOp>(
inits.append(newInitOperands.begin(), newInitOperands.end());
auto newLoop = scf::ForOp::create(
rewriter, loc, loopOp.getLowerBound(), loopOp.getUpperBound(),
- loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {});
+ loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
+ loopOp.getUnsignedInt());
// Move the loop body to the new op.
Block *loopBody = loopOp.getBody();
@@ -935,7 +936,8 @@ static LogicalResult addInitOperandsToLoopNest(
auto newLoop = scf::ForOp::create(
rewriter, forLoop.getLoc(), forLoop.getLowerBound(),
forLoop.getUpperBound(), forLoop.getStep(), newInits,
- [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {});
+ [&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {},
+ forLoop.getUnsignedInt());
// Merge the body of the new loop with the body of the old loops.
SmallVector<Value> sourceBlockArgs;
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 57317951d609c..51196960fc7bc 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1233,6 +1233,7 @@ static void getPerfectlyNestedLoopsImpl(
static Loops stripmineSink(scf::ForOp forOp, Value factor,
ArrayRef<scf::ForOp> targets) {
+ assert(!forOp.getUnsignedInt() && "unsigned loops are not supported");
auto originalStep = forOp.getStep();
auto iv = forOp.getInductionVar();
@@ -1241,6 +1242,8 @@ static Loops stripmineSink(scf::ForOp forOp, Value factor,
Loops innerLoops;
for (auto t : targets) {
+ assert(!t.getUnsignedInt() && "unsigned loops are not supported");
+
// Save information for splicing ops out of t when done
auto begin = t.getBody()->begin();
auto nOps = t.getBody()->getOperations().size();
@@ -1415,6 +1418,8 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
Rewriter...
[truncated]
|
rengolin
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good to me, I don't mind if this is an explicit or implicit attribute.
Won't approve to let other people review, as this changes syntax and semantics.
b5492f0 to
e2e71b3
Compare
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
e2e71b3 to
e021424
Compare
joker-eph
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, but let's wait a bit to let enough time for people to provide feedback.
kuhar
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Makes sense
banach-space
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks!
Add a new unit attribute to allow for unsigned integer comparison.
Example:
Discussion: https://discourse.llvm.org/t/scf-should-scf-for-support-unsigned-comparison/84655