[MLIR][affine-loop-fusion] Handle defining ops between the source and dest loops

[tungld]

Summary

This patch handles defining ops between the source and dest loop nests, and prevents loop nests with iter_args from being fused.

If there is any SSA value in the dest loop nest whose defining op has dependence from the source loop nest, we cannot fuse the loop nests.

If there is a affine.for with iter_args, prevent it from being fused.

[tungld]

Remove a redundant method

[tungld]

Clean the code

[tungld]

Check causal dependence

[dcaballe]

Hey tungld,

Thanks for the patch! There is a different but somehow related issue that is being address here: https://reviews.llvm.org/D97032.
I would like to know what @bondhugula and @andydavis1 think but I think we might want to add the proper %0 -> %i1 dependence to the MDG,
instead of special-casing the load scenario. Even though it's not strictly a memref dependence, it's a memref-related dependence between two graph nodes.

Interestingly, we should also consider more generic cases. For example:

affine.for %i0 = 0 to 10 {
  affine.store %cst, %b[] : memref<f32>
  affine.store %cst, %a[%i0] : memref<10xf32>
}

%0 = "unknown_op" (%b) : (memref<f32>) -> f32

affine.for %i1 = 0 to 10 {
  %1 = affine.load %a[%i1] : memref<10xf32>
  %2 = divf %0, %1 : f32
}

affine.for %i0 = 0 to 10 {
  affine.store %cst, %b[] : memref<f32>
  affine.store %cst, %a[%i0] : memref<10xf32>
}

%0 = affine.for ...

affine.for %i1 = 0 to 10 {
  %1 = affine.load %a[%i1] : memref<10xf32>
  %2 = divf %0, %1 : f32
}

I think the former would be addressed by https://reviews.llvm.org/D97032 and the latter would need support for iter_args in the fusion algorithm
so I wouldn't try to handle them right now. However, we have to make sure that this solution paves the way to support the more complex cases.

Thanks,
Diego

[bondhugula]

In D97030#2575203, @dcaballe wrote:

Hey tungld,

Thanks for the patch! There is a different but somehow related issue that is being address here: https://reviews.llvm.org/D97032.
I would like to know what @bondhugula and @andydavis1 think but I think we might want to add the proper %0 -> %i1 dependence to the MDG,
instead of special-casing the load scenario. Even though it's not strictly a memref dependence, it's a memref-related dependence between two graph nodes.

Not really. This one is just a pure SSA dependence and nothing to do with memory dependences. The value may be loaded into from a memref but it's a value. The memref user op is a load which may participate in separate/other memref dependences.

Interestingly, we should also consider more generic cases. For example:

affine.for %i0 = 0 to 10 {
  affine.store %cst, %b[] : memref<f32>
  affine.store %cst, %a[%i0] : memref<10xf32>
}

%0 = "unknown_op" (%b) : (memref<f32>) -> f32

affine.for %i1 = 0 to 10 {
  %1 = affine.load %a[%i1] : memref<10xf32>
  %2 = divf %0, %1 : f32
}

affine.for %i0 = 0 to 10 {
  affine.store %cst, %b[] : memref<f32>
  affine.store %cst, %a[%i0] : memref<10xf32>
}

%0 = affine.for ...

affine.for %i1 = 0 to 10 {
  %1 = affine.load %a[%i1] : memref<10xf32>
  %2 = divf %0, %1 : f32
}

I think the former would be addressed by https://reviews.llvm.org/D97032 and the latter would need support for iter_args in the fusion algorithm

The author of D97032 is already taking care of all side-effecting intervening ops (like writes, frees, unknown side-effecting ops like call) in two revisions in a generic way using memory effects / side-effect interfaces; those are pretty simple patches. The first of those is D97032 that takes care of memref writes/frees. But this patch addresses a different issue. This is not a memory dependence but an SSA dependence like other SSA edges.

so I wouldn't try to handle them right now. However, we have to make sure that this solution paves the way to support the more complex cases.

Thanks,
Diego

[bondhugula]

Some comments on clarifying doc / code comments.

mlir/lib/Transforms/LoopFusion.cpp

789

This line always needed a comment. Stores anyway don't define SSA values.

1457–1458

Nit on terminology: please either use "... depends on ..." or "has a dependence from" / "has a dependence to".

1463–1464

in between the loops -> in between the loop nests ? It's otherwise ambiguous.

"having dependence on the source loop" -> I think this needs to be rephrased for accuracy as well. "... with a user in .."?

