hodu_cpu_kernels 0.2.4

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

#include "ops_concat_split.h"
#include "types.h"
#include <stdbool.h>

// ============================================================================
// CONCATENATION OPERATIONS
// ============================================================================
//
// Metadata layout:
// - metadata[0]: num_els (total number of elements in output)
// - metadata[1]: num_dims (number of dimensions)
// - metadata[2..2+num_dims]: output_shape (shape of concatenated output)
// - metadata[2+num_dims]: concat_dim (dimension along which to concatenate)
// - metadata[2+num_dims+1]: num_inputs (number of input tensors)
// - metadata[2+num_dims+2..]: input_shapes (flattened: num_inputs * num_dims elements)
// - ...: input_strides (flattened: num_inputs * num_dims elements)
// - ...: input_offsets (num_inputs elements)
// - ...: input_buffer_offsets (num_inputs elements)
//
// Algorithm:
// For each output element, determine which input tensor it comes from by
// checking the index along concat_dim, then map to the corresponding input
// element using the input's shape and strides.

/// Macro to implement a concatenation operation
///
/// @param TYPENAME C type for the operation
/// @param FN_NAME Function name
#define CONCAT_OP(TYPENAME, FN_NAME)                                                               \
    void FN_NAME(const void *input_ptr, void *output_ptr, const size_t *metadata) {                \
        const TYPENAME *input = (const TYPENAME *)input_ptr;                                       \
        TYPENAME *output = (TYPENAME *)output_ptr;                                                 \
                                                                                                   \
        const size_t num_els = metadata[0];                                                        \
        const size_t num_dims = metadata[1];                                                       \
        const size_t *output_shape = metadata + 2;                                                 \
        const size_t concat_dim = metadata[2 + num_dims];                                          \
        const size_t num_inputs = metadata[2 + num_dims + 1];                                      \
        const size_t *input_shapes = metadata + 2 + num_dims + 2;                                  \
        const size_t *input_strides = input_shapes + num_inputs * num_dims;                        \
        const size_t *input_offsets = input_strides + num_inputs * num_dims;                       \
        const size_t *input_buffer_offsets = input_offsets + num_inputs;                           \
                                                                                                   \
        /* Process each output element */                                                          \
        for (size_t id = 0; id < num_els; id++) {                                                  \
            /* Calculate output indices from flat id */                                            \
            size_t output_indices[16];                                                             \
            size_t temp = id;                                                                      \
            for (int d = (int)num_dims - 1; d >= 0; d--) {                                         \
                output_indices[d] = temp % output_shape[d];                                        \
                temp /= output_shape[d];                                                           \
            }                                                                                      \
                                                                                                   \
            /* Determine which input tensor to read from */                                        \
            size_t dim_index = output_indices[concat_dim];                                         \
            size_t tensor_idx = 0;                                                                 \
            size_t cumulative_size = 0;                                                            \
                                                                                                   \
            for (size_t i = 0; i < num_inputs; i++) {                                              \
                size_t dim_size = input_shapes[i * num_dims + concat_dim];                         \
                if (dim_index < cumulative_size + dim_size) {                                      \
                    tensor_idx = i;                                                                \
                    break;                                                                         \
                }                                                                                  \
                cumulative_size += dim_size;                                                       \
            }                                                                                      \
                                                                                                   \
            /* Calculate input index */                                                            \
            const size_t *input_stride = input_strides + tensor_idx * num_dims;                    \
            size_t input_offset = input_offsets[tensor_idx];                                       \
            size_t buffer_offset = input_buffer_offsets[tensor_idx];                               \
                                                                                                   \
            size_t flat_index = input_offset;                                                      \
            for (size_t i = 0; i < num_dims; i++) {                                                \
                size_t idx =                                                                       \
                    (i == concat_dim) ? (dim_index - cumulative_size) : output_indices[i];         \
                flat_index += idx * input_stride[i];                                               \
            }                                                                                      \
                                                                                                   \
            output[id] = input[buffer_offset + flat_index];                                        \
        }                                                                                          \
    }

