Some higher level operations such as torch.max generates linalg generic
that returns both the index and the value of the max operation. However
sometimes not all information is being used. This however blocks
vectorization for certain cases which causes performance degradation.
This patch aims to fix this issue.
Details
Diff Detail
- Repository
- rG LLVM Github Monorepo
Event Timeline
Thanks Stanley! Have a few refactoring requests to make the code more understandable.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
1035 | No related to this patch, but just to make it more readable since the method is getting too big now. Can you change this to if (!generic.hasTensorSemantics()) { .... return dedupedOutputs; } /// body of the else goes here.... | |
1061–1063 | To make the logic easier to follow do you mind moving lines 1006 - 1017 and 1030 - 1074 into a helper function that just uses the outputOpOperand. So something like bool isOutputOpOperandDead(linalg::GenericOp genericOp, OpOperand &outOperand, OpResult result) { if (!result.use_empty()) return false; if(!genericOp.payloadUsesValueFromOperand(outputOpOperand)) return true; // Logic for cycle added in this change. } ... deduplicateOutputOperands ... { ... for (const auto &outputOperand : ...) { .... if (isOutputOpOperandDead(...)) { droppedOpOperands.push_back(...) if (genericOp.canOpOperandsBeDropped(droppedOpOperands)) continue; droppedOpOperands.pop_back(); } // Lines 1024 - 1029.... } If this wasnt easy to follow, feel free to punt to me. |
THanks for refactoring. More comments. Marking as requesting change till tests are added.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
876 | Nit: I think you can use genericOp.getBody(). | |
878 | This can be an assert I think | |
882 | Please use getRegionOutputArgs (https://github.com/llvm/llvm-project/blob/3b20765cf725e57d0f99c3d292fe4891b0ee801d/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L364) and get outArg from there. Even better, just take the BlockArgument as the argument instead of OpIndex. |
Exclude dynamic case which was unsupported to begin with and is unlikely to be outs anyways.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
1109 | I don't think we need to check this. The dynamic case is not unsupported. It is expected to be supported. Could you provide more context here? |
@mravishankar It was breaking one of IREE's e2e test https://github.com/iree-org/iree/blob/main/tests/e2e/regression/dynamic_torch_index_select_negative.mlir. It tried creating a placeholder/dummy tensor for a dynamic tensor and it broke.
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
885–895 | the operand index of the yield has to match block argument index otherwise it is wrong. it could be something like: yield %arg1, %arg0 |
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
885–895 | shouldn't this not matter for this part of the code since, this part is only checking for uses? |
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
885–895 | it's mean to check if the cycle is dead right? Doesn't this affect it? |
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
885–895 | Understood, Mahesh also helped clarify this, thanks for this suggestion :) |
Mostly looks good. Have a few nits. Should be ready to land after that!
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
861 | You can drop BlockArgument & for BlockArgument . Could we make the ABI a little simpler (we want to avoid the implicit assumption that outputs come after inputs in linalg ops). So make this bool isResultValueDead(linalg::GenericOp genericOp, OpResult result). Then outputOpOperand = genericOp.getOutputOperands(result.getResultNumber()); outputArg = genericOp.getRegionOutputArgs(result.getResultNumber()); Also below to check the usage in the yield you can just do if (yieldOp.getOperand(result.getResultNumber()) != outputArg) return false; | |
880 | You can drop this check. You are checking below argUserOp is a linalg::YieldOp. That has no uses by construction. | |
885 | Replace these lines with auto yieldOp = dyn_cast<linalg::YieldOp>(argUserOp); if (!yieldOp) return false; Then see above for changing the condition check. | |
1111 | Do we still need to do this? | |
1232 | I think you can replace this with if (!genericOp.hasTensorSemantics()) return failure() Also note, LLVM style is to not have braces around bodies that are a single statement (unlike IREE for example) |
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
1269 | looks like my comment shifted line after you latest update. My comment was about this line: |
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
1268 | Instead of returning from within the loop, can you just check if a modification was made or not in the loop and return success or failure after the loop |
Moved return out of the loop, brought back updateRootInPlace, and moved yieldOp to before mergeblocks s.t we don't need to replace blockarguments
Thanks!
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
1273 | Nit: return success(hasRemovedCycles); |
Looks like this broke the windows mlir buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/26882
Hey @stella.stamenova, thanks for the notifications, I am pushing up a fix here https://reviews.llvm.org/D135612 :)
mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | ||
---|---|---|
860 | done, thanks! | |
861 | done! |
this should be a static function