teeny_cuda/testing/
mod.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use teeny_core::device::{
18    context::{Context, DeviceInfo},
19    program::Kernel,
20};
21
22use crate::{
23    compiler::target::{Capability, capability_from_device_info},
24    device::context::Cuda,
25    device::program::CudaProgram,
26    device::{CudaDevice, CudaLaunchConfig},
27    errors::Result,
28};
29
30/// A device + capability pair set up for a test, via [`setup_cuda_env`].
31pub struct CudaTestEnv {
32    /// The opened test device.
33    pub device: CudaDevice<'static>,
34    /// The device's (possibly `TEENYC_CAPABILITY`-overridden) compute capability.
35    pub capability: Capability,
36}
37
38/// Asserts CUDA is available, opens the first device, and resolves its compute capability
39/// (overridable via the `TEENYC_CAPABILITY` env var, e.g. `sm_90`) — the standard setup for
40/// `teeny-cuda`'s own device-dependent tests.
41pub fn setup_cuda_env() -> Result<CudaTestEnv> {
42    let cuda_available = Cuda::is_available()?;
43    assert!(cuda_available, "CUDA is not available");
44    println!("[1/9] CUDA available");
45
46    let cuda = Cuda::try_new()?;
47    let devices = cuda.list_devices()?;
48    assert!(!devices.is_empty(), "No CUDA devices found");
49    println!("[2/9] found {} device(s)", devices.len());
50
51    let device = cuda.device(&devices[0].id())?;
52    let device_capability = capability_from_device_info(&device.info)?;
53
54    let capability = if let Ok(val) = std::env::var("TEENYC_CAPABILITY") {
55        let parsed = val
56            .strip_prefix("sm_")
57            .and_then(|s| s.parse::<i32>().ok())
58            .and_then(|n| Capability::from_major_minor(n / 10, n % 10))
59            .ok_or_else(|| {
60                crate::errors::Error::UnknownCapability(format!(
61                    "TEENYC_CAPABILITY={val:?} is not a recognised sm version"
62                ))
63            })?;
64        println!(
65            "[3/9] device: {} (capability overridden: {device_capability} → {parsed})",
66            device.info.name
67        );
68        parsed
69    } else {
70        println!(
71            "[3/9] device: {} (capability: {device_capability})",
72            device.info.name
73        );
74        device_capability
75    };
76
77    Ok(CudaTestEnv { device, capability })
78}
79
80/// Build a launch config using block size from kernel metadata.
81///
82/// `n_elements` — total number of elements to process.
83/// `program`    — compiled program; `metadata.threads_per_block()` is used as the block size.
84pub fn launch_config_from_program<K: Kernel>(
85    n_elements: usize,
86    program: &CudaProgram<'_, K>,
87) -> CudaLaunchConfig {
88    let threads = program.metadata.threads_per_block().max(1);
89    CudaLaunchConfig {
90        grid: [(n_elements as u32).div_ceil(threads), 1, 1],
91        block: [threads, 1, 1],
92        cluster: [program.metadata.num_ctas.max(1), 1, 1],
93    }
94}
95
96/// Build a launch config with an explicit block size.
97///
98/// Use when the block size is determined at the call site rather than from PTX metadata
99/// (e.g. when the grid is computed before the program is compiled).
100pub fn launch_config(n_elements: usize, block_size: i32) -> CudaLaunchConfig {
101    CudaLaunchConfig {
102        grid: [(n_elements as u32).div_ceil(block_size as u32), 1, 1],
103        block: [block_size as u32, 1, 1],
104        cluster: [1, 1, 1],
105    }
106}
107
108/// Build a launch config with a pre-computed grid and block from kernel metadata.
109///
110/// Use when the launch grid (number of CTAs) is known independently of element count —
111/// for example, tiled matmul kernels where grid = `ceil(M/TILE_M) * ceil(N/TILE_N)`.
112/// The block and cluster dimensions are read from `program.metadata`.
113pub fn launch_config_with_grid<K: Kernel>(
114    grid_x: usize,
115    program: &CudaProgram<'_, K>,
116) -> CudaLaunchConfig {
117    let threads = program.metadata.threads_per_block().max(1);
118    CudaLaunchConfig {
119        grid: [grid_x as u32, 1, 1],
120        block: [threads, 1, 1],
121        cluster: [program.metadata.num_ctas.max(1), 1, 1],
122    }
123}
124
125/// Loads a compiled program directly from raw PTX bytes via the driver's JIT loader (skipping
126/// `nvptxcompiler`) — useful in tests that already have PTX in hand.
127pub fn load_program_from_ptx<K: Kernel>(ptx: &[u8]) -> Result<CudaProgram<'static, K>> {
128    println!("      loading PTX directly via driver JIT...");
129    let program = CudaProgram::<K>::try_from_ptx(ptx)?;
130    println!(
131        "[7/9] loaded PTX: module={:#x} function={:#x} num_warps={} num_ctas={}",
132        program.module_ptr(),
133        program.function_ptr(),
134        program.metadata.num_warps,
135        program.metadata.num_ctas,
136    );
137    Ok(program)
138}
teeny_cuda/testing/mod.rs

teeny_cuda/testing/
mod.rs
