nms_cuda详解

微凉code

已于 2024-10-20 16:16:03 修改

阅读量1.3k

点赞数 20

分类专栏： cuda和pytorch算子文章标签：深度学习

于 2024-10-13 11:35:14 首次发布

本文链接：https://blog.youkuaiyun.com/qq_41921315/article/details/141435800

版权

cuda和pytorch算子专栏收录该内容

3 篇文章

订阅专栏

此篇为PyTorch 自定义算子：复现CPU和CUDA版的二维卷积的代码详解
这篇是为了展示setup在构建简单的cpp算子的使用,仓库地址:onnx_op

1.环境配置
整体结构如下图所示,架构和之前一样

2.nms.cpp

nms的思路还是很简单的，和之前python的版本一致
python版本

#include "pytorch_cpp_helper.hpp"
Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold, int offset);
Tensor nms_cpu(Tensor boxes,Tensor scores,float iou_threshold,int offset)
{
    /*
        基本思路：
        (1) 取出x1,y1,x2,y2,areas_t,scores
        (2) 循环将bbox和其他剩余的bbox进行交并比计算，将交并比大于阈值的bbox从这个集合中剔除出去,设为false
        (3) 继续循环
        为了提高效率，我们保留bbox不动，最终保留的也都是bbox在原集合中的索引(mask_select)
    */
    if(boxes.numel()==0){
        return at::empty({0},boxes.options().dtype(at::kLong));
    }
    // 获取boxes的各个维度
    auto x1_t =boxes.select(1,0).contiguous();
    auto y1_t = boxes.select(1,1).contiguous();
    auto x2_t = boxes.select(1,2).contiguous();
    auto y2_t = boxes.select(1,3).contiguous();
    // 获取每个box的面积
    Tensor areas_t = (x2_t - x1_t + offset) * (y2_t - y1_t + offset);
    auto order_t =std::get<1> (scores.sort(0,true));
    
    auto nboxes = boxes.size(0);
    Tensor select_t = at::ones({nboxes}, boxes.options().dtype(at::kBool));
    // 用data_ptr可以很方便的获取一个tensor的元素指针，从而访问tensor
    auto select = select_t.data_ptr<bool>();
    auto order = order_t.data_ptr<int64_t>();
    auto x1 = x1_t.data_ptr<float>();
    auto y1 = y1_t.data_ptr<float>();
    auto x2 = x2_t.data_ptr<float>();
    auto y2 = y2_t.data_ptr<float>();
    auto areas = areas_t.data_ptr<float>();

    for (int64_t _i = 0; _i < nboxes; _i++) {
        if (select[_i] == false) continue;
        auto i = order[_i];
        auto ix1 = x1[i];
        auto iy1 = y1[i];
        auto ix2 = x2[i];
        auto iy2 = y2[i];
        auto iarea = areas[i];

        for (int64_t _j = _i + 1; _j < nboxes; _j++) {
            if (select[_j] == false) continue;
            auto j = order[_j];
            auto xx1 = std::max(ix1, x1[j]);
            auto yy1 = std::max(iy1, y1[j]);
            auto xx2 = std::min(ix2, x2[j]);
            auto yy2 = std::min(iy2, y2[j]);

            auto w = std::max(0.f, xx2 - xx1 + offset);
            auto h = std::max(0.f, yy2 - yy1 + offset);
            auto inter = w * h;
            auto ovr = inter / (iarea + areas[j] - inter);
            if (ovr > iou_threshold) select[_j] = false;
        } 
    }
    return order_t.masked_select(select_t);
}


PYBIND11_MODULE(my_ops, m)
{
    m.def("nms", nms_cpu, "nms_compute",
        py::arg("boxes"), py::arg("scores"), py::arg("iou_threshold"),
        py::arg("offset"));
    m.def("nms_cuda", NMSCUDAKernelLauncher, "nms_compute_cuda",
        py::arg("boxes"), py::arg("scores"), py::arg("iou_threshold"),
        py::arg("offset"));
}

3.nms_cuda.cu

