[mlir][xegpu] Refine layout assignment in XeGPU SIMT distribution. #142687
Conversation
|
@adam-smnk sorry to bother you. :-P would you have some time to give me a general review on this or approve? :-) |
| nullptr); | ||
| terminator.getSuccessorRegions(operands, successors); | ||
|
|
||
| for (mlir::RegionSuccessor &successor : successors) { |
There was a problem hiding this comment.
Does this also handle the block arguments of forOp, which is also handled by BranchOpInterface? updateBranchOpInterface maybe not needed in this case.
There was a problem hiding this comment.
great point. I checked this again and seems like you are right. BranchOpInterface is doing some redundant work here.
I removed updateBranchOpInterface and also added some comment with example to clarify the logic.
one caveat is that we still need to a check for region ops and filter them out (done in updateOp).
There was a problem hiding this comment.
General formatting comment: please add more line breaks (both for IR and checks with CHECK-SAME). The longest line right now is almost 500 characters wide which really impacts readability.
There was a problem hiding this comment.
good point. I added more line breaks.
There was a problem hiding this comment.
Similarly, a small cleanup could be done here
There was a problem hiding this comment.
I found its not very severe in this case. so keep unchanged for now.
| }) | ||
| .Case<mlir::FunctionOpInterface>( | ||
| [&](mlir::FunctionOpInterface funcOp) { | ||
| r = updateFunctionOpInterface(builder, funcOp, |
There was a problem hiding this comment.
For completeness, are function calls covered too?
There was a problem hiding this comment.
currently no. we don't expect function calls inside a kernel in my understanding.
| /// clang-format on | ||
| /// In this example, at scf.yield, control-flow can transfer to successor | ||
| /// regions. One is the ^bb0 (for loop body) and the other is the scf.for op | ||
| /// itself (yield the results). So we update both the block arguments of the |
There was a problem hiding this comment.
This may need to double check. My understanding is that the other is not scf.for itself, it is the region right following the scf.for.
There was a problem hiding this comment.
after out discussion I updated the comment again. thanks for the suggestions.
| @@ -12,6 +12,8 @@ | |||
| #include "mlir/Dialect/GPU/IR/GPUDialect.h" | |||
| #include "mlir/Dialect/GPU/Utils/DistributionUtils.h" | |||
There was a problem hiding this comment.
nit: some headers seem not needed. Could you help to clean them? make the header includes as less as possible.
There was a problem hiding this comment.
thanks. I cleaned up unused headers shown by vscode. let me know if there is a better way to find this out.
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h" |
There was a problem hiding this comment.
clean up un-related header includes.
Changes: