-
Notifications
You must be signed in to change notification settings - Fork 791
[CPU] Tile reduction dimensions for non-root reduction ops. #21500
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
Conversation
| // Tile all the reduction ops for target vector sizes. It is a nop for | ||
| // rootOp because it is already tiled with the same tile sizes. It | ||
| // ensures that all the dimensions are tiled in all the reduction ops. | ||
| funcPassManager.addPass( | ||
| createLLVMCPUTilePass(static_cast<IREE::CPU::TilingLevel>(i))); |
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.
It is not NOP if the dimension size is not divisible by tile size, because it becomes dynamic shapes and the for loops can't be folded away.
I think we need to provide an option to LLVMCPUTilePass that allows only tiling on non-root ops.
maxbartel
left a comment
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.
Thanks @hanhanW! We will try this out on Monday and tell you if we still see some issues.
b37133e to
cd80291
Compare
In softmax dispatch, there are two reduction ops. Only the root op is tiled for reduction dimensions when we switched to LLVMCPUTileRootAndFuseInputOperandsPass. It results in large vector sizes in the other reduction op when `util.assume.hint` ops are present. We did not hit the issue in e2e tests because AnnotateDispatchAssumptions pass behaves differently. The value range is [0, 0] if the input is from `flow.tensor.dynamic_constant`. Fixes iree-org#21359 Signed-off-by: hanhanW <[email protected]>
Signed-off-by: hanhanW <[email protected]>
…generic` Signed-off-by: hanhanW <[email protected]>
…#21500) The revision adds an option to skip root op in LLVMCPUTile pass, and uses it in multi level tiling pipeline. In softmax dispatch, there are two reduction ops. Only the root op is tiled for reduction dimensions when we switched to LLVMCPUTileRootAndFuseInputOperandsPass. It results in large vector sizes in the other reduction op when `util.assume.hint` ops are present. We did not hit the issue in e2e tests because AnnotateDispatchAssumptions pass behaves differently. The value range is [0, 0] if the input is from `flow.tensor.dynamic_constant`. Fixes iree-org#21359 --------- Signed-off-by: hanhanW <[email protected]>
…#21500) The revision adds an option to skip root op in LLVMCPUTile pass, and uses it in multi level tiling pipeline. In softmax dispatch, there are two reduction ops. Only the root op is tiled for reduction dimensions when we switched to LLVMCPUTileRootAndFuseInputOperandsPass. It results in large vector sizes in the other reduction op when `util.assume.hint` ops are present. We did not hit the issue in e2e tests because AnnotateDispatchAssumptions pass behaves differently. The value range is [0, 0] if the input is from `flow.tensor.dynamic_constant`. Fixes iree-org#21359 --------- Signed-off-by: hanhanW <[email protected]> Signed-off-by: keshavvinayak01 <[email protected]>
The revision adds an option to skip root op in LLVMCPUTile pass, and uses it in multi level tiling pipeline.
In softmax dispatch, there are two reduction ops. Only the root op is tiled for reduction dimensions when we switched to LLVMCPUTileRootAndFuseInputOperandsPass. It results in large vector sizes in the other reduction op when
util.assume.hintops are present.We did not hit the issue in e2e tests because AnnotateDispatchAssumptions pass behaves differently. The value range is [0, 0] if the input is from
flow.tensor.dynamic_constant.Fixes #21359