[MLIR][Affine] Privatize certain escaping memrefs

[vinayaka-polymage]

Summary

During affine loop fusion, create private memrefs for escaping memrefs
too under the conditions that:

• the source is not removed after fusion, and
• the destination does not write to the memref.

This creates more fusion opportunities as illustrated in the test case.

[bondhugula]

Nice improvement. A couple of comments.

mlir/lib/Transforms/LoopFusion.cpp

1582

The "source is not removed after fusion" part of the condition needs an additional line of explanation I think.

mlir/test/Transforms/loop-fusion.mlir

1209

Nice!

3075

Typo: should be ... and has two consumers?

[vinayaka-polymage]

Thanks for the quick review, corrected the typos, and added additional
comments for clarification.

[vinayaka-polymage]

Addressed typos.

[ayzhuang]

Looks great!

[bondhugula]

Looks great - thanks!

mlir/lib/Transforms/LoopFusion.cpp

while (!worklist.empty()) {
           // account remaining dependences, escaping memrefs and the fusion
           // insertion point.
           bool removeSrcNode = canRemoveSrcNodeAfterFusion(
               srcId, dstId, bestSlice, fusedLoopInsPoint, srcEscapingMemRefs,
               mdg);
 
           DenseSet<Value> privateMemrefs;
           for (Value memref : producerConsumerMemrefs) {
-            // Don't create a private memref if 'srcNode' writes to escaping
-            // memrefs.
-            if (srcEscapingMemRefs.count(memref) > 0)
+            // If `memref` is an escaping one, do not create a private memref
+            // for the below scenarios, since doing so will leave the escaping
+            // memref unmodified as all the writes originally meant for the
+            // escaping memref would be performed on the private memref:
+            // 1. The source is to be removed after fusion,
+            // OR
+            // 2. The destination writes to `memref`.
+            if (srcEscapingMemRefs.count(memref) > 0 &&
+                (removeSrcNode || dstNode->getStoreOpCount(memref) > 0))
               continue;
 
             // Don't create a private memref if 'srcNode' has in edges on
             // 'memref' or 'dstNode' has out edges on 'memref'.
             if (mdg->getIncomingMemRefAccesses(srcId, memref) > 0 ||
                 mdg->getOutEdgeCount(dstId, memref) > 0)
               continue;
 

mlir/test/Transforms/loop-fusion.mlir

affine.for %i1 = 0 to 9 {
     %v0 = affine.load %arg0[%i1] : memref<10xf32>
   }
   // This tests that the loop nest '%i0' should not be removed after fusion
   // because it writes to memref argument '%arg0', and its read region
   // does not cover its write region (so fusion would shrink the write region
   // in the fused loop nest, so complete live out data region would not
   // be written).
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
-  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
-  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
-  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}} : memref<1xf32>
+  // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @should_fuse_live_out_arg(%{{.*}}: memref<10xf32>) {
func @should_fuse_escaping_memref_but_preserve_src_loop() -> memref<10xf32> {
   }
   affine.for %i1 = 0 to 9 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // This tests that the loop nest '%i0' should not be removed after fusion
   // because it writes to memref '%m', which is returned by the function, and
   // the '%i1' memory region does not cover '%i0' memory region.
 
-  // CHECK-DAG:   memref.alloc() : memref<10xf32>
+  // CHECK-DAG:   memref.alloc() : memref<1xf32>
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
-  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  affine.for %{{.*}} = 0 to 9 {
-  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
-  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}} : memref<1xf32>
+  // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return %{{.*}} : memref<10xf32>
   return %m : memref<10xf32>
 }
 // -----
 
 // This should fuse with the %in becoming a 1x1x1.
 func @R3_to_R2_reshape() {
func @should_fuse_multi_store_producer_and_privatize_memfefs() {
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
 
   return
 }
 
 // -----
 
-func @should_fuse_multi_store_producer_with_scaping_memrefs_and_remove_src(
+func @should_fuse_multi_store_producer_with_escaping_memrefs_and_remove_src(
     %a : memref<10xf32>, %b : memref<10xf32>) {
   %cst = constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
     affine.store %cst, %b[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 10 {
     %0 = affine.load %a[%i1] : memref<10xf32>
   }
 
   affine.for %i2 = 0 to 10 {
     %0 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Producer loop '%i0' should be removed after fusion since fusion is maximal.
-  // No memref should be privatized since they escape the function.
+  // No memref should be privatized since they escape the function, and the
+  // producer is removed after fusion.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // 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:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 
 // -----
 
-func @should_fuse_multi_store_producer_with_scaping_memrefs_and_preserve_src(
+func @should_fuse_multi_store_producer_with_escaping_memrefs_and_preserve_src(
     %a : memref<10xf32>, %b : memref<10xf32>) {
   %cst = constant 0.000000e+00 : f32
   affine.for %i0 = 0 to 10 {
     affine.store %cst, %a[%i0] : memref<10xf32>
     affine.store %cst, %b[%i0] : memref<10xf32>
   }
 
   affine.for %i1 = 0 to 5 {
     %0 = affine.load %a[%i1] : memref<10xf32>
   }
 
   affine.for %i2 = 0 to 10 {
     %0 = affine.load %b[%i2] : memref<10xf32>
   }
 
   // Loops '%i0' and '%i2' should be fused first and '%i0' should be removed
   // since fusion is maximal. Then the fused loop and '%i1' should be fused
   // and the fused loop shouldn't be removed since fusion is not maximal.
   // CHECK:       affine.for %{{.*}} = 0 to 10 {
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK:       affine.for %{{.*}} = 0 to 5 {
-  // 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<1xf32>
+  // CHECK-NEXT:    affine.store %{{.*}} : memref<10xf32>
+  // CHECK-NEXT:    affine.load %{{.*}} : memref<10xf32>
+  // CHECK-NEXT:    affine.load %{{.*}} : memref<1xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   return
 }
 
 // -----
 
func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){
   return
 }
 // CHECK-LABEL: func @call_op_does_not_prevent_fusion
 // CHECK:         affine.for
 // CHECK-NOT:     affine.for
 
 // -----
 
+// Test for source that writes to an escaping memref and has two consumers.
+// Fusion should create private memrefs in place of `%arg0` since the source is
+// not to be removed after fusion and the destinations do not write to `%arg0`.
+// This should enable both the consumers to benefit from fusion, which would not
+// be possible if private memrefs were not created.
+func @should_fuse_with_both_consumers_separately(%arg0: memref<10xf32>) {
+  %cf7 = constant 7.0 : f32
+  affine.for %i0 = 0 to 10 {
+    affine.store %cf7, %arg0[%i0] : memref<10xf32>
+  }
+  affine.for %i1 = 0 to 7 {
+    %v0 = affine.load %arg0[%i1] : memref<10xf32>
+  }
+  affine.for %i1 = 5 to 9 {
+    %v0 = affine.load %arg0[%i1] : memref<10xf32>
+  }
+  return
+}
+
+// CHECK-LABEL: func @should_fuse_with_both_consumers_separately
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.store
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.store
+// CHECK-NEXT:      affine.load
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.store
+// CHECK-NEXT:      affine.load
+
+// -----
+
 // Fusion is avoided when the slice computed is invalid. Comments below describe
 // incorrect backward slice computation. Similar logic applies for forward slice
 // as well.
 func @no_fusion_cannot_compute_valid_slice() {
   %A = memref.alloc() : memref<5xf32>
   %B = memref.alloc() : memref<6xf32>
   %C = memref.alloc() : memref<5xf32>
   %cst = constant 0. : f32