[Mlir-commits] [mlir] 98daa4e - [MLIR] Fix incorrect removal of source loop in loop fusion

[llvmlistbot at llvm.org]

Author: Groverkss
Date: 2021-11-23T02:54:09+05:30
New Revision: 98daa4e425b02a0c27d53992e80510bda65809c5

URL: https://github.com/llvm/llvm-project/commit/98daa4e425b02a0c27d53992e80510bda65809c5
DIFF: https://github.com/llvm/llvm-project/commit/98daa4e425b02a0c27d53992e80510bda65809c5.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Analysis/Utils.cpp
    mlir/test/Transforms/loop-fusion-4.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Utils.cpp b/mlir/lib/Analysis/Utils.cpp
index 9565441863018..17c499a0453bf 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 6c5119d18f950..de3894d74801b 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 
diff erence 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