oxidized-cuda-kernels 0.1.1

Additional CUDA kernels for Oxidized Transformers
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

use cudarc::driver::sys::{CUdeviceptr, CUstream};
use half::{bf16, f16};

macro_rules! nonzero {
    ($name:ident, $type:ty) => {
        extern "C" {
            pub(crate) fn $name(
                num_elems: usize,
                num_dims: usize,
                info: *const usize,
                input: *const $type,
                out: *mut CUdeviceptr,
                num_nonzeros: *mut usize,
                stream: CUstream,
            );
        }
    };
}

nonzero!(nonzero_bf16, bf16);
nonzero!(nonzero_f16, f16);
nonzero!(nonzero_f32, f32);
nonzero!(nonzero_f64, f64);
nonzero!(nonzero_u8, u8);
nonzero!(nonzero_u32, u32);
nonzero!(nonzero_i64, i64);