singe-npp 0.1.0-alpha.8

Safe Rust wrappers for NVIDIA Performance Primitives library (NPP).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! Scratch-buffer support for NPP wrappers.
//!
//! `ScratchBuffer` owns temporary device memory used by NPP operations that
//! expose host-side buffer-size queries.

use singe_cuda::memory::DeviceMemory;
use singe_npp_sys as sys;

use crate::{
    error::{Error, Result},
    types::BufferDescriptor,
    utility::to_i32,
};

#[derive(Debug)]
pub struct ScratchBuffer {
    memory: DeviceMemory<u8>,
}

impl ScratchBuffer {
    pub fn create(bytes: usize) -> Result<Self> {
        Ok(Self {
            memory: DeviceMemory::create(bytes)?,
        })
    }

    pub fn from_memory(memory: DeviceMemory<u8>) -> Self {
        Self { memory }
    }

    pub const fn len(&self) -> usize {
        self.memory.len()
    }

    pub const fn is_empty(&self) -> bool {
        self.memory.is_empty()
    }

    pub fn require(&self, bytes: usize) -> Result<()> {
        if self.len() < bytes {
            return Err(Error::LengthMismatch {
                name: "scratch buffer".into(),
                expected: bytes,
                actual: self.len(),
            });
        }

        Ok(())
    }

    pub fn as_mut_ptr(&mut self) -> *mut u8 {
        self.memory.as_mut_ptr()
    }

    pub fn into_device_memory(self) -> DeviceMemory<u8> {
        self.memory
    }
}

#[derive(Debug)]
pub(crate) struct BufferDescriptors {
    memory: DeviceMemory<sys::NppiBufferDescriptor>,
}

impl BufferDescriptors {
    pub fn as_mut_ptr(&mut self) -> *mut sys::NppiBufferDescriptor {
        self.memory.as_mut_ptr()
    }
}

pub(crate) fn create_buffer_descriptors(
    buffer_sizes: impl IntoIterator<Item = usize>,
) -> Result<(Vec<DeviceMemory<u8>>, BufferDescriptors)> {
    let buffer_sizes: Vec<_> = buffer_sizes.into_iter().collect();
    let mut buffers = Vec::with_capacity(buffer_sizes.len());
    let mut descriptors = Vec::with_capacity(buffer_sizes.len());

    for bytes in buffer_sizes {
        let descriptor_size = to_i32(bytes, "buffer size")?;
        let buffer = DeviceMemory::<u8>::create(bytes)?;
        descriptors.push(BufferDescriptor {
            data: buffer.as_mut_ptr().cast(),
            size: descriptor_size,
        });
        buffers.push(buffer);
    }

    let raw_descriptors = descriptors
        .into_iter()
        .map(sys::NppiBufferDescriptor::from)
        .collect::<Vec<_>>();

    Ok((
        buffers,
        BufferDescriptors {
            memory: DeviceMemory::from_slice(&raw_descriptors)?,
        },
    ))
}

pub(crate) fn create_repeated_buffer_descriptors(
    count: usize,
    bytes_per_buffer: usize,
) -> Result<(Vec<DeviceMemory<u8>>, BufferDescriptors)> {
    create_buffer_descriptors(std::iter::repeat_n(bytes_per_buffer, count))
}