// 版权声明
#include "nms_cuda_kernel.cuh"
#include "pytorch_cuda_helper.hpp"

/**
 * @brief CUDA核启动器，用于非最大抑制（NMS）。
 * 
 * @param boxes 候选框的坐标，格式为[y1, x1, y2, x2]。
 * @param scores 候选框的分数。
 * @param iou_threshold NMS算法中使用的IoU（交并比）阈值。
 * @param offset 用于索引调整的偏移量。
 * @return Tensor 包含经过NMS后保留的候选框索引的张量。
 */
Tensor NMSCUDAKernelLauncher(Tensor boxes, Tensor scores, float iou_threshold, int offset) {
  // 确保操作在正确的CUDA设备上进行
  at::cuda::CUDAGuard device_guard(boxes.device());

  // 如果没有候选框，则直接返回一个空张量
  if (boxes.numel() == 0) {
    return at::empty({0}, boxes.options().dtype(at::kLong));
  }

  // 根据分数对候选框进行排序，获取排序后的索引
  // std::get 是 C++ 标准库中的一个函数，它主要用于访问 std::tuple 或者其他聚合类型（如 std::pair）中的元素。
  auto order_t = std::get<1>(scores.sort(0, /*descending=*/true));
  // 使用排序后的索引对候选框进行排序
  auto boxes_sorted = boxes.index_select(0, order_t);

  // 计算候选框的数量
  int boxes_num = boxes.size(0);
  // 计算列块数，用于CUDA网格和块的配置
  const int col_blocks = (boxes_num + threadsPerBlock - 1) / threadsPerBlock;
  const int col_blocks_alloc = GET_BLOCKS(boxes_num, threadsPerBlock);
  // 创建一个掩码张量，用于存储NMS过程中每个候选框的状态
  Tensor mask = at::empty({boxes_num, col_blocks}, boxes.options().dtype(at::kLong));
  // 配置CUDA网格和块的尺寸
  dim3 blocks(col_blocks_alloc, col_blocks_alloc);
  dim3 threads(threadsPerBlock);
  // 获取当前的CUDA流
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  // 启动NMS CUDA核
  nms_cuda<<<blocks, threads, 0, stream>>>(boxes_num, iou_threshold, offset, boxes_sorted.data_ptr<float>(), (unsigned long long*)mask.data_ptr<int64_t>());
  
  // 创建一个张量，用于存储最终保留的候选框索引
  at::Tensor keep_t = at::zeros({boxes_num}, boxes.options().dtype(at::kBool).device(at::kCUDA));
  // 从掩码中收集保留的候选框索引
  gather_keep_from_mask<<<1, min(col_blocks, THREADS_PER_BLOCK), col_blocks * sizeof(unsigned long long), stream>>>(keep_t.data_ptr<bool>(), (unsigned long long*)mask.data_ptr<int64_t>(), boxes_num);
  // 检查CUDA操作是否有错误
  AT_CUDA_CHECK(cudaGetLastError());
  // 返回排序后的索引中对应保留候选框的索引
  return order_t.masked_select(keep_t);
}

现在我们重点看一下nms_cuda_kernel.cuh

3.nms_cuda_kernel.cuh

为了便于理解,可以自己运行下面的程序来看实际输出

#include <stdio.h>
#include "run.h"
__global__ void checkIndex(int n_boxes)
{
    int threadsPerBlock=64;
    // 获取当前线程块在二维网格中的行索引和列索引
    const int row_start = blockIdx.y;
    const int col_start = blockIdx.x;
    const int tid = threadIdx.x;
    
    // 如果行索引大于列索引，直接返回，不进行后续计算
    if (row_start > col_start) return;
    //可以看到只运行了右上角
    //printf("row_start:%d  col_start:%d\n",row_start,col_start);
    const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
    const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);
    if(row_start==2&&col_start==2){
        printf("row_size:%d,col_size:%d\n",row_size,col_size);
        if(tid<col_size){
            printf("block(2,2)线程%d搬运的数据是:%d-%d\n",tid,(threadsPerBlock * col_start + tid) * 4 + 0,(threadsPerBlock * col_start + tid) * 4 + 3);
        }
    }
    // 确保所有线程都完成了数据加载
    __syncthreads();

}

