Struct AtlasRegistry 

pub struct AtlasRegistry { /* private fields */ }

The PTX/CUBIN modules for one loaded model.

Model-scoped: the blob set comes from atlas_kernels::ptx_for_model, so it changes with the checkpoint. Previously this was fused into a process OnceLock singleton whose get_or_init(ordinal, kernel_blobs) silently discarded the second caller’s blobs — a swapped-in model would have run the previous model’s kernels with no error at all.

Obtain one with AtlasRegistry::load and propagate it (Arc<AtlasRegistry>); there is deliberately no global accessor. Dropping the last handle unloads the modules.

Implementations

impl AtlasRegistry

pub fn load( ordinal: usize, kernel_blobs: &[(&'static str, &'static [u8])], ) -> Result<Arc<Self>>

Load this model’s kernel modules into the process CUDA context.

Each call produces a fresh, independent module set; nothing is shared with a previously loaded model except the context and stream.

pub fn host(&self) -> &Arc<CudaHost>

The process CUDA context this registry’s modules live in.

pub fn ctx(&self) -> &Arc<CudaContext>

pub fn stream(&self) -> &Arc<CudaStream>

pub fn module_names(&self) -> impl Iterator<Item = &'static str> + '_

Module names this registry loaded, for diagnostics.

pub fn function( &self, module_name: &str, func_name: &str, ) -> Result<CudaFunction>

Look up a cached function handle (cudarc safe API).

pub fn function_cached( &self, cache: &OnceLock<CudaFunction>, module_name: &str, func_name: &str, ) -> Result<CudaFunction>

Look up a function handle with OnceLock caching (cudarc safe API).

pub fn raw_function_cached( &self, cache: &OnceLock<RawCudaFunc>, module_name: &str, func_name: &str, ) -> Result<RawCudaFunc>

Look up a raw CUfunction handle with OnceLock caching. Uses the raw CUDA driver API — no cudarc struct layout dependency.

pub fn raw_stream(&self) -> u64

Get the raw CUstream handle for Atlas’s own stream.

pub fn device_symbol( &self, module_name: &str, symbol: &str, ) -> Result<(u64, usize)>

Resolve a __device__ symbol in a loaded PTX module to its device pointer + byte length. Required for drivers that read/write device globals without launching a kernel (e.g. InnerQ calibration state). symbol must be the linker-visible name — C++ namespace symbols are Itanium-mangled (_ZN7tq_plus14d_innerq_scaleE).

pub unsafe fn copy_h2d_async( &self, dst: u64, src: *const c_void, bytes: usize, stream: u64, ) -> Result<()>

Async H2D copy into a previously-resolved device pointer.

Safety

Caller must ensure dst is a valid device pointer and the bytes pointed to by src outlive the copy (host buffers must persist until the next sync on stream).

pub unsafe fn copy_d2h_async( &self, dst: *mut c_void, src: u64, bytes: usize, stream: u64, ) -> Result<()>

Async D2H copy from a device pointer. Same lifetime caveats as the H2D variant.

Safety

Caller must keep dst alive until stream is synchronised.

pub fn stream_synchronize(&self, stream: u64) -> Result<()>

Block the calling thread until all prior work on stream completes.

pub unsafe fn launch_on_stream( &self, raw_func: RawCudaFunc, cfg: LaunchConfig, stream_ptr: u64, kernel_params: &mut [*mut c_void], ) -> Result<()>

Launch a kernel on a specified raw CUDA stream.

When stream_ptr comes from the caller (e.g. torch.cuda.current_stream().cuda_stream), this ensures kernels are captured during CUDA graph recording.

Safety

• kernel_params must contain valid pointers to arguments matching the kernel signature.
• stream_ptr must be a valid CUstream handle (or 0 to use Atlas’s own stream).
• raw_func must be a valid CUfunction obtained from raw_function_cached.

Trait Implementations

impl Drop for AtlasRegistry

fn drop(&mut self)

Unloads this model’s modules. Reached when the last Arc handle goes, which — because there is no global accessor — happens exactly when the owning run ends.

impl Send for AtlasRegistry

impl Sync for AtlasRegistry