[mlir] 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.

[merge_guards_bot]

Unit tests: pass. 62057 tests passed, 0 failed and 784 were skipped.

clang-tidy: unknown.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

258

nit: Add newlines before comments.

272

Remove trivial braces.

mlir/test/lib/Transforms/TestLoopFusion.cpp

193

Remove trivial braces.

[dcaballe]

Thanks for working on this Andy! Some minor comments

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

4

nit: remove this first ----- so that we don't process an empty chunk?

18

nit: for all these tests, please remove unnecessary lhs (e.g., %{{.*}} =) for loads, addf, and other ops if we don't capture their output in a variable

mlir/test/lib/Transforms/TestLoopFusion.cpp

180

Trye -> Try

187

Would it make sense to have only the fixed point optimization "mode"? Do we need to keep the old mode for some reason (other than, maybe, fixing some tests)?
If so, probably better to create an independent pass (TestFixedPointLoopFusion?) for the fixed point mode since it's independent from the other mode. In that way, we wouldn't need to use both flags -test-loop-fusion -test-loop-fusion-transformation for testing.

[bondhugula]

Adds affine loop fusion transformation function to LoopFusionUtils.

Can you clarify here if it's moving something as opposed to introducing something new?

[bondhugula]

mlir/include/mlir/Transforms/LoopFusionUtils.h

56

A bit strange to have the destination insertion point stored in srcSlice! Separate arg?

mlir/lib/Transforms/Utils/LoopFusionUtils.cpp

257

srcForOp should do?

mlir/test/lib/Transforms/TestLoopFusion.cpp

45

"test-loop-fusion" sounds fine?

[andydavis1]

thanks. Sending update ...

mlir/include/mlir/Transforms/LoopFusionUtils.h

56

They are inter-dependent: each source slice is computed for a specific dest insertion point.

mlir/test/lib/Transforms/TestLoopFusion.cpp

45

That name collides with another flag name. Eventually, we'll be able to get rid of the other test modes and collapse things down to one flag, but probably too soon right now.

187

Once things stabilize, we can move towards having a fixed point mode. For now, its nice to have specific tests that check slice computation and fusion dependence checks.

[andydavis1]

addresing comments

[merge_guards_bot]

Unit tests: pass. 62323 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[dcaballe]

LGTM. Just a minor suggestion.

mlir/test/lib/Transforms/TestLoopFusion.cpp

187

Ok, it makes sense to me, thanks! I would still suggest moving the transformation part (i.e, the code under the if condition (lines 187-200) to its own pass within this file so that we don't have to use both flags to test the transformation flavor. That code seems independent from the rest of the test pass.

[bondhugula]

LGTM - thanks!

[andydavis1]

rebasing with head

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[mehdi_amini]

Reverted in 2724ada8d2ddfec8f1e3d39f1d02f8b5d5224f1b as this broke the ASAN checker for a use-after-free in mlir::canFuseLoops(mlir::AffineForOp, mlir::AffineForOp, unsigned int, mlir::ComputationSliceState*) lib/Transforms/Utils/LoopFusionUtils.cpp:202:41

[andydavis1]

Thanks. A fix is on the way...

mlir/include/mlir/Transforms/LoopFusionUtils.h

 /// in the same block and dependences would not be violated). Otherwise
 /// 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/Analysis/Utils.h"
 #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;
FusionResult mlir::canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
   // from 'forOpA' and 'forOpB'.
   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 (auto 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(
     "test-loop-fusion-dependence-check",
     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
 
 std::unique_ptr<OpPassBase<FuncOp>> mlir::createTestLoopFusionPass() {
   return std::make_unique<TestLoopFusion>();
 }
 
 // Gathers all AffineForOps in 'block' at 'currLoopDepth' in 'depthToLoops'.
 static void
 gatherLoops(Block *block, unsigned currLoopDepth,
             DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops) {
   auto &loopsAtDepth = depthToLoops[currLoopDepth];
   for (auto &op : *block) {
     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,
+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 << '[';
     sliceUnion.lbs[k].print(os);
     os << ", ";
     sliceUnion.ubs[k].print(os);
     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,
+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);
+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'.
+static bool
+iterateLoops(DenseMap<unsigned, SmallVector<AffineForOp, 2>> &depthToLoops,
+             LoopFunc fn) {
+  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());
       }
     }
   }
+  return changed;
+}
+
+void TestLoopFusion::runOnFunction() {
+  DenseMap<unsigned, SmallVector<AffineForOp, 2>> depthToLoops;
+  if (clTestLoopFusionTransformation) {
+    // Run loop fusion until a fixed point is reached.
+    bool changed = true;
+    while (changed) {
+      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'.
+      changed = iterateLoops(depthToLoops, testLoopFusionTransformation);
+    }
+    return;
+  }
+
+  // Gather all AffineForOps by loop depth.
+  for (auto &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.");