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

[Abhishek Varma]

================
@@ -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;
----------------
Abhishek-Varma 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)?

Well, it is that but it'd bail out in the case of the IR mentioned in the link - this patch would

Essentially this would bail out cleanly if we have either `%4 = insert_slice %2` or `%5 = extract_slicce %4` in between - as the condition is only going to allow `scf.yield` or `tensor.parallel_insert_slice` as the other uses of producer. And by the nature of these allowed ops since they're terminator ops, we have a clear path to fusing the consumer and isn't adding to the complexity which might be the case in previous PR ?

>  I think it is not complete(or real) multi-consumers fusion

The scope of the PR as explained above is not to deal with multi-consumers fusion of the topology you've shared. The intention is to only extend the current support from :-
```
%a = scf.for/forall {
    %0 = scf.for/forall {
              tiledProducer
        }
    %1=consumerOp1 ins(%0) // ONLY consumer
    scf.yield/tensor.parallel_insert_slice %1 
}
```

to :-
```
%a = scf.for/forall {
    %0 = scf.for/forall {
              tiledProducer
        }
    %1=consumerOp ins(%0) // ONLY "real" consumer
    scf.yield/tensor.parallel_insert_slice %0 // Clearly this has been made OPTIONAL
    scf.yield/tensor.parallel_insert_slice %1 
}
```
Apologies for adding confusion with "real" - but I hope I've been able to communicate the intention of the PR and how this isn't adding to the complexity (as it isn't aiming to deal with the IR you mentioned).




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