Repository: pytorch
Issue: [docs] Idea collection of examples of custom ops / inline torch extensions
Labels: module: docs,triaged,module: custom-operators

### 🚀 The feature, motivation and pitch

Following up on https://github.com/pytorch/pytorch/issues/68066#issuecomment-1611836317:

<hr>

- a complete example of a custom op (maybe with forward/backward and registration) implemented via load_inline with only CUDA op or both CPU/SSE (how to properly do SSE impls?) and CUDA code - maybe for bilateral filter or batched base64 encoding/decoding or some bit magics like popcnt (https://github.com/pytorch/ao/issues/292 or groupbitsum from https://github.com/Felix-Petersen/difflogic)

- maybe general discussion of whether `cupy` api/features for this are worth being compatible with? pure nvrtc compilation? pure C compilation / marshalling (of pure pointers + numbers + C strings + ctype structs?)?

- support for launcher-less kernels via `cupy`-like api? (also a good fit for nvrtc) for instance, here https://github.com/pytorch/pytorch/blob/main/test/test_cpp_extensions_jit.py#L338 there's even an unintuitive need for a c++ source for function declaration. it would be nicer to do without and be able to directly call the kernel

- example of calling a C function consuming raw data pointers or DLPack structures. It often is much simpler to have these kinds of interfaces for thin C wrappers (e.g. audio decoding). It would be nice to have this example as inline, as often it's just dozen lines of C++ code in extern C function to call some underlying library like BLAS / cuBLAS / some media encoding-decoding.

- example of using extra (currently unbound) functions cuBLAS / cudnn / blas. Maybe an example here could be some code using `ctypes`? it should be explained how to correctly marshal CUDA pointers.

<hr>

More ideas from https://github.com/pytorch/pytorch/pull/105947#issuecomment-1655693163:


- also interesting would be examples on how to register ops currently missing in pytorch, e.g. adding a quantized impl or some missing sparse op, so that dispatcher would use it and let go the exception that some op impl can't be found

- also interesting an example of adding a custom backward impl via a decorator

- also interesting an example of cuda op, but without having to manually add a C++ launcher stub (e.g. so that pytorch would generate the kernel launcher and maybe that the kernel can be compiled by nvrtc - like in CuPy)

- and also the CuPy-like API could be worth a discussion

- also interesting would be an example of adding a pure C-op passing/consuming a dlpack tensor or even pure pointers - useful for interaction with C libraries of image/audio encoding/decoding

<hr>

Trying to distill from above concrete ideas on improving FFI story of PyTorch:
- Fixing interop issues: https://github.com/pytorch/pytorch/issues/34651 https://github.com/pytorch/pytorch/issues/69491
- CuPy-like features:
  - API like [RawKernel](https://docs.cupy.dev/en/stable/reference/generated/cupy.RawKernel.html) [RawModule](https://docs.cupy.dev/en/stable/reference/generated/cupy.RawModule.html)
  - `nvrtc`-compilation (supposed to be faster than dealing with full nvcc / working with disk)
  - kernel launch functionality from Python directly (like CuPy)
- Examples of interop with pure C functions:
  - consuming pure pointers
  - consuming DLPack tensors
  - Examples of using libopus loop-less API: https://github.com/jlaine/opuslib/blob/master/opuslib/api/decoder.py https://github.com/orion-labs/opuslib
  - Examples of interop PyTorch<> C code using ffmpeg libraries
  - Examples of interaction of interop of PyTorch tensors with ctypes-functions (e.g. how to properly deal with tensor memory ownership)
 - Examples of C++ extensions:
   - Examples of interop of PyTorch with opencv via C++ code snippets interacting with libopencv-dev
   - Examples of using directly CUDA kernel files from [FasterTransformer](https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/kernels/layernorm_kernels.cu ) or [OneFlow](https://github.com/Oneflow-Inc/oneflow/blob/2d24fe08be1b1bedcc22fb409c5d688924ce89fc/oneflow/user/kernels/layer_norm_gpu_kernel.cu). Somehow getting a compiled kernel without fully installing a python Package gives a nice feeling of experimentation and flexibility :)
 

cc @svekars @carljparker