[mlir][AffineOps] Adds affine loop fusion transformation function to LoopFusionUtils.

[andydavis1]

Summary

Adds affine loop fusion transformation function to LoopFusionUtils.
Updates TestLoopFusion utility to run loop fusion transformation until a fixed point is reached.
Adds unit tests to test the transformation.
Includes ASAN bug fix for D73190.

[dcaballe]

Thanks, Andy. LGTM. Please wait a bit so that others have time to comment.

[bondhugula]

Looks good to me; just a few minor doc comments.

mlir/test/lib/Transforms/TestLoopFusion.cpp

76–77

Document return value?

129–131

Document return value.

145–162

Function doc comment.

[rriddle]

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

274

nit: Spell out auto here.

mlir/test/lib/Transforms/TestLoopFusion.cpp

199

do {

} while (iterateLoops(...)); ?

213

nit: Spell out auto here.

214

Drop trivial braces.

[andydavis1]

addressing comments.

[andydavis1]

addressing comments.

[andydavis1]

addressing review comments

mlir/include/mlir/Transforms/LoopFusionUtils.h

 /// returns a FusionResult explaining why fusion is not feasible.
 /// NOTE: This function is not feature complete and should only be used in
 /// testing.
 /// TODO(andydavis) Update comments when this function is fully implemented.
 FusionResult canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
                           unsigned dstLoopDepth,
                           ComputationSliceState *srcSlice);
 
+/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
+/// and source slice loop bounds specified in 'srcSlice'.
+void fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
+               ComputationSliceState *srcSlice);
+
 /// LoopNestStats aggregates various per-loop statistics (eg. loop trip count
 /// and operation count) for a loop nest up until (and including) the innermost
 /// loop body.
 struct LoopNestStats {
   /// Map from AffineForOp to immediate child AffineForOps in its loop body.
   DenseMap<Operation *, SmallVector<AffineForOp, 2>> loopMap;
   /// Map from AffineForOp to count of operations in its loop body.
   DenseMap<Operation *, uint64_t> opCountMap;

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Transforms/LoopUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "loop-fusion-utils"
 
 using namespace mlir;
if (failed(mlir::computeSliceUnion(opsA, opsB, dstLoopDepth, numCommonLoops,
                                      isSrcForOpBeforeDstForOp, srcSlice))) {
     LLVM_DEBUG(llvm::dbgs() << "computeSliceUnion failed\n");
     return FusionResult::FailPrecondition;
   }
 
   return FusionResult::Success;
 }
 
