hodu_cpu_kernels 0.2.0

hodu cpu 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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210

#include "ops_matrix.h"
#include "types.h"
#include <stdint.h>
#include <string.h>

// ============================================================================
// BATCHED MATRIX MULTIPLICATION (MATMUL)
// ============================================================================
//
// Performs batched matrix multiplication with broadcasting support.
// Computes C[..., i, j] = sum_k A[..., i, k] * B[..., k, j]
//
// Metadata layout:
// - metadata[0]: num_els (total number of output elements)
// - metadata[1]: lhs_ndim (number of dimensions in lhs)
// - metadata[2]: rhs_ndim (number of dimensions in rhs)
// - metadata[3]: batch_ndim (number of batch dimensions in output)
// - metadata[4..4+lhs_ndim]: lhs_shape
// - metadata[4+lhs_ndim..4+lhs_ndim+rhs_ndim]: rhs_shape
// - metadata[4+lhs_ndim+rhs_ndim..4+lhs_ndim+rhs_ndim+batch_ndim]: batch_shape
// - metadata[...+lhs_ndim]: lhs_strides
// - metadata[...+rhs_ndim]: rhs_strides
// - metadata[...]: lhs_offset
// - metadata[...+1]: rhs_offset
// - metadata[...+2]: M (rows of lhs matrix)
// - metadata[...+3]: K (cols of lhs / rows of rhs)
// - metadata[...+4]: N (cols of rhs matrix)
//
// Algorithm:
// For each output element (batch_idx, i, j):
// 1. Decompose flat batch_idx into multi-dimensional batch indices
// 2. Map batch indices to lhs/rhs with broadcasting (size 1 dims stay at 0)
// 3. Compute dot product: sum_k lhs[batch, i, k] * rhs[batch, k, j]
//
// Broadcasting:
// Batch dimensions of size 1 are broadcast by using index 0 for that dimension.

/// Macro to implement batched matrix multiplication
///
/// @param TYPE C type for the operation
/// @param TYPE_SUFFIX Suffix for function naming
#define MATMUL_OP(TYPE, TYPE_SUFFIX)                                                               \
    void matmul_##TYPE_SUFFIX(const void *lhs_ptr, const void *rhs_ptr, void *output_ptr,          \
                              const size_t *metadata) {                                            \
        const TYPE *lhs = (const TYPE *)lhs_ptr;                                                   \
        const TYPE *rhs = (const TYPE *)rhs_ptr;                                                   \
        TYPE *output = (TYPE *)output_ptr;                                                         \
                                                                                                   \
        const size_t num_els = metadata[0];                                                        \
        const size_t lhs_ndim = metadata[1];                                                       \
        const size_t rhs_ndim = metadata[2];                                                       \
        const size_t batch_ndim = metadata[3];                                                     \
                                                                                                   \
        const size_t *lhs_shape = metadata + 4;                                                    \
        const size_t *rhs_shape = lhs_shape + lhs_ndim;                                            \
        const size_t *batch_shape = rhs_shape + rhs_ndim;                                          \
        const size_t *lhs_strides = batch_shape + batch_ndim;                                      \
        const size_t *rhs_strides = lhs_strides + lhs_ndim;                                        \
        const size_t lhs_offset = *(rhs_strides + rhs_ndim);                                       \
        const size_t rhs_offset = *(rhs_strides + rhs_ndim + 1);                                   \
        const size_t M = *(rhs_strides + rhs_ndim + 2);                                            \
        const size_t K = *(rhs_strides + rhs_ndim + 3);                                            \
        const size_t N = *(rhs_strides + rhs_ndim + 4);                                            \
                                                                                                   \
        for (size_t idx = 0; idx < num_els; idx++) {                                               \
            /* Calculate output position: batch_idx, i, j */                                       \
            size_t mn = idx % (M * N);                                                             \
            size_t batch_idx = idx / (M * N);                                                      \
            size_t i = mn / N;                                                                     \
            size_t j = mn % N;                                                                     \
                                                                                                   \
            /* Compute batch indices from flat batch_idx */                                        \
            size_t batch_indices[16];                                                              \
            size_t temp = batch_idx;                                                               \
            for (int d = (int)batch_ndim - 1; d >= 0; d--) {                                       \
                batch_indices[d] = temp % batch_shape[d];                                          \
                temp /= batch_shape[d];                                                            \
            }                                                                                      \
                                                                                                   \
            /* Map batch indices to lhs indices (with broadcasting) */                             \
            size_t lhs_batch_ndim = lhs_ndim - 2;                                                  \
            size_t lhs_batch_indices[16];                                                          \
            for (size_t d = 0; d < lhs_batch_ndim; d++) {                                          \
                size_t batch_dim_idx = batch_ndim - lhs_batch_ndim + d;                            \
                lhs_batch_indices[d] = (lhs_shape[d] == 1) ? 0 : batch_indices[batch_dim_idx];     \
            }                                                                                      \
                                                                                                   \
            /* Map batch indices to rhs indices (with broadcasting) */                             \
            size_t rhs_batch_ndim = rhs_ndim - 2;                                                  \
            size_t rhs_batch_indices[16];                                                          \
            for (size_t d = 0; d < rhs_batch_ndim; d++) {                                          \
                size_t batch_dim_idx = batch_ndim - rhs_batch_ndim + d;                            \
                rhs_batch_indices[d] = (rhs_shape[d] == 1) ? 0 : batch_indices[batch_dim_idx];     \
            }                                                                                      \
                                                                                                   \
            /* Compute matrix multiplication for this output element */                            \
            TYPE sum = 0;                                                                          \
            for (size_t k = 0; k < K; k++) {                                                       \
                /* Calculate lhs index: batch_indices + [i, k] */                                  \
                size_t lhs_idx = lhs_offset;                                                       \
                for (size_t d = 0; d < lhs_batch_ndim; d++) {                                      \
                    lhs_idx += lhs_batch_indices[d] * lhs_strides[d];                              \
                }                                                                                  \
                lhs_idx += i * lhs_strides[lhs_ndim - 2];                                          \
                lhs_idx += k * lhs_strides[lhs_ndim - 1];                                          \
                                                                                                   \
                /* Calculate rhs index: batch_indices + [k, j] */                                  \
                size_t rhs_idx = rhs_offset;                                                       \
                for (size_t d = 0; d < rhs_batch_ndim; d++) {                                      \
                    rhs_idx += rhs_batch_indices[d] * rhs_strides[d];                              \
                }                                                                                  \
                rhs_idx += k * rhs_strides[rhs_ndim - 2];                                          \
                rhs_idx += j * rhs_strides[rhs_ndim - 1];                                          \
                                                                                                   \
                sum += lhs[lhs_idx] * rhs[rhs_idx];                                                \
            }                                                                                      \
                                                                                                   \
            output[idx] = sum;                                                                     \
        }                                                                                          \
    }

