python调用底层c++算子示例

test.cpp

#include <torch/extension.h>

// 定义一个简单的加法函数
at::Tensor add(at::Tensor a, at::Tensor b) {
    return a + b;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("add", &add, "Add two tensors");
}

test.py

import torch
from torch.utils.cpp_extension import load

test_load = load(
    name='test_load', 
    sources=['test.cpp'],
    extra_cflags=['-O0', '-g'], #debug
    #extra_cflags=['-O1'], #no debug
    verbose=True,
)
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = test_load.add(a, b)
print(result)  # Should print tensor([5, 7, 9])

注意：load中name实际上传给test.cpp中TORCH_EXTENSION_NAME， 会编译生成一个test_load.so 动态库，首次运行会编译这个库，第二次运行直接不用编译，从缓存里调用库

test2.py:调用test.py中test_load模块

from test import test_load
import torch
a = torch.tensor([1, 2, 3])
b = torch.tensor([4, 5, 6])
result = test_load.add(a, b)
print(result)  # Should print tensor([5, 7, 9])

如何调试python种嵌入代码呢
launch.json

{
    // Use IntelliSense to learn about possible attributes.
    // Hover to view descriptions of existing attributes.
    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
    "version": "0.2.0",
    "configurations": [
        {
            "name": "Python test_load",
            //"type": "debugpy",
            "type": "cppdbg",
            "request": "launch",
            "program": "/root/miniconda3/bin/python",
            "console": "integratedTerminal",
            "args": ["test.py"],
            "cwd": "${workspaceFolder}",
            "MIMode": "gdb",
            "env": {
                "PATH": "/root/miniconda3/bin/:$PATH",
                "LD_LIBRARY_PATH": "$PWD:extensions/speech/lib/:$LD_LIBRARY_PATH",
                "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7"
            }
        }
    ]
}

同时需要修改extra_cflags的编译方式

#extra_cflags=['-O0', '-g'],

运行文件带有 with debug_info表示可以调试的版本
不带debug_info是不可调试版本


带调试信息的就可以在test.cpp中设置断点进行调试

如果是多进程调用代码，需要改成函数调用，在cpp文件注释才会跳转，否则debug时候不会跳转到注释，如

    data_loader = mp.Process(target=gen_data_stream,
                             daemon = True,
                             args=(world_size,
                                   data_stream,
                                   data_conf,
                                   10))

需要把gen_data_stream外衣玻璃，直接运行里面函数调用才能跳转到cpp