shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

[mlir] Delete unroll-full option for Affine/SCF unroll pass (#164658)

Browse Source 
Make the unroll-factor take -1 as "full" and avoid potential conflict
when passing both an explicit factor and unroll-full=true.
This commit is contained in:
lonely eagle 2025-10-24 02:45:39 +08:00 committed by GitHub
parent b08bbe5ada
commit e665f245f5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 28 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlir/include/mlir/Dialect/Affine/Passes.h
View File

@ -106,7 +106,6 @@ std::unique_ptr<OperationPass<func::FuncOp>> createLoopTilingPass();
/// all) or the default unroll factor is used (LoopUnroll:kDefaultUnrollFactor).
std::unique_ptr<InterfacePass<FunctionOpInterface>> createLoopUnrollPass(
int unrollFactor = -1, bool unrollUpToFactor = false,
bool unrollFull = false,
const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr);
/// Creates a loop unroll jam pass to unroll jam by the specified factor. A

4
mlir/include/mlir/Dialect/Affine/Passes.td
View File

@ -203,12 +203,10 @@ def AffineLoopUnroll : InterfacePass<"affine-loop-unroll", "FunctionOpInterface"
let summary = "Unroll affine loops";
let constructor = "mlir::affine::createLoopUnrollPass()";
let options = [
Option<"unrollFactor", "unroll-factor", "unsigned", /*default=*/"4",
Option<"unrollFactor", "unroll-factor", "int64_t", /*default=*/"4",
"Use this unroll factor for all loops being unrolled">,
Option<"unrollUpToFactor", "unroll-up-to-factor", "bool",
/*default=*/"false", "Allow unrolling up to the factor specified">,
Option<"unrollFull", "unroll-full", "bool", /*default=*/"false",
"Fully unroll loops">,
Option<"numRepetitions", "unroll-num-reps", "unsigned", /*default=*/"1",
"Unroll innermost loops repeatedly this many times">,
Option<"unrollFullThreshold", "unroll-full-threshold", "unsigned",

21
mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
View File

@ -45,18 +45,15 @@ struct LoopUnroll : public affine::impl::AffineLoopUnrollBase<LoopUnroll> {
const std::function<unsigned(AffineForOp)> getUnrollFactor;
LoopUnroll() : getUnrollFactor(nullptr) {}
LoopUnroll(const LoopUnroll &other)
= default;
LoopUnroll(const LoopUnroll &other) = default;
explicit LoopUnroll(
std::optional<unsigned> unrollFactor = std::nullopt,
bool unrollUpToFactor = false, bool unrollFull = false,
bool unrollUpToFactor = false,
const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr)
: getUnrollFactor(getUnrollFactor) {
if (unrollFactor)
this->unrollFactor = *unrollFactor;
this->unrollUpToFactor = unrollUpToFactor;
this->unrollFull = unrollFull;
}
void runOnOperation() override;
@ -85,11 +82,17 @@ static void gatherInnermostLoops(FunctionOpInterface f,
}
void LoopUnroll::runOnOperation() {
if (!(unrollFactor.getValue() > 0 || unrollFactor.getValue() == -1)) {
emitError(UnknownLoc::get(&getContext()),
"Invalid option: 'unroll-factor' should be greater than 0 or "
"equal to -1");
return signalPassFailure();
}
FunctionOpInterface func = getOperation();
if (func.isExternal())
return;
if (unrollFull && unrollFullThreshold.hasValue()) {
if (unrollFactor.getValue() == -1 && unrollFullThreshold.hasValue()) {
// Store short loops as we walk.
SmallVector<AffineForOp, 4> loops;
@ -130,7 +133,7 @@ LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
return loopUnrollByFactor(forOp, getUnrollFactor(forOp),
/*annotateFn=*/nullptr, cleanUpUnroll);
// Unroll completely if full loop unroll was specified.
if (unrollFull)
if (unrollFactor.getValue() == -1)
return loopUnrollFull(forOp);
// Otherwise, unroll by the given unroll factor.
if (unrollUpToFactor)
@ -141,9 +144,9 @@ LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
std::unique_ptr<InterfacePass<FunctionOpInterface>>
mlir::affine::createLoopUnrollPass(
int unrollFactor, bool unrollUpToFactor, bool unrollFull,
int unrollFactor, bool unrollUpToFactor,
const std::function<unsigned(AffineForOp)> &getUnrollFactor) {
return std::make_unique<LoopUnroll>(
unrollFactor == -1 ? std::nullopt : std::optional<unsigned>(unrollFactor),
unrollUpToFactor, unrollFull, getUnrollFactor);
unrollUpToFactor, getUnrollFactor);
}

6
mlir/test/Dialect/Affine/unroll.mlir
View File

@ -1,9 +1,9 @@
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true}))" | FileCheck %s --check-prefix UNROLL-FULL
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true unroll-full-threshold=2}))" | FileCheck %s --check-prefix SHORT
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=-1}))" | FileCheck %s --check-prefix UNROLL-FULL
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=-1 unroll-full-threshold=2}))" | FileCheck %s --check-prefix SHORT
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=4}))" | FileCheck %s --check-prefix UNROLL-BY-4
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=1}))" | FileCheck %s --check-prefix UNROLL-BY-1
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix UNROLL-CLEANUP-LOOP
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(gpu.module(gpu.func(affine-loop-unroll{unroll-full=true})))" | FileCheck %s --check-prefix GPU-UNROLL-FULL
// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(gpu.module(gpu.func(affine-loop-unroll{unroll-factor=-1})))" | FileCheck %s --check-prefix GPU-UNROLL-FULL
// UNROLL-FULL-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
// UNROLL-FULL-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)>