+/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
+/// and source slice loop bounds specified in 'srcSlice'.
+void mlir::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
+                     ComputationSliceState *srcSlice) {
+  // Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'.
+  OpBuilder b(srcSlice->insertPoint->getBlock(), srcSlice->insertPoint);
+  BlockAndValueMapping mapper;
+  b.clone(*srcForOp, mapper);
+
+  // Update 'sliceLoopNest' upper and lower bounds from computed 'srcSlice'.
+  SmallVector<AffineForOp, 4> sliceLoops;
+  for (unsigned i = 0, e = srcSlice->ivs.size(); i < e; ++i) {
+    auto loopIV = mapper.lookupOrNull(srcSlice->ivs[i]);
+    if (!loopIV)
+      continue;
+    auto forOp = getForInductionVarOwner(loopIV);
+    sliceLoops.push_back(forOp);
+    if (AffineMap lbMap = srcSlice->lbs[i])
+      forOp.setLowerBound(srcSlice->lbOperands[i], lbMap);
+    if (AffineMap ubMap = srcSlice->ubs[i])
+      forOp.setUpperBound(srcSlice->ubOperands[i], ubMap);
+  }
+
+  // Promote any single iteration slice loops.
+  for (AffineForOp forOp : sliceLoops)
+    promoteIfSingleIteration(forOp);
+}
+
 /// Collect loop nest statistics (eg. loop trip count and operation count)
 /// in 'stats' for loop nest rooted at 'forOp'. Returns true on success,
 /// returns false otherwise.
 bool mlir::getLoopNestStats(AffineForOp forOpRoot, LoopNestStats *stats) {
   auto walkResult = forOpRoot.walk([&](AffineForOp forOp) {
     auto *childForOp = forOp.getOperation();
     auto *parentForOp = forOp.getParentOp();
     if (!llvm::isa<FuncOp>(parentForOp)) {

mlir/test/Transforms/loop-fusion-transformation.mlir

+// RUN: mlir-opt %s -test-loop-fusion -test-loop-fusion-transformation -split-input-file -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @slice_depth1_loop_nest() {
+func @slice_depth1_loop_nest() {
+  %0 = alloc() : memref<100xf32>
+  %cst = constant 7.000000e+00 : f32
+  affine.for %i0 = 0 to 16 {
+    affine.store %cst, %0[%i0] : memref<100xf32>
+  }
+  affine.for %i1 = 0 to 5 {
+    %1 = affine.load %0[%i1] : memref<100xf32>
+  }
+  // CHECK:      affine.for %[[IV0:.*]] = 0 to 5 {
+  // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%[[IV0]]] : memref<100xf32>
+  // CHECK-NEXT:   affine.load %{{.*}}[%[[IV0]]] : memref<100xf32>
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @should_fuse_reduction_to_pointwise() {
+func @should_fuse_reduction_to_pointwise() {
+  %a = alloc() : memref<10x10xf32>
+  %b = alloc() : memref<10xf32>
+  %c = alloc() : memref<10xf32>
+
+  %cf7 = constant 7.0 : f32
+
+  affine.for %i0 = 0 to 10 {
+    affine.for %i1 = 0 to 10 {
+      %v0 = affine.load %b[%i0] : memref<10xf32>
+      %v1 = affine.load %a[%i0, %i1] : memref<10x10xf32>
+      %v3 = addf %v0, %v1 : f32
+      affine.store %v3, %b[%i0] : memref<10xf32>
+    }
+  }
+  affine.for %i2 = 0 to 10 {
+    %v4 = affine.load %b[%i2] : memref<10xf32>
+    affine.store %v4, %c[%i2] : memref<10xf32>
+  }
+
+  // Match on the fused loop nest.
+  // Should fuse in entire inner loop on %i1 from source loop nest, as %i1
+  // is not used in the access function of the store/load on %b.
+  // CHECK:       affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:      affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
+  // CHECK-NEXT:      addf %{{.*}}, %{{.*}} : f32
+  // CHECK-NEXT:      affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    }
+  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:  }
+  // CHECK-NEXT:  return
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @should_fuse_avoiding_dependence_cycle() {
+func @should_fuse_avoiding_dependence_cycle() {
+  %a = alloc() : memref<10xf32>
+  %b = alloc() : memref<10xf32>
+  %c = alloc() : memref<10xf32>
+
+  %cf7 = constant 7.0 : f32
+
+  // Set up the following dependences:
+  // 1) loop0 -> loop1 on memref '%{{.*}}'
+  // 2) loop0 -> loop2 on memref '%{{.*}}'
+  // 3) loop1 -> loop2 on memref '%{{.*}}'
+  affine.for %i0 = 0 to 10 {
+    %v0 = affine.load %a[%i0] : memref<10xf32>
+    affine.store %cf7, %b[%i0] : memref<10xf32>
+  }
+  affine.for %i1 = 0 to 10 {
+    affine.store %cf7, %a[%i1] : memref<10xf32>
+    %v1 = affine.load %c[%i1] : memref<10xf32>
+  }
+  affine.for %i2 = 0 to 10 {
+    %v2 = affine.load %b[%i2] : memref<10xf32>
+    affine.store %cf7, %c[%i2] : memref<10xf32>
+  }
+  // Fusing loop first loop into last would create a cycle:
+  //   {1} <--> {0, 2}
+  // However, we can avoid the dependence cycle if we first fuse loop0 into
+  // loop1:
+  //   {0, 1) --> {2}
+  // Then fuse this loop nest with loop2:
+  //   {0, 1, 2}
+  //
+  // CHECK:      affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:   affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return
+  return
+}

mlir/test/lib/Transforms/TestLoopFusion.cpp

static llvm::cl::opt<bool> clTestDependenceCheck(
     llvm::cl::desc("Enable testing of loop fusion dependence check"),
     llvm::cl::cat(clOptionsCategory));
 
 static llvm::cl::opt<bool> clTestSliceComputation(
     "test-loop-fusion-slice-computation",
     llvm::cl::desc("Enable testing of loop fusion slice computation"),
     llvm::cl::cat(clOptionsCategory));
 
+static llvm::cl::opt<bool> clTestLoopFusionTransformation(
+    "test-loop-fusion-transformation",
+    llvm::cl::desc("Enable testing of loop fusion transformation"),
+    llvm::cl::cat(clOptionsCategory));
+
 namespace {
 
 struct TestLoopFusion : public FunctionPass<TestLoopFusion> {
   void runOnFunction() override;
 };
 
 } // end anonymous namespace
 
if (auto forOp = dyn_cast<AffineForOp>(op)) {
       loopsAtDepth.push_back(forOp);
       gatherLoops(forOp.getBody(), currLoopDepth + 1, depthToLoops);
     }
   }
 }
 
 // Run fusion dependence check on 'loops[i]' and 'loops[j]' at loop depths
 // in range ['loopDepth' + 1, 'maxLoopDepth'].
 // Emits a remark on 'loops[i]' if a fusion-preventing dependence exists.
-static void testDependenceCheck(SmallVector<AffineForOp, 2> &loops, unsigned i,
-                                unsigned j, unsigned loopDepth,
+// Returns false as IR is not transformed.
+static bool testDependenceCheck(AffineForOp srcForOp, AffineForOp dstForOp,
+                                unsigned i, unsigned j, unsigned loopDepth,
                                 unsigned maxLoopDepth) {
-  AffineForOp srcForOp = loops[i];
-  AffineForOp dstForOp = loops[j];
   mlir::ComputationSliceState sliceUnion;
   for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
     FusionResult result =
         mlir::canFuseLoops(srcForOp, dstForOp, d, &sliceUnion);
     if (result.value == FusionResult::FailBlockDependence) {
       srcForOp.getOperation()->emitRemark("block-level dependence preventing"
                                           " fusion of loop nest ")
           << i << " into loop nest " << j << " at depth " << loopDepth;
     }
   }
+  return false;
 }
 
 // Returns the index of 'op' in its block.
 static unsigned getBlockIndex(Operation &op) {
   unsigned index = 0;
   for (auto &opX : *op.getBlock()) {
     if (&op == &opX)
       break;
for (unsigned k = 0, e = sliceUnion.lbs.size(); k < e; ++k) {
     os << "] ";
   }
   return os.str();
 }
 
 // Computes fusion slice union on 'loops[i]' and 'loops[j]' at loop depths
 // in range ['loopDepth' + 1, 'maxLoopDepth'].
 // Emits a string representation of the slice union as a remark on 'loops[j]'.
-static void testSliceComputation(SmallVector<AffineForOp, 2> &loops, unsigned i,
-                                 unsigned j, unsigned loopDepth,
+// Returns false as IR is not transformed.
+static bool testSliceComputation(AffineForOp forOpA, AffineForOp forOpB,
+                                 unsigned i, unsigned j, unsigned loopDepth,
                                  unsigned maxLoopDepth) {
-  AffineForOp forOpA = loops[i];
-  AffineForOp forOpB = loops[j];
   for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
     mlir::ComputationSliceState sliceUnion;
     FusionResult result = mlir::canFuseLoops(forOpA, forOpB, d, &sliceUnion);
     if (result.value == FusionResult::Success) {
       forOpB.getOperation()->emitRemark("slice (")
           << " src loop: " << i << ", dst loop: " << j << ", depth: " << d
           << " : " << getSliceStr(sliceUnion) << ")";
     }
   }
+  return false;
 }
 
-void TestLoopFusion::runOnFunction() {
-  // Gather all AffineForOps by loop depth.
-  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
-  for (auto &block : getFunction()) {
-    gatherLoops(&block, /*currLoopDepth=*/0, depthToLoops);
+// Attempts to fuse 'forOpA' into 'forOpB' at loop depths in range
+// ['loopDepth' + 1, 'maxLoopDepth'].
+// Returns true if loops were successfully fused, false otherwise.
+static bool testLoopFusionTransformation(AffineForOp forOpA, AffineForOp forOpB,
+                                         unsigned i, unsigned j,
+                                         unsigned loopDepth,
+                                         unsigned maxLoopDepth) {
+  for (unsigned d = loopDepth + 1; d <= maxLoopDepth; ++d) {
+    mlir::ComputationSliceState sliceUnion;
+    FusionResult result = mlir::canFuseLoops(forOpA, forOpB, d, &sliceUnion);
+    if (result.value == FusionResult::Success) {
+      mlir::fuseLoops(forOpA, forOpB, &sliceUnion);
+      // Note: 'forOpA' is removed to simplify test output. A proper loop
+      // fusion pass should check the data dependence graph and run memref
+      // region analysis to ensure removing 'forOpA' is safe.
+      forOpA.erase();
+      return true;
+    }
   }
+  return false;
+}
+
+using LoopFunc = function_ref<bool(AffineForOp, AffineForOp, unsigned, unsigned,
+                                   unsigned, unsigned)>;
 
-  // Run tests on all combinations of src/dst loop nests in 'depthToLoops'.
+// If 'return_on_change' is true, returns on first invocation of 'fn' which
+// returns true.
+static bool
+iterateLoops(DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops,
+             LoopFunc fn, bool return_on_change = false) {
+  bool changed = false;
   for (auto &depthAndLoops : depthToLoops) {
     unsigned loopDepth = depthAndLoops.first;
     auto &loops = depthAndLoops.second;
     unsigned numLoops = loops.size();
     for (unsigned j = 0; j < numLoops; ++j) {
       for (unsigned k = 0; k < numLoops; ++k) {
-        if (j == k)
-          continue;
-        if (clTestDependenceCheck)
-          testDependenceCheck(loops, j, k, loopDepth, depthToLoops.size());
-        if (clTestSliceComputation)
-          testSliceComputation(loops, j, k, loopDepth, depthToLoops.size());
+        if (j != k)
+          changed |=
+              fn(loops[j], loops[k], j, k, loopDepth, depthToLoops.size());
+        if (changed && return_on_change)
+          return true;
       }
     }
   }
+  return changed;
+}
+
+void TestLoopFusion::runOnFunction() {
+  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  if (clTestLoopFusionTransformation) {
+    // Run loop fusion until a fixed point is reached.
+    do {
+      depthToLoops.clear();
+      // Gather all AffineForOps by loop depth.
+      for (auto &block : getFunction())
+        gatherLoops(&block, /*currLoopDepth=*/0, depthToLoops);
+
+      // Try to fuse all combinations of src/dst loop nests in 'depthToLoops'.
+    } while (iterateLoops(depthToLoops, testLoopFusionTransformation,
+                          /*return_on_change=*/true));
+    return;
+  }
+
+  // Gather all AffineForOps by loop depth.
+  for (Block &block : getFunction())
+    gatherLoops(&block, /*currLoopDepth=*/0, depthToLoops);
+
+  // Run tests on all combinations of src/dst loop nests in 'depthToLoops'.
+  if (clTestDependenceCheck)
+    iterateLoops(depthToLoops, testDependenceCheck);
+  if (clTestSliceComputation)
+    iterateLoops(depthToLoops, testSliceComputation);
 }
 
 static PassRegistration<TestLoopFusion>
     pass("test-loop-fusion", "Tests loop fusion utility functions.");