Integration in PyTorch

[flecomet]

Hi,
Thank you for this work, it looks really promising.
Maybe my question is a bit naive, but what would be needed to integrate your optimized kernels in PyTorch directly ?
Have you looked into it already ?

[andravin]

Spio already works with PyTorch. It creates custom operations that worth with PyTorch Inductor.

The point of Spio was to show the power of just in time compiled kernels written in CUDA. PyTorch does not really enable this directly, so I built a framework to do so.

Spio achieves higher performance than the Triton Language kernels for grouped convolution, and does it with faster compilation and no auto-tuning.

As a matter of fact, PyTorch still uses cuDNN kernels for grouped convolution, because the Triton Language kernels are slower; and Spio is 3x faster than cuDNN.

Unlike other JIT frameworks, Spio does not use nvcc and does not need a CUDA installation. Like PyTorch, Spio simply depends on NVIDIA Python packages that provide CUDA libraries. This has already impacted projects like CuPy, which now has an open issue to do the same.

[flecomet]

Thanks for your reply !

What I meant by integrating in Pytorch was for your kernel to be used by default by pytorch, when possible, meaning it would become part of the pytorch package.

This would be really practical to seamlessly obtain performance gains without having to build the package, etc.
But as I understand it, since PyTorch does not enable JIT kernels directly, maybe it's not yet possible...

[andravin]

Will leave this issue open as a reminder to check for any new opportunities for PyTorch integration as they arise. Thanks for creating the issue.