This folder is an example of how to integrate your own custom kernels into ao such that
- They work on as many devices and operating systems as possible
- They compose with
torch.compile()
without graph breaks
The goal is that you can focus on just writing your custom CUDA or C++ kernel and we can package it up so it's available via torchao.ops.your_custom_kernel
.
We've integrated a test kernel which implements a non-maximum supression (NMS) op which you can use as a template for your own kernels.
- Install the cudatoolkit https://anaconda.org/conda-forge/cudatoolkit
- In
csrc/cuda
author your custom kernel and ensure you expose aTORCH_LIBRARY_IMPL
which will exposetorchao::your_custom_kernel
- In
csrc/
author acpp
stub which will include aTORCH_LIBRARY_FRAGMENT
which will place your custom kernel in thetorchao.ops
namespace and also expose a public function with the right arguments - In
torchao/ops.py
is where you'll expose the python API which your new end users will leverage - Write a new test in
test/test_ops.py
which most importantly needs to passopcheck()
, this ensures that your custom kernel composes out of the box withtorch.compile()
And that's it! Once CI passes and your code merged you'll be able to point people to torchao.ops.your_custom_kernel
. If you're working on an interesting kernel and would like someone else to handle the release and package management please feel free to open an issue.
If you'd like to learn more please check out torch.library
The important dependencies are already taken care of in our CI so feel free to test in CI directly
- cudatoolkit so you can build your own custom extensions locally. We highly recommend using https://anaconda.org/conda-forge/cudatoolkit for installation
- manylinux with CUDA support. In your own Github actions you can integrate this support using
uses: pytorch/test-infra/.github/workflows/linux_job.yml@main