int main(void)
{
    int boxes_num = 190;
    int threadsPerBlock=64;
    int col_blocks = (boxes_num + threadsPerBlock - 1) / threadsPerBlock;
    int col_blocks_alloc = GET_BLOCKS(boxes_num, threadsPerBlock);
    //int col_blocks_alloc = 2;
    printf("%d,%d,%d,%d\n",boxes_num,threadsPerBlock,col_blocks,col_blocks_alloc);
    dim3 block(col_blocks_alloc, col_blocks_alloc);
    dim3 threads(threadsPerBlock);

    printf("block: %d, %d, %d\n", block.x, block.y, block.z);
    printf("thread: %d, %d, %d\n", threads.x, threads.y, threads.z);
    //printf("grid: %d, %d, %d\n", grid.x, grid.y, grid.z);

    // //checkIndex<<<grid,block>>>();
    checkIndex<<<block,threads>>>(boxes_num);
    cudaDeviceReset();

    return 0;
}

上面的测试程序可以使用下面两布命令测试
nvcc test_main.cu -o test
./test

3.1 nms_cuda函数

// Copyright (c) OpenMMLab. All rights reserved
#ifndef NMS_CUDA_KERNEL_CUH
#define NMS_CUDA_KERNEL_CUH

#include <float.h>
#include "common_cuda_helper.hpp"
#include "pytorch_cuda_helper.hpp"

// 该语句定义了一个名为threadsPerBlock的常量，其值等于unsigned long long int类型的字节数乘以8。
// 这意味着它计算了一个线程块中可以容纳的线程数量。
int const threadsPerBlock = sizeof(unsigned long long int) * 8;

__device__ inline bool devIoU(float const *const a, float const *const b,
                              const int offset, const float threshold) {
  float left = fmaxf(a[0], b[0]), right = fminf(a[2], b[2]);
  float top = fmaxf(a[1], b[1]), bottom = fminf(a[3], b[3]);
  float width = fmaxf(right - left + offset, 0.f),
        height = fmaxf(bottom - top + offset, 0.f);
  float interS = width * height;
  float Sa = (a[2] - a[0] + offset) * (a[3] - a[1] + offset);
  float Sb = (b[2] - b[0] + offset) * (b[3] - b[1] + offset);
  return interS > threshold * (Sa + Sb - interS);
}
// 实现非最大抑制（NMS）算法的CUDA内核函数
// 用于对一系列边界框进行筛选，去除重叠度较高的框
// 参数说明：
// - n_boxes: 输入边界框的数量
// - iou_threshold: 用于判断两个框是否重叠的阈值
// - offset: 边界框的偏移量，用于调整框的位置
// - dev_boxes: 设备（GPU）上的边界框数据
// - dev_mask: 设备（GPU）上的掩码数据，用于记录每个边界框的抑制结果
__global__ static void nms_cuda(const int n_boxes, const float iou_threshold,
                                const int offset, const float *dev_boxes,
                                unsigned long long *dev_mask) {
  // 计算所需的块数量，以适应输入的边界框数量
  int blocks = (n_boxes + threadsPerBlock - 1) / threadsPerBlock;

  // 使用2D循环遍历所有块，处理NMS算法
  CUDA_2D_KERNEL_BLOCK_LOOP(col_start, blocks, row_start, blocks) {
    // 获取当前线程在块内的索引
    const int tid = threadIdx.x;
    /*
        这段代码可能用于对称矩阵操作或需要避免重复计算的场景。若row_start大于col_start，
        意味着处理的是矩阵中的上三角部分（对于下三角矩阵而言）。为了避免重复计算
        （比如只计算下三角或主对角线元素），当检测到row_start大于col_start时，代码提前返回，
        跳过这部分计算。这样可以确保每个相关元素仅被计算一次。
    */
    // 确保行号不大于列号，以避免重复计算 blockidx.y>blockidx.x,return
    if (row_start > col_start) return;

    // 计算当前块的行和列应处理的边界框数量
    // n_boxes-blockIdx.y*threadsPerBlock
    const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
    const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

    // 共享内存中存储当前块所需的边界框数据
    __shared__ float block_boxes[threadsPerBlock * 4];
    // 每个线程负责将全局内存中的边界框数据加载到共享内存中
    if (tid < col_size) {
      block_boxes[tid * 4 + 0] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 0];
      block_boxes[tid * 4 + 1] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 1];
      block_boxes[tid * 4 + 2] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 2];
      block_boxes[tid * 4 + 3] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 3];
    }
    // 确保所有线程都完成了数据加载
    __syncthreads();

    // 每个线程处理不同的行边界框
    if (tid < row_size) {
      // 获取当前线程应处理的边界框索引
      const int cur_box_idx = threadsPerBlock * row_start + tid;
      // 指向当前边界框的数据
      const float *cur_box = dev_boxes + cur_box_idx * 4;
      // 初始化相关变量
      int i = 0;
      unsigned long long int t = 0;
      int start = 0;
      // 如果行号等于列号，从下一个边界框开始处理
      if (row_start == col_start) {
        start = tid + 1;
      }
      // 遍历列边界框，计算重叠度并更新掩码
      for (i = start; i < col_size; i++) {
        // 如果当前边界框与其它边界框的重叠度超过阈值，则在掩码中做标记
        if (devIoU(cur_box, block_boxes + i * 4, offset, iou_threshold)) {
          t |= 1ULL << i;
        }
      }
      // 将计算结果存储到掩码数组中
      dev_mask[cur_box_idx * gridDim.y + col_start] = t;
    }
  }
}