[tungld]

Reflect reviewers' comments

[tungld]

Hi @bondhugula, @dcaballe,

Thanks for your comments! I just updated the patch.

This patch is to add one additional node in to MDG, which is from a defining op of an SSA value to its users. So, I think it can handle this

affine.for %i0 = 0 to 10 {
  affine.store %cst, %b[] : memref<f32>
  affine.store %cst, %a[%i0] : memref<10xf32>
}

%0 = affine.for ...

affine.for %i1 = 0 to 10 {
  %1 = affine.load %a[%i1] : memref<10xf32>
  %2 = divf %0, %1 : f32
}

because %0 = affine.for defines an SSA value.

I tried to run with the following program:

affine.for %arg2 = 0 to 10 {
   affine.store %cst, %arg1[] : memref<f32>
   affine.store %cst, %arg0[%arg2] : memref<10xf32>
}
%0 = affine.for %arg2 = 0 to 10 step 2 iter_args(%arg3 = %cst_0) -> (f32) {
   %1 = affine.load %arg0[%arg2] : memref<10xf32>
   %2 = addf %arg3, %1 : f32
   affine.yield %2 : f32
}
affine.for %arg2 = 0 to 10 {
    %1 = affine.load %arg0[%arg2] : memref<10xf32>
    %2 = divf %0, %1 : f32
}

and got

%0 = affine.for %arg2 = 0 to 10 step 2 iter_args(%arg3 = %cst_0) -> (f32) {
  affine.store %cst, %arg1[] : memref<f32>
  affine.store %cst, %arg0[%arg2] : memref<10xf32>
  %1 = affine.load %arg0[%arg2] : memref<10xf32>
  %2 = addf %arg3, %1 : f32
  affine.yield %2 : f32
}
affine.for %arg2 = 0 to 10 {
   %1 = affine.load %arg0[%arg2] : memref<10xf32>
   %2 = divf %0, %1 : f32
}

It seemed it could identify the defining op %0 = affine.for..., but the output program was not semantically similar to the input program due to step in affine.for. I think this is somewhat related to what @dcaballe mentioned

so I wouldn't try to handle them right now. However, we have to make sure that this solution paves the way to support the more complex cases.

[dcaballe]

I would like to know what @bondhugula and @andydavis1 think but I think we might want to add the proper
%0 -> %i1 dependence to the MDG,

Actually, as you could infer from my answer, I completely misunderstood the implementation of the patch! As
@tungld mentioned, it seems we were already adding edges for some SSA dependences:

Add dependence edges between nodes which produce SSA values and their
users. Load ops can be considered as the ones producing SSA values.

This is what I was referring to. I was not aware of it :).

However, the following code would still fail after the patch, right @tungld?

affine.for %arg2 = 0 to 10 {
  %2 = affine.load %arg1[] : memref<f32>
  affine.store %2, %arg0[%arg2] : memref<10xf32>
}
%0 = affine.load %arg1[] : memref<f32>
%1 = addf %0, %0 : f32
affine.for %arg2 = 0 to 10 {
  %2 = affine.load %arg0[%arg2] : memref<10xf32>
  %3 = divf %1, %2 : f32
}

Maybe we should create a single node that includes all the top-level ops between loops?
That would require further discussion, though. The current fix should be good for now.

It seemed it could identify the defining op %0 = affine.for..., but the output program was not
semantically similar to the input program due to step in affine.for. I think this is somewhat
related to what @dcaballe mentioned

Sorry, I should've mentioned that loops with iter_args are not supported. Affine analysis is only looking at
memory dependences and not taking SSA loop-carried dependences into account. We should prevent loops
with iter_args to be fused. IIRC, @bondhugula suggested that we could demote these SSA values to memrefs
(reg2mem) before fusion so that all the analysis happens through memory.

The patch looks great! Just one more comment.

mlir/lib/Transforms/LoopFusion.cpp

958

Thanks!

1466

We are just checking for a dependence that would prevent finding an insertion point for the fused loop. Could we please move this to getFusedLoopNestInsertionPoint? In that way, sibling fusion will also get it.

[tungld]

However, the following code would still fail after the patch, right @tungld?

affine.for %arg2 = 0 to 10 {
  %2 = affine.load %arg1[] : memref<f32>
  affine.store %2, %arg0[%arg2] : memref<10xf32>
}
%0 = affine.load %arg1[] : memref<f32>
%1 = addf %0, %0 : f32
affine.for %arg2 = 0 to 10 {
  %2 = affine.load %arg0[%arg2] : memref<10xf32>
  %3 = divf %1, %2 : f32
}

