atlas_core/
stream.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use cudarc::driver::CudaStream;
4use std::sync::Arc;
5
6use crate::device::AtlasDevice;
7use crate::error::Result;
8
9/// CUDA stream wrapper for asynchronous kernel execution.
10pub struct AtlasStream {
11    pub stream: Arc<CudaStream>,
12    pub device: AtlasDevice,
13}
14
15impl AtlasStream {
16    /// Create a new CUDA stream on the given device.
17    pub fn new(device: &AtlasDevice) -> Result<Self> {
18        let stream = device
19            .ctx
20            .new_stream()
21            .map_err(crate::error::AtlasError::CudaDriver)?;
22        Ok(Self {
23            stream,
24            device: device.clone(),
25        })
26    }
27}