PyTorch

When PyTorch is used for model training and inference, many operators are called for computation. To integrate and deploy your custom operators in PyTorch, the following methods are available:

Kernel launch: By registering custom operators through torch.library or Pybind, you can enable PyTorch to call the operator kernel program.
Single-operator API call: For details about the adaptation plugin development process and examples in this mode, see "OpPlugin-based Operator Adaptation" in Ascend Extension for PyTorch Framework Feature Guide.
Graph mode call: For details about the adaptation development guide for custom operators in PyTorch graph mode, see "Integrating Custom Operators to a Graph" in PyTorch Graph Mode User Guide (TorchAir).

Figure 1 Operator deployment in PyTorch

This section describes how to register a custom operator through torch.library and Pybind and enable PyTorch to call the operator kernel program.

torch.library provides a set of APIs for extending PyTorch's core operator library. It allows you to create new operators and provide custom implementations for them.
Pybind is an open-source tool that bridges C++ and Python, enabling seamless integration of C++ code into the Python environment.

Pybind is used to quickly expose C++ functions to Python for efficient API binding. However, the generated operators cannot be identified by the PyTorch operator system and do not support schema definition or graph tracing capability. Therefore, torch.compile is not supported. In contrast, torch.library provides a mechanism for in-depth integration with PyTorch's core operator system and supports operator registration, schema definition, and graph tracing, which are prerequisites for supporting torch.compile. You can choose the method that best fits your needs.

torch.library

The following code uses the add_custom operator as an example to describe how to call the operator kernel program through torch.library. Only the core steps are described here. For details about the complete sample, see torch.library sample.

Set up the environment.
In addition to installing the CANN software package by referring to Environment Setup, you also need to install the following dependencies:
- Installing PyTorch

Implement the custom operator on the NPU.

Implement the operator on the kernel and use <<<>>> to call the operator's kernel function to complete the specified computation. In the sample, the c10_npu::getCurrentNPUStream API is called to obtain the current NPU stream. The return type is NPUStream. For details about how to use this API, see "(beta) c10_npu::getCurrentNPUStream" in Ascend Extension for PyTorch Custom API Reference.

Note that the memory for inputs x and y in this sample is allocated in the outer Python call script.

namespace ascendc_ops {
at::Tensor ascendc_add(const at::Tensor& x, const at::Tensor& y)
{
    // Allocate runtime resources and obtain the current NPU stream by calling c10_npu::getCurrentNPUStream().
    auto aclStream = c10_npu::getCurrentNPUStream().stream(false);
    // Allocate the output buffer on the device.
    at::Tensor z = at::empty_like(x);
    uint32_t numBlocks = 8;
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }
    // Call the kernel function using <<<>>> to complete the specified computation.
    add_custom<<<numBlocks, nullptr, aclStream>>>((uint8_t*)(x.mutable_data_ptr()), (uint8_t*)(y.mutable_data_ptr()), (uint8_t*)(z.mutable_data_ptr()), totalLength);
    // Copy the computation result from the device to the host and free the allocated resources.
    return z;
}
} // namespace ascendc_ops

PyTorch provides the TORCH_LIBRARY macro as the core API for registering custom operators. This macro is used to create and initialize a custom operator library. After registration, the custom operator can be called in Python using torch.ops.namespace.op_name. TORCH_LIBRARY_IMPL is used to bind the operator logic to a specific DispatchKey (PyTorch's device scheduling identifier). For NPU devices, the operator implementation needs to be registered with the DispatchKey PrivateUse1.

// Register the operator with torch.library.
TORCH_LIBRARY(ascendc_ops, m)
{
    m.def("ascendc_add(Tensor x, Tensor y) -> Tensor");
}

// Register the implementation with PrivateUse1 on the NPU device.
TORCH_LIBRARY_IMPL(ascendc_ops, PrivateUse1, m)
{
    m.impl("ascendc_add", TORCH_FN(ascendc_ops::ascendc_add));
}

Compile and generate a dynamic operator library.
Use the Python test script to perform the test.
In add_custom_test.py, the generated custom operator library is loaded using torch.ops.load_library, the registered ascendc_add function is called, and the output on the NPU is compared with the CPU's standard addition result to verify the numerical correctness of the custom operator.

Pybind

The following code uses the add_custom operator as an example to describe how to call a custom operator in a PyTorch script using Pybind. Only the core steps are described here. For details about the complete sample, see Pybind sample.

Set up the environment.
In addition to installing the CANN software package by referring to Environment Setup, you also need to install the following dependencies:
- Installing PyTorch
- Install Pybind11.
```
pip3 install pybind11 expecttest
```

Implement the custom operator on the NPU.

Note that the memory for inputs x and y in this sample is allocated in the Python call script.

// Header files required for calling Pybind and PyTorch
#include <pybind11/pybind11.h>
#include <torch/extension.h>

#include "torch_npu/csrc/core/npu/NPUStream.h"
// Header file required for implementation on the kernel
#include "kernel_operator.h" 
...
namespace ascendc_ops {
at::Tensor ascendc_add(const at::Tensor& x, const at::Tensor& y)
{
    // Allocate runtime resources and obtain the current NPU stream by calling c10_npu::getCurrentNPUStream().
    auto aclStream = c10_npu::getCurrentNPUStream().stream(false);
    // Allocate the output buffer on the device.
    at::Tensor z = at::empty_like(x);
    uint32_t numBlocks = 8;
    uint32_t totalLength = 1;
    for (uint32_t size : x.sizes()) {
        totalLength *= size;
    }
    // Call the kernel function using <<<>>> to complete the specified computation.
    add_custom<<<numBlocks, nullptr, aclStream>>>((uint8_t*)(x.mutable_data_ptr()), (uint8_t*)(y.mutable_data_ptr()), (uint8_t*)(z.mutable_data_ptr()), totalLength);
    // Copy the computation result from the device to the host and free the allocated resources.
    return z;
}
} // namespace ascendc_ops

Define the Pybind module to encapsulate C++ functions into Python functions. PYBIND11_MODULE is a macro in the Pybind11 library and is used to define a Python module. It takes two parameters. The first parameter is the encapsulated module name, and the second parameter is a Pybind11 module object, which is used to define functions, classes, and constants in the module. By calling the m.def() method, you can convert the ascendc_ops::ascendc_add function from the previous step into the Python function ascendc_add, so that it can be called in Python code.

PYBIND11_MODULE(ascendc_ops, m)// ascendc_ops: module name; m: module object.
{
    m.doc() = "add_custom pybind11 interfaces";// optional module docstring
    m.def("ascendc_add", &ascendc_ops::ascendc_add, "");// Bind the ascendc_add function to the Pybind module.
}

Compile and generate a dynamic operator library.
In the Python call script, use the Torch API to generate random input data and allocate the memory, import the encapsulated custom module ascendc_ops, call the run_add_custom function in ascendc_ops, and execute operators on the NPU.

Parent topic: AI Framework Operator Adaptation