point-in-rbbox

功能

实现point-in-rbbox算子（支持CUDA和Triton），用于计算3D旋转框内的点云，相比PyTorch版本显存有明显优化

依赖

• CUDA：11.8
• PyTorch：2.0.0
• Triton：2.0.0（使用Triton算子时需要）

安装

• pip安装

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX" # 支持不同GPU架构【可选】
export ENABLE_BF16=1 # 支持BF16【可选，部分GPU可能不支持】
python3 -m pip install points-in-rbbox

• 源码安装

git clone https://github.com/lh9171338/points-in-rbbox.git
cd points-in-rbbox

export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0+PTX" # 支持不同GPU架构【可选】
export ENABLE_BF16=1 # 支持BF16【可选，部分GPU可能不支持】
python3 setup.py install

使用

• CUDA算子

from points_in_rbbox import points_in_rbbox_cuda

mask = points_in_rbbox_cuda(points, boxes)

• Triton算子（无需编译，自动适配GPU架构）

from points_in_rbbox import points_in_rbbox_triton

mask = points_in_rbbox_triton(points, boxes)

显存&耗时