-
Notifications
You must be signed in to change notification settings - Fork 32
[AMD] Added canonicalization pattern to propagate DotOp attrs #792
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: refine-ops-pass
Are you sure you want to change the base?
Conversation
This PR: makes Refine pass FuncOp based, instead of ModuleOp replaces walkers with PatternRewriters Co-authored-by: ravil-mobile <[email protected]>
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
Addressed Issue #https://github.com/ROCm/triton-internal/issues/727
Closes #https://github.com/ROCm/triton-internal/issues/727