nms_cuda的解析在这个链接
在这里插入图片描述

CUDA_2D_KERNEL_BLOCK_LOOP(col_start, blocks, row_start, blocks)
这个看起来是循环，但实际上还是一次执行,旧版的mmcv是下面的代码

  const int row_start = blockIdx.y;
  const int col_start = blockIdx.x;

if (row_start > col_start) return;
在block(0,1)中我们计算了bbox(0,63)和bbox(64,127)的iou
但是在bbox(1,0)中，我们又计算了一遍bbox(64,127)和bbox(0,63)的iou，所以我们跳过一些block即可,也就是只计算上三角行(加对角线)
row_size和col_size

    const int row_size =
        fminf(n_boxes - row_start * threadsPerBlock, threadsPerBlock);
    const int col_size =
        fminf(n_boxes - col_start * threadsPerBlock, threadsPerBlock);

对于block(0,0),我们是计算64个bbox的iou，但是假如说只有190和bbox，那么对于block(0,2)来说，我们计算的就不是64个iou了，我们计算的是bbox(0,63)和bbox(128,189)的iou，那么对应的来说，线程应该搬运(128,189)的数据到共享内存中。那么对于block(2,2)来说，我们计算的时候也要注意row_size。
4. 加载数据到共享内存

	__shared__ float block_boxes[threadsPerBlock * 4];
    if (tid < col_size) {
      block_boxes[tid * 4 + 0] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 0];
      block_boxes[tid * 4 + 1] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 1];
      block_boxes[tid * 4 + 2] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 2];
      block_boxes[tid * 4 + 3] =
          dev_boxes[(threadsPerBlock * col_start + tid) * 4 + 3];
    }
    // 确保所有线程都完成了数据加载
    __syncthreads();

每个线程只搬运自己的数据，例如thread0只搬运bbox0的数据,同时等待所以线程都完成数据加载
5.计算IOU

    // 每个线程处理不同的行边界框
    if (tid < row_size) {
      // 获取当前线程应处理的边界框索引
      const int cur_box_idx = threadsPerBlock * row_start + tid;
      // 指向当前边界框的数据
      const float *cur_box = dev_boxes + cur_box_idx * 4;
      // 初始化相关变量
      int i = 0;
      unsigned long long int t = 0;
      int start = 0;
      // 如果行号等于列号，从下一个边界框开始处理
      if (row_start == col_start) {
        start = tid + 1;
      }
      // 遍历列边界框，计算重叠度并更新掩码
      for (i = start; i < col_size; i++) {
        // 如果当前边界框与其它边界框的重叠度超过阈值，则在掩码中做标记
        if (devIoU(cur_box, block_boxes + i * 4, offset, iou_threshold)) {
          t |= 1ULL << i;
        }
      }
      // 将计算结果存储到掩码数组中
      dev_mask[cur_box_idx * gridDim.y + col_start] = t;
    }

