[mlir] Support hoisting whole affine for loops in LICM

[bondhugula]

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

73–78

Combine into a dyn_cast.

119–121

All comments to be terminated with a full stop.

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

84–92

Please indent for better readability.

84–92

No numbered SSA values in check lines please. Either use %{{.*}} or capture value and compare if that's relevant for the test.

253–263

Numbered SSA values shouldn't have been used here - these somehow escaped review. You could just ensure you don't introduce any more here.

286

Likewise.

458

Likewise here and below.

[ayzhuang]

Address code review comments to use dyn_cast, add full stop, indent, and change to %{{.*}}.

[bondhugula]

Looks good - thanks.

[ayzhuang]

@bondhugula Thank you for the review. @sgrechanik Please help land this.

[ayzhuang]

We'll merge it on Wednesday if there is no more comments. Thanks.

mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp

 bool isOpLoopInvariant(Operation &op, Value indVar,
                        SmallPtrSetImpl<Operation *> &opsWithUsers,
                        SmallPtrSetImpl<Operation *> &opsToHoist) {
   LLVM_DEBUG(llvm::dbgs() << "iterating on op: " << op;);
 
   if (isa<AffineIfOp>(op)) {
     if (!checkInvarianceOfNestedIfOps(&op, indVar, opsWithUsers, opsToHoist)) {
       return false;
     }
-  } else if (isa<AffineForOp>(op)) {
-    // If the body of a predicated region has a for loop, we don't hoist the
-    // 'affine.if'.
+  } else if (auto forOp = dyn_cast<AffineForOp>(op)) {
+    if (!areAllOpsInTheBlockListInvariant(forOp.getLoopBody(), indVar,
+                                          opsWithUsers, opsToHoist)) {
       return false;
+    }
   } else if (isa<AffineDmaStartOp, AffineDmaWaitOp>(op)) {
     // TODO: Support DMA ops.
     return false;
   } else if (!isa<ConstantOp>(op)) {
     // Register op in the set of ops that have users.
     opsWithUsers.insert(&op);
     if (isa<AffineMapAccessInterface>(op)) {
       Value memref = isa<AffineReadOpInterface>(op)
if (isa<AffineMapAccessInterface>(op)) {
         }
       }
     }
 
     if (op.getNumOperands() == 0 && !isa<AffineYieldOp>(op)) {
       LLVM_DEBUG(llvm::dbgs() << "\nNon-constant op with 0 operands\n");
       return false;
     }
-    for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
-      auto *operandSrc = op.getOperand(i).getDefiningOp();
+  }
+
+  // Check operands.
 
-      LLVM_DEBUG(
-          op.getOperand(i).print(llvm::dbgs() << "\nIterating on operand\n"));
 
-      // If the loop IV is the operand, this op isn't loop invariant.
-      if (indVar == op.getOperand(i)) {
-        LLVM_DEBUG(llvm::dbgs() << "\nLoop IV is the operand\n");
-        return false;
-      }
 
