[TKW] Move IGEMM conv impl to common place. #295
Conversation
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
iree/turbine/kernel/kernels.py
Outdated
| @@ -0,0 +1,158 @@ | |||
| # Copyright 2024 The IREE Authors | |||
There was a problem hiding this comment.
Maybe we can put it in iree/turbine/kernel/wave/templates/igemm.py or iree/turbine/kernel/wave/templates/conv.py. Here are some of my thoughts/ideas behind this:
- Reasoning under
iree/turbine/kernel/wave/because it is specifically awavekernel and with the current stateiree/turbine/kernel/is not wave exclusive. - Under
iree/turbine/kernel/wave/templatebecause it will be somewhat uniform to the way shark-ai stack is organizing their kernels/template kernels. (example link) - Should be in it's own file
igemm.py/conv.py, because in the future, we may store different variants of igemm or conv (such as igemm, direct convolution, backward convs, etc). In a similar fashion that in the future, we may haveattention.pythat has different variants such as vanilla, paged, quantized, evoform.
There was a problem hiding this comment.
Good idea, done.
iree/turbine/kernel/kernels.py
Outdated
| repeat, out, mapping=out_mapping, elements_per_thread=ELEMS_PER_THREAD | ||
| ) | ||
|
|
||
| symbols = { |
There was a problem hiding this comment.
Can we move the symbols: dict[Symbol, int] and other non-symbolic/materialized variables out of the kernels.py file into the driver file?
Looking at it from a perspective of a library writer, I'd probably want to expose ratios, block_m, block_n, block_k and all these other stuff in the driver such that I can come up with my own heuristics/tunings/configs for igemm depending on the input size.
There was a problem hiding this comment.
Without this there is no way for user to know which symbols are actually used inside the kernel. My idea was to return symbols dict with some sensible defaults and user then can tweak it on their side.
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Move TKW IGEMM conv impl from the test folder to some common place to allow it to be reused outside the tests (e.g. in iree-kernel-benchmark). Not sure what the proper place for it, suggestions are welcome. --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com> Signed-off-by: xintin <vermagaurav9408@gmail.com>
Move TKW IGEMM conv impl from the test folder to some common place to allow it to be reused outside the tests (e.g. in iree-kernel-benchmark).
Not sure what the proper place for it, suggestions are welcome.