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

[vinayaka-polymage]

Summary

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.

[dcaballe]

Hey Vinayaka!

Thank you so much for the fix! Some comments below.

mlir/lib/Transforms/LoopFusion.cpp

772

This looks good to me! I think we have to be more conservative and also add any non-affine operation that produces or uses a memref. Otherwise, hasNonAffineUsersOnThePath won't return the expected output if we have a read, call or alloc-like op or any unknown operation that could potentially make the memref escape, right?

mlir/test/Transforms/loop-fusion.mlir

2849

Would you mind adding also tests with a call op, a load op and an unknown op using the memref, if we don't have it already? Thanks!

2860

I think we should check loads and stores also because we may fuse and do not remove the producer loop.

[dcaballe]

https://reviews.llvm.org/D97030 is somehow related. I think the common solution is to represent these missing dependences in the MDG. We may have to coordinate the fixes.

[bondhugula]

LGTM.

mlir/lib/Transforms/LoopFusion.cpp

772

I feel those can be handled in a next patch. This doesn't fix everything but is catching writing and freeing users the right way.

mlir/test/Transforms/loop-fusion.mlir

2849

But this patch isn't meant to fix those. It's only for the users which have a write side-effect or a free side-effect.

[dcaballe]

If you could improve the lit test so that loads and stores are also checked, that should be all. Thanks!

mlir/lib/Transforms/LoopFusion.cpp

772

Ok, fair enough!

[vinayaka-polymage]

Test case fix.

[vinayaka-polymage]

mlir/lib/Transforms/LoopFusion.cpp

772

Thanks for the suggestions!
I will create a new patch with the above suggested changes and paste the link here.

mlir/test/Transforms/loop-fusion.mlir

2860

That's correct, I had missed this, thanks! I have corrected this.

[dcaballe]

Thanks for addressing the comments!

[vinayaka-polymage]

Suggested additional checks and test-cases submitted at https://reviews.llvm.org/D97252

mlir/lib/Transforms/LoopFusion.cpp

if (auto forOp = dyn_cast<AffineForOp>(op)) {
       Node node(nextNodeId++, &op);
       nodes.insert({node.id, node});
     } else if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
       // Create graph node for top-level op, which could have a memory write
       // side effect.
       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});
       }
     }
   }
 
   for (auto &idAndNode : nodes) {

mlir/test/Transforms/loop-fusion.mlir

func @should_fuse_multi_store_producer_with_scaping_memrefs_and_preserve_src(
   // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
   // CHECK-NEXT:  }
   // CHECK-NOT:   affine.for
 
   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