It worked because we had load in the first loop which has no write-read dependence to %0 = affine.load %arg1[] : memref<f32>. I got this output using this patch:

%0 = affine.load %arg1[] : memref<f32>
%1 = addf %0, %0 : f32
affine.for %arg2 = 0 to 10 {
  %2 = affine.load %arg1[] : memref<f32>
  affine.store %2, %arg0[%arg2] : memref<10xf32>
  %3 = affine.load %arg0[%arg2] : memref<10xf32>
  %4 = divf %1, %3 : f32
}

To make dependence, I changed to use store as follows:

%cst = constant 0.000000e+00 : f32
  affine.for %arg2 = 0 to 10 {
    affine.store %cst, %arg1[] : memref<f32>
    affine.store %cst, %arg0[%arg2] : memref<10xf32>
  }
  %0 = affine.load %arg1[] : memref<f32>
  %1 = addf %0, %0 : f32
  affine.for %arg2 = 0 to 10 {
      %2 = affine.load %arg0[%arg2] : memref<10xf32>
        %3 = divf %1, %2 : f32
  }

and the output was correct:

%cst = constant 0.000000e+00 : f32
    affine.for %arg2 = 0 to 10 {
      affine.store %cst, %arg1[] : memref<f32>
      affine.store %cst, %arg0[%arg2] : memref<10xf32>
    }
    %0 = affine.load %arg1[] : memref<f32>
    %1 = addf %0, %0 : f32
    affine.for %arg2 = 0 to 10 {
      %2 = affine.load %arg0[%arg2] : memref<10xf32>
      %3 = divf %1, %2 : f32
    }

I think the add between affine.load and affine.for was handled by hasNonAffineUsersOnThePath.

Could we please move this to getFusedLoopNestInsertionPoint? In that way, sibling fusion will also get it.

Yes, let me do this, and come back soon.

[tungld]

We should prevent loops with iter_args to be fused

Hi Diego, do you prefer to have it in this patch? I could do it by simply checking the source and dest affine.for, if one of them uses iter_args (or returns a SSA value), we prevent them to be fused.

[tungld]

Move the check into getFusedLoopNestInsertionPoint and prevent fusing loop nests that return values

[tungld]

I moved the check of defining ops into getFusedLoopNestInsertionPoint and prevented loop nests with values from being fused. Also added TODO for loop nests with values. Thanks!

[bondhugula]

In D97030#2578354, @tungld wrote:

We should prevent loops with iter_args to be fused

Hi Diego, do you prefer to have it in this patch? I could do it by simply checking the source and dest affine.for, if one of them uses iter_args (or returns a SSA value), we prevent them to be fused.

Several affine passes need to be updated to correctly handle / bail out on iter_args. Loop unroll was recently fixed in its presence. I think it could be done in another patch.

[bondhugula]

LGTM.

mlir/lib/Transforms/LoopFusion.cpp

373

Should this be named gatherNonMemRefDefiningNodes? I'm a bit confused. Or is it clear from other context?

mlir/test/Transforms/loop-fusion.mlir

2959

This TODO could be confusing. It's not clear whether handling return values/iter_args in fusion is worthwhile; it may be better/simpler to do a reg2mem, do the fusion without any iter_args, and then again do mem2reg. You can either update the comment to reflect this or simply drop the comment.

[bondhugula]

If you have fixed for the iter_args handling, please also update the commit summary to reflect that.

[dcaballe]

I think the add between affine.load and affine.for was handled by hasNonAffineUsersOnThePath.

Great!

LGTM. It should be good to go after addressing the remaining minor comments. Thanks!

mlir/test/Transforms/loop-fusion.mlir

2938

typo

2952

Please, add checks for ops in the loop body. Loops could be fused without removing the src loop.

2959

+1

[tungld]

Edit comments and lit tests

[tungld]

Hi Uday and Diego,

Thank you so much for your comments! I just updated the patch.

I don't have permission to land this, please help me to do it. My "name (email)" is "Tung D. Le ([e-mail])".

Thanks,
Tung.

mlir/lib/Transforms/LoopFusion.cpp

373

By definition MDG edge, non-memref value will refer to the dependence that one op defines values used in another op. Anyway, I updated comments here to make it clearer.

[bondhugula]

In D97030#2581436, @tungld wrote:

Hi Uday and Diego,

Thank you so much for your comments! I just updated the patch.

I don't have permission to land this, please help me to do it. My "name (email)" is "Tung D. Le ([e-mail])".

Thanks,
Tung.

@tungld could you please rebase and fix a conflict due to another revision that updated fusion?

[tungld]

Rebase and fix a conflict

[tungld]

@tungld could you please rebase and fix a conflict due to another revision that updated fusion?

Done. Thanks!

[tungld]

@dcaballe @bondhugula any further comments? Thanks!

[dcaballe]

LGTM! Just minor comments. I think Uday planned to commit it. Otherwise, I could commit it tomorrow together with one of mine.

mlir/test/Transforms/loop-fusion.mlir

2953

You can remove the %{{.*}} = if you are not capturing the return value.

2955

Just matching the ops here would suffice:

// CHECK-NEXT:    affine.load
// CHECK-NEXT:    affine.yield

[tungld]

Edit lit tests according to @dcaballe's comments.

[tungld]

Just minor comments.

Done. Thanks!

mlir/lib/Transforms/LoopFusion.cpp

struct Edge {
     // If this edge is stored in Edge = Node.outEdges[i], then
     // 'Node.outEdges[i].id' is the identifier of the dest node of the edge.
     unsigned id;
     // The SSA value on which this edge represents a dependence.
     // If the value is a memref, then the dependence is between graph nodes
     // which contain accesses to the same memref 'value'. If the value is a
     // non-memref value, then the dependence is between a graph node which
     // defines an SSA value and another graph node which uses the SSA value
-    // (e.g. a constant operation defining a value which is used inside a loop
-    // nest).
+    // (e.g. a constant or load operation defining a value which is used inside
+    // a loop nest).
     Value value;
   };
 
   // Map from node id to Node.
   DenseMap<unsigned, Node> nodes;
   // Map from node id to list of input edges.
   DenseMap<unsigned, SmallVector<Edge, 2>> inEdges;
   // Map from node id to list of output edges.
