custom operator
**Custom operators** in PyTorch enable **extending the framework with specialized operations** — implementing functionality not available in standard libraries, optimizing performance-critical code with CUDA kernels, or integrating external libraries for domain-specific needs.
**What Are Custom Operators?**
- **Definition**: User-defined operations extending PyTorch.
- **Use Cases**: Missing ops, CUDA optimization, library integration.
- **Levels**: Python functions, C++ extensions, CUDA kernels.
- **Integration**: Works with autograd, torch.compile, export.
**Why Custom Operators**
- **Performance**: Fused operations, CUDA optimization.
- **Functionality**: Operations not in standard PyTorch.
- **Integration**: Connect external C++/CUDA libraries.
- **Research**: Implement novel operations.
**Custom Op Levels**
**Complexity Spectrum**:
```
Level | Performance | Complexity | Use Case
----------------|-------------|------------|------------------
Python function | Low | Easy | Prototyping
torch.autograd | Medium | Easy | Custom backward
C++ extension | High | Medium | CPU optimization
CUDA extension | Highest | Hard | GPU optimization
Triton kernel | High | Medium | GPU, Python-like
```
**Python Custom Function**
**With Custom Backward**:
```python
import torch
from torch.autograd import Function
class MyReLU(Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input.clamp(min=0)
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
grad_input = grad_output.clone()
grad_input[input < 0] = 0
return grad_input
# Usage
my_relu = MyReLU.apply
output = my_relu(input_tensor)
```
**C++ Extension**
**Setup** (setup.py):
```python
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension
setup(
name="my_ops",
ext_modules=[
CppExtension(
"my_ops",
["my_ops.cpp"],
),
],
cmdclass={"build_ext": BuildExtension},
)
```
**C++ Implementation** (my_ops.cpp):
```cpp
#include
torch::Tensor my_add(torch::Tensor a, torch::Tensor b) {
TORCH_CHECK(a.sizes() == b.sizes(), "Size mismatch");
return a + b; // Simple example
}
torch::Tensor fused_gelu(torch::Tensor x) {
// Fused GELU: x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
auto x3 = x * x * x;
auto inner = 0.79788456 * (x + 0.044715 * x3);
return x * 0.5 * (1.0 + torch::tanh(inner));
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("my_add", &my_add, "Element-wise addition");
m.def("fused_gelu", &fused_gelu, "Fused GELU activation");
}
```
**Usage**:
```python
import torch
import my_ops
x = torch.randn(1000, 1000)
y = my_ops.fused_gelu(x)
```
**CUDA Extension**
**CUDA Kernel** (my_ops_cuda.cu):
```cuda
#include
#include
#include
template
__global__ void fused_gelu_kernel(
const scalar_t* __restrict__ input,
scalar_t* __restrict__ output,
size_t size
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < size) {
scalar_t x = input[idx];
scalar_t x3 = x * x * x;
scalar_t inner = 0.79788456f * (x + 0.044715f * x3);
output[idx] = x * 0.5f * (1.0f + tanhf(inner));
}
}
torch::Tensor fused_gelu_cuda(torch::Tensor input) {
auto output = torch::empty_like(input);
const int threads = 256;
const int blocks = (input.numel() + threads - 1) / threads;
AT_DISPATCH_FLOATING_TYPES(input.type(), "fused_gelu", ([&] {
fused_gelu_kernel<<>>(
input.data_ptr(),
output.data_ptr(),
input.numel()
);
}));
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("fused_gelu", &fused_gelu_cuda, "Fused GELU (CUDA)");
}
```
**Triton Alternative**
**Easier GPU Kernels**:
```python
import triton
import triton.language as tl
import torch
@triton.jit
def gelu_kernel(x_ptr, output_ptr, n_elements, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offsets = pid * BLOCK + tl.arange(0, BLOCK)
mask = offsets < n_elements
x = tl.load(x_ptr + offsets, mask=mask)
# GELU computation
x3 = x * x * x
inner = 0.79788456 * (x + 0.044715 * x3)
output = x * 0.5 * (1.0 + tl.libdevice.tanh(inner))
tl.store(output_ptr + offsets, output, mask=mask)
def fused_gelu_triton(x):
output = torch.empty_like(x)
n = x.numel()
gelu_kernel[(n // 1024 + 1,)](x, output, n, BLOCK=1024)
return output
```
**Registering for torch.compile**
```python
import torch
from torch.library import Library
# Define custom library
my_lib = Library("myops", "DEF")
# Register schema
my_lib.define("fused_gelu(Tensor x) -> Tensor")
# Register implementation
@torch.library.impl(my_lib, "fused_gelu", "CUDA")
def fused_gelu_impl(x):
return fused_gelu_cuda(x)
# Now works with torch.compile
@torch.compile
def model(x):
return torch.ops.myops.fused_gelu(x)
```
Custom operators are **essential for pushing PyTorch performance boundaries** — when standard operations aren't sufficient, custom ops enable the optimizations and integrations that production ML systems require.