VScode 里面使用 python 去直接调用 CUDA

上一个帖子主要分享了如何去将 C++ 程序打包成一个package。我们最后的目的实际上是想把 CUDA 的程序打包成一个 Package ， C++ 程序只是起到了桥梁的作用：

首先：CUDA 程序和 C++ 的程序一样，都有一个 .cu 的源文件和一个 .h 的头文件。

我们的文件包含 Cpp 文件组成，负责当作 CUDA 和 Python 的桥梁。还有对应的 CUDA 的源代码文件和头文件。将这个cpp 文件命名成 ext.cpp.

#include <torch/extension.h>
#include "interpolation_kernel.h"  ## CUDA 的头文件PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){m.def("trilinear_interpolation",&trilinear_interpolation);
}

cpp_properities.json 配置文件

{"configurations": [{"name": "Linux","includePath": ["${workspaceFolder}/**","/home/smiao/anaconda3/envs/Gen_3DGS/lib/python3.8","/home/smiao/anaconda3/envs/Gen_3DGS/lib/python3.8/site-packages/torch/include/","/home/smiao/anaconda3/envs/Gen_3DGS/lib/python3.8/site-packages/torch/include/torch/csrc/api/include/"],"defines": [],"compilerPath": "/usr/bin/gcc","cStandard": "c17","cppStandard": "gnu++14","intelliSenseMode": "linux-gcc-x64"}],"version": 4

CUDA 部分：

CUDA 的头文件 *** interpolation_kernel.h ***

#include <torch/extension.h>#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x "must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)torch::Tensor trilinear_interpolation(torch::Tensor feats, torch::Tensor point);

对应的源代码文件*** interpolation_kernel.cu ***

include 的头文件和源代码文件尽量放在同一级的目录

#include <torch/extension.h>
#include "interpolation_kernel.h"torch::Tensor trilinear_interpolation(torch::Tensor feats, torch::Tensor points){CHECK_CUDA(feats);CHECK_CUDA(points);return feats;
}

配置文件 setup.py 部分：

配置文件的包含 ** *.cpp 文件和 *.cu 文件 **
其他的部分应该尽量不去改变。

from setuptools import setup
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
import os
import globos.path.dirname(os.path.abspath(__file__))setup(name="cuda_tutorial",version='1.0',ext_modules=[CUDAExtension(name='cuda_tutorial',sources=["interpolation_kernel.cu","ext.cpp"], extra_compile_args={'cxx': ['-O2'],'nvcc': ['-O2']})],cmdclass={'build_ext': BuildExtension}
)