[MLIR][Affine] Fix/improve affine sibling fusion

[bondhugula]

Summary

The sibling fusion profitability checks shouldn't rely on the presence
of a store op in the sibling. The reuse is between the loads.

Fixes issues raised at https://discourse.llvm.org/t/understanding-the-affine-loop-fusion-pass/69452

[bondhugula]

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

1878–1884

@dcaballe @andydavis1 This check can be completely removed. The single store was being used in the profitability check and it wasn't needed/correct to do that.

[dcaballe]

it makes sense. I think sibling fusion fell behind when I generalized producer-consumer fusion as I didn't look too much into it at that time.

[ayzhuang]

mlir/test/Transforms/loop-fusion.mlir

1205–1206

Please update the comment.

[bondhugula]

Address review comment.

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

1878–1884

Removing this will be addressed in the next PR since it exposes bugs due to missing checks for sibling fusion.

mlir/test/Transforms/loop-fusion.mlir

1205–1206

Thanks, done.

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

while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
       // Check if fusion would be profitable and at what depth.
 
       // Get unique 'sibNode' load op to 'memref'.
       SmallVector<Operation *, 2> sibLoadOpInsts;
       sibNode->getLoadOpsForMemref(memref, &sibLoadOpInsts);
       // Currently findSiblingNodeToFuse searches for siblings with one load.
       assert(sibLoadOpInsts.size() == 1);
       Operation *sibLoadOpInst = sibLoadOpInsts[0];
-      assert(!sibNode->stores.empty());
-      // TODO: Choose the store which postdominates all other stores.
-      auto *sibStoreOpInst = sibNode->stores.back();
 
       // Gather 'dstNode' load ops to 'memref'.
       SmallVector<Operation *, 2> dstLoadOpInsts;
       dstNode->getLoadOpsForMemref(memref, &dstLoadOpInsts);
 
       SmallVector<AffineForOp, 4> dstLoopIVs;
       getAffineForIVs(*dstLoadOpInsts[0], &dstLoopIVs);
       unsigned dstLoopDepthTest = dstLoopIVs.size();
