[Wave] Fun projects for beginnners

[raikonenfnu]

Are you interested in learning more about GPU programming and developing cool optimizations? Do you want to help build next generation and state-of-the-art machine learning models and layers? Do you want to define the future programming paradigm of machine learning and GPU layers? Look no further, come join us in building "Wave"!

Here are some fun starter tasks to look at:

Core infrastructure

• Implement backwards attention kernel on TKW
• Scalar support #285 - @nithinsubbiah
• Auto detection of hardware architecture in runtime - [TKW] Add CDNA2 + CDNA3 Int8 intrinsics and refactor intrinsic enums #279
• Support more architecture in codegen aside from CDNA.
• Refactor Wave ops definition and make it autogen documentations
• Replace compile context with compile command

Useful Integration/Deployment on LLM and GenAI model

• Implement Pytorch Linear Layer (forward, backward) that uses custom TKW kernel
• Implement Pytorch Conv2d Layer (forward, backward) that uses custom TKW kernel
• Implement Pytorch Attention Layer (forward, backward) that uses custom TKW kernel
• Integrate Wave kernel into sharktank for e2e S.T we can inject IR from TKW into model IR.

Useful Operations for Quantized LLM and GenAI workload

• Add support for tkw.tanh (lowering to math::TanhOp)
• Add support for tkw.abs #322 - @NoumanAmir657
• Add support for tkw.round_even (math::RoundEvenOp) - @egebeysel
• Add support for tkw.minimum (elementwise/ arith::MinimumFOp) #339 - @NoumanAmir657
• Add support for tkw.clamp (takes upper bound, lowerbound arguments and combines tkw.minimum and tkw.maximum)
• Add support for tkw.min (reduction) dependent on tkw.minimum (elementwise)

[NoumanAmir657]

Would this be open for someone who has almost no knowledge of GPU programming except for some basic understanding.

[raikonenfnu]

Hi @NoumanAmir657 thanks for reaching out! There are plenty of tasks to go around, depending on your experience level, would probably assign easier tasks to start with.

[NoumanAmir657]

Hi @NoumanAmir657 thanks for reaching out! There are plenty of tasks to go around, depending on your experience level, would probably assign easier tasks to start with.

Hi, yes would love any pointers to which task would help me getting started. You can assign me whichever you see fit and I can work towards it

[raikonenfnu]

@NoumanAmir657 I think a good starter task is to implement tkw.abs and it's lowering into math.absfOp or math.absIOp in math dialect of MLIR. Here is a sample of a similar task on reciprocal (82852d1).

LMK what you think. :)

Disclaimer: if this(or any task assigned to external collaborator) becomes high priority, someone internal may take on this task to get it out asap.

[NoumanAmir657]

@NoumanAmir657 I think a good starter task is to implement tkw.abs and it's lowering into math.absfOp or math.absIOp in math dialect of MLIR. Here is a sample of a similar task on reciprocal (82852d1).

LMK what you think. :)

Disclaimer: if this(or any task assigned to external collaborator) becomes high priority, someone internal may take on this task to get it out asap.

Yes this seems a good starting point. I shall make a PR soon. Thanks!

[raikonenfnu]

thanks @NoumanAmir657 for the awesome work with tkw.abs! let me know if you're interested in picking up other issues. The next important thing for us is adding support for tkw.minimum (elementwise) and extending it's support to tkw.min(reduction) but feel free to pick other things more aligned with your interest of course.

[NoumanAmir657]

thanks @NoumanAmir657 for the awesome work with tkw.abs! let me know if you're interested in picking up other issues. The next important thing for us is adding support for tkw.minimum (elementwise) and extending it's support to tkw.min(reduction) but feel free to pick other things more aligned with your interest of course.

Thanks for getting me started. I want to contribute more. Over the weekend, I will go over issues and decide on one. Thanks!

[NoumanAmir657]

@raikonenfnu Hi. You can assign me the tkw.minimum issue.
Thanks!

[raikonenfnu]

Hi. You can assign me the tkw.minimum issue.

@NoumanAmir657 SG! thanks :)

[egebeysel]

Hi, can I volunteer for the tkw.round_even support?, thanks :)

[raikonenfnu]

Hi, can I volunteer for the tkw.round_even support?, thanks :)

Sounds great! :)

[raikonenfnu]

@egebeysel Here are some sample PRs that does similar thing to what you'd be doing: (82852d1, 71eb1c8)

But in this case, you'd want to lower tkw.round_to_even to math.round_even from the math standard dialect in MLIR https://mlir.llvm.org/docs/Dialects/MathOps/#mathroundeven-mathroundevenop. Do reach out if you have any questions! :)

[ziereis]

@raikonenfnu Hello, I have some questions regarding "Support more architectures in codegen aside from CDNA." Currently, the codegen directly emits ops from the AMDGPU dialect, like barriers, MMA ops, etc. What is the plan for supporting GPUs from different vendors? Are these ops going to be turned into a higher-level dialect, for instance a vector.contract instead of a amdgpu.mma and let iree handle the lowering? Or is the plan to directly emit the ops for the target architecture in codegen.py?

[raikonenfnu]

Hey @ziereis SFLR, just got back from holidays!

What is the plan for supporting GPUs from different vendors?

We do plan on adding support for different GPUs. However, currently we are putting a little more prioity in polishing variance of attention, low-bit precisions/quantizations, and also performance but if someone from the community is interested in adding support on different HWs, we are happy to talk RFCs and take PRs!

Are these ops going to be turned into a higher-level dialect, for instance a vector.contract instead of a amdgpu.mma and let iree handle the lowering? Or is the plan to directly emit the ops for the target architecture in codegen.py?

We are planning to keep most of the emitting of ops in iree-turbine (not iree proper). Hence the latter, or a variant of the latter option would most likely be the way we are going down. If you are interested in adding support on a new target architecture, do let me or us know! Would be happy to brainstorm over an GH issue or VC. :)

[NoumanAmir657]

@raikonenfnu you can assign me that tkw.min issue. Thanks!

[NoumanAmir657]

@raikonenfnu Can you help me pick another issue?

[raikonenfnu]

Hey @NoumanAmir657, sure thing! I think now that you are quite familiar with the code base there are two potential interesting projects here:

1. Refactor wave kernel compilation from context based to function call based. Super useful for future wave caching and runtime performance as well as code quality.
2. Scalar codegen, how do we pass in scalar arguments as inputs to kernels, super useful for quantization stuff

[NoumanAmir657]

@raikonenfnu I would like to work on the 1st one. Can you point me to how to get started on this?