Adaptation Development and Usage (pybind-based)

This document describes the complete adaptation development workflow for calling custom operators through the TorchNPU framework by using C++ extensions and single-operator APIs. The workflow covers operator definition, operator adaptation, and ATen IR registration and binding, enabling calls to custom operators. Unlike the common TORCH_LIBRARY method, this sample uses pybind for binding and registration to achieve more flexible input type support.

Operator Adaptation Development

Prerequisites

Before getting started, ensure that you have completed the installation of the following environments:

Install the NPU driver, firmware, and CANN software (Toolkit, ops, and NNAL) by referring to CANN Software Installation.
Install the PyTorch framework by referring to Installation Guide.

Adaptation File Structure

├── build_and_run.sh                // Script for compiling and installing the custom operator wheel package and executing the test case
├── csrc                            // Directory of the C++ code at the operator adaptation layer
│   └── add_custom.cpp              // Forward and backward adaptation code, ATen IR registration, and binding for the custom operator
├── cpp_extension_pybind            // Python-side code for the custom operator package
│   ├── __init__.py                 // Python initialization file
│   └── ops                         // Defines operator APIs
│       └── __init__.py             // Python initialization file
├── setup.py                        // Compilation file of the wheel package
└── test                            // Directory for test cases
    └── test_add_custom.py          // Script for executing operator test cases in eager mode

Procedure

Implement the C++ operator code, adaptation layer, custom operator schema registration, and implementation binding in the add_custom.cpp file under the csrc directory. The code sample is as follows:

The following is the code sample for single-device scenarios. In these scenarios, the const c10::OptionalDeviceGuard device_guard(device_of(self)); statement is optional. In multi-device scenarios, this statement must be included in the adaptation code.

// Register forward implementation for NPU devices
at::Tensor add_custom_impl_npu(const at::Tensor& self, const at::Tensor& other)
{
    const c10::OptionalDeviceGuard device_guard(device_of(self));
    // Allocate output memory
    at::Tensor result = at::empty_like(self);

    at::Scalar alpha = 1.0;

    // Call the ACLNN API for computation
    EXEC_NPU_CMD_EXT(aclnnAdd, self, other, alpha, result);
    return result;
}

PYBIND11_MODULE(custom_ops_lib, m)
{
    m.def("add_custom", &add_custom_impl_npu, "");
}

Add the operator call logic and load the .so file in the __init__.py and ops/__init__.py files under the cpp_extension_base directory.

# __init__.py
__all__ = ['ops', 'add_custom']
from .ops import add_custom

# ops/__init__.py
__all__ = ["add_custom"]
from cpp_extension_pybind.custom_ops_lib import add_custom

Usage Example

After completing the operator adaptation development, you can call the custom operator through C++ extensions.

Complete the creation, development, compilation, and deployment workflow for the custom operator project. For details, see CANN Ascend C Operator Development.

Download the code sample.

# Download the code sample
git clone https://gitcode.com/Ascend/op-plugin
# Go to the code directory
cd examples/cpp_extension_pybind

Complete operator adaptation. For details, see Operator Adaptation Development.
Run the following command to compile, install, and execute the test script:
```
bash build_and_run.sh
```
The following output indicates successful execution:
```
Ran xx tests in xx s
OK
```