[MLIR][affine] Certain Call Ops to prevent fusion

[vinayaka-polymage]

Summary

Fixes a bug in affine fusion pipeline where an incorrect fusion is performed
despite a Call Op that potentially modifies memrefs under consideration
exists between source and target.

[dcaballe]

It looks great! Thank you so much!

[dcaballe]

mlir/lib/Transforms/LoopFusion.cpp

774

Giving this another thought, should we be more conservative and add a node for any operation that uses a memref, even if it's unknown? In other words, we could remove this check and add a check like:

else if (llvm::any_of(op.getOperands(), [&](Value v) {
            return v.getType().isa<MemRefType>();
          })) {
                ...

after the MemoryEffectOpInterface check below. It would catch any CallOpInterface op and any other unknown operation using a memref.

[vinayaka-polymage]

mlir/lib/Transforms/LoopFusion.cpp

774

That sounds good. We might not even need checking MemoryEffectOpInterface then?

[dcaballe]

mlir/lib/Transforms/LoopFusion.cpp

774

That sounds good. We might not even need checking MemoryEffectOpInterface then?

That was my impression when looking at your first patch. Since this is a correctness issue, I think it's important to be overconservative and then add exceptions as we see fit.

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

774

I don't think that would be the right approach. For eg., we can't be assuming that side-effect free ops using memrefs are writing to memrefs! For eg. if you just assumed any op could, you'd run into cases where things like dim op would start blocking fusion. It's not just overconservative but a limitation and not the right thing. If a registered op is really writing to memrefs or memory, it is wrong for the op to not declare the appropriate memory effects. If it doesn't declare those, it's not the transformation's responsibility to still work correctly. You can of course be conservative for unregistered ops perhaps. I believe a number of transformations/utilities make use of the effects interface and they rely on the ops correctly setting those effects. One of the reasons the effects interfaces exist is to not have to whitelist/hardcode a specific set of ops or not be forced to be overconservative.

[dcaballe]

mlir/lib/Transforms/LoopFusion.cpp

774

Ok, that sounds reasonable. I was more worried about ops with arbitrary side effects. I think unregistered ops would fall into that category, right? Anyways, we could add support for those cases in a separate patch. This patch should be ready to go. Thanks!

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

774

That's right - I think we can check for unregistered ops and assume the most conservative case for it. For registered ops, I think we are fine with memory effects.

[vinayaka-polymage]

Thanks a lot for the insights @bondhugula and @dcaballe ! No more changes needed for this patch then.

[bondhugula]

One minor comment.

mlir/lib/Transforms/LoopFusion.cpp

769

You could just use getOperandTypes() here.

[vinayaka-polymage]

Addressing the getOperandTypes comment

[vinayaka-polymage]

mlir/lib/Transforms/LoopFusion.cpp

769

Thanks for the tip! I have made the change.

mlir/lib/Transforms/LoopFusion.cpp

for (auto &op : f.front()) {
     } else if (op.getNumRegions() != 0) {
       // Return false if another region is found (not currently supported).
       return false;
     } else if (op.getNumResults() > 0 && !op.use_empty()) {
       // Create graph node for top-level producer of SSA values, which
       // could be used by loop nest nodes.
       Node node(nextNodeId++, &op);
       nodes.insert({node.id, node});
+    } else if (isa<CallOpInterface>(op)) {
+      // Create graph node for top-level Call Op that takes any argument of
+      // memref type. Call Op that returns one or more memref type results
+      // is already taken care of, by the previous conditions.
+      if (llvm::any_of(op.getOperands(), [&](Value v) {
+            return v.getType().isa<MemRefType>();
+          })) {
+        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, MemoryEffects::Free>(
                 it.getEffect());

mlir/test/Transforms/loop-fusion.mlir

func @should_fuse_multi_store_producer_with_scaping_memrefs_and_preserve_src(
   // 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>

 // 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
+
+// -----
+
+func private @some_function(memref<16xf32>)
+func @call_op_prevents_fusion(%arg0: memref<16xf32>){
+  %A = 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>
+  }
+  call @some_function(%A) : (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>
+  }
+  return
+}
+// CHECK-LABEL: func @call_op_prevents_fusion
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      affine.store
+// CHECK:         call
+// CHECK:         affine.for
+// CHECK-NEXT:      affine.load
+// CHECK-NEXT:      addf
+// CHECK-NEXT:      affine.store
+
+// -----
+
+func private @some_function()
+func @call_op_does_not_prevent_fusion(%arg0: memref<16xf32>){
+  %A = 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>
+  }
+  call @some_function() : () -> ()
+  %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>
+  }
+  return
+}
+// CHECK-LABEL: func @call_op_does_not_prevent_fusion
+// CHECK:         affine.for
+// CHECK-NOT:     affine.for