// Define concat operations for all types
CONCAT_OP(bool, concat_bool)
CONCAT_OP(f8e4m3_t, concat_f8e4m3)
CONCAT_OP(f8e5m2_t, concat_f8e5m2)
CONCAT_OP(bf16_t, concat_bf16)
CONCAT_OP(f16_t, concat_f16)
CONCAT_OP(float, concat_f32)
CONCAT_OP(double, concat_f64)
CONCAT_OP(int8_t, concat_i8)
CONCAT_OP(int16_t, concat_i16)
CONCAT_OP(int32_t, concat_i32)
CONCAT_OP(int64_t, concat_i64)
CONCAT_OP(uint8_t, concat_u8)
CONCAT_OP(uint16_t, concat_u16)
CONCAT_OP(uint32_t, concat_u32)
CONCAT_OP(uint64_t, concat_u64)

// ============================================================================
// SPLIT OPERATIONS
// ============================================================================
//
// Metadata layout:
// - metadata[0]: num_els (total number of elements in output)
// - metadata[1]: num_dims (number of dimensions)
// - metadata[2..2+num_dims]: input_shape (shape of input tensor)
// - metadata[2+num_dims..2+2*num_dims]: input_strides (strides of input tensor)
// - metadata[2+2*num_dims]: input_offset (starting offset in input tensor)
// - metadata[2+2*num_dims+1]: split_dim (dimension along which to split)
// - metadata[2+2*num_dims+2]: output_size_on_dim (size of output along split dimension)
// - metadata[2+2*num_dims+3]: split_offset (offset along split dimension where output starts)
//
// Algorithm:
// For each output element, compute its indices in the output space, then
// map to the input space by adding split_offset to the index along split_dim.

/// Macro to implement a split operation
///
/// @param TYPENAME C type for the operation
/// @param FN_NAME Function name
#define SPLIT_OP(TYPENAME, FN_NAME)                                                                \
    void FN_NAME(const void *input_ptr, void *output_ptr, const size_t *metadata) {                \
        const TYPENAME *input = (const TYPENAME *)input_ptr;                                       \
        TYPENAME *output = (TYPENAME *)output_ptr;                                                 \
                                                                                                   \
        const size_t num_els = metadata[0];                                                        \
        const size_t num_dims = metadata[1];                                                       \
        const size_t *input_shape = metadata + 2;                                                  \
        const size_t *input_strides = metadata + 2 + num_dims;                                     \
        const size_t input_offset = metadata[2 + 2 * num_dims];                                    \
        const size_t split_dim = metadata[2 + 2 * num_dims + 1];                                   \
        const size_t output_size_on_dim = metadata[2 + 2 * num_dims + 2];                          \
        const size_t split_offset = metadata[2 + 2 * num_dims + 3];                                \
                                                                                                   \
        /* Calculate output shape */                                                               \
        size_t output_shape[16];                                                                   \
        for (size_t i = 0; i < num_dims; i++) {                                                    \
            output_shape[i] = (i == split_dim) ? output_size_on_dim : input_shape[i];              \
        }                                                                                          \
                                                                                                   \
        /* Process each output element */                                                          \
        for (size_t id = 0; id < num_els; id++) {                                                  \
            /* Calculate output indices from flat id */                                            \
            size_t output_indices[16];                                                             \
            size_t temp = id;                                                                      \
            for (int d = (int)num_dims - 1; d >= 0; d--) {                                         \
                output_indices[d] = temp % output_shape[d];                                        \
                temp /= output_shape[d];                                                           \
            }                                                                                      \
                                                                                                   \
            /* Calculate input flat index */                                                       \
            size_t flat_index = input_offset;                                                      \
            for (size_t i = 0; i < num_dims; i++) {                                                \
                size_t idx =                                                                       \
                    (i == split_dim) ? (output_indices[i] + split_offset) : output_indices[i];     \
                flat_index += idx * input_strides[i];                                              \
            }                                                                                      \
                                                                                                   \
            output[id] = input[flat_index];                                                        \
        }                                                                                          \
    }

// Define split operations for all types
SPLIT_OP(bool, split_bool)
SPLIT_OP(f8e4m3_t, split_f8e4m3)
SPLIT_OP(f8e5m2_t, split_f8e5m2)
SPLIT_OP(bf16_t, split_bf16)
SPLIT_OP(f16_t, split_f16)
SPLIT_OP(float, split_f32)
SPLIT_OP(double, split_f64)
SPLIT_OP(int8_t, split_i8)
SPLIT_OP(int16_t, split_i16)
SPLIT_OP(int32_t, split_i32)
SPLIT_OP(int64_t, split_i64)
SPLIT_OP(uint8_t, split_u8)
SPLIT_OP(uint16_t, split_u16)
SPLIT_OP(uint32_t, split_u32)
SPLIT_OP(uint64_t, split_u64)