Writing a backend for PyTorch in Kompute

[han-minhee]

I'm planning to implement a PyTorch backend in Kompute, first aiming at a few basic operators and moving onto the full set.
Is there any project like this, or what are your thoughts?

I have some experience in porting PyTorch/Tensorflow models to C++ CUDA/OpenCL programs, but don't have any experience with modifying ATen/ or PyTorch code itself. I also have to admit that I've just started trying Vulkan / Kompute, but I already feel attached to it :)

[axsaucedo]

This sounds interesting @han-minhee ,there has been previous implementations of Kompute as backend for ml frameworks such as jax, you can have a look at that to get an idea as well. Writing a backend for PyTorch would be a large undertaking but feel free to share here any questions (or design docs) and we'll be happy to support and point you in the right direction.

[han-minhee]

@axsaucedo Thank you for the reply! Did you mean vkJAX? This would be a good reference. In addition, do you think I would need to add much codes to Kompute itself, other than using the functions already implemented in Kompute? I currently don't have any id of how much operators needed for the porting are available in the current Kompute. It would be my pleasure if I can add some more codes to the Kompute itself.

[han-minhee]

Reopening the issue as I accidentally closed it

[axsaucedo]

Yes that's the project I was referencing, there may be others that could be relevant to identify

do you think I would need to add much codes to Kompute itself, other than using the functions already implemented in Kompute?

I can't think of anything specific, but if we come across something then it could certainly be explored

I currently don't have any id of how much operators needed for the porting are available in the current Kompute.

Operators are extensible so that shouldn't be the main issue, but any limitations of the framework can be explored

It would be my pleasure if I can add some more codes to the Kompute itself.

Great that can be explored

[han-minhee]

I see! I'll start working on it, though I don't know how much it will take 🙃

[sorasoras]

I see! I'll start working on it, though I don't know how much it will take 🙃

hmm, I was thinking about turning Metal Performance Shaders (MPS) into Vulkan via Kompute, then it could function as plugin like pytorch-directml. Nomic-backend for ggml is interesting reference as well.
Do you have any progress on your part?

[axsaucedo]

@sorasoras that's a super interesting proposal - I would be quite interested to support this as it could indeed provide quite a flexible backend. I believe there's not been much movement here, so it could be a great one to explore especially as now the Nomic-backend for ggml is being integrated and can be used as reference. We'll also be looking at extending the examples sections of the Kompute docs as well as potentially adding tests to ensure we're testing upstream integrations to flag any deltas with changes.