[Mlir-commits] [mlir] [MLIR][TilingInterface] Extend consumer fusion for multi-use of producer (PR #110105)

[llvmlistbot at llvm.org]

================
@@ -1481,21 +1481,33 @@ checkAssumptionForFusingConsumer(tensor::InsertSliceOp candidateSliceOp) {
 /// failure otherwise.
 static FailureOr<OpOperand *> getConsumerFromUses(Value val,
                                                   Block *containingOpBlock) {
-  // Step 1. Check that the value has exactly one use.
-  if (!llvm::hasSingleElement(val.getUses()))
-    return failure();
-  // Step 2. Get uses.
-  OpOperand &operand = (*val.getUses().begin());
-  Operation *consumerOp = operand.getOwner();
-  // TODO: We have to init result of consumer before scf.for, use
-  //       DestinationStyleOpInterface to get result shape from init for now.
-  //       Add support for other op such as op has InferTypeOpInterface.
-  if (!isa<TilingInterface>(consumerOp) ||
-      !isa<DestinationStyleOpInterface>(consumerOp))
-    return failure();
-  if (containingOpBlock != consumerOp->getBlock())
-    return failure();
-  return &operand;
+  // Check that the value has exactly one use which isn't a scf.yield or a
+  // tensor.parallel_insert_slice op.
+  Operation *visitedConsumerOp = nullptr;
----------------
Yun-Fly wrote:

Hi. First of all, IIUC, this patch actually intends to fuse consumer of the tiled producer with multiple user, one of which is normal `consumerOp` and the rests are `scf.yield` or `tensor.parallel_insert_slice`, just like what we talked [here](https://github.com/llvm/llvm-project/pull/94190/commits/dfeccecb599853fc14fb9f2a20ba98d0cbc15017#r1735809986)?

If so, I think it is not complete(or real) multi-consumers fusion, which should at least cover following topology:

```
%0 = scf.for(){
  tiledProducer ....
}
%1=consumerOp1 ins(%0)
%2=consumerOp2 ins(%0)
```
 
BTW: I have added this support in previous PR in fact, but finally reverted as suggested considering of review complexity.

https://github.com/llvm/llvm-project/pull/110105