unsigned getOutEdgeCount(unsigned id, Value memref = nullptr) {
     if (outEdges.count(id) > 0)
       for (auto &outEdge : outEdges[id])
         if (!memref || outEdge.value == memref)
           ++outEdgeCount;
     return outEdgeCount;
   }
 
   // Computes and returns an insertion point operation, before which the
   // the fused <srcId, dstId> loop nest can be inserted while preserving
   // dependences. Returns nullptr if no such insertion point is found.
   Operation *getFusedLoopNestInsertionPoint(unsigned srcId, unsigned dstId) {
     if (outEdges.count(srcId) == 0)
       return getNode(dstId)->op;
 
     // Build set of insts in range (srcId, dstId) which depend on 'srcId'.
     SmallPtrSet<Operation *, 2> srcDepInsts;
     for (auto &outEdge : outEdges[srcId])
bool MemRefDependenceGraph::init(FuncOp f) {
 
   for (auto &idAndNode : nodes) {
     LLVM_DEBUG(llvm::dbgs() << "Create node " << idAndNode.first << " for:\n"
                             << *(idAndNode.second.op) << "\n");
     (void)idAndNode;
   }
 
   // Add dependence edges between nodes which produce SSA values and their
-  // users.
+  // users. Load ops can be considered as the ones producing SSA values.
   for (auto &idAndNode : nodes) {
     const Node &node = idAndNode.second;
-    if (!node.loads.empty() || !node.stores.empty())
+    if (!node.stores.empty())
       continue;
     auto *opInst = node.op;
     for (auto value : opInst->getResults()) {
       for (auto *user : value.getUsers()) {
         SmallVector<AffineForOp, 4> loops;
         getLoopIVs(*user, &loops);
         if (loops.empty())
           continue;
static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   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.
+/// true. Otherwise, return false.
 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
while (!worklist.empty()) {
                 << "Can't fuse: non-affine users in between the loops\n.");
             continue;
           }
 
           // Compute an operation list insertion point for the fused loop
           // nest which preserves dependences.
           Operation *fusedLoopInsPoint =
               mdg->getFusedLoopNestInsertionPoint(srcNode->id, dstNode->id);
           if (fusedLoopInsPoint == nullptr)
             continue;
 
           // Compute the innermost common loop depth for dstNode
           // producer-consumer loads/stores.
           SmallVector<Operation *, 2> dstMemrefOps;
           for (Operation *op : dstNode->loads)
             if (producerConsumerMemrefs.count(
                     cast<AffineReadOpInterface>(op).getMemRef()) > 0)
               dstMemrefOps.push_back(op);
           for (Operation *op : dstNode->stores)
             if (producerConsumerMemrefs.count(
                     cast<AffineWriteOpInterface>(op).getMemRef()))
               dstMemrefOps.push_back(op);
           unsigned dstLoopDepthTest = getInnermostCommonLoopDepth(dstMemrefOps);
 
           // Check the feasibility of fusing src loop nest into dst loop nest
           // at loop depths in range [1, dstLoopDepthTest].

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
 }
+
+// -----
+
+// CHECK-LABEL: func @should_fuse_defining_node_has_no_dependence_on_source_node
+func @should_fuse_defining_node_has_no_dependence_on_source_node(
+    %a : memref<10xf32>, %b : memref<f32>) -> () {
+  affine.for %i0 = 0 to 10 {
+    %0 = affine.load %b[] : memref<f32>
+    affine.store %0, %a[%i0] : memref<10xf32>
+  }
+  %0 = affine.load %b[] : memref<f32>
+  affine.for %i1 = 0 to 10 {
+    %1 = affine.load %a[%i1] : memref<10xf32>
+    %2 = divf %0, %1 : f32
+  }
+
+  // Loops '%i0' and '%i1' should be fused even though there is a defining
+  // node between the loops. It is because the node has no dependence on '%i0'.
+  // CHECK:       affine.load %{{.*}}[] : memref<f32>
+  // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.load %{{.*}}[] : memref<f32>
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    divf
+  // CHECK-NEXT:  }
+  // CHECK-NOT:   affine.for
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @should_not_fuse_defining_node_has_dependence_on_source_loop
+func @should_not_fuse_defining_node_has_dependence_on_source_loop(
+    %a : memref<10xf32>, %b : memref<f32>) -> () {
+  %cst = constant 0.000000e+00 : f32
+  affine.for %i0 = 0 to 10 {
+    affine.store %cst, %b[] : memref<f32>
+    affine.store %cst, %a[%i0] : memref<10xf32>
+  }
+  %0 = affine.load %b[] : memref<f32>
+  affine.for %i1 = 0 to 10 {
+    %1 = affine.load %a[%i1] : memref<10xf32>
+    %2 = divf %0, %1 : f32
+  }
+
+  // Loops '%i0' and '%i1' should not be fused because the defining node
+  // of '%0' used in '%i1' has dependence on loop '%i0'.
+  // CHECK:       affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[] : memref<f32>
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:  }
+  // CHECK-NEXT:  affine.load %{{.*}}[] : memref<f32>
+  // CHECK:       affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    divf
+  // CHECK-NEXT:  }
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @should_not_fuse_defining_node_has_transitive_dependence_on_source_loop
+func @should_not_fuse_defining_node_has_transitive_dependence_on_source_loop(
+    %a : memref<10xf32>, %b : memref<10xf32>, %c : memref<f32>) -> () {
+  %cst = constant 0.000000e+00 : f32
+  affine.for %i0 = 0 to 10 {
+    affine.store %cst, %a[%i0] : memref<10xf32>
+    affine.store %cst, %b[%i0] : memref<10xf32>
+  }
+  affine.for %i1 = 0 to 10 {
+    %1 = affine.load %b[%i1] : memref<10xf32>
+    affine.store %1, %c[] : memref<f32>
+  }
+  %0 = affine.load %c[] : memref<f32>
+  affine.for %i2 = 0 to 10 {
+    %1 = affine.load %a[%i2] : memref<10xf32>
+    %2 = divf %0, %1 : f32
+  }
+
+  // When loops '%i0' and '%i2' are evaluated first, they should not be
+  // fused. The defining node of '%0' in loop '%i2' has transitive dependence on
+  // loop '%i0'. After that, loops '%i0' and '%i1' are evaluated, and they will
+  // be fused as usual.
+  // CHECK:       affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    affine.store %{{.*}}, %{{.*}}[] : memref<f32>
+  // CHECK-NEXT:  }
+  // CHECK-NEXT:  affine.load %{{.*}}[] : memref<f32>
+  // CHECK:       affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    affine.load %{{.*}}[%{{.*}}] : memref<10xf32>
+  // CHECK-NEXT:    divf
+  // CHECK-NEXT:  }
+  // CHECK-NOT:   affine.for
+  return
+}