// Define matmul for all types
MATMUL_OP(f8e4m3_t, f8e4m3)
MATMUL_OP(f8e5m2_t, f8e5m2)
MATMUL_OP(bf16_t, bf16)
MATMUL_OP(f16_t, f16)
MATMUL_OP(f32_t, f32)
MATMUL_OP(f64_t, f64)
MATMUL_OP(int8_t, i8)
MATMUL_OP(int16_t, i16)
MATMUL_OP(int32_t, i32)
MATMUL_OP(int64_t, i64)
MATMUL_OP(uint8_t, u8)
MATMUL_OP(uint16_t, u16)
MATMUL_OP(uint32_t, u32)
MATMUL_OP(uint64_t, u64)

// ============================================================================
// 2D MATRIX MULTIPLICATION (DOT)
// ============================================================================
//
// Performs simple 2D matrix multiplication without batching.
// Computes C[i, j] = sum_k A[i, k] * B[k, j]
//
// Metadata layout:
// - metadata[0]: M (number of rows in lhs)
// - metadata[1]: K (number of cols in lhs / rows in rhs)
// - metadata[2]: N (number of cols in rhs)
// - metadata[3]: lhs_stride_m (stride for lhs rows)
// - metadata[4]: lhs_stride_k (stride for lhs cols)
// - metadata[5]: rhs_stride_k (stride for rhs rows)
// - metadata[6]: rhs_stride_n (stride for rhs cols)
// - metadata[7]: lhs_offset (starting offset in lhs)
// - metadata[8]: rhs_offset (starting offset in rhs)
//
// Algorithm:
// Simple triple-nested loop over (row, col, k) for matrix multiplication.
// No tiling or cache optimization as CPU doesn't benefit from shared memory
// tiling like GPUs do.

/// Macro to implement 2D matrix multiplication
///
/// @param TYPE C type for the operation
/// @param TYPE_SUFFIX Suffix for function naming
#define DOT_OP(TYPE, TYPE_SUFFIX)                                                                  \
    void dot_##TYPE_SUFFIX(const void *lhs_ptr, const void *rhs_ptr, void *output_ptr,             \
                           const size_t *metadata) {                                               \
        const TYPE *lhs = (const TYPE *)lhs_ptr;                                                   \
        const TYPE *rhs = (const TYPE *)rhs_ptr;                                                   \
        TYPE *output = (TYPE *)output_ptr;                                                         \
                                                                                                   \
        const size_t M = metadata[0];                                                              \
        const size_t K = metadata[1];                                                              \
        const size_t N = metadata[2];                                                              \
        const size_t lhs_stride_m = metadata[3];                                                   \
        const size_t lhs_stride_k = metadata[4];                                                   \
        const size_t rhs_stride_k = metadata[5];                                                   \
        const size_t rhs_stride_n = metadata[6];                                                   \
        const size_t lhs_offset = metadata[7];                                                     \
        const size_t rhs_offset = metadata[8];                                                     \
                                                                                                   \
        /* Simple matrix multiplication without tiling (CPU doesn't benefit from shared memory) */ \
        for (size_t row = 0; row < M; row++) {                                                     \
            for (size_t col = 0; col < N; col++) {                                                 \
                TYPE sum = 0;                                                                      \
                for (size_t k = 0; k < K; k++) {                                                   \
                    size_t lhs_idx = lhs_offset + row * lhs_stride_m + k * lhs_stride_k;           \
                    size_t rhs_idx = rhs_offset + k * rhs_stride_k + col * rhs_stride_n;           \
                    sum += lhs[lhs_idx] * rhs[rhs_idx];                                            \
                }                                                                                  \
                output[row * N + col] = sum;                                                       \
            }                                                                                      \
        }                                                                                          \
    }

// Define dot for all types
DOT_OP(f8e4m3_t, f8e4m3)
DOT_OP(f8e5m2_t, f8e5m2)
DOT_OP(bf16_t, bf16)
DOT_OP(f16_t, f16)
DOT_OP(f32_t, f32)
DOT_OP(f64_t, f64)
DOT_OP(int8_t, i8)
DOT_OP(int16_t, i16)
DOT_OP(int32_t, i32)
DOT_OP(int64_t, i64)
DOT_OP(uint8_t, u8)
DOT_OP(uint16_t, u16)
DOT_OP(uint32_t, u32)
DOT_OP(uint64_t, u64)