[MLIR] [affine-loop-fusion] Fix a bug about non-result ops in affine-loop-fusion

[tungld]

Summary

This patch fixes the following bug when calling --affine-loop-fusion

Input program:

mlir
func @should_not_fuse_since_top_level_non_affine_non_result_users(
    %in0 : memref<32xf32>, %in1 : memref<32xf32>) {
  %c0 = constant 0 : index
  %cst_0 = constant 0.000000e+00 : f32

  affine.for %d = 0 to 32 {
    %lhs = affine.load %in0[%d] : memref<32xf32>
    %rhs = affine.load %in1[%d] : memref<32xf32>
    %add = addf %lhs, %rhs : f32
    affine.store %add, %in0[%d] : memref<32xf32>
  }
  store %cst_0, %in0[%c0] : memref<32xf32>
  affine.for %d = 0 to 32 {
    %lhs = affine.load %in0[%d] : memref<32xf32>
    %rhs = affine.load %in1[%d] : memref<32xf32>
    %add = addf %lhs, %rhs: f32
    affine.store %add, %in0[%d] : memref<32xf32>
  }
  return
}

call --affine-loop-fusion, we got an incorrect output:

mlir
func @should_not_fuse_since_top_level_non_affine_non_result_users(%arg0: memref<32xf32>, %arg1: memref<32xf32>) {
  %c0 = constant 0 : index
  %cst = constant 0.000000e+00 : f32
  store %cst, %arg0[%c0] : memref<32xf32>
  affine.for %arg2 = 0 to 32 {
    %0 = affine.load %arg0[%arg2] : memref<32xf32>
    %1 = affine.load %arg1[%arg2] : memref<32xf32>
    %2 = addf %0, %1 : f32
    affine.store %2, %arg0[%arg2] : memref<32xf32>
    %3 = affine.load %arg0[%arg2] : memref<32xf32>
    %4 = affine.load %arg1[%arg2] : memref<32xf32>
    %5 = addf %3, %4 : f32
    affine.store %5, %arg0[%arg2] : memref<32xf32>
  }
  return
}

This happened because when analyzing the source and destination nodes,
affine loop fusion ignored non-result ops sandwitched between them. In
other words, the MemRefDependencyGraph in the affine loop fusion ignored
these non-result ops.

This patch solves the issue by adding these non-result ops to the
MemRefDependencyGraph.

[bondhugula]

Thanks for noticing this issue. A couple of comments.

mlir/lib/Transforms/LoopFusion.cpp

772

"which could be memory side-effect" -> "which could have a memory side effect"

773

You should really be checking for the Write MemoryEffect via the MemoryEffectOpInterface. Using the interface, you can even get hold of the actual memref Value. Otherwise, this is just overly conservative. If an op doesn't implement the MemoryEffectOpInterface, we have the guarantee to assume it's not writing to a memref.

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

773

-> which could have a memory write side effect

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

771

There is nothing special about zero result ops for what you are addressing. An op with results can also have memory write side effects. The case right above op.getNumResults() > 0 && !op.use_empty() will miss it if the results don't have uses, i.e., you could have an op with memory write side effects, one or more results, and with those results being unused: such an op can't be ignored.

What you need is:

else if (auto memOp = dyn_cast<MemoryEffectOpInterface>(op)) {
   // Get hold of the effects; check for write effect.
   ...
}

[tungld]

Use MemWrite side effect to precisely capture the situation.

[tungld]

@bondhugula Thank you for your comments! I updated the patch.

[tungld]

Readd the lit test into the new revision

[tungld]

Change the lit test name

[tungld]

@bondhugula any additional modification needed? thanks!

[bondhugula]

In D95668#2538220, @tungld wrote:

@bondhugula any additional modification needed? thanks!

Can you please mark the comments that are done as "Done"?

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

776–780

Btw, for a write effect, doing a getValue() on it also gives you the operand Value that is actually written to. Would this be useful in any way? The previous case was about things that generate SSA values and so such SSA edges are accounted for. But I can't recall what the impact of just adding a node like this would be. Would the memref being written to get looked at among the operands?

[tungld]

mlir/lib/Transforms/LoopFusion.cpp

776–780

By adding this node, it is used when checking whether a src and dest loop can be fused or not (i.e. inside canFuseSrcWhichWritesToLiveOut), which calls hasNonAffineUsersOnThePath to check non-affine users.

[bondhugula]

LGTM

[bondhugula]

Thanks for improving this. Let me know if you want me to commit this for you.

[tungld]

Thank you! Please help me to land it. My "name (email)" : "Tung D. Le ([e-mail])".

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 (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());
+          })) {
+        Node node(nextNodeId++, &op);
+        nodes.insert({node.id, node});
+      }
     }
   }
 
   for (auto &idAndNode : nodes) {
     LLVM_DEBUG(llvm::dbgs() << "Create node " << idAndNode.first << " for:\n"
                             << *(idAndNode.second.op) << "\n");
     (void)idAndNode;
   }

mlir/test/Transforms/loop-fusion.mlir

 // CHECK:  affine.for
 // CHECK:    addf
 // CHECK:  affine.for
 // CHECK:    mulf
 // CHECK:    subf
 
 // -----
 
+// CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_mem_write_users
+func @should_not_fuse_since_top_level_non_affine_mem_write_users(
+    %in0 : memref<32xf32>, %in1 : memref<32xf32>) {
+  %c0 = constant 0 : index
+  %cst_0 = constant 0.000000e+00 : f32
+
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = addf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  store %cst_0, %in0[%c0] : memref<32xf32>
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = addf %lhs, %rhs: f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  return
+}
+
+// CHECK:  affine.for
+// CHECK:    addf
+// CHECK:  affine.for
+// CHECK:    addf
+
+// -----
+
 // MAXIMAL-LABEL: func @fuse_minor_affine_map
 func @fuse_minor_affine_map(%in: memref<128xf32>, %out: memref<20x512xf32>) {
   %tmp = alloc() : memref<128xf32>
 
   affine.for %arg4 = 0 to 128 {
     %ld = affine.load %in[%arg4] : memref<128xf32>
     affine.store %ld, %tmp[%arg4] : memref<128xf32>
   }