candle-metal-kernels/src/reduce.metal


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
139
140
141
142
143

#include <metal_stdlib>
using namespace metal;

#define MAX(x, y) ((x) > (y) ? (x) : (y))

METAL_FUNC uint get_strided_index(
    uint idx,
    constant size_t &num_dims,
    constant size_t *dims,
    constant size_t *strides
) {
    uint strided_i = 0;
    for (uint d = 0; d < num_dims; d++) {
        uint dim_idx = num_dims - 1 - d;
        strided_i += (idx % dims[dim_idx]) * strides[dim_idx];
        idx /= dims[dim_idx];
    }
    return strided_i;
}

constant int THREADGROUP_SIZE = 2048;

# define REDUCE(FN, NAME, T) \
kernel void NAME( \
    constant size_t &src_numel, \
    constant size_t &el_to_sum_per_block, \
    device const T *src,  \
    device T *dst, \
    uint id [[ thread_position_in_grid ]], \
    uint tid [[ thread_index_in_threadgroup ]], \
    uint dst_id [[ threadgroup_position_in_grid ]], \
    uint block_dim [[ threads_per_threadgroup ]] \
) { \
     \
   threadgroup float shared_memory[THREADGROUP_SIZE]; \
      \
   shared_memory[tid] = 0; \
   /* \
   // Elements summed in this block range from dst_id * el_to_sum_per_block  \
   // to (dst_id + 1) * el_to_sum_per_block. \
   */ \
   size_t start_idx = dst_id * el_to_sum_per_block; \
   size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); \
   size_t idx = start_idx + tid; \
   while (idx < stop_idx) { \
     /* \
     // TODO: Fast version for the contiguous case. \
     // size_t strided_i = get_strided_index(idx, num_dims, dims, strides); \
     */ \
     T x = shared_memory[tid]; \
     T y = src[idx]; \
     shared_memory[tid] = FN; \
     idx += block_dim; \
   } \
      \
   threadgroup_barrier(mem_flags::mem_none); \
    \
   /* \
   // reduction in shared memory \
   */ \
   for (uint s = block_dim / 2; s > 0; s >>= 1) { \
       if (tid < s) { \
           T x = shared_memory[tid]; \
           T y = shared_memory[tid + s]; \
           shared_memory[tid] = FN; \
       } \
       threadgroup_barrier(mem_flags::mem_none); \
   } \
    \
   dst[dst_id] = shared_memory[0]; \
} \


REDUCE(x + y, fast_sum_float, float)
REDUCE(x * y, fast_mul_float, float)
REDUCE(max(x, y), fast_max_float, float)

#define SOFTMAX(NAME, T)                                                          \
kernel void NAME(                                                                 \
    constant size_t &src_numel,                                                   \
    constant size_t &el_to_sum_per_block,                                         \
    device const T *src,                                                          \
    device T *dst,                                                                \
                                                                                  \
    uint id [[ thread_position_in_grid ]],                                        \
    uint tid [[ thread_index_in_threadgroup ]],                                   \
    uint dst_id [[ threadgroup_position_in_grid ]],                               \
    uint block_dim [[ threads_per_threadgroup ]]                                  \
) {                                                                               \
    threadgroup float shared_memory[THREADGROUP_SIZE];                                \
    shared_memory[tid] = -INFINITY;                                            \
    size_t start_idx = dst_id * el_to_sum_per_block;                              \
    size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel);            \
    size_t idx = start_idx + tid;                                                 \
                                                                                  \
    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
                                                                                  \
    while (idx < stop_idx) {                                                      \
        shared_memory[tid] = MAX(shared_memory[tid], src[idx]);                   \
        idx += block_dim;                                                         \
    }                                                                             \
                                                                                  \
    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
                                                                                  \
    for (uint s = block_dim / 2; s > 0; s >>= 1) {                                \
        if (tid < s) {                                                            \
            shared_memory[tid] = MAX(shared_memory[tid], shared_memory[tid + s]); \
        }                                                                         \
    }                                                                             \
                                                                                  \
    threadgroup_barrier(mem_flags::mem_threadgroup);                              \
                                                                                  \
    float _max = shared_memory[0];                                                    \
                                                                                  \
    shared_memory[tid] = 0;                                                       \
                                                                                  \
    idx = start_idx + tid;                                                        \
    while (idx < stop_idx) {                                                      \
        const T val = T(exp(src[idx] - _max));                                    \
        dst[idx] = val;                                                           \
        shared_memory[tid] += val;                                                \
        idx += block_dim;                                                         \
    }                                                                             \
    for (uint s = block_dim / 2; s > 0; s >>= 1) {                                \
        if (tid < s) {                                                            \
            shared_memory[tid] += shared_memory[tid + s];                         \
        }                                                                         \
        threadgroup_barrier(mem_flags::mem_threadgroup);                          \
    }                                                                             \
                                                                                  \
    const T inv_acc = T(1/shared_memory[0]);                                         \
    idx = start_idx + tid;                                                        \
    while (idx < stop_idx) {                                                      \
        dst[idx] *= inv_acc;                                                      \
        idx += block_dim;                                                         \
    }                                                                             \
}                                                                                 \

SOFTMAX(softmax_float, float)
SOFTMAX(softmax_half, half)
#if __METAL_VERSION__ >= 310
SOFTMAX(softmax_bfloat, bfloat)
#endif