[MLIR] Fix incorrect removal of source loop in loop fusion

[Groverkss]

Summary

This patch fixes a bug in loop fusion pass where the source loop was removed
even when the fused loop did not cover all iterations of the source loop.

This was because the fast hueristic check for checking if source loop and
fused loop have same iterations did not take into account steps in loop.

[dcaballe]

Thanks a lot for the fix! LGTM. Just a minor comment.

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

165

Could you please check a bit more operations to make sure that fusion is actually happening in the second loop and that the first loop is the original one?

[bondhugula]

LGTM - just one minor NFC comment to add to @dcaballe's comments above.

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

149

Can you add this test case slightly above where it's more appropriate (related to surrounding test cases) instead of at the end? Always appending to the file in general leads to more conflicts (downstream or otherwise) and spreads around various things arbitrarily in the file.

[Groverkss]

• Added more checks for test case and moved it to up to related tests

[Groverkss]

@dcaballe Is the test case fine now?

[dcaballe]

Thanks for addressing the feedback. LGTM. Just one minor comment. Feel free to fix it a commit it directly.

Thanks!
Diego

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

100

since you are capturing idx_1, probably good to match here. That or not capture it.

[Groverkss]

• Match captured index in test case.

mlir/lib/Analysis/Utils.cpp

for (unsigned i = 0, end = lbs.size(); i < end; ++i) {
     if (!srcLbResult.isa<AffineConstantExpr>() ||
         !srcUbResult.isa<AffineConstantExpr>() ||
         !dstLbResult.isa<AffineConstantExpr>() ||
         !dstUbResult.isa<AffineConstantExpr>())
       return llvm::None;
 
     // Check if src and dst loop bounds are the same. If not, we can guarantee
     // that the slice is not maximal.
-    if (srcLbResult != dstLbResult || srcUbResult != dstUbResult)
+    if (srcLbResult != dstLbResult || srcUbResult != dstUbResult ||
+        srcLoop.getStep() != dstLoop.getStep())
       return false;
   }
 
   return true;
 }
 
 /// Returns true if it is deterministically verified that the original iteration
 /// space of the slice is contained within the new iteration space that is

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

 
 // PRODUCER-CONSUMER:        affine.for
 // PRODUCER-CONSUMER-NEXT:     affine.for
 // PRODUCER-CONSUMER-NOT:    affine.for
 // PRODUCER-CONSUMER: return
 
 // -----
 
+// Expects fusion of producer into consumer at depth 1 and source loop to not
+// be removed due to difference in loop steps.
+// PRODUCER-CONSUMER-LABEL: func @check_src_dst_step
+func @check_src_dst_step(%m : memref<100xf32>,
+                         %src: memref<100xf32>,
+                         %out: memref<100xf32>) {
+  affine.for %i0 = 0 to 100 {
+    %r1 = affine.load %src[%i0]: memref<100xf32>
+    affine.store %r1, %m[%i0] : memref<100xf32>
+  }
+  affine.for %i2 = 0 to 100 step 2 {
+    %r2 = affine.load %m[%i2] : memref<100xf32>
+    affine.store %r2, %out[%i2] : memref<100xf32>
+  }
+  return
+}
+
+// Check if the fusion did take place as well as that the source loop was
+// not removed. To check if fusion took place, the read instruction from the
+// original source loop is checked to be in the fused loop.
+//
+// PRODUCER-CONSUMER:        affine.for %[[idx_0:.*]] = 0 to 100 {
+// PRODUCER-CONSUMER-NEXT:     %[[result_0:.*]] = affine.load %[[arr1:.*]][%[[idx_0]]] : memref<100xf32>
+// PRODUCER-CONSUMER-NEXT:     affine.store %[[result_0]], %{{.*}}[%[[idx_0]]] : memref<100xf32>
+// PRODUCER-CONSUMER-NEXT:   }
+// PRODUCER-CONSUMER:        affine.for %[[idx_1:.*]] = 0 to 100 step 2 {
+// PRODUCER-CONSUMER:          affine.load %[[arr1]][%[[idx_1]]] : memref<100xf32>
+// PRODUCER-CONSUMER:        }
+// PRODUCER-CONSUMER:        return
+
+// -----
+
 // SIBLING-MAXIMAL-LABEL:   func @reduce_add_non_maximal_f32_f32(
 func @reduce_add_non_maximal_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1 : memref<1x64xf32, 1>, %arg2 : memref<1x64xf32, 1>) {
     %cst_0 = arith.constant 0.000000e+00 : f32
     %cst_1 = arith.constant 1.000000e+00 : f32
     affine.for %arg3 = 0 to 1 {
       affine.for %arg4 = 0 to 64 {
         %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
           %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>

 // Test checks the loop structure is preserved after sibling fusion
 // since the destination loop and source loop trip counts do not
 // match.
 // SIBLING-MAXIMAL:        %[[cst_0:.*]] = arith.constant 0.000000e+00 : f32
 // 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) {