[AMDAIEFusePackIntoLoop] Handle case where there is no extract_slice #983
Conversation
newling
requested review from
MaheshRavishankar,
nirvedhmeshram,
yzhang93,
Abhishek-Varma and
jtuyls
as code owners
newling
force-pushed
the
fuse_pack_without_slice
branch
2 times, most recently
from
fff5f23 to
ad87490
Compare
| /// A utility function specific to this pass which, given a value, would | ||
| /// traverse the def-chain till it either finds a tensor.extract_slice op or a | ||
| /// BlockArgument. | ||
| FailureOr<Operation *> getParentBeforeLoop(BlockArgument blkArg) { |
There was a problem hiding this comment.
Add explanations for this utility function.
There was a problem hiding this comment.
I've inlined the function
| FailureOr<Operation *> getParentBeforeLoop(BlockArgument blkArg) { | ||
| Operation *parent = blkArg.getOwner()->getParentOp(); | ||
| LoopLikeOpInterface loop = dyn_cast<LoopLikeOpInterface>(parent); | ||
| if (!loop) return failure(); |
There was a problem hiding this comment.
No failure here isn't necessary
| operand = defOp->getOperand(0); | ||
| // roll back through all the packs immediately preceding `operand`. | ||
| while (operand.getDefiningOp() && | ||
| isa<tensor::PackOp>(operand.getDefiningOp())) { |
There was a problem hiding this comment.
I think you can use while (isa_and_present<tensor::PackOp>(operand.getDefiningOp())).
| Operation *defOp = operand.getDefiningOp(); | ||
| if (!defOp) return failure(); | ||
| tensor::ExtractSliceOp sliceOp = dyn_cast<tensor::ExtractSliceOp>(defOp); | ||
| if (!sliceOp) return failure(); |
There was a problem hiding this comment.
Like commented above, you can simply check dyn_cast_if_present<tensor::ExtractSliceOp>(operand.getDefiningOp()).
| else { | ||
| if (auto parent = | ||
| dyn_cast_if_present<tensor::PackOp>(operand.getDefiningOp())) { | ||
| Block *genericBlock = genericOp->getBlock(); |
There was a problem hiding this comment.
I think here you need to use else if, and then add another else in which adding some debug message and return.
There was a problem hiding this comment.
I'll change to else-if, but the logic is correct I think so won't return and don't think debug print message is useful here.
There was a problem hiding this comment.
Logically the else condition is not needed, but I was thinking it may good to add some debug message so people know why the pack is not fused into the loop in some specific situation. You can decide whether to add it.
There was a problem hiding this comment.
The operand might not even come from a pack though, it could be anything...
newling
force-pushed
the
fuse_pack_without_slice
branch
from
ad87490 to
7aacd04
Compare
| } | ||
|
|
||
| // Case 2 outlined above. | ||
| else if (auto blkArg = dyn_cast<BlockArgument>(sliceOp.getSource())) { |
There was a problem hiding this comment.
We don't want it to enter both statements which might happen if this is if
There was a problem hiding this comment.
How would it enter both statements? After it enters the first it will return.
This is very similar to #976
When the AIE grid we're using is
m x nwhere one or both ofmandnis1, for a matmul we get pack ops that produce operands for matmuls directly. i.e. as opposed topack->extract_slice->matmulwe havepack->matmul.TODO(newling) add a test and see if the pass can be simplified further.