[AMD] Added canonicalization pattern to propagate DotOp attrs#792
[AMD] Added canonicalization pattern to propagate DotOp attrs#792ravil-mobile wants to merge 50 commits into
Conversation
This PR:
makes Refine pass FuncOp based, instead of ModuleOp
replaces walkers with PatternRewriters
Co-authored-by: ravil-mobile <ravil.aviva.com@gmail.com>
Inlined refine ops functions
[AMD] Added `local_alloc` refinement
- Added a lit-test to `elementwise.mlir` which test refinement of `convert_layout` from one `mma` to another `mma` layoyts - Added a bug fix
|
Alright, I'm getting back to this after our May sprint. For this PR can you add a clarifying comment as to what direction attributes are being propagated and why. You could add something like the below explanation, but fixup the syntax a bit. |
|
|
||
| auto definingOp = operand.getDefiningOp(); | ||
| if (definingOp->getBlock() != currBlock) | ||
| continue; |
There was a problem hiding this comment.
I'm not sure what this is prohibiting; we do want to allow these attributes to be carries across basic blocks. For example, they could be loop carried, or we could prefetch a tile of data in the prologue, and we want to get the order of loads in the prologue to match the dot ordering in the loop.
| allowedDialects |= | ||
| mlir::isa<amdgpu::TritonAMDGPUDialect>(definingOp->getDialect()); | ||
| if (!allowedDialects) | ||
| continue; |
There was a problem hiding this comment.
What is this prohibiting and allowing?
| SmallVector<Attribute> updatedAttrs(operandArrayAttr.getValue()); | ||
| auto result = operandArrayAttr.walk([&targetAttr](AttrType itemAttr) { | ||
| return itemAttr == targetAttr ? WalkResult::interrupt() | ||
| : WalkResult::advance(); | ||
| }); | ||
|
|
||
| if (!result.wasInterrupted()) { | ||
| updatedAttrs.push_back(targetAttr); | ||
| } | ||
|
|
||
| if (!updatedAttrs.empty()) | ||
| definingOp->setAttr( | ||
| attrName, mlir::ArrayAttr::get(op->getContext(), updatedAttrs)); |
There was a problem hiding this comment.
Is this block of code ensuring we only add a dotAttr if one doesn't already exist?
If so, this can also be used as a recursion exit criteria. For example, after we've labeled all the refined local_loads with their dotAttr, then we recursively label all their respective subviews with dot attribute, which is good.
But then when we recur from all the subview to the same parent, we can't place all the children't dotAttr onto the same parent (since the labels only refer to refinement), so we can exit the recursion at that point. Meaning when an op already has a dotAttr, don't recur to it's parent.
3070646 to
0341d75
Compare
ca0a8c2 to
58798c7
Compare
Addressed Issue #https://github.com/ROCm/triton-internal/issues/727
Closes #https://github.com/ROCm/triton-internal/issues/727