Check that every kernel has a parity case

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);

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) {
@@ -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<<<blocks, threads>>>(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<<<blocks, threads>>>(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];

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;}

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.

src/cuda/kernels.rs

@@ -208,29 +208,6 @@ pub fn matmul_batched_tn(
     })
 }
 
-pub fn transpose_2d(x: &CudaBuffer, rows: usize, cols: usize) -> Result<CudaBuffer> {
-    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<CudaBuffer> {
-    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(

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<String> = 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:?}"
+    );
+}