MMRoate—CUDA实现旋转目标检测中的PointsInPolygon()函数

// 包含了数据格式的定义 相当于一个数据接口文件, 后面使用的at接口，以及PYBIND11_MODULE()函数
// 都在这个头文件中
#include <torch/extension.h> 
#include <cmath>
#include <iostream>
#include <vector>
// 对`PointsJFLauncher()`函数先进行声明，该函数的定义在points_justify_kernel.cu文件中；
// 否则g++编译器编译到文件中下面的语句会报错；
// PointsJFLaucher(points, polygons, rows, cols, output);
int PointsJFLaucher(const at::Tensor points, const at::Tensor polygons,
                const int rows, const int cols, at::Tensor output);
// 宏定义用于变量检查，且注意宏定义必须在一行内写完，写不完可以使用\进行换行
// CHECK_CUDA：张量类型检查；
// CHECK_CONTIGUOUS：张量连续性检查，后续会用到参数的指针获取数据，所以数据存放必须是连续的；
#define CHECK_CUDA(x) AT_CHECK(x.type().is_cuda(), #x, " must be a CUDAtensor ")
#define CHECK_CONTIGUOUS(x) \
  AT_CHECK(x.is_contiguous(), #x, " must be contiguous ")
#define CHECK_INPUT(x) \
  CHECK_CUDA(x);       \
  CHECK_CONTIGUOUS(x)
// Python文件中调用的函数名称
int pointsJf(at::Tensor points, at::Tensor polygons, at::Tensor output) {
    //输入变量的类型检查
    CHECK_INPUT(points);
    CHECK_INPUT(polygons);
    CHECK_INPUT(output);
    //检查points格式 [ ,2]
    int size_points = points.size(1);
    if(size_points != 2){
        printf("wrong points size\n");
        return 0;
    //检查polygons格式 [ ,8]
    int size_polygons = points.size(1);
    if(size_polygons != 2){
        printf("wrong polygons size\n");
        return 0;
    int rows = points.size(0);
    int cols = polygons.size(0);
    // 调用CUDA文件中的Launcher函数
    PointsJFLaucher(points, polygons, rows, cols, output);
    return 1;
// PYBIND11_MODULE提供python调用接口
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("pointsJf", &pointsJf, "Justify a point is in polygon or not");  //该对象对应的方法名   函数指针  说明  
}

#include <ATen/ATen.h>
#include <THC/THCAtomics.cuh>
#include <math.h>
#define EPS 1e-8
// 用于线程循环分配任务
#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)
// 宏定义，每个BLOCK内最多有1024个线程
#define THREADS_PER_BLOCK 1024
// 获取当前计算需要多少个BLOCK，最多可以使用65000个BLOCK
inline int GET_BLOCKS(const int N) {
    int optimal_block_num = (N + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
    int max_block_num = 65000;
    return min(optimal_block_num, max_block_num);
// Point的结构体
struct point
    float x,y;
// template范型
template <typename scalar_t>
__global__ void PointsJF(const int nthreads, const scalar_t *vertex1, const scalar_t *vertex2, 
                    const int rows, const int cols, scalar_t *inside_flag) {
    CUDA_1D_KERNEL_LOOP(index, nthreads) {
        //确定现在计算第row个点 和 第col个多边形间的关系
        // 这里的index就是 cols * rows，代表当前是第几个任务
        int row = index / cols;
        int col = index % cols;
        const scalar_t *offset_vertex1 = vertex1 + row * 2;
        const scalar_t *offset_vertex2 = vertex2 + col * 8;
        // 定义point结构体数组
        point point_[1];  // 一个点
        point polygon[4];  // 一个polygon由4个点构成
        point_[0].x = offset_vertex1[0];
        point_[0].y = offset_vertex1[1];
        polygon[0].x = offset_vertex2[0];
        polygon[0].y = offset_vertex2[1];
        polygon[1].x = offset_vertex2[2];
        polygon[1].y = offset_vertex2[3];
        polygon[2].x = offset_vertex2[4];
        polygon[2].y = offset_vertex2[5];
        polygon[3].x = offset_vertex2[6];
        polygon[3].y = offset_vertex2[7];
        // 指示当前点与多边形的所属关系
        int nCross = 0;
        int i, j;
        float sx, sy, tx, ty, px, py, x;
        for(i=0,j=3;i<4;j=i,i++)    // i,j分别是当前节点和下一个点
            sx = polygon[i].x;
            sy = polygon[i].y;
            tx = polygon[j].x;
            ty = polygon[j].y;
            px = point_[0].x;
            py = point_[0].y;
            if ( py < min(sy, ty)) //如果目标点低于这个线段，跳过
                continue; 
            if ( py > max(sy, ty)) //如果目标点高于这个线段，跳过
                continue; 
            // 顶点情况
            if((sx == px && sy == py) || (tx == px && ty == py))
                break;
            else //射线法
                if((sy < py && ty >= py) || (sy >= py && ty < py)) 
                //如果过p1画水平线，过p2画水平线，目标点在这两条线中间
                    x = sx + (py - sy) * (tx - sx) / (ty - sy);
                    // 在边界线上
                    if(x == px)
                        break;
                    if(x > px)
                        nCross++;
        if (nCross % 2 == 1) {
            inside_flag[index] = 1.0; //如果是奇数，说明在多边形里
        else {
            inside_flag[index] = 0.0; //否则在多边形外 或 边上
        return;
// 在cpp文件中声明的函数的函数体实现，其中变量points, polygons, rows, cols用const修饰，
// 我的理解是不希望这些变量的内容在函数内部被修改;
int PointsJFLaucher(const at::Tensor points, const at::Tensor polygons,
                const int rows, const int cols, at::Tensor output) {
    const int output_size = rows * cols;
    // AT_DISPATCH_FLOATING_TYPES_AND_HALF()函数:
    // 第一个参数points.type(): 第一个参数是输入数据类型;
    // 第二个参数"PointsJFLaucher": 第二个参数是标识符，用于显示报错信息；
    // 第三个参数是一个匿名函数(匿名函数就是[&]{…………}这部分内容)
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
        points.type(), "PointsJFLaucher", ([&] {
        const scalar_t *vertex1 = points.data<scalar_t>();  // 取数据对应的地址
        const scalar_t *vertex2 = polygons.data<scalar_t>();  // 取数据对应的地址
        scalar_t *inside_flag = output.data<scalar_t>();   // 取数据对应的地址
        PointsJF<scalar_t>
            <<<GET_BLOCKS(output_size), THREADS_PER_BLOCK>>>(