From 5020e498440b0016adef7e99806aa55c4837b441 Mon Sep 17 00:00:00 2001 From: Sam Date: Wed, 5 Jun 2024 13:08:22 -0500 Subject: [PATCH 1/8] 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, DeclareOpInterfaceMethods]> { 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, AllTypesMatch<["lowerBound", "upperBound", "step"]>, @@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [ AttrSizedOperandSegments, AutomaticAllocationScope, DeclareOpInterfaceMethods, RecursiveMemoryEffects, SingleBlockImplicitTerminator<"scf::InParallelOp">, @@ -510,24 +510,6 @@ def ForallOp : SCF_Op<"forall", [ ]; let extraClassDeclaration = [{ - // Get lower bounds as OpFoldResult. - SmallVector getMixedLowerBound() { - Builder b(getOperation()->getContext()); - return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b); - } - - // Get upper bounds as OpFoldResult. - SmallVector getMixedUpperBound() { - Builder b(getOperation()->getContext()); - return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b); - } - - // Get steps as OpFoldResult. - SmallVector getMixedStep() { - Builder b(getOperation()->getContext()); - return getMixedValues(getStaticStep(), getDynamicStep(), b); - } - /// Get lower bounds as values. SmallVector 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, + DeclareOpInterfaceMethods, RecursiveMemoryEffects, DeclareOpInterfaceMethods, 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 AffineForOp::getLoopRegions() { return {&getRegion()}; } -std::optional AffineForOp::getSingleInductionVar() { - return getInductionVar(); -} +ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; } -std::optional AffineForOp::getSingleLowerBound() { +SmallVector AffineForOp::getMixedLowerBound() { if (!hasConstantLowerBound()) - return std::nullopt; + return {}; OpBuilder b(getContext()); - return OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound())); + return {OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))}; } -std::optional AffineForOp::getSingleStep() { +SmallVector AffineForOp::getMixedStep() { OpBuilder b(getContext()); - return OpFoldResult(b.getI64IntegerAttr(getStepAsInt())); + return {OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))}; } -std::optional AffineForOp::getSingleUpperBound() { +SmallVector AffineForOp::getMixedUpperBound() { if (!hasConstantUpperBound()) - return std::nullopt; + return {}; OpBuilder b(getContext()); - return OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound())); + return {OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))}; } FailureOr 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 ForOp::getSingleInductionVar() { - return getInductionVar(); -} +ValueRange ForOp::getInductionVars() { return {getInductionVar()}; } -std::optional ForOp::getSingleLowerBound() { - return OpFoldResult(getLowerBound()); +SmallVector ForOp::getMixedLowerBound() { + return {OpFoldResult(getLowerBound())}; } -std::optional ForOp::getSingleStep() { - return OpFoldResult(getStep()); +SmallVector ForOp::getMixedStep() { + return {OpFoldResult(getStep())}; } -std::optional ForOp::getSingleUpperBound() { - return OpFoldResult(getUpperBound()); +SmallVector ForOp::getMixedUpperBound() { + return {OpFoldResult(getUpperBound())}; } std::optional ForOp::getLoopResults() { return getResults(); } @@ -1428,28 +1426,26 @@ SmallVector ForallOp::getCombiningOps(BlockArgument bbArg) { return storeOps; } -std::optional ForallOp::getSingleInductionVar() { - if (getRank() != 1) - return std::nullopt; - return getInductionVar(0); +ValueRange ForallOp::getInductionVars() { + return getBody()->getArguments().take_front(getRank()); } -std::optional ForallOp::getSingleLowerBound() { - if (getRank() != 1) - return std::nullopt; - return getMixedLowerBound()[0]; +// Get lower bounds as OpFoldResult. +SmallVector ForallOp::getMixedLowerBound() { + Builder b(getOperation()->getContext()); + return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b); } -std::optional ForallOp::getSingleUpperBound() { - if (getRank() != 1) - return std::nullopt; - return getMixedUpperBound()[0]; +// Get upper bounds as OpFoldResult. +SmallVector ForallOp::getMixedUpperBound() { + Builder b(getOperation()->getContext()); + return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b); } -std::optional ForallOp::getSingleStep() { - if (getRank() != 1) - return std::nullopt; - return getMixedStep()[0]; +// Get steps as OpFoldResult. +SmallVector 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 ParallelOp::getLoopRegions() { return {&getRegion()}; } -std::optional ParallelOp::getSingleInductionVar() { - if (getNumLoops() != 1) - return std::nullopt; - return getBody()->getArgument(0); -} +ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); } -std::optional ParallelOp::getSingleLowerBound() { - if (getNumLoops() != 1) - return std::nullopt; - return getLowerBound()[0]; +SmallVector ParallelOp::getMixedLowerBound() { + return getLowerBound(); } -std::optional ParallelOp::getSingleUpperBound() { - if (getNumLoops() != 1) - return std::nullopt; - return getUpperBound()[0]; +SmallVector ParallelOp::getMixedUpperBound() { + return getUpperBound(); } -std::optional ParallelOp::getSingleStep() { - if (getNumLoops() != 1) - return std::nullopt; - return getStep()[0]; -} +SmallVector ParallelOp::getMixedStep() { return getStep(); } ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) { auto ivArg = llvm::dyn_cast(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 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 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 Date: Wed, 5 Jun 2024 14:37:57 -0500 Subject: [PATCH 2/8] 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 Date: Wed, 5 Jun 2024 20:56:02 -0500 Subject: [PATCH 3/8] 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, DeclareOpInterfaceMethods]> { @@ -671,7 +671,7 @@ def AffineParallelOp : Affine_Op<"parallel", I32ElementsAttr:$lowerBoundsGroups, AffineMapAttr:$upperBoundsMap, I32ElementsAttr:$upperBoundsGroups, - I64SmallVectorArrayAttr:$steps, + I64SmallVectorArrayAttr:$step, Variadic:$mapOperands); let results = (outs Variadic:$results); let regions = (region SizedRegion<1>:$region); @@ -682,7 +682,7 @@ def AffineParallelOp : Affine_Op<"parallel", OpBuilder<(ins "TypeRange":$resultTypes, "ArrayRef":$reductions, "ArrayRef":$lbMaps, "ValueRange":$lbArgs, "ArrayRef":$ubMaps, "ValueRange":$ubArgs, - "ArrayRef":$steps)> + "ArrayRef":$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, AllTypesMatch<["lowerBound", "upperBound", "step"]>, @@ -302,7 +302,7 @@ def ForallOp : SCF_Op<"forall", [ AutomaticAllocationScope, DeclareOpInterfaceMethods, RecursiveMemoryEffects, SingleBlockImplicitTerminator<"scf::InParallelOp">, @@ -510,6 +510,27 @@ def ForallOp : SCF_Op<"forall", [ ]; let extraClassDeclaration = [{ + // Get lower bounds as OpFoldResult. + SmallVector getMixedLowerBound() { + auto maybeLowerBounds = getLowerBounds(); + assert(maybeLowerBounds.has_value() && "expected values"); + return *maybeLowerBounds; + } + + // Get upper bounds as OpFoldResult. + SmallVector getMixedUpperBound() { + auto maybeUpperBounds = getUpperBounds(); + assert(maybeUpperBounds.has_value() && "expected values"); + return *maybeUpperBounds; + } + + // Get steps as OpFoldResult. + SmallVector getMixedStep() { + auto maybeSteps = getSteps(); + assert(maybeSteps.has_value() && "expected values"); + return *maybeSteps; + } + /// Get lower bounds as values. SmallVector getLowerBound(OpBuilder &b) { return getValueOrCreateConstantIndexOp(b, getLoc(), getMixedLowerBound()); @@ -744,7 +765,7 @@ def ParallelOp : SCF_Op<"parallel", [AutomaticAllocationScope, AttrSizedOperandSegments, DeclareOpInterfaceMethods, + "getLowerBounds", "getUpperBounds", "getSteps"]>, RecursiveMemoryEffects, DeclareOpInterfaceMethods, 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 { 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(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 lbs = getValueOrCreateConstantIndexOp( - rewriter, loc, forallOp.getMixedLowerBound()); - SmallVector ubs = getValueOrCreateConstantIndexOp( - rewriter, loc, forallOp.getMixedUpperBound()); - SmallVector steps = - getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep()); + SmallVector lbs = forallOp.getLowerBound(rewriter); + SmallVector ubs = forallOp.getUpperBound(rewriter); + SmallVector steps = forallOp.getStep(rewriter); // Create empty scf.parallel op. auto parallelOp = rewriter.create(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 AffineForOp::getLoopRegions() { return {&getRegion()}; } ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; } -SmallVector AffineForOp::getMixedLowerBound() { +std::optional> AffineForOp::getLowerBounds() { if (!hasConstantLowerBound()) - return {}; + return std::nullopt; OpBuilder b(getContext()); - return {OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))}; + return SmallVector{ + OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))}; } -SmallVector AffineForOp::getMixedStep() { +std::optional> AffineForOp::getSteps() { OpBuilder b(getContext()); - return {OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))}; + return SmallVector{ + OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))}; } -SmallVector AffineForOp::getMixedUpperBound() { +std::optional> AffineForOp::getUpperBounds() { if (!hasConstantUpperBound()) return {}; OpBuilder b(getContext()); - return {OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))}; + return SmallVector{ + OpFoldResult(b.getI64IntegerAttr(getConstantUpperBound()))}; } FailureOr AffineForOp::replaceWithAdditionalYields( @@ -3753,7 +3756,7 @@ SmallVector 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 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 steps = getSteps(); + SmallVector 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 steps = op.getSteps(); + SmallVector 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 ForOp::getMixedLowerBound() { - return {OpFoldResult(getLowerBound())}; +std::optional> ForOp::getLowerBounds() { + return SmallVector{OpFoldResult(getLowerBound())}; } -SmallVector ForOp::getMixedStep() { - return {OpFoldResult(getStep())}; +std::optional> ForOp::getSteps() { + return SmallVector{OpFoldResult(getStep())}; } -SmallVector ForOp::getMixedUpperBound() { - return {OpFoldResult(getUpperBound())}; +std::optional> ForOp::getUpperBounds() { + return SmallVector{OpFoldResult(getUpperBound())}; } std::optional ForOp::getLoopResults() { return getResults(); } @@ -1431,19 +1431,19 @@ ValueRange ForallOp::getInductionVars() { } // Get lower bounds as OpFoldResult. -SmallVector ForallOp::getMixedLowerBound() { +std::optional> ForallOp::getLowerBounds() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b); } // Get upper bounds as OpFoldResult. -SmallVector ForallOp::getMixedUpperBound() { +std::optional> ForallOp::getUpperBounds() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b); } // Get steps as OpFoldResult. -SmallVector ForallOp::getMixedStep() { +std::optional> ForallOp::getSteps() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticStep(), getDynamicStep(), b); } @@ -3006,15 +3006,17 @@ SmallVector ParallelOp::getLoopRegions() { return {&getRegion()}; } ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); } -SmallVector ParallelOp::getMixedLowerBound() { +std::optional> ParallelOp::getLowerBounds() { return getLowerBound(); } -SmallVector ParallelOp::getMixedUpperBound() { +std::optional> ParallelOp::getUpperBounds() { return getUpperBound(); } -SmallVector ParallelOp::getMixedStep() { return getStep(); } +std::optional> ParallelOp::getSteps() { + return getStep(); +} ParallelOp mlir::scf::getParallelForInductionVarOwner(Value val) { auto ivArg = llvm::dyn_cast(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 lbs = getValueOrCreateConstantIndexOp( - rewriter, loc, forallOp.getMixedLowerBound()); - SmallVector ubs = getValueOrCreateConstantIndexOp( - rewriter, loc, forallOp.getMixedUpperBound()); - SmallVector steps = - getValueOrCreateConstantIndexOp(rewriter, loc, forallOp.getMixedStep()); + SmallVector lbs = forallOp.getLowerBound(rewriter); + SmallVector ubs = forallOp.getUpperBound(rewriter); + SmallVector steps = forallOp.getStep(rewriter); LoopNest loopNest = scf::buildLoopNest(rewriter, loc, lbs, ubs, steps); SmallVector 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 maybeLb = loopLikeOp.getSingleLowerBound(); + std::optional maybeSingleLb = + loopLikeOp.getSingleLowerBound(); + EXPECT_TRUE(maybeSingleLb.has_value()); + std::optional maybeSingleUb = + loopLikeOp.getSingleUpperBound(); + EXPECT_TRUE(maybeSingleUb.has_value()); + std::optional maybeSingleStep = loopLikeOp.getSingleStep(); + EXPECT_TRUE(maybeSingleStep.has_value()); + std::optional maybeSingleIndVar = + loopLikeOp.getSingleInductionVar(); + EXPECT_TRUE(maybeSingleIndVar.has_value()); + + std::optional> maybeLb = + loopLikeOp.getLowerBounds(); EXPECT_TRUE(maybeLb.has_value()); - std::optional maybeUb = loopLikeOp.getSingleUpperBound(); + EXPECT_EQ((*maybeLb).size(), 1u); + std::optional> maybeUb = + loopLikeOp.getUpperBounds(); EXPECT_TRUE(maybeUb.has_value()); - std::optional maybeStep = loopLikeOp.getSingleStep(); + EXPECT_EQ((*maybeUb).size(), 1u); + std::optional> maybeStep = loopLikeOp.getSteps(); EXPECT_TRUE(maybeStep.has_value()); - std::optional 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 maybeLb = loopLikeOp.getSingleLowerBound(); - EXPECT_FALSE(maybeLb.has_value()); - std::optional maybeUb = loopLikeOp.getSingleUpperBound(); - EXPECT_FALSE(maybeUb.has_value()); - std::optional maybeStep = loopLikeOp.getSingleStep(); - EXPECT_FALSE(maybeStep.has_value()); - std::optional maybeIndVar = + std::optional maybeSingleLb = + loopLikeOp.getSingleLowerBound(); + EXPECT_FALSE(maybeSingleLb.has_value()); + std::optional maybeSingleUb = + loopLikeOp.getSingleUpperBound(); + EXPECT_FALSE(maybeSingleUb.has_value()); + std::optional maybeSingleStep = loopLikeOp.getSingleStep(); + EXPECT_FALSE(maybeSingleStep.has_value()); + std::optional maybeSingleIndVar = loopLikeOp.getSingleInductionVar(); - EXPECT_FALSE(maybeIndVar.has_value()); + EXPECT_FALSE(maybeSingleIndVar.has_value()); + + std::optional> maybeLb = + loopLikeOp.getLowerBounds(); + EXPECT_TRUE(maybeLb.has_value()); + EXPECT_EQ((*maybeLb).size(), 2u); + std::optional> maybeUb = + loopLikeOp.getUpperBounds(); + EXPECT_TRUE(maybeUb.has_value()); + EXPECT_EQ((*maybeUb).size(), 2u); + std::optional> 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 Date: Thu, 6 Jun 2024 11:31:11 -0500 Subject: [PATCH 4/8] change return type of getInductionVars to SmallVector --- 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 AffineForOp::getLoopRegions() { return {&getRegion()}; } -ValueRange AffineForOp::getInductionVars() { return {getInductionVar()}; } +SmallVector AffineForOp::getInductionVars() { + return {getInductionVar()}; +} std::optional> 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(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 nonZeroNumThreads = llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) { return !isConstantIntValue(ofr, 0); @@ -746,7 +746,7 @@ FailureOr linalg::tileReductionUsingForall( b.getIndexAttr(0)); SmallVector 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( loc, cast(initOperand.getType()), @@ -814,7 +814,7 @@ FailureOr 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 ForOp::getInductionVars() { return {getInductionVar()}; } std::optional> ForOp::getLowerBounds() { return SmallVector{OpFoldResult(getLowerBound())}; @@ -1426,8 +1426,8 @@ SmallVector ForallOp::getCombiningOps(BlockArgument bbArg) { return storeOps; } -ValueRange ForallOp::getInductionVars() { - return getBody()->getArguments().take_front(getRank()); +SmallVector ForallOp::getInductionVars() { + return SmallVector(getBody()->getArguments().take_front(getRank())); } // Get lower bounds as OpFoldResult. @@ -3004,7 +3004,9 @@ void ParallelOp::print(OpAsmPrinter &p) { SmallVector ParallelOp::getLoopRegions() { return {&getRegion()}; } -ValueRange ParallelOp::getInductionVars() { return getBody()->getArguments(); } +SmallVector ParallelOp::getInductionVars() { + return SmallVector(getBody()->getArguments()); +} std::optional> ParallelOp::getLowerBounds() { return getLowerBound(); From 009fd15ab8abefd56afe6424e27f912a4166329d Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 6 Jun 2024 14:02:52 -0500 Subject: [PATCH 5/8] address maks's comments --- mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp | 2 +- mlir/lib/Dialect/SCF/IR/SCF.cpp | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp b/mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp index 4eacaa8d1e327..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)); - auto threadIds = forallOp.getInductionVars(); + SmallVector threadIds = forallOp.getInductionVars(); SmallVector nonZeroNumThreads = llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) { return !isConstantIntValue(ofr, 0); diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index 0ce10ebdad3e2..a930f8c71454c 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -1427,7 +1427,7 @@ SmallVector ForallOp::getCombiningOps(BlockArgument bbArg) { } SmallVector ForallOp::getInductionVars() { - return SmallVector(getBody()->getArguments().take_front(getRank())); + return SmallVector{getBody()->getArguments().take_front(getRank())}; } // Get lower bounds as OpFoldResult. @@ -3005,7 +3005,7 @@ void ParallelOp::print(OpAsmPrinter &p) { SmallVector ParallelOp::getLoopRegions() { return {&getRegion()}; } SmallVector ParallelOp::getInductionVars() { - return SmallVector(getBody()->getArguments()); + return SmallVector{getBody()->getArguments()}; } std::optional> ParallelOp::getLowerBounds() { From d34ad95aba669b5700976f0d2ed4d68b4902e9be Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 6 Jun 2024 15:26:06 -0500 Subject: [PATCH 6/8] change interface method names again and revert steps operand change --- .../mlir/Dialect/Affine/IR/AffineOps.td | 10 +++---- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 27 ++++++++++++------- .../mlir/Interfaces/LoopLikeInterface.td | 16 +++++------ .../AffineToStandard/AffineToStandard.cpp | 4 +-- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 18 ++++++------- mlir/lib/Dialect/Affine/Utils/Utils.cpp | 2 +- mlir/lib/Dialect/SCF/IR/SCF.cpp | 24 ++++++++--------- .../Dialect/SCF/LoopLikeSCFOpsTest.cpp | 18 +++++++------ 8 files changed, 64 insertions(+), 55 deletions(-) diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index 4c032e66f7a83..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, DeclareOpInterfaceMethods]> { @@ -671,7 +671,7 @@ def AffineParallelOp : Affine_Op<"parallel", I32ElementsAttr:$lowerBoundsGroups, AffineMapAttr:$upperBoundsMap, I32ElementsAttr:$upperBoundsGroups, - I64SmallVectorArrayAttr:$step, + I64SmallVectorArrayAttr:$steps, Variadic:$mapOperands); let results = (outs Variadic:$results); let regions = (region SizedRegion<1>:$region); @@ -682,7 +682,7 @@ def AffineParallelOp : Affine_Op<"parallel", OpBuilder<(ins "TypeRange":$resultTypes, "ArrayRef":$reductions, "ArrayRef":$lbMaps, "ValueRange":$lbArgs, "ArrayRef":$ubMaps, "ValueRange":$ubArgs, - "ArrayRef":$step)> + "ArrayRef":$steps)> ]; let extraClassDeclaration = [{ @@ -727,7 +727,7 @@ def AffineParallelOp : Affine_Op<"parallel", static StringRef getUpperBoundsGroupsAttrStrName() { return "upperBoundsGroups"; } - static StringRef getStepsAttrStrName() { return "step"; } + static StringRef getStepsAttrStrName() { return "steps"; } /// 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 66b478f141b32..3704b15972278 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, AllTypesMatch<["lowerBound", "upperBound", "step"]>, @@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [ AttrSizedOperandSegments, AutomaticAllocationScope, DeclareOpInterfaceMethods, RecursiveMemoryEffects, SingleBlockImplicitTerminator<"scf::InParallelOp">, @@ -510,23 +510,26 @@ def ForallOp : SCF_Op<"forall", [ ]; let extraClassDeclaration = [{ + SmallVector getInductionVars() { + return getLoopInductionVars(); + } // Get lower bounds as OpFoldResult. SmallVector getMixedLowerBound() { - auto maybeLowerBounds = getLowerBounds(); + auto maybeLowerBounds = getLoopLowerBounds(); assert(maybeLowerBounds.has_value() && "expected values"); return *maybeLowerBounds; } // Get upper bounds as OpFoldResult. SmallVector getMixedUpperBound() { - auto maybeUpperBounds = getUpperBounds(); + auto maybeUpperBounds = getLoopUpperBounds(); assert(maybeUpperBounds.has_value() && "expected values"); return *maybeUpperBounds; } // Get steps as OpFoldResult. SmallVector getMixedStep() { - auto maybeSteps = getSteps(); + auto maybeSteps = getLoopSteps(); assert(maybeSteps.has_value() && "expected values"); return *maybeSteps; } @@ -588,7 +591,7 @@ def ForallOp : SCF_Op<"forall", [ } ::mlir::Value getInductionVar(int64_t idx) { - return getInductionVars()[idx]; + return getLoopInductionVars()[idx]; } ::mlir::Block::BlockArgListType getRegionOutArgs() { @@ -764,8 +767,8 @@ def IfOp : SCF_Op<"if", [DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, RecursiveMemoryEffects, DeclareOpInterfaceMethods, SingleBlockImplicitTerminator<"scf::ReduceOp">, @@ -845,6 +848,10 @@ def ParallelOp : SCF_Op<"parallel", ]; let extraClassDeclaration = [{ + // Get induction variables. + SmallVector getInductionVars() { + return getLoopInductionVars(); + } 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 bace8f8384d44..5312ace4db68e 100644 --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td @@ -96,7 +96,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { Return all induction variables. }], /*retTy=*/"::llvm::SmallVector<::mlir::Value>", - /*methodName=*/"getInductionVars", + /*methodName=*/"getLoopInductionVars", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ @@ -107,7 +107,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { Return all lower bounds. }], /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>", - /*methodName=*/"getLowerBounds", + /*methodName=*/"getLoopLowerBounds", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ @@ -118,7 +118,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { Return all steps. }], /*retTy=*/"std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>", - /*methodName=*/"getSteps", + /*methodName=*/"getLoopSteps", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ @@ -129,7 +129,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { Return all upper bounds. }], /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>", - /*methodName=*/"getUpperBounds", + /*methodName=*/"getLoopUpperBounds", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ @@ -234,7 +234,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { /// If there is a single induction variable return it, otherwise return /// std::nullopt. ::std::optional<::mlir::Value> getSingleInductionVar() { - auto inductionVars = this->getInductionVars(); + auto inductionVars = this->getLoopInductionVars(); if (inductionVars.size() == 1) return inductionVars[0]; return std::nullopt; @@ -242,7 +242,7 @@ 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->getLowerBounds(); + auto lowerBounds = this->getLoopLowerBounds(); if (lowerBounds.has_value() && (*lowerBounds).size() == 1) return (*lowerBounds)[0]; return std::nullopt; @@ -250,7 +250,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { /// Return the single step value or attribute if it exists, otherwise /// return std::nullopt. ::std::optional<::mlir::OpFoldResult> getSingleStep() { - auto steps = this->getSteps(); + auto steps = this->getLoopSteps(); if (steps.has_value() && (*steps).size() == 1) return (*steps)[0]; return std::nullopt; @@ -258,7 +258,7 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { /// Return the single upper bound value or attribute if it exists, otherwise /// return std::nullopt. ::std::optional<::mlir::OpFoldResult> getSingleUpperBound() { - auto upperBounds = this->getUpperBounds(); + auto upperBounds = this->getLoopUpperBounds(); if (upperBounds.has_value() && (*upperBounds).size() == 1) return (*upperBounds)[0]; return std::nullopt; diff --git a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp index 20487b32e3fe0..10ccd5c97783b 100644 --- a/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp +++ b/mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp @@ -196,8 +196,8 @@ class AffineParallelLowering : public OpRewritePattern { return rewriter.notifyMatchFailure(op, "couldn't convert upper bounds"); upperBoundTuple.push_back(upper); } - steps.reserve(op.getStep().size()); - for (int64_t step : op.getStep()) + steps.reserve(op.getSteps().size()); + for (int64_t step : op.getSteps()) steps.push_back(rewriter.create(loc, step)); // Get the terminator op. diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index 5467c60242664..d5cb04743dfb9 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2454,11 +2454,11 @@ bool AffineForOp::matchingBoundOperandList() { SmallVector AffineForOp::getLoopRegions() { return {&getRegion()}; } -SmallVector AffineForOp::getInductionVars() { +SmallVector AffineForOp::getLoopInductionVars() { return {getInductionVar()}; } -std::optional> AffineForOp::getLowerBounds() { +std::optional> AffineForOp::getLoopLowerBounds() { if (!hasConstantLowerBound()) return std::nullopt; OpBuilder b(getContext()); @@ -2466,13 +2466,13 @@ std::optional> AffineForOp::getLowerBounds() { OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))}; } -std::optional> AffineForOp::getSteps() { +std::optional> AffineForOp::getLoopSteps() { OpBuilder b(getContext()); return SmallVector{ OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))}; } -std::optional> AffineForOp::getUpperBounds() { +std::optional> AffineForOp::getLoopUpperBounds() { if (!hasConstantUpperBound()) return {}; OpBuilder b(getContext()); @@ -3758,7 +3758,7 @@ SmallVector AffineParallelOp::getLoopRegions() { return {&getRegion()}; } -unsigned AffineParallelOp::getNumDims() { return getStep().size(); } +unsigned AffineParallelOp::getNumDims() { return getSteps().size(); } AffineParallelOp::operand_range AffineParallelOp::getLowerBoundsOperands() { return getOperands().take_front(getLowerBoundsMap().getNumInputs()); @@ -3843,7 +3843,7 @@ void AffineParallelOp::setUpperBounds(ValueRange ubOperands, AffineMap map) { } void AffineParallelOp::setSteps(ArrayRef newSteps) { - setStepAttr(getBodyBuilder().getI64ArrayAttr(newSteps)); + setStepsAttr(getBodyBuilder().getI64ArrayAttr(newSteps)); } // check whether resultType match op or not in affine.parallel @@ -3893,14 +3893,14 @@ LogicalResult AffineParallelOp::verify() { auto numDims = getNumDims(); if (getLowerBoundsGroups().getNumElements() != numDims || getUpperBoundsGroups().getNumElements() != numDims || - getStep().size() != numDims || getBody()->getNumArguments() != numDims) { + getSteps().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 (" << getStep().size() + << "), and the number of steps (" << getSteps().size() << ") must all match"; } @@ -4018,7 +4018,7 @@ void AffineParallelOp::print(OpAsmPrinter &p) { printMinMaxBound(p, getUpperBoundsMapAttr(), getUpperBoundsGroupsAttr(), getUpperBoundsOperands(), "min"); p << ')'; - SmallVector steps = getStep(); + SmallVector steps = getSteps(); 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 a652ee4a488d1..f46381403bc52 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 steps = op.getStep(); + SmallVector steps = op.getSteps(); // 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 a930f8c71454c..e921177f73215 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -378,17 +378,17 @@ LogicalResult ForOp::verifyRegions() { return success(); } -SmallVector ForOp::getInductionVars() { return {getInductionVar()}; } +SmallVector ForOp::getLoopInductionVars() { return {getInductionVar()}; } -std::optional> ForOp::getLowerBounds() { +std::optional> ForOp::getLoopLowerBounds() { return SmallVector{OpFoldResult(getLowerBound())}; } -std::optional> ForOp::getSteps() { +std::optional> ForOp::getLoopSteps() { return SmallVector{OpFoldResult(getStep())}; } -std::optional> ForOp::getUpperBounds() { +std::optional> ForOp::getLoopUpperBounds() { return SmallVector{OpFoldResult(getUpperBound())}; } @@ -1426,24 +1426,24 @@ SmallVector ForallOp::getCombiningOps(BlockArgument bbArg) { return storeOps; } -SmallVector ForallOp::getInductionVars() { +SmallVector ForallOp::getLoopInductionVars() { return SmallVector{getBody()->getArguments().take_front(getRank())}; } // Get lower bounds as OpFoldResult. -std::optional> ForallOp::getLowerBounds() { +std::optional> ForallOp::getLoopLowerBounds() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b); } // Get upper bounds as OpFoldResult. -std::optional> ForallOp::getUpperBounds() { +std::optional> ForallOp::getLoopUpperBounds() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b); } // Get steps as OpFoldResult. -std::optional> ForallOp::getSteps() { +std::optional> ForallOp::getLoopSteps() { Builder b(getOperation()->getContext()); return getMixedValues(getStaticStep(), getDynamicStep(), b); } @@ -3004,19 +3004,19 @@ void ParallelOp::print(OpAsmPrinter &p) { SmallVector ParallelOp::getLoopRegions() { return {&getRegion()}; } -SmallVector ParallelOp::getInductionVars() { +SmallVector ParallelOp::getLoopInductionVars() { return SmallVector{getBody()->getArguments()}; } -std::optional> ParallelOp::getLowerBounds() { +std::optional> ParallelOp::getLoopLowerBounds() { return getLowerBound(); } -std::optional> ParallelOp::getUpperBounds() { +std::optional> ParallelOp::getLoopUpperBounds() { return getUpperBound(); } -std::optional> ParallelOp::getSteps() { +std::optional> ParallelOp::getLoopSteps() { return getStep(); } diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp index 07504a99fecd3..75cd2bfb01de0 100644 --- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp +++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp @@ -40,17 +40,18 @@ class SCFLoopLikeTest : public ::testing::Test { EXPECT_TRUE(maybeSingleIndVar.has_value()); std::optional> maybeLb = - loopLikeOp.getLowerBounds(); + loopLikeOp.getLoopLowerBounds(); EXPECT_TRUE(maybeLb.has_value()); EXPECT_EQ((*maybeLb).size(), 1u); std::optional> maybeUb = - loopLikeOp.getUpperBounds(); + loopLikeOp.getLoopUpperBounds(); EXPECT_TRUE(maybeUb.has_value()); EXPECT_EQ((*maybeUb).size(), 1u); - std::optional> maybeStep = loopLikeOp.getSteps(); + std::optional> maybeStep = + loopLikeOp.getLoopSteps(); EXPECT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 1u); - EXPECT_EQ(loopLikeOp.getInductionVars().size(), 1u); + EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 1u); } void checkMultidimensional(LoopLikeOpInterface loopLikeOp) { @@ -67,17 +68,18 @@ class SCFLoopLikeTest : public ::testing::Test { EXPECT_FALSE(maybeSingleIndVar.has_value()); std::optional> maybeLb = - loopLikeOp.getLowerBounds(); + loopLikeOp.getLoopLowerBounds(); EXPECT_TRUE(maybeLb.has_value()); EXPECT_EQ((*maybeLb).size(), 2u); std::optional> maybeUb = - loopLikeOp.getUpperBounds(); + loopLikeOp.getLoopUpperBounds(); EXPECT_TRUE(maybeUb.has_value()); EXPECT_EQ((*maybeUb).size(), 2u); - std::optional> maybeStep = loopLikeOp.getSteps(); + std::optional> maybeStep = + loopLikeOp.getLoopSteps(); EXPECT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 2u); - EXPECT_EQ(loopLikeOp.getInductionVars().size(), 2u); + EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 2u); } MLIRContext context; From e0e526210a5ab6ce28fbc5fa5ee24f79cb1ee9a8 Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 6 Jun 2024 16:04:47 -0500 Subject: [PATCH 7/8] return option induction vars --- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 10 +++++++--- mlir/include/mlir/Interfaces/LoopLikeInterface.td | 8 ++++---- mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 4 ++-- mlir/lib/Dialect/SCF/IR/SCF.cpp | 8 +++++--- mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp | 10 ++++++++-- 5 files changed, 26 insertions(+), 14 deletions(-) diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td index 3704b15972278..d425c1c2a47b4 100644 --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -511,7 +511,9 @@ def ForallOp : SCF_Op<"forall", [ let extraClassDeclaration = [{ SmallVector getInductionVars() { - return getLoopInductionVars(); + auto maybeInductionVars = getLoopInductionVars();; + assert(maybeInductionVars.has_value() && "expected values"); + return *maybeInductionVars; } // Get lower bounds as OpFoldResult. SmallVector getMixedLowerBound() { @@ -591,7 +593,7 @@ def ForallOp : SCF_Op<"forall", [ } ::mlir::Value getInductionVar(int64_t idx) { - return getLoopInductionVars()[idx]; + return getInductionVars()[idx]; } ::mlir::Block::BlockArgListType getRegionOutArgs() { @@ -850,7 +852,9 @@ def ParallelOp : SCF_Op<"parallel", let extraClassDeclaration = [{ // Get induction variables. SmallVector getInductionVars() { - return getLoopInductionVars(); + auto 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 5312ace4db68e..2e6aabda30b07 100644 --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td @@ -95,12 +95,12 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { InterfaceMethod<[{ Return all induction variables. }], - /*retTy=*/"::llvm::SmallVector<::mlir::Value>", + /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::Value>>", /*methodName=*/"getLoopInductionVars", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ - return {}; + return std::nullopt; }] >, InterfaceMethod<[{ @@ -235,8 +235,8 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { /// std::nullopt. ::std::optional<::mlir::Value> getSingleInductionVar() { auto inductionVars = this->getLoopInductionVars(); - if (inductionVars.size() == 1) - return inductionVars[0]; + 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 diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index d5cb04743dfb9..0a58d2fdb02f5 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2454,8 +2454,8 @@ bool AffineForOp::matchingBoundOperandList() { SmallVector AffineForOp::getLoopRegions() { return {&getRegion()}; } -SmallVector AffineForOp::getLoopInductionVars() { - return {getInductionVar()}; +std::optional> AffineForOp::getLoopInductionVars() { + return SmallVector{getInductionVar()}; } std::optional> AffineForOp::getLoopLowerBounds() { diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index e921177f73215..c00579443ea29 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -378,7 +378,9 @@ LogicalResult ForOp::verifyRegions() { return success(); } -SmallVector ForOp::getLoopInductionVars() { return {getInductionVar()}; } +std::optional> ForOp::getLoopInductionVars() { + return SmallVector{getInductionVar()}; +} std::optional> ForOp::getLoopLowerBounds() { return SmallVector{OpFoldResult(getLowerBound())}; @@ -1426,7 +1428,7 @@ SmallVector ForallOp::getCombiningOps(BlockArgument bbArg) { return storeOps; } -SmallVector ForallOp::getLoopInductionVars() { +std::optional> ForallOp::getLoopInductionVars() { return SmallVector{getBody()->getArguments().take_front(getRank())}; } @@ -3004,7 +3006,7 @@ void ParallelOp::print(OpAsmPrinter &p) { SmallVector ParallelOp::getLoopRegions() { return {&getRegion()}; } -SmallVector ParallelOp::getLoopInductionVars() { +std::optional> ParallelOp::getLoopInductionVars() { return SmallVector{getBody()->getArguments()}; } diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp index 75cd2bfb01de0..20dbc8d362d27 100644 --- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp +++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp @@ -51,7 +51,10 @@ class SCFLoopLikeTest : public ::testing::Test { loopLikeOp.getLoopSteps(); EXPECT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 1u); - EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 1u); + std::optional> maybeInductionVars = + loopLikeOp.getLoopInductionVars(); + EXPECT_TRUE(maybeInductionVars.has_value()); + EXPECT_EQ((*maybeInductionVars).size(), 1u); } void checkMultidimensional(LoopLikeOpInterface loopLikeOp) { @@ -79,7 +82,10 @@ class SCFLoopLikeTest : public ::testing::Test { loopLikeOp.getLoopSteps(); EXPECT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 2u); - EXPECT_EQ(loopLikeOp.getLoopInductionVars().size(), 2u); + std::optional> maybeInductionVars = + loopLikeOp.getLoopInductionVars(); + EXPECT_TRUE(maybeInductionVars.has_value()); + EXPECT_EQ((*maybeInductionVars).size(), 2u); } MLIRContext context; From 7115a6e08bba43fe9750f8cef5c73f6be1b373fd Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 7 Jun 2024 11:45:44 -0500 Subject: [PATCH 8/8] address review comments --- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 19 ++++++------ .../mlir/Interfaces/LoopLikeInterface.td | 30 +++++++++++++------ mlir/lib/Dialect/SCF/IR/SCF.cpp | 8 ++--- .../Dialect/SCF/LoopLikeSCFOpsTest.cpp | 16 +++++----- 4 files changed, 43 insertions(+), 30 deletions(-) diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td index d425c1c2a47b4..f35ea962bea16 100644 --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -510,28 +510,29 @@ def ForallOp : SCF_Op<"forall", [ ]; let extraClassDeclaration = [{ + /// Get induction variables. SmallVector getInductionVars() { - auto maybeInductionVars = getLoopInductionVars();; + std::optional> maybeInductionVars = getLoopInductionVars(); assert(maybeInductionVars.has_value() && "expected values"); return *maybeInductionVars; } - // Get lower bounds as OpFoldResult. + /// Get lower bounds as OpFoldResult. SmallVector getMixedLowerBound() { - auto maybeLowerBounds = getLoopLowerBounds(); + std::optional> maybeLowerBounds = getLoopLowerBounds(); assert(maybeLowerBounds.has_value() && "expected values"); return *maybeLowerBounds; } - // Get upper bounds as OpFoldResult. + /// Get upper bounds as OpFoldResult. SmallVector getMixedUpperBound() { - auto maybeUpperBounds = getLoopUpperBounds(); + std::optional> maybeUpperBounds = getLoopUpperBounds(); assert(maybeUpperBounds.has_value() && "expected values"); return *maybeUpperBounds; } - // Get steps as OpFoldResult. + /// Get steps as OpFoldResult. SmallVector getMixedStep() { - auto maybeSteps = getLoopSteps(); + std::optional> maybeSteps = getLoopSteps(); assert(maybeSteps.has_value() && "expected values"); return *maybeSteps; } @@ -850,9 +851,9 @@ def ParallelOp : SCF_Op<"parallel", ]; let extraClassDeclaration = [{ - // Get induction variables. + /// Get induction variables. SmallVector getInductionVars() { - auto maybeInductionVars = getLoopInductionVars();; + std::optional> maybeInductionVars = getLoopInductionVars();; assert(maybeInductionVars.has_value() && "expected values"); return *maybeInductionVars; } diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td index 2e6aabda30b07..b748d5e29114a 100644 --- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td +++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td @@ -93,47 +93,59 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> { }] >, InterfaceMethod<[{ - Return all induction variables. + 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<::llvm::SmallVector<::mlir::Value>>", /*methodName=*/"getLoopInductionVars", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ - return std::nullopt; + return ::std::nullopt; }] >, InterfaceMethod<[{ - Return all lower bounds. + 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<::llvm::SmallVector<::mlir::OpFoldResult>>", /*methodName=*/"getLoopLowerBounds", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ - return std::nullopt; + return ::std::nullopt; }] >, InterfaceMethod<[{ - Return all steps. + 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<::llvm::SmallVector<::mlir::OpFoldResult>>", + /*retTy=*/"::std::optional<::llvm::SmallVector<::mlir::OpFoldResult>>", /*methodName=*/"getLoopSteps", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ - return std::nullopt; + return ::std::nullopt; }] >, InterfaceMethod<[{ - Return all upper bounds. + 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<::llvm::SmallVector<::mlir::OpFoldResult>>", /*methodName=*/"getLoopUpperBounds", /*args=*/(ins), /*methodBody=*/"", /*defaultImplementation=*/[{ - return std::nullopt; + return ::std::nullopt; }] >, InterfaceMethod<[{ diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index c00579443ea29..5e94f4dc612a7 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -379,19 +379,19 @@ LogicalResult ForOp::verifyRegions() { } std::optional> ForOp::getLoopInductionVars() { - return SmallVector{getInductionVar()}; + return SmallVector{getInductionVar()}; } std::optional> ForOp::getLoopLowerBounds() { - return SmallVector{OpFoldResult(getLowerBound())}; + return SmallVector{OpFoldResult(getLowerBound())}; } std::optional> ForOp::getLoopSteps() { - return SmallVector{OpFoldResult(getStep())}; + return SmallVector{OpFoldResult(getStep())}; } std::optional> ForOp::getLoopUpperBounds() { - return SmallVector{OpFoldResult(getUpperBound())}; + return SmallVector{OpFoldResult(getUpperBound())}; } std::optional ForOp::getLoopResults() { return getResults(); } diff --git a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp index 20dbc8d362d27..53a4af14d119a 100644 --- a/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp +++ b/mlir/unittests/Dialect/SCF/LoopLikeSCFOpsTest.cpp @@ -41,19 +41,19 @@ class SCFLoopLikeTest : public ::testing::Test { std::optional> maybeLb = loopLikeOp.getLoopLowerBounds(); - EXPECT_TRUE(maybeLb.has_value()); + ASSERT_TRUE(maybeLb.has_value()); EXPECT_EQ((*maybeLb).size(), 1u); std::optional> maybeUb = loopLikeOp.getLoopUpperBounds(); - EXPECT_TRUE(maybeUb.has_value()); + ASSERT_TRUE(maybeUb.has_value()); EXPECT_EQ((*maybeUb).size(), 1u); std::optional> maybeStep = loopLikeOp.getLoopSteps(); - EXPECT_TRUE(maybeStep.has_value()); + ASSERT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 1u); std::optional> maybeInductionVars = loopLikeOp.getLoopInductionVars(); - EXPECT_TRUE(maybeInductionVars.has_value()); + ASSERT_TRUE(maybeInductionVars.has_value()); EXPECT_EQ((*maybeInductionVars).size(), 1u); } @@ -72,19 +72,19 @@ class SCFLoopLikeTest : public ::testing::Test { std::optional> maybeLb = loopLikeOp.getLoopLowerBounds(); - EXPECT_TRUE(maybeLb.has_value()); + ASSERT_TRUE(maybeLb.has_value()); EXPECT_EQ((*maybeLb).size(), 2u); std::optional> maybeUb = loopLikeOp.getLoopUpperBounds(); - EXPECT_TRUE(maybeUb.has_value()); + ASSERT_TRUE(maybeUb.has_value()); EXPECT_EQ((*maybeUb).size(), 2u); std::optional> maybeStep = loopLikeOp.getLoopSteps(); - EXPECT_TRUE(maybeStep.has_value()); + ASSERT_TRUE(maybeStep.has_value()); EXPECT_EQ((*maybeStep).size(), 2u); std::optional> maybeInductionVars = loopLikeOp.getLoopInductionVars(); - EXPECT_TRUE(maybeInductionVars.has_value()); + ASSERT_TRUE(maybeInductionVars.has_value()); EXPECT_EQ((*maybeInductionVars).size(), 2u); }