使用Pytorch的cpp_extension调用外部C函数

Python中调用C/C++函数的方式

常见的在Python内调用C/C++函数的方式有：

• 直接使用ctypes；
• 使用Pybind11/Boost.Python等Package包装交互代码后完成调用；

一般来说，如果需要在Python中调用cpp，需要了解并熟悉c/c++的编译、链接以及动态库导出等步骤。在通过Cmake+编译器将cpp项目编译成shared library(.so/.dll文件)后，在Python内使用ctypes便可以进行调用。以此种方式进行调用需要手动进行c文件编译，并需要手动处理python与ctype之间的数据格式转换、指针问题等。

诸如Numpy这类package的运作方式便是使用大量ctypes从编译后的动态库中调用函数(这也是其whl需要指明平台架构与机器字长的原因之一，不同平台不同机器字长的系统会编译出不同的动态库)。

当然，手动编译、手动处理数据格式转换、手动包装指针与内存分配，意味着对于现代cpp类型(数组、vector、STL容器、面向对象编程设计)，需要有大量样本代码去处理其在python与c之间的数据传递，使得工作量大幅增加。

Pybind11是近几年比较热门的一个Python Package（用于取代Boost.Python），目的在于在C和Python之间建立起数据类型转换的桥梁，并基于此来达到两种语言相互调用的目的。使用Pybind11可以显著减少工作量。不过用户依旧需要自己使用cmake这类工具去处理编译问题。

如果我们在使用Pytorch的过程中，需要调用外部cpp函数(甚至是cuda代码)，当然也可以使用上述的ctypes与Pybind11这两种方式，但是由于Pytorch有自己的Tensor数据格式与及其c实现，我们在调用自定义kenrel的时候也免不了使用Tensor，如果使用上述原生方式不仅需要手动处理编译还需要去链接Torch的动态库以访问Tensor等数据结构，显然这非常的不现实。

**因此Pytorch为程序员提供了自己调用C/C++外部拓展的方式即torch.utils.cpp_extension。**简单来说，该module有两种运作模式，分别为：

• JIT(Just-in-time)模式：在Python代码运行时直接加载C/C++外部拓展，一般来说在首次加载的时候会将cpp/c/cuda源文件编译为动态库随后再被当前的Python进程调用；
• Setuptool模式：配合Python的setuptools外部库配合使用，并在构建Python Package时按照需求调用编译器对c/cpp/cu源文件进行编译并得到对应的动态链接库。

torch.utils.cpp_extension会自动链接libtorch等torch相关动态链接库，在简单引用torch相关头文件后便可以在c/cpp内使用Tensor数据结构，无需控制内存分配、指针传递等问题。

什么是setuptools

使用过Python外部Package的大家都知道，安装一个三方库的指令为:

pip install somepackages

pip会帮你处理该包的依赖，下载对应平台、对应Py版本、对应包版本、对应操作系统的whl文件并安装。但有时候我们可能会使用源码安装一个python包，当然一样也可以使用pip进行安装，不同的是我们可以看到一个三方包的源码及其文件编排格式。我们大概率会在源码根目录发现一个setup.py的文件，有时候还会伴随有setup.cfg与pyproject.toml文件。这些都是告诉Python及其toolkit如何正确编译及安装该package的文件。我们一般安装好Python后使用venv(或者conda)，默认的venv内使用pip list指令会发现里面有个安装的三方包叫setuptools：

$ pip list

Package Version
-------     -------
pip         25.2
setuptools  80.9.0

这里的setuptools是一个第三方包，专门用于处理那些使用源码安装的包的安装过程，其主要功能是为这些三方包提供非纯Python(non-pure-python)源码安装支持。当一个三方包的源码内包含c或cpp文件，并且其会作为Cpython的extension工作，我们就需要使用setuotools来尝试对这些非python code进行安装时编译。

显然，配合setuptools使用的目的在于提前对source code进行编译，以便隐藏源码并在后续运行时直接调用。当然在调试或开发时，使用JIT模式是个更方便的选择。

在JIT模式下调用C++函数并进行调试

先创建一个简单的cuda kernel以及其C++与对应的Tensor wrapper:

//File name: kernel.cu

#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>

template <typename scalar_t>
__global__ void custom_op_kernel(const scalar_t* __restrict__ input,
                                 scalar_t* __restrict__ output,
                                 const int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        output[idx] = input[idx] * input[idx]; // Example operation: squaring the input
    }
}


torch::Tensor custom_op_forward(torch::Tensor input) {
    const auto size = input.numel();
    auto output = torch::empty_like(input);

    const int threads = 256;
    const int blocks = (size + threads - 1) / threads;

    AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "custom_op_forward", ([&] {
        custom_op_kernel<scalar_t><<<blocks, threads>>>(
            input.data_ptr<scalar_t>(),
            output.data_ptr<scalar_t>(),
            size);
    }));

    return output;
}

并用一个cpp文件来暴露其API给Torch，注意其中的PYBIND11_MODULE宏的使用，必须使用该宏函数才能令可调用的函数暴露给Python：

//File name: kernel.cpp
#include <torch/extension.h>

torch::Tensor custom_op_forward(torch::Tensor input);

// Macro privided by torch. Expose cpp function to Python.
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("custom_op_forward", &custom_op_forward, "Custom CUDA Op forward");
}

