CUDA: handle OW > 65535 in im2col (2D and 3D) (#22944)

`im2col_cuda` and `im2col_3d_cuda` both dispatch with
`block_nums.y = OW`. CUDA caps grid Y at 65535. Conv1d encoders on
raw 16 kHz audio with T > 65535 (~ 4 s) trip the limit -- e.g. SEANet
at 11 s lands at OW = 176000 -- and the launch returns
`invalid configuration argument`.

Clamp `block_nums.y` to `MIN(OW, MAX_GRIDDIM_Y)` and loop inside the
kernel with stride `MAX_GRIDDIM_Y`. Same in-kernel stride pattern
already used for the z axis (`MAX_GRIDDIM_Z`). Both 2D `im2col_kernel`
and 3D `im2col_3d_kernel` need the same fix. Bit-identical for
OW <= 65535 (single iteration of the new outer loop).

Tested on T4 / Jetson Orin with a SEANet encoder running on 11 s /
16 kHz audio (im2col reaching OW ~ 176000); pre-fix launch returns
`invalid configuration argument`, post-fix runs to completion.
Existing test-backend-ops im2col cases unchanged.
This commit is contained in:
CrispStrobe
2026-05-11 19:48:29 +02:00
committed by GitHub
parent e936660760
commit 8e1f9d0834
+9 -6
View File
@@ -1,5 +1,6 @@
#include "im2col.cuh" #include "im2col.cuh"
#define MAX_GRIDDIM_Y 65535
#define MAX_GRIDDIM_Z 65535 #define MAX_GRIDDIM_Z 65535
template <typename T> template <typename T>
@@ -18,8 +19,8 @@ static __global__ void im2col_kernel(
const int64_t ikh = rem / KW; const int64_t ikh = rem / KW;
const int64_t ikw = rem - ikh * KW; const int64_t ikw = rem - ikh * KW;
const int64_t iow = blockIdx.y; for (int64_t iow = blockIdx.y; iow < OW; iow += MAX_GRIDDIM_Y) {
for (int64_t iz = blockIdx.z; iz < N_OH; iz+=MAX_GRIDDIM_Z) { for (int64_t iz = blockIdx.z; iz < N_OH; iz += MAX_GRIDDIM_Z) {
const int64_t in = iz / OH; const int64_t in = iz / OH;
const int64_t ioh = iz - in * OH; const int64_t ioh = iz - in * OH;
@@ -36,6 +37,7 @@ static __global__ void im2col_kernel(
dst[offset_dst] = x[offset_src + iih * IW + iiw]; dst[offset_dst] = x[offset_src + iih * IW + iiw];
} }
} }
}
GGML_UNUSED(IC); GGML_UNUSED(IC);
GGML_UNUSED(KH); GGML_UNUSED(KH);
@@ -51,7 +53,7 @@ static void im2col_cuda(const float * x, T* dst,
const int64_t num_blocks = (IC_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE; const int64_t num_blocks = (IC_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE;
const int64_t N_OH = N * OH; const int64_t N_OH = N * OH;
const int64_t KH_KW = KW*KH; const int64_t KH_KW = KW*KH;
dim3 block_nums(num_blocks, OW, MIN(N_OH, MAX_GRIDDIM_Z)); dim3 block_nums(num_blocks, MIN(OW, MAX_GRIDDIM_Y), MIN(N_OH, MAX_GRIDDIM_Z));
im2col_kernel<<<block_nums, MIN(IC_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(x, dst, IC, IW, IH, OH, OW, KW, KH, im2col_kernel<<<block_nums, MIN(IC_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(x, dst, IC, IW, IH, OH, OW, KW, KH,
IC_IH_IW, IH_IW, N_OH, KH_KW, IC_KH_KW, IC_IH_IW, IH_IW, N_OH, KH_KW, IC_KH_KW,
s0, s1, p0, p1, d0, d1); s0, s1, p0, p1, d0, d1);
@@ -136,8 +138,8 @@ static __global__ void im2col_3d_kernel(
const int64_t ikh = (i - iic * KD_KH_KW - ikd * KH_KW) / KW; const int64_t ikh = (i - iic * KD_KH_KW - ikd * KH_KW) / KW;
const int64_t ikw = i % KW; const int64_t ikw = i % KW;
const int64_t iow = blockIdx.y; for (int64_t iow = blockIdx.y; iow < OW; iow += MAX_GRIDDIM_Y) {
for (int64_t iz = blockIdx.z; iz < N_OD_OH; iz+=MAX_GRIDDIM_Z) { for (int64_t iz = blockIdx.z; iz < N_OD_OH; iz += MAX_GRIDDIM_Z) {
const int64_t in = iz / OD_OH; const int64_t in = iz / OD_OH;
const int64_t iod = (iz - in*OD_OH) / OH; const int64_t iod = (iz - in*OD_OH) / OH;
const int64_t ioh = iz % OH; const int64_t ioh = iz % OH;
@@ -155,6 +157,7 @@ static __global__ void im2col_3d_kernel(
dst[offset_dst] = src[offset_src]; dst[offset_dst] = src[offset_src];
} }
} }
}
} }
// [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW] // [N*IC, ID, IH, IW] => [N*OD, OH, OW, IC * KD * KH * KW]
@@ -178,7 +181,7 @@ static void im2col_3d_cuda(const float * src, T* dst,
const int64_t OH_OW_IC_KD_KH_KW = OH*OW*IC*KD*KH*KW; const int64_t OH_OW_IC_KD_KH_KW = OH*OW*IC*KD*KH*KW;
const int64_t OW_IC_KD_KH_KW = OW*IC*KD*KH*KW; const int64_t OW_IC_KD_KH_KW = OW*IC*KD*KH*KW;
const int64_t num_blocks = (IC_KD_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE; const int64_t num_blocks = (IC_KD_KH_KW + CUDA_IM2COL_BLOCK_SIZE - 1) / CUDA_IM2COL_BLOCK_SIZE;
dim3 block_nums(num_blocks, OW, MIN(N_OD_OH, MAX_GRIDDIM_Z)); dim3 block_nums(num_blocks, MIN(OW, MAX_GRIDDIM_Y), MIN(N_OD_OH, MAX_GRIDDIM_Z));
im2col_3d_kernel<<<block_nums, MIN(IC_KD_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW, im2col_3d_kernel<<<block_nums, MIN(IC_KD_KH_KW, CUDA_IM2COL_BLOCK_SIZE) , 0, stream>>>(src, dst, N, IC, ID, IH, IW, OC, KD, KH, KW, OD, OH, OW,
OH_OW, KD_KH_KW, ID_IH_IW, KH_KW, IH_IW, IC_ID_IH_IW, OH_OW, KD_KH_KW, ID_IH_IW, KH_KW, IH_IW, IC_ID_IH_IW,
IC_KD_KH_KW, OW_KD_KH_KW, OD_OH_OW_IC_KD_KH_KW, IC_KD_KH_KW, OW_KD_KH_KW, OD_OH_OW_IC_KD_KH_KW,