From 41e3ec1e89114bd8d5f3b7c3aa714d109fb9f919 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Mon, 2 Nov 2026 13:42:00 +0300 Subject: [PATCH 1/3] remove the transpose kernels nothing calls --- cuda/include/kernels.h | 2 -- cuda/kernels.cu | 56 ------------------------------------------ cuda/stubs.c | 2 -- src/cuda/ffi.rs | 8 ------ src/cuda/kernels.rs | 23 ----------------- 5 files changed, 91 deletions(-) diff --git a/cuda/include/kernels.h b/cuda/include/kernels.h index 09c066b..6b64052 100644 --- a/cuda/include/kernels.h +++ b/cuda/include/kernels.h @@ -60,8 +60,6 @@ int fastnn_cuda_matmul_nt(const float* a, const float* b, float* c, int m, int n int fastnn_cuda_matmul_tn(const float* a, const float* b, float* c, int m, int n, int k); int fastnn_cuda_matmul_batched_nt(const float* a, const float* b, float* c, int m, int n, int k, int batch_size); int fastnn_cuda_matmul_batched_tn(const float* a, const float* b, float* c, int m, int n, int k, int batch_size); -int fastnn_cuda_transpose(const float* input, float* output, int rows, int cols); -int fastnn_cuda_transpose_batched(const float* input, float* output, int batch_size, int rows, int cols); int fastnn_cuda_permute_nd(const float* input, float* output, const int* out_strides, const int* in_strides, const int* perm, int ndim, int numel); int fastnn_cuda_sum(const float* input, float* output, size_t n); diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 5cc64c7..6b98cfb 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -494,62 +494,6 @@ extern "C" int fastnn_cuda_matmul_batched(const float* a, const float* b, float* return 0; } -__global__ void kernel_transpose(const float* input, float* output, int rows, int cols) { - __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; // +1 pad avoids bank conflicts - - int x = blockIdx.x * TILE_SIZE + threadIdx.x; - int y = blockIdx.y * TILE_SIZE + threadIdx.y; - - if (x < cols && y < rows) { - tile[threadIdx.y][threadIdx.x] = input[y * cols + x]; - } - __syncthreads(); - - x = blockIdx.y * TILE_SIZE + threadIdx.x; - y = blockIdx.x * TILE_SIZE + threadIdx.y; - - if (x < rows && y < cols) { - output[y * rows + x] = tile[threadIdx.x][threadIdx.y]; - } -} - -extern "C" int fastnn_cuda_transpose(const float* input, float* output, int rows, int cols) { - dim3 threads(TILE_SIZE, TILE_SIZE); - dim3 blocks(div_ceil(cols, TILE_SIZE), div_ceil(rows, TILE_SIZE)); - kernel_transpose<<>>(input, output, rows, cols); - CUDA_CHECK(cudaGetLastError()); - return 0; -} - -__global__ void kernel_transpose_batched(const float* input, float* output, int batch_size, int rows, int cols) { - int b = blockIdx.z; - if (b >= batch_size) return; - - __shared__ float tile[TILE_SIZE][TILE_SIZE + 1]; - int offset = b * rows * cols; - - int x = blockIdx.x * TILE_SIZE + threadIdx.x; - int y = blockIdx.y * TILE_SIZE + threadIdx.y; - - if (x < cols && y < rows) - tile[threadIdx.y][threadIdx.x] = input[offset + y * cols + x]; - __syncthreads(); - - x = blockIdx.y * TILE_SIZE + threadIdx.x; - y = blockIdx.x * TILE_SIZE + threadIdx.y; - - if (x < rows && y < cols) - output[offset + y * rows + x] = tile[threadIdx.x][threadIdx.y]; -} - -extern "C" int fastnn_cuda_transpose_batched(const float* input, float* output, int batch_size, int rows, int cols) { - dim3 threads(TILE_SIZE, TILE_SIZE); - dim3 blocks(div_ceil(cols, TILE_SIZE), div_ceil(rows, TILE_SIZE), batch_size); - kernel_transpose_batched<<>>(input, output, batch_size, rows, cols); - CUDA_CHECK(cudaGetLastError()); - return 0; -} - struct PermuteNDParams { int out_strides[PERMUTE_MAX_DIMS]; int in_strides[PERMUTE_MAX_DIMS]; diff --git a/cuda/stubs.c b/cuda/stubs.c index ce66987..1eba133 100644 --- a/cuda/stubs.c +++ b/cuda/stubs.c @@ -52,8 +52,6 @@ int fastnn_cuda_matmul_nt(const float*a,const float*b,float*c,int m,int n,int k) int fastnn_cuda_matmul_tn(const float*a,const float*b,float*c,int m,int n,int k){(void)a;(void)b;(void)c;(void)m;(void)n;(void)k;STUB;} int fastnn_cuda_matmul_batched_nt(const float*a,const float*b,float*c,int m,int n,int k,int batch){(void)a;(void)b;(void)c;(void)m;(void)n;(void)k;(void)batch;STUB;} int fastnn_cuda_matmul_batched_tn(const float*a,const float*b,float*c,int m,int n,int k,int batch){(void)a;(void)b;(void)c;(void)m;(void)n;(void)k;(void)batch;STUB;} -int fastnn_cuda_transpose(const float*i,float*o,int rows,int cols) {(void)i;(void)o;(void)rows;(void)cols;STUB;} -int fastnn_cuda_transpose_batched(const float*i,float*o,int b,int rows,int cols){(void)i;(void)o;(void)b;(void)rows;(void)cols;STUB;} int fastnn_cuda_permute_nd(const float*i,float*o,const int*os,const int*is,const int*p,int nd,int n){(void)i;(void)o;(void)os;(void)is;(void)p;(void)nd;(void)n;STUB;} int fastnn_cuda_init(int device_id) {(void)device_id;STUB;} int fastnn_cuda_device_count(void) {return 0;} diff --git a/src/cuda/ffi.rs b/src/cuda/ffi.rs index da59b41..159c1f4 100644 --- a/src/cuda/ffi.rs +++ b/src/cuda/ffi.rs @@ -219,14 +219,6 @@ unsafe extern "C" { batch: c_int, ) -> c_int; - pub fn fastnn_cuda_transpose(x: *const f32, out: *mut f32, rows: c_int, cols: c_int) -> c_int; - pub fn fastnn_cuda_transpose_batched( - x: *const f32, - out: *mut f32, - batch: c_int, - rows: c_int, - cols: c_int, - ) -> c_int; /// General gather: `out[i] = x[dot(coords(i, out_strides), in_strides[perm])]`. /// A stride of 0 in `in_strides` broadcasts that dimension, which is how /// `expand` reuses this kernel. diff --git a/src/cuda/kernels.rs b/src/cuda/kernels.rs index 117a42a..459cd56 100644 --- a/src/cuda/kernels.rs +++ b/src/cuda/kernels.rs @@ -208,29 +208,6 @@ pub fn matmul_batched_tn( }) } -pub fn transpose_2d(x: &CudaBuffer, rows: usize, cols: usize) -> Result { - launch(rows * cols, "transpose", |out| unsafe { - ffi::fastnn_cuda_transpose(x.as_ptr(), out.as_mut_ptr(), rows as c_int, cols as c_int) - }) -} - -pub fn transpose_batched( - x: &CudaBuffer, - batch: usize, - rows: usize, - cols: usize, -) -> Result { - launch(batch * rows * cols, "transpose_batched", |out| unsafe { - ffi::fastnn_cuda_transpose_batched( - x.as_ptr(), - out.as_mut_ptr(), - batch as c_int, - rows as c_int, - cols as c_int, - ) - }) -} - /// Gathers `numel` elements by decoding each output index with `out_strides` and re-encoding with /// `in_strides` permuted by `perm`. A stride of 0 broadcasts. pub fn permute_nd( From 19a93f3bca182eb66fd5bd2c073efe4f3b6f9854 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Mon, 2 Nov 2026 13:42:00 +0300 Subject: [PATCH 2/3] run cublas in full fp32 instead of tf32 --- cuda/kernels.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cuda/kernels.cu b/cuda/kernels.cu index 6b98cfb..446b0da 100644 --- a/cuda/kernels.cu +++ b/cuda/kernels.cu @@ -32,7 +32,7 @@ extern "C" int fastnn_cuda_init(int device_id) { if (g_cublas_handle == nullptr) { CUBLAS_CHECK(cublasCreate(&g_cublas_handle)); - CUBLAS_CHECK(cublasSetMathMode(g_cublas_handle, CUBLAS_TF32_TENSOR_OP_MATH)); + CUBLAS_CHECK(cublasSetMathMode(g_cublas_handle, CUBLAS_DEFAULT_MATH)); } if (g_curand_gen == nullptr) { From 2a7df0e6a0a426bfc26f821897a255a5843446b8 Mon Sep 17 00:00:00 2001 From: Canberk Pitirli Date: Mon, 2 Nov 2026 13:42:00 +0300 Subject: [PATCH 3/3] check that every kernel has a parity case --- tests/cuda_parity.rs | 141 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 140 insertions(+), 1 deletion(-) diff --git a/tests/cuda_parity.rs b/tests/cuda_parity.rs index a933394..f130761 100644 --- a/tests/cuda_parity.rs +++ b/tests/cuda_parity.rs @@ -107,12 +107,28 @@ fn check_backward(name: &str, inputs: &[Tensor], f: impl Fn(&[Tensor]) -> Tensor fn elementwise() { let a = sample(&[4, 16]); let b = sample(&[4, 16]); + // Strictly positive, for the ops that need it. + let p = Tensor::from_vec(spread(64, 0.1, 1.9), &[4, 16]); check_forward("add", &[a.clone(), b.clone()], |x| x[0].add(&x[1])); + check_forward("sub", &[a.clone(), b.clone()], |x| x[0].sub(&x[1])); check_forward("mul", &[a.clone(), b.clone()], |x| x[0].mul(&x[1])); + check_forward("div", &[a.clone(), p.clone()], |x| x[0].div(&x[1])); + check_forward("add_scalar", &[a.clone()], |x| x[0].add_scalar(0.25)); + check_forward("mul_scalar", &[a.clone()], |x| x[0].mul_scalar(-1.5)); + check_forward("neg", &[a.clone()], |x| x[0].neg()); + check_forward("abs", &[a.clone()], |x| x[0].abs()); check_forward("exp", &[a.clone()], |x| x[0].exp()); + check_forward("sqrt", &[p.clone()], |x| x[0].sqrt()); + check_forward("log", &[p.clone()], |x| x[0].log()); + check_forward("pow_scalar", &[p.clone()], |x| x[0].powf(1.5)); check_forward("clamp", &[a.clone()], |x| x[0].clamp(-0.5, 0.5)); - check_backward("mul", &[a, b], |x| x[0].mul(&x[1]).sum()); + + check_backward("mul", &[a.clone(), b], |x| x[0].mul(&x[1]).sum()); + check_backward("div", &[a, p.clone()], |x| x[0].div(&x[1]).sum()); + check_backward("sqrt", &[p.clone()], |x| x[0].sqrt().sum()); + check_backward("log", &[p.clone()], |x| x[0].log().sum()); + check_backward("pow_scalar", &[p], |x| x[0].powf(1.5).sum()); } #[test] @@ -195,6 +211,16 @@ fn gemm_batched() { &[sample(&[4, 8, 12]), sample(&[4, 6, 12])], |x| x[0].matmul_nt(&x[1]), ); + check_forward( + "batched tn", + &[sample(&[4, 12, 8]), sample(&[4, 12, 6])], + |x| x[0].matmul_tn(&x[1]), + ); + check_backward( + "batched", + &[sample(&[4, 8, 12]), sample(&[4, 12, 6])], + |x| x[0].matmul(&x[1]).sum(), + ); // A weight meeting a batch of inputs. check_forward( "shared weight", @@ -222,6 +248,8 @@ fn reductions() { check_forward("sum", &[x.clone()], |x| x[0].sum()); check_forward("mean", &[x.clone()], |x| x[0].mean()); + check_forward("max", &[x.clone()], |x| Tensor::scalar(x[0].max())); + check_forward("min", &[x.clone()], |x| Tensor::scalar(x[0].min())); check_forward("sum_axis 0", &[x.clone()], |x| x[0].sum_axis(0)); check_forward("sum_axis 1", &[x.clone()], |x| x[0].sum_axis(1)); check_backward("sum_axis", &[x, sample(&[10])], |x| { @@ -389,3 +417,114 @@ fn dilated_convolution_lowering() { folded.mul(&weight).sum() }); } + +/// Every kernel symbol in `ffi.rs`, with the test here that runs it on both +/// devices. Plumbing (init, allocation, copies, sync) runs under every case. +/// +/// This is a list, not a trace. A row says where a kernel is covered; the test +/// below only checks that the list and `ffi.rs` agree and that each named test +/// exists. Keep a row honest when you add one. +const KERNEL_COVERAGE: &[(&str, &str)] = &[ + ("fastnn_cuda_init", RUNTIME), + ("fastnn_cuda_device_count", RUNTIME), + ("fastnn_cuda_synchronize", RUNTIME), + ("fastnn_cuda_malloc", RUNTIME), + ("fastnn_cuda_free", RUNTIME), + ("fastnn_cuda_memcpy_h2d", RUNTIME), + ("fastnn_cuda_memcpy_d2h", RUNTIME), + ("fastnn_cuda_memcpy_d2d", RUNTIME), + ("fastnn_cuda_memset", RUNTIME), + ("fastnn_cuda_get_memory_info", RUNTIME), + ("fastnn_cuda_add", "elementwise"), + ("fastnn_cuda_sub", "elementwise"), + ("fastnn_cuda_mul", "elementwise"), + ("fastnn_cuda_div", "elementwise"), + ("fastnn_cuda_add_scalar", "elementwise"), + ("fastnn_cuda_mul_scalar", "elementwise"), + ("fastnn_cuda_neg", "elementwise"), + ("fastnn_cuda_abs", "elementwise"), + ("fastnn_cuda_sqrt", "elementwise"), + ("fastnn_cuda_exp", "elementwise"), + ("fastnn_cuda_log", "elementwise"), + ("fastnn_cuda_pow_scalar", "elementwise"), + ("fastnn_cuda_clamp", "elementwise"), + ("fastnn_cuda_relu", "activations"), + ("fastnn_cuda_relu_backward", "activations"), + ("fastnn_cuda_sigmoid", "activations"), + ("fastnn_cuda_sigmoid_backward", "activations"), + ("fastnn_cuda_tanh_forward", "activations"), + ("fastnn_cuda_tanh_backward", "activations"), + ("fastnn_cuda_gelu", "activations"), + ("fastnn_cuda_gelu_backward", "activations"), + ("fastnn_cuda_silu", "activations"), + ("fastnn_cuda_silu_backward", "activations"), + ("fastnn_cuda_leaky_relu", "activations"), + ("fastnn_cuda_softmax", "softmax"), + ("fastnn_cuda_softmax_backward", "softmax"), + ("fastnn_cuda_log_softmax", "softmax"), + ("fastnn_cuda_matmul", "gemm_layouts"), + ("fastnn_cuda_matmul_nt", "gemm_layouts"), + ("fastnn_cuda_matmul_tn", "gemm_layouts"), + ("fastnn_cuda_matmul_batched", "gemm_batched"), + ("fastnn_cuda_matmul_batched_nt", "gemm_batched"), + ("fastnn_cuda_matmul_batched_tn", "gemm_batched"), + ("fastnn_cuda_permute_nd", "layout_ops"), + ("fastnn_cuda_sum", "reductions"), + ("fastnn_cuda_sum_axis", "reductions"), + ("fastnn_cuda_mean", "reductions"), + ("fastnn_cuda_max", "reductions"), + ("fastnn_cuda_min", "reductions"), + ("fastnn_cuda_layer_norm_forward", "layer_norm"), + ("fastnn_cuda_layer_norm_backward", "layer_norm"), + ("fastnn_cuda_im2col", "convolution_lowering"), + ("fastnn_cuda_col2im", "convolution_lowering"), + ("fastnn_cuda_embedding_forward", "embedding_lookup"), + ("fastnn_cuda_embedding_backward", "embedding_lookup"), +]; + +const RUNTIME: &str = "runtime"; + +/// Runs without a GPU: it reads `ffi.rs` and this file, nothing else. +#[test] +fn every_kernel_has_a_parity_case() { + let declared: Vec<&str> = include_str!("../src/cuda/ffi.rs") + .lines() + .filter_map(|line| line.trim_start().strip_prefix("pub fn ")) + .map(|rest| { + rest.split(|c: char| !c.is_alphanumeric() && c != '_') + .next() + .unwrap_or("") + }) + .filter(|name| name.starts_with("fastnn_cuda_")) + .collect(); + let this_file = include_str!("cuda_parity.rs"); + + let unlisted: Vec<&str> = declared + .iter() + .copied() + .filter(|symbol| !KERNEL_COVERAGE.iter().any(|(listed, _)| listed == symbol)) + .collect(); + let stale: Vec<&str> = KERNEL_COVERAGE + .iter() + .map(|(symbol, _)| *symbol) + .filter(|symbol| !declared.contains(symbol)) + .collect(); + let no_such_test: Vec = KERNEL_COVERAGE + .iter() + .filter(|(_, test)| *test != RUNTIME && !this_file.contains(&format!("\nfn {test}()"))) + .map(|(symbol, test)| format!("{symbol} -> {test}")) + .collect(); + + assert!( + unlisted.is_empty(), + "kernels with no parity row: {unlisted:?}" + ); + assert!( + stale.is_empty(), + "rows for kernels that no longer exist: {stale:?}" + ); + assert!( + no_such_test.is_empty(), + "rows naming a test that is not in this file: {no_such_test:?}" + ); +}