随后在Python文件内使用torch.utils.cpp_extension.load使用JIT模式加载该cpp函数，调用过上述宏的函数便可以直接从module内load出：

"""
File name: custom_op_load.py
"""

import torch
import os
from torch.utils.cpp_extension import load

# get the shared library module 
kernel_module = load(
    name="custom_ops",
    sources=["kernel.cpp", "kernel.cu"],
    verbose=True
)

input_tensor = torch.randn(10).cuda()
print("Input Tensor with shape:", input_tensor.shape)
print("Input Tensor values:", input_tensor)

# call the exposed api(using macro) from compiled shared library
output_tensor = kernel_module.custom_op_forward(input_tensor)
print("Output Tensor with shape:", output_tensor.shape)
print("Output Tensor values:", output_tensor)

文件目录为:

$ tree ./

./
├── custom_op_load.py
├── kernel.cpp
└── kernel.cu

1 directory, 3 files

随后我们便可以直接编译运行该Python文件，注意Python环境需要包含torch库，我这里使用了conda环境：

(my_env)$ python ./custom_op_load.py

正常情况下，我们会发现程序输出了正确的结果，即对一个一维Tensor内的元素做了平方。同时也可以看到torch输出了一些调试信息，告诉我们其是如何对cpp代码进行编译的：

[1/3] c++ -MMD -MF kernel.o.d -DTORCH_EXTENSION_NAME=custom_ops -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/include -isystem /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /opt/cuda/include -isystem /home/blackcat/program/anaconda3/envs/modelopt/include/python3.12 -fPIC -std=c++17 -c /home/blackcat/workspace/python/resnet/custom_ops/kernel.cpp -o kernel.o

[2/3] /opt/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output kernel.cuda.o.d -DTORCH_EXTENSION_NAME=custom_ops -DTORCH_API_INCLUDE_EXTENSION_H -isystem /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/include -isystem /home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /opt/cuda/include -isystem /home/blackcat/program/anaconda3/envs/modelopt/include/python3.12 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_75,code=compute_75 -gencode=arch=compute_75,code=sm_75 --compiler-options '-fPIC' -std=c++17 -c /home/blackcat/workspace/python/resnet/custom_ops/kernel.cu -o kernel.cuda.o 

[3/3] c++ kernel.o kernel.cuda.o -shared -L/home/blackcat/program/anaconda3/envs/modelopt/lib/python3.12/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/opt/cuda/lib64 -lcudart -o custom_ops.so

显然，前两步是调用c++(g++)与nvcc编译，第三步是链接编译结果得到动态链接库(.so文件)。但我们并没有在代码目录下发现编译结果与链接结果，这是因为torch会将编译结果放在cache目录下，如果我们在python代码内的load函数前添加一行：

os.environ["TORCH_EXTENSIONS_DIR"] = os.path.join(os.getcwd(), "jit_build")

便可以让编译结果输出在当前work_dir，最后我们得到的目录结构如下：

./
├── custom_op_load.py
├── jit_build
│   └── custom_ops
│       ├── build.ninja
│       ├── custom_ops.so
│       ├── kernel.cuda.o
│       └── kernel.o
├── kernel.cpp
└── kernel.cu

3 directories, 7 files

使用ldd指令去列出编译得到的custom_op.so链接的其他动态库，得到：

$ ldd ./jit_build/custom_ops/custom_ops.so
    linux-vdso.so.1 (0x00007f2a53a7a000)
    libc10.so => not found
    libc10_cuda.so => not found
    libtorch_cpu.so => not found
    libtorch_cuda.so => not found
    libtorch.so => not found
    libtorch_python.so => not found
    libcudart.so.13 => /opt/cuda/lib64/libcudart.so.13 (0x00007f2a53600000)
    libstdc++.so.6 => /usr/lib/libstdc++.so.6 (0x00007f2a53200000)
    libm.so.6 => /usr/lib/libm.so.6 (0x00007f2a534f2000)
    libgcc_s.so.1 => /usr/lib/libgcc_s.so.1 (0x00007f2a5396a000)
    libc.so.6 => /usr/lib/libc.so.6 (0x00007f2a52e00000)
    /usr/lib64/ld-linux-x86-64.so.2 (0x00007f2a53a7c000)
    libdl.so.2 => /usr/lib/libdl.so.2 (0x00007f2a53963000)
    libpthread.so.0 => /usr/lib/libpthread.so.0 (0x00007f2a5395e000)
    librt.so.1 => /usr/lib/librt.so.1 (0x00007f2a53959000)

可见，使用torch.utils.cpp_extension会自动链接libtorch等torch相关动态库。其实不论我们是否在cpp内使用了<torch/extension.h>，torch.utils.cpp_extension都会为我们链接这些动态库，因此假如我们的目标只是在Python内调用C++函数(与Pytroch无关)，最好还是使用Pybind11或ctypes。

总结

使用torch.utils.cpp_extension在Python调用C/C++/Cuda拓展的方式有用两种：

• 使用load等api的JIT模式；
• 配合setuptools使用的构建Package模式；

JIT模式更方便我们对kernel进行调试，每次修改后运行都会重新进行编译。配合setuptools更适合分发代码，可以一次性构建编译代码以便后续代码的快速运行，省去了重复编译时间。