CUDA TensorRT Python智能提示补全解决方案

最近在使用python 测试一些 tensorrt 和 cuda 特性，发现没有提示和补全，于是尝试解决

解决方案

解决方案来自 https://github.com/NVIDIA/TensorRT/issues/1714

pybind11-stubgen --ignore-all-errors tensorrt_bindings

这个是针对特定版本 8.6.1 才有的，由NVIDIA提供

在pypi上NVIDIA还提供了一些 bindings 应该是用来支持更新版本的TensorRT

安装 pybind11-stubgen 和 tensorrt_bindings

pip install pybind11-stubgen tensorrt_bindings

然后执行

pybind11-stubgen --ignore-all-errors tensorrt_bindings

会在执行路径下生成一个 stubs/tensorrt_bindings/tensorrt.pyi

我测试过，其实不安装bingdings包也可以，这个工具是基于提供的 动态库 tenssort.so 来生成类型文件，直接执行pybind11-stubgen --ignore-all-errors tensorrt 即可 ，关键是底层动态库

有问题，欢迎大家留言、进群讨论或私聊：【群号：392784757】

然后将这个文件路径加入到 settings.json中（拿vscode举例）

{
  "python.defaultInterpreterPath": "yourpath/envs/llm-quant/bin/python",
  "python.analysis.extraPaths": [
    "yourpath/stubs",
  ]
}

效果

未生成前

生成后

同样的针对 pycuda 和 cuda-python 也是如此操作

题外话：pycuda 和 cuda-python 区别

Both CUDA-Python and pyCUDA allow you to write GPU kernels using CUDA C++. The kernel is presented as a string to the python code to compile and run.

The key difference is that the host-side code in one case is coming from the community (Andreas K and others) whereas in the CUDA Python case it is coming from NVIDIA. There are syntactical differences of course, but you should be able to perform basic operations using either methodology.

CUDA Python allows for the possibility to have a “standardized” host api/interface, while still being able to use other methodologies such as Numba to enable (for example) the writing of kernel code in python.

This blog and the questions that follow it may be of interest.

FWIW there are other python/CUDA methodologies. numba, cupy, CUDA python, and pycuda are some of the available approaches to tap into CUDA acceleration from Python. CUDA Python can interoperate with most or all of them.

https://forums.developer.nvidia.com/t/cuda-python-vs-pycuda/216751

Unifying the CUDA Python Ecosystem | NVIDIA Technical Blog https://developer.nvidia.com/blog/unifying-the-cuda-python-ecosystem/

pybind11-stubgen --ignore-all-errors cuda.cudart
pybind11-stubgen --ignore-all-errors cuda.cuda

未生成前

生成后

知识点

1. .so / .pyd 代码底层实现

2. Stub（存根） 是一个通用术语，指“只提供接口定义，不包含实现”的代码文件

3. .pyi 是 Python 的“存根文件”（Stub File），它是一个纯类型注解文件，用于为 .py 或 .pyd（C 扩展）模块提供 类型提示信息。 为模块提供“类型签名”，供 IDE 和类型检查工具使用，从而实现智能提示，补全