while (findSiblingNodeToFuse(dstNode, &visitedSibNodeIds, &idAndMemref)) {
       }
 
       // Skip if fusion is not feasible at any loop depths.
       if (maxLegalFusionDepth == 0)
         continue;
 
       unsigned bestDstLoopDepth = maxLegalFusionDepth;
       if (!maximalFusion) {
-        // Check if fusion would be profitable.
-        if (!isFusionProfitable(sibLoadOpInst, sibStoreOpInst, dstAffineForOp,
+        // Check if fusion would be profitable. For sibling fusion, the sibling
+        // load op is treated as the src "store" op for fusion profitability
+        // purposes. The footprint of the load in the slice relative to the
+        // unfused source's determines reuse.
+        if (!isFusionProfitable(sibLoadOpInst, sibLoadOpInst, dstAffineForOp,
                                 depthSliceUnions, maxLegalFusionDepth,
                                 &bestDstLoopDepth, computeToleranceThreshold))
           continue;
       }
 
       assert(bestDstLoopDepth > 0 && "Unexpected loop fusion depth");
       assert(!depthSliceUnions[bestDstLoopDepth - 1].isEmpty() &&
              "Fusion depth has no computed slice union");
auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
       // which it also has an input dependence edge.
       DenseSet<Value> loadAndStoreMemrefSet;
       sibNode->getLoadAndStoreMemrefSet(&loadAndStoreMemrefSet);
       if (llvm::any_of(loadAndStoreMemrefSet, [=](Value memref) {
             return mdg->getIncomingMemRefAccesses(sibNode->id, memref) > 0;
           }))
         return false;
 
-      // Check that all stores are to the same memref.
+      // Check that all stores are to the same memref if any.
       DenseSet<Value> storeMemrefs;
       for (auto *storeOpInst : sibNode->stores) {
         storeMemrefs.insert(
             cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
       }
-      if (storeMemrefs.size() != 1)
+      if (storeMemrefs.size() > 1)
         return false;
 
       // Skip if a memref value in one node is used by a non-affine memref
       // access that lies between 'dstNode' and 'sibNode'.
       if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
           hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
         return false;
       return true;

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

 // MAXIMAL-DAG: [[$MAP7:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 // MAXIMAL-DAG: [[$MAP8:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 * 16 - d1 + 15)>
 // MAXIMAL-LABEL: func @fuse_across_varying_dims_complex
 // MAXIMAL-NEXT:  memref.alloc() : memref<64x1xf32>
 // MAXIMAL-NEXT:  arith.constant 0 : index
 // MAXIMAL-NEXT:  memref.alloc() : memref<2x2x3x3x16x1xf32>
 // MAXIMAL-NEXT:  memref.alloc() : memref<144x4xf32>
 // MAXIMAL-NEXT:  affine.for %{{.*}} = 0 to 9 {
-// MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 9 {
-// MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 4 {
-// MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 16 {
-// MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 64 {
-// MAXIMAL-NEXT:            affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:            affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:            affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:            affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:            affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:            affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x1xf32>
-// MAXIMAL-NEXT:            affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
-// MAXIMAL-NEXT:          }
-// MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 4 {
-// MAXIMAL-NEXT:            affine.for %{{.*}} = 0 to 16 {
-// MAXIMAL-NEXT:              affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:              affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
-// MAXIMAL-NEXT:            }
+// MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 4 {
+// MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 16 {
+// MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 64 {
+// MAXIMAL-NEXT:          affine.apply [[$MAP0]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply [[$MAP1]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply [[$MAP2]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply [[$MAP3]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply [[$MAP4]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.load %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] : memref<2x2x3x3x16x1xf32>
+// MAXIMAL-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}}, 0] : memref<64x1xf32>
+// MAXIMAL-NEXT:        }
+// MAXIMAL-NEXT:        affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:        affine.load %{{.*}}[%{{.*}} * 16 + %{{.*}}, 0] : memref<64x1xf32>
+// MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 9 {
+// MAXIMAL-NEXT:          affine.for %{{.*}} = 0 to 4 {
 // MAXIMAL-NEXT:            affine.for %{{.*}} = 0 to 16 {
-// MAXIMAL-NEXT:              affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:              affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<144x4xf32>
-// MAXIMAL-NEXT:            }
-// MAXIMAL-NEXT:          }
-// MAXIMAL-NEXT:          affine.apply [[$MAP8]](%{{.*}}, %{{.*}})
-// MAXIMAL-NEXT:          affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, 0] : memref<64x1xf32>
-// MAXIMAL-NEXT:        }
-// MAXIMAL-NEXT:      }
-// MAXIMAL-NEXT:    }
-// MAXIMAL-NEXT:  }
+// MAXIMAL-NEXT:              affine.apply [[$MAP8]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:              affine.load %{{.*}}[%{{.*}} * 16 - %{{.*}} + 15, 0] : memref<64x1xf32>
+// MAXIMAL-NEXT:            }
+// MAXIMAL-NEXT:          }
+// MAXIMAL-NEXT:        }
+// MAXIMAL-NEXT:      }
+// MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 16 {
+// MAXIMAL-NEXT:        affine.apply [[$MAP7]](%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:        affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<144x4xf32>
+// MAXIMAL-NEXT:      }
+// MAXIMAL-NEXT:    }
+// MAXIMAL-NEXT:  }
 
 // -----
 
 func.func @should_fuse_with_slice_union() {
   %a = memref.alloc() : memref<100xf32>
   %c0 = arith.constant 0 : index
   %cf0 = arith.constant 0.0 : f32
 

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

 // SIBLING-MAXIMAL-NEXT:        %[[cst_1:.*]] = arith.constant 1.000000e+00 : f32
 // SIBLING-MAXIMAL-NEXT:           affine.for %[[idx_0:.*]]= 0 to 1 {
 // SIBLING-MAXIMAL-NEXT:             affine.for %[[idx_1:.*]] = 0 to 64 {
 // SIBLING-MAXIMAL-NEXT:               %[[result_1:.*]] = affine.for %[[idx_2:.*]] = 0 to 32 iter_args(%[[iter_0:.*]] = %[[cst_1]]) -> (f32) {
 // SIBLING-MAXIMAL-NEXT:                 %[[result_0:.*]] = affine.for %[[idx_3:.*]] = 0 to 64 iter_args(%[[iter_1:.*]] = %[[cst_0]]) -> (f32) {
 
 // -----
 
+// SIBLING-MAXIMAL-LABEL: func @sibling_load_only
+func.func @sibling_load_only(%arg0: memref<10xf32>) {
+  affine.for %arg1 = 0 to 10 {
+    %0 = affine.load %arg0[%arg1] : memref<10xf32>
+  }
+  affine.for %arg1 = 0 to 10 {
+    %0 = affine.load %arg0[%arg1] : memref<10xf32>
+  }
+  // SIBLING-MAXIMAL-NEXT: affine.for
+  // SIBLING-MAXIMAL-NEXT:   affine.load
+  // SIBLING-MAXIMAL-NEXT:   affine.load
+  return
+}
+
+// -----
+
 // PRODUCER-CONSUMER-LABEL: func @fusion_for_multiple_blocks() {
 func.func @fusion_for_multiple_blocks() {
 ^bb0:
   %m = memref.alloc() : memref<10xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 10 {
     affine.store %cf7, %m[%i0] : memref<10xf32>

mlir/test/Transforms/loop-fusion.mlir

func.func @should_fuse_at_depth1_with_trip_count_19() {
   // CHECK-NEXT:  }
   // CHECK-NEXT:  return
   return
 }
 
 
 // -----
 
-// CHECK-LABEL: func @should_fuse_with_private_memrefs_with_diff_shapes() {
-func.func @should_fuse_with_private_memrefs_with_diff_shapes() {
+// CHECK-LABEL: func @should_fuse_with_private_memref() {
+func.func @should_fuse_with_private_memref() {
   %m = memref.alloc() : memref<100xf32>
   %cf7 = arith.constant 7.0 : f32
 
   affine.for %i0 = 0 to 100 {
     affine.store %cf7, %m[%i0] : memref<100xf32>
   }
   affine.for %i1 = 0 to 17 {
     %v0 = affine.load %m[%i1] : memref<100xf32>
   }
   affine.for %i2 = 0 to 82 {
     %v1 = affine.load %m[%i2] : memref<100xf32>
   }
-  // Should create two new private memrefs customized to the shapes accessed
-  // by loops %{{.*}} and %{{.*}}.
-  // CHECK-DAG:  memref.alloc() : memref<1xf32>
+  // Should create a new private memref.
   // CHECK-DAG:  memref.alloc() : memref<1xf32>
   // CHECK:      affine.for %{{.*}} = 0 to 17 {
   // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
-  // CHECK-NEXT: }
-  // CHECK-NEXT: affine.for %{{.*}} = 0 to 82 {
-  // CHECK-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // CHECK-NEXT: }
   // CHECK-NEXT: return
   return
 }
 
 // -----