From 9c8a8a7d0da011f1570733265020a40079127b05 Mon Sep 17 00:00:00 2001 From: River Riddle Date: Fri, 30 Aug 2019 16:49:01 -0700 Subject: [PATCH] Add a canonicalization to erase empty AffineForOps. AffineForOp themselves are pure and can be removed if there are no internal operations. PiperOrigin-RevId: 266481293 --- mlir/lib/Dialect/AffineOps/AffineOps.cpp | 17 +++++- mlir/lib/Transforms/LoopInvariantCodeMotion.cpp | 6 -- mlir/test/AffineOps/canonicalize.mlir | 10 ++++ .../Transforms/loop-invariant-code-motion.mlir | 68 ++++++++-------------- 4 files changed, 49 insertions(+), 52 deletions(-) diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp index 9d1fd9e..1c46b77 100644 --- a/mlir/lib/Dialect/AffineOps/AffineOps.cpp +++ b/mlir/lib/Dialect/AffineOps/AffineOps.cpp @@ -1283,6 +1283,21 @@ void print(OpAsmPrinter *p, AffineForOp op) { } namespace { +/// This is a pattern to fold trivially empty loops. +struct AffineForEmptyLoopFolder : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + PatternMatchResult matchAndRewrite(AffineForOp forOp, + PatternRewriter &rewriter) const override { + // Check that the body only contains a terminator. + auto *body = forOp.getBody(); + if (std::next(body->begin()) != body->end()) + return matchFailure(); + rewriter.replaceOp(forOp, llvm::None); + return matchSuccess(); + } +}; + /// This is a pattern to fold constant loop bounds. struct AffineForLoopBoundFolder : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -1343,7 +1358,7 @@ struct AffineForLoopBoundFolder : public OpRewritePattern { void AffineForOp::getCanonicalizationPatterns(OwningRewritePatternList &results, MLIRContext *context) { - results.insert(context); + results.insert(context); } AffineBound AffineForOp::getLowerBound() { diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp index be39297..6150996 100644 --- a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp +++ b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp @@ -228,12 +228,6 @@ void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) { } LLVM_DEBUG(forOp.getOperation()->print(llvm::dbgs() << "Modified loop\n")); - - // If the for loop body has a single operation (the terminator), erase it. - if (forOp.getBody()->getOperations().size() == 1) { - assert(isa(forOp.getBody()->front())); - forOp.erase(); - } } void LoopInvariantCodeMotion::runOnFunction() { diff --git a/mlir/test/AffineOps/canonicalize.mlir b/mlir/test/AffineOps/canonicalize.mlir index 8b0c5e6..077f97e 100644 --- a/mlir/test/AffineOps/canonicalize.mlir +++ b/mlir/test/AffineOps/canonicalize.mlir @@ -414,3 +414,13 @@ func @constant_fold_bounds(%N : index) { } return } + +// ----- + +// CHECK-LABEL: func @fold_empty_loop() { +func @fold_empty_loop() { + // CHECK-NOT: affine.for + affine.for %i = 0 to 10 { + } + return +} diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir index 6fc03ae..f7143b7 100644 --- a/mlir/test/Transforms/loop-invariant-code-motion.mlir +++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir @@ -18,8 +18,7 @@ func @nested_loops_both_having_invariant_code() { // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32 // CHECK-NEXT: affine.for %arg0 = 0 to 10 { // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -58,7 +57,7 @@ func @nested_loops_code_invariant_to_both() { // CHECK-NEXT: %cst = constant 7.000000e+00 : f32 // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32 // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32 - // CHECK-NEXT: return + return } @@ -79,12 +78,10 @@ func @single_loop_nothing_invariant() { // CHECK-NEXT: %3 = affine.load %1[%arg0] : memref<10xf32> // CHECK-NEXT: %4 = addf %2, %3 : f32 // CHECK-NEXT: affine.store %4, %0[%arg0] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: return + return } - func @invariant_code_inside_affine_if() { %m = alloc() : memref<10xf32> %cf8 = constant 8.0 : f32 @@ -106,8 +103,7 @@ func @invariant_code_inside_affine_if() { // CHECK-NEXT: %2 = addf %cst, %cst : f32 // CHECK-NEXT: affine.store %2, %0[%arg0] : memref<10xf32> // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -133,13 +129,11 @@ func @dependent_stores() { // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32 // CHECK-NEXT: %2 = addf %cst, %cst : f32 // CHECK-NEXT: affine.for %arg0 = 0 to 10 { - + // CHECK-NEXT: affine.for %arg1 = 0 to 10 { // CHECK-NEXT: affine.store %2, %0[%arg1] : memref<10xf32> // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -164,11 +158,10 @@ func @independent_stores() { // CHECK-NEXT: %2 = addf %cst, %cst : f32 // CHECK-NEXT: affine.for %arg0 = 0 to 10 { // CHECK-NEXT: affine.for %arg1 = 0 to 10 { - // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32> + // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32> // CHECK-NEXT: affine.store %2, %0[%arg1] : memref<10xf32> // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -194,10 +187,8 @@ func @load_dependent_store() { // CHECK-NEXT: affine.for %arg0 = 0 to 10 { // CHECK-NEXT: affine.for %arg1 = 0 to 10 { // CHECK-NEXT: affine.store %1, %0[%arg1] : memref<10xf32> - // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32> + return } @@ -221,12 +212,10 @@ func @load_after_load() { // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32 // CHECK-NEXT: %2 = addf %cst, %cst : f32 // CHECK-NEXT: affine.for %arg0 = 0 to 10 { - // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32> + // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32> // CHECK-NEXT: affine.for %arg1 = 0 to 10 { - // CHECK-NEXT: %4 = affine.load %0[%arg1] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + // CHECK-NEXT: %4 = affine.load %0[%arg1] : memref<10xf32> + return } @@ -250,8 +239,7 @@ func @invariant_affine_if() { // CHECK-NEXT: %1 = addf %cst, %cst : f32 // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32> // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -278,8 +266,7 @@ func @invariant_affine_if2() { // CHECK-NEXT: affine.store %1, %0[%arg1] : memref<10xf32> // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -311,8 +298,7 @@ func @invariant_affine_nested_if() { // CHECK-NEXT: } // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -348,8 +334,7 @@ func @invariant_affine_nested_if_else() { // CHECK-NEXT: } // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -385,8 +370,7 @@ func @invariant_affine_nested_if_else2() { // CHECK-NEXT: %4 = affine.load %0[%arg0] : memref<10xf32> // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -417,8 +401,7 @@ func @invariant_affine_nested_if2() { // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32> // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -450,8 +433,7 @@ func @invariant_affine_for_inside_affine_if() { // CHECK-NEXT: } // CHECK-NEXT: } // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -472,8 +454,7 @@ func @invariant_constant_and_load() { // CHECK-NEXT: %2 = affine.load %1[%c0] : memref<100xf32> // CHECK-NEXT: affine.for %arg0 = 0 to 5 { // CHECK-NEXT: affine.store %2, %0[%arg0] : memref<100xf32> - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -497,9 +478,7 @@ func @nested_load_store_same_memref() { // CHECK-NEXT: %1 = affine.load %0[%c0] : memref<10xf32> // CHECK-NEXT: affine.for %arg1 = 0 to 10 { // CHECK-NEXT: affine.store %cst, %0[%arg1] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: } - // CHECK-NEXT: return + return } @@ -522,8 +501,7 @@ func @nested_load_store_same_memref2() { // CHECK-NEXT: affine.for %arg0 = 0 to 10 { // CHECK-NEXT: affine.store %cst, %0[%c0] : memref<10xf32> // CHECK-NEXT: %1 = affine.load %0[%arg0] : memref<10xf32> - // CHECK-NEXT: } - // CHECK-NEXT: return + return } -- 2.7.4