[Mlir-commits] [mlir] 1533298 - [MLIR][affine] Prevent fusion when ops with memory effect free are present between producer and consumer

[Uday Bondhugula]

Author: Vinayaka Bandishti
Date: 2021-02-22T23:21:02+05:30
New Revision: 15332982c3d8d1c9e8bd7ae9f76f22f77adb51ee

URL: https://github.com/llvm/llvm-project/commit/15332982c3d8d1c9e8bd7ae9f76f22f77adb51ee
DIFF: https://github.com/llvm/llvm-project/commit/15332982c3d8d1c9e8bd7ae9f76f22f77adb51ee.diff

LOG: [MLIR][affine] Prevent fusion when ops with memory effect free are present between producer and consumer

This commit fixes a bug in affine fusion pipeline where an
incorrect fusion is performed despite a dealloc op is present
between a producer and a consumer. This is done by creating a
node for dealloc op in the MDG.

Reviewed By: bondhugula, dcaballe

Differential Revision: https://reviews.llvm.org/D97032

Added: 
    

Modified: 
    mlir/lib/Transforms/LoopFusion.cpp
    mlir/test/Transforms/loop-fusion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index e13616e1d71c..d6d26a85215b 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -768,7 +768,8 @@ bool MemRefDependenceGraph::init(FuncOp f) {
       SmallVector<MemoryEffects::EffectInstance, 1> effects;
       effectInterface.getEffects(effects);
       if (llvm::any_of(effects, [](const MemoryEffects::EffectInstance &it) {
-            return isa<MemoryEffects::Write>(it.getEffect());
+            return isa<MemoryEffects::Write, MemoryEffects::Free>(
+                it.getEffect());
           })) {
         Node node(nextNodeId++, &op);
         nodes.insert({node.id, node});

diff  --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index 0c20ea46ad5e..21584fa6f7f1 100644
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -2835,3 +2835,34 @@ func @should_fuse_multi_store_producer_with_scaping_memrefs_and_preserve_src(
 
   return
 }
+
+// -----
+func @should_not_fuse_due_to_dealloc(%arg0: memref<16xf32>){
+  %A = alloc() : memref<16xf32>
+  %C = alloc() : memref<16xf32>
+  %cst_1 = constant 1.000000e+00 : f32
+  affine.for %arg1 = 0 to 16 {
+    %a = affine.load %arg0[%arg1] : memref<16xf32>
+    affine.store %a, %A[%arg1] : memref<16xf32>
+    affine.store %a, %C[%arg1] : memref<16xf32>
+  }
+  dealloc %C : memref<16xf32>
+  %B = alloc() : memref<16xf32>
+  affine.for %arg1 = 0 to 16 {
+    %a = affine.load %A[%arg1] : memref<16xf32>
+    %b = addf %cst_1, %a : f32
+    affine.store %b, %B[%arg1] : memref<16xf32>
+  }
+  dealloc %A : memref<16xf32>
+  return
+}
+// CHECK-LABEL: func @should_not_fuse_due_to_dealloc
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      affine.store
+// CHECK-NEXT:      affine.store
+// CHECK:         dealloc
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      addf
+// CHECK-NEXT:      affine.store