mircuda 0.1.0

Native, explicit-stream Rust gateway to NVIDIA CUDA
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

use std::marker::PhantomData;

use crate::{Context, DeviceBuffer, DeviceElement, Error, Result, Stream, bf16, f16};

mod sealed {
    pub trait Sealed {}
}

/// Storage type accepted as the output of a native dense CUTLASS plan.
pub trait DenseMatmulOutput: DeviceElement + sealed::Sealed {
    #[doc(hidden)]
    const NATIVE_TYPE: mircuda_sys::DenseMatmulDataType;
}

/// Storage type accepted as an input of a native dense CUTLASS plan.
pub trait DenseMatmulElement: DenseMatmulOutput {
    #[doc(hidden)]
    const INPUT_NATIVE_TYPE: mircuda_sys::DenseMatmulDataType;
}

impl sealed::Sealed for f16 {}
impl sealed::Sealed for bf16 {}
impl sealed::Sealed for f32 {}

impl DenseMatmulElement for f16 {
    const INPUT_NATIVE_TYPE: mircuda_sys::DenseMatmulDataType =
        mircuda_sys::DenseMatmulDataType::F16;
}

impl DenseMatmulElement for bf16 {
    const INPUT_NATIVE_TYPE: mircuda_sys::DenseMatmulDataType =
        mircuda_sys::DenseMatmulDataType::Bf16;
}

impl DenseMatmulOutput for f16 {
    const NATIVE_TYPE: mircuda_sys::DenseMatmulDataType = mircuda_sys::DenseMatmulDataType::F16;
}

impl DenseMatmulOutput for bf16 {
    const NATIVE_TYPE: mircuda_sys::DenseMatmulDataType = mircuda_sys::DenseMatmulDataType::Bf16;
}

impl DenseMatmulOutput for f32 {
    const NATIVE_TYPE: mircuda_sys::DenseMatmulDataType = mircuda_sys::DenseMatmulDataType::F32;
}

/// Fixed geometry for `C[M,N] = alpha * A[M,K] * B[N,K]^T + beta * C[M,N]`.
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub struct DenseMatmulSpec {
    m: usize,
    n: usize,
    k: usize,
}

impl DenseMatmulSpec {
    /// Creates a validated row-major dense projection shape.
    pub fn new(m: usize, n: usize, k: usize) -> Result<Self> {
        if m == 0 || n == 0 || k == 0 {
            return Err(Error::InvalidMatmulShape);
        }
        for value in [m, n, k] {
            let _ = i32::try_from(value)?;
        }
        let _ = m.checked_mul(k).ok_or(Error::InvalidMatmulShape)?;
        let _ = n.checked_mul(k).ok_or(Error::InvalidMatmulShape)?;
        let _ = m.checked_mul(n).ok_or(Error::InvalidMatmulShape)?;
        Ok(Self { m, n, k })
    }

    /// Output row count.
    #[must_use]
    pub const fn m(self) -> usize {
        self.m
    }

    /// Output column count.
    #[must_use]
    pub const fn n(self) -> usize {
        self.n
    }

    /// Reduction dimension.
    #[must_use]
    pub const fn k(self) -> usize {
        self.k
    }
}

/// Stream-bound dense CUTLASS plan with persistent native workspace.
#[derive(Debug)]
pub struct DenseMatmulPlan<T: DenseMatmulElement, O: DenseMatmulOutput = T> {
    native: mircuda_sys::DenseMatmulPlan,
    spec: DenseMatmulSpec,
    marker: PhantomData<(T, O)>,
}

impl<T: DenseMatmulElement, O: DenseMatmulOutput> DenseMatmulPlan<T, O> {
    /// Creates a fixed-shape plan without synchronizing the stream.
    pub fn new(context: &Context, stream: &Stream, spec: DenseMatmulSpec) -> Result<Self> {
        let native_spec = mircuda_sys::DenseMatmulSpec {
            m: spec.m,
            n: spec.n,
            k: spec.k,
            input_type: T::INPUT_NATIVE_TYPE,
            output_type: O::NATIVE_TYPE,
        };
        Ok(Self {
            native: context.native.create_dense_matmul_plan(&stream.native, native_spec)?,
            spec,
            marker: PhantomData,
        })
    }

    /// Returns the immutable plan geometry.
    #[must_use]
    pub const fn spec(&self) -> DenseMatmulSpec {
        self.spec
    }

    /// Native workspace retained for repeated launches.
    #[must_use]
    #[allow(clippy::missing_const_for_fn)]
    pub fn workspace_bytes(&self) -> usize {
        self.native.workspace_bytes()
    }

    /// Enqueues the fixed-shape multiplication on the plan stream.
    pub fn execute(
        &mut self,
        stream: &Stream,
        a: &DeviceBuffer<T>,
        b: &DeviceBuffer<T>,
        c: &mut DeviceBuffer<O>,
        alpha: f32,
        beta: f32,
    ) -> Result<()> {
        validate_len("A", self.spec.m * self.spec.k, a.len())?;
        validate_len("B", self.spec.n * self.spec.k, b.len())?;
        validate_len("C", self.spec.m * self.spec.n, c.len())?;
        Ok(self
            .native
            .execute(&stream.native, &a.native, &b.native, &c.native, alpha, beta)?)
    }
}

const fn validate_len(operand: &'static str, expected: usize, actual: usize) -> Result<()> {
    if expected == actual {
        Ok(())
    } else {
        Err(Error::MatmulLengthMismatch { operand, expected, actual })
    }
}