diff --git a/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu b/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu index 2206941..2fc4da7 100644 --- a/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu +++ b/crates/tropical-gemm-cuda/kernels/tropical_gemm.cu @@ -40,8 +40,16 @@ #define INF_I64 (1LL << 60) #define NEG_INF_I64 (-(1LL << 60)) -// Memory layout helpers +// Memory layout helpers. +// OFFSET_COL is `int` arithmetic: use it ONLY for shared-memory tiles (As/Bs) +// and per-thread accumulator indices, whose leading dimension is a tile size, so +// the product can never reach 2^31. OFFSET_COL64 forces the multiply into 64-bit +// (32x32->64, a single IMAD.WIDE on the device) for GLOBAL-memory accesses, where +// a single matrix with >= 2^31 elements (col*ld) would otherwise overflow `int` +// and corrupt addressing (issue #63). The cast wraps the multiply itself, not the +// finished sum -- `(long long)(col*ld+row)` would already have overflowed. #define OFFSET_COL(row, col, ld) ((col) * (ld) + (row)) +#define OFFSET_COL64(row, col, ld) ((long long)(col) * (ld) + (row)) // Integer max/min functions __device__ __forceinline__ int max_i32(int a, int b) { return a > b ? a : b; } @@ -131,9 +139,9 @@ __device__ double atomicAddDouble(double* address, double val) { int col = A_TILE_COL + i + tile_idx; \ int dst = OFFSET_COL(A_TILE_ROW, i + A_TILE_COL, BLOCK_SIZE_M); \ if (BLOCK_IDX == DIM_GRID_X - 1 || tile_idx >= K - BLOCK_SIZE_K) \ - As[dst] = (row < M && col < K) ? SRC[OFFSET_COL(row, col, M)] : (PAD); \ + As[dst] = (row < M && col < K) ? SRC[OFFSET_COL64(row, col, M)] : (PAD); \ else \ - As[dst] = SRC[OFFSET_COL(row, col, M)]; \ + As[dst] = SRC[OFFSET_COL64(row, col, M)]; \ } // B tile -> Bs. B is column-major with leading dimension K. @@ -144,9 +152,9 @@ __device__ double atomicAddDouble(double* address, double val) { int col = BLOCK_SIZE_N * BLOCK_IDY + i + B_TILE_COL; \ int dst = OFFSET_COL(B_TILE_ROW, i + B_TILE_COL, BLOCK_SIZE_K); \ if (tile_idx >= K - BLOCK_SIZE_K || BLOCK_IDY == DIM_GRID_Y - 1) \ - Bs[dst] = (row < K && col < N) ? SRC[OFFSET_COL(row, col, K)] : (PAD); \ + Bs[dst] = (row < K && col < N) ? SRC[OFFSET_COL64(row, col, K)] : (PAD); \ else \ - Bs[dst] = SRC[OFFSET_COL(row, col, K)]; \ + Bs[dst] = SRC[OFFSET_COL64(row, col, K)]; \ } // Edge block iff last block-row or last block-col; interior blocks are fully in @@ -162,7 +170,7 @@ __device__ double atomicAddDouble(double* address, double val) { int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \ int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \ if (!TILE_IS_EDGE || (row < M && col < N)) \ - DST[OFFSET_COL(row, col, M)] = accum[OFFSET_COL(tm, tn, THREAD_SIZE_M)]; \ + DST[OFFSET_COL64(row, col, M)] = accum[OFFSET_COL(tm, tn, THREAD_SIZE_M)]; \ } \ } @@ -175,7 +183,7 @@ __device__ double atomicAddDouble(double* address, double val) { int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \ int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \ if (!TILE_IS_EDGE || (row < M && col < N)) { \ - int out_idx = OFFSET_COL(row, col, M); \ + long long out_idx = OFFSET_COL64(row, col, M); \ int local_idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ DST[out_idx] = accum[local_idx]; \ ARGMAX_DST[out_idx] = accum_idx[local_idx]; \ @@ -193,7 +201,7 @@ __device__ double atomicAddDouble(double* address, double val) { int row = BLOCK_SIZE_M * BLOCK_IDX + THREAD_SIZE_M * threadIdx.x + tm; \ int col = BLOCK_SIZE_N * BLOCK_IDY + THREAD_SIZE_N * threadIdx.y + tn; \ if (!TILE_IS_EDGE || (row < M && col < N)) { \ - int out_idx = OFFSET_COL(row, col, M); \ + long long out_idx = OFFSET_COL64(row, col, M); \ int local_idx = OFFSET_COL(tm, tn, THREAD_SIZE_M); \ DST[out_idx] = accum[local_idx]; \ ARGMAX_DST[out_idx] = (ZERO_FN)(accum[local_idx]) ? 0 : accum_idx[local_idx]; \ @@ -678,13 +686,13 @@ extern "C" __global__ void KERNEL_NAME( \ TYPE* __restrict__ grad_a, \ int M, int N, int K \ ) { \ - int idx = blockIdx.x * blockDim.x + threadIdx.x; \ - int total = M * N; \ + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; \ + long long total = (long long)M * N; \ if (idx < total) { \ - int i = idx % M; \ + int i = (int)(idx % M); \ int k = (int)argmax[idx]; \ if (k >= 0 && k < K) { \ - ATOMIC_ADD(&grad_a[i + k * M], grad_c[idx]); \ + ATOMIC_ADD(&grad_a[i + (long long)k * M], grad_c[idx]); \ } \ } \ } @@ -696,13 +704,13 @@ extern "C" __global__ void KERNEL_NAME( \ TYPE* __restrict__ grad_b, \ int M, int N, int K \ ) { \ - int idx = blockIdx.x * blockDim.x + threadIdx.x; \ - int total = M * N; \ + long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x; \ + long long total = (long long)M * N; \ if (idx < total) { \ - int j = idx / M; \ + int j = (int)(idx / M); \ int k = (int)argmax[idx]; \ if (k >= 0 && k < K) { \ - ATOMIC_ADD(&grad_b[k + j * K], grad_c[idx]); \ + ATOMIC_ADD(&grad_b[k + (long long)j * K], grad_c[idx]); \ } \ } \ } @@ -771,7 +779,7 @@ extern "C" __global__ void KERNEL_NAME( \ float* __restrict__ C, \ unsigned int* __restrict__ argmax, \ int M, int N, int K, \ - int strideA, int strideB, int strideC \ + long long strideA, long long strideB, long long strideC \ ) { \ const int BLOCK_SIZE_M = 64; \ const int BLOCK_SIZE_K = 32; \ @@ -783,7 +791,7 @@ extern "C" __global__ void KERNEL_NAME( \ const int bszn = BLOCK_SIZE_N / THREAD_SIZE_N; \ const int THREAD_NUM_PER_BLOCK = bszm * bszn; \ \ - int batch_idx = blockIdx.z; \ + long long batch_idx = blockIdx.z; \ int DIM_GRID_X = (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M; \ int DIM_GRID_Y = (N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N; \ int BLOCK_IDX = blockIdx.x % DIM_GRID_X; \ @@ -867,56 +875,58 @@ TROPICAL_GEMM_BATCHED_ARGMAX_F32(tropical_maxmul_f32_nn_batched_with_argmax, 0. // Same shared body as the single-matrix kernels; only the operand base pointers // differ -- each advances by its per-batch stride. This replaces omeinsum's // host-side per-slice clone_dtod loop with a single launch over already- -// contiguous device buffers. `(int)blockIdx.z` matches the original `int -// batch_idx = blockIdx.z` indexing; the host chunks the batch under the -// gridDim.z cap and pre-offsets each chunk's base, so blockIdx.z stays in range. +// contiguous device buffers. The per-batch stride is `long long` and the base +// offset `(long long)blockIdx.z * stride` is computed in 64-bit, so a single +// matrix of >= 2^31 elements addresses correctly (issue #63); the host chunks +// the batch under the gridDim.z cap and pre-offsets each chunk's base, so +// blockIdx.z stays in range. #define TROPICAL_GEMM_BATCHED_F32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ extern "C" __global__ void KERNEL_NAME( \ const float* __restrict__ A, const float* __restrict__ B, \ float* __restrict__ C, int M, int N, int K, \ - int strideA, int strideB, int strideC \ + long long strideA, long long strideB, long long strideC \ ) { \ TROPICAL_GEMM_BODY(float, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, \ - (A + (int)blockIdx.z * strideA), \ - (B + (int)blockIdx.z * strideB), \ - (C + (int)blockIdx.z * strideC)) \ + (A + (long long)blockIdx.z * strideA), \ + (B + (long long)blockIdx.z * strideB), \ + (C + (long long)blockIdx.z * strideC)) \ } #define TROPICAL_GEMM_BATCHED_F64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ extern "C" __global__ void KERNEL_NAME( \ const double* __restrict__ A, const double* __restrict__ B, \ double* __restrict__ C, int M, int N, int K, \ - int strideA, int strideB, int strideC \ + long long strideA, long long strideB, long long strideC \ ) { \ TROPICAL_GEMM_BODY(double, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, \ - (A + (int)blockIdx.z * strideA), \ - (B + (int)blockIdx.z * strideB), \ - (C + (int)blockIdx.z * strideC)) \ + (A + (long long)blockIdx.z * strideA), \ + (B + (long long)blockIdx.z * strideB), \ + (C + (long long)blockIdx.z * strideC)) \ } #define TROPICAL_GEMM_BATCHED_I32(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ extern "C" __global__ void KERNEL_NAME( \ const int* __restrict__ A, const int* __restrict__ B, \ int* __restrict__ C, int M, int N, int K, \ - int strideA, int strideB, int strideC \ + long long strideA, long long strideB, long long strideC \ ) { \ TROPICAL_GEMM_BODY(int, 64, 32, 64, INIT_VAL, COMPARE_FN, MUL_FN, \ - (A + (int)blockIdx.z * strideA), \ - (B + (int)blockIdx.z * strideB), \ - (C + (int)blockIdx.z * strideC)) \ + (A + (long long)blockIdx.z * strideA), \ + (B + (long long)blockIdx.z * strideB), \ + (C + (long long)blockIdx.z * strideC)) \ } #define TROPICAL_GEMM_BATCHED_I64(KERNEL_NAME, INIT_VAL, COMPARE_FN, MUL_FN) \ extern "C" __global__ void KERNEL_NAME( \ const long long* __restrict__ A, const long long* __restrict__ B, \ long long* __restrict__ C, int M, int N, int K, \ - int strideA, int strideB, int strideC \ + long long strideA, long long strideB, long long strideC \ ) { \ TROPICAL_GEMM_BODY(long long, 32, 16, 32, INIT_VAL, COMPARE_FN, MUL_FN, \ - (A + (int)blockIdx.z * strideA), \ - (B + (int)blockIdx.z * strideB), \ - (C + (int)blockIdx.z * strideC)) \ + (A + (long long)blockIdx.z * strideA), \ + (B + (long long)blockIdx.z * strideB), \ + (C + (long long)blockIdx.z * strideC)) \ } // --- Forward batched GEMM kernel instantiations (mirror the non-batched set) --- diff --git a/crates/tropical-gemm-cuda/src/kernels.rs b/crates/tropical-gemm-cuda/src/kernels.rs index 23fb666..a2e8abf 100644 --- a/crates/tropical-gemm-cuda/src/kernels.rs +++ b/crates/tropical-gemm-cuda/src/kernels.rs @@ -30,17 +30,17 @@ use tropical_gemm::types::{TropicalMaxMul, TropicalMaxPlus, TropicalMinPlus, Tro /// most this size (see `launch_kernel_batched_impl`). const MAX_GRID_DIM_Z: usize = 65535; -/// Convert a per-batch element stride to the `i32` the kernels take, failing +/// Convert a per-batch element stride to the `i64` the kernels take, failing /// instead of truncating. The chunked launches derive operand base offsets from -/// the `usize` stride while the kernel reads the stride as `i32`; an `as i32` -/// cast that wrapped would desynchronise the two and silently corrupt -/// addressing. The kernel signature is `i32`, so a stride past `i32::MAX` is -/// unrepresentable on the device regardless — reject it at the boundary. -fn stride_to_i32(stride: usize, what: &str) -> Result { - i32::try_from(stride).map_err(|_| { +/// the `usize` stride while the kernel reads the stride as `i64` (`long long`); +/// keeping the conversion fallible guarantees the two can never desynchronise. +/// The kernels do 64-bit global addressing (issue #63), so there is no longer an +/// `i32::MAX` ceiling — only a `usize` value past `i64::MAX` (≥ 2⁶³ elements, +/// physically impossible to allocate) would fail here. +fn stride_to_i64(stride: usize, what: &str) -> Result { + i64::try_from(stride).map_err(|_| { CudaError::DimensionMismatch(format!( - "batched GEMM {what} stride {stride} exceeds i32::MAX; \ - matrix too large for the strided-batched kernel" + "batched GEMM {what} stride {stride} exceeds i64::MAX" )) }) } @@ -172,13 +172,13 @@ fn launch_kernel_batched_impl let n_i32 = n as i32; let k_i32 = k as i32; // Per-batch element extents, used both to offset each chunk's operand base - // (as `usize`) and as the kernel's `i32` stride. Derive the `i32` form - // fallibly so a stride past `i32::MAX` can't truncate out of sync with the - // `usize` offsets below (a.len()/b.len()/c.len() already fit in `usize`). + // (as `usize`) and as the kernel's `i64` stride. Derive the `i64` form + // fallibly so it stays in sync with the `usize` offsets below (a.len()/ + // b.len()/c.len() already fit in `usize`). let (sa, sb, sc) = (m * k, k * n, m * n); - let stride_a = stride_to_i32(sa, "operand A")?; - let stride_b = stride_to_i32(sb, "operand B")?; - let stride_c = stride_to_i32(sc, "output C")?; + let stride_a = stride_to_i64(sa, "operand A")?; + let stride_b = stride_to_i64(sb, "operand B")?; + let stride_c = stride_to_i64(sc, "output C")?; let stream = ctx.stream(); // The batch maps to `blockIdx.z`, which CUDA caps at `MAX_GRID_DIM_Z`. Launch @@ -764,15 +764,14 @@ pub unsafe fn launch_gemm_external_batched_with_argmax_f32( // Per-batch element strides. A/B carry the (possibly padded) stride declared // by the external DLPack tensor; C/argmax are freshly allocated contiguous, - // so their stride is rows*cols. Derive the kernel's `i32` strides fallibly - // so a stride past `i32::MAX` can't truncate out of sync with the `usize` - // base-pointer offsets below. + // so their stride is rows*cols. Derive the kernel's `i64` strides fallibly + // so they stay in sync with the `usize` base-pointer offsets below. let stride_a = a.stride(); let stride_b = b.stride(); let stride_c = c.tensor.stride(); - let stride_a_i32 = stride_to_i32(stride_a, "operand A")?; - let stride_b_i32 = stride_to_i32(stride_b, "operand B")?; - let stride_c_i32 = stride_to_i32(stride_c, "output C")?; + let stride_a_i64 = stride_to_i64(stride_a, "operand A")?; + let stride_b_i64 = stride_to_i64(stride_b, "operand B")?; + let stride_c_i64 = stride_to_i64(stride_c, "output C")?; let n_i32 = n as i32; // Swapped: N becomes "M" let m_i32 = m as i32; // Swapped: M becomes "N" let k_i32 = k as i32; @@ -814,9 +813,9 @@ pub unsafe fn launch_gemm_external_batched_with_argmax_f32( .arg(&n_i32) .arg(&m_i32) .arg(&k_i32) - .arg(&stride_b_i32) // strideA (B's stride in our swap) - .arg(&stride_a_i32) // strideB (A's stride in our swap) - .arg(&stride_c_i32); // strideC + .arg(&stride_b_i64) // strideA (B's stride in our swap) + .arg(&stride_a_i64) // strideB (A's stride in our swap) + .arg(&stride_c_i64); // strideC builder.launch(cfg)?; start += chunk; } diff --git a/crates/tropical-gemm-cuda/src/lib.rs b/crates/tropical-gemm-cuda/src/lib.rs index ab3a7a8..e0b4a0f 100644 --- a/crates/tropical-gemm-cuda/src/lib.rs +++ b/crates/tropical-gemm-cuda/src/lib.rs @@ -928,7 +928,11 @@ pub fn tropical_backward_a_gpu_kernel( let total = m * n; let block_size = 256u32; - let grid_size = ((total as u32) + block_size - 1) / block_size; + // Compute the block count in 64-bit: `total as u32` would truncate when + // m*n >= 2^32 and launch far too few blocks (issue #63). The kernel covers + // every output element with a 64-bit linear index, so one block per 256 + // elements is correct; the count itself stays well within gridDim.x. + let grid_size = (total as u64).div_ceil(block_size as u64) as u32; let cfg = cudarc::driver::LaunchConfig { grid_dim: (grid_size, 1, 1), @@ -993,7 +997,11 @@ pub fn tropical_backward_b_gpu_kernel( let total = m * n; let block_size = 256u32; - let grid_size = ((total as u32) + block_size - 1) / block_size; + // Compute the block count in 64-bit: `total as u32` would truncate when + // m*n >= 2^32 and launch far too few blocks (issue #63). The kernel covers + // every output element with a 64-bit linear index, so one block per 256 + // elements is correct; the count itself stays well within gridDim.x. + let grid_size = (total as u64).div_ceil(block_size as u64) as u32; let cfg = cudarc::driver::LaunchConfig { grid_dim: (grid_size, 1, 1), @@ -1285,6 +1293,193 @@ mod tests { } } + // ======================================================================== + // Issue #63 regression: 64-bit global addressing past the i32 ceiling. + // ======================================================================== + // These cross the 2^31-element boundary, so they need a large-VRAM GPU and + // are #[ignore]d (CI skips CUDA entirely; the user validates on HPC A800s). + // Run with, e.g.: + // cargo test -p tropical-gemm-cuda --features cuda -- --ignored --nocapture + // + // 46341^2 = 2_147_488_281 > 2^31 (2_147_483_648), so the corner cell's + // column-major offset (m*n - 1) overflows i32. K is kept tiny so the A/B + // operands and the GEMM work stay small -- only the m*n output is large. + // Host memory stays minimal: operands are O(m*k + k*n), and the huge output + // is never copied back wholesale -- individual cells are sampled device-side. + + const BIG_DIM: usize = 46341; // sqrt of just over 2^31; m*n crosses i32::MAX + + // Compile-time guarantee that the shared shape crosses i32::MAX, so the + // corner cell's column-major offset (m*n - 1) exercises 64-bit addressing. + const _: () = assert!(BIG_DIM * BIG_DIM > i32::MAX as usize + 1); + + // Shared column-major operand patterns for the forward regressions: simple + // per-index values so any sampled cell's max-plus result is host-recomputable. + fn big_a_val(i: usize, p: usize) -> f32 { + ((i % 97) as f32) + (p as f32) * 0.5 - 40.0 + } + fn big_b_val(p: usize, j: usize) -> f32 { + ((j % 89) as f32) - (p as f32) * 0.25 - 40.0 + } + + /// Build column-major A (m×k) and B (k×n) from [`big_a_val`]/[`big_b_val`]. + fn build_big_operands(m: usize, k: usize, n: usize) -> (Vec, Vec) { + let mut a = vec![0f32; m * k]; + for p in 0..k { + for i in 0..m { + a[p * m + i] = big_a_val(i, p); + } + } + let mut b = vec![0f32; k * n]; + for j in 0..n { + for p in 0..k { + b[j * k + p] = big_b_val(p, j); + } + } + (a, b) + } + + /// Sample each `(i, j)` from column-major device output `c` and assert it + /// equals the max-plus reference. Reads one element per cell (never the full + /// m*n output), so host memory stays tiny even when `c` is multiple GB. + fn assert_big_maxplus_cells( + ctx: &CudaContext, + c: &cudarc::driver::CudaSlice, + m: usize, + k: usize, + cells: &[(usize, usize)], + ) { + let stream = ctx.stream(); + for &(i, j) in cells { + let off = j * m + i; // column-major linear offset + let host = stream.clone_dtoh(&c.slice(off..off + 1)).unwrap(); + stream.synchronize().unwrap(); + let mut exp = f32::NEG_INFINITY; + for p in 0..k { + exp = exp.max(big_a_val(i, p) + big_b_val(p, j)); + } + assert!( + (host[0] - exp).abs() < 1e-3, + "cell ({i},{j}) off {off}: got {}, want {}", + host[0], + exp + ); + } + } + + /// Forward strided-batched (batch=1): the exact issue-#63 repro. The output + /// C stride (= m*n) exceeds i32::MAX, and the corner cell's within-matrix + /// offset overflows i32 -- both must address correctly under 64-bit indexing. + #[test] + #[ignore = "allocates ~8.6 GB of VRAM; run manually on a large-memory GPU"] + fn test_batched_forward_offset_over_i32_max() { + use cudarc::driver::CudaSlice; + let Some(ctx) = cuda_context_or_skip() else { + return; + }; + let stream = ctx.stream(); + + let (m, n, k) = (BIG_DIM, BIG_DIM, 2usize); + + let (a_host, b_host) = build_big_operands(m, k, n); + let a_dev: CudaSlice = stream.clone_htod(&a_host).unwrap(); + let b_dev: CudaSlice = stream.clone_htod(&b_host).unwrap(); + // Uninitialized output: the kernel fully writes every cell. + let mut c_dev = unsafe { stream.alloc::(m * n) }.unwrap(); + + as CudaKernel>::launch_gemm_batched( + ctx, &a_dev, &b_dev, &mut c_dev, 1, m, k, n, + ) + .unwrap(); + + // Sample cells, including the corner whose offset = m*n - 1 > i32::MAX. + let cells = [ + (0usize, 0usize), + (m - 1, n - 1), + (m / 2, n / 2), + (m - 1, 0), + (0, n - 1), + (12345, 40000), + ]; + assert_big_maxplus_cells(ctx, &c_dev, m, k, &cells); + } + + /// Single-matrix forward (`launch_gemm`): the issue note that a *non*-batched + /// matrix of >= 2^31 elements also overflows the within-matrix offset. Shares + /// the STORE_C_TILE / OFFSET_COL64 path with the batched kernel. + #[test] + #[ignore = "allocates ~8.6 GB of VRAM; run manually on a large-memory GPU"] + fn test_single_matrix_offset_over_i32_max() { + let Some(ctx) = cuda_context_or_skip() else { + return; + }; + + let (m, n, k) = (BIG_DIM, BIG_DIM, 2usize); + + let (a_host, b_host) = build_big_operands(m, k, n); + let a = GpuMatrix::from_host(ctx, &a_host, m, k).unwrap(); + let b = GpuMatrix::from_host(ctx, &b_host, k, n).unwrap(); + let mut c = GpuMatrix::::alloc(ctx, m, n).unwrap(); + as CudaKernel>::launch_gemm(ctx, &a, &b, &mut c).unwrap(); + + let cells = [(0usize, 0usize), (m - 1, n - 1), (m - 1, 0), (0, n - 1)]; + assert_big_maxplus_cells(ctx, c.as_slice(), m, k, &cells); + } + + /// Backward scatter past i32::MAX: with m*n >= 2^31 the kernel's `total`/`idx` + /// and the host grid_size must all be 64-bit. Bulk grad_c/argmax are zeroed + /// on-device (no 8 GB host buffers); only the corner cell (idx = m*n - 1) is + /// written, and its contribution must land at the right grad_a/grad_b index. + #[test] + #[ignore = "allocates ~16 GB of VRAM; run manually on a large-memory GPU"] + fn test_backward_index_over_i32_max() { + let Some(ctx) = cuda_context_or_skip() else { + return; + }; + let stream = ctx.stream(); + + let (m, n, k) = (BIG_DIM, BIG_DIM, 4usize); + + // Device-zeroed inputs (no host-side 8 GB allocations). + let mut grad_c = GpuMatrix::::alloc(ctx, m, n).unwrap(); + let mut argmax = stream.alloc_zeros::(m * n).unwrap(); + + // Route the corner cell's gradient to k-index `kk`. off = m*n - 1. + let (i, j) = (m - 1, n - 1); + let off = j * m + i; + let kk: ArgmaxIndex = 3; // < k + let grad_val = 2.5f32; + stream + .memcpy_htod( + &[grad_val], + &mut grad_c.as_slice_mut().slice_mut(off..off + 1), + ) + .unwrap(); + stream + .memcpy_htod(&[kk], &mut argmax.slice_mut(off..off + 1)) + .unwrap(); + + let grad_a = tropical_backward_a_gpu_kernel(ctx, &grad_c, &argmax, m, k, n).unwrap(); + let grad_b = tropical_backward_b_gpu_kernel(ctx, &grad_c, &argmax, m, k, n).unwrap(); + + // grad_a[i + kk*m] and grad_b[kk + j*k] must each receive exactly grad_val + // (all other cells contributed 0). Outputs are small (m×k, k×n). + let ga = grad_a.to_host(ctx).unwrap(); + let gb = grad_b.to_host(ctx).unwrap(); + let ga_idx = i + (kk as usize) * m; + let gb_idx = (kk as usize) + j * k; + assert!( + (ga[ga_idx] - grad_val).abs() < 1e-4, + "grad_a[{ga_idx}] = {}, want {grad_val}", + ga[ga_idx] + ); + assert!( + (gb[gb_idx] - grad_val).abs() < 1e-4, + "grad_b[{gb_idx}] = {}, want {grad_val}", + gb[gb_idx] + ); + } + #[test] fn test_tropical_matmul_gpu_with_argmax_maxplus() { if cuda_context_or_skip().is_none() {