[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

LogicalResult res =
                                /*symbolOperands=*/{},
                                /*domInstFilter=*/&*forOp.getBody()->begin());
   assert(succeeded(res) &&
          "replaceAllMemrefUsesWith should always succeed here");
   (void)res;
   return newMemRef;
 }
 
+/// Walking from node 'srcId' to node 'dstId' (exclusive of 'srcId' and
+/// 'dstId'), if there is any non-affine operation accessing 'memref', return
+/// false. Otherwise, return true.
+static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
+                                       Value memref,
+                                       MemRefDependenceGraph *mdg) {
+  auto *srcNode = mdg->getNode(srcId);
+  auto *dstNode = mdg->getNode(dstId);
+  Value::user_range users = memref.getUsers();
+  // For each MemRefDependenceGraph's node that is between 'srcNode' and
+  // 'dstNode' (exclusive of 'srcNodes' and 'dstNode'), check whether any
+  // non-affine operation in the node accesses the 'memref'.
+  for (auto &idAndNode : mdg->nodes) {
+    Operation *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.
+      // Interrupt the walk if found.
+      auto walkResult = op->walk([&](Operation *user) {
+        // Skip affine ops.
+        if (isMemRefDereferencingOp(*user))
+          return WalkResult::advance();
+        // Find a non-affine op that uses the memref.
+        if (llvm::is_contained(users, user))
+          return WalkResult::interrupt();
+        return WalkResult::advance();
+      });
+      if (walkResult.wasInterrupted())
+        return true;
+    }
+  }
+  return false;
+}
+
+/// Check whether a memref value in node 'srcId' has a non-affine that
+/// is between node 'srcId' and node 'dstId' (exclusive of 'srcNode' and
+/// 'dstNode').
+static bool hasNonAffineUsersOnThePath(unsigned srcId, unsigned dstId,
+                                       MemRefDependenceGraph *mdg) {
+  // Collect memref values in node 'srcId'.
+  auto *srcNode = mdg->getNode(srcId);
+  llvm::SmallDenseSet<Value, 2> memRefValues;
+  srcNode->op->walk([&](Operation *op) {
+    // Skip affine ops.
+    if (isa<AffineForOp>(op))
+      return WalkResult::advance();
+    for (Value v : op->getOperands())
+      // Collect memref values only.
+      if (v.getType().isa<MemRefType>())
+        memRefValues.insert(v);
+    return WalkResult::advance();
+  });
+  // Looking for users between node 'srcId' and node 'dstId'.
+  for (Value memref : memRefValues)
+    if (hasNonAffineUsersOnThePath(srcId, dstId, memref, 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.
 // Returns true if 'dstNode's read/write region to 'memref' is a super set of
 // 'srcNode's write region to 'memref' and 'srcId' has only one output edge.
 // TODO(andydavis) Generalize this to handle more live in/out cases.
 static bool
 canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
                                AffineWriteOpInterface srcLiveOutStoreOp,
                                MemRefDependenceGraph *mdg) {
canFuseSrcWhichWritesToLiveOut(unsigned srcId, unsigned dstId,
   if (!dstNumElements.hasValue())
     return false;
 
   // Return false if write region is not a superset of 'srcNodes' write
   // region to 'memref'.
   // 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;
 }
 
 // Checks the profitability of fusing a backwards slice of the loop nest
 // surrounding 'srcOpInst' into the loop nest surrounding 'dstLoadOpInsts'.
 // The argument 'srcStoreOpInst' is used to calculate the storage reduction on
 // the memref being produced and consumed, which is an input to the cost model.
 // For producer-consumer fusion, 'srcStoreOpInst' will be the same as
auto canFuseWithSibNode = [&](Node *sibNode, Value memref) {
       // Check that all stores are to the same memref.
       DenseSet<Value> storeMemrefs;
       for (auto *storeOpInst : sibNode->stores) {
         storeMemrefs.insert(
             cast<AffineWriteOpInterface>(storeOpInst).getMemRef());
       }
       if (storeMemrefs.size() != 1)
         return false;
+
+      // Skip if a memref value in one node is used by a non-affine memref
+      // access that lies between 'dstNode' and 'sibNode'.
+      if (hasNonAffineUsersOnThePath(dstNode->id, sibNode->id, mdg) ||
+          hasNonAffineUsersOnThePath(sibNode->id, dstNode->id, mdg))
+        return false;
       return true;
     };
 
     // Search for siblings which load the same memref function argument.
     auto fn = dstNode->op->getParentOfType<FuncOp>();
     for (unsigned i = 0, e = fn.getNumArguments(); i != e; ++i) {
       for (auto *user : fn.getArgument(i).getUsers()) {
         if (auto loadOp = dyn_cast<AffineReadOpInterface>(user)) {
           // Gather loops surrounding 'use'.
           SmallVector<AffineForOp, 4> loops;
           getLoopIVs(*user, &loops);
           // Skip 'use' if it is not within a loop nest.

mlir/test/Transforms/loop-fusion.mlir

 // CHECK-NEXT:    addf
 // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    affine.load %{{.*}}[0] : memref<1xf32>
 // CHECK-NEXT:    affine.load %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
 // CHECK-NEXT:    mulf
 // CHECK-NEXT:    affine.store %{{.*}}, %arg{{.*}}[%arg{{.*}}] : memref<?xf32>
 // CHECK-NEXT:  }
 // CHECK-NEXT:  return
+
+// -----
+
+// 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