[Mlir-commits] [mlir] e239424 - Add an option for unrolling loops up to a factor.
Lubomir Litchev
llvmlistbot at llvm.org
Tue Sep 8 09:23:52 PDT 2020
Author: Lubomir Litchev
Date: 2020-09-08T09:23:38-07:00
New Revision: e2394245eb28695d5eed5d7c015e99141993c723
URL: https://github.com/llvm/llvm-project/commit/e2394245eb28695d5eed5d7c015e99141993c723
DIFF: https://github.com/llvm/llvm-project/commit/e2394245eb28695d5eed5d7c015e99141993c723.diff
LOG: Add an option for unrolling loops up to a factor.
Currently, there is no option to allow for unrolling a loop up to a specific factor (specified by the user).
The code for doing that is there and there are benefits when unrolling is done to smaller loops (smaller than the factor specified).
Reviewed By: bondhugula
Differential Revision: https://reviews.llvm.org/D87111
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Passes.h
mlir/include/mlir/Dialect/Affine/Passes.td
mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
mlir/lib/Transforms/Utils/LoopUtils.cpp
mlir/test/Dialect/SCF/loop-unroll.mlir
mlir/test/lib/Transforms/TestLoopUnrolling.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index db1c3bfead94..580fbf53ae4f 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Passes.h
@@ -61,7 +61,8 @@ std::unique_ptr<OperationPass<FuncOp>> createLoopTilingPass();
/// and no callback is provided, anything passed from the command-line (if at
/// all) or the default unroll factor is used (LoopUnroll:kDefaultUnrollFactor).
std::unique_ptr<OperationPass<FuncOp>> createLoopUnrollPass(
- int unrollFactor = -1, bool unrollFull = false,
+ 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
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td
index 0e7f3e43661e..7515dbaa33d8 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Passes.td
@@ -71,6 +71,8 @@ def AffineLoopUnroll : FunctionPass<"affine-loop-unroll"> {
let options = [
Option<"unrollFactor", "unroll-factor", "unsigned", /*default=*/"4",
"Use this unroll factor for all loops being unrolled">,
+ Option<"unrollUpToFactor", "unroll-up-to-factor", "bool", /*default=*/"false",
+ "Allow unroling up to the factor specicied">,
Option<"unrollFull", "unroll-full", "bool", /*default=*/"false",
"Fully unroll loops">,
Option<"numRepetitions", "unroll-num-reps", "unsigned", /*default=*/"1",
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
index edb21384080f..3dc236f3c068 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp
@@ -9,7 +9,6 @@
// This file implements loop unrolling.
//
//===----------------------------------------------------------------------===//
-
#include "PassDetail.h"
#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
@@ -45,11 +44,13 @@ struct LoopUnroll : public AffineLoopUnrollBase<LoopUnroll> {
: AffineLoopUnrollBase<LoopUnroll>(other),
getUnrollFactor(other.getUnrollFactor) {}
explicit LoopUnroll(
- Optional<unsigned> unrollFactor = None, bool unrollFull = false,
+ Optional<unsigned> unrollFactor = None, bool unrollUpToFactor = false,
+ bool unrollFull = false,
const std::function<unsigned(AffineForOp)> &getUnrollFactor = nullptr)
: getUnrollFactor(getUnrollFactor) {
if (unrollFactor)
this->unrollFactor = *unrollFactor;
+ this->unrollUpToFactor = unrollUpToFactor;
this->unrollFull = unrollFull;
}
@@ -126,13 +127,16 @@ LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) {
if (unrollFull)
return loopUnrollFull(forOp);
// Otherwise, unroll by the given unroll factor.
+ if (unrollUpToFactor) {
+ return loopUnrollUpToFactor(forOp, unrollFactor);
+ }
return loopUnrollByFactor(forOp, unrollFactor);
}
std::unique_ptr<OperationPass<FuncOp>> mlir::createLoopUnrollPass(
- int unrollFactor, bool unrollFull,
+ int unrollFactor, bool unrollUpToFactor, bool unrollFull,
const std::function<unsigned(AffineForOp)> &getUnrollFactor) {
return std::make_unique<LoopUnroll>(
- unrollFactor == -1 ? None : Optional<unsigned>(unrollFactor), unrollFull,
- getUnrollFactor);
+ unrollFactor == -1 ? None : Optional<unsigned>(unrollFactor),
+ unrollUpToFactor, unrollFull, getUnrollFactor);
}
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index db6a071367d6..7ae45171ddbd 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -469,7 +469,6 @@ LogicalResult mlir::loopUnrollFull(AffineForOp forOp) {
LogicalResult mlir::loopUnrollUpToFactor(AffineForOp forOp,
uint64_t unrollFactor) {
Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
-
if (mayBeConstantTripCount.hasValue() &&
mayBeConstantTripCount.getValue() < unrollFactor)
return loopUnrollByFactor(forOp, mayBeConstantTripCount.getValue());
diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir
index 775188bf0ed9..134daa303ed8 100644
--- a/mlir/test/Dialect/SCF/loop-unroll.mlir
+++ b/mlir/test/Dialect/SCF/loop-unroll.mlir
@@ -2,6 +2,7 @@
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=3' | FileCheck %s --check-prefix UNROLL-BY-3
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2
// RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2
+// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO
func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index,
%arg3: memref<?xf32>) {
@@ -248,3 +249,24 @@ func @static_loop_unroll_by_3_promote_epilogue(%arg0 : memref<?xf32>) {
// UNROLL-BY-3-NEXT: }
// UNROLL-BY-3-NEXT: store %{{.*}}, %[[MEM]][%[[C9]]] : memref<?xf32>
// UNROLL-BY-3-NEXT: return
+
+
+// Test unroll-up-to functionality.
+func @static_loop_unroll_up_to_factor(%arg0 : memref<?xf32>) {
+ %0 = constant 7.0 : f32
+ %lb = constant 0 : index
+ %ub = constant 2 : index
+ affine.for %i0 = %lb to %ub {
+ store %0, %arg0[%i0] : memref<?xf32>
+ }
+ return
+}
+// UNROLL-UP-TO-LABEL: func @static_loop_unroll_up_to_factor
+// UNROLL-UP-TO-SAME: %[[MEM:.*0]]: memref<?xf32>
+// UNROLL-UP-TO-DAG: %[[C0:.*]] = constant 0 : index
+// UNROLL-UP-TO-DAG: %[[C2:.*]] = constant 2 : index
+// UNROLL-UP-TO-NEXT: %[[V0:.*]] = affine.apply {{.*}}
+// UNROLL-UP-TO-NEXT: store %{{.*}}, %[[MEM]][%[[V0]]] : memref<?xf32>
+// UNROLL-UP-TO-NEXT: %[[V1:.*]] = affine.apply {{.*}}
+// UNROLL-UP-TO-NEXT: tore %{{.*}}, %[[MEM]][%[[V1]]] : memref<?xf32>
+// UNROLL-UP-TO-NEXT: return
diff --git a/mlir/test/lib/Transforms/TestLoopUnrolling.cpp b/mlir/test/lib/Transforms/TestLoopUnrolling.cpp
index 712fddb97028..396f08b2cba3 100644
--- a/mlir/test/lib/Transforms/TestLoopUnrolling.cpp
+++ b/mlir/test/lib/Transforms/TestLoopUnrolling.cpp
@@ -55,6 +55,9 @@ class TestLoopUnrollingPass
Option<uint64_t> unrollFactor{*this, "unroll-factor",
llvm::cl::desc("Loop unroll factor."),
llvm::cl::init(1)};
+ Option<bool> unrollUpToFactor{*this, "unroll-up-to-factor",
+ llvm::cl::desc("Loop unroll up to factor."),
+ llvm::cl::init(false)};
Option<unsigned> loopDepth{*this, "loop-depth", llvm::cl::desc("Loop depth."),
llvm::cl::init(0)};
};
More information about the Mlir-commits
mailing list