Skip to content

fix: FP16 crash on CUDA 12.x + kernel optimizations #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zaki699-blip wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: FP16 crash on CUDA 12.x + kernel optimizations #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
176 changes: 102 additions & 74 deletions src/grid_sample_3d.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,16 +15,16 @@ inline int get_num_blocks(int n) {

template <typename scalar_t>
__global__ void grid_sample_3d_nearest_kernel(
const scalar_t* input,
const scalar_t* grid,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ grid,
size_t N, size_t C, size_t D_in, size_t H_in, size_t W_in,
size_t input_stride_N, size_t input_stride_C, size_t input_stride_D, size_t input_stride_H, size_t input_stride_W,
size_t D_grid, size_t H_grid, size_t W_grid,
size_t grid_stride_N, size_t grid_stride_D, size_t grid_stride_H, size_t grid_stride_W, size_t grid_stride_XYZ,
size_t output_stride_N, size_t output_stride_C, size_t output_stride_D, size_t output_stride_H, size_t output_stride_W,
bool align_corners,
GridSample3DPaddingMode padding_mode,
scalar_t* output
scalar_t* __restrict__ output
) {
unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;

Expand All @@ -46,37 +46,42 @@ __global__ void grid_sample_3d_nearest_kernel(
const scalar_t y = *(grid_NDHW_offset + grid_stride_XYZ);
const scalar_t z = *(grid_NDHW_offset + 2 * grid_stride_XYZ);

scalar_t ix = compute_index(x, W_in, padding_mode, align_corners);
scalar_t iy = compute_index(y, H_in, padding_mode, align_corners);
scalar_t iz = compute_index(z, D_in, padding_mode, align_corners);
// compute_index returns float — stay in float for rounding
float fix = compute_index(x, W_in, padding_mode, align_corners);
float fiy = compute_index(y, H_in, padding_mode, align_corners);
float fiz = compute_index(z, D_in, padding_mode, align_corners);

int ix_nearest = static_cast<int>(::roundf(ix));
int iy_nearest = static_cast<int>(::roundf(iy));
int iz_nearest = static_cast<int>(::roundf(iz));
int ix_nearest = static_cast<int>(roundf(fix));
int iy_nearest = static_cast<int>(roundf(fiy));
int iz_nearest = static_cast<int>(roundf(fiz));

scalar_t *input_NC_offset = const_cast<scalar_t *>(input_N_offset);
// Hoist boundary check and precompute offset outside channel loop
bool valid = (ix_nearest >= 0 && ix_nearest < static_cast<int>(W_in) &&
iy_nearest >= 0 && iy_nearest < static_cast<int>(H_in) &&
iz_nearest >= 0 && iz_nearest < static_cast<int>(D_in));
size_t spatial_offset = ix_nearest * input_stride_W + iy_nearest * input_stride_H + iz_nearest * input_stride_D;

const scalar_t *input_NC_offset = input_N_offset;
scalar_t *output_NCDHW_offset = output_N_offset + d * output_stride_D + h * output_stride_H + w * output_stride_W;
for (auto c = 0; c < C; c++) {
if(ix_nearest >= 0 && ix_nearest < W_in && iy_nearest >= 0 && iy_nearest < H_in && iz_nearest >= 0 && iz_nearest < D_in) {
*output_NCDHW_offset = input_NC_offset[ix_nearest * input_stride_W + iy_nearest * input_stride_H + iz_nearest * input_stride_D];
} else {
*output_NCDHW_offset = static_cast<scalar_t>(0);
}
for (size_t c = 0; c < C; c++) {
*output_NCDHW_offset = valid ? __ldg(&input_NC_offset[spatial_offset]) : static_cast<scalar_t>(0);
input_NC_offset += input_stride_C;
output_NCDHW_offset += output_stride_C;
}
}

template <typename scalar_t>
__global__ void grid_sample_3d_bilinear_kernel(
const scalar_t* input,
const scalar_t* grid,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ grid,
size_t N, size_t C, size_t D_in, size_t H_in, size_t W_in,
size_t input_stride_N, size_t input_stride_C, size_t input_stride_D, size_t input_stride_H, size_t input_stride_W,
size_t D_grid, size_t H_grid, size_t W_grid,
size_t grid_stride_N, size_t grid_stride_D, size_t grid_stride_H, size_t grid_stride_W, size_t grid_stride_XYZ,
size_t output_stride_N, size_t output_stride_C, size_t output_stride_D, size_t output_stride_H, size_t output_stride_W,
bool align_corners,
GridSample3DPaddingMode padding_mode,
scalar_t* output
scalar_t* __restrict__ output
) {

unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
Expand All @@ -99,73 +104,96 @@ __global__ void grid_sample_3d_bilinear_kernel(
const scalar_t y = *(grid_NDHW_offset + grid_stride_XYZ);
const scalar_t z = *(grid_NDHW_offset + 2 * grid_stride_XYZ);

scalar_t ix = compute_index(x, W_in, padding_mode, align_corners);
scalar_t iy = compute_index(y, H_in, padding_mode, align_corners);
scalar_t iz = compute_index(z, D_in, padding_mode, align_corners);

int x0 = static_cast<int>(floor(ix));
int y0 = static_cast<int>(floor(iy));
int z0 = static_cast<int>(floor(iz));
// compute_index now returns float directly — no conversion needed
float fix = compute_index(x, W_in, padding_mode, align_corners);
float fiy = compute_index(y, H_in, padding_mode, align_corners);
float fiz = compute_index(z, D_in, padding_mode, align_corners);

int x0 = static_cast<int>(floorf(fix));
int y0 = static_cast<int>(floorf(fiy));
int z0 = static_cast<int>(floorf(fiz));
int x1 = x0 + 1;
int y1 = y0 + 1;
int z1 = z0 + 1;

scalar_t v000 = (ix - x0) * (iy - y0) * (iz - z0);
scalar_t v100 = (static_cast<scalar_t>(x1) - ix) * (iy - y0) * (iz - z0);
scalar_t v010 = (ix - x0) * (static_cast<scalar_t>(y1) - iy) * (iz - z0);
scalar_t v110 = (static_cast<scalar_t>(x1) - ix) * (static_cast<scalar_t>(y1) - iy) * (iz - z0);
scalar_t v001 = (ix - x0) * (iy - y0) * (static_cast<scalar_t>(z1) - iz);
scalar_t v101 = (static_cast<scalar_t>(x1) - ix) * (iy - y0) * (static_cast<scalar_t>(z1) - iz);
scalar_t v011 = (ix - x0) * (static_cast<scalar_t>(y1) - iy) * (static_cast<scalar_t>(z1) - iz);
scalar_t v111 = (static_cast<scalar_t>(x1) - ix) * (static_cast<scalar_t>(y1) - iy) * (static_cast<scalar_t>(z1) - iz);

scalar_t *input_NC_offset = const_cast<scalar_t *>(input_N_offset);
// Compute trilinear weights in float32
float wx = fix - static_cast<float>(x0);
float wy = fiy - static_cast<float>(y0);
float wz = fiz - static_cast<float>(z0);
float wx1 = 1.0f - wx;
float wy1 = 1.0f - wy;
float wz1 = 1.0f - wz;

// NOTE: weight naming matches corner positions (opposite of original)
// w_abc = weight for corner (xa, yb, zc) where a,b,c index into {0,1}
float w000 = wx1 * wy1 * wz1;
float w100 = wx * wy1 * wz1;
float w010 = wx1 * wy * wz1;
float w110 = wx * wy * wz1;
float w001 = wx1 * wy1 * wz;
float w101 = wx * wy1 * wz;
float w011 = wx1 * wy * wz;
float w111 = wx * wy * wz;

// Hoist boundary checks outside channel loop — coordinates are the same for all C
bool in_x0 = (x0 >= 0 && x0 < static_cast<int>(W_in));
bool in_x1 = (x1 >= 0 && x1 < static_cast<int>(W_in));
bool in_y0 = (y0 >= 0 && y0 < static_cast<int>(H_in));
bool in_y1 = (y1 >= 0 && y1 < static_cast<int>(H_in));
bool in_z0 = (z0 >= 0 && z0 < static_cast<int>(D_in));
bool in_z1 = (z1 >= 0 && z1 < static_cast<int>(D_in));

bool valid_000 = in_x0 && in_y0 && in_z0;
bool valid_100 = in_x1 && in_y0 && in_z0;
bool valid_010 = in_x0 && in_y1 && in_z0;
bool valid_110 = in_x1 && in_y1 && in_z0;
bool valid_001 = in_x0 && in_y0 && in_z1;
bool valid_101 = in_x1 && in_y0 && in_z1;
bool valid_011 = in_x0 && in_y1 && in_z1;
bool valid_111 = in_x1 && in_y1 && in_z1;

// Precompute spatial offsets within each channel plane (constant across C)
size_t off_000 = x0 * input_stride_W + y0 * input_stride_H + z0 * input_stride_D;
size_t off_100 = x1 * input_stride_W + y0 * input_stride_H + z0 * input_stride_D;
size_t off_010 = x0 * input_stride_W + y1 * input_stride_H + z0 * input_stride_D;
size_t off_110 = x1 * input_stride_W + y1 * input_stride_H + z0 * input_stride_D;
size_t off_001 = x0 * input_stride_W + y0 * input_stride_H + z1 * input_stride_D;
size_t off_101 = x1 * input_stride_W + y0 * input_stride_H + z1 * input_stride_D;
size_t off_011 = x0 * input_stride_W + y1 * input_stride_H + z1 * input_stride_D;
size_t off_111 = x1 * input_stride_W + y1 * input_stride_H + z1 * input_stride_D;

const scalar_t *input_NC_offset = input_N_offset;
scalar_t *output_NCDHW_offset = output_N_offset + d * output_stride_D + h * output_stride_H + w * output_stride_W;

for(auto c = 0; c < C; c++) {
scalar_t value = static_cast<scalar_t>(0);
if(x1 >= 0 && x1 < W_in && y1 >= 0 && y1 < H_in && z1 >= 0 && z1 < D_in) {
value += v000 * input_NC_offset[x1 * input_stride_W + y1 * input_stride_H + z1 * input_stride_D];
}
if(x0 >= 0 && x0 < W_in && y1 >= 0 && y1 < H_in && z1 >= 0 && z1 < D_in) {
value += v100 * input_NC_offset[x0 * input_stride_W + y1 * input_stride_H + z1 * input_stride_D];
}
if(x1 >= 0 && x1 < W_in && y0 >= 0 && y0 < H_in && z1 >= 0 && z1 < D_in) {
value += v010 * input_NC_offset[x1 * input_stride_W + y0 * input_stride_H + z1 * input_stride_D];
}
if(x0 >= 0 && x0 < W_in && y0 >= 0 && y0 < H_in && z1 >= 0 && z1 < D_in) {
value += v110 * input_NC_offset[x0 * input_stride_W + y0 * input_stride_H + z1 * input_stride_D];
}
if(x1 >= 0 && x1 < W_in && y1 >= 0 && y1 < H_in && z0 >= 0 && z0 < D_in) {
value += v001 * input_NC_offset[x1 * input_stride_W + y1 * input_stride_H + z0 * input_stride_D];
}
if(x0 >= 0 && x0 < W_in && y1 >= 0 && y1 < H_in && z0 >= 0 && z0 < D_in) {
value += v101 * input_NC_offset[x0 * input_stride_W + y1 * input_stride_H + z0 * input_stride_D];
}
if(x1 >= 0 && x1 < W_in && y0 >= 0 && y0 < H_in && z0 >= 0 && z0 < D_in) {
value += v011 * input_NC_offset[x1 * input_stride_W + y0 * input_stride_H + z0 * input_stride_D];
}
if(x0 >= 0 && x0 < W_in && y0 >= 0 && y0 < H_in && z0 >= 0 && z0 < D_in) {
value += v111 * input_NC_offset[x0 * input_stride_W + y0 * input_stride_H + z0 * input_stride_D];
}
*output_NCDHW_offset = value;
for(size_t c = 0; c < C; c++) {
// Accumulate in float32 for precision, use __ldg for read-only cache path
float value = 0.0f;
if(valid_000) value += w000 * static_cast<float>(__ldg(&input_NC_offset[off_000]));
if(valid_100) value += w100 * static_cast<float>(__ldg(&input_NC_offset[off_100]));
if(valid_010) value += w010 * static_cast<float>(__ldg(&input_NC_offset[off_010]));
if(valid_110) value += w110 * static_cast<float>(__ldg(&input_NC_offset[off_110]));
if(valid_001) value += w001 * static_cast<float>(__ldg(&input_NC_offset[off_001]));
if(valid_101) value += w101 * static_cast<float>(__ldg(&input_NC_offset[off_101]));
if(valid_011) value += w011 * static_cast<float>(__ldg(&input_NC_offset[off_011]));
if(valid_111) value += w111 * static_cast<float>(__ldg(&input_NC_offset[off_111]));

*output_NCDHW_offset = static_cast<scalar_t>(value);
input_NC_offset += input_stride_C;
output_NCDHW_offset += output_stride_C;

}

}

template <typename scalar_t>
int grid_sample_3d_cuda(
const scalar_t* input,
const scalar_t* grid,
const scalar_t* __restrict__ input,
const scalar_t* __restrict__ grid,
size_t N, size_t C, size_t D_in, size_t H_in, size_t W_in,
size_t D_grid, size_t H_grid, size_t W_grid,
bool align_corners,
GridSample3DInterpolationMode interpolationMode,
GridSample3DPaddingMode paddingMode,
scalar_t* output,
scalar_t* __restrict__ output,
cudaStream_t stream
) {

Expand Down Expand Up @@ -234,25 +262,25 @@ int grid_sample_3d_cuda(

// template specialization
template int grid_sample_3d_cuda<float>(
const float* input,
const float* grid,
const float* __restrict__ input,
const float* __restrict__ grid,
size_t N, size_t C, size_t D_in, size_t H_in, size_t W_in,
size_t D_grid, size_t H_grid, size_t W_grid,
bool align_corners,
GridSample3DInterpolationMode interpolationMode,
GridSample3DPaddingMode paddingMode,
float* output,
float* __restrict__ output,
cudaStream_t stream
);

template int grid_sample_3d_cuda<half>(
const half* input,
const half* grid,
const half* __restrict__ input,
const half* __restrict__ grid,
size_t N, size_t C, size_t D_in, size_t H_in, size_t W_in,
size_t D_grid, size_t H_grid, size_t W_grid,
bool align_corners,
GridSample3DInterpolationMode interpolationMode,
GridSample3DPaddingMode paddingMode,
half* output,
half* __restrict__ output,
cudaStream_t stream
);
Loading