// 每个线程要算col_size次iou
t |= 1ULL << i;
通过下面的测试程序可以看到，t的作用

#include<iostream>
#include<cstdint>
#include <bitset>
using namespace std;

std::string intToBinaryString(int num) {
	std::bitset<32> bits(num);
	return bits.to_string();
}
int main()
{
	unsigned long long int t = 0;
	cout << "1ULL:" << 1ULL << endl;
	for (int i = 0; i < 5; i++) {
		int tmp = 1ULL << i;
		//可以看到<<i,i是几，对应的二进制的字符串位置上就是1
		cout << "1ULL<<i::" << tmp <<"  binary:" <<intToBinaryString(tmp)<< endl;
	}
	t += (1ULL) << 1;
	t += (1ULL) << 3;
	t += (1ULL) << 6;
	cout << "str(t):" << intToBinaryString(t);
	return 0;
}

在这里插入图片描述
5. 创建keep，将mask数据从CUDA上存到keep中

 // 创建一个张量，用于存储最终保留的候选框索引
  at::Tensor keep_t = at::zeros({boxes_num}, boxes.options().dtype(at::kBool).device(at::kCUDA));
  // 从掩码中收集保留的候选框索引
  gather_keep_from_mask<<<1, min(col_blocks, THREADS_PER_BLOCK), col_blocks * sizeof(unsigned long long), stream>>>(keep_t.data_ptr<bool>(), (unsigned long long*)mask.data_ptr<int64_t>(), boxes_num);

在 CUDA 中，<<<gridDim, blockDim, sharedMemSize, stream>>> 这种语法用于指定内核函数的执行配置参数，具体解释如下：

gridDim（第一个参数）：
这是一个 dim3 类型的变量，用于指定网格（grid）的维度，即启动的线程块的数量和布局。例如，dim3(10, 5, 1) 表示在三维网格中有 10×5×1 个线程块。在你的例子中，这里是 1，表示只启动一个一维的网格，可能因为这个操作只需要一个块的维度就足够完成任务。
blockDim（第二个参数）：
同样是一个 dim3 类型的变量，用于指定每个线程块（block）中的线程数量和布局。例如，dim3(32, 16, 1) 表示每个线程块中有 32×16×1 个线程。在你的例子中，这里是 min(col_blocks, THREADS_PER_BLOCK)，表示根据 col_blocks 和 THREADS_PER_BLOCK 的较小值来确定每个线程块中的线程数量，这样可以确保不超过一定的线程数量限制。这样表述mask一个块的数据分配一个线程来取回。
sharedMemSize（第三个参数）：
这是一个以字节为单位的整数，表示每个线程块所分配的共享内存大小。在你的例子中，这里是 col_blocks * sizeof(unsigned long long)，表示根据 col_blocks 的数量乘以 unsigned long long 类型的大小来确定每个线程块分配的共享内存大小。共享内存可以被同一线程块中的所有线程快速访问，用于存储临时数据或中间结果，以提高内核函数的性能。
stream（第四个参数）：
这是一个 CUDA 流对象，用于指定内核函数在哪个流上执行。流可以让多个 GPU 操作并发执行，提高 GPU 的利用率。在你的例子中，这个流参数用于确保该内核函数在特定的流上执行，以便与其他操作进行协调和同步。如果不指定流，默认使用零号流。

3.2 gather_keep_from_mask函数

