Skip to main content

singe_npp/
context.rs

1use singe_cuda::{
2    context::Context as CudaContext,
3    stream::{BorrowedStream, Stream},
4};
5use singe_npp_sys as sys;
6
7use crate::error::Result;
8
9/// CUDA stream metadata passed to context-aware NPP functions.
10pub struct StreamContext {
11    stream: BorrowedStream,
12    raw: sys::NppStreamContext,
13}
14
15impl StreamContext {
16    /// Creates an NPP stream context from a CUDA stream.
17    ///
18    /// The returned context borrows the stream and caches the device properties
19    /// required by NPP context-aware entry points.
20    ///
21    /// # Errors
22    ///
23    /// Returns an error if the stream's CUDA context cannot be bound, or if CUDA
24    /// cannot query the stream device, device properties, or stream flags.
25    pub fn create(stream: &Stream) -> Result<Self> {
26        stream.context().bind()?;
27        let device = stream.device()?;
28        let properties = device.properties()?;
29        let flags = stream.flags()?;
30        let stream = stream.to_borrowed();
31
32        Ok(Self {
33            raw: sys::NppStreamContext {
34                hStream: stream.as_raw(),
35                nCudaDeviceId: device.id(),
36                nMultiProcessorCount: properties.multi_processor_count,
37                nMaxThreadsPerMultiProcessor: properties.max_threads_per_multi_processor,
38                nMaxThreadsPerBlock: properties.max_threads_per_block,
39                nSharedMemPerBlock: properties.shared_mem_per_block as _,
40                nCudaDevAttrComputeCapabilityMajor: properties.major,
41                nCudaDevAttrComputeCapabilityMinor: properties.minor,
42                nStreamFlags: flags.bits(),
43                nReserved0: 0,
44            },
45            stream,
46        })
47    }
48
49    pub fn stream(&self) -> &BorrowedStream {
50        &self.stream
51    }
52
53    pub fn cuda_context(&self) -> &CudaContext {
54        self.stream.context()
55    }
56
57    pub const fn device_id(&self) -> i32 {
58        self.raw.nCudaDeviceId
59    }
60
61    pub const fn as_raw(&self) -> sys::NppStreamContext {
62        sys::NppStreamContext {
63            hStream: self.stream.as_raw(),
64            nCudaDeviceId: self.raw.nCudaDeviceId,
65            nMultiProcessorCount: self.raw.nMultiProcessorCount,
66            nMaxThreadsPerMultiProcessor: self.raw.nMaxThreadsPerMultiProcessor,
67            nMaxThreadsPerBlock: self.raw.nMaxThreadsPerBlock,
68            nSharedMemPerBlock: self.raw.nSharedMemPerBlock,
69            nCudaDevAttrComputeCapabilityMajor: self.raw.nCudaDevAttrComputeCapabilityMajor,
70            nCudaDevAttrComputeCapabilityMinor: self.raw.nCudaDevAttrComputeCapabilityMinor,
71            nStreamFlags: self.raw.nStreamFlags,
72            nReserved0: self.raw.nReserved0,
73        }
74    }
75
76    pub fn synchronize(&self) -> Result<()> {
77        self.stream.synchronize().map_err(Into::into)
78    }
79}