Share TILE_SIZE.

One of the nice things about using Rust for both CPU and GPU code is the
ability to share things between them. So let's do that in `gemm`.

Author: nnethercote

Cargo.lock

@@ -8,6 +8,7 @@ blastoff = { path = "../../../crates/blastoff" }
 cuda_std = { path = "../../../crates/cuda_std" }
 cust = { path = "../../../crates/cust" }
 cust_raw = { path = "../../../crates/cust_raw", features = ["driver"] }
+gemm-kernels = { path = "kernels" }
 ndarray = { version = "0.16", features = ["approx"] }
 ndarray-rand = "0.15.0"
 rand = "0.9"

examples/cuda/gemm/Cargo.toml

@@ -3,6 +3,8 @@ use cuda_std::address_space;
 use cuda_std::kernel;
 use cuda_std::thread;
 
+pub const TILE_SIZE: usize = 16;
+
 #[kernel]
 #[allow(improper_ctypes_definitions)]
 /// Tiled GEMM kernel for C = alpha * A * B + beta * C.
@@ -38,7 +40,6 @@ pub unsafe fn gemm_tiled(
     alpha: f32,
     beta: f32,
 ) {
-    const TILE_SIZE: usize = 16;
     const TILE_SIZE_2D: usize = TILE_SIZE * TILE_SIZE;
 
     // Shared GPU memory is modelled with `#[address_space(shared)] static mut`. Unlike normal

examples/cuda/gemm/kernels/src/gemm_tiled.rs

@@ -2,4 +2,4 @@ mod gemm_naive;
 mod gemm_tiled;
 
 pub use crate::gemm_naive::gemm_naive;
-pub use crate::gemm_tiled::gemm_tiled;
+pub use crate::gemm_tiled::{TILE_SIZE, gemm_tiled};

examples/cuda/gemm/kernels/src/lib.rs

@@ -13,6 +13,7 @@ use cust::memory::CopyDestination as _;
 use cust::module;
 use cust::stream;
 use cust::util::SliceExt as _;
+use gemm_kernels::TILE_SIZE;
 use ndarray::Array;
 use ndarray_rand::RandomExt as _;
 use ndarray_rand::rand_distr::Uniform;
@@ -430,9 +431,6 @@ pub fn gemm_tiled(
     assert_eq!(mat_b.len(), k * n);
     assert_eq!(mat_c.len(), m * n);
 
-    // These values must be aligned with the kernel code.
-    const TILE_SIZE: usize = 16;
-
     let kernel_cell = cell::LazyCell::new(|| {
         module
             .get_function("gemm_tiled")