tract_cuda/kernels/
mod.rs

1#![allow(unused)]
2
3pub mod array;
4pub mod binary;
5pub mod causal_conv1d_update;
6pub mod conv;
7pub mod conv_cudnn;
8mod conv_source;
9pub mod element_wise;
10pub mod fft;
11pub mod flash_attn;
12pub mod gdn_recurrent;
13mod iff;
14pub(crate) mod launch_args;
15pub mod matmul;
16pub mod nn;
17pub(crate) mod utils;
18
19use std::env;
20use std::path::{Path, PathBuf};
21use std::sync::OnceLock;
22
23use crate::ops::GgmlQuantQ81Fact;
24use crate::tensor::{CudaBuffer, CudaTensor};
25use anyhow::{bail, ensure};
26use cudarc::driver::{CudaView, CudaViewMut};
27pub use iff::Iff;
28use tract_core::internal::ExoticFact;
29use tract_core::prelude::{TDim, TractResult};
30use tract_core::tract_linalg::block_quant::{BlockQuant, BlockQuantFact, Q4_0, Q8_1};
31use tract_gpu::tensor::{DeviceTensor, OwnedDeviceTensor};
32use tract_gpu::utils::as_q40_tensor;
33
34const MAX_THREADS: usize = 1024;
35const WARP_SIZE: usize = 32;
36
37static CUBIN_FOLDER: OnceLock<PathBuf> = OnceLock::new();
38
39/// The user's cache directory: `%LOCALAPPDATA%` on Windows, `Library/Caches`
40/// on Apple targets, `XDG_CACHE_HOME` or `~/.cache` elsewhere.
41///
42/// Per the XDG spec a relative `XDG_CACHE_HOME` is ignored rather than resolved
43/// against the working directory.
44fn user_cache_dir() -> Option<PathBuf> {
45    if cfg!(target_os = "windows") {
46        return env::var_os("LOCALAPPDATA").filter(|v| !v.is_empty()).map(PathBuf::from);
47    }
48    let home = || env::var_os("HOME").filter(|v| !v.is_empty()).map(PathBuf::from);
49    if cfg!(target_vendor = "apple") {
50        return home().map(|h| h.join("Library/Caches"));
51    }
52    env::var_os("XDG_CACHE_HOME")
53        .map(PathBuf::from)
54        .filter(|p| p.is_absolute())
55        .or_else(|| home().map(|h| h.join(".cache")))
56}
57
58pub fn cubin_dir() -> &'static Path {
59    CUBIN_FOLDER
60        .get_or_init(|| {
61            user_cache_dir()
62                .unwrap_or_else(|| ".cache".into())
63                .join("tract")
64                .join(env!("CARGO_PKG_VERSION"))
65                .join("cuda")
66                .join(crate::utils::REQUIRED_CUDA_API.to_string())
67                .join("cubins")
68        })
69        .as_path()
70}
71
72const ELEMENT_WISE_OPS: &str = include_str!("cu/element_wise.cu");
73const BINARY_OPS: &str = include_str!("cu/binary.cu");
74const ARRAY_OPS: &str = include_str!("cu/array.cu");
75const NN_OPS: &str = include_str!("cu/nn.cu");
76const GGML_MM_MV: &str = include_str!("cu/mm_mv.cu");
77const GGML_MM_MV_Q: &str = include_str!("cu/mm_mv_q.cu");
78const GGML_QUANTIZE: &str = include_str!("cu/quantize.cu");
79const FLASH_ATTN: &str = include_str!("cu/flash_attn.cu");
80const GDN_RECURRENT: &str = include_str!("cu/gdn_recurrent.cu");
81const FFT_OPS: &str = include_str!("cu/fft.cu");
82pub const COMMON_H: &str = include_str!("cu/common.cuh");
83
84/// The convolution kernels are rendered rather than read: every instantiation the Cnn
85/// library carries, concatenated once.
86fn cnn_ops() -> &'static str {
87    static CNN_OPS: OnceLock<String> = OnceLock::new();
88    CNN_OPS.get_or_init(conv_source::conv_library_source)
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
92pub enum LibraryName {
93    ElementWise,
94    Binary,
95    Array,
96    NN,
97    Cnn,
98    GdnRecurrent,
99    Ggml,
100    GgmlQ,
101    Quant,
102    FlashAttn,
103    Fft,
104}
105
106fn fnv1a64(text: &str) -> u64 {
107    const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
108    const FNV_PRIME: u64 = 0x00000100000001B3;
109
110    let mut hash = FNV_OFFSET_BASIS;
111    for b in text.as_bytes() {
112        hash ^= *b as u64;
113        hash = hash.wrapping_mul(FNV_PRIME);
114    }
115    hash
116}
117
118impl LibraryName {
119    // Libraries compiled when a CUDA context is created. Gated DeltaNet is
120    // intentionally omitted: it is model-specific and compiled on first use.
121    pub const EAGER: [LibraryName; 10] = [
122        Self::FlashAttn,
123        Self::ElementWise,
124        Self::Binary,
125        Self::Array,
126        Self::NN,
127        Self::Cnn,
128        Self::Ggml,
129        Self::GgmlQ,
130        Self::Quant,
131        Self::Fft,
132    ];
133
134    pub fn content(&self) -> &str {
135        match self {
136            Self::ElementWise => ELEMENT_WISE_OPS,
137            Self::Binary => BINARY_OPS,
138            Self::Array => ARRAY_OPS,
139            Self::NN => NN_OPS,
140            Self::Cnn => cnn_ops(),
141            Self::GdnRecurrent => GDN_RECURRENT,
142            Self::Ggml => GGML_MM_MV,
143            Self::GgmlQ => GGML_MM_MV_Q,
144            Self::Quant => GGML_QUANTIZE,
145            Self::FlashAttn => FLASH_ATTN,
146            Self::Fft => FFT_OPS,
147        }
148    }
149
150    pub fn cubin_path(&self) -> PathBuf {
151        let basename = match self {
152            Self::ElementWise => "element_wise",
153            Self::Binary => "binary",
154            Self::Array => "array",
155            Self::NN => "nn",
156            Self::Cnn => "cnn",
157            Self::GdnRecurrent => "gdn_recurrent",
158            Self::Ggml => "mm_mv",
159            Self::GgmlQ => "mm_mv_q",
160            Self::Quant => "quantize",
161            Self::FlashAttn => "flash_attn",
162            Self::Fft => "fft",
163        };
164        let hash = fnv1a64(self.content());
165        cubin_dir().join(format!("{}_{}.cubin", basename, hash))
166    }
167}
168
169pub use tract_gpu::utils::BroadcastKind;
170
171fn tensor_size(t: &DeviceTensor) -> usize {
172    let exotic_fact: Option<&dyn ExoticFact> = match t {
173        DeviceTensor::Owned(ot) => {
174            let cuda_tensor =
175                ot.downcast_ref::<CudaTensor>().expect("Non Cuda-Tensor in a Cuda Context");
176            cuda_tensor.exotic_fact()
177        }
178        DeviceTensor::ArenaView(av) => av.exotic_fact(),
179    };
180
181    if let Some(of) = exotic_fact {
182        of.buffer_sizes()
183            .iter()
184            .sum::<TDim>()
185            .as_i64()
186            .expect("Symbols should be resolved at this point") as usize
187    } else {
188        t.len() * t.datum_type().size_of()
189    }
190}
191
192pub fn get_cuda_view(t: &DeviceTensor) -> CudaView<'_, u8> {
193    let size = tensor_size(t);
194    get_sliced_cuda_view(t, 0, size).unwrap()
195}
196
197// NOTE: offset and len are in bytes
198pub fn get_sliced_cuda_view(
199    t: &DeviceTensor,
200    offset: usize,
201    len: usize,
202) -> TractResult<CudaView<'_, u8>> {
203    ensure!(offset + len <= tensor_size(t));
204    let buffer = t.device_buffer().downcast_ref::<CudaBuffer>().unwrap();
205    let offset = t.buffer_offset::<usize>() + offset;
206    Ok(buffer.slice(offset..(offset + len)))
207}
208
209pub fn get_cuda_view_mut(t: &DeviceTensor) -> CudaViewMut<'_, u8> {
210    let size = t.len() * t.datum_type().size_of();
211    get_sliced_cuda_view_mut(t, 0, size).unwrap()
212}
213
214// NOTE: offset and len are in bytes
215pub fn get_sliced_cuda_view_mut(
216    t: &DeviceTensor,
217    offset: usize,
218    len: usize,
219) -> TractResult<CudaViewMut<'_, u8>> {
220    ensure!(offset + len <= t.len() * t.datum_type().size_of());
221    let buffer: &CudaBuffer = t.device_buffer().downcast_ref::<CudaBuffer>().unwrap();
222    let offset = t.buffer_offset::<usize>() + offset;
223    let ptr: *const CudaBuffer = buffer;
224    let mut_buffer: &mut CudaBuffer = unsafe { (ptr as *mut CudaBuffer).as_mut().unwrap() };
225    Ok(mut_buffer.inner.slice_mut(offset..(offset + len)))
226}
tract_cuda/kernels/mod.rs

tract_cuda/kernels/
mod.rs
