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
|
use crate::{CpuStorage, DType, Result, Shape};
use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
pub type CudaError = cudarc::driver::DriverError;
#[derive(Debug, Clone)]
pub struct CudaDevice(std::sync::Arc<cudarc::driver::CudaDevice>);
// TODO: Switch to pre-compiled PTX kernels rather than compiling on the fly.
const AFFINE_CU: &str = r#"
extern "C" __global__ void affine_f32(
const size_t numel,
const float *x,
float *y,
const float mul,
const float add
) {
unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= numel) {
return;
}
y[i] = x[i] * mul + add;
}
"#;
impl CudaDevice {
pub(crate) fn new(ordinal: usize) -> Result<Self> {
let device = cudarc::driver::CudaDevice::new(ordinal)?;
Ok(Self(device))
}
pub(crate) fn ordinal(&self) -> usize {
self.0.ordinal()
}
pub(crate) fn zeros_impl(&self, shape: &Shape, dtype: DType) -> Result<CudaStorage> {
let elem_count = shape.elem_count();
match dtype {
DType::F32 => {
let data = self.0.alloc_zeros::<f32>(elem_count)?;
Ok(CudaStorage::F32(data))
}
DType::F64 => {
let data = self.0.alloc_zeros::<f64>(elem_count)?;
Ok(CudaStorage::F64(data))
}
}
}
pub(crate) fn cuda_from_cpu_storage(&self, storage: &CpuStorage) -> Result<CudaStorage> {
match storage {
CpuStorage::F32(storage) => {
let data = self.0.htod_sync_copy(storage)?;
Ok(CudaStorage::F32(data))
}
CpuStorage::F64(storage) => {
let data = self.0.htod_sync_copy(storage)?;
Ok(CudaStorage::F64(data))
}
}
}
}
#[derive(Debug, Clone)]
pub enum CudaStorage {
F32(CudaSlice<f32>),
F64(CudaSlice<f64>),
}
impl CudaStorage {
pub fn dtype(&self) -> DType {
match self {
Self::F32(_) => DType::F32,
Self::F64(_) => DType::F64,
}
}
pub fn device(&self) -> CudaDevice {
match self {
Self::F32(slice) => CudaDevice(slice.device()),
Self::F64(slice) => CudaDevice(slice.device()),
}
}
pub(crate) fn affine_impl(
&self,
shape: &Shape,
stride: &[usize],
mul: f64,
add: f64,
) -> Result<Self> {
match self {
Self::F32(arg) => {
if !shape.is_contiguous(stride) {
todo!("affine is only implemented for the contiguous case")
}
let dev = arg.device();
let module_name = "affine_f32";
if !dev.has_func(module_name, module_name) {
let ptx = cudarc::nvrtc::compile_ptx(AFFINE_CU).unwrap();
dev.load_ptx(ptx, module_name, &[module_name])?;
}
let elem_count = shape.elem_count();
let fwd_fn = dev.get_func(module_name, module_name).unwrap();
let cfg = LaunchConfig::for_num_elems(elem_count as u32);
// SAFETY: if this function returns Ok(..), the kernel has been applied
// and has set the initially unset memory.
let out = unsafe { dev.alloc::<f32>(elem_count) }?;
let params = (elem_count, arg, &out, mul as f32, add as f32);
// SAFETY: well, well, well...
unsafe { fwd_fn.launch(cfg, params) }?;
Ok(Self::F32(out))
}
Self::F64(_) => {
todo!()
}
}
}
pub(crate) fn to_cpu_storage(&self) -> Result<CpuStorage> {
match self {
Self::F32(slice) => {
let dev = slice.device();
let cpu_storage = dev.dtoh_sync_copy(slice)?;
Ok(CpuStorage::F32(cpu_storage))
}
Self::F64(slice) => {
let dev = slice.device();
let cpu_storage = dev.dtoh_sync_copy(slice)?;
Ok(CpuStorage::F64(cpu_storage))
}
}
}
}
|
