[MLIR][Linalg] Transform pass to hoist vector transfer operations for a batch reduce matmul operation #121748
Conversation
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
arun-thmn
changed the title
arun-thmn
changed the title
|
Hi @rolfmorel, please review the draft PR and let me know the initial comments. |
There was a problem hiding this comment.
Thanks for working on these patterns, @arun-thmn!
I left some comments here and there, though only on the hoisting transform for now.
My main points are as follows:
-
Split into two PRs, one for the hoisting transform and another for the FMA rewrite.
- I suggest moving the FMA rewrite to its own PR.
-
Move files to
mlir/{include,lib}/mlir/Dialect/Vector. -
Try to generalize the hoisting transform?
- that is, are the outer m- and n- loops really necessary for valid hoisting?
- if we can generalize, might be worthwhile to expose the hoisting as a Vector dialect C++ utility function as well.
-
We should check that the relevant matching failures are all exercised. Are there interesting edge cases to consider (e.g. if the loop order is not as expected)?
| "apply_patterns.vector.hoist_vector_transfer", | ||
| [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> { | ||
| let description = [{ | ||
| Finds pattern to hoist the possible vector transfer reads/writes outside the reduction and k-loop |
| void populateFoldAddIntoDestPatterns(RewritePatternSet &patterns); | ||
|
|
||
|
|
||
| /// Pattern to hoists the vector transfer reads/writes outside the reduction and |
There was a problem hiding this comment.
nit: "hoists the" -> "hoist"
| @@ -0,0 +1,267 @@ | |||
| //===- HoistVectorTransfers.cpp ---------------------------------------*- C++-*-===// | |||
There was a problem hiding this comment.
nit: right justification appears off.
|
|
||
| using namespace mlir; | ||
|
|
||
| // Function to retrives vector transfer read operations (Acc, Lhs, and Rhs) from contraction operation. |
There was a problem hiding this comment.
nit: I believe comments should also keep to the 80-char limit. Also holds for comments further down.
Does clang-format leave this comment line (and the others) as is?
There was a problem hiding this comment.
nit1: "retrieves" and "(RHS, LHS, Acc)"
nit2: we probably want to write out op names as "vector.transfer_read".
| return list; | ||
| } | ||
|
|
||
| // Function to retrive subview from vector transfer read operation. |
There was a problem hiding this comment.
nit: "retrive" -> "retrieve" (also in other places)
| reductionForOp.getLoc(), reductionForOp.getLowerBound(), | ||
| reductionForOp.getUpperBound(), reductionForOp.getStep(), | ||
| ValueRange{vectorReadOpValue}, | ||
| [&](OpBuilder &rewriterNewReductionForOp, Location locNewReductionForOp, |
There was a problem hiding this comment.
In case of such nested lambdas, can we please pull out these lambdas into a name-bound lambda or, equivalently, into a method? This is to reduce the rightward drift.
|
|
||
| // Code to re-create the reduction and k loop with iter args | ||
| auto vectorReadOpValue = cloneVectorReadOp->getResult(0); | ||
| auto newReductionForOp = rewriter.create<scf::ForOp>( |
There was a problem hiding this comment.
I would guess LoopLikeOpInterface::replaceWithAdditionalYields is well suited to what you are trying to do here. Please have a look if that would help simplify this section.
| /// scf.for %arg3 = %c0 to %c32 step %c4 // m-loop | ||
| /// scf.for %arg4 = %c0 to %c64 step %c64 // n-loop | ||
| /// %subview_2 = memref.subview %subview[%arg3, %arg4] [4, 64] [1, 1] | ||
| /// scf.for %arg5 = %c0 to %c24 step %c1 // reduction-loop |
|
|
||
| // Code to hoist vector transfer write after reduction loop and also to | ||
| // update the yield of k loop | ||
| auto newKForOp = |
There was a problem hiding this comment.
Here as well: I am not sure why this section appears a bit involved. Per the example in the docs, you are just hoisting one transfer_read and "threading" that vector through the iter_args of two loops, right? Doesn't two applications of LoopLikeOpInterface::replaceWithAdditionalYields suffice for this?
| userOp->erase(); | ||
| } | ||
| } | ||
| contractOp.erase(); |
There was a problem hiding this comment.
Try to make use of the replaceAllUsesWith API. Alternatively, modify the relevant operand once the region containing the contract has been transplanted to the loops with the new iter_args (though this approach is not always safe so probably best avoided as well).
arun-thmn
changed the title
2 participants
This transformation pass hoists the vector transfer reads/writes outside the reduction and k-loop for
batch reduce matmuloperation iflicmfails.