Add initial splitk reduce #2908
Conversation
Why would I create PR to move functions to separate headers when they are only used once?
Why would I create a PR for an operator that is not used anywhere?
This will slow down getting splitk reduce implemented as there are a lot more changes to come for this. I stopped here to avoid having a very large PR and avoiding dealing with merge conflicts. |
| {"reduced", "decltype(" + generate_make_shape(reduce_output_shape) + ")"}, | ||
| {"lambda", v.at("lambda").to<std::string>()}, | ||
| {"transformers", make_transformer_args(vec)}, | ||
| {"preamble", v.get("preamble", std::string{})}}); |
There was a problem hiding this comment.
Nitpick: Not sure why all of this is shifted instead of one change for adding assign.
| })) | ||
| continue; | ||
| auto v = ins->get_operator().to_value(); | ||
| auto axes = v["axes"].to_vector<std::int64_t>(); |
There was a problem hiding this comment.
Add coverage for continues; non float and empty splits() above.
There was a problem hiding this comment.
I would prefer not to add tests for those cases, since I hope to support those other cases in the future, and will add test cases for that then.
There was a problem hiding this comment.
Then add a TODO there or make it an issue we can track for later improvement.
Co-authored-by: Umang Yadav <[email protected]>
This adds a pass to do splitk reductions. It also updates the kernel to handle splitk reduciton as well. Right now, this only works for a single reduce, future work will enable this for multiple reductions. This is only enabled when using
MIGRAPHX_ENABLE_SPLIT_REDUCE=1until we get more testing for this.I do see a 2x improvement when using this on large reductions like 327680.