Add a CPU lowering pass

microsoft / triton-shared

Shared Middle-Layer for Triton Compilation

MIT License

142 stars 27 forks source link

Add a CPU lowering pass #26

Closed shintaro-iwasaki closed 9 months ago

shintaro-iwasaki commented 9 months ago

Per discussion in #24, this PR creates a reference-CPU backend for triton-shared using the standard MLIR->LLVM lowering passes.

Without any changes in the Triton runtime, we can run triton-shared on CPU from Python-Triton code.
- Potentially it is useful for testing.
A very basic kernel (see reduce.py) runs and succeeds.
Many kernels do not run because of a lack of LLVM lowering and other problems.
- For example, memref::tensorStore->LLVM seems not implemented, so many kernels are not supported at this point.

Help from the community and triton-shared core developers is needed to improve and maintain it. I'd appreciate your feedback and suggestions.

Note: compared with the PR I used in the discussion, this PR version can run a compute kernel (reduce.py, as an example) by using MLIR's CRunnerUtils.h.

(EDIT: let me understand if I can sign this CLA)

nhat-nguyen commented 9 months ago

Thank you so much @shintaro-iwasaki! I'll start taking a look 😄

nhat-nguyen commented 9 months ago

@kshama-msft Would you mind helping Shintaro with the CLA question? Thank you!

kshama-msft commented 9 months ago

@shintaro-iwasaki if the CLA has been generated by the github bot then feel free to sign it. You should review with your internal open-source contribution legal teams to confirm that it is okay for you to sign.

shintaro-iwasaki commented 9 months ago

Thanks a lot for your review, @nhat-nguyen! I am still figuring out the CLA thing with the legal team. Please give me some more time. Once it's cleared, let me fix all the issues pointed out.

shintaro-iwasaki commented 9 months ago

@microsoft-github-policy-service agree [company="Meta Platforms, Inc."]

shintaro-iwasaki commented 9 months ago

@microsoft-github-policy-service agree company="Meta Platforms, Inc."

shintaro-iwasaki commented 9 months ago

The legal thing has been cleared on my side. Let me fix the issues pointed out and update the PR by the end of today.

shintaro-iwasaki commented 9 months ago

Thanks! I updated the commit as well as rebased onto the latest main. Please let me know any further suggestions!

(Do you want the PR author (me) to resolve the comments? Or should it be done by reviewers?)

nhat-nguyen commented 9 months ago

Thanks! I updated the commit as well as rebased onto the latest main. Please let me know any further suggestions!

(Do you want the PR author (me) to resolve the comments? Or should it be done by reviewers?)

Thank you @shintaro-iwasaki, everything looks good! We usually let the PR authors resolve the comments but let me know if you'd prefer otherwise. Once the pipeline runs finish, I can merge in the changes. 😄

nhat-nguyen commented 9 months ago

Merged! Thank @shintaro-iwasaki again!