[Mlir-commits] [mlir] 6b4c122 - [mlir][loops] Add getters for multi dim loop variables in `LoopLikeOpInterface` (#94516)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jun 7 16:25:46 PDT 2024
Author: srcarroll
Date: 2024-06-07T18:25:43-05:00
New Revision: 6b4c12284795a3030e37b17047271a47a69bb587
URL: https://github.com/llvm/llvm-project/commit/6b4c12284795a3030e37b17047271a47a69bb587
DIFF: https://github.com/llvm/llvm-project/commit/6b4c12284795a3030e37b17047271a47a69bb587.diff
LOG: [mlir][loops] Add getters for multi dim loop variables in `LoopLikeOpInterface` (#94516)
This patch adds `getLoopInductionVars`, `getLoopLowerBounds`,
`getLoopBounds`, `getLoopSteps` interface methods to
`LoopLIkeOpInterface`. The corresponding single value versions have been
moved to shared class declaration and have been implemented based on the
new interface methods.
Added:
Modified:
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
mlir/include/mlir/Interfaces/LoopLikeInterface.td
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
mlir/lib/Dialect/SCF/IR/SCF.cpp
mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 3640055ea8da8..dbec741cf1b1f 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -118,8 +118,8 @@ def AffineForOp : Affine_Op<"for",
[AttrSizedOperandSegments, AutomaticAllocationScope,
ImplicitAffineTerminator, ConditionallySpeculatable,
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound", "getYieldedValuesMutable",
+ ["getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
+ "getLoopUpperBounds", "getYieldedValuesMutable",
"replaceWithAdditionalYields"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>]> {
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 0b063aa772bab..f35ea962bea16 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -136,8 +136,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
def ForOp : SCF_Op<"for",
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInitsMutable", "getLoopResults", "getRegionIterArgs",
- "getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
- "getSingleUpperBound", "getYieldedValuesMutable",
+ "getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
+ "getLoopUpperBounds", "getYieldedValuesMutable",
"promoteIfSingleIteration", "replaceWithAdditionalYields",
"yieldTiledValuesAndReplace"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
@@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getInitsMutable", "getRegionIterArgs", "getSingleInductionVar",
- "getSingleLowerBound", "getSingleUpperBound", "getSingleStep",
+ ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,22 +510,31 @@ def ForallOp : SCF_Op<"forall", [
];
let extraClassDeclaration = [{
- // Get lower bounds as OpFoldResult.
+ /// Get induction variables.
+ SmallVector<Value> getInductionVars() {
+ std::optional<SmallVector<Value>> maybeInductionVars = getLoopInductionVars();
+ assert(maybeInductionVars.has_value() && "expected values");
+ return *maybeInductionVars;
+ }
+ /// Get lower bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedLowerBound() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
+ std::optional<SmallVector<OpFoldResult>> maybeLowerBounds = getLoopLowerBounds();
+ assert(maybeLowerBounds.has_value() && "expected values");
+ return *maybeLowerBounds;
}
- // Get upper bounds as OpFoldResult.
+ /// Get upper bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedUpperBound() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
+ std::optional<SmallVector<OpFoldResult>> maybeUpperBounds = getLoopUpperBounds();
+ assert(maybeUpperBounds.has_value() && "expected values");
+ return *maybeUpperBounds;
}
- // Get steps as OpFoldResult.
+ /// Get steps as OpFoldResult.
SmallVector<OpFoldResult> getMixedStep() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticStep(), getDynamicStep(), b);
+ std::optional<SmallVector<OpFoldResult>> maybeSteps = getLoopSteps();
+ assert(maybeSteps.has_value() && "expected values");
+ return *maybeSteps;
}
/// Get lower bounds as values.
@@ -584,10 +593,6 @@ def ForallOp : SCF_Op<"forall", [
getNumDynamicControlOperands() + getRank());
}
- ::mlir::ValueRange getInductionVars() {
- return getBody()->getArguments().take_front(getRank());
- }
-
::mlir::Value getInductionVar(int64_t idx) {
return getInductionVars()[idx];
}
@@ -765,8 +770,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
def ParallelOp : SCF_Op<"parallel",
[AutomaticAllocationScope,
AttrSizedOperandSegments,
- DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getSingleInductionVar",
- "getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
+ DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps"]>,
RecursiveMemoryEffects,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::ReduceOp">,
@@ -846,8 +851,11 @@ def ParallelOp : SCF_Op<"parallel",
];
let extraClassDeclaration = [{
- ValueRange getInductionVars() {
- return getBody()->getArguments();
+ /// Get induction variables.
+ SmallVector<Value> getInductionVars() {
+ std::optional<SmallVector<Value>> maybeInductionVars = getLoopInductionVars();;
+ assert(maybeInductionVars.has_value() && "expected values");
+ return *maybeInductionVars;
}
unsigned getNumLoops() { return getStep().size(); }
unsigned getNumReductions() { return getInitVals().size(); }
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index f0dc6e60eba58..b748d5e29114a 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -93,51 +93,59 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}]
>,
InterfaceMethod<[{
- If there is a single induction variable return it, otherwise return
- std::nullopt.
+ Return all induction variables, if they exist. If the op has no notion of
+ induction variable, then return std::nullopt. If it does have
+ a notion but an instance doesn't have induction variables, then
+ return empty vector.
}],
- /*retTy=*/"::std::optional<::mlir::Value>",
- /*methodName=*/"getSingleInductionVar",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::Value>>",
+ /*methodName=*/"getLoopInductionVars",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return ::std::nullopt;
}]
>,
InterfaceMethod<[{
- Return the single lower bound value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all lower bounds, if they exist. If the op has no notion of
+ lower bounds, then return std::nullopt. If it does have
+ a notion but an instance doesn't have lower bounds, then
+ return empty vector.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleLowerBound",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getLoopLowerBounds",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return ::std::nullopt;
}]
>,
InterfaceMethod<[{
- Return the single step value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all steps, if they exist. If the op has no notion of
+ steps, then return std::nullopt. If it does have
+ a notion but an instance doesn't have steps, then
+ return empty vector.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleStep",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getLoopSteps",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return ::std::nullopt;
}]
>,
InterfaceMethod<[{
- Return the single upper bound value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all upper bounds, if they exist. If the op has no notion of
+ lower bounds, then return std::nullopt. If it does have
+ a notion but an instance doesn't have lower bounds, then
+ return empty vector.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleUpperBound",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getLoopUpperBounds",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return ::std::nullopt;
}]
>,
InterfaceMethod<[{
@@ -235,6 +243,39 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}];
let extraSharedClassDeclaration = [{
+ /// If there is a single induction variable return it, otherwise return
+ /// std::nullopt.
+ ::std::optional<::mlir::Value> getSingleInductionVar() {
+ auto inductionVars = this->getLoopInductionVars();
+ if (inductionVars.has_value() && (*inductionVars).size() == 1)
+ return (*inductionVars)[0];
+ return std::nullopt;
+ }
+ /// Return the single lower bound value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
+ auto lowerBounds = this->getLoopLowerBounds();
+ if (lowerBounds.has_value() && (*lowerBounds).size() == 1)
+ return (*lowerBounds)[0];
+ return std::nullopt;
+ }
+ /// Return the single step value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleStep() {
+ auto steps = this->getLoopSteps();
+ if (steps.has_value() && (*steps).size() == 1)
+ return (*steps)[0];
+ return std::nullopt;
+ }
+ /// Return the single upper bound value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleUpperBound() {
+ auto upperBounds = this->getLoopUpperBounds();
+ if (upperBounds.has_value() && (*upperBounds).size() == 1)
+ return (*upperBounds)[0];
+ return std::nullopt;
+ }
+
/// Append the specified additional "init" operands: replace this loop with
/// a new loop that has the additional init operands. The loop body of this
/// loop is moved over to the new loop.
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 2e31487bd55a0..0a58d2fdb02f5 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2454,27 +2454,30 @@ bool AffineForOp::matchingBoundOperandList() {
SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
-std::optional<Value> AffineForOp::getSingleInductionVar() {
- return getInductionVar();
+std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() {
+ return SmallVector<Value>{getInductionVar()};
}
-std::optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() {
if (!hasConstantLowerBound())
return std::nullopt;
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()));
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
}
-std::optional<OpFoldResult> AffineForOp::getSingleStep() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopSteps() {
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getStepAsInt()));
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
}
-std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() {
if (!hasConstantUpperBound())
- return std::nullopt;
+ return {};
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
}
FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
index b0a4de2da1e86..8b0e04fb61b1b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Loops.cpp
@@ -184,8 +184,7 @@ static void replaceIndexOpsByInductionVariables(RewriterBase &rewriter,
for (Operation *loopOp : loopOps) {
llvm::TypeSwitch<Operation *>(loopOp)
.Case([&](scf::ParallelOp parallelOp) {
- allIvs.append(parallelOp.getInductionVars().begin(),
- parallelOp.getInductionVars().end());
+ allIvs.append(parallelOp.getInductionVars());
})
.Case([&](scf::ForOp forOp) {
allIvs.push_back(forOp.getInductionVar());
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
index fd314ef9f8134..a0a0e11a6903d 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp
@@ -243,7 +243,7 @@ static void calculateTileOffsetsAndSizes(
OpBuilder::InsertionGuard g(b);
b.setInsertionPointToStart(forallOp.getBody(0));
- ValueRange threadIds = forallOp.getInductionVars();
+ SmallVector<Value> threadIds = forallOp.getInductionVars();
SmallVector<OpFoldResult> nonZeroNumThreads =
llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
return !isConstantIntValue(ofr, 0);
@@ -746,7 +746,7 @@ FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
b.getIndexAttr(0));
SmallVector<OpFoldResult> sizes = tiledSizes;
sizes[reductionDim] = b.getIndexAttr(1);
- outOffsets[reductionDim] = forallOp.getInductionVars().front();
+ outOffsets[reductionDim] = forallOp.getInductionVars()[0];
// TODO: use SubsetExtractOpInterface once it is available.
tiledDpsInitOperands.push_back(b.create<tensor::ExtractSliceOp>(
loc, cast<RankedTensorType>(initOperand.getType()),
@@ -814,7 +814,7 @@ FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
int64_t sizeIdx = 0;
for (int64_t i = 0, e = numThreads.size(); i < e; ++i) {
if (i == reductionDim) {
- resultOffsetsRank.push_back(forallOp.getInductionVars().front());
+ resultOffsetsRank.push_back(forallOp.getInductionVars()[0]);
resultSizesRank.push_back(b.getIndexAttr(1));
continue;
}
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 107fd0690f193..5e94f4dc612a7 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -378,20 +378,20 @@ LogicalResult ForOp::verifyRegions() {
return success();
}
-std::optional<Value> ForOp::getSingleInductionVar() {
- return getInductionVar();
+std::optional<SmallVector<Value>> ForOp::getLoopInductionVars() {
+ return SmallVector<Value>{getInductionVar()};
}
-std::optional<OpFoldResult> ForOp::getSingleLowerBound() {
- return OpFoldResult(getLowerBound());
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {
+ return SmallVector<OpFoldResult>{OpFoldResult(getLowerBound())};
}
-std::optional<OpFoldResult> ForOp::getSingleStep() {
- return OpFoldResult(getStep());
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopSteps() {
+ return SmallVector<OpFoldResult>{OpFoldResult(getStep())};
}
-std::optional<OpFoldResult> ForOp::getSingleUpperBound() {
- return OpFoldResult(getUpperBound());
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopUpperBounds() {
+ return SmallVector<OpFoldResult>{OpFoldResult(getUpperBound())};
}
std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }
@@ -1428,28 +1428,26 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-std::optional<Value> ForallOp::getSingleInductionVar() {
- if (getRank() != 1)
- return std::nullopt;
- return getInductionVar(0);
+std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {
+ return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
-std::optional<OpFoldResult> ForallOp::getSingleLowerBound() {
- if (getRank() != 1)
- return std::nullopt;
- return getMixedLowerBound()[0];
+// Get lower bounds as OpFoldResult.
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopLowerBounds() {
+ Builder b(getOperation()->getContext());
+ return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
}
-std::optional<OpFoldResult> ForallOp::getSingleUpperBound() {
- if (getRank() != 1)
- return std::nullopt;
- return getMixedUpperBound()[0];
+// Get upper bounds as OpFoldResult.
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopUpperBounds() {
+ Builder b(getOperation()->getContext());
+ return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
}
-std::optional<OpFoldResult> ForallOp::getSingleStep() {
- if (getRank() != 1)
- return std::nullopt;
- return getMixedStep()[0];
+// Get steps as OpFoldResult.
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopSteps() {
+ Builder b(getOperation()->getContext());
+ return getMixedValues(getStaticStep(), getDynamicStep(), b);
}
ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) {
@@ -3008,28 +3006,20 @@ void ParallelOp::print(OpAsmPrinter &p) {
SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
-std::optional<Value> ParallelOp::getSingleInductionVar() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getBody()->getArgument(0);
+std::optional<SmallVector<Value>> ParallelOp::getLoopInductionVars() {
+ return SmallVector<Value>{getBody()->getArguments()};
}
-std::optional<OpFoldResult> ParallelOp::getSingleLowerBound() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getLowerBound()[0];
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopLowerBounds() {
+ return getLowerBound();
}
-std::optional<OpFoldResult> ParallelOp::getSingleUpperBound() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getUpperBound()[0];
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopUpperBounds() {
+ return getUpperBound();
}
-std::optional<OpFoldResult> ParallelOp::getSingleStep() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getStep()[0];
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopSteps() {
+ return getStep();
}
ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
index 198cb2e6cc69e..5da1b76e929be 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForallToFor.cpp
@@ -34,12 +34,9 @@ mlir::scf::forallToForLoop(RewriterBase &rewriter, scf::ForallOp forallOp,
rewriter.setInsertionPoint(forallOp);
Location loc = forallOp.getLoc();
- SmallVector<Value> lbs = getValueOrCreateConstantIndexOp(
- rewriter, loc, forallOp.getMixedLowerBound());
- SmallVector<Value> ubs = getValueOrCreateConstantIndexOp(
- rewriter, loc, forallOp.getMixedUpperBound());
- SmallVector<Value> steps =
- getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep());
+ SmallVector<Value> lbs = forallOp.getLowerBound(rewriter);
+ SmallVector<Value> ubs = forallOp.getUpperBound(rewriter);
+ SmallVector<Value> steps = forallOp.getStep(rewriter);
LoopNest loopNest = scf::buildLoopNest(rewriter, loc, lbs, ubs, steps);
SmallVector<Value> ivs = llvm::map_to_vector(
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
index 1fc0331300379..44e6840b03a3d 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
@@ -35,12 +35,9 @@ LogicalResult mlir::scf::forallToParallelLoop(RewriterBase &rewriter,
"only fully bufferized scf.forall ops can be lowered to scf.parallel");
// Convert mixed bounds and steps to SSA values.
- SmallVector<Value> lbs = getValueOrCreateConstantIndexOp(
- rewriter, loc, forallOp.getMixedLowerBound());
- SmallVector<Value> ubs = getValueOrCreateConstantIndexOp(
- rewriter, loc, forallOp.getMixedUpperBound());
- SmallVector<Value> steps =
- getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep());
+ SmallVector<Value> lbs = forallOp.getLowerBound(rewriter);
+ SmallVector<Value> ubs = forallOp.getUpperBound(rewriter);
+ SmallVector<Value> steps = forallOp.getStep(rewriter);
// Create empty scf.parallel op.
auto parallelOp = rewriter.create<scf::ParallelOp>(loc, lbs, ubs, steps);
diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
index 6bc0fd6113b9b..53a4af14d119a 100644
--- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
+++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -27,27 +27,65 @@ class SCFLoopLikeTest : public ::testing::Test {
}
void checkUnidimensional(LoopLikeOpInterface loopLikeOp) {
- std::optional<OpFoldResult> maybeLb = loopLikeOp.getSingleLowerBound();
- EXPECT_TRUE(maybeLb.has_value());
- std::optional<OpFoldResult> maybeUb = loopLikeOp.getSingleUpperBound();
- EXPECT_TRUE(maybeUb.has_value());
- std::optional<OpFoldResult> maybeStep = loopLikeOp.getSingleStep();
- EXPECT_TRUE(maybeStep.has_value());
- std::optional<OpFoldResult> maybeIndVar =
+ std::optional<OpFoldResult> maybeSingleLb =
+ loopLikeOp.getSingleLowerBound();
+ EXPECT_TRUE(maybeSingleLb.has_value());
+ std::optional<OpFoldResult> maybeSingleUb =
+ loopLikeOp.getSingleUpperBound();
+ EXPECT_TRUE(maybeSingleUb.has_value());
+ std::optional<OpFoldResult> maybeSingleStep = loopLikeOp.getSingleStep();
+ EXPECT_TRUE(maybeSingleStep.has_value());
+ std::optional<OpFoldResult> maybeSingleIndVar =
loopLikeOp.getSingleInductionVar();
- EXPECT_TRUE(maybeIndVar.has_value());
+ EXPECT_TRUE(maybeSingleIndVar.has_value());
+
+ std::optional<SmallVector<OpFoldResult>> maybeLb =
+ loopLikeOp.getLoopLowerBounds();
+ ASSERT_TRUE(maybeLb.has_value());
+ EXPECT_EQ((*maybeLb).size(), 1u);
+ std::optional<SmallVector<OpFoldResult>> maybeUb =
+ loopLikeOp.getLoopUpperBounds();
+ ASSERT_TRUE(maybeUb.has_value());
+ EXPECT_EQ((*maybeUb).size(), 1u);
+ std::optional<SmallVector<OpFoldResult>> maybeStep =
+ loopLikeOp.getLoopSteps();
+ ASSERT_TRUE(maybeStep.has_value());
+ EXPECT_EQ((*maybeStep).size(), 1u);
+ std::optional<SmallVector<Value>> maybeInductionVars =
+ loopLikeOp.getLoopInductionVars();
+ ASSERT_TRUE(maybeInductionVars.has_value());
+ EXPECT_EQ((*maybeInductionVars).size(), 1u);
}
void checkMultidimensional(LoopLikeOpInterface loopLikeOp) {
- std::optional<OpFoldResult> maybeLb = loopLikeOp.getSingleLowerBound();
- EXPECT_FALSE(maybeLb.has_value());
- std::optional<OpFoldResult> maybeUb = loopLikeOp.getSingleUpperBound();
- EXPECT_FALSE(maybeUb.has_value());
- std::optional<OpFoldResult> maybeStep = loopLikeOp.getSingleStep();
- EXPECT_FALSE(maybeStep.has_value());
- std::optional<OpFoldResult> maybeIndVar =
+ std::optional<OpFoldResult> maybeSingleLb =
+ loopLikeOp.getSingleLowerBound();
+ EXPECT_FALSE(maybeSingleLb.has_value());
+ std::optional<OpFoldResult> maybeSingleUb =
+ loopLikeOp.getSingleUpperBound();
+ EXPECT_FALSE(maybeSingleUb.has_value());
+ std::optional<OpFoldResult> maybeSingleStep = loopLikeOp.getSingleStep();
+ EXPECT_FALSE(maybeSingleStep.has_value());
+ std::optional<OpFoldResult> maybeSingleIndVar =
loopLikeOp.getSingleInductionVar();
- EXPECT_FALSE(maybeIndVar.has_value());
+ EXPECT_FALSE(maybeSingleIndVar.has_value());
+
+ std::optional<SmallVector<OpFoldResult>> maybeLb =
+ loopLikeOp.getLoopLowerBounds();
+ ASSERT_TRUE(maybeLb.has_value());
+ EXPECT_EQ((*maybeLb).size(), 2u);
+ std::optional<SmallVector<OpFoldResult>> maybeUb =
+ loopLikeOp.getLoopUpperBounds();
+ ASSERT_TRUE(maybeUb.has_value());
+ EXPECT_EQ((*maybeUb).size(), 2u);
+ std::optional<SmallVector<OpFoldResult>> maybeStep =
+ loopLikeOp.getLoopSteps();
+ ASSERT_TRUE(maybeStep.has_value());
+ EXPECT_EQ((*maybeStep).size(), 2u);
+ std::optional<SmallVector<Value>> maybeInductionVars =
+ loopLikeOp.getLoopInductionVars();
+ ASSERT_TRUE(maybeInductionVars.has_value());
+ EXPECT_EQ((*maybeInductionVars).size(), 2u);
}
MLIRContext context;
More information about the Mlir-commits
mailing list