fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
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
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

//! Tensor ops, one module per family. Each computes on the inputs' device and
//! attaches a rule from [`autograd::ops`](crate::autograd::ops).

pub mod activation;
pub mod arith;
pub mod conv;
pub mod index;
pub mod matmul;
pub mod norm;
pub mod pool;
pub mod reduce;
pub mod unary;
pub mod view;

use rayon::prelude::*;

use crate::cuda::{ffi, kernels};

use super::Tensor;
use super::shape;
use super::storage::Storage;

/// Below this many elements rayon costs more than it saves.
const PARALLEL_THRESHOLD: usize = 1 << 15;

pub(crate) fn unary_op(
    x: &Tensor,
    cpu: impl Fn(f32) -> f32 + Send + Sync,
    gpu: ffi::UnaryKernel,
) -> Tensor {
    match x.storage() {
        Storage::Cpu(data) => Tensor::from_vec(map(data, cpu), x.shape()),
        Storage::Cuda(buf) => {
            let out = kernels::unary(buf, x.numel(), gpu).expect("cuda unary op");
            Tensor::raw(Storage::Cuda(out), x.shape().to_vec(), x.device())
        }
    }
}

pub(crate) fn scalar_op(
    x: &Tensor,
    s: f32,
    cpu: impl Fn(f32, f32) -> f32 + Send + Sync,
    gpu: ffi::ScalarKernel,
) -> Tensor {
    match x.storage() {
        Storage::Cpu(data) => Tensor::from_vec(map(data, |v| cpu(v, s)), x.shape()),
        Storage::Cuda(buf) => {
            let out = kernels::scalar(buf, s, x.numel(), gpu).expect("cuda scalar op");
            Tensor::raw(Storage::Cuda(out), x.shape().to_vec(), x.device())
        }
    }
}

pub(crate) fn binary_op(
    a: &Tensor,
    b: &Tensor,
    cpu: impl Fn(f32, f32) -> f32 + Send + Sync,
    gpu: ffi::BinaryKernel,
) -> Tensor {
    assert_eq!(
        a.device(),
        b.device(),
        "device mismatch: {} and {}; move one with .to()",
        a.device(),
        b.device()
    );

    if a.shape() != b.shape() {
        return broadcast_binary_op(a, b, cpu, gpu);
    }

    match (a.storage(), b.storage()) {
        (Storage::Cpu(x), Storage::Cpu(y)) => Tensor::from_vec(zip(x, y, cpu), a.shape()),
        (Storage::Cuda(x), Storage::Cuda(y)) => {
            let out = kernels::binary(x, y, a.numel(), gpu).expect("cuda binary op");
            Tensor::raw(Storage::Cuda(out), a.shape().to_vec(), a.device())
        }
        _ => unreachable!("devices already checked equal"),
    }
}

/// GPU expands both sides first, CPU walks broadcast strides and skips the copy.
fn broadcast_binary_op(
    a: &Tensor,
    b: &Tensor,
    cpu: impl Fn(f32, f32) -> f32 + Send + Sync,
    gpu: ffi::BinaryKernel,
) -> Tensor {
    let out_shape = shape::broadcast(a.shape(), b.shape());

    if a.is_cuda() {
        let a_full = view::broadcast_to(a, &out_shape);
        let b_full = view::broadcast_to(b, &out_shape);
        return binary_op(&a_full, &b_full, cpu, gpu);
    }

    let ndim = out_shape.len();
    let out_strides = shape::strides_for(&out_shape);
    let a_strides = shape::broadcast_strides(&shape::pad_left(a.shape(), ndim), &out_shape);
    let b_strides = shape::broadcast_strides(&shape::pad_left(b.shape(), ndim), &out_shape);
    let (a_data, b_data) = (a.to_vec(), b.to_vec());

    let data: Vec<f32> = (0..shape::numel(&out_shape))
        .map(|flat| {
            let (mut rest, mut ai, mut bi) = (flat, 0usize, 0usize);
            for d in 0..ndim {
                let coord = rest / out_strides[d];
                rest %= out_strides[d];
                ai += coord * a_strides[d];
                bi += coord * b_strides[d];
            }
            cpu(a_data[ai], b_data[bi])
        })
        .collect();

    Tensor::from_vec(data, &out_shape)
}

pub(crate) fn map(data: &[f32], f: impl Fn(f32) -> f32 + Send + Sync) -> Vec<f32> {
    if data.len() >= PARALLEL_THRESHOLD {
        data.par_iter().map(|&v| f(v)).collect()
    } else {
        data.iter().map(|&v| f(v)).collect()
    }
}

pub(crate) fn zip(a: &[f32], b: &[f32], f: impl Fn(f32, f32) -> f32 + Send + Sync) -> Vec<f32> {
    debug_assert_eq!(a.len(), b.len());
    if a.len() >= PARALLEL_THRESHOLD {
        a.par_iter()
            .zip(b.par_iter())
            .map(|(&x, &y)| f(x, y))
            .collect()
    } else {
        a.iter().zip(b.iter()).map(|(&x, &y)| f(x, y)).collect()
    }
}