[Mlir-commits] [mlir] Reland Refactor LoopFuseSiblingOp and support parallel fusion #94391 (PR #97607)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jul 3 10:02:32 PDT 2024
https://github.com/srcarroll updated https://github.com/llvm/llvm-project/pull/97607
>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 01/34] 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 50852d570440e0041c8b2b38925c4af05fac0636 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Tue, 4 Jun 2024 14:44:53 -0500
Subject: [PATCH 02/34] Refactor LoopFuseSiblingOp and support parallel fusion
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 16 ++
.../SCF/TransformOps/SCFTransformOps.cpp | 53 +++--
.../SCF/Transforms/ParallelLoopFusion.cpp | 204 +----------------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 208 ++++++++++++++++++
.../SCF/transform-loop-fuse-sibling.mlir | 53 +++++
5 files changed, 304 insertions(+), 230 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index bc09cc7f7fa5e..2944d8ffac022 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -156,6 +156,12 @@ Loops tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef<Value> sizes);
void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
scf::ForOp root);
+/// Prepends operations of firstPloop's body into secondPloop's body.
+/// Updates secondPloop with new loop.
+void fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,
+ OpBuilder builder,
+ llvm::function_ref<bool(Value, Value)> mayAlias);
+
/// Given two scf.forall loops, `target` and `source`, fuses `target` into
/// `source`. Assumes that the given loops are siblings and are independent of
/// each other.
@@ -177,6 +183,16 @@ scf::ForallOp fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp fuseIndependentSiblingForLoops(scf::ForOp target, scf::ForOp source,
RewriterBase &rewriter);
+/// Given two scf.parallel loops, `target` and `source`, fuses `target` into
+/// `source`. Assumes that the given loops are siblings and are independent of
+/// each other.
+///
+/// This function does not perform any legality checks and simply fuses the
+/// loops. The caller is responsible for ensuring that the loops are legal to
+/// fuse.
+scf::ParallelOp fuseIndependentSiblingParallelLoops(scf::ParallelOp target,
+ scf::ParallelOp source,
+ RewriterBase &rewriter);
} // namespace mlir
#endif // MLIR_DIALECT_SCF_UTILS_UTILS_H_
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 69f83d8bd70da..1c53e89d69040 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -442,39 +442,32 @@ static DiagnosedSilenceableFailure isOpSibling(Operation *target,
return DiagnosedSilenceableFailure::success();
}
-/// Check if `target` scf.forall can be fused into `source` scf.forall.
+/// Check if `target` scf loop can be fused into `source` scf loop.
+/// Applies for scf.for, scf.forall, and scf.parallel.
///
/// This simply checks if both loops have the same bounds, steps and mapping.
/// No attempt is made at checking that the side effects of `target` and
/// `source` are independent of each other.
-static bool isForallWithIdenticalConfiguration(Operation *target,
- Operation *source) {
- auto targetOp = dyn_cast<scf::ForallOp>(target);
- auto sourceOp = dyn_cast<scf::ForallOp>(source);
- if (!targetOp || !sourceOp)
- return false;
-
- return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
- targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
- targetOp.getMixedStep() == sourceOp.getMixedStep() &&
- targetOp.getMapping() == sourceOp.getMapping();
-}
-
-/// Check if `target` scf.for can be fused into `source` scf.for.
-///
-/// This simply checks if both loops have the same bounds and steps. No attempt
-/// is made at checking that the side effects of `target` and `source` are
-/// independent of each other.
-static bool isForWithIdenticalConfiguration(Operation *target,
- Operation *source) {
- auto targetOp = dyn_cast<scf::ForOp>(target);
- auto sourceOp = dyn_cast<scf::ForOp>(source);
+template <typename LoopTy>
+static bool isLoopWithIdenticalConfiguration(Operation *target,
+ Operation *source) {
+ static_assert(llvm::is_one_of<LoopTy, scf::ForallOp, scf::ForOp,
+ scf::ParallelOp>::value,
+ "applies to only `forall`, `for` and `parallel`");
+ auto targetOp = dyn_cast<LoopTy>(target);
+ auto sourceOp = dyn_cast<LoopTy>(source);
if (!targetOp || !sourceOp)
return false;
- return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
- targetOp.getUpperBound() == sourceOp.getUpperBound() &&
- targetOp.getStep() == sourceOp.getStep();
+ if constexpr (std::is_same_v<LoopTy, scf::ForallOp>)
+ return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
+ targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
+ targetOp.getMixedStep() == sourceOp.getMixedStep() &&
+ targetOp.getMapping() == sourceOp.getMapping();
+ else
+ return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
+ targetOp.getUpperBound() == sourceOp.getUpperBound() &&
+ targetOp.getStep() == sourceOp.getStep();
}
DiagnosedSilenceableFailure
@@ -502,12 +495,16 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
Operation *fusedLoop;
/// TODO: Support fusion for loop-like ops besides scf.for and scf.forall.
- if (isForWithIdenticalConfiguration(target, source)) {
+ if (isLoopWithIdenticalConfiguration<scf::ForOp>(target, source)) {
fusedLoop = fuseIndependentSiblingForLoops(
cast<scf::ForOp>(target), cast<scf::ForOp>(source), rewriter);
- } else if (isForallWithIdenticalConfiguration(target, source)) {
+ } else if (isLoopWithIdenticalConfiguration<scf::ForallOp>(target, source)) {
fusedLoop = fuseIndependentSiblingForallLoops(
cast<scf::ForallOp>(target), cast<scf::ForallOp>(source), rewriter);
+ } else if (isLoopWithIdenticalConfiguration<scf::ParallelOp>(target,
+ source)) {
+ fusedLoop = fuseIndependentSiblingParallelLoops(
+ cast<scf::ParallelOp>(target), cast<scf::ParallelOp>(source), rewriter);
} else
return emitSilenceableFailure(target->getLoc())
<< "operations cannot be fused";
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 5934d85373b03..abac91cfaf7d9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -16,6 +16,7 @@
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/OpDefinition.h"
@@ -30,207 +31,6 @@ namespace mlir {
using namespace mlir;
using namespace mlir::scf;
-/// Verify there are no nested ParallelOps.
-static bool hasNestedParallelOp(ParallelOp ploop) {
- auto walkResult =
- ploop.getBody()->walk([](ParallelOp) { return WalkResult::interrupt(); });
- return walkResult.wasInterrupted();
-}
-
-/// Verify equal iteration spaces.
-static bool equalIterationSpaces(ParallelOp firstPloop,
- ParallelOp secondPloop) {
- if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
- return false;
-
- auto matchOperands = [&](const OperandRange &lhs,
- const OperandRange &rhs) -> bool {
- // TODO: Extend this to support aliases and equal constants.
- return std::equal(lhs.begin(), lhs.end(), rhs.begin());
- };
- return matchOperands(firstPloop.getLowerBound(),
- secondPloop.getLowerBound()) &&
- matchOperands(firstPloop.getUpperBound(),
- secondPloop.getUpperBound()) &&
- matchOperands(firstPloop.getStep(), secondPloop.getStep());
-}
-
-/// Checks if the parallel loops have mixed access to the same buffers. Returns
-/// `true` if the first parallel loop writes to the same indices that the second
-/// loop reads.
-static bool haveNoReadsAfterWriteExceptSameIndex(
- ParallelOp firstPloop, ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;
- SmallVector<Value> bufferStoresVec;
- firstPloop.getBody()->walk([&](memref::StoreOp store) {
- bufferStores[store.getMemRef()].push_back(store.getIndices());
- bufferStoresVec.emplace_back(store.getMemRef());
- });
- auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {
- Value loadMem = load.getMemRef();
- // Stop if the memref is defined in secondPloop body. Careful alias analysis
- // is needed.
- auto *memrefDef = loadMem.getDefiningOp();
- if (memrefDef && memrefDef->getBlock() == load->getBlock())
- return WalkResult::interrupt();
-
- for (Value store : bufferStoresVec)
- if (store != loadMem && mayAlias(store, loadMem))
- return WalkResult::interrupt();
-
- auto write = bufferStores.find(loadMem);
- if (write == bufferStores.end())
- return WalkResult::advance();
-
- // Check that at last one store was retrieved
- if (!write->second.size())
- return WalkResult::interrupt();
-
- auto storeIndices = write->second.front();
-
- // Multiple writes to the same memref are allowed only on the same indices
- for (const auto &othStoreIndices : write->second) {
- if (othStoreIndices != storeIndices)
- return WalkResult::interrupt();
- }
-
- // Check that the load indices of secondPloop coincide with store indices of
- // firstPloop for the same memrefs.
- auto loadIndices = load.getIndices();
- if (storeIndices.size() != loadIndices.size())
- return WalkResult::interrupt();
- for (int i = 0, e = storeIndices.size(); i < e; ++i) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
- loadIndices[i]) {
- auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
- auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
- if (storeIndexDefOp && loadIndexDefOp) {
- if (!isMemoryEffectFree(storeIndexDefOp))
- return WalkResult::interrupt();
- if (!isMemoryEffectFree(loadIndexDefOp))
- return WalkResult::interrupt();
- if (!OperationEquivalence::isEquivalentTo(
- storeIndexDefOp, loadIndexDefOp,
- [&](Value storeIndex, Value loadIndex) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
- firstToSecondPloopIndices.lookupOrDefault(loadIndex))
- return failure();
- else
- return success();
- },
- /*markEquivalent=*/nullptr,
- OperationEquivalence::Flags::IgnoreLocations)) {
- return WalkResult::interrupt();
- }
- } else
- return WalkResult::interrupt();
- }
- }
- return WalkResult::advance();
- });
- return !walkResult.wasInterrupted();
-}
-
-/// Analyzes dependencies in the most primitive way by checking simple read and
-/// write patterns.
-static LogicalResult
-verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- if (!haveNoReadsAfterWriteExceptSameIndex(
- firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))
- return failure();
-
- IRMapping secondToFirstPloopIndices;
- secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
- firstPloop.getBody()->getArguments());
- return success(haveNoReadsAfterWriteExceptSameIndex(
- secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));
-}
-
-static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- return !hasNestedParallelOp(firstPloop) &&
- !hasNestedParallelOp(secondPloop) &&
- equalIterationSpaces(firstPloop, secondPloop) &&
- succeeded(verifyDependencies(firstPloop, secondPloop,
- firstToSecondPloopIndices, mayAlias));
-}
-
-/// Prepends operations of firstPloop's body into secondPloop's body.
-/// Updates secondPloop with new loop.
-static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
- OpBuilder builder,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- Block *block1 = firstPloop.getBody();
- Block *block2 = secondPloop.getBody();
- IRMapping firstToSecondPloopIndices;
- firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
-
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias))
- return;
-
- DominanceInfo dom;
- // We are fusing first loop into second, make sure there are no users of the
- // first loop results between loops.
- for (Operation *user : firstPloop->getUsers())
- if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
- return;
-
- ValueRange inits1 = firstPloop.getInitVals();
- ValueRange inits2 = secondPloop.getInitVals();
-
- SmallVector<Value> newInitVars(inits1.begin(), inits1.end());
- newInitVars.append(inits2.begin(), inits2.end());
-
- IRRewriter b(builder);
- b.setInsertionPoint(secondPloop);
- auto newSecondPloop = b.create<ParallelOp>(
- secondPloop.getLoc(), secondPloop.getLowerBound(),
- secondPloop.getUpperBound(), secondPloop.getStep(), newInitVars);
-
- Block *newBlock = newSecondPloop.getBody();
- auto term1 = cast<ReduceOp>(block1->getTerminator());
- auto term2 = cast<ReduceOp>(block2->getTerminator());
-
- b.inlineBlockBefore(block2, newBlock, newBlock->begin(),
- newBlock->getArguments());
- b.inlineBlockBefore(block1, newBlock, newBlock->begin(),
- newBlock->getArguments());
-
- ValueRange results = newSecondPloop.getResults();
- if (!results.empty()) {
- b.setInsertionPointToEnd(newBlock);
-
- ValueRange reduceArgs1 = term1.getOperands();
- ValueRange reduceArgs2 = term2.getOperands();
- SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
- newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());
-
- auto newReduceOp = b.create<scf::ReduceOp>(term2.getLoc(), newReduceArgs);
-
- for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
- term1.getReductions(), term2.getReductions()))) {
- Block &oldRedBlock = reg.front();
- Block &newRedBlock = newReduceOp.getReductions()[i].front();
- b.inlineBlockBefore(&oldRedBlock, &newRedBlock, newRedBlock.begin(),
- newRedBlock.getArguments());
- }
-
- firstPloop.replaceAllUsesWith(results.take_front(inits1.size()));
- secondPloop.replaceAllUsesWith(results.take_back(inits2.size()));
- }
- term1->erase();
- term2->erase();
- firstPloop.erase();
- secondPloop.erase();
- secondPloop = newSecondPloop;
-}
-
void mlir::scf::naivelyFuseParallelOps(
Region ®ion, llvm::function_ref<bool(Value, Value)> mayAlias) {
OpBuilder b(region);
@@ -259,7 +59,7 @@ void mlir::scf::naivelyFuseParallelOps(
}
for (MutableArrayRef<ParallelOp> ploops : ploopChains) {
for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
- fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);
+ mlir::fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);
}
}
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 6658cca03eba7..d85339f32dbe3 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
@@ -1070,6 +1071,206 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
return tileLoops;
}
+/// Checks if the parallel loops have mixed access to the same buffers. Returns
+/// `true` if the first parallel loop writes to the same indices that the second
+/// loop reads.
+static bool haveNoReadsAfterWriteExceptSameIndex(
+ scf::ParallelOp firstPloop, scf::ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;
+ SmallVector<Value> bufferStoresVec;
+ firstPloop.getBody()->walk([&](memref::StoreOp store) {
+ bufferStores[store.getMemRef()].push_back(store.getIndices());
+ bufferStoresVec.emplace_back(store.getMemRef());
+ });
+ auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {
+ Value loadMem = load.getMemRef();
+ // Stop if the memref is defined in secondPloop body. Careful alias analysis
+ // is needed.
+ auto *memrefDef = loadMem.getDefiningOp();
+ if (memrefDef && memrefDef->getBlock() == load->getBlock())
+ return WalkResult::interrupt();
+
+ for (Value store : bufferStoresVec)
+ if (store != loadMem && mayAlias(store, loadMem))
+ return WalkResult::interrupt();
+
+ auto write = bufferStores.find(loadMem);
+ if (write == bufferStores.end())
+ return WalkResult::advance();
+
+ // Check that at last one store was retrieved
+ if (!write->second.size())
+ return WalkResult::interrupt();
+
+ auto storeIndices = write->second.front();
+
+ // Multiple writes to the same memref are allowed only on the same indices
+ for (const auto &othStoreIndices : write->second) {
+ if (othStoreIndices != storeIndices)
+ return WalkResult::interrupt();
+ }
+
+ // Check that the load indices of secondPloop coincide with store indices of
+ // firstPloop for the same memrefs.
+ auto loadIndices = load.getIndices();
+ if (storeIndices.size() != loadIndices.size())
+ return WalkResult::interrupt();
+ for (int i = 0, e = storeIndices.size(); i < e; ++i) {
+ if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
+ loadIndices[i]) {
+ auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
+ auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
+ if (storeIndexDefOp && loadIndexDefOp) {
+ if (!isMemoryEffectFree(storeIndexDefOp))
+ return WalkResult::interrupt();
+ if (!isMemoryEffectFree(loadIndexDefOp))
+ return WalkResult::interrupt();
+ if (!OperationEquivalence::isEquivalentTo(
+ storeIndexDefOp, loadIndexDefOp,
+ [&](Value storeIndex, Value loadIndex) {
+ if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
+ firstToSecondPloopIndices.lookupOrDefault(loadIndex))
+ return failure();
+ else
+ return success();
+ },
+ /*markEquivalent=*/nullptr,
+ OperationEquivalence::Flags::IgnoreLocations)) {
+ return WalkResult::interrupt();
+ }
+ } else
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+ return !walkResult.wasInterrupted();
+}
+
+/// Analyzes dependencies in the most primitive way by checking simple read and
+/// write patterns.
+static LogicalResult
+verifyDependencies(scf::ParallelOp firstPloop, scf::ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ if (!haveNoReadsAfterWriteExceptSameIndex(
+ firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))
+ return failure();
+
+ IRMapping secondToFirstPloopIndices;
+ secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
+ firstPloop.getBody()->getArguments());
+ return success(haveNoReadsAfterWriteExceptSameIndex(
+ secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));
+}
+
+/// Verify equal iteration spaces.
+static bool equalIterationSpaces(scf::ParallelOp firstPloop,
+ scf::ParallelOp secondPloop) {
+ if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
+ return false;
+
+ auto matchOperands = [&](const OperandRange &lhs,
+ const OperandRange &rhs) -> bool {
+ // TODO: Extend this to support aliases and equal constants.
+ return std::equal(lhs.begin(), lhs.end(), rhs.begin());
+ };
+ return matchOperands(firstPloop.getLowerBound(),
+ secondPloop.getLowerBound()) &&
+ matchOperands(firstPloop.getUpperBound(),
+ secondPloop.getUpperBound()) &&
+ matchOperands(firstPloop.getStep(), secondPloop.getStep());
+}
+
+/// Verify there are no nested ParallelOps.
+static bool hasNestedParallelOp(scf::ParallelOp ploop) {
+ auto walkResult = ploop.getBody()->walk(
+ [](scf::ParallelOp) { return WalkResult::interrupt(); });
+ return walkResult.wasInterrupted();
+}
+
+static bool isFusionLegal(scf::ParallelOp firstPloop,
+ scf::ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ return !hasNestedParallelOp(firstPloop) &&
+ !hasNestedParallelOp(secondPloop) &&
+ equalIterationSpaces(firstPloop, secondPloop) &&
+ succeeded(verifyDependencies(firstPloop, secondPloop,
+ firstToSecondPloopIndices, mayAlias));
+}
+
+void mlir::fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,
+ OpBuilder builder,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ Block *block1 = firstPloop.getBody();
+ Block *block2 = secondPloop.getBody();
+ IRMapping firstToSecondPloopIndices;
+ firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
+
+ if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
+ mayAlias))
+ return;
+
+ DominanceInfo dom;
+ // We are fusing first loop into second, make sure there are no users of the
+ // first loop results between loops.
+ for (Operation *user : firstPloop->getUsers())
+ if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
+ return;
+
+ ValueRange inits1 = firstPloop.getInitVals();
+ ValueRange inits2 = secondPloop.getInitVals();
+
+ SmallVector<Value> newInitVars(inits1.begin(), inits1.end());
+ newInitVars.append(inits2.begin(), inits2.end());
+
+ IRRewriter b(builder);
+ b.setInsertionPoint(secondPloop);
+ auto newSecondPloop = b.create<scf::ParallelOp>(
+ secondPloop.getLoc(), secondPloop.getLowerBound(),
+ secondPloop.getUpperBound(), secondPloop.getStep(), newInitVars);
+
+ Block *newBlock = newSecondPloop.getBody();
+ auto term1 = cast<scf::ReduceOp>(block1->getTerminator());
+ auto term2 = cast<scf::ReduceOp>(block2->getTerminator());
+
+ b.inlineBlockBefore(block2, newBlock, newBlock->begin(),
+ newBlock->getArguments());
+ b.inlineBlockBefore(block1, newBlock, newBlock->begin(),
+ newBlock->getArguments());
+
+ ValueRange results = newSecondPloop.getResults();
+ if (!results.empty()) {
+ b.setInsertionPointToEnd(newBlock);
+
+ ValueRange reduceArgs1 = term1.getOperands();
+ ValueRange reduceArgs2 = term2.getOperands();
+ SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
+ newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());
+
+ auto newReduceOp = b.create<scf::ReduceOp>(term2.getLoc(), newReduceArgs);
+
+ for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
+ term1.getReductions(), term2.getReductions()))) {
+ Block &oldRedBlock = reg.front();
+ Block &newRedBlock = newReduceOp.getReductions()[i].front();
+ b.inlineBlockBefore(&oldRedBlock, &newRedBlock, newRedBlock.begin(),
+ newRedBlock.getArguments());
+ }
+
+ firstPloop.replaceAllUsesWith(results.take_front(inits1.size()));
+ secondPloop.replaceAllUsesWith(results.take_back(inits2.size()));
+ }
+ term1->erase();
+ term2->erase();
+ firstPloop.erase();
+ secondPloop.erase();
+ secondPloop = newSecondPloop;
+}
+
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
@@ -1171,3 +1372,10 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
return fusedLoop;
}
+
+scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
+ scf::ParallelOp target, scf::ParallelOp source, RewriterBase &rewriter) {
+ auto mayAlias = [&](Value val1, Value val2) -> bool { return false; };
+ mlir::fuseIfLegal(target, source, rewriter, mayAlias);
+ return source;
+}
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 0f51b1cdbe0cf..46c6be36c3271 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -47,6 +47,59 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: func @fuse_two_parallel
+// CHECK-SAME: ([[A:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}) {
+func.func @fuse_two_parallel(%A: memref<2x2xf32>, %B: memref<2x2xf32>) {
+// CHECK-DAG: [[C2:%.*]] = arith.constant 2 : index
+// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index
+// CHECK-DAG: [[C1:%.*]] = arith.constant 1 : index
+// CHECK-DAG: [[C1FP:%.*]] = arith.constant 1.
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c1fp = arith.constant 1.0 : f32
+// CHECK: [[SUM:%.*]] = memref.alloc()
+ %sum = memref.alloc() : memref<2x2xf32>
+// CHECK: scf.parallel ([[I:%.*]], [[J:%.*]]) = ([[C0]], [[C0]])
+// CHECK-SAME: to ([[C2]], [[C2]]) step ([[C1]], [[C1]]) {
+// CHECK: [[B_ELEM:%.*]] = memref.load [[B]]{{\[}}[[I]], [[J]]]
+// CHECK: [[SUM_ELEM:%.*]] = arith.addf [[B_ELEM]], [[C1FP]]
+// CHECK: memref.store [[SUM_ELEM]], [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK-NOT: scf.parallel
+// CHECK: [[SUM_ELEM_:%.*]] = memref.load [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK: [[A_ELEM:%.*]] = memref.load [[A]]{{\[}}[[I]], [[J]]]
+// CHECK: [[PRODUCT_ELEM:%.*]] = arith.mulf [[SUM_ELEM_]], [[A_ELEM]]
+// CHECK: memref.store [[PRODUCT_ELEM]], [[B]]{{\[}}[[I]], [[J]]]
+// CHECK: scf.reduce
+// CHECK: }
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
+ %sum_elem = arith.addf %B_elem, %c1fp : f32
+ memref.store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
+ scf.reduce
+ }
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %sum_elem = memref.load %sum[%i, %j] : memref<2x2xf32>
+ %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
+ %product_elem = arith.mulf %sum_elem, %A_elem : f32
+ memref.store %product_elem, %B[%i, %j] : memref<2x2xf32>
+ scf.reduce
+ }
+// CHECK: memref.dealloc [[SUM]]
+ memref.dealloc %sum : memref<2x2xf32>
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.parallel"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %parallel:2 = transform.split_handle %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %fused = transform.loop.fuse_sibling %parallel#0 into %parallel#1 : (!transform.any_op,!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK: func.func @fuse_2nd_for_into_1st([[A:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
func.func @fuse_2nd_for_into_1st(%A: tensor<128xf32>, %B: tensor<128xf32>) -> (tensor<128xf32>, tensor<128xf32>) {
// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index
>From b73238a9472b0682f250e37848ad504d21a57059 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 5 Jun 2024 10:47:00 -0500
Subject: [PATCH 03/34] add checkFusionStructuralLegality
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 7 ++++++
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 26 +++++++++++++++++++++
2 files changed, 33 insertions(+)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 2944d8ffac022..834857f177cdf 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -156,6 +156,13 @@ Loops tilePerfectlyNested(scf::ForOp rootForOp, ArrayRef<Value> sizes);
void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
scf::ForOp root);
+//===----------------------------------------------------------------------===//
+// Fusion related helpers
+//===----------------------------------------------------------------------===//
+
+template <typename LoopTy>
+bool checkFusionStructuralLegality(Operation *target, Operation *source);
+
/// Prepends operations of firstPloop's body into secondPloop's body.
/// Updates secondPloop with new loop.
void fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index d85339f32dbe3..c490983335470 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1184,6 +1184,10 @@ static bool equalIterationSpaces(scf::ParallelOp firstPloop,
matchOperands(firstPloop.getStep(), secondPloop.getStep());
}
+//===----------------------------------------------------------------------===//
+// Fusion related helpers
+//===----------------------------------------------------------------------===//
+
/// Verify there are no nested ParallelOps.
static bool hasNestedParallelOp(scf::ParallelOp ploop) {
auto walkResult = ploop.getBody()->walk(
@@ -1191,6 +1195,28 @@ static bool hasNestedParallelOp(scf::ParallelOp ploop) {
return walkResult.wasInterrupted();
}
+template <typename LoopTy>
+static bool checkFusionStructuralLegality(Operation *target,
+ Operation *source) {
+ static_assert(llvm::is_one_of<LoopTy, scf::ForallOp, scf::ForOp,
+ scf::ParallelOp>::value,
+ "applies to only `forall`, `for` and `parallel`");
+ auto targetOp = dyn_cast<LoopTy>(target);
+ auto sourceOp = dyn_cast<LoopTy>(source);
+ if (!targetOp || !sourceOp)
+ return false;
+
+ if constexpr (std::is_same_v<LoopTy, scf::ForallOp>)
+ return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
+ targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
+ targetOp.getMixedStep() == sourceOp.getMixedStep() &&
+ targetOp.getMapping() == sourceOp.getMapping();
+ else
+ return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
+ targetOp.getUpperBound() == sourceOp.getUpperBound() &&
+ targetOp.getStep() == sourceOp.getStep();
+}
+
static bool isFusionLegal(scf::ParallelOp firstPloop,
scf::ParallelOp secondPloop,
const IRMapping &firstToSecondPloopIndices,
>From f5bbd131bb7713ae47a58587b2d9acf82dc3b12f Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 5 Jun 2024 14:19:56 -0500
Subject: [PATCH 04/34] replace isLoopWithIdenticalConfiguration with
checkFusionStructuralLegality
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 4 +-
.../SCF/TransformOps/SCFTransformOps.cpp | 50 ++++++-------------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 32 +++++-------
3 files changed, 29 insertions(+), 57 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index 834857f177cdf..ab9d154aa480d 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -160,8 +160,8 @@ void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
// Fusion related helpers
//===----------------------------------------------------------------------===//
-template <typename LoopTy>
-bool checkFusionStructuralLegality(Operation *target, Operation *source);
+bool checkFusionStructuralLegality(LoopLikeOpInterface &target,
+ LoopLikeOpInterface &source);
/// Prepends operations of firstPloop's body into secondPloop's body.
/// Updates secondPloop with new loop.
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 1c53e89d69040..9f541b94af474 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -442,34 +442,6 @@ static DiagnosedSilenceableFailure isOpSibling(Operation *target,
return DiagnosedSilenceableFailure::success();
}
-/// Check if `target` scf loop can be fused into `source` scf loop.
-/// Applies for scf.for, scf.forall, and scf.parallel.
-///
-/// This simply checks if both loops have the same bounds, steps and mapping.
-/// No attempt is made at checking that the side effects of `target` and
-/// `source` are independent of each other.
-template <typename LoopTy>
-static bool isLoopWithIdenticalConfiguration(Operation *target,
- Operation *source) {
- static_assert(llvm::is_one_of<LoopTy, scf::ForallOp, scf::ForOp,
- scf::ParallelOp>::value,
- "applies to only `forall`, `for` and `parallel`");
- auto targetOp = dyn_cast<LoopTy>(target);
- auto sourceOp = dyn_cast<LoopTy>(source);
- if (!targetOp || !sourceOp)
- return false;
-
- if constexpr (std::is_same_v<LoopTy, scf::ForallOp>)
- return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
- targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
- targetOp.getMixedStep() == sourceOp.getMixedStep() &&
- targetOp.getMapping() == sourceOp.getMapping();
- else
- return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
- targetOp.getUpperBound() == sourceOp.getUpperBound() &&
- targetOp.getStep() == sourceOp.getStep();
-}
-
DiagnosedSilenceableFailure
transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
transform::TransformResults &results,
@@ -485,29 +457,37 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
<< "source handle (got " << llvm::range_size(sourceOps) << ")";
}
- Operation *target = *targetOps.begin();
- Operation *source = *sourceOps.begin();
+ LoopLikeOpInterface target =
+ dyn_cast<LoopLikeOpInterface>(*targetOps.begin());
+ LoopLikeOpInterface source =
+ dyn_cast<LoopLikeOpInterface>(*sourceOps.begin());
+ if (!target || !source)
+ return emitSilenceableFailure(target->getLoc())
+ << "target or source is not a loop op";
// Check if the target and source are siblings.
DiagnosedSilenceableFailure diag = isOpSibling(target, source);
if (!diag.succeeded())
return diag;
+ if (!mlir::checkFusionStructuralLegality(target, source))
+ return emitSilenceableFailure(target->getLoc())
+ << "operations cannot be fused";
+
Operation *fusedLoop;
/// TODO: Support fusion for loop-like ops besides scf.for and scf.forall.
- if (isLoopWithIdenticalConfiguration<scf::ForOp>(target, source)) {
+ if (isa<scf::ForOp>(target) && isa<scf::ForOp>(source)) {
fusedLoop = fuseIndependentSiblingForLoops(
cast<scf::ForOp>(target), cast<scf::ForOp>(source), rewriter);
- } else if (isLoopWithIdenticalConfiguration<scf::ForallOp>(target, source)) {
+ } else if (isa<scf::ForallOp>(target) && isa<scf::ForallOp>(source)) {
fusedLoop = fuseIndependentSiblingForallLoops(
cast<scf::ForallOp>(target), cast<scf::ForallOp>(source), rewriter);
- } else if (isLoopWithIdenticalConfiguration<scf::ParallelOp>(target,
- source)) {
+ } else if (isa<scf::ParallelOp>(target) && isa<scf::ParallelOp>(source)) {
fusedLoop = fuseIndependentSiblingParallelLoops(
cast<scf::ParallelOp>(target), cast<scf::ParallelOp>(source), rewriter);
} else
return emitSilenceableFailure(target->getLoc())
- << "operations cannot be fused";
+ << "unsupported loop type for fusion";
assert(fusedLoop && "failed to fuse operations");
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index c490983335470..ce20730459c2a 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1195,26 +1195,18 @@ static bool hasNestedParallelOp(scf::ParallelOp ploop) {
return walkResult.wasInterrupted();
}
-template <typename LoopTy>
-static bool checkFusionStructuralLegality(Operation *target,
- Operation *source) {
- static_assert(llvm::is_one_of<LoopTy, scf::ForallOp, scf::ForOp,
- scf::ParallelOp>::value,
- "applies to only `forall`, `for` and `parallel`");
- auto targetOp = dyn_cast<LoopTy>(target);
- auto sourceOp = dyn_cast<LoopTy>(source);
- if (!targetOp || !sourceOp)
- return false;
-
- if constexpr (std::is_same_v<LoopTy, scf::ForallOp>)
- return targetOp.getMixedLowerBound() == sourceOp.getMixedLowerBound() &&
- targetOp.getMixedUpperBound() == sourceOp.getMixedUpperBound() &&
- targetOp.getMixedStep() == sourceOp.getMixedStep() &&
- targetOp.getMapping() == sourceOp.getMapping();
- else
- return targetOp.getLowerBound() == sourceOp.getLowerBound() &&
- targetOp.getUpperBound() == sourceOp.getUpperBound() &&
- targetOp.getStep() == sourceOp.getStep();
+bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
+ LoopLikeOpInterface &source) {
+ auto iterSpaceEq =
+ target.getMixedLowerBound() == source.getMixedLowerBound() &&
+ target.getMixedUpperBound() == source.getMixedUpperBound() &&
+ target.getMixedStep() == source.getMixedStep();
+ auto forAllTarget = dyn_cast<scf::ForallOp>(*target);
+ auto forAllSource = dyn_cast<scf::ForallOp>(*source);
+ if (forAllTarget && forAllSource)
+ return iterSpaceEq &&
+ forAllTarget.getMapping() == forAllSource.getMapping();
+ return iterSpaceEq;
}
static bool isFusionLegal(scf::ParallelOp firstPloop,
>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 05/34] 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 06/34] 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 07/34] 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();
>From 009fd15ab8abefd56afe6424e27f912a4166329d Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 6 Jun 2024 14:02:52 -0500
Subject: [PATCH 08/34] 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<Value> threadIds = forallOp.getInductionVars();
SmallVector<OpFoldResult> 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<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
}
SmallVector<Value> ForallOp::getInductionVars() {
- return SmallVector<Value>(getBody()->getArguments().take_front(getRank()));
+ return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
// Get lower bounds as OpFoldResult.
@@ -3005,7 +3005,7 @@ void ParallelOp::print(OpAsmPrinter &p) {
SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
SmallVector<Value> ParallelOp::getInductionVars() {
- return SmallVector<Value>(getBody()->getArguments());
+ return SmallVector<Value>{getBody()->getArguments()};
}
std::optional<SmallVector<OpFoldResult>> ParallelOp::getLowerBounds() {
>From d34ad95aba669b5700976f0d2ed4d68b4902e9be Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 6 Jun 2024 15:26:06 -0500
Subject: [PATCH 09/34] 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<LoopLikeOpInterface,
- ["getInductionVars", "getLowerBounds", "getSteps",
- "getUpperBounds", "getYieldedValuesMutable",
+ ["getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
+ "getLoopUpperBounds", "getYieldedValuesMutable",
"replaceWithAdditionalYields"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands"]>]> {
@@ -671,7 +671,7 @@ def AffineParallelOp : Affine_Op<"parallel",
I32ElementsAttr:$lowerBoundsGroups,
AffineMapAttr:$upperBoundsMap,
I32ElementsAttr:$upperBoundsGroups,
- I64SmallVectorArrayAttr:$step,
+ I64SmallVectorArrayAttr:$steps,
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>":$step)>
+ "ArrayRef<int64_t>":$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<LoopLikeOpInterface,
["getInitsMutable", "getLoopResults", "getRegionIterArgs",
- "getInductionVars", "getLowerBounds", "getSteps",
- "getUpperBounds", "getYieldedValuesMutable",
+ "getLoopInductionVars", "getLoopLowerBounds", "getLoopSteps",
+ "getLoopUpperBounds", "getYieldedValuesMutable",
"promoteIfSingleIteration", "replaceWithAdditionalYields",
"yieldTiledValuesAndReplace"]>,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
@@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getInitsMutable", "getRegionIterArgs", "getInductionVars",
- "getLowerBounds", "getUpperBounds", "getSteps",
+ ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
@@ -510,23 +510,26 @@ def ForallOp : SCF_Op<"forall", [
];
let extraClassDeclaration = [{
+ SmallVector<Value> getInductionVars() {
+ return getLoopInductionVars();
+ }
// Get lower bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedLowerBound() {
- auto maybeLowerBounds = getLowerBounds();
+ auto maybeLowerBounds = getLoopLowerBounds();
assert(maybeLowerBounds.has_value() && "expected values");
return *maybeLowerBounds;
}
// Get upper bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedUpperBound() {
- auto maybeUpperBounds = getUpperBounds();
+ auto maybeUpperBounds = getLoopUpperBounds();
assert(maybeUpperBounds.has_value() && "expected values");
return *maybeUpperBounds;
}
// Get steps as OpFoldResult.
SmallVector<OpFoldResult> 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<RegionBranchOpInterface, [
def ParallelOp : SCF_Op<"parallel",
[AutomaticAllocationScope,
AttrSizedOperandSegments,
- DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getInductionVars",
- "getLowerBounds", "getUpperBounds", "getSteps"]>,
+ DeclareOpInterfaceMethods<LoopLikeOpInterface, ["getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps"]>,
RecursiveMemoryEffects,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
SingleBlockImplicitTerminator<"scf::ReduceOp">,
@@ -845,6 +848,10 @@ def ParallelOp : SCF_Op<"parallel",
];
let extraClassDeclaration = [{
+ // Get induction variables.
+ SmallVector<Value> 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<AffineParallelOp> {
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<arith::ConstantIndexOp>(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<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
-SmallVector<Value> AffineForOp::getInductionVars() {
+SmallVector<Value> AffineForOp::getLoopInductionVars() {
return {getInductionVar()};
}
-std::optional<SmallVector<OpFoldResult>> AffineForOp::getLowerBounds() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopLowerBounds() {
if (!hasConstantLowerBound())
return std::nullopt;
OpBuilder b(getContext());
@@ -2466,13 +2466,13 @@ std::optional<SmallVector<OpFoldResult>> AffineForOp::getLowerBounds() {
OpFoldResult(b.getI64IntegerAttr(getConstantLowerBound()))};
}
-std::optional<SmallVector<OpFoldResult>> AffineForOp::getSteps() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopSteps() {
OpBuilder b(getContext());
return SmallVector<OpFoldResult>{
OpFoldResult(b.getI64IntegerAttr(getStepAsInt()))};
}
-std::optional<SmallVector<OpFoldResult>> AffineForOp::getUpperBounds() {
+std::optional<SmallVector<OpFoldResult>> AffineForOp::getLoopUpperBounds() {
if (!hasConstantUpperBound())
return {};
OpBuilder b(getContext());
@@ -3758,7 +3758,7 @@ SmallVector<Region *> 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<int64_t> 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<int64_t, 8> steps = getStep();
+ SmallVector<int64_t, 8> 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<int64_t, 8> steps = op.getStep();
+ SmallVector<int64_t, 8> 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<Value> ForOp::getInductionVars() { return {getInductionVar()}; }
+SmallVector<Value> ForOp::getLoopInductionVars() { return {getInductionVar()}; }
-std::optional<SmallVector<OpFoldResult>> ForOp::getLowerBounds() {
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {
return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
}
-std::optional<SmallVector<OpFoldResult>> ForOp::getSteps() {
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopSteps() {
return SmallVector<OpFoldResult, 1>{OpFoldResult(getStep())};
}
-std::optional<SmallVector<OpFoldResult>> ForOp::getUpperBounds() {
+std::optional<SmallVector<OpFoldResult>> ForOp::getLoopUpperBounds() {
return SmallVector<OpFoldResult, 1>{OpFoldResult(getUpperBound())};
}
@@ -1426,24 +1426,24 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-SmallVector<Value> ForallOp::getInductionVars() {
+SmallVector<Value> ForallOp::getLoopInductionVars() {
return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
// Get lower bounds as OpFoldResult.
-std::optional<SmallVector<OpFoldResult>> ForallOp::getLowerBounds() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopLowerBounds() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticLowerBound(), getDynamicLowerBound(), b);
}
// Get upper bounds as OpFoldResult.
-std::optional<SmallVector<OpFoldResult>> ForallOp::getUpperBounds() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopUpperBounds() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticUpperBound(), getDynamicUpperBound(), b);
}
// Get steps as OpFoldResult.
-std::optional<SmallVector<OpFoldResult>> ForallOp::getSteps() {
+std::optional<SmallVector<OpFoldResult>> ForallOp::getLoopSteps() {
Builder b(getOperation()->getContext());
return getMixedValues(getStaticStep(), getDynamicStep(), b);
}
@@ -3004,19 +3004,19 @@ void ParallelOp::print(OpAsmPrinter &p) {
SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
-SmallVector<Value> ParallelOp::getInductionVars() {
+SmallVector<Value> ParallelOp::getLoopInductionVars() {
return SmallVector<Value>{getBody()->getArguments()};
}
-std::optional<SmallVector<OpFoldResult>> ParallelOp::getLowerBounds() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopLowerBounds() {
return getLowerBound();
}
-std::optional<SmallVector<OpFoldResult>> ParallelOp::getUpperBounds() {
+std::optional<SmallVector<OpFoldResult>> ParallelOp::getLoopUpperBounds() {
return getUpperBound();
}
-std::optional<SmallVector<OpFoldResult>> ParallelOp::getSteps() {
+std::optional<SmallVector<OpFoldResult>> 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<SmallVector<OpFoldResult>> maybeLb =
- loopLikeOp.getLowerBounds();
+ loopLikeOp.getLoopLowerBounds();
EXPECT_TRUE(maybeLb.has_value());
EXPECT_EQ((*maybeLb).size(), 1u);
std::optional<SmallVector<OpFoldResult>> maybeUb =
- loopLikeOp.getUpperBounds();
+ loopLikeOp.getLoopUpperBounds();
EXPECT_TRUE(maybeUb.has_value());
EXPECT_EQ((*maybeUb).size(), 1u);
- std::optional<SmallVector<OpFoldResult>> maybeStep = loopLikeOp.getSteps();
+ std::optional<SmallVector<OpFoldResult>> 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<SmallVector<OpFoldResult>> maybeLb =
- loopLikeOp.getLowerBounds();
+ loopLikeOp.getLoopLowerBounds();
EXPECT_TRUE(maybeLb.has_value());
EXPECT_EQ((*maybeLb).size(), 2u);
std::optional<SmallVector<OpFoldResult>> maybeUb =
- loopLikeOp.getUpperBounds();
+ loopLikeOp.getLoopUpperBounds();
EXPECT_TRUE(maybeUb.has_value());
EXPECT_EQ((*maybeUb).size(), 2u);
- std::optional<SmallVector<OpFoldResult>> maybeStep = loopLikeOp.getSteps();
+ std::optional<SmallVector<OpFoldResult>> 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 <srcarroll314 at gmail.com>
Date: Thu, 6 Jun 2024 16:04:47 -0500
Subject: [PATCH 10/34] 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<Value> getInductionVars() {
- return getLoopInductionVars();
+ auto maybeInductionVars = getLoopInductionVars();;
+ assert(maybeInductionVars.has_value() && "expected values");
+ return *maybeInductionVars;
}
// Get lower bounds as OpFoldResult.
SmallVector<OpFoldResult> 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<Value> 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<Region *> AffineForOp::getLoopRegions() { return {&getRegion()}; }
-SmallVector<Value> AffineForOp::getLoopInductionVars() {
- return {getInductionVar()};
+std::optional<SmallVector<Value>> AffineForOp::getLoopInductionVars() {
+ return SmallVector<Value>{getInductionVar()};
}
std::optional<SmallVector<OpFoldResult>> 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<Value> ForOp::getLoopInductionVars() { return {getInductionVar()}; }
+std::optional<SmallVector<Value>> ForOp::getLoopInductionVars() {
+ return SmallVector<Value, 1>{getInductionVar()};
+}
std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {
return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
@@ -1426,7 +1428,7 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-SmallVector<Value> ForallOp::getLoopInductionVars() {
+std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {
return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
@@ -3004,7 +3006,7 @@ void ParallelOp::print(OpAsmPrinter &p) {
SmallVector<Region *> ParallelOp::getLoopRegions() { return {&getRegion()}; }
-SmallVector<Value> ParallelOp::getLoopInductionVars() {
+std::optional<SmallVector<Value>> ParallelOp::getLoopInductionVars() {
return SmallVector<Value>{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<SmallVector<Value>> 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<SmallVector<Value>> 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 <srcarroll314 at gmail.com>
Date: Fri, 7 Jun 2024 11:45:44 -0500
Subject: [PATCH 11/34] 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<Value> getInductionVars() {
- auto maybeInductionVars = getLoopInductionVars();;
+ std::optional<SmallVector<Value>> maybeInductionVars = getLoopInductionVars();
assert(maybeInductionVars.has_value() && "expected values");
return *maybeInductionVars;
}
- // Get lower bounds as OpFoldResult.
+ /// Get lower bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedLowerBound() {
- auto maybeLowerBounds = getLoopLowerBounds();
+ std::optional<SmallVector<OpFoldResult>> maybeLowerBounds = getLoopLowerBounds();
assert(maybeLowerBounds.has_value() && "expected values");
return *maybeLowerBounds;
}
- // Get upper bounds as OpFoldResult.
+ /// Get upper bounds as OpFoldResult.
SmallVector<OpFoldResult> getMixedUpperBound() {
- auto maybeUpperBounds = getLoopUpperBounds();
+ std::optional<SmallVector<OpFoldResult>> maybeUpperBounds = getLoopUpperBounds();
assert(maybeUpperBounds.has_value() && "expected values");
return *maybeUpperBounds;
}
- // Get steps as OpFoldResult.
+ /// Get steps as OpFoldResult.
SmallVector<OpFoldResult> getMixedStep() {
- auto maybeSteps = getLoopSteps();
+ std::optional<SmallVector<OpFoldResult>> 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<Value> getInductionVars() {
- auto maybeInductionVars = getLoopInductionVars();;
+ std::optional<SmallVector<Value>> 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<SmallVector<Value>> ForOp::getLoopInductionVars() {
- return SmallVector<Value, 1>{getInductionVar()};
+ return SmallVector<Value>{getInductionVar()};
}
std::optional<SmallVector<OpFoldResult>> ForOp::getLoopLowerBounds() {
- return SmallVector<OpFoldResult, 1>{OpFoldResult(getLowerBound())};
+ return SmallVector<OpFoldResult>{OpFoldResult(getLowerBound())};
}
std::optional<SmallVector<OpFoldResult>> ForOp::getLoopSteps() {
- return SmallVector<OpFoldResult, 1>{OpFoldResult(getStep())};
+ return SmallVector<OpFoldResult>{OpFoldResult(getStep())};
}
std::optional<SmallVector<OpFoldResult>> ForOp::getLoopUpperBounds() {
- return SmallVector<OpFoldResult, 1>{OpFoldResult(getUpperBound())};
+ return SmallVector<OpFoldResult>{OpFoldResult(getUpperBound())};
}
std::optional<ResultRange> 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<SmallVector<OpFoldResult>> maybeLb =
loopLikeOp.getLoopLowerBounds();
- EXPECT_TRUE(maybeLb.has_value());
+ ASSERT_TRUE(maybeLb.has_value());
EXPECT_EQ((*maybeLb).size(), 1u);
std::optional<SmallVector<OpFoldResult>> maybeUb =
loopLikeOp.getLoopUpperBounds();
- EXPECT_TRUE(maybeUb.has_value());
+ ASSERT_TRUE(maybeUb.has_value());
EXPECT_EQ((*maybeUb).size(), 1u);
std::optional<SmallVector<OpFoldResult>> maybeStep =
loopLikeOp.getLoopSteps();
- EXPECT_TRUE(maybeStep.has_value());
+ ASSERT_TRUE(maybeStep.has_value());
EXPECT_EQ((*maybeStep).size(), 1u);
std::optional<SmallVector<Value>> 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<SmallVector<OpFoldResult>> maybeLb =
loopLikeOp.getLoopLowerBounds();
- EXPECT_TRUE(maybeLb.has_value());
+ ASSERT_TRUE(maybeLb.has_value());
EXPECT_EQ((*maybeLb).size(), 2u);
std::optional<SmallVector<OpFoldResult>> maybeUb =
loopLikeOp.getLoopUpperBounds();
- EXPECT_TRUE(maybeUb.has_value());
+ ASSERT_TRUE(maybeUb.has_value());
EXPECT_EQ((*maybeUb).size(), 2u);
std::optional<SmallVector<OpFoldResult>> maybeStep =
loopLikeOp.getLoopSteps();
- EXPECT_TRUE(maybeStep.has_value());
+ ASSERT_TRUE(maybeStep.has_value());
EXPECT_EQ((*maybeStep).size(), 2u);
std::optional<SmallVector<Value>> maybeInductionVars =
loopLikeOp.getLoopInductionVars();
- EXPECT_TRUE(maybeInductionVars.has_value());
+ ASSERT_TRUE(maybeInductionVars.has_value());
EXPECT_EQ((*maybeInductionVars).size(), 2u);
}
>From 6336fdf28f06c7525bcf6822a386f8c4cabe3c2d Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Fri, 7 Jun 2024 15:25:40 -0500
Subject: [PATCH 12/34] update after rebase
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index ce20730459c2a..e3660e89fb684 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1198,9 +1198,9 @@ static bool hasNestedParallelOp(scf::ParallelOp ploop) {
bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
LoopLikeOpInterface &source) {
auto iterSpaceEq =
- target.getMixedLowerBound() == source.getMixedLowerBound() &&
- target.getMixedUpperBound() == source.getMixedUpperBound() &&
- target.getMixedStep() == source.getMixedStep();
+ target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
+ target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
+ target.getLoopSteps() == source.getLoopSteps();
auto forAllTarget = dyn_cast<scf::ForallOp>(*target);
auto forAllSource = dyn_cast<scf::ForallOp>(*source);
if (forAllTarget && forAllSource)
>From 86406c335dd216ac91a941102c060bc680af10b1 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Sat, 8 Jun 2024 23:31:57 -0500
Subject: [PATCH 13/34] refactor main parallel fusion logic from fuseIfLegal to
util func
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 6 -
.../SCF/Transforms/ParallelLoopFusion.cpp | 159 +++++++++++++++++-
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 140 ++++++---------
3 files changed, 208 insertions(+), 97 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index ab9d154aa480d..ac4434b337890 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -163,12 +163,6 @@ void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
bool checkFusionStructuralLegality(LoopLikeOpInterface &target,
LoopLikeOpInterface &source);
-/// Prepends operations of firstPloop's body into secondPloop's body.
-/// Updates secondPloop with new loop.
-void fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,
- OpBuilder builder,
- llvm::function_ref<bool(Value, Value)> mayAlias);
-
/// Given two scf.forall loops, `target` and `source`, fuses `target` into
/// `source`. Assumes that the given loops are siblings and are independent of
/// each other.
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index abac91cfaf7d9..326a8f93162b9 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -31,6 +31,163 @@ namespace mlir {
using namespace mlir;
using namespace mlir::scf;
+/// Verify there are no nested ParallelOps.
+static bool hasNestedParallelOp(ParallelOp ploop) {
+ auto walkResult =
+ ploop.getBody()->walk([](ParallelOp) { return WalkResult::interrupt(); });
+ return walkResult.wasInterrupted();
+}
+
+/// Verify equal iteration spaces.
+static bool equalIterationSpaces(ParallelOp firstPloop,
+ ParallelOp secondPloop) {
+ if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
+ return false;
+
+ auto matchOperands = [&](const OperandRange &lhs,
+ const OperandRange &rhs) -> bool {
+ // TODO: Extend this to support aliases and equal constants.
+ return std::equal(lhs.begin(), lhs.end(), rhs.begin());
+ };
+ return matchOperands(firstPloop.getLowerBound(),
+ secondPloop.getLowerBound()) &&
+ matchOperands(firstPloop.getUpperBound(),
+ secondPloop.getUpperBound()) &&
+ matchOperands(firstPloop.getStep(), secondPloop.getStep());
+}
+
+/// Checks if the parallel loops have mixed access to the same buffers. Returns
+/// `true` if the first parallel loop writes to the same indices that the second
+/// loop reads.
+static bool haveNoReadsAfterWriteExceptSameIndex(
+ ParallelOp firstPloop, ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;
+ SmallVector<Value> bufferStoresVec;
+ firstPloop.getBody()->walk([&](memref::StoreOp store) {
+ bufferStores[store.getMemRef()].push_back(store.getIndices());
+ bufferStoresVec.emplace_back(store.getMemRef());
+ });
+ auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {
+ Value loadMem = load.getMemRef();
+ // Stop if the memref is defined in secondPloop body. Careful alias analysis
+ // is needed.
+ auto *memrefDef = loadMem.getDefiningOp();
+ if (memrefDef && memrefDef->getBlock() == load->getBlock())
+ return WalkResult::interrupt();
+
+ for (Value store : bufferStoresVec)
+ if (store != loadMem && mayAlias(store, loadMem))
+ return WalkResult::interrupt();
+
+ auto write = bufferStores.find(loadMem);
+ if (write == bufferStores.end())
+ return WalkResult::advance();
+
+ // Check that at last one store was retrieved
+ if (!write->second.size())
+ return WalkResult::interrupt();
+
+ auto storeIndices = write->second.front();
+
+ // Multiple writes to the same memref are allowed only on the same indices
+ for (const auto &othStoreIndices : write->second) {
+ if (othStoreIndices != storeIndices)
+ return WalkResult::interrupt();
+ }
+
+ // Check that the load indices of secondPloop coincide with store indices of
+ // firstPloop for the same memrefs.
+ auto loadIndices = load.getIndices();
+ if (storeIndices.size() != loadIndices.size())
+ return WalkResult::interrupt();
+ for (int i = 0, e = storeIndices.size(); i < e; ++i) {
+ if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
+ loadIndices[i]) {
+ auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
+ auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
+ if (storeIndexDefOp && loadIndexDefOp) {
+ if (!isMemoryEffectFree(storeIndexDefOp))
+ return WalkResult::interrupt();
+ if (!isMemoryEffectFree(loadIndexDefOp))
+ return WalkResult::interrupt();
+ if (!OperationEquivalence::isEquivalentTo(
+ storeIndexDefOp, loadIndexDefOp,
+ [&](Value storeIndex, Value loadIndex) {
+ if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
+ firstToSecondPloopIndices.lookupOrDefault(loadIndex))
+ return failure();
+ else
+ return success();
+ },
+ /*markEquivalent=*/nullptr,
+ OperationEquivalence::Flags::IgnoreLocations)) {
+ return WalkResult::interrupt();
+ }
+ } else
+ return WalkResult::interrupt();
+ }
+ }
+ return WalkResult::advance();
+ });
+ return !walkResult.wasInterrupted();
+}
+
+/// Analyzes dependencies in the most primitive way by checking simple read and
+/// write patterns.
+static LogicalResult
+verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ if (!haveNoReadsAfterWriteExceptSameIndex(
+ firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))
+ return failure();
+
+ IRMapping secondToFirstPloopIndices;
+ secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
+ firstPloop.getBody()->getArguments());
+ return success(haveNoReadsAfterWriteExceptSameIndex(
+ secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));
+}
+
+static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
+ const IRMapping &firstToSecondPloopIndices,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ return !hasNestedParallelOp(firstPloop) &&
+ !hasNestedParallelOp(secondPloop) &&
+ equalIterationSpaces(firstPloop, secondPloop) &&
+ succeeded(verifyDependencies(firstPloop, secondPloop,
+ firstToSecondPloopIndices, mayAlias));
+}
+
+/// Prepends operations of firstPloop's body into secondPloop's body.
+/// Updates secondPloop with new loop.
+static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
+ OpBuilder builder,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
+ Block *block1 = firstPloop.getBody();
+ Block *block2 = secondPloop.getBody();
+ IRMapping firstToSecondPloopIndices;
+ firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
+
+ if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
+ mayAlias))
+ return;
+
+ DominanceInfo dom;
+ // We are fusing first loop into second, make sure there are no users of the
+ // first loop results between loops.
+ for (Operation *user : firstPloop->getUsers())
+ if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
+ return;
+
+ IRRewriter rewriter(builder);
+ secondPloop = mlir::fuseIndependentSiblingParallelLoops(
+ firstPloop, secondPloop, rewriter);
+ ;
+}
+
void mlir::scf::naivelyFuseParallelOps(
Region ®ion, llvm::function_ref<bool(Value, Value)> mayAlias) {
OpBuilder b(region);
@@ -59,7 +216,7 @@ void mlir::scf::naivelyFuseParallelOps(
}
for (MutableArrayRef<ParallelOp> ploops : ploopChains) {
for (int i = 0, e = ploops.size(); i + 1 < e; ++i)
- mlir::fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);
+ fuseIfLegal(ploops[i], ploops[i + 1], b, mayAlias);
}
}
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index e3660e89fb684..5f58767be409d 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1188,13 +1188,6 @@ static bool equalIterationSpaces(scf::ParallelOp firstPloop,
// Fusion related helpers
//===----------------------------------------------------------------------===//
-/// Verify there are no nested ParallelOps.
-static bool hasNestedParallelOp(scf::ParallelOp ploop) {
- auto walkResult = ploop.getBody()->walk(
- [](scf::ParallelOp) { return WalkResult::interrupt(); });
- return walkResult.wasInterrupted();
-}
-
bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
LoopLikeOpInterface &source) {
auto iterSpaceEq =
@@ -1209,86 +1202,6 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
return iterSpaceEq;
}
-static bool isFusionLegal(scf::ParallelOp firstPloop,
- scf::ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- return !hasNestedParallelOp(firstPloop) &&
- !hasNestedParallelOp(secondPloop) &&
- equalIterationSpaces(firstPloop, secondPloop) &&
- succeeded(verifyDependencies(firstPloop, secondPloop,
- firstToSecondPloopIndices, mayAlias));
-}
-
-void mlir::fuseIfLegal(scf::ParallelOp firstPloop, scf::ParallelOp &secondPloop,
- OpBuilder builder,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- Block *block1 = firstPloop.getBody();
- Block *block2 = secondPloop.getBody();
- IRMapping firstToSecondPloopIndices;
- firstToSecondPloopIndices.map(block1->getArguments(), block2->getArguments());
-
- if (!isFusionLegal(firstPloop, secondPloop, firstToSecondPloopIndices,
- mayAlias))
- return;
-
- DominanceInfo dom;
- // We are fusing first loop into second, make sure there are no users of the
- // first loop results between loops.
- for (Operation *user : firstPloop->getUsers())
- if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
- return;
-
- ValueRange inits1 = firstPloop.getInitVals();
- ValueRange inits2 = secondPloop.getInitVals();
-
- SmallVector<Value> newInitVars(inits1.begin(), inits1.end());
- newInitVars.append(inits2.begin(), inits2.end());
-
- IRRewriter b(builder);
- b.setInsertionPoint(secondPloop);
- auto newSecondPloop = b.create<scf::ParallelOp>(
- secondPloop.getLoc(), secondPloop.getLowerBound(),
- secondPloop.getUpperBound(), secondPloop.getStep(), newInitVars);
-
- Block *newBlock = newSecondPloop.getBody();
- auto term1 = cast<scf::ReduceOp>(block1->getTerminator());
- auto term2 = cast<scf::ReduceOp>(block2->getTerminator());
-
- b.inlineBlockBefore(block2, newBlock, newBlock->begin(),
- newBlock->getArguments());
- b.inlineBlockBefore(block1, newBlock, newBlock->begin(),
- newBlock->getArguments());
-
- ValueRange results = newSecondPloop.getResults();
- if (!results.empty()) {
- b.setInsertionPointToEnd(newBlock);
-
- ValueRange reduceArgs1 = term1.getOperands();
- ValueRange reduceArgs2 = term2.getOperands();
- SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
- newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());
-
- auto newReduceOp = b.create<scf::ReduceOp>(term2.getLoc(), newReduceArgs);
-
- for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
- term1.getReductions(), term2.getReductions()))) {
- Block &oldRedBlock = reg.front();
- Block &newRedBlock = newReduceOp.getReductions()[i].front();
- b.inlineBlockBefore(&oldRedBlock, &newRedBlock, newRedBlock.begin(),
- newRedBlock.getArguments());
- }
-
- firstPloop.replaceAllUsesWith(results.take_front(inits1.size()));
- secondPloop.replaceAllUsesWith(results.take_back(inits2.size()));
- }
- term1->erase();
- term2->erase();
- firstPloop.erase();
- secondPloop.erase();
- secondPloop = newSecondPloop;
-}
-
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
@@ -1393,7 +1306,54 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
scf::ParallelOp target, scf::ParallelOp source, RewriterBase &rewriter) {
- auto mayAlias = [&](Value val1, Value val2) -> bool { return false; };
- mlir::fuseIfLegal(target, source, rewriter, mayAlias);
- return source;
+ Block *block1 = target.getBody();
+ Block *block2 = source.getBody();
+ auto term1 = cast<scf::ReduceOp>(block1->getTerminator());
+ auto term2 = cast<scf::ReduceOp>(block2->getTerminator());
+
+ ValueRange inits1 = target.getInitVals();
+ ValueRange inits2 = source.getInitVals();
+
+ SmallVector<Value> newInitVars(inits1.begin(), inits1.end());
+ newInitVars.append(inits2.begin(), inits2.end());
+
+ rewriter.setInsertionPoint(source);
+ auto fusedLoop = rewriter.create<scf::ParallelOp>(
+ source.getLoc(), source.getLowerBound(), source.getUpperBound(),
+ source.getStep(), newInitVars);
+ Block *newBlock = fusedLoop.getBody();
+ rewriter.inlineBlockBefore(block2, newBlock, newBlock->begin(),
+ newBlock->getArguments());
+ rewriter.inlineBlockBefore(block1, newBlock, newBlock->begin(),
+ newBlock->getArguments());
+
+ ValueRange results = fusedLoop.getResults();
+ if (!results.empty()) {
+ rewriter.setInsertionPointToEnd(newBlock);
+
+ ValueRange reduceArgs1 = term1.getOperands();
+ ValueRange reduceArgs2 = term2.getOperands();
+ SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
+ newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());
+
+ auto newReduceOp =
+ rewriter.create<scf::ReduceOp>(term2.getLoc(), newReduceArgs);
+
+ for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
+ term1.getReductions(), term2.getReductions()))) {
+ Block &oldRedBlock = reg.front();
+ Block &newRedBlock = newReduceOp.getReductions()[i].front();
+ rewriter.inlineBlockBefore(&oldRedBlock, &newRedBlock,
+ newRedBlock.begin(),
+ newRedBlock.getArguments());
+ }
+ target.replaceAllUsesWith(results.take_front(inits1.size()));
+ source.replaceAllUsesWith(results.take_back(inits2.size()));
+ }
+ term1->erase();
+ term2->erase();
+ target.erase();
+ source.erase();
+
+ return fusedLoop;
}
>From 694d589dc535892f3dda9d27c2a43052fc0b445e Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Sat, 8 Jun 2024 23:35:30 -0500
Subject: [PATCH 14/34] remove unused functions
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 113 ---------------------------
1 file changed, 113 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 5f58767be409d..e6cb88c427da8 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1071,119 +1071,6 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
return tileLoops;
}
-/// Checks if the parallel loops have mixed access to the same buffers. Returns
-/// `true` if the first parallel loop writes to the same indices that the second
-/// loop reads.
-static bool haveNoReadsAfterWriteExceptSameIndex(
- scf::ParallelOp firstPloop, scf::ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- DenseMap<Value, SmallVector<ValueRange, 1>> bufferStores;
- SmallVector<Value> bufferStoresVec;
- firstPloop.getBody()->walk([&](memref::StoreOp store) {
- bufferStores[store.getMemRef()].push_back(store.getIndices());
- bufferStoresVec.emplace_back(store.getMemRef());
- });
- auto walkResult = secondPloop.getBody()->walk([&](memref::LoadOp load) {
- Value loadMem = load.getMemRef();
- // Stop if the memref is defined in secondPloop body. Careful alias analysis
- // is needed.
- auto *memrefDef = loadMem.getDefiningOp();
- if (memrefDef && memrefDef->getBlock() == load->getBlock())
- return WalkResult::interrupt();
-
- for (Value store : bufferStoresVec)
- if (store != loadMem && mayAlias(store, loadMem))
- return WalkResult::interrupt();
-
- auto write = bufferStores.find(loadMem);
- if (write == bufferStores.end())
- return WalkResult::advance();
-
- // Check that at last one store was retrieved
- if (!write->second.size())
- return WalkResult::interrupt();
-
- auto storeIndices = write->second.front();
-
- // Multiple writes to the same memref are allowed only on the same indices
- for (const auto &othStoreIndices : write->second) {
- if (othStoreIndices != storeIndices)
- return WalkResult::interrupt();
- }
-
- // Check that the load indices of secondPloop coincide with store indices of
- // firstPloop for the same memrefs.
- auto loadIndices = load.getIndices();
- if (storeIndices.size() != loadIndices.size())
- return WalkResult::interrupt();
- for (int i = 0, e = storeIndices.size(); i < e; ++i) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndices[i]) !=
- loadIndices[i]) {
- auto *storeIndexDefOp = storeIndices[i].getDefiningOp();
- auto *loadIndexDefOp = loadIndices[i].getDefiningOp();
- if (storeIndexDefOp && loadIndexDefOp) {
- if (!isMemoryEffectFree(storeIndexDefOp))
- return WalkResult::interrupt();
- if (!isMemoryEffectFree(loadIndexDefOp))
- return WalkResult::interrupt();
- if (!OperationEquivalence::isEquivalentTo(
- storeIndexDefOp, loadIndexDefOp,
- [&](Value storeIndex, Value loadIndex) {
- if (firstToSecondPloopIndices.lookupOrDefault(storeIndex) !=
- firstToSecondPloopIndices.lookupOrDefault(loadIndex))
- return failure();
- else
- return success();
- },
- /*markEquivalent=*/nullptr,
- OperationEquivalence::Flags::IgnoreLocations)) {
- return WalkResult::interrupt();
- }
- } else
- return WalkResult::interrupt();
- }
- }
- return WalkResult::advance();
- });
- return !walkResult.wasInterrupted();
-}
-
-/// Analyzes dependencies in the most primitive way by checking simple read and
-/// write patterns.
-static LogicalResult
-verifyDependencies(scf::ParallelOp firstPloop, scf::ParallelOp secondPloop,
- const IRMapping &firstToSecondPloopIndices,
- llvm::function_ref<bool(Value, Value)> mayAlias) {
- if (!haveNoReadsAfterWriteExceptSameIndex(
- firstPloop, secondPloop, firstToSecondPloopIndices, mayAlias))
- return failure();
-
- IRMapping secondToFirstPloopIndices;
- secondToFirstPloopIndices.map(secondPloop.getBody()->getArguments(),
- firstPloop.getBody()->getArguments());
- return success(haveNoReadsAfterWriteExceptSameIndex(
- secondPloop, firstPloop, secondToFirstPloopIndices, mayAlias));
-}
-
-/// Verify equal iteration spaces.
-static bool equalIterationSpaces(scf::ParallelOp firstPloop,
- scf::ParallelOp secondPloop) {
- if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
- return false;
-
- auto matchOperands = [&](const OperandRange &lhs,
- const OperandRange &rhs) -> bool {
- // TODO: Extend this to support aliases and equal constants.
- return std::equal(lhs.begin(), lhs.end(), rhs.begin());
- };
- return matchOperands(firstPloop.getLowerBound(),
- secondPloop.getLowerBound()) &&
- matchOperands(firstPloop.getUpperBound(),
- secondPloop.getUpperBound()) &&
- matchOperands(firstPloop.getStep(), secondPloop.getStep());
-}
-
//===----------------------------------------------------------------------===//
// Fusion related helpers
//===----------------------------------------------------------------------===//
>From 67cb64f1a773795bfb2d4e9f0c981dd502572676 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Sun, 9 Jun 2024 14:45:57 -0500
Subject: [PATCH 15/34] refactor fuseIndependentSiblingForLoops to reuse
replaceWithAdditionalYields
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 4 +-
mlir/lib/Dialect/SCF/IR/SCF.cpp | 54 ++++++++
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 151 +++++++++++++++------
3 files changed, 168 insertions(+), 41 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index f35ea962bea16..e7b9665f797fa 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -301,8 +301,8 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
- "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
+ ["getInitsMutable", "getRegionIterArgs", "getLoopResults", "getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps", "replaceWithAdditionalYields",
"promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 5e94f4dc612a7..6ad181e2f3d77 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -616,8 +616,62 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,
regions.push_back(RegionSuccessor(getResults()));
}
+std::optional<ResultRange> ForallOp::getLoopResults() { return getResults(); }
+
SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }
+FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
+ RewriterBase &rewriter, ValueRange newInitOperands,
+ bool replaceInitOperandUsesInLoop,
+ const NewYieldValuesFn &newYieldValuesFn) {
+ // Create a new loop before the existing one, with the extra operands.
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(getOperation());
+ auto inits = llvm::to_vector(getOutputs());
+ inits.append(newInitOperands.begin(), newInitOperands.end());
+ scf::ForallOp newLoop = rewriter.create<scf::ForallOp>(
+ getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
+ inits, getMapping());
+
+ // Generate the new yield values and append them to the scf.yield operation.
+ auto yieldOp = cast<scf::InParallelOp>(getTerminator());
+ ArrayRef<BlockArgument> newIterArgs =
+ newLoop.getBody()->getArguments().take_back(newInitOperands.size());
+ {
+ OpBuilder::InsertionGuard g(rewriter);
+ rewriter.setInsertionPoint(yieldOp);
+ SmallVector<Value> newYieldedValues =
+ newYieldValuesFn(rewriter, getLoc(), newIterArgs);
+ assert(newInitOperands.size() == newYieldedValues.size() &&
+ "expected as many new yield values as new iter operands");
+ // rewriter.modifyOpInPlace(yieldOp, [&]() {
+ // yieldOp.getResultsMutable().append(newYieldedValues);
+ // });
+ }
+
+ // Move the loop body to the new op.
+ rewriter.mergeBlocks(getBody(), newLoop.getBody(),
+ newLoop.getBody()->getArguments().take_front(
+ getBody()->getNumArguments()));
+
+ if (replaceInitOperandUsesInLoop) {
+ // Replace all uses of `newInitOperands` with the corresponding basic block
+ // arguments.
+ for (auto it : llvm::zip(newInitOperands, newIterArgs)) {
+ rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),
+ [&](OpOperand &use) {
+ Operation *user = use.getOwner();
+ return newLoop->isProperAncestor(user);
+ });
+ }
+ }
+
+ // Replace the old loop.
+ rewriter.replaceOp(getOperation(),
+ newLoop->getResults().take_front(getNumResults()));
+ return cast<LoopLikeOpInterface>(newLoop.getOperation());
+}
+
/// Promotes the loop body of a forallOp to its containing block if it can be
/// determined that the loop has a single iteration.
LogicalResult scf::ForallOp::promoteIfSingleIteration(RewriterBase &rewriter) {
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index e6cb88c427da8..a61428208c405 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1089,6 +1089,92 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
return iterSpaceEq;
}
+template <typename LoopTy>
+void fuseTerminator(RewriterBase &rewriter, LoopTy target, LoopTy source,
+ LoopTy &fused, IRMapping &mapping) {}
+
+template <>
+void fuseTerminator(RewriterBase &rewriter, scf::ForallOp target,
+ scf::ForallOp source, scf::ForallOp &fused,
+ IRMapping &mapping) {
+ // Fuse the old terminator in_parallel ops into the new one.
+ scf::InParallelOp targetTerm = target.getTerminator();
+ scf::InParallelOp sourceTerm = source.getTerminator();
+ scf::InParallelOp fusedTerm = fused.getTerminator();
+ rewriter.setInsertionPointToStart(fusedTerm.getBody());
+ for (Operation &op : targetTerm.getYieldingOps())
+ rewriter.clone(op, mapping);
+ for (Operation &op : sourceTerm.getYieldingOps())
+ rewriter.clone(op, mapping);
+}
+
+template <>
+void fuseTerminator(RewriterBase &rewriter, scf::ForOp target,
+ scf::ForOp source, scf::ForOp &fused, IRMapping &mapping) {
+ // Build fused yield results by appropriately mapping original yield operands.
+ SmallVector<Value> yieldResults;
+ for (Value operand : target.getBody()->getTerminator()->getOperands())
+ yieldResults.push_back(mapping.lookupOrDefault(operand));
+ for (Value operand : source.getBody()->getTerminator()->getOperands())
+ yieldResults.push_back(mapping.lookupOrDefault(operand));
+ if (!yieldResults.empty())
+ rewriter.create<scf::YieldOp>(source.getLoc(), yieldResults);
+}
+
+template <typename LoopTy>
+LoopTy createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
+ RewriterBase &rewriter) {
+ auto targetResults = target.getLoopResults();
+ auto sourceResults = source.getLoopResults();
+ int64_t numTargetOuts = (*targetResults).size();
+ int64_t numSourceOuts = (*sourceResults).size();
+ printf("numTargetOuts %ld\n", numTargetOuts);
+
+ // Create fused shared_outs.
+ SmallVector<Value> fusedOuts;
+ llvm::append_range(fusedOuts, *targetResults);
+ llvm::append_range(fusedOuts, *sourceResults);
+
+ // Create a new scf.forall op after the source loop.
+ rewriter.setInsertionPointAfter(source);
+ // LoopTy fusedLoop = builder.create<LoopTy>(
+ // source.getLoc(), source.getLoopLowerBounds(),
+ // source.getLoopUpperBounds(), source.getLoopSteps(), fusedOuts,
+ // source->getAttrs());
+ LoopTy fusedLoop = rewriter.cloneWithoutRegions(cast<LoopTy>(source));
+
+ // Map control operands.
+ IRMapping mapping;
+ mapping.map(*target.getLoopInductionVars(),
+ *fusedLoop.getLoopInductionVars());
+ mapping.map(*source.getLoopInductionVars(),
+ *fusedLoop.getLoopInductionVars());
+
+ // Map shared outs.
+ mapping.map(target.getRegionIterArgs(),
+ fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
+ mapping.map(source.getRegionIterArgs(),
+ fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
+
+ // Append everything except the terminator into the fused operation.
+ rewriter.setInsertionPointToStart(
+ &fusedLoop.getLoopRegions().front()->front());
+ for (Operation &op :
+ target.getLoopRegions().front()->front().without_terminator())
+ rewriter.clone(op, mapping);
+ for (Operation &op :
+ source.getLoopRegions().front()->front().without_terminator())
+ rewriter.clone(op, mapping);
+
+ fuseTerminator<LoopTy>(rewriter, cast<LoopTy>(target), cast<LoopTy>(source),
+ cast<LoopTy>(fusedLoop), mapping);
+
+ rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
+ rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
+
+ return fusedLoop;
+}
+
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
@@ -1144,50 +1230,37 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
RewriterBase &rewriter) {
- unsigned numTargetOuts = target.getNumResults();
- unsigned numSourceOuts = source.getNumResults();
-
- // Create fused init_args, with target's init_args before source's init_args.
- SmallVector<Value> fusedInitArgs;
- llvm::append_range(fusedInitArgs, target.getInitArgs());
- llvm::append_range(fusedInitArgs, source.getInitArgs());
-
- // Create a new scf.for op after the source loop (with scf.yield terminator
- // (without arguments) only in case its init_args is empty).
- rewriter.setInsertionPointAfter(source);
- scf::ForOp fusedLoop = rewriter.create<scf::ForOp>(
- source.getLoc(), source.getLowerBound(), source.getUpperBound(),
- source.getStep(), fusedInitArgs);
-
+ auto targetIterArgs = target.getRegionIterArgs();
+ auto targetInductionVar = target.getInductionVar();
+ SmallVector<Value> targetYieldOperands(source.getYieldedValues());
+ auto sourceIterArgs = source.getRegionIterArgs();
+ auto sourceInductionVar = source.getInductionVar();
+ SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
+ scf::ForOp fusedLoop = cast<scf::ForOp>(*target.replaceWithAdditionalYields(
+ rewriter, source.getInitArgs(), /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
+ return sourceYieldOperands;
+ }));
// Map original induction variables and operands to those of the fused loop.
IRMapping mapping;
- mapping.map(target.getInductionVar(), fusedLoop.getInductionVar());
- mapping.map(target.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
- mapping.map(source.getInductionVar(), fusedLoop.getInductionVar());
- mapping.map(source.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
-
+ mapping.map(targetInductionVar, fusedLoop.getInductionVar());
+ mapping.map(targetIterArgs,
+ fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
+ mapping.map(targetYieldOperands,
+ fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
+ mapping.map(sourceInductionVar, fusedLoop.getInductionVar());
+ mapping.map(sourceIterArgs,
+ fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
+ mapping.map(sourceYieldOperands,
+ fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
// Merge target's body into the new (fused) for loop and then source's body.
- rewriter.setInsertionPointToStart(fusedLoop.getBody());
- for (Operation &op : target.getBody()->without_terminator())
- rewriter.clone(op, mapping);
+ rewriter.setInsertionPoint(fusedLoop.getBody()->getTerminator());
for (Operation &op : source.getBody()->without_terminator())
rewriter.clone(op, mapping);
-
- // Build fused yield results by appropriately mapping original yield operands.
- SmallVector<Value> yieldResults;
- for (Value operand : target.getBody()->getTerminator()->getOperands())
- yieldResults.push_back(mapping.lookupOrDefault(operand));
- for (Value operand : source.getBody()->getTerminator()->getOperands())
- yieldResults.push_back(mapping.lookupOrDefault(operand));
- if (!yieldResults.empty())
- rewriter.create<scf::YieldOp>(source.getLoc(), yieldResults);
-
- // Replace old loops by substituting their uses by results of the fused loop.
- rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
- rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
-
+ auto newTerm = rewriter.clone(*fusedLoop.getBody()->getTerminator(), mapping);
+ rewriter.replaceOp(fusedLoop.getBody()->getTerminator(), newTerm);
+ rewriter.replaceOp(source,
+ fusedLoop.getResults().take_back(source.getNumResults()));
return fusedLoop;
}
>From cc8599f69a90f9b460bbff950505004c214ac72e Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Sun, 9 Jun 2024 17:20:46 -0500
Subject: [PATCH 16/34] refactor fuseIndependentSiblingForallLoops to reuse
replaceWithAdditionalYields
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 12 +---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 69 +++++++++----------
.../SCF/transform-loop-fuse-sibling.mlir | 3 +-
3 files changed, 36 insertions(+), 48 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 6ad181e2f3d77..6850d632f10d0 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -637,17 +637,7 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
auto yieldOp = cast<scf::InParallelOp>(getTerminator());
ArrayRef<BlockArgument> newIterArgs =
newLoop.getBody()->getArguments().take_back(newInitOperands.size());
- {
- OpBuilder::InsertionGuard g(rewriter);
- rewriter.setInsertionPoint(yieldOp);
- SmallVector<Value> newYieldedValues =
- newYieldValuesFn(rewriter, getLoc(), newIterArgs);
- assert(newInitOperands.size() == newYieldedValues.size() &&
- "expected as many new yield values as new iter operands");
- // rewriter.modifyOpInPlace(yieldOp, [&]() {
- // yieldOp.getResultsMutable().append(newYieldedValues);
- // });
- }
+ newLoop.getTerminator().erase();
// Move the loop body to the new op.
rewriter.mergeBlocks(getBody(), newLoop.getBody(),
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index a61428208c405..a822b4199fe9d 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1128,7 +1128,6 @@ LoopTy createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
auto sourceResults = source.getLoopResults();
int64_t numTargetOuts = (*targetResults).size();
int64_t numSourceOuts = (*sourceResults).size();
- printf("numTargetOuts %ld\n", numTargetOuts);
// Create fused shared_outs.
SmallVector<Value> fusedOuts;
@@ -1178,51 +1177,49 @@ LoopTy createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
- unsigned numTargetOuts = target.getNumResults();
- unsigned numSourceOuts = source.getNumResults();
-
- // Create fused shared_outs.
- SmallVector<Value> fusedOuts;
- llvm::append_range(fusedOuts, target.getOutputs());
- llvm::append_range(fusedOuts, source.getOutputs());
-
- // Create a new scf.forall op after the source loop.
- rewriter.setInsertionPointAfter(source);
- scf::ForallOp fusedLoop = rewriter.create<scf::ForallOp>(
- source.getLoc(), source.getMixedLowerBound(), source.getMixedUpperBound(),
- source.getMixedStep(), fusedOuts, source.getMapping());
-
+ auto targetIterArgs = target.getRegionIterArgs();
+ auto targetInductionVar = target.getInductionVars();
+ SmallVector<Value> targetYieldOperands(target.getYieldedValues());
+ auto sourceIterArgs = source.getRegionIterArgs();
+ auto sourceInductionVar = source.getInductionVars();
+ scf::InParallelOp sourceTerm = source.getTerminator();
+ auto sourceYieldOps = sourceTerm.getYieldingOps();
+ auto sourceBody = source.getBody();
+ SmallVector<Value> sourceYieldOperands(llvm::map_range(
+ sourceTerm.getDests(), [](auto arg) { return cast<Value>(arg); }));
+ scf::ForallOp fusedLoop =
+ cast<scf::ForallOp>(*target.replaceWithAdditionalYields(
+ rewriter, source.getOutputs(), /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
+ for (Operation &op : sourceYieldOps)
+ b.clone(op);
+ return sourceYieldOperands;
+ }));
// Map control operands.
IRMapping mapping;
- mapping.map(target.getInductionVars(), fusedLoop.getInductionVars());
- mapping.map(source.getInductionVars(), fusedLoop.getInductionVars());
-
- // Map shared outs.
- mapping.map(target.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
- mapping.map(source.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
-
+ mapping.map(targetInductionVar, fusedLoop.getInductionVars());
+ mapping.map(targetIterArgs,
+ fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
+ mapping.map(targetYieldOperands,
+ fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
+ mapping.map(sourceInductionVar, fusedLoop.getInductionVars());
+ mapping.map(sourceIterArgs,
+ fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
+ mapping.map(sourceYieldOperands,
+ fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
// Append everything except the terminator into the fused operation.
- rewriter.setInsertionPointToStart(fusedLoop.getBody());
- for (Operation &op : target.getBody()->without_terminator())
- rewriter.clone(op, mapping);
- for (Operation &op : source.getBody()->without_terminator())
+ rewriter.setInsertionPoint(fusedLoop.getBody()->getTerminator());
+ for (Operation &op : sourceBody->without_terminator())
rewriter.clone(op, mapping);
// Fuse the old terminator in_parallel ops into the new one.
- scf::InParallelOp targetTerm = target.getTerminator();
- scf::InParallelOp sourceTerm = source.getTerminator();
scf::InParallelOp fusedTerm = fusedLoop.getTerminator();
- rewriter.setInsertionPointToStart(fusedTerm.getBody());
- for (Operation &op : targetTerm.getYieldingOps())
- rewriter.clone(op, mapping);
+ rewriter.setInsertionPointToEnd(fusedTerm.getBody());
for (Operation &op : sourceTerm.getYieldingOps())
rewriter.clone(op, mapping);
- // Replace old loops by substituting their uses by results of the fused loop.
- rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
- rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
+ rewriter.replaceOp(source,
+ fusedLoop.getResults().take_back(source.getNumResults()));
return fusedLoop;
}
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 46c6be36c3271..47bfe0baa7651 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -189,7 +189,8 @@ module attributes {transform.with_named_sequence} {
// -----
-// CHECK: func.func @matmul_fuse_2nd_forall_into_1st([[A1:%.*]]: {{.*}}, [[A2:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
+// CHECK-LABEL: func.func @matmul_fuse_2nd_forall_into_1st
+// CHECK-SAME: [[A1:%.*]]: {{.*}}, [[A2:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
func.func @matmul_fuse_2nd_forall_into_1st(%A1 : tensor<128x128xf32>, %A2 : tensor<128x128xf32>, %B : tensor<128x128xf32>) -> (tensor<128x128xf32>, tensor<128x128xf32>) {
%zero = arith.constant 0.0 : f32
%out_alloc = tensor.empty() : tensor<128x128xf32>
>From 48b1af9cb4392b8ccad748e17ce40fa997db6a59 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Sun, 9 Jun 2024 19:12:40 -0500
Subject: [PATCH 17/34] wip
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 4 +-
mlir/lib/Dialect/SCF/IR/SCF.cpp | 4 +
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 170 ++++++---------------
3 files changed, 50 insertions(+), 128 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index e7b9665f797fa..b9345f6ecdbb2 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -302,8 +302,8 @@ def ForallOp : SCF_Op<"forall", [
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
["getInitsMutable", "getRegionIterArgs", "getLoopResults", "getLoopInductionVars",
- "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps", "replaceWithAdditionalYields",
- "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps", "getYieldedValuesMutable",
+ "replaceWithAdditionalYields", "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 6850d632f10d0..b4a16e519a15a 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -1472,6 +1472,10 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
+std::optional<MutableArrayRef<OpOperand>> ForallOp::getYieldedValuesMutable() {
+ return getOutputsMutable();
+}
+
std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {
return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index a822b4199fe9d..fb2d1d11fb6ae 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1090,134 +1090,76 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
}
template <typename LoopTy>
-void fuseTerminator(RewriterBase &rewriter, LoopTy target, LoopTy source,
- LoopTy &fused, IRMapping &mapping) {}
+void fuseTerminator(RewriterBase &rewriter, LoopTy source, LoopTy &fused,
+ IRMapping &mapping) {}
template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForallOp target,
- scf::ForallOp source, scf::ForallOp &fused,
- IRMapping &mapping) {
+void fuseTerminator(RewriterBase &rewriter, scf::ForallOp source,
+ scf::ForallOp &fused, IRMapping &mapping) {
// Fuse the old terminator in_parallel ops into the new one.
- scf::InParallelOp targetTerm = target.getTerminator();
- scf::InParallelOp sourceTerm = source.getTerminator();
scf::InParallelOp fusedTerm = fused.getTerminator();
- rewriter.setInsertionPointToStart(fusedTerm.getBody());
- for (Operation &op : targetTerm.getYieldingOps())
- rewriter.clone(op, mapping);
- for (Operation &op : sourceTerm.getYieldingOps())
+ rewriter.setInsertionPointToEnd(fusedTerm.getBody());
+ for (Operation &op : source.getTerminator().getYieldingOps())
rewriter.clone(op, mapping);
}
template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForOp target,
- scf::ForOp source, scf::ForOp &fused, IRMapping &mapping) {
+void fuseTerminator(RewriterBase &rewriter, scf::ForOp source,
+ scf::ForOp &fused, IRMapping &mapping) {
// Build fused yield results by appropriately mapping original yield operands.
- SmallVector<Value> yieldResults;
- for (Value operand : target.getBody()->getTerminator()->getOperands())
- yieldResults.push_back(mapping.lookupOrDefault(operand));
- for (Value operand : source.getBody()->getTerminator()->getOperands())
- yieldResults.push_back(mapping.lookupOrDefault(operand));
- if (!yieldResults.empty())
- rewriter.create<scf::YieldOp>(source.getLoc(), yieldResults);
+ auto newTerm = rewriter.clone(*fused.getBody()->getTerminator(), mapping);
+ rewriter.replaceOp(fused.getBody()->getTerminator(), newTerm);
}
template <typename LoopTy>
-LoopTy createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
- RewriterBase &rewriter) {
- auto targetResults = target.getLoopResults();
- auto sourceResults = source.getLoopResults();
- int64_t numTargetOuts = (*targetResults).size();
- int64_t numSourceOuts = (*sourceResults).size();
-
- // Create fused shared_outs.
- SmallVector<Value> fusedOuts;
- llvm::append_range(fusedOuts, *targetResults);
- llvm::append_range(fusedOuts, *sourceResults);
-
- // Create a new scf.forall op after the source loop.
- rewriter.setInsertionPointAfter(source);
- // LoopTy fusedLoop = builder.create<LoopTy>(
- // source.getLoc(), source.getLoopLowerBounds(),
- // source.getLoopUpperBounds(), source.getLoopSteps(), fusedOuts,
- // source->getAttrs());
- LoopTy fusedLoop = rewriter.cloneWithoutRegions(cast<LoopTy>(source));
-
- // Map control operands.
- IRMapping mapping;
- mapping.map(*target.getLoopInductionVars(),
- *fusedLoop.getLoopInductionVars());
- mapping.map(*source.getLoopInductionVars(),
- *fusedLoop.getLoopInductionVars());
-
- // Map shared outs.
- mapping.map(target.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_front(numTargetOuts));
- mapping.map(source.getRegionIterArgs(),
- fusedLoop.getRegionIterArgs().take_back(numSourceOuts));
-
- // Append everything except the terminator into the fused operation.
- rewriter.setInsertionPointToStart(
- &fusedLoop.getLoopRegions().front()->front());
- for (Operation &op :
- target.getLoopRegions().front()->front().without_terminator())
- rewriter.clone(op, mapping);
- for (Operation &op :
- source.getLoopRegions().front()->front().without_terminator())
- rewriter.clone(op, mapping);
-
- fuseTerminator<LoopTy>(rewriter, cast<LoopTy>(target), cast<LoopTy>(source),
- cast<LoopTy>(fusedLoop), mapping);
-
- rewriter.replaceOp(target, fusedLoop.getResults().take_front(numTargetOuts));
- rewriter.replaceOp(source, fusedLoop.getResults().take_back(numSourceOuts));
-
- return fusedLoop;
-}
-
-scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
- scf::ForallOp source,
- RewriterBase &rewriter) {
+LoopLikeOpInterface
+createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
+ RewriterBase &rewriter, NewYieldValuesFn newYieldValuesFn) {
auto targetIterArgs = target.getRegionIterArgs();
- auto targetInductionVar = target.getInductionVars();
+ auto targetInductionVar = *target.getLoopInductionVars();
SmallVector<Value> targetYieldOperands(target.getYieldedValues());
auto sourceIterArgs = source.getRegionIterArgs();
- auto sourceInductionVar = source.getInductionVars();
- scf::InParallelOp sourceTerm = source.getTerminator();
- auto sourceYieldOps = sourceTerm.getYieldingOps();
- auto sourceBody = source.getBody();
- SmallVector<Value> sourceYieldOperands(llvm::map_range(
- sourceTerm.getDests(), [](auto arg) { return cast<Value>(arg); }));
- scf::ForallOp fusedLoop =
- cast<scf::ForallOp>(*target.replaceWithAdditionalYields(
- rewriter, source.getOutputs(), /*replaceInitOperandUsesInLoop=*/false,
- [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
- for (Operation &op : sourceYieldOps)
- b.clone(op);
- return sourceYieldOperands;
- }));
+ auto sourceInductionVar = *source.getLoopInductionVars();
+ SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
+ auto sourceRegion = source.getLoopRegions().front();
+ LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
+ rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
+ newYieldValuesFn);
+
// Map control operands.
IRMapping mapping;
- mapping.map(targetInductionVar, fusedLoop.getInductionVars());
+ mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
mapping.map(targetIterArgs,
fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
mapping.map(targetYieldOperands,
fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
- mapping.map(sourceInductionVar, fusedLoop.getInductionVars());
+ mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
mapping.map(sourceIterArgs,
fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
mapping.map(sourceYieldOperands,
fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
// Append everything except the terminator into the fused operation.
- rewriter.setInsertionPoint(fusedLoop.getBody()->getTerminator());
- for (Operation &op : sourceBody->without_terminator())
+ rewriter.setInsertionPoint(
+ fusedLoop.getLoopRegions().front()->front().getTerminator());
+ for (Operation &op : sourceRegion->front().without_terminator())
rewriter.clone(op, mapping);
- // Fuse the old terminator in_parallel ops into the new one.
- scf::InParallelOp fusedTerm = fusedLoop.getTerminator();
- rewriter.setInsertionPointToEnd(fusedTerm.getBody());
- for (Operation &op : sourceTerm.getYieldingOps())
- rewriter.clone(op, mapping);
+ fuseTerminator<LoopTy>(rewriter, cast<LoopTy>(source),
+ cast<LoopTy>(fusedLoop), mapping);
+
+ return fusedLoop;
+}
+scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
+ scf::ForallOp source,
+ RewriterBase &rewriter) {
+ scf::ForallOp fusedLoop = cast<scf::ForallOp>(createFused<scf::ForallOp>(
+ target, source, rewriter,
+ [&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
+ for (Operation &op : source.getTerminator().getYieldingOps())
+ b.clone(op);
+ return source.getYieldedValues();
+ }));
rewriter.replaceOp(source,
fusedLoop.getResults().take_back(source.getNumResults()));
@@ -1227,35 +1169,11 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
RewriterBase &rewriter) {
- auto targetIterArgs = target.getRegionIterArgs();
- auto targetInductionVar = target.getInductionVar();
- SmallVector<Value> targetYieldOperands(source.getYieldedValues());
- auto sourceIterArgs = source.getRegionIterArgs();
- auto sourceInductionVar = source.getInductionVar();
- SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
- scf::ForOp fusedLoop = cast<scf::ForOp>(*target.replaceWithAdditionalYields(
- rewriter, source.getInitArgs(), /*replaceInitOperandUsesInLoop=*/false,
+ scf::ForOp fusedLoop = cast<scf::ForOp>(createFused<scf::ForOp>(
+ target, source, rewriter,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
- return sourceYieldOperands;
+ return source.getYieldedValues();
}));
- // Map original induction variables and operands to those of the fused loop.
- IRMapping mapping;
- mapping.map(targetInductionVar, fusedLoop.getInductionVar());
- mapping.map(targetIterArgs,
- fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
- mapping.map(targetYieldOperands,
- fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
- mapping.map(sourceInductionVar, fusedLoop.getInductionVar());
- mapping.map(sourceIterArgs,
- fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
- mapping.map(sourceYieldOperands,
- fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
- // Merge target's body into the new (fused) for loop and then source's body.
- rewriter.setInsertionPoint(fusedLoop.getBody()->getTerminator());
- for (Operation &op : source.getBody()->without_terminator())
- rewriter.clone(op, mapping);
- auto newTerm = rewriter.clone(*fusedLoop.getBody()->getTerminator(), mapping);
- rewriter.replaceOp(fusedLoop.getBody()->getTerminator(), newTerm);
rewriter.replaceOp(source,
fusedLoop.getResults().take_back(source.getNumResults()));
return fusedLoop;
>From 7a51cb34afd5d8a2b67cceaef457f50c032affbd Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Mon, 17 Jun 2024 14:36:44 -0500
Subject: [PATCH 18/34] Decouple concrete loop type from `createFused` function
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 35 +++++++++++++++++++++-------
1 file changed, 27 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index fb2d1d11fb6ae..910c41b3e3d54 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1111,10 +1111,29 @@ void fuseTerminator(RewriterBase &rewriter, scf::ForOp source,
rewriter.replaceOp(fused.getBody()->getTerminator(), newTerm);
}
-template <typename LoopTy>
-LoopLikeOpInterface
-createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
- RewriterBase &rewriter, NewYieldValuesFn newYieldValuesFn) {
+// TODO: We should maybe add this as a method to LoopLikeOpInterface.
+// For now, this acts as a placeholder.
+template <>
+void fuseTerminator(RewriterBase &rewriter, LoopLikeOpInterface source,
+ LoopLikeOpInterface &fused, IRMapping &mapping) {
+ if (isa<scf::ForOp>(source) && isa<scf::ForOp>(fused)) {
+ fuseTerminator(rewriter, cast<scf::ForOp>(source), cast<scf::ForOp>(fused),
+ mapping);
+ } else if (isa<scf::ForallOp>(source) && isa<scf::ForallOp>(fused)) {
+ fuseTerminator(rewriter, cast<scf::ForallOp>(source),
+ cast<scf::ForallOp>(fused), mapping);
+ } else if (isa<scf::ParallelOp>(source) && isa<scf::ParallelOp>(fused)) {
+ fuseTerminator(rewriter, cast<scf::ParallelOp>(source),
+ cast<scf::ParallelOp>(fused), mapping);
+ } else {
+ return;
+ }
+}
+
+LoopLikeOpInterface createFused(LoopLikeOpInterface target,
+ LoopLikeOpInterface source,
+ RewriterBase &rewriter,
+ NewYieldValuesFn newYieldValuesFn) {
auto targetIterArgs = target.getRegionIterArgs();
auto targetInductionVar = *target.getLoopInductionVars();
SmallVector<Value> targetYieldOperands(target.getYieldedValues());
@@ -1144,8 +1163,8 @@ createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
for (Operation &op : sourceRegion->front().without_terminator())
rewriter.clone(op, mapping);
- fuseTerminator<LoopTy>(rewriter, cast<LoopTy>(source),
- cast<LoopTy>(fusedLoop), mapping);
+ // TODO: Replace with interface method if added
+ fuseTerminator(rewriter, source, fusedLoop, mapping);
return fusedLoop;
}
@@ -1153,7 +1172,7 @@ createFused(LoopLikeOpInterface target, LoopLikeOpInterface source,
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
- scf::ForallOp fusedLoop = cast<scf::ForallOp>(createFused<scf::ForallOp>(
+ scf::ForallOp fusedLoop = cast<scf::ForallOp>(createFused(
target, source, rewriter,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
for (Operation &op : source.getTerminator().getYieldingOps())
@@ -1169,7 +1188,7 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
scf::ForOp source,
RewriterBase &rewriter) {
- scf::ForOp fusedLoop = cast<scf::ForOp>(createFused<scf::ForOp>(
+ scf::ForOp fusedLoop = cast<scf::ForOp>(createFused(
target, source, rewriter,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
return source.getYieldedValues();
>From 30873263faaab18267109231094af408b819059a Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Mon, 17 Jun 2024 15:10:44 -0500
Subject: [PATCH 19/34] Refactor ForallOp::replaceWithAdditionalYields
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 9 +++------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 5 ++---
2 files changed, 5 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index b4a16e519a15a..c5a9e18e2610c 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -633,12 +633,7 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
inits, getMapping());
- // Generate the new yield values and append them to the scf.yield operation.
- auto yieldOp = cast<scf::InParallelOp>(getTerminator());
- ArrayRef<BlockArgument> newIterArgs =
- newLoop.getBody()->getArguments().take_back(newInitOperands.size());
newLoop.getTerminator().erase();
-
// Move the loop body to the new op.
rewriter.mergeBlocks(getBody(), newLoop.getBody(),
newLoop.getBody()->getArguments().take_front(
@@ -647,7 +642,9 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
if (replaceInitOperandUsesInLoop) {
// Replace all uses of `newInitOperands` with the corresponding basic block
// arguments.
- for (auto it : llvm::zip(newInitOperands, newIterArgs)) {
+ for (auto it :
+ llvm::zip(newInitOperands, newLoop.getBody()->getArguments().take_back(
+ newInitOperands.size()))) {
rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),
[&](OpOperand &use) {
Operation *user = use.getOwner();
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 910c41b3e3d54..5ef6718bc5346 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1175,9 +1175,8 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp fusedLoop = cast<scf::ForallOp>(createFused(
target, source, rewriter,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
- for (Operation &op : source.getTerminator().getYieldingOps())
- b.clone(op);
- return source.getYieldedValues();
+ // `ForallOp` does not have yields, rather an `InParallelOp` terminator.
+ return ValueRange{};
}));
rewriter.replaceOp(source,
fusedLoop.getResults().take_back(source.getNumResults()));
>From bcf3d4aaed9e425f3a3b2d97660c6e816e333abe Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Mon, 17 Jun 2024 15:50:55 -0500
Subject: [PATCH 20/34] revert unnecessary changes
---
mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 7 ++++---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 6 ------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 1 -
3 files changed, 4 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index b9345f6ecdbb2..bf95fbe6721cf 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -301,9 +301,10 @@ def ForallOp : SCF_Op<"forall", [
AttrSizedOperandSegments,
AutomaticAllocationScope,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getInitsMutable", "getRegionIterArgs", "getLoopResults", "getLoopInductionVars",
- "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps", "getYieldedValuesMutable",
- "replaceWithAdditionalYields", "promoteIfSingleIteration", "yieldTiledValuesAndReplace"]>,
+ ["getInitsMutable", "getRegionIterArgs", "getLoopInductionVars",
+ "getLoopLowerBounds", "getLoopUpperBounds", "getLoopSteps",
+ "replaceWithAdditionalYields", "promoteIfSingleIteration",
+ "yieldTiledValuesAndReplace"]>,
RecursiveMemoryEffects,
SingleBlockImplicitTerminator<"scf::InParallelOp">,
DeclareOpInterfaceMethods<RegionBranchOpInterface>,
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index c5a9e18e2610c..deface43028b1 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -616,8 +616,6 @@ void ForOp::getSuccessorRegions(RegionBranchPoint point,
regions.push_back(RegionSuccessor(getResults()));
}
-std::optional<ResultRange> ForallOp::getLoopResults() { return getResults(); }
-
SmallVector<Region *> ForallOp::getLoopRegions() { return {&getRegion()}; }
FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
@@ -1469,10 +1467,6 @@ SmallVector<Operation *> ForallOp::getCombiningOps(BlockArgument bbArg) {
return storeOps;
}
-std::optional<MutableArrayRef<OpOperand>> ForallOp::getYieldedValuesMutable() {
- return getOutputsMutable();
-}
-
std::optional<SmallVector<Value>> ForallOp::getLoopInductionVars() {
return SmallVector<Value>{getBody()->getArguments().take_front(getRank())};
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 5ef6718bc5346..2e61f9998a7d8 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -15,7 +15,6 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/IRMapping.h"
>From 0cb3c4ea08b22eea318fa47634914f921f08f7f2 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Tue, 18 Jun 2024 10:35:53 -0500
Subject: [PATCH 21/34] cleanup
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 4 ++--
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 6 +++---
mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir | 3 +--
3 files changed, 6 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index deface43028b1..2baef9ca45db1 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -625,13 +625,13 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
// Create a new loop before the existing one, with the extra operands.
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(getOperation());
- auto inits = llvm::to_vector(getOutputs());
+ SmallVector<Value> inits(getOutputs());
inits.append(newInitOperands.begin(), newInitOperands.end());
scf::ForallOp newLoop = rewriter.create<scf::ForallOp>(
getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
inits, getMapping());
- newLoop.getTerminator().erase();
+ rewriter.eraseOp(newLoop.getTerminator());
// Move the loop body to the new op.
rewriter.mergeBlocks(getBody(), newLoop.getBody(),
newLoop.getBody()->getArguments().take_front(
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 2e61f9998a7d8..09da6e6233ffc 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1110,8 +1110,8 @@ void fuseTerminator(RewriterBase &rewriter, scf::ForOp source,
rewriter.replaceOp(fused.getBody()->getTerminator(), newTerm);
}
-// TODO: We should maybe add this as a method to LoopLikeOpInterface.
-// For now, this acts as a placeholder.
+// TODO: We should maybe add a method to LoopLikeOpInterface that will
+// facilitate this transformation. For now, this acts as a placeholder.
template <>
void fuseTerminator(RewriterBase &rewriter, LoopLikeOpInterface source,
LoopLikeOpInterface &fused, IRMapping &mapping) {
@@ -1162,7 +1162,7 @@ LoopLikeOpInterface createFused(LoopLikeOpInterface target,
for (Operation &op : sourceRegion->front().without_terminator())
rewriter.clone(op, mapping);
- // TODO: Replace with interface method if added
+ // TODO: Replace with corresponding interface method if added
fuseTerminator(rewriter, source, fusedLoop, mapping);
return fusedLoop;
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 47bfe0baa7651..46c6be36c3271 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -189,8 +189,7 @@ module attributes {transform.with_named_sequence} {
// -----
-// CHECK-LABEL: func.func @matmul_fuse_2nd_forall_into_1st
-// CHECK-SAME: [[A1:%.*]]: {{.*}}, [[A2:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
+// CHECK: func.func @matmul_fuse_2nd_forall_into_1st([[A1:%.*]]: {{.*}}, [[A2:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
func.func @matmul_fuse_2nd_forall_into_1st(%A1 : tensor<128x128xf32>, %A2 : tensor<128x128xf32>, %B : tensor<128x128xf32>) -> (tensor<128x128xf32>, tensor<128x128xf32>) {
%zero = arith.constant 0.0 : f32
%out_alloc = tensor.empty() : tensor<128x128xf32>
>From 7e41a549f966956204f6f0971831e0423a9aeb9d Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Fri, 21 Jun 2024 15:52:08 -0500
Subject: [PATCH 22/34] address some review comments
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 13 ++++++-------
.../Dialect/SCF/TransformOps/SCFTransformOps.cpp | 9 ++++++---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 13 ++++++-------
3 files changed, 18 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 2baef9ca45db1..0c967ac68a081 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -626,7 +626,7 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(getOperation());
SmallVector<Value> inits(getOutputs());
- inits.append(newInitOperands.begin(), newInitOperands.end());
+ llvm::append_range(inits, newInitOperands);
scf::ForallOp newLoop = rewriter.create<scf::ForallOp>(
getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
inits, getMapping());
@@ -640,14 +640,13 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
if (replaceInitOperandUsesInLoop) {
// Replace all uses of `newInitOperands` with the corresponding basic block
// arguments.
- for (auto it :
+ for (auto &&[newOperand, oldOperand] :
llvm::zip(newInitOperands, newLoop.getBody()->getArguments().take_back(
newInitOperands.size()))) {
- rewriter.replaceUsesWithIf(std::get<0>(it), std::get<1>(it),
- [&](OpOperand &use) {
- Operation *user = use.getOwner();
- return newLoop->isProperAncestor(user);
- });
+ rewriter.replaceUsesWithIf(newOperand, oldOperand, [&](OpOperand &use) {
+ Operation *user = use.getOwner();
+ return newLoop->isProperAncestor(user);
+ });
}
}
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 99f92d7e24840..0e13b503098f0 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -261,8 +261,10 @@ loopScheduling(scf::ForOp forOp,
return 1;
};
- std::optional<int64_t> ubConstant = getConstantIntValue(forOp.getUpperBound());
- std::optional<int64_t> lbConstant = getConstantIntValue(forOp.getLowerBound());
+ std::optional<int64_t> ubConstant =
+ getConstantIntValue(forOp.getUpperBound());
+ std::optional<int64_t> lbConstant =
+ getConstantIntValue(forOp.getLowerBound());
DenseMap<Operation *, unsigned> opCycles;
std::map<unsigned, std::vector<Operation *>> wrappedSchedule;
for (Operation &op : forOp.getBody()->getOperations()) {
@@ -528,7 +530,8 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
<< "operations cannot be fused";
Operation *fusedLoop;
- /// TODO: Support fusion for loop-like ops besides scf.for and scf.forall.
+ // TODO: Support fusion for loop-like ops besides scf.for, scf.forall
+ // and scf.parallel.
if (isa<scf::ForOp>(target) && isa<scf::ForOp>(source)) {
fusedLoop = fuseIndependentSiblingForLoops(
cast<scf::ForOp>(target), cast<scf::ForOp>(source), rewriter);
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 09da6e6233ffc..dc15015e9bec2 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1076,7 +1076,7 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
LoopLikeOpInterface &source) {
- auto iterSpaceEq =
+ bool iterSpaceEq =
target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
target.getLoopSteps() == source.getLoopSteps();
@@ -1125,6 +1125,7 @@ void fuseTerminator(RewriterBase &rewriter, LoopLikeOpInterface source,
fuseTerminator(rewriter, cast<scf::ParallelOp>(source),
cast<scf::ParallelOp>(fused), mapping);
} else {
+ llvm_unreachable("unsupported loop types.");
return;
}
}
@@ -1239,13 +1240,11 @@ scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
newRedBlock.begin(),
newRedBlock.getArguments());
}
- target.replaceAllUsesWith(results.take_front(inits1.size()));
- source.replaceAllUsesWith(results.take_back(inits2.size()));
}
- term1->erase();
- term2->erase();
- target.erase();
- source.erase();
+ rewriter.replaceOp(target, results.take_front(inits1.size()));
+ rewriter.replaceOp(source, results.take_back(inits2.size()));
+ rewriter.eraseOp(term1);
+ rewriter.eraseOp(term2);
return fusedLoop;
}
>From cc95d75d2cc09f8a33850f3867c8313e374a0dfd Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Mon, 24 Jun 2024 14:56:48 -0500
Subject: [PATCH 23/34] move `createFused` to `LoopLikeInterface.h`
---
.../mlir/Interfaces/LoopLikeInterface.h | 20 ++++
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 101 ++++--------------
mlir/lib/Interfaces/LoopLikeInterface.cpp | 42 ++++++++
3 files changed, 82 insertions(+), 81 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 42609e824c86a..d862439a07790 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -57,4 +57,24 @@ class HasParallelRegion : public TraitBase<ConcreteType, HasParallelRegion> {
/// Include the generated interface declarations.
#include "mlir/Interfaces/LoopLikeInterface.h.inc"
+namespace mlir {
+/// A function that rewrites `target`'s terminator as a teminator obtained by
+/// fusing `source` into `target`.
+using FuseTerminatorFn =
+ std::function<void(RewriterBase &rewriter, LoopLikeOpInterface source,
+ LoopLikeOpInterface &target, IRMapping mapping)>;
+
+/// Returns a fused `LoopLikeOpInterface` created by fusing `source` to
+/// `target`. The `NewYieldValuesFn` callback is used to pass to the
+/// `replaceWithAdditionalYields` interface method to replace the loop with a
+/// new loop with (possibly) additional yields, while the `FuseTerminatorFn`
+/// callback is repsonsible for updating the fused loop terminator.
+LoopLikeOpInterface createFused(LoopLikeOpInterface target,
+ LoopLikeOpInterface source,
+ RewriterBase &rewriter,
+ NewYieldValuesFn newYieldValuesFn,
+ FuseTerminatorFn fuseTerminatorFn);
+
+} // namespace mlir
+
#endif // MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index dc15015e9bec2..93e7a40845b2e 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1082,93 +1082,14 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
target.getLoopSteps() == source.getLoopSteps();
auto forAllTarget = dyn_cast<scf::ForallOp>(*target);
auto forAllSource = dyn_cast<scf::ForallOp>(*source);
+ // TODO: Decouple checks on concrete loop types and move this function
+ // somewhere for general utility for `LoopLikeOpInterface`
if (forAllTarget && forAllSource)
return iterSpaceEq &&
forAllTarget.getMapping() == forAllSource.getMapping();
return iterSpaceEq;
}
-template <typename LoopTy>
-void fuseTerminator(RewriterBase &rewriter, LoopTy source, LoopTy &fused,
- IRMapping &mapping) {}
-
-template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForallOp source,
- scf::ForallOp &fused, IRMapping &mapping) {
- // Fuse the old terminator in_parallel ops into the new one.
- scf::InParallelOp fusedTerm = fused.getTerminator();
- rewriter.setInsertionPointToEnd(fusedTerm.getBody());
- for (Operation &op : source.getTerminator().getYieldingOps())
- rewriter.clone(op, mapping);
-}
-
-template <>
-void fuseTerminator(RewriterBase &rewriter, scf::ForOp source,
- scf::ForOp &fused, IRMapping &mapping) {
- // Build fused yield results by appropriately mapping original yield operands.
- auto newTerm = rewriter.clone(*fused.getBody()->getTerminator(), mapping);
- rewriter.replaceOp(fused.getBody()->getTerminator(), newTerm);
-}
-
-// TODO: We should maybe add a method to LoopLikeOpInterface that will
-// facilitate this transformation. For now, this acts as a placeholder.
-template <>
-void fuseTerminator(RewriterBase &rewriter, LoopLikeOpInterface source,
- LoopLikeOpInterface &fused, IRMapping &mapping) {
- if (isa<scf::ForOp>(source) && isa<scf::ForOp>(fused)) {
- fuseTerminator(rewriter, cast<scf::ForOp>(source), cast<scf::ForOp>(fused),
- mapping);
- } else if (isa<scf::ForallOp>(source) && isa<scf::ForallOp>(fused)) {
- fuseTerminator(rewriter, cast<scf::ForallOp>(source),
- cast<scf::ForallOp>(fused), mapping);
- } else if (isa<scf::ParallelOp>(source) && isa<scf::ParallelOp>(fused)) {
- fuseTerminator(rewriter, cast<scf::ParallelOp>(source),
- cast<scf::ParallelOp>(fused), mapping);
- } else {
- llvm_unreachable("unsupported loop types.");
- return;
- }
-}
-
-LoopLikeOpInterface createFused(LoopLikeOpInterface target,
- LoopLikeOpInterface source,
- RewriterBase &rewriter,
- NewYieldValuesFn newYieldValuesFn) {
- auto targetIterArgs = target.getRegionIterArgs();
- auto targetInductionVar = *target.getLoopInductionVars();
- SmallVector<Value> targetYieldOperands(target.getYieldedValues());
- auto sourceIterArgs = source.getRegionIterArgs();
- auto sourceInductionVar = *source.getLoopInductionVars();
- SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
- auto sourceRegion = source.getLoopRegions().front();
- LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
- rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
- newYieldValuesFn);
-
- // Map control operands.
- IRMapping mapping;
- mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
- mapping.map(targetIterArgs,
- fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
- mapping.map(targetYieldOperands,
- fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
- mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
- mapping.map(sourceIterArgs,
- fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
- mapping.map(sourceYieldOperands,
- fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
- // Append everything except the terminator into the fused operation.
- rewriter.setInsertionPoint(
- fusedLoop.getLoopRegions().front()->front().getTerminator());
- for (Operation &op : sourceRegion->front().without_terminator())
- rewriter.clone(op, mapping);
-
- // TODO: Replace with corresponding interface method if added
- fuseTerminator(rewriter, source, fusedLoop, mapping);
-
- return fusedLoop;
-}
-
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
scf::ForallOp source,
RewriterBase &rewriter) {
@@ -1177,6 +1098,15 @@ scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
// `ForallOp` does not have yields, rather an `InParallelOp` terminator.
return ValueRange{};
+ },
+ [&](RewriterBase &b, LoopLikeOpInterface source,
+ LoopLikeOpInterface &target, IRMapping mapping) {
+ auto sourceForall = cast<scf::ForallOp>(source);
+ auto targetForall = cast<scf::ForallOp>(target);
+ scf::InParallelOp fusedTerm = targetForall.getTerminator();
+ b.setInsertionPointToEnd(fusedTerm.getBody());
+ for (Operation &op : sourceForall.getTerminator().getYieldingOps())
+ b.clone(op, mapping);
}));
rewriter.replaceOp(source,
fusedLoop.getResults().take_back(source.getNumResults()));
@@ -1191,12 +1121,21 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
target, source, rewriter,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBBArgs) {
return source.getYieldedValues();
+ },
+ [&](RewriterBase &b, LoopLikeOpInterface source,
+ LoopLikeOpInterface &target, IRMapping mapping) {
+ auto sourceFor = cast<scf::ForOp>(source);
+ auto targetFor = cast<scf::ForOp>(target);
+ auto newTerm = b.clone(*targetFor.getBody()->getTerminator(), mapping);
+ b.replaceOp(targetFor.getBody()->getTerminator(), newTerm);
}));
rewriter.replaceOp(source,
fusedLoop.getResults().take_back(source.getNumResults()));
return fusedLoop;
}
+// TODO: Finish refactoring this a la the above, but likely requires additional
+// interface methods.
scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
scf::ParallelOp target, scf::ParallelOp source, RewriterBase &rewriter) {
Block *block1 = target.getBody();
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 1e0e87b64e811..aefd388461570 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -8,6 +8,8 @@
#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "llvm/ADT/DenseSet.h"
@@ -113,3 +115,43 @@ LogicalResult detail::verifyLoopLikeOpInterface(Operation *op) {
return success();
}
+
+LoopLikeOpInterface mlir::createFused(LoopLikeOpInterface target,
+ LoopLikeOpInterface source,
+ RewriterBase &rewriter,
+ NewYieldValuesFn newYieldValuesFn,
+ FuseTerminatorFn fuseTerminatorFn) {
+ auto targetIterArgs = target.getRegionIterArgs();
+ auto targetInductionVar = *target.getLoopInductionVars();
+ SmallVector<Value> targetYieldOperands(target.getYieldedValues());
+ auto sourceIterArgs = source.getRegionIterArgs();
+ auto sourceInductionVar = *source.getLoopInductionVars();
+ SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
+ auto sourceRegion = source.getLoopRegions().front();
+ LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
+ rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
+ newYieldValuesFn);
+
+ // Map control operands.
+ IRMapping mapping;
+ mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
+ mapping.map(targetIterArgs,
+ fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
+ mapping.map(targetYieldOperands,
+ fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
+ mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
+ mapping.map(sourceIterArgs,
+ fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
+ mapping.map(sourceYieldOperands,
+ fusedLoop.getYieldedValues().take_back(sourceIterArgs.size()));
+ // Append everything except the terminator into the fused operation.
+ rewriter.setInsertionPoint(
+ fusedLoop.getLoopRegions().front()->front().getTerminator());
+ for (Operation &op : sourceRegion->front().without_terminator())
+ rewriter.clone(op, mapping);
+
+ // TODO: Replace with corresponding interface method if added
+ fuseTerminatorFn(rewriter, source, fusedLoop, mapping);
+
+ return fusedLoop;
+}
>From 3430a36fda3c53d466550a7d8fd13b331f96f005 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 26 Jun 2024 13:51:34 -0500
Subject: [PATCH 24/34] address more review comments
---
mlir/lib/Dialect/SCF/IR/SCF.cpp | 4 ++--
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 10 +++++-----
2 files changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp
index 0c967ac68a081..1e42376ce58ca 100644
--- a/mlir/lib/Dialect/SCF/IR/SCF.cpp
+++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp
@@ -629,9 +629,9 @@ FailureOr<LoopLikeOpInterface> ForallOp::replaceWithAdditionalYields(
llvm::append_range(inits, newInitOperands);
scf::ForallOp newLoop = rewriter.create<scf::ForallOp>(
getLoc(), getMixedLowerBound(), getMixedUpperBound(), getMixedStep(),
- inits, getMapping());
+ inits, getMapping(),
+ /*bodyBuilderFn =*/[](OpBuilder &, Location, ValueRange) {});
- rewriter.eraseOp(newLoop.getTerminator());
// Move the loop body to the new op.
rewriter.mergeBlocks(getBody(), newLoop.getBody(),
newLoop.getBody()->getArguments().take_front(
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 93e7a40845b2e..e7496cd97cd63 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1124,7 +1124,6 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
},
[&](RewriterBase &b, LoopLikeOpInterface source,
LoopLikeOpInterface &target, IRMapping mapping) {
- auto sourceFor = cast<scf::ForOp>(source);
auto targetFor = cast<scf::ForOp>(target);
auto newTerm = b.clone(*targetFor.getBody()->getTerminator(), mapping);
b.replaceOp(targetFor.getBody()->getTerminator(), newTerm);
@@ -1151,8 +1150,9 @@ scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
rewriter.setInsertionPoint(source);
auto fusedLoop = rewriter.create<scf::ParallelOp>(
- source.getLoc(), source.getLowerBound(), source.getUpperBound(),
- source.getStep(), newInitVars);
+ rewriter.getFusedLoc(target.getLoc(), source.getLoc()),
+ source.getLowerBound(), source.getUpperBound(), source.getStep(),
+ newInitVars);
Block *newBlock = fusedLoop.getBody();
rewriter.inlineBlockBefore(block2, newBlock, newBlock->begin(),
newBlock->getArguments());
@@ -1168,8 +1168,8 @@ scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
SmallVector<Value> newReduceArgs(reduceArgs1.begin(), reduceArgs1.end());
newReduceArgs.append(reduceArgs2.begin(), reduceArgs2.end());
- auto newReduceOp =
- rewriter.create<scf::ReduceOp>(term2.getLoc(), newReduceArgs);
+ auto newReduceOp = rewriter.create<scf::ReduceOp>(
+ rewriter.getFusedLoc(term1.getLoc(), term2.getLoc()), newReduceArgs);
for (auto &&[i, reg] : llvm::enumerate(llvm::concat<Region>(
term1.getReductions(), term2.getReductions()))) {
>From 8447c121b95279a283b5e7b25f094f6abb062216 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 26 Jun 2024 20:06:43 -0500
Subject: [PATCH 25/34] switch to function_ref
---
mlir/include/mlir/Interfaces/LoopLikeInterface.h | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index d862439a07790..cfe2c14b838f6 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -61,8 +61,8 @@ namespace mlir {
/// A function that rewrites `target`'s terminator as a teminator obtained by
/// fusing `source` into `target`.
using FuseTerminatorFn =
- std::function<void(RewriterBase &rewriter, LoopLikeOpInterface source,
- LoopLikeOpInterface &target, IRMapping mapping)>;
+ function_ref<void(RewriterBase &rewriter, LoopLikeOpInterface source,
+ LoopLikeOpInterface &target, IRMapping mapping)>;
/// Returns a fused `LoopLikeOpInterface` created by fusing `source` to
/// `target`. The `NewYieldValuesFn` callback is used to pass to the
>From fbd7b72bb44c7833a683d93fccaa9d992856ee8b Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 26 Jun 2024 22:52:20 -0500
Subject: [PATCH 26/34] check optional values
---
mlir/lib/Interfaces/LoopLikeInterface.cpp | 27 +++++++++++++++++------
1 file changed, 20 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index aefd388461570..6f0ebec0519be 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -122,24 +122,37 @@ LoopLikeOpInterface mlir::createFused(LoopLikeOpInterface target,
NewYieldValuesFn newYieldValuesFn,
FuseTerminatorFn fuseTerminatorFn) {
auto targetIterArgs = target.getRegionIterArgs();
- auto targetInductionVar = *target.getLoopInductionVars();
+ std::optional<SmallVector<Value>> targetInductionVar =
+ target.getLoopInductionVars();
SmallVector<Value> targetYieldOperands(target.getYieldedValues());
auto sourceIterArgs = source.getRegionIterArgs();
- auto sourceInductionVar = *source.getLoopInductionVars();
+ std::optional<SmallVector<Value>> sourceInductionVar =
+ *source.getLoopInductionVars();
SmallVector<Value> sourceYieldOperands(source.getYieldedValues());
auto sourceRegion = source.getLoopRegions().front();
- LoopLikeOpInterface fusedLoop = *target.replaceWithAdditionalYields(
- rewriter, source.getInits(), /*replaceInitOperandUsesInLoop=*/false,
- newYieldValuesFn);
+
+ FailureOr<LoopLikeOpInterface> maybeFusedLoop =
+ target.replaceWithAdditionalYields(rewriter, source.getInits(),
+ /*replaceInitOperandUsesInLoop=*/false,
+ newYieldValuesFn);
+ if (failed(maybeFusedLoop))
+ llvm_unreachable("failed to replace loop");
+ LoopLikeOpInterface fusedLoop = *maybeFusedLoop;
// Map control operands.
IRMapping mapping;
- mapping.map(targetInductionVar, *fusedLoop.getLoopInductionVars());
+ std::optional<SmallVector<Value>> fusedInductionVar =
+ fusedLoop.getLoopInductionVars();
+ if (fusedInductionVar) {
+ if (!targetInductionVar || !sourceInductionVar)
+ llvm_unreachable("expected target and source loops to have induction vars");
+ mapping.map(*targetInductionVar, *fusedInductionVar);
+ mapping.map(*sourceInductionVar, *fusedInductionVar);
+ }
mapping.map(targetIterArgs,
fusedLoop.getRegionIterArgs().take_front(targetIterArgs.size()));
mapping.map(targetYieldOperands,
fusedLoop.getYieldedValues().take_front(targetIterArgs.size()));
- mapping.map(sourceInductionVar, *fusedLoop.getLoopInductionVars());
mapping.map(sourceIterArgs,
fusedLoop.getRegionIterArgs().take_back(sourceIterArgs.size()));
mapping.map(sourceYieldOperands,
>From ffb73a7a76b382414f8f8295f6d6dc14a3edfa99 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 26 Jun 2024 23:34:38 -0500
Subject: [PATCH 27/34] replace equalIterationSpaces with
checkFusionStructuredLegality
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 5 +++--
.../SCF/Transforms/ParallelLoopFusion.cpp | 20 +------------------
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 4 ++--
3 files changed, 6 insertions(+), 23 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index ac4434b337890..ca3ab0aeae1de 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -160,8 +160,9 @@ void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
// Fusion related helpers
//===----------------------------------------------------------------------===//
-bool checkFusionStructuralLegality(LoopLikeOpInterface &target,
- LoopLikeOpInterface &source);
+/// Check structural compatibility between two loops such as iteration space.
+bool checkFusionStructuralLegality(LoopLikeOpInterface target,
+ LoopLikeOpInterface source);
/// Given two scf.forall loops, `target` and `source`, fuses `target` into
/// `source`. Assumes that the given loops are siblings and are independent of
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 326a8f93162b9..fd57a9228186e 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -38,24 +38,6 @@ static bool hasNestedParallelOp(ParallelOp ploop) {
return walkResult.wasInterrupted();
}
-/// Verify equal iteration spaces.
-static bool equalIterationSpaces(ParallelOp firstPloop,
- ParallelOp secondPloop) {
- if (firstPloop.getNumLoops() != secondPloop.getNumLoops())
- return false;
-
- auto matchOperands = [&](const OperandRange &lhs,
- const OperandRange &rhs) -> bool {
- // TODO: Extend this to support aliases and equal constants.
- return std::equal(lhs.begin(), lhs.end(), rhs.begin());
- };
- return matchOperands(firstPloop.getLowerBound(),
- secondPloop.getLowerBound()) &&
- matchOperands(firstPloop.getUpperBound(),
- secondPloop.getUpperBound()) &&
- matchOperands(firstPloop.getStep(), secondPloop.getStep());
-}
-
/// Checks if the parallel loops have mixed access to the same buffers. Returns
/// `true` if the first parallel loop writes to the same indices that the second
/// loop reads.
@@ -156,7 +138,7 @@ static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
llvm::function_ref<bool(Value, Value)> mayAlias) {
return !hasNestedParallelOp(firstPloop) &&
!hasNestedParallelOp(secondPloop) &&
- equalIterationSpaces(firstPloop, secondPloop) &&
+ checkFusionStructuralLegality(firstPloop, secondPloop) &&
succeeded(verifyDependencies(firstPloop, secondPloop,
firstToSecondPloopIndices, mayAlias));
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index e7496cd97cd63..fab6592d9eb2a 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1074,8 +1074,8 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
// Fusion related helpers
//===----------------------------------------------------------------------===//
-bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface &target,
- LoopLikeOpInterface &source) {
+bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
+ LoopLikeOpInterface source) {
bool iterSpaceEq =
target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
>From a6d0588da17170b1d3653efb51704b10d770dc58 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 11:31:03 -0500
Subject: [PATCH 28/34] check if isOpSibling in checkFusionStructuralLegality
---
mlir/include/mlir/Dialect/SCF/Utils/Utils.h | 6 +-
.../SCF/TransformOps/SCFTransformOps.cpp | 84 +----------------
.../SCF/Transforms/ParallelLoopFusion.cpp | 3 +-
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 91 ++++++++++++++++++-
.../SCF/transform-loop-fuse-sibling.mlir | 3 +-
5 files changed, 99 insertions(+), 88 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
index ca3ab0aeae1de..59aeff2da14ea 100644
--- a/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/SCF/Utils/Utils.h
@@ -160,9 +160,11 @@ void getPerfectlyNestedLoops(SmallVectorImpl<scf::ForOp> &nestedLoops,
// Fusion related helpers
//===----------------------------------------------------------------------===//
-/// Check structural compatibility between two loops such as iteration space.
+/// Check structural compatibility between two loops such as iteration space
+/// and dominance.
bool checkFusionStructuralLegality(LoopLikeOpInterface target,
- LoopLikeOpInterface source);
+ LoopLikeOpInterface source,
+ Diagnostic &diag);
/// Given two scf.forall loops, `target` and `source`, fuses `target` into
/// `source`. Assumes that the given loops are siblings and are independent of
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 0e13b503098f0..3e0a483615a3d 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -425,78 +425,6 @@ void transform::TakeAssumedBranchOp::getEffects(
// LoopFuseSiblingOp
//===----------------------------------------------------------------------===//
-/// Check if `target` and `source` are siblings, in the context that `target`
-/// is being fused into `source`.
-///
-/// This is a simple check that just checks if both operations are in the same
-/// block and some checks to ensure that the fused IR does not violate
-/// dominance.
-static DiagnosedSilenceableFailure isOpSibling(Operation *target,
- Operation *source) {
- // Check if both operations are same.
- if (target == source)
- return emitSilenceableFailure(source)
- << "target and source need to be different loops";
-
- // Check if both operations are in the same block.
- if (target->getBlock() != source->getBlock())
- return emitSilenceableFailure(source)
- << "target and source are not in the same block";
-
- // Check if fusion will violate dominance.
- DominanceInfo domInfo(source);
- if (target->isBeforeInBlock(source)) {
- // Since `target` is before `source`, all users of results of `target`
- // need to be dominated by `source`.
- for (Operation *user : target->getUsers()) {
- if (!domInfo.properlyDominates(source, user, /*enclosingOpOk=*/false)) {
- return emitSilenceableFailure(target)
- << "user of results of target should be properly dominated by "
- "source";
- }
- }
- } else {
- // Since `target` is after `source`, all values used by `target` need
- // to dominate `source`.
-
- // Check if operands of `target` are dominated by `source`.
- for (Value operand : target->getOperands()) {
- Operation *operandOp = operand.getDefiningOp();
- // Operands without defining operations are block arguments. When `target`
- // and `source` occur in the same block, these operands dominate `source`.
- if (!operandOp)
- continue;
-
- // Operand's defining operation should properly dominate `source`.
- if (!domInfo.properlyDominates(operandOp, source,
- /*enclosingOpOk=*/false))
- return emitSilenceableFailure(target)
- << "operands of target should be properly dominated by source";
- }
-
- // Check if values used by `target` are dominated by `source`.
- bool failed = false;
- OpOperand *failedValue = nullptr;
- visitUsedValuesDefinedAbove(target->getRegions(), [&](OpOperand *operand) {
- Operation *operandOp = operand->get().getDefiningOp();
- if (operandOp && !domInfo.properlyDominates(operandOp, source,
- /*enclosingOpOk=*/false)) {
- // `operand` is not an argument of an enclosing block and the defining
- // op of `operand` is outside `target` but does not dominate `source`.
- failed = true;
- failedValue = operand;
- }
- });
-
- if (failed)
- return emitSilenceableFailure(failedValue->getOwner())
- << "values used inside regions of target should be properly "
- "dominated by source";
- }
-
- return DiagnosedSilenceableFailure::success();
-}
-
DiagnosedSilenceableFailure
transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
transform::TransformResults &results,
@@ -520,14 +448,10 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
return emitSilenceableFailure(target->getLoc())
<< "target or source is not a loop op";
- // Check if the target and source are siblings.
- DiagnosedSilenceableFailure diag = isOpSibling(target, source);
- if (!diag.succeeded())
- return diag;
-
- if (!mlir::checkFusionStructuralLegality(target, source))
- return emitSilenceableFailure(target->getLoc())
- << "operations cannot be fused";
+ // Check if loops can be fused
+ Diagnostic diag(target.getLoc(), DiagnosticSeverity::Error);
+ if (!mlir::checkFusionStructuralLegality(target, source, diag))
+ return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
Operation *fusedLoop;
// TODO: Support fusion for loop-like ops besides scf.for, scf.forall
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index fd57a9228186e..b46535078dd8b 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -136,9 +136,10 @@ verifyDependencies(ParallelOp firstPloop, ParallelOp secondPloop,
static bool isFusionLegal(ParallelOp firstPloop, ParallelOp secondPloop,
const IRMapping &firstToSecondPloopIndices,
llvm::function_ref<bool(Value, Value)> mayAlias) {
+ Diagnostic diag(firstPloop.getLoc(), DiagnosticSeverity::Remark);
return !hasNestedParallelOp(firstPloop) &&
!hasNestedParallelOp(secondPloop) &&
- checkFusionStructuralLegality(firstPloop, secondPloop) &&
+ checkFusionStructuralLegality(firstPloop, secondPloop, diag) &&
succeeded(verifyDependencies(firstPloop, secondPloop,
firstToSecondPloopIndices, mayAlias));
}
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index fab6592d9eb2a..b1a367281a6ca 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -17,6 +17,7 @@
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dominance.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -1074,8 +1075,86 @@ TileLoops mlir::extractFixedOuterLoops(scf::ForOp rootForOp,
// Fusion related helpers
//===----------------------------------------------------------------------===//
+/// Check if `target` and `source` are siblings, in the context that `target`
+/// is being fused into `source`.
+///
+/// This is a simple check that just checks if both operations are in the same
+/// block and some checks to ensure that the fused IR does not violate
+/// dominance.
+static bool isOpSibling(Operation *target, Operation *source,
+ Diagnostic &diag) {
+ // Check if both operations are same.
+ if (target == source) {
+ diag << "target and source need to be different loops";
+ return false;
+ }
+
+ // Check if both operations are in the same block.
+ if (target->getBlock() != source->getBlock()) {
+ diag << "target and source are not in the same block";
+ return false;
+ }
+
+ // Check if fusion will violate dominance.
+ DominanceInfo domInfo(source);
+ if (target->isBeforeInBlock(source)) {
+ // Since `target` is before `source`, all users of results of `target`
+ // need to be dominated by `source`.
+ for (Operation *user : target->getUsers()) {
+ if (!domInfo.properlyDominates(source, user, /*enclosingOpOk=*/false)) {
+ diag << "user of results of target should "
+ "be properly dominated by source";
+ return false;
+ }
+ }
+ } else {
+ // Since `target` is after `source`, all values used by `target` need
+ // to dominate `source`.
+
+ // Check if operands of `target` are dominated by `source`.
+ for (Value operand : target->getOperands()) {
+ Operation *operandOp = operand.getDefiningOp();
+ // Operands without defining operations are block arguments. When `target`
+ // and `source` occur in the same block, these operands dominate `source`.
+ if (!operandOp)
+ continue;
+
+ // Operand's defining operation should properly dominate `source`.
+ if (!domInfo.properlyDominates(operandOp, source,
+ /*enclosingOpOk=*/false)) {
+ diag << "operands of target should be properly dominated by source";
+ return false;
+ }
+ }
+
+ // Check if values used by `target` are dominated by `source`.
+ bool failed = false;
+ OpOperand *failedValue = nullptr;
+ visitUsedValuesDefinedAbove(target->getRegions(), [&](OpOperand *operand) {
+ Operation *operandOp = operand->get().getDefiningOp();
+ if (operandOp && !domInfo.properlyDominates(operandOp, source,
+ /*enclosingOpOk=*/false)) {
+ // `operand` is not an argument of an enclosing block and the defining
+ // op of `operand` is outside `target` but does not dominate `source`.
+ failed = true;
+ failedValue = operand;
+ }
+ });
+
+ if (failed) {
+ diag << "values used inside regions of target should be properly "
+ "dominated by source";
+ diag.attachNote(failedValue->getOwner()->getLoc()) << "see operation";
+ return false;
+ }
+ }
+
+ return true;
+}
+
bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
- LoopLikeOpInterface source) {
+ LoopLikeOpInterface source,
+ Diagnostic &diag) {
bool iterSpaceEq =
target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
@@ -1085,9 +1164,13 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
// TODO: Decouple checks on concrete loop types and move this function
// somewhere for general utility for `LoopLikeOpInterface`
if (forAllTarget && forAllSource)
- return iterSpaceEq &&
- forAllTarget.getMapping() == forAllSource.getMapping();
- return iterSpaceEq;
+ iterSpaceEq =
+ iterSpaceEq && forAllTarget.getMapping() == forAllSource.getMapping();
+ if (!iterSpaceEq) {
+ diag << "target and source iteration spaces must be equal";
+ return false;
+ }
+ return isOpSibling(target, source, diag);
}
scf::ForallOp mlir::fuseIndependentSiblingForallLoops(scf::ForallOp target,
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 46c6be36c3271..b03aa5cf38bfa 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -335,8 +335,9 @@ func.func @target_for_region_uses_result_of_source_for_err(%A: tensor<128xf32>,
%6 = vector.transfer_write %5, %arg4[%arg3] {in_bounds = [true]} : vector<16xf32>, tensor<128xf32>
scf.yield %6 : tensor<128xf32>
}
- %dup1 = scf.for %arg3 = %c0 to %c128 step %c16 iter_args(%arg4 = %B) -> (tensor<128xf32>) {
// expected-error @below {{values used inside regions of target should be properly dominated by source}}
+ %dup1 = scf.for %arg3 = %c0 to %c128 step %c16 iter_args(%arg4 = %B) -> (tensor<128xf32>) {
+ // expected-note @below {{see operation}}
%dup2 = vector.transfer_read %1[%arg3], %cst {in_bounds = [true]} : tensor<128xf32>, vector<16xf32>
%dup3 = vector.transfer_read %arg4[%arg3], %cst {in_bounds = [true]} : tensor<128xf32>, vector<16xf32>
%dup5 = arith.addf %dup3, %dup2 : vector<16xf32>
>From ff47980d71330f65ecf05451f4d2345145a24e21 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 11:46:02 -0500
Subject: [PATCH 29/34] remove extra dominance check
---
mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp | 7 -------
1 file changed, 7 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index b46535078dd8b..95ec8861aee2b 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -158,13 +158,6 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
mayAlias))
return;
- DominanceInfo dom;
- // We are fusing first loop into second, make sure there are no users of the
- // first loop results between loops.
- for (Operation *user : firstPloop->getUsers())
- if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
- return;
-
IRRewriter rewriter(builder);
secondPloop = mlir::fuseIndependentSiblingParallelLoops(
firstPloop, secondPloop, rewriter);
>From c6847ec9212aa1754ad27c16f568a8d16346197d Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 12:34:20 -0500
Subject: [PATCH 30/34] address more review comments
---
mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp | 6 ++----
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 1 +
2 files changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
index 3e0a483615a3d..8c93554f4016e 100644
--- a/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
+++ b/mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp
@@ -440,10 +440,8 @@ transform::LoopFuseSiblingOp::apply(transform::TransformRewriter &rewriter,
<< "source handle (got " << llvm::range_size(sourceOps) << ")";
}
- LoopLikeOpInterface target =
- dyn_cast<LoopLikeOpInterface>(*targetOps.begin());
- LoopLikeOpInterface source =
- dyn_cast<LoopLikeOpInterface>(*sourceOps.begin());
+ auto target = dyn_cast<LoopLikeOpInterface>(*targetOps.begin());
+ auto source = dyn_cast<LoopLikeOpInterface>(*sourceOps.begin());
if (!target || !source)
return emitSilenceableFailure(target->getLoc())
<< "target or source is not a loop op";
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index b1a367281a6ca..666b67517f4d4 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1220,6 +1220,7 @@ scf::ForOp mlir::fuseIndependentSiblingForLoops(scf::ForOp target,
// interface methods.
scf::ParallelOp mlir::fuseIndependentSiblingParallelLoops(
scf::ParallelOp target, scf::ParallelOp source, RewriterBase &rewriter) {
+ OpBuilder::InsertionGuard guard(rewriter);
Block *block1 = target.getBody();
Block *block2 = source.getBody();
auto term1 = cast<scf::ReduceOp>(block1->getTerminator());
>From f50c6aa14b36836950cc47909d4cca03d5ede8e3 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 12:55:08 -0500
Subject: [PATCH 31/34] add more lit tests for scf.parallel
---
.../SCF/transform-loop-fuse-sibling.mlir | 144 ++++++++++++++++++
1 file changed, 144 insertions(+)
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index b03aa5cf38bfa..1d46a3d88f47d 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -100,6 +100,116 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: func @fuse_two_parallel_reverse
+// CHECK-SAME: ([[A:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}) {
+func.func @fuse_two_parallel_reverse(%A: memref<2x2xf32>, %B: memref<2x2xf32>) {
+// CHECK-DAG: [[C2:%.*]] = arith.constant 2 : index
+// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index
+// CHECK-DAG: [[C1:%.*]] = arith.constant 1 : index
+// CHECK-DAG: [[C1FP:%.*]] = arith.constant 1.
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c1fp = arith.constant 1.0 : f32
+// CHECK: [[SUM:%.*]] = memref.alloc()
+ %sum = memref.alloc() : memref<2x2xf32>
+// CHECK: scf.parallel ([[I:%.*]], [[J:%.*]]) = ([[C0]], [[C0]])
+// CHECK-SAME: to ([[C2]], [[C2]]) step ([[C1]], [[C1]]) {
+// CHECK: [[SUM_ELEM_:%.*]] = memref.load [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK: [[A_ELEM:%.*]] = memref.load [[A]]{{\[}}[[I]], [[J]]]
+// CHECK: [[PRODUCT_ELEM:%.*]] = arith.mulf [[SUM_ELEM_]], [[A_ELEM]]
+// CHECK: memref.store [[PRODUCT_ELEM]], [[B]]{{\[}}[[I]], [[J]]]
+// CHECK-NOT: scf.parallel
+// CHECK: [[B_ELEM:%.*]] = memref.load [[B]]{{\[}}[[I]], [[J]]]
+// CHECK: [[SUM_ELEM:%.*]] = arith.addf [[B_ELEM]], [[C1FP]]
+// CHECK: memref.store [[SUM_ELEM]], [[SUM]]{{\[}}[[I]], [[J]]]
+// CHECK: scf.reduce
+// CHECK: }
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
+ %sum_elem = arith.addf %B_elem, %c1fp : f32
+ memref.store %sum_elem, %sum[%i, %j] : memref<2x2xf32>
+ scf.reduce
+ }
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %sum_elem = memref.load %sum[%i, %j] : memref<2x2xf32>
+ %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
+ %product_elem = arith.mulf %sum_elem, %A_elem : f32
+ memref.store %product_elem, %B[%i, %j] : memref<2x2xf32>
+ scf.reduce
+ }
+// CHECK: memref.dealloc [[SUM]]
+ memref.dealloc %sum : memref<2x2xf32>
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.parallel"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %parallel:2 = transform.split_handle %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %fused = transform.loop.fuse_sibling %parallel#1 into %parallel#0 : (!transform.any_op,!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+// CHECK-LABEL: func @fuse_reductions_two
+// CHECK-SAME: (%[[A:.*]]: memref<2x2xf32>, %[[B:.*]]: memref<2x2xf32>) -> (f32, f32)
+func.func @fuse_reductions_two(%A: memref<2x2xf32>, %B: memref<2x2xf32>) -> (f32, f32) {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+// CHECK-DAG: %[[INIT1:.*]] = arith.constant 1.000000e+00 : f32
+// CHECK-DAG: %[[INIT2:.*]] = arith.constant 2.000000e+00 : f32
+// CHECK: %[[RES:.*]]:2 = scf.parallel (%[[I:.*]], %[[J:.*]]) = (%[[C0]], %[[C0]])
+// CHECK-SAME: to (%[[C2]], %[[C2]]) step (%[[C1]], %[[C1]])
+// CHECK-SAME: init (%[[INIT1]], %[[INIT2]]) -> (f32, f32)
+// CHECK: %[[VAL_A:.*]] = memref.load %[[A]][%[[I]], %[[J]]]
+// CHECK: %[[VAL_B:.*]] = memref.load %[[B]][%[[I]], %[[J]]]
+// CHECK: scf.reduce(%[[VAL_A]], %[[VAL_B]] : f32, f32) {
+// CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
+// CHECK: %[[R:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32
+// CHECK: scf.reduce.return %[[R]] : f32
+// CHECK: }
+// CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
+// CHECK: %[[R:.*]] = arith.mulf %[[LHS]], %[[RHS]] : f32
+// CHECK: scf.reduce.return %[[R]] : f32
+// CHECK: }
+// CHECK: return %[[RES]]#0, %[[RES]]#1 : f32, f32
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %init1 = arith.constant 1.0 : f32
+ %init2 = arith.constant 2.0 : f32
+ %res1 = scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) init(%init1) -> f32 {
+ %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
+ scf.reduce(%A_elem : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %1 = arith.addf %lhs, %rhs : f32
+ scf.reduce.return %1 : f32
+ }
+ }
+ %res2 = scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) init(%init2) -> f32 {
+ %B_elem = memref.load %B[%i, %j] : memref<2x2xf32>
+ scf.reduce(%B_elem : f32) {
+ ^bb0(%lhs: f32, %rhs: f32):
+ %1 = arith.mulf %lhs, %rhs : f32
+ scf.reduce.return %1 : f32
+ }
+ }
+ return %res1, %res2 : f32, f32
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.parallel"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %parallel:2 = transform.split_handle %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %fused = transform.loop.fuse_sibling %parallel#0 into %parallel#1 : (!transform.any_op,!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK: func.func @fuse_2nd_for_into_1st([[A:%.*]]: {{.*}}, [[B:%.*]]: {{.*}}
func.func @fuse_2nd_for_into_1st(%A: tensor<128xf32>, %B: tensor<128xf32>) -> (tensor<128xf32>, tensor<128xf32>) {
// CHECK-DAG: [[C0:%.*]] = arith.constant 0 : index
@@ -382,3 +492,37 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+func.func @non_matching_iteration_spaces_err(%A: memref<2x2xf32>, %B: memref<2x2xf32>) {
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c1fp = arith.constant 1.0 : f32
+ %sum = memref.alloc() : memref<2x2xf32>
+ // expected-error @below {{target and source iteration spaces must be equal}}
+ scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
+ %B_elem = memref.load %B[%i, %c0] : memref<2x2xf32>
+ %sum_elem = arith.addf %B_elem, %c1fp : f32
+ memref.store %sum_elem, %sum[%i, %c0] : memref<2x2xf32>
+ scf.reduce
+ }
+ scf.parallel (%i, %j) = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+ %sum_elem = memref.load %sum[%i, %j] : memref<2x2xf32>
+ %A_elem = memref.load %A[%i, %j] : memref<2x2xf32>
+ %product_elem = arith.mulf %sum_elem, %A_elem : f32
+ memref.store %product_elem, %B[%i, %j] : memref<2x2xf32>
+ scf.reduce
+ }
+ memref.dealloc %sum : memref<2x2xf32>
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.parallel"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %parallel:2 = transform.split_handle %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %fused = transform.loop.fuse_sibling %parallel#0 into %parallel#1 : (!transform.any_op,!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
>From 6dd68c1b8408f05acaff9d040d4a686044295fcc Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 13:24:11 -0500
Subject: [PATCH 32/34] check for equal loop types in
checkFusionStructuralLegality
---
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 5 +++
.../SCF/transform-loop-fuse-sibling.mlir | 33 +++++++++++++++++++
2 files changed, 38 insertions(+)
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 666b67517f4d4..0c966bf182cbd 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1155,6 +1155,11 @@ static bool isOpSibling(Operation *target, Operation *source,
bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
LoopLikeOpInterface source,
Diagnostic &diag) {
+ if (target->getName() != source->getName()) {
+ diag << "target and source must be same loop type";
+ return false;
+ }
+
bool iterSpaceEq =
target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 1d46a3d88f47d..505013d328962 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -526,3 +526,36 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+func.func @non_matching_loop_types_err(%A: memref<2xf32>, %B: memref<2xf32>) {
+ %c2 = arith.constant 2 : index
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c1fp = arith.constant 1.0 : f32
+ %sum = memref.alloc() : memref<2xf32>
+ // expected-error @below {{target and source must be same loop type}}
+ scf.for %i = %c0 to %c2 step %c1 {
+ %B_elem = memref.load %B[%i] : memref<2xf32>
+ %sum_elem = arith.addf %B_elem, %c1fp : f32
+ memref.store %sum_elem, %sum[%i] : memref<2xf32>
+ }
+ scf.parallel (%i) = (%c0) to (%c2) step (%c1) {
+ %sum_elem = memref.load %sum[%i] : memref<2xf32>
+ %A_elem = memref.load %A[%i] : memref<2xf32>
+ %product_elem = arith.mulf %sum_elem, %A_elem : f32
+ memref.store %product_elem, %B[%i] : memref<2xf32>
+ scf.reduce
+ }
+ memref.dealloc %sum : memref<2xf32>
+ return
+}
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["scf.for"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.structured.match ops{["scf.parallel"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %fused = transform.loop.fuse_sibling %0 into %1 : (!transform.any_op,!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
>From 99d821b47cb731ac7a12b60c44d88af5ad2fb0d1 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Thu, 27 Jun 2024 13:42:20 -0500
Subject: [PATCH 33/34] address more comments
---
mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp | 1 -
mlir/lib/Dialect/SCF/Utils/Utils.cpp | 8 +++-----
2 files changed, 3 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
index 95ec8861aee2b..b775f988576e3 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ParallelLoopFusion.cpp
@@ -161,7 +161,6 @@ static void fuseIfLegal(ParallelOp firstPloop, ParallelOp &secondPloop,
IRRewriter rewriter(builder);
secondPloop = mlir::fuseIndependentSiblingParallelLoops(
firstPloop, secondPloop, rewriter);
- ;
}
void mlir::scf::naivelyFuseParallelOps(
diff --git a/mlir/lib/Dialect/SCF/Utils/Utils.cpp b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
index 0c966bf182cbd..a79aef34e48b1 100644
--- a/mlir/lib/Dialect/SCF/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/SCF/Utils/Utils.cpp
@@ -1164,13 +1164,11 @@ bool mlir::checkFusionStructuralLegality(LoopLikeOpInterface target,
target.getLoopLowerBounds() == source.getLoopLowerBounds() &&
target.getLoopUpperBounds() == source.getLoopUpperBounds() &&
target.getLoopSteps() == source.getLoopSteps();
- auto forAllTarget = dyn_cast<scf::ForallOp>(*target);
- auto forAllSource = dyn_cast<scf::ForallOp>(*source);
// TODO: Decouple checks on concrete loop types and move this function
// somewhere for general utility for `LoopLikeOpInterface`
- if (forAllTarget && forAllSource)
- iterSpaceEq =
- iterSpaceEq && forAllTarget.getMapping() == forAllSource.getMapping();
+ if (auto forAllTarget = dyn_cast<scf::ForallOp>(*target))
+ iterSpaceEq = iterSpaceEq && forAllTarget.getMapping() ==
+ cast<scf::ForallOp>(*source).getMapping();
if (!iterSpaceEq) {
diag << "target and source iteration spaces must be equal";
return false;
>From 7f9c172662717be33e8d340b5eb744c51c6137b6 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Wed, 3 Jul 2024 11:54:39 -0500
Subject: [PATCH 34/34] Fix bug in fusion refactor and add test
---
mlir/lib/Interfaces/LoopLikeInterface.cpp | 4 +-
.../SCF/transform-loop-fuse-sibling.mlir | 56 +++++++++++++++++++
2 files changed, 59 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Interfaces/LoopLikeInterface.cpp b/mlir/lib/Interfaces/LoopLikeInterface.cpp
index 6f0ebec0519be..226815f79227c 100644
--- a/mlir/lib/Interfaces/LoopLikeInterface.cpp
+++ b/mlir/lib/Interfaces/LoopLikeInterface.cpp
@@ -138,6 +138,7 @@ LoopLikeOpInterface mlir::createFused(LoopLikeOpInterface target,
if (failed(maybeFusedLoop))
llvm_unreachable("failed to replace loop");
LoopLikeOpInterface fusedLoop = *maybeFusedLoop;
+ rewriter.moveOpBefore(fusedLoop, source);
// Map control operands.
IRMapping mapping;
@@ -145,7 +146,8 @@ LoopLikeOpInterface mlir::createFused(LoopLikeOpInterface target,
fusedLoop.getLoopInductionVars();
if (fusedInductionVar) {
if (!targetInductionVar || !sourceInductionVar)
- llvm_unreachable("expected target and source loops to have induction vars");
+ llvm_unreachable(
+ "expected target and source loops to have induction vars");
mapping.map(*targetInductionVar, *fusedInductionVar);
mapping.map(*sourceInductionVar, *fusedInductionVar);
}
diff --git a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
index 91ed2a5269d74..f8246b74a5744 100644
--- a/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
+++ b/mlir/test/Dialect/SCF/transform-loop-fuse-sibling.mlir
@@ -371,6 +371,62 @@ module attributes {transform.with_named_sequence} {
}
}
+
+// -----
+
+// CHECK: #[[$MAP:.+]] = affine_map<(d0) -> (d0 * 32)
+#map = affine_map<(d0) -> (d0 * 32)>
+#map1 = affine_map<(d0, d1) -> (d0, d1)>
+module {
+ // CHECK: func.func @loop_sibling_fusion(%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}
+ func.func @loop_sibling_fusion(%arg0: tensor<128xf32>, %arg1: tensor<128x128xf16>, %arg2: tensor<128x64xf32>, %arg3: tensor<128x128xf32>) -> (tensor<128xf32>, tensor<128x128xf16>) {
+ // CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<128x128xf16>
+ // CHECK-NEXT: %[[RESULTS:.*]]:2 = scf.forall (%[[I:.*]]) in (4) shared_outs(%[[S1:.*]] = %[[ARG0]], %[[S2:.*]] = %[[ARG1]]) -> (tensor<128xf32>, tensor<128x128xf16>) {
+ // CHECK-NEXT: %[[IDX:.*]] = affine.apply #[[$MAP]](%[[I]])
+ // CHECK-NEXT: %[[SLICE0:.*]] = tensor.extract_slice %[[ARG3]][%[[IDX]], 0] [32, 1] [1, 1] : tensor<128x128xf32> to tensor<32xf32>
+ // CHECK-NEXT: %[[SLICE1:.*]] = tensor.extract_slice %[[ARG3]][%[[IDX]], 0] [32, 128] [1, 1] : tensor<128x128xf32> to tensor<32x128xf32>
+ // CHECK-NEXT: %[[SLICE2:.*]] = tensor.extract_slice %[[EMPTY]][%[[IDX]], 0] [32, 128] [1, 1] : tensor<128x128xf16> to tensor<32x128xf16>
+ // CHECK-NEXT: %[[GENERIC:.*]] = linalg.generic {{.*}} ins(%[[SLICE1]] : {{.*}}) outs(%[[SLICE2]] : {{.*}})
+ // CHECK: scf.forall.in_parallel {
+ // CHECK-NEXT: tensor.parallel_insert_slice %[[SLICE0]] into %[[S1]][%[[IDX]]] [32] [1] : tensor<32xf32> into tensor<128xf32>
+ // CHECK-NEXT: tensor.parallel_insert_slice %[[GENERIC]] into %[[S2]][%[[IDX]], 0] [32, 128] [1, 1] : tensor<32x128xf16> into tensor<128x128xf16>
+ // CHECK-NEXT: }
+ // CHECK-NEXT: } {mapping = [#gpu.warp<linear_dim_0>]}
+ // CHECK-NEXT: return %[[RESULTS]]#0, %[[RESULTS]]#1
+ %0 = scf.forall (%arg4) in (4) shared_outs(%arg5 = %arg0) -> (tensor<128xf32>) {
+ %3 = affine.apply #map(%arg4)
+ %extracted_slice = tensor.extract_slice %arg3[%3, 0] [32, 1] [1, 1] : tensor<128x128xf32> to tensor<32xf32>
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %extracted_slice into %arg5[%3] [32] [1] : tensor<32xf32> into tensor<128xf32>
+ }
+ } {mapping = [#gpu.warp<linear_dim_0>]}
+ %1 = tensor.empty() : tensor<128x128xf16>
+ %2 = scf.forall (%arg4) in (4) shared_outs(%arg5 = %arg1) -> (tensor<128x128xf16>) {
+ %3 = affine.apply #map(%arg4)
+ %extracted_slice = tensor.extract_slice %arg3[%3, 0] [32, 128] [1, 1] : tensor<128x128xf32> to tensor<32x128xf32>
+ %extracted_slice_0 = tensor.extract_slice %1[%3, 0] [32, 128] [1, 1] : tensor<128x128xf16> to tensor<32x128xf16>
+ %4 = linalg.generic {indexing_maps = [#map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice : tensor<32x128xf32>) outs(%extracted_slice_0 : tensor<32x128xf16>) {
+ ^bb0(%in: f32, %out: f16):
+ %5 = arith.truncf %in : f32 to f16
+ linalg.yield %5 : f16
+ } -> tensor<32x128xf16>
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %4 into %arg5[%3, 0] [32, 128] [1, 1] : tensor<32x128xf16> into tensor<128x128xf16>
+ }
+ } {mapping = [#gpu.warp<linear_dim_0>]}
+ return %0, %2 : tensor<128xf32>, tensor<128x128xf16>
+ }
+}
+
+module attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%root: !transform.any_op) {
+ %loops = transform.structured.match ops{["scf.forall"]} in %root : (!transform.any_op) -> !transform.any_op
+ %loop1, %loop2 = transform.split_handle %loops : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+ %loop3 = transform.loop.fuse_sibling %loop1 into %loop2 : (!transform.any_op, !transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
// -----
func.func @source_for_uses_result_of_target_for_err(%A: tensor<128xf32>, %B: tensor<128xf32>) -> (tensor<128xf32>, tensor<128xf32>) {
More information about the Mlir-commits
mailing list