PyTorch Custom Operation

05-10-2026 05-10-2026 blog 23 minutes read (About 3501 words)  visitsIntroductionUsing PyTorch custom operations is common in PyTorch models. PyTorch custom operations can be custom classes and custom functions implemented in C++ and CUDA and used in both Python and C++ inference programs. In this blog post, I would like to share how to implement PyTorch custom operations in C++ and CUDA, and how to use them in PyTorch models and AOTInductor compiled inference programs, using a simple identity convolution example. PyTorch Custom FunctionPyTorch custom functions can be implemented in C++ and CUDA and registered using the TORCH_LIBRARY_IMPL macro. Both the CPU and CUDA implementations can be provided, and PyTorch will dispatch to the correct implementation based on the device of the input tensors. custom_ops.cpp12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152torch::Tensor identity_conv_cpu_impl(const torch::Tensor& input){ TORCH_CHECK(!input.is_cuda(), "identity_conv_cpu_impl: input must be a CPU tensor"); return input.clone();}torch::Tensor identity_conv_cuda_impl(const torch::Tensor& input){ TORCH_CHECK(input.is_cuda(), "identity_conv_cuda_impl: input must be a CUDA tensor");

auto output = torch::empty_like(input); const int64_t numel = input.numel(); if (numel == 0) return output; const int ndim = input.dim(); const auto opts = torch::TensorOptions().dtype(torch::kInt64).device(input.device()); const auto shape_dev = torch::tensor( std::vector<int64_t>(input.sizes().begin(), input.sizes().end()), opts); const auto strides_dev = torch::tensor( std::vector<int64_t>(input.strides().begin(), input.strides().end()), opts); constexpr int kThreads = 256; const int blocks = static_cast<int>((numel + kThreads - 1) / kThreads); AT_DISPATCH_FLOATING_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, input.scalar_type(), "identity_conv_cuda_impl", [&]() { identity_kernel<scalar_t><<<blocks, kThreads>>>( input.data_ptr<scalar_t>(), output.data_ptr<scalar_t>(), shape_dev.data_ptr<int64_t>(), strides_dev.data_ptr<int64_t>(), ndim, numel); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return output;}

custom_op_registration.cpp1234567891011TORCH_LIBRARY_IMPL(my_ops, CUDA, m){ m.impl("identity_conv_op", identity_conv_cuda_impl);}TORCH_LIBRARY_IMPL(my_ops, CPU, m){ m.impl("identity_conv_op", identity_conv_cpu_impl);}

PyTorch Custom ClassPyTorch custom functions are stateless and cannot hold any parameters. If we would like to implement a custom class that holds some parameters and has a forward() method that can be called from Python, we can use torch::CustomClassHolder to define a custom class in C++ and register it with TORCH_LIBRARY macro.

custom_class.cpp12345678910111213141516171819202122232425struct IdentityConvClass : torch::CustomClassHolder{ int64_t channels_; explicit IdentityConvClass(int64_t channels) : channels_(channels) {} torch::Tensor forward(const torch::Tensor& x) { return x.is_cuda() ? identity_conv_cuda_impl(x) : identity_conv_cpu_impl(x); } int64_t get_channels() const { return channels_; }};

custom_class_registration.cpp12345678910111213141516171819202122232425262728293031323334353637TORCH_LIBRARY(my_ops, m){ m.class_<IdentityConvClass>("IdentityConvClass") .def(torch::init<int64_t>()) .def("forward", &IdentityConvClass::forward) .def("get_channels", &IdentityConvClass::get_channels) .def("__obj_flatten__", [](const c10::intrusive_ptr<IdentityConvClass>& self) { return std::make_tuple( std::make_tuple(std::string("channels"), self->channels_)); }) .def_pickle( [](const c10::intrusive_ptr<IdentityConvClass>& self) -> int64_t { return self->channels_; }, [](int64_t channels) -> c10::intrusive_ptr<IdentityConvClass> { return c10::make_intrusive<IdentityConvClass>(channels); }); m.def("identity_conv_op(Tensor x) -> Tensor");}

Using Custom Operations and Classes In PyTorchThe PyTorch custom classes, functions, and their registrations in C++ are built into a shared library (libidentity_conv_ops.so) that can be loaded and registered in PyTorch using torch.ops.load_library.

For torch.compile and torch.export compatibility, we also need to register “fake” (abstract) versions of the custom classes and functions in PyTorch using @register_fake_class and @torch.library.register_fake so that the FakeTensor-based symbolic tracing can work correctly without having to execute the actual C++/CUDA code during tracing. custom_ops.py12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394"""custom_ops.py=============Loads the C++ / CUDA shared library and sets up all custom PyTorch operationsused by the IdentityModel: 1. torch.classes.my_ops.IdentityConvClass (registered by the shared library) - A fake/abstract version is registered here so that torch.export can trace through module attributes that hold an instance of this class. 2. my_ops::identity_conv_op (schema + CPU + CUDA registered by the shared library) - register_fake: abstract implementation for torch.export / FakeTensor.

"""import osimport torchimport torch.library_default_lib = os.path.join( os.path.dirname(os.path.abspath(__file__)), "..", "ext", "libidentity_conv_ops.so")_lib_path = os.path.abspath( os.environ.get("IDENTITY_CONV_OPS_LIB", _default_lib))torch.ops.load_library(_lib_path)from torch._library.fake_class_registry import register_fake_class@register_fake_class("my_ops::IdentityConvClass")class FakeIdentityConvClass: """Abstract counterpart of IdentityConvClass used during torch.export.""" def __init__(self, channels: int) -> None: self.channels_ = channels def __obj_flatten__(self): return (("channels", self.channels_), ) @classmethod def __obj_unflatten__(cls, flat): return cls(dict(flat)["channels"]) def forward(self, x: torch.Tensor) -> torch.Tensor: return torch.empty_like(x) def get_channels(self) -> int: return self.channels_@torch.library.register_fake("my_ops::identity_conv_op")def _identity_conv_op_fake(x: torch.Tensor) -> torch.Tensor: """Abstract implementation used by torch.export / FakeTensor tracing.""" return torch.empty_like(x)identity_conv_op = torch.ops.my_ops.identity_conv_op

PyTorch custom classes can be loaded using torch.classes and PyTorch custom functions can be loaded using torch.ops after the shared library is loaded. model.py1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071

72737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119"""model.py========Defines the four-layer IdentityModel used in the AOTInductor demo.Layer layout------------ layer1 : IdentityConv - native PyTorch operators layer2 : IdentityConvCustomClass - torch.classes C++/CUDA custom class layer3 : IdentityConvCustomOp - torch.library.custom_op C++/CUDA op layer4 : IdentityConv - native PyTorch operatorsEvery layer is an identity transformation, so model(x) == x for any input x."""import torchimport torch.nn as nnfrom custom_ops import identity_conv_op class IdentityConv(nn.Module): """Identity convolution implemented with native PyTorch operators. Uses a depthwise Conv2d with kernel_size=1 and weight=1.0, which is equivalent to a no-op (output == input). This layer is compatible with torch.export and AOTInductor out of the box. """ def __init__(self, channels: int) -> None: super().__init__() self.conv = nn.Conv2d( in_channels=channels, out_channels=channels, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=channels, bias=False, ) self.conv.weight.data = torch.ones(channels, 1, 1, 1) self.conv.weight.requires_grad = False def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x)class IdentityConvCustomClass(nn.Module): """Identity convolution backed by a torch.classes C++/CUDA custom class.

At runtime the forward call is dispatched to the CUDA kernel registered inside IdentityConvClass (csrc/identity_conv.cpp + .cu). For torch.export compatibility a FakeIdentityConvClass is registered in custom_ops.py via @register_fake_class so that symbolic tracing works. """ def __init__(self, channels: int) -> None: super().__init__() self.obj = torch.classes.my_ops.IdentityConvClass(channels) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.obj.forward(x)class IdentityConvCustomOp(nn.Module): """Identity convolution backed by a torch.library.custom_op C++/CUDA op. The op (my_ops::identity_conv_op) is defined in custom_ops.py with: • a register_fake implementation for torch.export tracing • a register_kernel("cuda") implementation that calls the CUDA kernel """ def __init__(self, channels: int) -> None: super().__init__() self.channels = channels def forward(self, x: torch.Tensor) -> torch.Tensor: return identity_conv_op(x)class IdentityModel(nn.Module): """Four-layer identity model for AOTInductor demo.""" def __init__(self, channels: int) -> None: super().__init__() self.layer1 = IdentityConv(channels) self.layer2 = IdentityConvCustomClass(channels) self.layer3 = IdentityConvCustomOp(channels) self.layer4 = IdentityConv(channels) def forward(self, x: torch.