2
mlir/test/Transforms/scf-loop-unroll.mlir
View File

@ -1,6 +1,6 @@
// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=3" -split-input-file -canonicalize | FileCheck %s
// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=1" -split-input-file -canonicalize | FileCheck %s --check-prefix UNROLL-BY-1
// RUN: mlir-opt %s --test-loop-unrolling="unroll-full=true" -split-input-file -canonicalize | FileCheck %s --check-prefix UNROLL-FULL
// RUN: mlir-opt %s --test-loop-unrolling="unroll-factor=-1" -split-input-file -canonicalize | FileCheck %s --check-prefix UNROLL-FULL
// CHECK-LABEL: scf_loop_unroll_single
func.func @scf_loop_unroll_single(%arg0 : f32, %arg1 : f32) -> f32 {

20
mlir/test/lib/Dialect/SCF/TestLoopUnrolling.cpp
View File

@ -42,11 +42,10 @@ struct TestLoopUnrollingPass
TestLoopUnrollingPass(const TestLoopUnrollingPass &) {}
explicit TestLoopUnrollingPass(uint64_t unrollFactorParam,
unsigned loopDepthParam,
bool annotateLoopParam, bool unrollFullParam) {
bool annotateLoopParam) {
unrollFactor = unrollFactorParam;
loopDepth = loopDepthParam;
annotateLoop = annotateLoopParam;
unrollFull = unrollFactorParam;
}
void getDependentDialects(DialectRegistry &registry) const override {
@ -54,6 +53,12 @@ struct TestLoopUnrollingPass
}
void runOnOperation() override {
if (!(unrollFactor.getValue() > 0 || unrollFactor.getValue() == -1)) {
emitError(UnknownLoc::get(&getContext()),
"Invalid option: 'unroll-factor' should be greater than 0 or "
"equal to -1");
return signalPassFailure();
}
SmallVector<scf::ForOp, 4> loops;
getOperation()->walk([&](scf::ForOp forOp) {
if (getNestingDepth(forOp) == loopDepth)
@ -65,15 +70,15 @@ struct TestLoopUnrollingPass
}
};
for (auto loop : loops) {
if (unrollFull)
if (unrollFactor.getValue() == -1)
(void)loopUnrollFull(loop);
else
(void)loopUnrollByFactor(loop, unrollFactor, annotateFn);
}
}
Option<uint64_t> unrollFactor{*this, "unroll-factor",
llvm::cl::desc("Loop unroll factor."),
llvm::cl::init(1)};
Option<int64_t> unrollFactor{*this, "unroll-factor",
llvm::cl::desc("Loop unroll factor."),
llvm::cl::init(1)};
Option<bool> annotateLoop{*this, "annotate",
llvm::cl::desc("Annotate unrolled iterations."),
llvm::cl::init(false)};
@ -82,9 +87,6 @@ struct TestLoopUnrollingPass
llvm::cl::init(false)};
Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
llvm::cl::init(0)};
Option<bool> unrollFull{*this, "unroll-full",
llvm::cl::desc("Full unroll loops."),
llvm::cl::init(false)};
};
} // namespace