[mlir][Affine] Revisit fusion candidates after successful fusion

[dcaballe]

Summary

Current fusion algorithm doesn't have full support for multi-store producers.
For some scenarios, this means that some loop nests may or may not be fused
depending on the order they are visited. For example, if we have A->B, A->C, B->C
and try to fuse A+B first, it'll fail due to A's multiple outgoing edges. However,
after fusing B+C, A would have a single outgoing edge and could be fused. This is
currently not happening if we visit A before fusing B+C because we do not visit A
again after B+C fusion.

This patch changes the fusion algorithm so that after fusing two loop nests (B+C)
we revisit previously visited nodes (A) so that they are considered again for
fusion in the context of the new fused loop nest.

[ftynse]

The reasoning is sound to me, but deferring approval to Andy who understands this algorithm.

[bondhugula]

Thanks for improving this. Agree that there may be an opportunity to make things more robust.

mlir/test/Transforms/loop-fusion.mlir

2528–2537

Nice.

[andydavis1]

Thanks for improving this. Looks good to me...

mlir/lib/Transforms/LoopFusion.cpp

while (!worklist.empty()) {
 
             // Collect dst loop stats after memref privatization transformation.
             LoopNestStateCollector dstLoopCollector;
             dstLoopCollector.collect(dstAffineForOp.getOperation());
 
             // Add new load ops to current Node load op list 'loads' to
             // continue fusing based on new operands.
             for (auto *loadOpInst : dstLoopCollector.loadOpInsts) {
-              auto loadMemRef =
-                  cast<AffineReadOpInterface>(loadOpInst).getMemRef();
               // NOTE: Change 'loads' to a hash set in case efficiency is an
               // issue. We still use a vector since it's expected to be small.
-              if (visitedMemrefs.count(loadMemRef) == 0 &&
-                  !llvm::is_contained(loads, loadOpInst))
+              if (!llvm::is_contained(loads, loadOpInst))
                 loads.push_back(loadOpInst);
             }
+            // Clear visited memrefs after fusion so that previously visited src
+            // nodes are considered for fusion again in the context of the new
+            // fused node.
+            // TODO: This shouldn't be necessary if we visited candidates in the
+            // dependence graph in post-order or once we fully support
+            // multi-store producers. Currently, in a multi-store producer
+            // scenario such as A->B, A->C, B->C, we fail to fuse A+B due to the
+            // multiple outgoing edges. However, after fusing B+C, A has a
+            // single outgoing edge and can be fused if we revisit it in the
+            // context of the new fused B+C node.
+            visitedMemrefs.clear();
 
             // Clear and add back loads and stores.
             mdg->clearNodeLoadAndStores(dstNode->id);
             mdg->addToNode(dstId, dstLoopCollector.loadOpInsts,
                            dstLoopCollector.storeOpInsts);
             // Remove old src loop nest if it no longer has outgoing dependence
             // edges, and if it does not write to a memref which escapes the
             // function. If 'writesToLiveInOrOut' is true, then 'srcNode' has

mlir/test/Transforms/loop-fusion.mlir

 }
 // CHECK:      affine.for
 // CHECK-NEXT:   affine.for
 // CHECK-NEXT:     affine.vector_load
 // CHECK-NEXT:     affine.vector_store
 // CHECK-NEXT:     affine.vector_load
 // CHECK-NEXT:     affine.vector_store
 // CHECK-NOT:  affine.for
+
+// -----
+
+// CHECK-LABEL: func @multi_outgoing_edges
+func @multi_outgoing_edges(%in0 : memref<32xf32>,
+                      %in1 : memref<32xf32>) {
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = addf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = subf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = mulf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = divf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  return
+}
+
+// CHECK:      affine.for
+// CHECK-NOT:  affine.for
+// CHECK:        addf
+// CHECK-NOT:  affine.for
+// CHECK:        subf
+// CHECK-NOT:  affine.for
+// CHECK:        mulf
+// CHECK-NOT:  affine.for
+// CHECK:        divf