grid_sampler.cu

#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <cuda_fp16.h>
#include <iostream>
#include <vector>

// 边界模式枚举
enum class GridSampleMode {
    BILINEAR = 0,
    NEAREST = 1
};

enum class GridSamplePaddingMode {
    ZEROS = 0,
    BORDER = 1,
    REFLECTION = 2
};

// 双线性插值辅助函数
__device__ __forceinline__ float bilinear_interpolate(
    const float* input,
    int height, int width,
    float y, float x,
    int c, int n,
    int input_height, int input_width, int input_channels, int input_batch
) {
    // 计算输入索引（用于边界检查）
    // int input_idx = n * input_height * input_width * input_channels +
    //                c * input_height * input_width +
    //                (int)y * input_width + (int)x;

    // 边界检查
    if (x < 0 || x >= width || y < 0 || y >= height) {
        return 0.0f;
    }

    // 获取四个角点坐标
    int y_low = (int)y;
    int x_low = (int)x;
    int y_high = y_low + 1;
    int x_high = x_low + 1;

    // 边界处理
    y_low = max(0, min(y_low, height - 1));
    x_low = max(0, min(x_low, width - 1));
    y_high = max(0, min(y_high, height - 1));
    x_high = max(0, min(x_high, width - 1));

    // 计算权重
    float ly = y - y_low;
    float lx = x - x_low;
    float hy = 1.0f - ly;
    float hx = 1.0f - lx;

    // 获取四个角点的值
    float v1 = input[n * input_height * input_width * input_channels +
                     c * input_height * input_width +
                     y_low * input_width + x_low];
    float v2 = input[n * input_height * input_width * input_channels +
                     c * input_height * input_width +
                     y_low * input_width + x_high];
    float v3 = input[n * input_height * input_width * input_channels +
                     c * input_height * input_width +
                     y_high * input_width + x_low];
    float v4 = input[n * input_height * input_width * input_channels +
                     c * input_height * input_width +
                     y_high * input_width + x_high];

    // 双线性插值
    float w1 = hy * hx;
    float w2 = hy * lx;
    float w3 = ly * hx;
    float w4 = ly * lx;

    return w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4;
}

// 最近邻插值辅助函数
__device__ __forceinline__ float nearest_interpolate(
    const float* input,
    int height, int width,
    float y, float x,
    int c, int n,
    int input_height, int input_width, int input_channels, int input_batch
) {
    // 边界检查
    if (x < 0 || x >= width || y < 0 || y >= height) {
        return 0.0f;
    }

    // 四舍五入到最近的整数
    int y_nearest = (int)roundf(y);
    int x_nearest = (int)roundf(x);

    // 边界处理
    y_nearest = max(0, min(y_nearest, height - 1));
    x_nearest = max(0, min(x_nearest, width - 1));

    int input_idx = n * input_height * input_width * input_channels +
                   c * input_height * input_width +
                   y_nearest * input_width + x_nearest;

    return input[input_idx];
}

// 边界处理函数
__device__ __forceinline__ void apply_padding_mode(
    float& x, float& y, int width, int height, GridSamplePaddingMode padding_mode
) {
    switch (padding_mode) {
        case GridSamplePaddingMode::BORDER:
            x = max(-0.5f, min(x, width - 0.5f));
            y = max(-0.5f, min(y, height - 0.5f));
            break;
        case GridSamplePaddingMode::REFLECTION:
            if (x < 0) x = -x;
            if (x >= width) x = 2 * width - 2 - x;
            if (y < 0) y = -y;
            if (y >= height) y = 2 * height - 2 - y;
            break;
        case GridSamplePaddingMode::ZEROS:
        default:
            // 不需要修改坐标，在插值函数中处理边界
            break;
    }
}

// 主kernel函数
__global__ void grid_sampler_kernel(
    const float* input,           // [N, C, H, W]
    const float* grid,            // [N, H_out, W_out, 2]
    float* output,                // [N, C, H_out, W_out]
    int input_batch, int input_channels, int input_height, int input_width,
    int output_height, int output_width,
    GridSampleMode mode,
    GridSamplePaddingMode padding_mode,
    bool align_corners
) {
    int n = blockIdx.z;
    int c = blockIdx.y;
    int h = blockIdx.x * blockDim.x + threadIdx.x;
    int w = (blockIdx.x * blockDim.x + threadIdx.x) % output_width;
    h = (blockIdx.x * blockDim.x + threadIdx.x) / output_width;

    if (n >= input_batch || c >= input_channels ||
        h >= output_height || w >= output_width) {
        return;
    }

    // 计算输出索引
    int output_idx = n * input_channels * output_height * output_width +
                    c * output_height * output_width +
                    h * output_width + w;

    // 获取grid坐标 [x, y]
    int grid_idx = n * output_height * output_width * 2 +
                   h * output_width * 2 + w * 2;
    float grid_x = grid[grid_idx];
    float grid_y = grid[grid_idx + 1];

    // 坐标变换：从[-1, 1]映射到[0, input_width-1]和[0, input_height-1]
    float x, y;
    if (align_corners) {
        x = (grid_x + 1.0f) * (input_width - 1) / 2.0f;
        y = (grid_y + 1.0f) * (input_height - 1) / 2.0f;
    } else {
        x = (grid_x + 1.0f) * input_width / 2.0f - 0.5f;
        y = (grid_y + 1.0f) * input_height / 2.0f - 0.5f;
    }

    // 应用边界处理
    apply_padding_mode(x, y, input_width, input_height, padding_mode);

    // 根据模式进行插值
    float result;
    switch (mode) {
        case GridSampleMode::BILINEAR:
            result = bilinear_interpolate(input, input_height, input_width,
                                        y, x, c, n, input_height, input_width,
                                        input_channels, input_batch);
            break;
        case GridSampleMode::NEAREST:
            result = nearest_interpolate(input, input_height, input_width,
                                       y, x, c, n, input_height, input_width,
                                       input_channels, input_batch);
            break;
        default:
            result = 0.0f;
            break;
    }

    output[output_idx] = result;
}

// 主机端包装函数
extern "C" {
    void grid_sampler_cuda(
        const float* input,
        const float* grid,
        float* output,
        int input_batch, int input_channels, int input_height, int input_width,
        int output_height, int output_width,
        int mode, int padding_mode, int align_corners,
        cudaStream_t stream = 0
    ) {
    // 计算网格和块大小
    dim3 block_size(256);
    int total_output_pixels = output_height * output_width;
    dim3 grid_size(
        (total_output_pixels + block_size.x - 1) / block_size.x,
        input_channels,
        input_batch
    );

        // 启动kernel
        grid_sampler_kernel<<<grid_size, block_size, 0, stream>>>(
            input, grid, output,
            input_batch, input_channels, input_height, input_width,
            output_height, output_width,
            static_cast<GridSampleMode>(mode),
            static_cast<GridSamplePaddingMode>(padding_mode),
            static_cast<bool>(align_corners)
        );

        // 检查错误
        cudaError_t err = cudaGetLastError();
        if (err != cudaSuccess) {
            std::cerr << "CUDA kernel launch failed: " << cudaGetErrorString(err) << std::endl;
        }
    }
}