From 98daa4e425b02a0c27d53992e80510bda65809c5 Mon Sep 17 00:00:00 2001
From: Groverkss <groverkss@gmail.com>
Date: Tue, 23 Nov 2021 02:48:03 +0530
Subject: [PATCH] [MLIR] Fix incorrect removal of source loop in loop fusion

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.

Reviewed By: dcaballe, bondhugula

Differential Revision: https://reviews.llvm.org/D114164
---
 mlir/lib/Analysis/Utils.cpp             |  3 ++-
 mlir/test/Transforms/loop-fusion-4.mlir | 32 +++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index 956544186301..17c499a0453b 100644
--- a/mlir/lib/Analysis/Utils.cpp
+++ b/mlir/lib/Analysis/Utils.cpp
@@ -207,7 +207,8 @@ Optional<bool> ComputationSliceState::isSliceMaximalFastCheck() const {
 
     // 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;
   }
 
diff --git a/mlir/test/Transforms/loop-fusion-4.mlir b/mlir/test/Transforms/loop-fusion-4.mlir
index 6c5119d18f95..de3894d74801 100644
--- a/mlir/test/Transforms/loop-fusion-4.mlir
+++ b/mlir/test/Transforms/loop-fusion-4.mlir
@@ -71,6 +71,38 @@ func @unflatten2d_with_transpose(%arg1: memref<8x7xf32>) {
 
 // -----
 
+// 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