Can we rewrite pytorch backend in triton? about triton

[ptillet]

I think that the majority of PyTorch GPU kernels could be rewritten much more concisely with Triton. This would also make it possible for PyTorch to no longer depend so heavily on cuBLAS, cuDNN, etc. and eventually make the kernels more portable across different hardware architectures. But one thing that the PyTorch team would probably want before they do that is an offline Triton compiler, so that DNN models can be properly deployed; now, Triton is only JIT-compiled

[ptillet]

@chsasank AMD has actually been looking at Triton for a while and I think they have started working on adding support. It sounds like a pretty hard project for someone unfamiliar with LLVM.

I think an off-line compiler project is probably more realistic. There's been a few efforts already, but none aimed at being production-ready. I think there is value in a proper offline compiler that can be used to generate binaries that can somehow be called from Python like a standard jit'd function :)

@S3Y3D I looked into it a while back but I think it looks exceedingly difficult to me. There are some papers that have added auto-diff support for CUDA-like code, but honestly I'm not confident this can be done without any performance hit. I think a better place for a JIT would be in a JAX/TorchScript-like JIT compiler based on Triton

[ptillet]

I'll close this for now, as at the moment it seems to be more an issue for PyTorch than for Triton :)

[chsasank]

That's great to hear and it would be an interesting project to rewrite some of the backend. Moving to triton kind of project allows adding support new hardware lot easier.

I am interested in AMD GPU support #46. I see the issue was raised about 6 months back. I wonder if there was any progress made already? If not, I can possibly contribute but would be nice to know how much effort it might be to add new backend.

[chsasank]

LLVM noob here but I can also contribute offline triton compiler if that's something not too hard.

[S3Y3D]

Hi,
What about 'Automatic Differentiation' in Triton like in PyTorch and JAX?

[adsharma]

I think there is value in a proper offline compiler that can be used to generate binaries

Are you envisioning a python AST -> systems programming language like C++/Rust?

I work on https://github.com/adsharma/py2many, which is currently focused on producing executables. It works the same way as this code:

https://github.com/openai/triton/blob/master/python/triton/code_gen.py

But the AST visitor spits out C++, Rust or 6 other languages.

There is some experimental support for generating python extensions from python AST via pyO3.

[ptillet]

@adsharma I think we want to stay focused on generating Triton-IR only. However, we'd love to have more frontends. A rust or C++ equivalent of [Python] @triton.jit would be really valuable.