From 70308adfc10c50fba8104461f90fbc53f4bd98e8 Mon Sep 17 00:00:00 2001 From: Zaki Ahmed Date: Fri, 6 Feb 2026 07:46:51 +0100 Subject: [PATCH] fix: FP16 crash on CUDA 12.x + kernel optimizations Fix critical FP16 (half-precision) crash (CUDA error 700: illegal memory access) when running on CUDA 12.x, and apply kernel performance optimizations. Bug Fixes --------- 1. Remove conflicting __half operator overloads (grid_sample_3d.cuh) CUDA 12.x ships native __half arithmetic operators. The 12 custom operator overloads caused ODR violations at link time, leading to undefined behavior. One overload was inherently broken: __half operator+=(const float&, const half&) attempts to modify a const reference. All custom overloads removed. A minimal to_float() helper handles __half to float conversion. 2. Fix compute_index() precision loss for __half (grid_sample_3d.cuh) compute_index() was templated in scalar_t, so when instantiated with __half (~3.3 decimal digits), unnormalization and coordinate math produced garbage indices, causing out-of-bounds memory access. compute_index() now works entirely in float32 and returns float. reflect_coordinates() rewritten as reflect_coordinates_f() in pure float32. 3. Fix nearest kernel channel stride bug (grid_sample_3d.cu) The nearest interpolation kernel never advanced input/output channel pointers inside the channel loop (missing += stride_C). This caused channel-0 data to be read for all channels. Performance Optimizations ------------------------- - __restrict__ on all kernel and launcher pointer parameters - Hoist boundary checks outside channel loops (constant across C) - Precompute spatial offsets outside channel loops - Use __ldg() for read-only global memory loads (texture cache path) - Compute trilinear weights in float32 (avoids repeated conversions) - Accumulate bilinear interpolation in float32, cast on output Tested: CUDA 12.6 / TensorRT 10.5 / SM80+ (Ampere) FP32: max_diff vs PyTorch = 0.0 (exact match) FP16: max_diff vs PyTorch = 9.76e-4 (expected for half) FP16 latency: 93.9us vs FP32 117.5us = 1.25x speedup --- src/grid_sample_3d.cu | 176 ++++++++++++++++++++++++----------------- src/grid_sample_3d.cuh | 169 +++++++++++---------------------------- src/grid_sample_3d.h | 6 +- 3 files changed, 152 insertions(+), 199 deletions(-) diff --git a/src/grid_sample_3d.cu b/src/grid_sample_3d.cu index 5e3aeda..18a18b2 100644 --- a/src/grid_sample_3d.cu +++ b/src/grid_sample_3d.cu @@ -15,8 +15,8 @@ inline int get_num_blocks(int n) { template __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, @@ -24,7 +24,7 @@ __global__ void grid_sample_3d_nearest_kernel( 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; @@ -46,29 +46,34 @@ __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(::roundf(ix)); - int iy_nearest = static_cast(::roundf(iy)); - int iz_nearest = static_cast(::roundf(iz)); + int ix_nearest = static_cast(roundf(fix)); + int iy_nearest = static_cast(roundf(fiy)); + int iz_nearest = static_cast(roundf(fiz)); - scalar_t *input_NC_offset = const_cast(input_N_offset); + // Hoist boundary check and precompute offset outside channel loop + bool valid = (ix_nearest >= 0 && ix_nearest < static_cast(W_in) && + iy_nearest >= 0 && iy_nearest < static_cast(H_in) && + iz_nearest >= 0 && iz_nearest < static_cast(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(0); - } + for (size_t c = 0; c < C; c++) { + *output_NCDHW_offset = valid ? __ldg(&input_NC_offset[spatial_offset]) : static_cast(0); + input_NC_offset += input_stride_C; + output_NCDHW_offset += output_stride_C; } } template __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, @@ -76,7 +81,7 @@ __global__ void grid_sample_3d_bilinear_kernel( 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; @@ -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(floor(ix)); - int y0 = static_cast(floor(iy)); - int z0 = static_cast(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(floorf(fix)); + int y0 = static_cast(floorf(fiy)); + int z0 = static_cast(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(x1) - ix) * (iy - y0) * (iz - z0); - scalar_t v010 = (ix - x0) * (static_cast(y1) - iy) * (iz - z0); - scalar_t v110 = (static_cast(x1) - ix) * (static_cast(y1) - iy) * (iz - z0); - scalar_t v001 = (ix - x0) * (iy - y0) * (static_cast(z1) - iz); - scalar_t v101 = (static_cast(x1) - ix) * (iy - y0) * (static_cast(z1) - iz); - scalar_t v011 = (ix - x0) * (static_cast(y1) - iy) * (static_cast(z1) - iz); - scalar_t v111 = (static_cast(x1) - ix) * (static_cast(y1) - iy) * (static_cast(z1) - iz); - - scalar_t *input_NC_offset = const_cast(input_N_offset); + // Compute trilinear weights in float32 + float wx = fix - static_cast(x0); + float wy = fiy - static_cast(y0); + float wz = fiz - static_cast(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(W_in)); + bool in_x1 = (x1 >= 0 && x1 < static_cast(W_in)); + bool in_y0 = (y0 >= 0 && y0 < static_cast(H_in)); + bool in_y1 = (y1 >= 0 && y1 < static_cast(H_in)); + bool in_z0 = (z0 >= 0 && z0 < static_cast(D_in)); + bool in_z1 = (z1 >= 0 && z1 < static_cast(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(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(__ldg(&input_NC_offset[off_000])); + if(valid_100) value += w100 * static_cast(__ldg(&input_NC_offset[off_100])); + if(valid_010) value += w010 * static_cast(__ldg(&input_NC_offset[off_010])); + if(valid_110) value += w110 * static_cast(__ldg(&input_NC_offset[off_110])); + if(valid_001) value += w001 * static_cast(__ldg(&input_NC_offset[off_001])); + if(valid_101) value += w101 * static_cast(__ldg(&input_NC_offset[off_101])); + if(valid_011) value += w011 * static_cast(__ldg(&input_NC_offset[off_011])); + if(valid_111) value += w111 * static_cast(__ldg(&input_NC_offset[off_111])); + + *output_NCDHW_offset = static_cast(value); input_NC_offset += input_stride_C; output_NCDHW_offset += output_stride_C; - } - + } template 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 ) { @@ -234,25 +262,25 @@ int grid_sample_3d_cuda( // template specialization template int grid_sample_3d_cuda( - 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( - 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 ); \ No newline at end of file diff --git a/src/grid_sample_3d.cuh b/src/grid_sample_3d.cuh index 98584fd..a72a562 100644 --- a/src/grid_sample_3d.cuh +++ b/src/grid_sample_3d.cuh @@ -1,149 +1,74 @@ -// from https://github.com/TrojanXu/onnxparser-trt-plugin-sample/blob/master/TensorRT/plugin/gridSamplerPlugin/gridSampler.cuh +// Originally from https://github.com/TrojanXu/onnxparser-trt-plugin-sample/blob/master/TensorRT/plugin/gridSamplerPlugin/gridSampler.cuh +// Rewritten: all coordinate math now runs in float32 to avoid __half precision +// issues and to remove conflicting operator overloads (CUDA 12.6 provides native +// __half arithmetic). #pragma once -#include - #include #include #include "grid_sample_3d.h" -static __forceinline__ __device__ -__half operator*(const __half& a, const int& b) -{ - return a * __int2half_rn(b); -} - -static __forceinline__ __device__ -__half operator*(const float& a, const half& b) -{ - return __float2half(a) * b; -} - -static __forceinline__ __device__ -__half operator+=(const float& a, const half& b) -{ - return __float2half(a) += b; -} - -static __forceinline__ __device__ -__half operator/(const __half& a, const int& b) -{ - return a / __int2half_rn(b); -} - -static __forceinline__ __device__ -__half operator+(const __half& a, const float& b) -{ - return a + __float2half(b); -} - -static __forceinline__ __device__ -__half operator-(const __half& a, const int& b) -{ - return a - __int2half_rn(b); -} - -static __forceinline__ __device__ -__half operator-(const int& a, const __half& b) -{ - return __int2half_rn(a) - b; -} - -static __forceinline__ __device__ -__half operator+=(const __half& a, const __half& b) -{ - return a + b; -} - -static __forceinline__ __device__ -__half min(const __half& a, const half& b) -{ - return __float2half(min(__half2float(a), __half2float(b))); -} - -static __forceinline__ __device__ -__half max(const __half& a, const half& b) -{ - return __float2half(max(__half2float(a), __half2float(b))); -} - -static __forceinline__ __device__ -__half fabs(const __half& a) -{ - //TODO return __habs(a); what happened. - return __float2half(fabs(__half2float(a))); -} - -static __forceinline__ __device__ -__half floor(const __half& a) -{ - return hfloor(a); -} - -static __forceinline__ __device__ -__half roundf(const __half& a) -{ - return hrint(a); -} +// ── Helpers to convert any scalar_t to float and back ──────────────────────── +template +static __forceinline__ __device__ float to_float(scalar_t v) { return static_cast(v); } -static __forceinline__ __device__ -__half fmod(const __half& a, const __half& b) -{ - return __float2half(fmodf(__half2float(a), __half2float(b))); -} +template <> +__forceinline__ __device__ float to_float<__half>(__half v) { return __half2float(v); } // borrow from pytorch aten/src/Aten/native/GridSampler.h -template +// All arithmetic is in float32 – safe for both float and __half callers. static __forceinline__ __device__ -scalar_t reflect_coordinates(scalar_t in, int twice_low, - int twice_high) { - if (twice_low == twice_high) { - return static_cast(0); - } - scalar_t min = static_cast(twice_low) / 2; - scalar_t span = static_cast(twice_high - twice_low) / 2; - in = ::fabs(in - min); - // `fmod` returns same sign as `in`, which is positive after the `fabs` above. - scalar_t extra = ::fmod(in, span); - int flips = static_cast(::floor(in / span)); - if (flips % 2 == 0) { - return extra + min; - } else { - return span - extra + min; - } +float reflect_coordinates_f(float in, int twice_low, int twice_high) { + if (twice_low == twice_high) { + return 0.0f; + } + float min_val = static_cast(twice_low) * 0.5f; + float span = static_cast(twice_high - twice_low) * 0.5f; + in = fabsf(in - min_val); + // `fmodf` returns same sign as `in`, which is positive after the `fabsf` above. + float extra = fmodf(in, span); + int flips = static_cast(floorf(in / span)); + if (flips % 2 == 0) { + return extra + min_val; + } else { + return span - extra + min_val; + } } +// compute_index: maps normalized coordinate [-1, 1] → pixel index. +// Works entirely in float32 regardless of scalar_t. +// Returns float so the caller can stay in float32 for weight computation. template static __forceinline__ __device__ -scalar_t compute_index( +float compute_index( const scalar_t coord, const int size, - const GridSample3DPaddingMode padding_mode, + const GridSample3DPaddingMode padding_mode, const bool align_corners ) { - // unnormalize coord from [-1, 1] to [0, size - 1] if align_corners = False - // else unnormalize coord from [-1, 1] to [-0.5, size - 0.5] - scalar_t coord_; - if(align_corners) { - coord_ = ((coord + 1.f) / 2) * (size - 1); - }else { - coord_ = ((coord + 1.f) * size - 1) / 2; + float fcoord = to_float(coord); + + // unnormalize: [-1, 1] → [0, size-1] (align_corners) or [-0.5, size-0.5] + float idx; + if (align_corners) { + idx = (fcoord + 1.0f) * 0.5f * static_cast(size - 1); + } else { + idx = ((fcoord + 1.0f) * static_cast(size) - 1.0f) * 0.5f; } - // check if the coord_ is out of input boundary, - // if so, make it back to the boundary based on padding_mode - if (padding_mode == GridSample3DPaddingMode::Border) { // border mode, clip to [0, size-1] - coord_ = ::min(static_cast(size-1), ::max(coord_, static_cast(0))); - } else if (padding_mode == GridSample3DPaddingMode::Reflection) { // reflection mode - if(align_corners) { - coord_ = reflect_coordinates(coord_, 0, 2 * (size - 1)); + // clamp / reflect based on padding mode + if (padding_mode == GridSample3DPaddingMode::Border) { + idx = fminf(static_cast(size - 1), fmaxf(idx, 0.0f)); + } else if (padding_mode == GridSample3DPaddingMode::Reflection) { + if (align_corners) { + idx = reflect_coordinates_f(idx, 0, 2 * (size - 1)); } else { - coord_ = reflect_coordinates(coord_, -1, 2 * size - 1); + idx = reflect_coordinates_f(idx, -1, 2 * size - 1); } - coord_ = ::min(static_cast(size-1), ::max(coord_, static_cast(0))); + idx = fminf(static_cast(size - 1), fmaxf(idx, 0.0f)); } - return coord_; + return idx; } \ No newline at end of file diff --git a/src/grid_sample_3d.h b/src/grid_sample_3d.h index 6c4d5e7..f4de6d1 100644 --- a/src/grid_sample_3d.h +++ b/src/grid_sample_3d.h @@ -11,14 +11,14 @@ enum class GridSample3DDataType {GFLOAT, GHALF}; template 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 );