[Mlir-commits] [mlir] f49d069 - [MLIR] Make promote single iteration optional on affine-loop-normalize
Uday Bondhugula
llvmlistbot at llvm.org
Mon Nov 28 18:29:16 PST 2022
Author: Uday Bondhugula
Date: 2022-11-29T07:58:10+05:30
New Revision: f49d069ac03baffd64b22792001ec40f973b3178
URL: https://github.com/llvm/llvm-project/commit/f49d069ac03baffd64b22792001ec40f973b3178
DIFF: https://github.com/llvm/llvm-project/commit/f49d069ac03baffd64b22792001ec40f973b3178.diff
LOG: [MLIR] Make promote single iteration optional on affine-loop-normalize
Make promote single iteration optional on affine-loop-normalize:
introduce a command-line flag and an argument on the utility. Disable it
by default since such a promotion isn't normally expected with loop
normalization: it could drop certain structure or information on the
loops that a user wanted to preserve.
Reviewed By: dcaballe
Differential Revision: https://reviews.llvm.org/D137605
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Passes.h
mlir/include/mlir/Dialect/Affine/Passes.td
mlir/include/mlir/Dialect/Affine/Utils.h
mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/test/Dialect/Affine/affine-loop-normalize.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h
index 8cc6f8c277d88..a735fc0474ee0 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.h
+++ b/mlir/include/mlir/Dialect/Affine/Passes.h
@@ -46,8 +46,11 @@ createAffineLoopInvariantCodeMotionPass();
/// ops.
std::unique_ptr<OperationPass<func::FuncOp>> createAffineParallelizePass();
-/// Apply normalization transformations to affine loop-like ops.
-std::unique_ptr<OperationPass<func::FuncOp>> createAffineLoopNormalizePass();
+/// Apply normalization transformations to affine loop-like ops. If
+/// `promoteSingleIter` is true, single iteration loops are promoted (i.e., the
+/// loop is replaced by its loop body).
+std::unique_ptr<OperationPass<func::FuncOp>>
+createAffineLoopNormalizePass(bool promoteSingleIter = false);
/// Performs packing (or explicit copying) of accessed memref regions into
/// buffers in the specified faster memory space through either pointwise copies
diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td
index 2faa3b9dc44dd..67b7eb32d1e9b 100644
--- a/mlir/include/mlir/Dialect/Affine/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Passes.td
@@ -381,6 +381,10 @@ def AffineParallelize : Pass<"affine-parallelize", "func::FuncOp"> {
def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> {
let summary = "Apply normalization transformations to affine loop-like ops";
let constructor = "mlir::createAffineLoopNormalizePass()";
+ let options = [
+ Option<"promoteSingleIter", "promote-single-iter", "bool",
+ /*default=*/"true", "Promote single iteration loops">,
+ ];
}
def LoopCoalescing : Pass<"affine-loop-coalescing", "func::FuncOp"> {
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index cea613dc97d7a..d5f02845715bd 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -159,14 +159,15 @@ vectorizeAffineLoopNest(std::vector<SmallVector<AffineForOp, 2>> &loops,
/// early if the op is already in a normalized form.
void normalizeAffineParallel(AffineParallelOp op);
-/// Normalize an affine.for op. If the affine.for op has only a single iteration
-/// only then it is simply promoted, else it is normalized in the traditional
-/// way, by converting the lower bound to zero and loop step to one. The upper
-/// bound is set to the trip count of the loop. Original loops must have a
-/// lower bound with only a single result. There is no such restriction on upper
-/// bounds. Returns success if the loop has been normalized (or is already in
-/// the normal form).
-LogicalResult normalizeAffineFor(AffineForOp op);
+/// Normalize an affine.for op. An affine.for op is normalized by converting the
+/// lower bound to zero and loop step to one. The upper bound is set to the trip
+/// count of the loop. Original loops must have a lower bound with only a single
+/// result. There is no such restriction on upper bounds. Returns success if the
+/// loop has been normalized (or is already in the normal form). If
+/// `promoteSingleIter` is true, the loop is simply promoted if it has a single
+/// iteration.
+LogicalResult normalizeAffineFor(AffineForOp op,
+ bool promoteSingleIter = false);
/// Traverse `e` and return an AffineExpr where all occurrences of `dim` have
/// been replaced by either:
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
index 93d7d17d4d9e2..f760866b96f27 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp
@@ -30,13 +30,16 @@ namespace {
/// that are already in a normalized form.
struct AffineLoopNormalizePass
: public impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> {
+ explicit AffineLoopNormalizePass(bool promoteSingleIter) {
+ this->promoteSingleIter = promoteSingleIter;
+ }
void runOnOperation() override {
- getOperation().walk([](Operation *op) {
+ getOperation().walk([&](Operation *op) {
if (auto affineParallel = dyn_cast<AffineParallelOp>(op))
normalizeAffineParallel(affineParallel);
else if (auto affineFor = dyn_cast<AffineForOp>(op))
- (void)normalizeAffineFor(affineFor);
+ (void)normalizeAffineFor(affineFor, promoteSingleIter);
});
}
};
@@ -44,6 +47,6 @@ struct AffineLoopNormalizePass
} // namespace
std::unique_ptr<OperationPass<func::FuncOp>>
-mlir::createAffineLoopNormalizePass() {
- return std::make_unique<AffineLoopNormalizePass>();
+mlir::createAffineLoopNormalizePass(bool promoteSingleIter) {
+ return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter);
}
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 0973e188b3dda..6bdebddeee2ae 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -546,15 +546,8 @@ void mlir::normalizeAffineParallel(AffineParallelOp op) {
ubExprs, op.getContext());
op.setUpperBounds(ranges.getOperands(), newUpperMap);
}
-
-/// Normalizes affine.for ops. If the affine.for op has only a single iteration
-/// only then it is simply promoted, else it is normalized in the traditional
-/// way, by converting the lower bound to zero and loop step to one. The upper
-/// bound is set to the trip count of the loop. For now, original loops must
-/// have lower bound with a single result only. There is no such restriction on
-/// upper bounds.
-LogicalResult mlir::normalizeAffineFor(AffineForOp op) {
- if (succeeded(promoteIfSingleIteration(op)))
+LogicalResult mlir::normalizeAffineFor(AffineForOp op, bool promoteSingleIter) {
+ if (promoteSingleIter && succeeded(promoteIfSingleIteration(op)))
return success();
// Check if the forop is already normalized.
diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
index 0c509748d144b..35ab7994cdf2f 100644
--- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir
@@ -1,4 +1,5 @@
// RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -affine-loop-normalize='promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=PROMOTE-SINGLE-ITER
// Normalize steps to 1 and lower bounds to 0.
@@ -39,8 +40,9 @@ func.func @relative_bounds(%arg: index) {
// Check that single iteration loop is removed and its body is promoted to the
// parent block.
-// CHECK-LABEL: func @single_iteration_loop
-func.func @single_iteration_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
+// CHECK-LABEL: func @promote_single_iter_loop
+// PROMOTE-SINGLE-ITER-LABEL: func @promote_single_iter_loop
+func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
affine.for %i = 0 to 1 {
%1 = affine.load %in[%i] : memref<1xf32>
affine.store %1, %out[%i] : memref<1xf32>
@@ -48,10 +50,10 @@ func.func @single_iteration_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
return
}
-// CHECK-NOT: affine.for
-// CHECK: affine.load
-// CHECK-NEXT: affine.store
-// CHECK-NEXT: return
+// PROMOTE-SINGLE-ITER-NEXT: arith.constant
+// PROMOTE-SINGLE-ITER-NEXT: affine.load
+// PROMOTE-SINGLE-ITER-NEXT: affine.store
+// PROMOTE-SINGLE-ITER-NEXT: return
// -----
More information about the Mlir-commits
mailing list