#include <stdio.h>
#include "run.h"
__global__ void checkIndex(int n_boxes)
{
    int threadsPerBlock=64;
    // 计算列块的数量，用于处理共享内存中的去除标记。
    const int col_blocks = 3;
    // 获取线程在块中的索引。
    const int tid = threadIdx.x;

    // 共享内存数组，用于标记哪些框应该被去除。
    // 注意：数组大小在运行时确定，由CUDA运行时分配。
    extern __shared__ unsigned long long removed[];

    // 初始化去除标记数组。
    for (int i = tid; i < col_blocks; i += blockDim.x) {
        removed[i] = 0;
    }
    // 确保所有线程都完成了数据加载
    __syncthreads();

  // 遍历列块，处理每个块中的线程块。
  for (int nblock = 0; nblock < col_blocks; ++nblock) {
    // 获取当前块的去除值。
    auto removed_val = removed[nblock];
    // 确保所有线程都完成当前块的处理后再继续。
    __syncthreads();
    // 计算当前块的起始偏移量。
    const int i_offset = nblock * threadsPerBlock;
    printf("i_offset:%d\n",i_offset);
    //每个线程执行64次循环
    for (int inblock = 0; inblock < threadsPerBlock; ++inblock) {
      // 计算当前线程在块中的索引。
      const int i = i_offset + inblock;
      printf("i:%d,线程:%d\n",i,tid);
      // 确保索引不超过框的数量。
      if (i >= n_boxes) break;
      // 如果当前线程对应的框未被标记为去除，则进行处理。
      if (!(removed_val & (1ULL << inblock))) {
        // 线程0负责标记该输出框为保留。
        if (tid == 0) {
          //keep[i] = true;
        }
        // 获取当前框的掩码地址。
        
        //auto p = dev_mask + i * col_blocks;
        // 去除所有与当前框重叠的框。
        for (int j = tid; j < col_blocks; j += blockDim.x) {
          //if (j >= nblock) removed[j] |= p[j];
          printf("此处循环执行了%d,blocdim.x:%d,线程为:%d\n",j,blockDim.x,tid);
        }
        // 确保所有线程都完成当前块的更新后再继续。
        __syncthreads();
        // 更新当前块的去除值。
        removed_val = removed[nblock];
      }
    }
  }
}

int main(void)
{
    int boxes_num = 192;
    int threadsPerBlock=64;
    int col_blocks = (boxes_num + threadsPerBlock - 1) / threadsPerBlock;
    int col_blocks_alloc = GET_BLOCKS(boxes_num, threadsPerBlock);
    //int col_blocks_alloc = 2;
    printf("%d,%d,%d,%d\n",boxes_num,threadsPerBlock,col_blocks,col_blocks_alloc);
    dim3 block(1);
    dim3 threads(3);

    printf("block: %d, %d, %d\n", block.x, block.y, block.z);
    printf("thread: %d, %d, %d\n", threads.x, threads.y, threads.z);
    
    //printf("grid: %d, %d, %d\n", grid.x, grid.y, grid.z);

    // //checkIndex<<<grid,block>>>();
    checkIndex<<<1,col_blocks,col_blocks * sizeof(unsigned long long)>>>(boxes_num);
    cudaDeviceReset();

    return 0;
}

