Skip to content

Commit 6b4c122

Browse files
srcarrollsrcarroll
authored
[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.
1 parent 33f4a77 commit 6b4c122

File tree

10 files changed

+198
-125
lines changed

10 files changed

+198
-125
lines changed

‎mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

+2-2
Original file line numberDiff line numberDiff line change
@@ -118,8 +118,8 @@ def AffineForOp : Affine_Op<"for",
118118
[AttrSizedOperandSegments, AutomaticAllocationScope,
119119
ImplicitAffineTerminator, ConditionallySpeculatable,
120120
RecursiveMemoryEffects, DeclareOpInterfaceMethods<LoopLikeOpInterface,
121-
["getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
122-
"getSingleUpperBound", "getYieldedValuesMutable",
121+
["getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
122+
"getLoopUpperBounds", "getYieldedValuesMutable",
123123
"replaceWithAdditionalYields"]>,
124124
DeclareOpInterfaceMethods<RegionBranchOpInterface,
125125
["getEntrySuccessorOperands"]>]> {

‎mlir/include/mlir/Dialect/SCF/IR/SCFOps.td

+29-21
Original file line numberDiff line numberDiff line change
@@ -136,8 +136,8 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
136136
def ForOp : SCF_Op<"for",
137137
[AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
138138
["getInitsMutable", "getLoopResults", "getRegionIterArgs",
139-
"getSingleInductionVar", "getSingleLowerBound", "getSingleStep",
140-
"getSingleUpperBound", "getYieldedValuesMutable",
139+
"getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
140+
"getLoopUpperBounds", "getYieldedValuesMutable",
141141
"promoteIfSingleIteration", "replaceWithAdditionalYields",
142142
"yieldTiledValuesAndReplace"]>,
143143
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
@@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [
301301
AttrSizedOperandSegments,
302302
AutomaticAllocationScope,
303303
DeclareOpInterfaceMethods<LoopLikeOpInterface,
304-
["getInitsMutable", "getRegionIterArgs", "getSingleInductionVar",
305-
"getSingleLowerBound", "getSingleUpperBound", "getSingleStep",
304+
["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
305+
"getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
306306
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
307307
RecursiveMemoryEffects,
308308
SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,22 +510,31 @@ def ForallOp : SCF_Op<"forall", [
510510
];
511511

512512
let extraClassDeclaration = [{
513-
// Get lower bounds as OpFoldResult.
513+
/// Get induction variables.
514+
SmallVector<Value> getInductionVars() {
515+
std::optional<SmallVector<Value>> maybeInductionVars = getLoopInductionVars();
516+
assert(maybeInductionVars.has_value() && "expected values");
517+
return *maybeInductionVars;
518+
}
519+
/// Get lower bounds as OpFoldResult.
514520
SmallVector<OpFoldResult> getMixedLowerBound() {
515-
Builder b(getOperation()->getContext());
516-
return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
521+
std::optional<SmallVector<OpFoldResult>> maybeLowerBounds = getLoopLowerBounds();
522+
assert(maybeLowerBounds.has_value() && "expected values");
523+
return *maybeLowerBounds;
517524
}
518525

519-
// Get upper bounds as OpFoldResult.
526+
/// Get upper bounds as OpFoldResult.
520527
SmallVector<OpFoldResult> getMixedUpperBound() {
521-
Builder b(getOperation()->getContext());
522-
return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
528+
std::optional<SmallVector<OpFoldResult>> maybeUpperBounds = getLoopUpperBounds();
529+
assert(maybeUpperBounds.has_value() && "expected values");
530+
return *maybeUpperBounds;
523531
}
524532

525-
// Get steps as OpFoldResult.
533+
/// Get steps as OpFoldResult.
526534
SmallVector<OpFoldResult> getMixedStep() {
527-
Builder b(getOperation()->getContext());
528-
return getMixedValues(getStaticStep(), getDynamicStep(), b);
535+
std::optional<SmallVector<OpFoldResult>> maybeSteps = getLoopSteps();
536+
assert(maybeSteps.has_value() && "expected values");
537+
return *maybeSteps;
529538
}
530539

531540
/// Get lower bounds as values.
@@ -584,10 +593,6 @@ def ForallOp : SCF_Op<"forall", [
584593
getNumDynamicControlOperands() + getRank());
585594
}
586595

587-
::mlir::ValueRange getInductionVars() {
588-
return getBody()->getArguments().take_front(getRank());
589-
}
590-
591596
::mlir::Value getInductionVar(int64_t idx) {
592597
return getInductionVars()[idx];
593598
}
@@ -765,8 +770,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
765770
def ParallelOp : SCF_Op<"parallel",
766771
[AutomaticAllocationScope,
767772
AttrSizedOperandSegments,
768-
DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getSingleInductionVar",
769-
"getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
773+
DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getLoopInductionVars",
774+
"getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps"]>,
770775
RecursiveMemoryEffects,
771776
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
772777
SingleBlockImplicitTerminator<"scf::ReduceOp">,
@@ -846,8 +851,11 @@ def ParallelOp : SCF_Op<"parallel",
846851
];
847852

848853
let extraClassDeclaration = [{
849-
ValueRange getInductionVars() {
850-
return getBody()->getArguments();
854+
/// Get induction variables.
855+
SmallVector<Value> getInductionVars() {
856+
std::optional<SmallVector<Value>> maybeInductionVars = getLoopInductionVars();;
857+
assert(maybeInductionVars.has_value() && "expected values");
858+
return *maybeInductionVars;
851859
}
852860
unsigned getNumLoops() { return getStep().size(); }
853861
unsigned getNumReductions() { return getInitVals().size(); }

‎mlir/include/mlir/Interfaces/LoopLikeInterface.td

+61-20
Original file line numberDiff line numberDiff line change
@@ -93,51 +93,59 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
9393
}]
9494
>,
9595
InterfaceMethod<[{
96-
If there is a single induction variable return it, otherwise return
97-
std::nullopt.
96+
Return all induction variables, if they exist. If the op has no notion of
97+
induction variable, then return std::nullopt. If it does have
98+
a notion but an instance doesn't have induction variables, then
99+
return empty vector.
98100
}],
99-
/*retTy=*/"::std::optional<::mlir::Value>",
100-
/*methodName=*/"getSingleInductionVar",
101+
/*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::Value>>",
102+
/*methodName=*/"getLoopInductionVars",
101103
/*args=*/(ins),
102104
/*methodBody=*/"",
103105
/*defaultImplementation=*/[{
104-
return std::nullopt;
106+
return ::std::nullopt;
105107
}]
106108
>,
107109
InterfaceMethod<[{
108-
Return the single lower bound value or attribute if it exists, otherwise
109-
return std::nullopt.
110+
Return all lower bounds, if they exist. If the op has no notion of
111+
lower bounds, then return std::nullopt. If it does have
112+
a notion but an instance doesn't have lower bounds, then
113+
return empty vector.
110114
}],
111-
/*retTy=*/"::std::optional<::mlir::OpFoldResult>",
112-
/*methodName=*/"getSingleLowerBound",
115+
/*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
116+
/*methodName=*/"getLoopLowerBounds",
113117
/*args=*/(ins),
114118
/*methodBody=*/"",
115119
/*defaultImplementation=*/[{
116-
return std::nullopt;
120+
return ::std::nullopt;
117121
}]
118122
>,
119123
InterfaceMethod<[{
120-
Return the single step value or attribute if it exists, otherwise
121-
return std::nullopt.
124+
Return all steps, if they exist. If the op has no notion of
125+
steps, then return std::nullopt. If it does have
126+
a notion but an instance doesn't have steps, then
127+
return empty vector.
122128
}],
123-
/*retTy=*/"::std::optional<::mlir::OpFoldResult>",
124-
/*methodName=*/"getSingleStep",
129+
/*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
130+
/*methodName=*/"getLoopSteps",
125131
/*args=*/(ins),
126132
/*methodBody=*/"",
127133
/*defaultImplementation=*/[{
128-
return std::nullopt;
134+
return ::std::nullopt;
129135
}]
130136
>,
131137
InterfaceMethod<[{
132-
Return the single upper bound value or attribute if it exists, otherwise
133-
return std::nullopt.
138+
Return all upper bounds, if they exist. If the op has no notion of
139+
lower bounds, then return std::nullopt. If it does have
140+
a notion but an instance doesn't have lower bounds, then
141+
return empty vector.
134142
}],
135-
/*retTy=*/"::std::optional<::mlir::OpFoldResult>",
136-
/*methodName=*/"getSingleUpperBound",
143+
/*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
144+
/*methodName=*/"getLoopUpperBounds",
137145
/*args=*/(ins),
138146
/*methodBody=*/"",
139147
/*defaultImplementation=*/[{
140-
return std::nullopt;
148+
return ::std::nullopt;
141149
}]
142150
>,
143151
InterfaceMethod<[{
@@ -235,6 +243,39 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
235243
}];
236244

237245
let extraSharedClassDeclaration = [{
246+
/// If there is a single induction variable return it, otherwise return
247+
/// std::nullopt.
248+
::std::optional<::mlir::Value> getSingleInductionVar() {
249+
auto inductionVars = this->getLoopInductionVars();
250+
if (inductionVars.has_value() && (*inductionVars).size() == 1)
251+
return (*inductionVars)[0];
252+
return std::nullopt;
253+
}
254+
/// Return the single lower bound value or attribute if it exists, otherwise
255+
/// return std::nullopt.
256+
::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
257+
auto lowerBounds = this->getLoopLowerBounds();
258+
if (lowerBounds.has_value() && (*lowerBounds).size() == 1)
259+
return (*lowerBounds)[0];
260+
return std::nullopt;
261+
}
262+
/// Return the single step value or attribute if it exists, otherwise
263+
/// return std::nullopt.
264+
::std::optional<::mlir::OpFoldResult> getSingleStep() {
265+
auto steps = this->getLoopSteps();
266+
if (steps.has_value() && (*steps).size() == 1)
267+
return (*steps)[0];
268+
return std::nullopt;
269+
}
270+
/// Return the single upper bound value or attribute if it exists, otherwise
271+
/// return std::nullopt.
272+
::std::optional<::mlir::OpFoldResult> getSingleUpperBound() {
273+
auto upperBounds = this->getLoopUpperBounds();
274+
if (upperBounds.has_value() && (*upperBounds).size() == 1)
275+
return (*upperBounds)[0];
276+
return std::nullopt;
277+
}
278+
238279
/// Append the specified additional "init" operands: replace this loop with
239280
/// a new loop that has the additional init operands. The loop body of this
240281
/// loop is moved over to the new loop.

‎mlir/lib/Dialect/Affine/IR/AffineOps.cpp

+12-9
Original file line numberDiff line numberDiff line change
@@ -2454,27 +2454,30 @@ bool AffineForOp::matchingBoundOperandList() {
24542454

24552455
SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
24562456

2457-
std::optional<Value> AffineForOp::getSingleInductionVar() {
2458-
return getInductionVar();
2457+
std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() {
2458+
return SmallVector<Value>{getInductionVar()};
24592459
}
24602460

2461-
std::optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
2461+
std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() {
24622462
if (!hasConstantLowerBound())
24632463
return std::nullopt;
24642464
OpBuilder b(getContext());
2465-
return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()));
2465+
return SmallVector<OpFoldResult>{
2466+
OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
24662467
}
24672468

2468-
std::optional<OpFoldResult> AffineForOp::getSingleStep() {
2469+
std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopSteps() {
24692470
OpBuilder b(getContext());
2470-
return OpFoldResult(b.getI64IntegerAttr(getStepAsInt()));
2471+
return SmallVector<OpFoldResult>{
2472+
OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
24712473
}
24722474

2473-
std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
2475+
std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() {
24742476
if (!hasConstantUpperBound())
2475-
return std::nullopt;
2477+
return {};
24762478
OpBuilder b(getContext());
2477-
return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
2479+
return SmallVector<OpFoldResult>{
2480+
OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
24782481
}
24792482

24802483
FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(

‎mlir/lib/Dialect/Linalg/Transforms/Loops.cpp

+1-2
Original file line numberDiff line numberDiff line change
@@ -184,8 +184,7 @@ static void replaceIndexOpsByInductionVariables(RewriterBase &rewriter,
184184
for (Operation *loopOp : loopOps) {
185185
llvm::TypeSwitch<Operation *>(loopOp)
186186
.Case([&](scf::ParallelOp parallelOp) {
187-
allIvs.append(parallelOp.getInductionVars().begin(),
188-
parallelOp.getInductionVars().end());
187+
allIvs.append(parallelOp.getInductionVars());
189188
})
190189
.Case([&](scf::ForOp forOp) {
191190
allIvs.push_back(forOp.getInductionVar());

‎mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

+3-3
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,7 @@ static void calculateTileOffsetsAndSizes(
243243
OpBuilder::InsertionGuard g(b);
244244
b.setInsertionPointToStart(forallOp.getBody(0));
245245

246-
ValueRange threadIds = forallOp.getInductionVars();
246+
SmallVector<Value> threadIds = forallOp.getInductionVars();
247247
SmallVector<OpFoldResult> nonZeroNumThreads =
248248
llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
249249
return !isConstantIntValue(ofr, 0);
@@ -746,7 +746,7 @@ FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
746746
b.getIndexAttr(0));
747747
SmallVector<OpFoldResult> sizes = tiledSizes;
748748
sizes[reductionDim] = b.getIndexAttr(1);
749-
outOffsets[reductionDim] = forallOp.getInductionVars().front();
749+
outOffsets[reductionDim] = forallOp.getInductionVars()[0];
750750
// TODO: use SubsetExtractOpInterface once it is available.
751751
tiledDpsInitOperands.push_back(b.create<tensor::ExtractSliceOp>(
752752
loc, cast<RankedTensorType>(initOperand.getType()),
@@ -814,7 +814,7 @@ FailureOr<linalg::ForallReductionTilingResult> linalg::tileReductionUsingForall(
814814
int64_t sizeIdx = 0;
815815
for (int64_t i = 0, e = numThreads.size(); i < e; ++i) {
816816
if (i == reductionDim) {
817-
resultOffsetsRank.push_back(forallOp.getInductionVars().front());
817+
resultOffsetsRank.push_back(forallOp.getInductionVars()[0]);
818818
resultSizesRank.push_back(b.getIndexAttr(1));
819819
continue;
820820
}

0 commit comments

Comments
 (0)