Add a canonicalization to erase empty AffineForOps.

AffineForOp themselves are pure and can be removed if there are no internal operations.

PiperOrigin-RevId: 266481293

diff --git a/mlir/lib/Dialect/AffineOps/AffineOps.cpp b/mlir/lib/Dialect/AffineOps/AffineOps.cpp
index 9d1fd9e..1c46b77 100644 (file)
--- 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<AffineForOp> {
+  using OpRewritePattern<AffineForOp>::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<AffineForOp> {
   using OpRewritePattern<AffineForOp>::OpRewritePattern;
@@ -1343,7 +1358,7 @@ struct AffineForLoopBoundFolder : public OpRewritePattern<AffineForOp> {
 
 void AffineForOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                               MLIRContext *context) {
-  results.insert<AffineForLoopBoundFolder>(context);
+  results.insert<AffineForEmptyLoopFolder, AffineForLoopBoundFolder>(context);
 }
 
 AffineBound AffineForOp::getLowerBound() {

diff --git a/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp b/mlir/lib/Transforms/LoopInvariantCodeMotion.cpp
index be39297..6150996 100644 (file)
--- 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<AffineTerminatorOp>(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 (file)
--- 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 (file)
--- 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
 }