// 使用CUDA __global__关键字声明一个全局函数，用于从掩码中收集并保留特定的框。
// 这个函数旨在处理一组框（如边界框），根据给定的掩码确定哪些框应该被保留。
// 参数:
// - keep: 一个指针，用于标记哪些框应该被保留。
// - dev_mask: 一个指向设备（GPU）上掩码数组的指针，该数组指示哪些框应该被保留。
// - n_boxes: 表示框的数量。
__global__ static void gather_keep_from_mask(bool *keep,
                                             const unsigned long long *dev_mask,
                                             const int n_boxes) {
  // 计算列块的数量，用于处理共享内存中的去除标记。
  const int col_blocks = (n_boxes + threadsPerBlock - 1) / threadsPerBlock;
  // 获取线程在块中的索引。
  const int tid = threadIdx.x;

  // 共享内存数组，用于标记哪些框应该被去除。
  // 注意：数组大小在运行时确定，由CUDA运行时分配。
  extern __shared__ unsigned long long removed[];

  // 初始化去除标记数组。
  for (int i = tid; i < col_blocks; i += blockDim.x) {
    removed[i] = 0;
  }
  // 确保所有线程都完成初始化后再继续。
  __syncthreads();

  // 遍历列块，处理每个块中的线程块。
  for (int nblock = 0; nblock < col_blocks; ++nblock) {
    // 获取当前块的去除值。
    auto removed_val = removed[nblock];
    // 确保所有线程都完成当前块的处理后再继续。
    __syncthreads();
    // 计算当前块的起始偏移量。
    const int i_offset = nblock * threadsPerBlock;
    // 使用#pragma unroll进行循环展开，以提高性能。
#pragma unroll
    for (int inblock = 0; inblock < threadsPerBlock; ++inblock) {
      // 计算当前线程在块中的索引。
      const int i = i_offset + inblock;
      // 确保索引不超过框的数量。
      if (i >= n_boxes) break;
      // 如果当前线程对应的框未被标记为去除，则进行处理。
      // removed_val为1表述该框被去除，则跳过。
      // removed_val为0表示该框保留，则保留。
      if (!(removed_val & (1ULL << inblock))) {
        // 线程0负责标记该输出框为保留。
        //全局只有这一句修改keep的地方，不用担心其他线程修改
        if (tid == 0) {
          keep[i] = true;
        }
        // 获取当前框的掩码地址。
        auto p = dev_mask + i * col_blocks;
        // 去除所有与当前框重叠的框。
        // 前面的框已经处理过了，只需要比较后面的block里面的框就可以
        for (int j = tid; j < col_blocks; j += blockDim.x) {
          if (j >= nblock) removed[j] |= p[j];
        }
        // 确保所有线程都完成当前块的更新后再继续。
        __syncthreads();
        // 更新当前块的去除值。
        removed_val = removed[nblock];
      }
    }
  }
}

removed为unsigned long long类型的，长度为col_blocks的一维数组
每个数字表述当前块的一个mask情况
仅使用线程0进行修改
removed[j] |= p[j]
若mask[i][j]为1，则remv[j]也设置为1
通过下面的测试，可以很清楚的看到 | 操作的效果.

#include<iostream>
#include<cstdint>
#include <bitset>
using namespace std;


std::string intToBinaryString(int num) {
	std::bitset<32> bits(num);
	return bits.to_string();
}

void test2()
{
	int tmp = 1ULL << 4;
	cout << "tmp:" << intToBinaryString(tmp) << endl;
	int remv = 2;
	int ans = remv | tmp;
	cout << "remv|tmp:" << (remv | tmp) << endl;
	cout << "str::" << intToBinaryString(ans) << endl;
}
int main()
{
	test2();
	return 0;
}

4.setup.py

使用 python setup.py install 安装算子库,python test.py进行测试

from setuptools import setup
from torch.utils import cpp_extension
import os

src_root = './'
cpp_src = ['nms.cpp','nms_cuda.cu']

if __name__ == '__main__':
    include_dirs = ['./']
    cpp_path = [os.path.join(src_root, src) for src in cpp_src]

    setup(
        name='panoflow',
        ext_modules=[
            cpp_extension.CUDAExtension(
                'my_ops', cpp_path, include_dirs=include_dirs)
        ],
        cmdclass={'build_ext': cpp_extension.BuildExtension})

test.py

import torch
import torch.nn as nn
# from my_conv import MyConv2d
import numpy as np
import my_ops
def test_3():
    bboxes = np.array([[100, 100, 210, 210, 0.72],
                [250, 250, 420, 420, 0.8],
                [220, 220, 320, 330, 0.92],
                [100, 100, 210, 210, 0.72],
                [230, 240, 325, 330, 0.81],
                [220, 230, 315, 340, 0.9]])
    bboxes=torch.from_numpy(bboxes)
    scores=bboxes[:,4].to(torch.float32).cuda()
    bboxes=bboxes[:,:4].to(torch.float32).cuda()
    keep=my_ops.nms_cuda(bboxes,scores,0.5,0)
    print(keep)
if __name__ == '__main__':
    test_3()