-      if (operandSrc != nullptr) {
-        LLVM_DEBUG(llvm::dbgs()
-                   << *operandSrc << "\nIterating on operand src\n");
+      LLVM_DEBUG(llvm::dbgs() << *operandSrc << "\nIterating on operand src\n");
 
-        // If the value was defined in the loop (outside of the
-        // if/else region), and that operation itself wasn't meant to
-        // be hoisted, then mark this operation loop dependent.
-        if (opsWithUsers.count(operandSrc) &&
-            opsToHoist.count(operandSrc) == 0) {
-          return false;
-        }
-      }
-    }
-  }
+      if (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0) {
 
   // If no operand was loop variant, mark this op for motion.
   opsToHoist.insert(&op);
   return true;
 }
 
 // Checks if all ops in a region (i.e. list of blocks) are loop invariant.
 bool areAllOpsInTheBlockListInvariant(
void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
   SmallPtrSet<Operation *, 8> opsWithUsers;
 
   for (auto &op : *loopBody) {
     // Register op in the set of ops that have users. This set is used
     // to prevent hoisting ops that depend on these ops that are
     // not being hoisted.
     if (!op.use_empty())
       opsWithUsers.insert(&op);
-    // We don't hoist for loops.
-    if (!isa<AffineForOp>(op)) {
-      if (!isa<AffineYieldOp>(op)) {
-        if (isOpLoopInvariant(op, indVar, opsWithUsers, opsToHoist)) {
-          opsToMove.push_back(&op);
-        }
-      }
-    }
-  }
 
   // For all instructions that we found to be invariant, place sequentially
   // right before the for loop.
   for (auto *op : opsToMove) {
     op->moveBefore(forOp);
   }
 
   LLVM_DEBUG(forOp->print(llvm::dbgs() << "Modified loop\n"));

mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir

affine.for %arg0 = 0 to 10 {
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // 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: affine.for %arg0 = 0 to 10 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
 
   return
 }
 
 // -----
 
 // The store-load forwarding can see through affine apply's since it relies on
func @nested_loops_code_invariant_to_both() {
   // CHECK-NEXT: %cst_0 = constant 8.000000e+00 : f32
   // CHECK-NEXT: %1 = addf %cst, %cst_0 : f32
 
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @nested_loops_inner_loops_invariant_to_outermost_loop
+func @nested_loops_inner_loops_invariant_to_outermost_loop(%m : memref<10xindex>) {
+  affine.for %arg0 = 0 to 20 {
+    affine.for %arg1 = 0 to 30 {
+      %v0 = affine.for %arg2 = 0 to 10 iter_args (%prevAccum = %arg1) -> index {
+        %v1 = affine.load %m[%arg2] : memref<10xindex>
+        %newAccum = addi %prevAccum, %v1 : index
+        affine.yield %newAccum : index
+      }
+    }
+  }
+
+  // CHECK:      affine.for %{{.*}} = 0 to 30 {
+  // CHECK-NEXT:   %{{.*}}  = affine.for %{{.*}}  = 0 to 10 iter_args(%{{.*}} = %{{.*}}) -> (index) {
+  // CHECK-NEXT:     %{{.*}}  = affine.load %{{.*}}[%{{.*}}  : memref<10xindex>
+  // CHECK-NEXT:     %{{.*}}  = addi %{{.*}}, %{{.*}} : index
+  // CHECK-NEXT:     affine.yield %{{.*}} : index
+  // CHECK-NEXT:   }
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 20 {
+  // CHECK-NEXT: }
+
+  return
+}
+
+// -----
+
 func @single_loop_nothing_invariant() {
   %m1 = memref.alloc() : memref<10xf32>
   %m2 = memref.alloc() : memref<10xf32>
   affine.for %arg0 = 0 to 10 {
     %v0 = affine.load %m1[%arg0] : memref<10xf32>
     %v1 = affine.load %m2[%arg0] : memref<10xf32>
     %v2 = addf %v0, %v1 : f32
     affine.store %v2, %m1[%arg0] : memref<10xf32>
affine.for %arg0 = 0 to 10 {
     %v0 = addf %cf7, %cf8 : f32
     affine.for %arg1 = 0 to 10 {
       %v1 = addf %cf7, %cf7 : f32
       %v3 = affine.load %m[%arg1] : memref<10xf32>
       %v2 = affine.load %m[%arg0] : memref<10xf32>
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // 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: %2 = addf %cst, %cst : f32
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
   // 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: affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT: %{{.*}} = affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
 
   return
 }
 
 // -----
 
 func @invariant_affine_if() {
   %m = memref.alloc() : memref<10xf32>
   %cf8 = constant 8.0 : f32
   affine.for %arg0 = 0 to 10 {
     affine.for %arg1 = 0 to 10 {
       affine.if affine_set<(d0, d1) : (d1 - d0 >= 0)> (%arg0, %arg0) {
           %cf9 = addf %cf8, %cf8 : f32
           affine.store %cf9, %m[%arg0] : memref<10xf32>
 
       }
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT: affine.if #set(%arg0, %arg0) {
   // CHECK-NEXT: %1 = addf %cst, %cst : f32
   // CHECK-NEXT: affine.store %1, %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: }
 
 
   return
 }
affine.for %arg1 = 0 to 10 {
       }
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // CHECK-NEXT: %1 = memref.alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT: affine.if #set(%arg0, %arg0) {
   // CHECK-NEXT: %2 = addf %cst, %cst : f32
   // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: affine.if #set(%arg0, %arg0) {
   // CHECK-NEXT: affine.store %2, %1[%arg0] : memref<10xf32>
   // CHECK-NEXT: } else {
   // CHECK-NEXT: %4 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: }
affine.for %arg1 = 0 to 10 {
           }
       }
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT: affine.if #set(%arg0, %arg0) {
   // CHECK-NEXT: %1 = addf %cst, %cst : f32
   // CHECK-NEXT: %2 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: affine.if #set(%arg0, %arg0) {
   // CHECK-NEXT: %3 = affine.load %0[%arg0] : memref<10xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: }
 
affine.store %cst, %m[%c0] : memref<10xf32>
         %v0 = affine.load %m[%arg0] : memref<10xf32>
     }
   }
 
   // CHECK: %0 = memref.alloc() : memref<10xf32>
   // CHECK-NEXT: %cst = constant 8.000000e+00 : f32
   // CHECK-NEXT: %c0 = constant 0 : index
   // CHECK-NEXT: affine.for %arg0 = 0 to 10 {
+  // CHECK-NEXT: }
+  // CHECK-NEXT: affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:   affine.store %cst, %0[%c0] : memref<10xf32>
   // CHECK-NEXT:   %1 = affine.load %0[%arg0] : memref<10xf32>
 
 
   return
 }
 
 // -----