[MLIR][Affine-loop-fusion] Fix a bug in affine-loop-fusion pass when there are non-affine operations

[tungld]

Summary

When there is a mix of affine load/store and non-affine operations (e.g. std.load, std.store),
affine-loop-fusion ignores the present of non-affine ops, thus changing the program semantics.

E.g. we have a program of three affine loops operating on the same memref in which one of them uses std.load and std.store, as follows.

affine.for
  affine.store %1
affine.for
  std.load %1
  std.store %1
affine.for
  affine.load %1
  affine.store %1

affine-loop-fusion will produce the following result which changed the program semantics:

affine.for
  std.load %1
  std.store %1
affine.for
  affine.store %1
  affine.load %1
  affine.store %1

This patch is to fix the above problem by checking non-affine users of the memref that are between the source and destination nodes of interest.

[tungld]

Add missing changes for mlir/lib/Transforms/LoopFusion.cpp

[bondhugula]

Thanks very much for addressing this issue. Some initial comments.

mlir/lib/Transforms/LoopFusion.cpp

961

Change to /// comments please.

966–969

Avoid auto here for readability please.

969

Comment here please on what this iteration is doing.

974

You don't need to use the found flag. Instead use .wasInterrupted() on walk to check if the walk was interrupted.

993

/// comments please.

994

Nit: mention whether between is inclusive or exclusive of srcId and dstId

[bondhugula]

mlir/lib/Transforms/LoopFusion.cpp

1001

Comment fix.

1004

Aren't you only interested in memref values here?

1004–1006

Use a SmallDenseSet? This will be quite expensive when nodes become larger.

1011

Instead of iterating over all values - iterate over only the memref values?

1870

non-affine memref access instead of non-affine operation?

[tungld]

Update the code according to Uday Bondhugula's comments.

[tungld]

@bondhugula Thank you for your comments! I updated the code based on your comments.

[tungld]

Add unit tests.

[andydavis1]

Thanks for addressing this issue! Let's see if Uday has any more comments before committing, but your latest changes look good to me...

[bondhugula]

Looks good, thanks. A few more comments to address.

mlir/lib/Transforms/LoopFusion.cpp

969

Avoid auto here please.

984

You can instead use: llvm::is_contained(users, user) and also drop the trivial braces around.

[tungld]

Thanks Uday! I updated the code based on your comments.

[bondhugula]

Thanks, please do commit. Let me know if you need me to.

[tungld]

@bondhugula it seems I don't have permission to push a commit. Could you please help me to land this patch? Many thanks!

[bondhugula]

In D82158#2110621, @tungld wrote:

@bondhugula it seems I don't have permission to push a commit. Could you please help me to land this patch? Many thanks!

Sure, will commit in a few minutes.

[bondhugula]

In D82158#2111900, @bondhugula wrote:

In D82158#2110621, @tungld wrote:

@bondhugula it seems I don't have permission to push a commit. Could you please help me to land this patch? Many thanks!

Sure, will commit in a few minutes.

Could you let me know what "name <email>" info you want on the commit?

[tungld]

Could you let me know what "name <email>" info you want on the commit?

Please use "Tung D. Le <[e-mail]>". Thanks!

mlir/lib/Transforms/LoopFusion.cpp

   return newMemRef;
 }
 
