[Mlir-commits] [mlir] [mlir][loops] Add getters for multi dim loop variables in `LoopLikeOpInterface` (PR #94516)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jun 6 09:32:30 PDT 2024
https://github.com/srcarroll updated https://github.com/llvm/llvm-project/pull/94516
>From 5020e498440b0016adef7e99806aa55c4837b441 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 5 Jun 2024 13:08:22 -0500
Subject: [PATCH 1/4] Add getters for multi dim loop variables in
LoopLikeOpInterface
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 4 +-
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 37 ++--------
.../mlir/Interfaces/LoopLikeInterface.td | 65 +++++++++++------
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 20 +++---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 70 +++++++------------
.../Dialect/SCF/LoopLikeSCFOpsTest.cpp | 8 +++
6 files changed, 97 insertions(+), 107 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index 3640055ea8da8..bb2c29b5733b8 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",
+ ["getInductionVars", "getMixedLowerBound", "getMixedStep",
+ "getMixedUpperBound", "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..3b28ca8b21d0f 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",
+ "getInductionVars", "getMixedLowerBound", "getMixedStep",
+ "getMixedUpperBound", "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", "getInductionVars",
+ "getMixedLowerBound", "getMixedUpperBound", "getMixedStep",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,24 +510,6 @@ def ForallOp : SCF_Op<"forall", [
];
let extraClassDeclaration = [{
- // Get lower bounds as OpFoldResult.
- SmallVector<OpFoldResult> getMixedLowerBound() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
- }
-
- // Get upper bounds as OpFoldResult.
- SmallVector<OpFoldResult> getMixedUpperBound() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
- }
-
- // Get steps as OpFoldResult.
- SmallVector<OpFoldResult> getMixedStep() {
- Builder b(getOperation()->getContext());
- return getMixedValues(getStaticStep(), getDynamicStep(), b);
- }
-
/// Get lower bounds as values.
SmallVector<Value> getLowerBound(OpBuilder &b) {
return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedLowerBound());
@@ -584,10 +566,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 +743,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods<RegionBranchOpInterface, [
def ParallelOp : SCF_Op<"parallel",
[AutomaticAllocationScope,
AttrSizedOperandSegments,
- DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getSingleInductionVar",
- "getSingleLowerBound", "getSingleUpperBound", "getSingleStep"]>,
+ DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getInductionVars",
+ "getMixedLowerBound", "getMixedUpperBound", "getMixedStep"]>,
RecursiveMemoryEffects,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::ReduceOp">,
@@ -846,9 +824,6 @@ def ParallelOp : SCF_Op<"parallel",
];
let extraClassDeclaration = [{
- ValueRange getInductionVars() {
- return getBody()->getArguments();
- }
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..813779c852027 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -93,51 +93,47 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}]
>,
InterfaceMethod<[{
- If there is a single induction variable return it, otherwise return
- std::nullopt.
+ Return all induction variables.
}],
- /*retTy=*/"::std::optional<::mlir::Value>",
- /*methodName=*/"getSingleInductionVar",
+ /*retTy=*/"::mlir::ValueRange",
+ /*methodName=*/"getInductionVars",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return {};
}]
>,
InterfaceMethod<[{
- Return the single lower bound value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all lower bounds.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleLowerBound",
+ /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedLowerBound",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return {};
}]
>,
InterfaceMethod<[{
- Return the single step value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all steps.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleStep",
+ /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedStep",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return {};
}]
>,
InterfaceMethod<[{
- Return the single upper bound value or attribute if it exists, otherwise
- return std::nullopt.
+ Return all upper bounds.
}],
- /*retTy=*/"::std::optional<::mlir::OpFoldResult>",
- /*methodName=*/"getSingleUpperBound",
+ /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
+ /*methodName=*/"getMixedUpperBound",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return std::nullopt;
+ return {};
}]
>,
InterfaceMethod<[{
@@ -235,6 +231,35 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
}];
let extraSharedClassDeclaration = [{
+ /// If there is a single induction variable return it, otherwise return
+ /// std::nullopt.
+ ::std::optional<::mlir::Value> getSingleInductionVar() {
+ if (this->getInductionVars().size() == 1)
+ return this->getInductionVars()[0];
+ return std::nullopt;
+ }
+ /// Return the single lower bound value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
+ if (this->getMixedLowerBound().size() == 1)
+ return this->getMixedLowerBound()[0];
+ return std::nullopt;
+ }
+ /// Return the single step value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleStep() {
+ if (this->getMixedStep().size() == 1)
+ return this->getMixedStep()[0];
+ return std::nullopt;
+ }
+ /// Return the single upper bound value or attribute if it exists, otherwise
+ /// return std::nullopt.
+ ::std::optional<::mlir::OpFoldResult> getSingleUpperBound() {
+ if (this->getMixedUpperBound().size() == 1)
+ return this->getMixedUpperBound()[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..746a9c919560c 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2454,27 +2454,25 @@ bool AffineForOp::matchingBoundOperandList() {
SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
-std::optional<Value> AffineForOp::getSingleInductionVar() {
- return getInductionVar();
-}
+ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; }
-std::optional<OpFoldResult> AffineForOp::getSingleLowerBound() {
+SmallVector<OpFoldResult> AffineForOp::getMixedLowerBound() {
if (!hasConstantLowerBound())
- return std::nullopt;
+ return {};
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()));
+ return {OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
}
-std::optional<OpFoldResult> AffineForOp::getSingleStep() {
+SmallVector<OpFoldResult> AffineForOp::getMixedStep() {
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getStepAsInt()));
+ return {OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
}
-std::optional<OpFoldResult> AffineForOp::getSingleUpperBound() {
+SmallVector<OpFoldResult> AffineForOp::getMixedUpperBound() {
if (!hasConstantUpperBound())
- return std::nullopt;
+ return {};
OpBuilder b(getContext());
- return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()));
+ return {OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
}
FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 107fd0690f193..e275ff1849c10 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -378,20 +378,18 @@ LogicalResult ForOp::verifyRegions() {
return success();
}
-std::optional<Value> ForOp::getSingleInductionVar() {
- return getInductionVar();
-}
+ValueRange ForOp::getInductionVars() { return {getInductionVar()}; }
-std::optional<OpFoldResult> ForOp::getSingleLowerBound() {
- return OpFoldResult(getLowerBound());
+SmallVector<OpFoldResult> ForOp::getMixedLowerBound() {
+ return {OpFoldResult(getLowerBound())};
}
-std::optional<OpFoldResult> ForOp::getSingleStep() {
- return OpFoldResult(getStep());
+SmallVector<OpFoldResult> ForOp::getMixedStep() {
+ return {OpFoldResult(getStep())};
}
-std::optional<OpFoldResult> ForOp::getSingleUpperBound() {
- return OpFoldResult(getUpperBound());
+SmallVector<OpFoldResult> ForOp::getMixedUpperBound() {
+ return {OpFoldResult(getUpperBound())};
}
std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }
@@ -1428,28 +1426,26 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-std::optional<Value> ForallOp::getSingleInductionVar() {
- if (getRank() != 1)
- return std::nullopt;
- return getInductionVar(0);
+ValueRange ForallOp::getInductionVars() {
+ return getBody()->getArguments().take_front(getRank());
}
-std::optional<OpFoldResult> ForallOp::getSingleLowerBound() {
- if (getRank() != 1)
- return std::nullopt;
- return getMixedLowerBound()[0];
+// Get lower bounds as OpFoldResult.
+SmallVector<OpFoldResult> ForallOp::getMixedLowerBound() {
+ 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.
+SmallVector<OpFoldResult> ForallOp::getMixedUpperBound() {
+ 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.
+SmallVector<OpFoldResult> ForallOp::getMixedStep() {
+ Builder b(getOperation()->getContext());
+ return getMixedValues(getStaticStep(), getDynamicStep(), b);
}
ForallOp mlir::scf::getForallOpThreadIndexOwner(Value val) {
@@ -3008,29 +3004,17 @@ 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);
-}
+ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); }
-std::optional<OpFoldResult> ParallelOp::getSingleLowerBound() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getLowerBound()[0];
+SmallVector<OpFoldResult> ParallelOp::getMixedLowerBound() {
+ return getLowerBound();
}
-std::optional<OpFoldResult> ParallelOp::getSingleUpperBound() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getUpperBound()[0];
+SmallVector<OpFoldResult> ParallelOp::getMixedUpperBound() {
+ return getUpperBound();
}
-std::optional<OpFoldResult> ParallelOp::getSingleStep() {
- if (getNumLoops() != 1)
- return std::nullopt;
- return getStep()[0];
-}
+SmallVector<OpFoldResult> ParallelOp::getMixedStep() { return getStep(); }
ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
auto ivArg = llvm::dyn_cast<BlockArgument>(val);
diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
index 6bc0fd6113b9b..d8cdb213070da 100644
--- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
+++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -36,6 +36,10 @@ class SCFLoopLikeTest : public ::testing::Test {
std::optional<OpFoldResult> maybeIndVar =
loopLikeOp.getSingleInductionVar();
EXPECT_TRUE(maybeIndVar.has_value());
+ EXPECT_EQ(loopLikeOp.getInductionVars().size(), 1u);
+ EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 1u);
+ EXPECT_EQ(loopLikeOp.getMixedStep().size(), 1u);
+ EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 1u);
}
void checkMultidimensional(LoopLikeOpInterface loopLikeOp) {
@@ -48,6 +52,10 @@ class SCFLoopLikeTest : public ::testing::Test {
std::optional<OpFoldResult> maybeIndVar =
loopLikeOp.getSingleInductionVar();
EXPECT_FALSE(maybeIndVar.has_value());
+ EXPECT_EQ(loopLikeOp.getInductionVars().size(), 2u);
+ EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 2u);
+ EXPECT_EQ(loopLikeOp.getMixedStep().size(), 2u);
+ EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 2u);
}
MLIRContext context;
>From 7d995815064cb25e47ec8e400de3692fbe5fdfba Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 5 Jun 2024 14:37:57 -0500
Subject: [PATCH 2/4] address review comment
---
.../mlir/Interfaces/LoopLikeInterface.td | 20 +++++++++++--------
1 file changed, 12 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 813779c852027..5cf3eba0bd9ed 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -234,29 +234,33 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
/// If there is a single induction variable return it, otherwise return
/// std::nullopt.
::std::optional<::mlir::Value> getSingleInductionVar() {
- if (this->getInductionVars().size() == 1)
- return this->getInductionVars()[0];
+ auto inductionVars = this->getInductionVars();
+ if (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() {
- if (this->getMixedLowerBound().size() == 1)
- return this->getMixedLowerBound()[0];
+ auto lowerBounds = this->getMixedLowerBound();
+ if (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() {
- if (this->getMixedStep().size() == 1)
- return this->getMixedStep()[0];
+ auto steps = this->getMixedStep();
+ if (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() {
- if (this->getMixedUpperBound().size() == 1)
- return this->getMixedUpperBound()[0];
+ auto upperBounds = this->getMixedUpperBound();
+ if (upperBounds.size() == 1)
+ return upperBounds[0];
return std::nullopt;
}
>From a5fa3b3c4903c344847ee544cd9812b6f0c70571 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 5 Jun 2024 20:56:02 -0500
Subject: [PATCH 3/4] Make return types optional and change names
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 10 +--
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 29 +++++++--
.../mlir/Interfaces/LoopLikeInterface.td | 42 ++++++-------
.../AffineToStandard/AffineToStandard.cpp | 4 +-
.../SCFToControlFlow/SCFToControlFlow.cpp | 9 +--
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 27 ++++----
mlir/lib/Dialect/Affine/Utils/Utils.cpp | 2 +-
mlir/lib/Dialect/SCF/IR/SCF.cpp | 26 ++++----
.../Dialect/SCF/Transforms/ForallToFor.cpp | 9 +--
.../Dialect/SCF/LoopLikeSCFOpsTest.cpp | 62 +++++++++++++------
10 files changed, 131 insertions(+), 89 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index bb2c29b5733b8..4c032e66f7a83 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,
- ["getInductionVars", "getMixedLowerBound", "getMixedStep",
- "getMixedUpperBound", "getYieldedValuesMutable",
+ ["getInductionVars", "getLowerBounds", "getSteps",
+ "getUpperBounds", "getYieldedValuesMutable",
"replaceWithAdditionalYields"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>]> {
@@ -671,7 +671,7 @@ def AffineParallelOp : Affine_Op<"parallel",
I32ElementsAttr:$lowerBoundsGroups,
AffineMapAttr:$upperBoundsMap,
I32ElementsAttr:$upperBoundsGroups,
- I64SmallVectorArrayAttr:$steps,
+ I64SmallVectorArrayAttr:$step,
Variadic<Index>:$mapOperands);
let results = (outs Variadic<AnyType>:$results);
let regions = (region SizedRegion<1>:$region);
@@ -682,7 +682,7 @@ def AffineParallelOp : Affine_Op<"parallel",
OpBuilder<(ins "TypeRange":$resultTypes,
"ArrayRef<arith::AtomicRMWKind>":$reductions, "ArrayRef<AffineMap>":$lbMaps,
"ValueRange":$lbArgs, "ArrayRef<AffineMap>":$ubMaps, "ValueRange":$ubArgs,
- "ArrayRef<int64_t>":$steps)>
+ "ArrayRef<int64_t>":$step)>
];
let extraClassDeclaration = [{
@@ -727,7 +727,7 @@ def AffineParallelOp : Affine_Op<"parallel",
static StringRef getUpperBoundsGroupsAttrStrName() {
return "upperBoundsGroups";
}
- static StringRef getStepsAttrStrName() { return "steps"; }
+ static StringRef getStepsAttrStrName() { return "step"; }
/// Returns `true` if the loop bounds have min/max expressions.
bool hasMinMaxBounds() {
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 3b28ca8b21d0f..66b478f141b32 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",
- "getInductionVars", "getMixedLowerBound", "getMixedStep",
- "getMixedUpperBound", "getYieldedValuesMutable",
+ "getInductionVars", "getLowerBounds", "getSteps",
+ "getUpperBounds", "getYieldedValuesMutable",
"promoteIfSingleIteration", "replaceWithAdditionalYields",
"yieldTiledValuesAndReplace"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
@@ -302,7 +302,7 @@ def ForallOp : SCF_Op<"forall", [
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInitsMutable", "getRegionIterArgs", "getInductionVars",
- "getMixedLowerBound", "getMixedUpperBound", "getMixedStep",
+ "getLowerBounds", "getUpperBounds", "getSteps",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,6 +510,27 @@ def ForallOp : SCF_Op<"forall", [
];
let extraClassDeclaration = [{
+ // Get lower bounds as OpFoldResult.
+ SmallVector<OpFoldResult> getMixedLowerBound() {
+ auto maybeLowerBounds = getLowerBounds();
+ assert(maybeLowerBounds.has_value() && "expected values");
+ return *maybeLowerBounds;
+ }
+
+ // Get upper bounds as OpFoldResult.
+ SmallVector<OpFoldResult> getMixedUpperBound() {
+ auto maybeUpperBounds = getUpperBounds();
+ assert(maybeUpperBounds.has_value() && "expected values");
+ return *maybeUpperBounds;
+ }
+
+ // Get steps as OpFoldResult.
+ SmallVector<OpFoldResult> getMixedStep() {
+ auto maybeSteps = getSteps();
+ assert(maybeSteps.has_value() && "expected values");
+ return *maybeSteps;
+ }
+
/// Get lower bounds as values.
SmallVector<Value> getLowerBound(OpBuilder &b) {
return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedLowerBound());
@@ -744,7 +765,7 @@ def ParallelOp : SCF_Op<"parallel",
[AutomaticAllocationScope,
AttrSizedOperandSegments,
DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getInductionVars",
- "getMixedLowerBound", "getMixedUpperBound", "getMixedStep"]>,
+ "getLowerBounds", "getUpperBounds", "getSteps"]>,
RecursiveMemoryEffects,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::ReduceOp">,
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 5cf3eba0bd9ed..cc79d026c8d4e 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -106,34 +106,34 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
InterfaceMethod<[{
Return all lower bounds.
}],
- /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
- /*methodName=*/"getMixedLowerBound",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getLowerBounds",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return {};
+ return std::nullopt;
}]
>,
InterfaceMethod<[{
Return all steps.
}],
- /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
- /*methodName=*/"getMixedStep",
+ /*retTy=*/"std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getSteps",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return {};
+ return std::nullopt;
}]
>,
InterfaceMethod<[{
Return all upper bounds.
}],
- /*retTy=*/"::llvm::SmallVector<::mlir::OpFoldResult>",
- /*methodName=*/"getMixedUpperBound",
+ /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>",
+ /*methodName=*/"getUpperBounds",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/[{
- return {};
+ return std::nullopt;
}]
>,
InterfaceMethod<[{
@@ -242,26 +242,26 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
/// Return the single lower bound value or attribute if it exists, otherwise
/// return std::nullopt.
::std::optional<::mlir::OpFoldResult> getSingleLowerBound() {
- auto lowerBounds = this->getMixedLowerBound();
- if (lowerBounds.size() == 1)
- return lowerBounds[0];
- return std::nullopt;
+ auto lowerBounds = this->getLowerBounds();
+ 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->getMixedStep();
- if (steps.size() == 1)
- return steps[0];
- return std::nullopt;
+ auto steps = this->getSteps();
+ 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->getMixedUpperBound();
- if (upperBounds.size() == 1)
- return upperBounds[0];
- return std::nullopt;
+ auto upperBounds = this->getUpperBounds();
+ if (upperBounds.has_value() && (*upperBounds).size() == 1)
+ return (*upperBounds)[0];
+ return std::nullopt;
}
/// Append the specified additional "init" operands: replace this loop with
diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
index 10ccd5c97783b..20487b32e3fe0 100644
--- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
+++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
@@ -196,8 +196,8 @@ class AffineParallelLowering : public OpRewritePattern<AffineParallelOp> {
return rewriter.notifyMatchFailure(op, "couldn't convert upper bounds");
upperBoundTuple.push_back(upper);
}
- steps.reserve(op.getSteps().size());
- for (int64_t step : op.getSteps())
+ steps.reserve(op.getStep().size());
+ for (int64_t step : op.getStep())
steps.push_back(rewriter.create<arith::ConstantIndexOp>(loc, step));
// Get the terminator op.
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 9eb8a289d7d65..48e1d88c1c75e 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -695,12 +695,9 @@ LogicalResult ForallLowering::matchAndRewrite(ForallOp forallOp,
"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<ParallelOp>(loc, lbs, ubs, steps);
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 746a9c919560c..d3f034a0660ba 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2456,23 +2456,26 @@ SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; }
-SmallVector<OpFoldResult> AffineForOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLowerBounds() {
if (!hasConstantLowerBound())
- return {};
+ return std::nullopt;
OpBuilder b(getContext());
- return {OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
}
-SmallVector<OpFoldResult> AffineForOp::getMixedStep() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getSteps() {
OpBuilder b(getContext());
- return {OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
}
-SmallVector<OpFoldResult> AffineForOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getUpperBounds() {
if (!hasConstantUpperBound())
return {};
OpBuilder b(getContext());
- return {OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
+ return SmallVector<OpFoldResult>{
+ OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))};
}
FailureOr<LoopLikeOpInterface> AffineForOp::replaceWithAdditionalYields(
@@ -3753,7 +3756,7 @@ SmallVector<Region *> AffineParallelOp::getLoopRegions() {
return {&getRegion()};
}
-unsigned AffineParallelOp::getNumDims() { return getSteps().size(); }
+unsigned AffineParallelOp::getNumDims() { return getStep().size(); }
AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() {
return getOperands().take_front(getLowerBoundsMap().getNumInputs());
@@ -3838,7 +3841,7 @@ void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) {
}
void AffineParallelOp::setSteps(ArrayRef<int64_t> newSteps) {
- setStepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
+ setStepAttr(getBodyBuilder().getI64ArrayAttr(newSteps));
}
// check whether resultType match op or not in affine.parallel
@@ -3888,14 +3891,14 @@ LogicalResult AffineParallelOp::verify() {
auto numDims = getNumDims();
if (getLowerBoundsGroups().getNumElements() != numDims ||
getUpperBoundsGroups().getNumElements() != numDims ||
- getSteps().size() != numDims || getBody()->getNumArguments() != numDims) {
+ getStep().size() != numDims || getBody()->getNumArguments() != numDims) {
return emitOpError() << "the number of region arguments ("
<< getBody()->getNumArguments()
<< ") and the number of map groups for lower ("
<< getLowerBoundsGroups().getNumElements()
<< ") and upper bound ("
<< getUpperBoundsGroups().getNumElements()
- << "), and the number of steps (" << getSteps().size()
+ << "), and the number of steps (" << getStep().size()
<< ") must all match";
}
@@ -4013,7 +4016,7 @@ void AffineParallelOp::print(OpAsmPrinter &p) {
printMinMaxBound(p, getUpperBoundsMapAttr(), getUpperBoundsGroupsAttr(),
getUpperBoundsOperands(), "min");
p << ')';
- SmallVector<int64_t, 8> steps = getSteps();
+ SmallVector<int64_t, 8> steps = getStep();
bool elideSteps = llvm::all_of(steps, [](int64_t step) { return step == 1; });
if (!elideSteps) {
p << " step (";
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index f46381403bc52..a652ee4a488d1 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -494,7 +494,7 @@ void mlir::affine::normalizeAffineParallel(AffineParallelOp op) {
return;
AffineMap lbMap = op.getLowerBoundsMap();
- SmallVector<int64_t, 8> steps = op.getSteps();
+ SmallVector<int64_t, 8> steps = op.getStep();
// No need to do any work if the parallel op is already normalized.
bool isAlreadyNormalized =
llvm::all_of(llvm::zip(steps, lbMap.getResults()), [](auto tuple) {
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index e275ff1849c10..281d73afee4a8 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -380,16 +380,16 @@ LogicalResult ForOp::verifyRegions() {
ValueRange ForOp::getInductionVars() { return {getInductionVar()}; }
-SmallVector<OpFoldResult> ForOp::getMixedLowerBound() {
- return {OpFoldResult(getLowerBound())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getLowerBounds() {
+ return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
}
-SmallVector<OpFoldResult> ForOp::getMixedStep() {
- return {OpFoldResult(getStep())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getSteps() {
+ return SmallVector<OpFoldResult, 1>{OpFoldResult(getStep())};
}
-SmallVector<OpFoldResult> ForOp::getMixedUpperBound() {
- return {OpFoldResult(getUpperBound())};
+std::optional<SmallVector<OpFoldResult>> ForOp::getUpperBounds() {
+ return SmallVector<OpFoldResult, 1>{OpFoldResult(getUpperBound())};
}
std::optional<ResultRange> ForOp::getLoopResults() { return getResults(); }
@@ -1431,19 +1431,19 @@ ValueRange ForallOp::getInductionVars() {
}
// Get lower bounds as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLowerBounds() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
}
// Get upper bounds as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getUpperBounds() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
}
// Get steps as OpFoldResult.
-SmallVector<OpFoldResult> ForallOp::getMixedStep() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getSteps() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticStep(), getDynamicStep(), b);
}
@@ -3006,15 +3006,17 @@ SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); }
-SmallVector<OpFoldResult> ParallelOp::getMixedLowerBound() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLowerBounds() {
return getLowerBound();
}
-SmallVector<OpFoldResult> ParallelOp::getMixedUpperBound() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getUpperBounds() {
return getUpperBound();
}
-SmallVector<OpFoldResult> ParallelOp::getMixedStep() { return getStep(); }
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getSteps() {
+ return getStep();
+}
ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) {
auto ivArg = llvm::dyn_cast<BlockArgument>(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/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
index d8cdb213070da..07504a99fecd3 100644
--- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
+++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp
@@ -27,35 +27,57 @@ class SCFLoopLikeTest : public ::testing::Test {
}
void checkUnidimensional(LoopLikeOpInterface loopLikeOp) {
- std::optional<OpFoldResult> maybeLb = loopLikeOp.getSingleLowerBound();
+ 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(maybeSingleIndVar.has_value());
+
+ std::optional<SmallVector<OpFoldResult>> maybeLb =
+ loopLikeOp.getLowerBounds();
EXPECT_TRUE(maybeLb.has_value());
- std::optional<OpFoldResult> maybeUb = loopLikeOp.getSingleUpperBound();
+ EXPECT_EQ((*maybeLb).size(), 1u);
+ std::optional<SmallVector<OpFoldResult>> maybeUb =
+ loopLikeOp.getUpperBounds();
EXPECT_TRUE(maybeUb.has_value());
- std::optional<OpFoldResult> maybeStep = loopLikeOp.getSingleStep();
+ EXPECT_EQ((*maybeUb).size(), 1u);
+ std::optional<SmallVector<OpFoldResult>> maybeStep = loopLikeOp.getSteps();
EXPECT_TRUE(maybeStep.has_value());
- std::optional<OpFoldResult> maybeIndVar =
- loopLikeOp.getSingleInductionVar();
- EXPECT_TRUE(maybeIndVar.has_value());
+ EXPECT_EQ((*maybeStep).size(), 1u);
EXPECT_EQ(loopLikeOp.getInductionVars().size(), 1u);
- EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 1u);
- EXPECT_EQ(loopLikeOp.getMixedStep().size(), 1u);
- EXPECT_EQ(loopLikeOp.getMixedLowerBound().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.getLowerBounds();
+ EXPECT_TRUE(maybeLb.has_value());
+ EXPECT_EQ((*maybeLb).size(), 2u);
+ std::optional<SmallVector<OpFoldResult>> maybeUb =
+ loopLikeOp.getUpperBounds();
+ EXPECT_TRUE(maybeUb.has_value());
+ EXPECT_EQ((*maybeUb).size(), 2u);
+ std::optional<SmallVector<OpFoldResult>> maybeStep = loopLikeOp.getSteps();
+ EXPECT_TRUE(maybeStep.has_value());
+ EXPECT_EQ((*maybeStep).size(), 2u);
EXPECT_EQ(loopLikeOp.getInductionVars().size(), 2u);
- EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 2u);
- EXPECT_EQ(loopLikeOp.getMixedStep().size(), 2u);
- EXPECT_EQ(loopLikeOp.getMixedLowerBound().size(), 2u);
}
MLIRContext context;
>From 1babe681d7858a4992303c62e22684cb73d82472 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 6 Jun 2024 11:31:11 -0500
Subject: [PATCH 4/4] change return type of getInductionVars to
SmallVector<Value>
---
mlir/include/mlir/Interfaces/LoopLikeInterface.td | 2 +-
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 4 +++-
mlir/lib/Dialect/Linalg/Transforms/Loops.cpp | 3 +--
mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp | 6 +++---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 10 ++++++----
5 files changed, 14 insertions(+), 11 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index cc79d026c8d4e..bace8f8384d44 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -95,7 +95,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
InterfaceMethod<[{
Return all induction variables.
}],
- /*retTy=*/"::mlir::ValueRange",
+ /*retTy=*/"::llvm::SmallVector<::mlir::Value>",
/*methodName=*/"getInductionVars",
/*args=*/(ins),
/*methodBody=*/"",
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index d3f034a0660ba..5467c60242664 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2454,7 +2454,9 @@ bool AffineForOp::matchingBoundOperandList() {
SmallVector<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
-ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; }
+SmallVector<Value> AffineForOp::getInductionVars() {
+ return {getInductionVar()};
+}
std::optional<SmallVector<OpFoldResult>> AffineForOp::getLowerBounds() {
if (!hasConstantLowerBound())
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..4eacaa8d1e327 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();
+ auto 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 281d73afee4a8..0ce10ebdad3e2 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -378,7 +378,7 @@ LogicalResult ForOp::verifyRegions() {
return success();
}
-ValueRange ForOp::getInductionVars() { return {getInductionVar()}; }
+SmallVector<Value> ForOp::getInductionVars() { return {getInductionVar()}; }
std::optional<SmallVector<OpFoldResult>> ForOp::getLowerBounds() {
return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
@@ -1426,8 +1426,8 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-ValueRange ForallOp::getInductionVars() {
- return getBody()->getArguments().take_front(getRank());
+SmallVector<Value> ForallOp::getInductionVars() {
+ return SmallVector<Value>(getBody()->getArguments().take_front(getRank()));
}
// Get lower bounds as OpFoldResult.
@@ -3004,7 +3004,9 @@ void ParallelOp::print(OpAsmPrinter &p) {
SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
-ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); }
+SmallVector<Value> ParallelOp::getInductionVars() {
+ return SmallVector<Value>(getBody()->getArguments());
+}
std::optional<SmallVector<OpFoldResult>> ParallelOp::getLowerBounds() {
return getLowerBound();
More information about the Mlir-commits
mailing list