metrale_gpu_runtime/
kernel.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `KernelModule`, a PTX module loaded through cudarc, plus a 1-D
4//! `launch_config` helper and `launch_vector_add`.
5//!
6//! Owner: gpu-runtime.
7//! Invariants: none beyond the types.
8
9use std::sync::Arc;
10
11use cudarc::driver::{
12    CudaContext, CudaFunction, CudaModule, CudaStream, LaunchConfig, PushKernelArg,
13};
14use cudarc::nvrtc::Ptx;
15
16use metrale_core::error::{MetraleError, Result};
17
18/// 2026-09-25: A PTX module loaded into a CUDA context through cudarc.
19pub struct KernelModule {
20    module: Arc<CudaModule>,
21}
22
23impl KernelModule {
24    /// 2026-09-25: Load PTX source text into `ctx`.
25    pub fn from_ptx_src(ctx: &Arc<CudaContext>, ptx_src: &str) -> Result<Self> {
26        let ptx = Ptx::from_src(ptx_src);
27        let module = ctx
28            .load_module(ptx)
29            .map_err(|e| MetraleError::ModuleLoad(format!("PTX load failed: {e}")))?;
30        Ok(Self { module })
31    }
32
33    /// 2026-09-25: Look up a kernel function by name.
34    pub fn get_function(&self, name: &str) -> Result<CudaFunction> {
35        self.module
36            .load_function(name)
37            .map_err(|e| MetraleError::ModuleLoad(format!("Function '{name}' not found: {e}")))
38    }
39}
40
41/// 2026-09-25: A 1-D launch of `ceil(n / block_size)` blocks of `block_size`
42/// threads, with no dynamic shared memory.
43pub fn launch_config(n: u32, block_size: u32) -> LaunchConfig {
44    LaunchConfig {
45        grid_dim: (n.div_ceil(block_size), 1, 1),
46        block_dim: (block_size, 1, 1),
47        shared_mem_bytes: 0,
48    }
49}
50
51/// 2026-09-25: Launch `func` with `(a_ptr, b_ptr, c_ptr, n)` in `ceil(n / 256)`
52/// blocks of 256 threads, through cudarc's launch builder.
53///
54/// # Safety
55///
56/// All pointers must be valid CUDA device pointers to f32 arrays of length >= n.
57pub unsafe fn launch_vector_add(
58    stream: &Arc<CudaStream>,
59    func: &CudaFunction,
60    a_ptr: u64,
61    b_ptr: u64,
62    c_ptr: u64,
63    n: u32,
64) -> Result<()> {
65    let cfg = launch_config(n, 256);
66    unsafe {
67        stream
68            .launch_builder(func)
69            .arg(&a_ptr)
70            .arg(&b_ptr)
71            .arg(&c_ptr)
72            .arg(&n)
73            .launch(cfg)
74            .map_err(|e| MetraleError::KernelLaunch(format!("vector_add launch failed: {e}")))?;
75    }
76    Ok(())
77}
metrale_gpu_runtime/kernel.rs

metrale_gpu_runtime/
kernel.rs
