dfdx 0.13.0

Ergonomic auto differentiation in Rust, with pytorch like apis.
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

#include "cuda_utils.cuh"

#define LONG_BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, FUNC, DFDX, DFDY) \
extern "C" __global__ void FORWARD( \
    const OP_STRUCT op, \
    const size_t numel, \
    const size_t num_dims, \
    const size_t *info, \
    const TYPENAME *lhs, \
    const TYPENAME *rhs, \
    TYPENAME *out \
) { \
    const size_t *dims = info; \
    const size_t *lhs_strides = info + num_dims; \
    const size_t *rhs_strides = info + 2 * num_dims; \
    for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
        unsigned int tmp_i = i; \
        unsigned int lhs_i = 0; \
        unsigned int rhs_i = 0; \
        for (int d = num_dims - 1; d >= 0; d--) { \
            unsigned int i_dim = tmp_i % dims[d]; \
            lhs_i += i_dim * lhs_strides[d]; \
            rhs_i += i_dim * rhs_strides[d]; \
            tmp_i /= dims[d]; \
        } \
        TYPENAME x = lhs ? lhs[lhs_i] : out[i]; \
        TYPENAME y = rhs ? rhs[rhs_i] : out[i]; \
        TYPENAME fx; \
        FUNC\
        out[i] = fx; \
    } \
} \
\
extern "C" __global__ void BACKWARD_LHS( \
    const OP_STRUCT op, \
    const size_t numel, \
    const size_t num_dims, \
    const size_t *info, \
    const TYPENAME *lhs, \
    TYPENAME *grad_lhs, \
    const size_t chunk_len, \
    const TYPENAME *rhs, \
    const TYPENAME *grad_out \
) { \
    const size_t *dims = info + 0 * num_dims; \
    const size_t *out_strides = info + 1 * num_dims; \
    const size_t *rhs_strides = info + 2 * num_dims; \
    TYPENAME zero = 0.0; \
    for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
        unsigned int tmp_i = i; \
        unsigned int out_i = 0; \
        unsigned int rhs_i = 0; \
        for (int d = num_dims - 1; d >= 0; d--) { \
            unsigned int i_dim = tmp_i % dims[d]; \
            out_i += i_dim * out_strides[d]; \
            rhs_i += i_dim * rhs_strides[d]; \
            tmp_i /= dims[d]; \
        } \
        TYPENAME x = lhs ? lhs[i / chunk_len] : zero; \
        TYPENAME y = rhs ? rhs[rhs_i] : zero; \
        TYPENAME go = grad_out[out_i]; \
        TYPENAME dfdx = (DFDX); \
        chunk_sum(chunk_len, dfdx * go, grad_lhs); \
    } \
} \
\
extern "C" __global__ void BACKWARD_RHS( \
    const OP_STRUCT op, \
    const size_t numel, \
    const size_t num_dims, \
    const size_t *info, \
    const TYPENAME *lhs, \
    const TYPENAME *rhs, \
    TYPENAME *grad_rhs, \
    const size_t chunk_len, \
    const TYPENAME *grad_out \
) { \
    const size_t *dims = info + 3 * num_dims; \
    const size_t *out_strides = info + 4 * num_dims; \
    const size_t *lhs_strides = info + 5 * num_dims; \
    TYPENAME zero = 0.0; \
    for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \
        unsigned int tmp_i = i; \
        unsigned int lhs_i = 0; \
        unsigned int out_i = 0; \
        for (int d = num_dims - 1; d >= 0; d--) { \
            unsigned int i_dim = tmp_i % dims[d]; \
            lhs_i += i_dim * lhs_strides[d]; \
            out_i += i_dim * out_strides[d]; \
            tmp_i /= dims[d]; \
        } \
        TYPENAME x = lhs ? lhs[lhs_i] : zero; \
        TYPENAME y = rhs ? rhs[i / chunk_len] : zero; \
        TYPENAME go = grad_out[out_i]; \
        TYPENAME dfdy = (DFDY); \
        chunk_sum(chunk_len, dfdy * go, grad_rhs); \
    } \
}

#define BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, FUNC, DFDX, DFDY) \
    LONG_BINARY_OP(TYPENAME, FORWARD, BACKWARD_LHS, BACKWARD_RHS, OP_STRUCT, fx = (FUNC);, DFDX, DFDY)