+// Walking from node 'srcId' to node 'dstId', if there is a use of 'memref' that
+// is a non-affine operation, return false.
+static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
+                                       Value memref,
+                                       MemRefDependenceGraph *mdg) {
+  auto *srcNode = mdg->getNode(srcId);
+  auto *dstNode = mdg->getNode(dstId);
+  auto users = memref.getUsers();
+  for (auto &idAndNode : mdg->nodes) {
+    auto *op = idAndNode.second.op;
+    // Take care of operations between srcNode and dstNode.
+    if (srcNode->op->isBeforeInBlock(op) && op->isBeforeInBlock(dstNode->op)) {
+      // Walk inside the operation to find any use of the memref.
+      bool found = false;
+      op->walk([&](Operation *user) {
+        // Skip affine ops.
+        if (isMemRefDereferencingOp(*user))
+          return WalkResult::advance();
+        // Find a non-affine op that uses the memref.
+        if (std::find(users.begin(), users.end(), user) != users.end()) {
+          found = true;
+          return WalkResult::interrupt();
+        }
+        return WalkResult::advance();
+      });
+      if (found)
+        return true;
+    }
+  }
+  return false;
+}
+
+// Check whether a value in node 'srcId' has a non-affine consumer that is
+// between node 'srcId' and node 'dstId'.
+static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
+                                       MemRefDependenceGraph *mdg) {
+  // Collect values in node 'srcId'.
+  auto *srcNode = mdg->getNode(srcId);
+  SmallVector<Value, 2> values;
+  srcNode->op->walk([&](Operation *op) {
+    // Skip affine ops
+    if (isa<AffineForOp>(op))
+      return WalkResult::advance();
+    for (Value v : op->getOperands()) {
+      if (std::find(values.begin(), values.end(), v) == values.end())
+        values.push_back(v);
+    }
+    return WalkResult::advance();
+  });
+  // Looking for users between node 'srcId' and node 'dstId'.
+  for (Value v : values) {
+    if (hasNonAffineUsersOnThePath(srcId, dstId, v, mdg))
+      return true;
+  }
+  return false;
+}
+
 // Checks if node 'srcId' can be safely fused into node 'dstId'. Node 'srcId'
 // may write to multiple memrefs but it is required that only one of them,
 // 'srcLiveOutStoreOp', has output edges.

   // TODO(andydavis) Check the shape and lower bounds here too.
   if (srcNumElements != dstNumElements)
     return false;
+
+  // Return false if 'memref' is used by a non-affine operation that is
+  // between node 'srcId' and node 'dstId'.
+  if (hasNonAffineUsersOnThePath(srcId, dstId, mdg))
+    return false;
+
   return true;
 }
 

       }
       if (storeMemrefs.size() != 1)
         return false;
+
+      // Skip if a value in one node is used by a non-affine operation that lies
+      // between 'dstNode' and 'sibNode'.
+      if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
+          hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
+        return false;
       return true;
     };
 

mlir/test/Transforms/loop-fusion.mlir

 // CHECK:        mulf
 // CHECK-NOT:  affine.for
 // CHECK:        divf
+
+// -----
+
+// CHECK-LABEL: func @should_not_fuse_since_non_affine_users
+func @should_not_fuse_since_non_affine_users(%in0 : memref<32xf32>,
+                      %in1 : 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>
+  }
+  affine.for %d = 0 to 32 {
+    %lhs = load %in0[%d] : memref<32xf32>
+    %rhs = load %in1[%d] : memref<32xf32>
+    %add = subf %lhs, %rhs : f32
+    store %add, %in0[%d] : memref<32xf32>
+  }
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = mulf %lhs, %rhs : f32
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  return
+}
+
+// CHECK:  affine.for
+// CHECK:    addf
+// CHECK:  affine.for
+// CHECK:    subf
+// CHECK:  affine.for
+// CHECK:    mulf
+
+// -----
+
+// CHECK-LABEL: func @should_not_fuse_since_top_level_non_affine_users
+func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,
+                      %in1 : memref<32xf32>) {
+  %sum = alloc() : memref<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
+    store %add, %sum[] : memref<f32>
+    affine.store %add, %in0[%d] : memref<32xf32>
+  }
+  %load_sum = load %sum[] : memref<f32>
+  affine.for %d = 0 to 32 {
+    %lhs = affine.load %in0[%d] : memref<32xf32>
+    %rhs = affine.load %in1[%d] : memref<32xf32>
+    %add = mulf %lhs, %rhs : f32
+    %sub = subf %add, %load_sum: f32
+    affine.store %sub, %in0[%d] : memref<32xf32>
+  }
+  dealloc %sum : memref<f32>
+  return
+}
+
+// CHECK:  affine.for
+// CHECK:    addf
+// CHECK:  affine.for
+// CHECK:    mulf
+// CHECK:    subf