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
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
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306

//! Unary tensor operations
//!
//! This module provides element-wise unary operations on tensors, including:
//! - Basic arithmetic: neg, abs, sign, square, sqrt, recip
//! - Activation functions: relu, sigmoid, tanh, gelu, softplus, silu, mish
//! - Trigonometric: sin, cos, tan
//! - Exponential/logarithmic: exp, exp2, exp10, ln, log2, log10
//! - Logical: logical_not
//! - Scalar operations: arithmetic and comparison operations with a scalar value
//!
//! All operations support strided tensor access and multiple data types.

use crate::{error::Result, kernels::macros::ops};
use core::ffi::c_void;

// Define all unary operations using the macro
ops!(
    neg,
    abs,
    sign,
    square,
    sqrt,
    recip,
    relu,
    sigmoid,
    tanh,
    gelu,
    softplus,
    silu,
    mish,
    sin,
    cos,
    tan,
    exp,
    exp2,
    exp10,
    ln,
    log2,
    log10,
    logical_not,
    add_scalar,
    sub_scalar,
    mul_scalar,
    div_scalar,
    pow_scalar,
    maximum_scalar,
    minimum_scalar,
    eq_scalar,
    ne_scalar,
    lt_scalar,
    le_scalar,
    gt_scalar,
    ge_scalar
);

/// Call unary operation by kernel name
///
/// Applies an element-wise unary operation to the input tensor.
///
/// # Metadata layout
/// - metadata[0]: num_els (total number of elements)
/// - metadata[1]: num_dims (number of dimensions)
/// - metadata[2..2+num_dims]: shape
/// - metadata[2+num_dims..2+2*num_dims]: strides
/// - metadata[2+2*num_dims]: offset
///
/// # Safety
/// - `input` must point to valid tensor data of the appropriate type
/// - `output` must point to a valid output buffer with sufficient capacity
/// - Metadata must accurately describe the tensor layout
pub fn call_ops_unary(
    kernel_name: crate::kernels::macros::Kernel,
    input: *const c_void,
    output: *mut c_void,
    metadata: &[usize],
) -> Result<()> {
    unsafe {
        dispatch_unary(kernel_name.0, input, output, metadata.as_ptr());
    }

    Ok(())
}

/// Call unary operation with scalar by kernel name
///
/// Applies an element-wise operation combining each tensor element with a scalar value.
/// Used for operations like add_scalar, mul_scalar, and comparison operations.
///
/// # Metadata layout (same as call_ops_unary)
/// - metadata[0]: num_els (total number of elements)
/// - metadata[1]: num_dims (number of dimensions)
/// - metadata[2..2+num_dims]: shape
/// - metadata[2+num_dims..2+2*num_dims]: strides
/// - metadata[2+2*num_dims]: offset
///
/// # Safety
/// - `input` must point to valid tensor data of the appropriate type
/// - `output` must point to a valid output buffer with sufficient capacity
/// - `scalar_val` must be of the same type as the tensor elements
/// - Metadata must accurately describe the tensor layout
pub fn call_ops_unary_scalar<T>(
    kernel_name: crate::kernels::macros::Kernel,
    input: *const c_void,
    output: *mut c_void,
    metadata: &[usize],
    scalar_val: T,
) -> Result<()> {
    unsafe {
        dispatch_unary_scalar(
            kernel_name.0,
            input,
            output,
            metadata.as_ptr(),
            &scalar_val as *const T as *const c_void,
        );
    }

    Ok(())
}

// Macro to automatically generate extern declarations and dispatch for all operations and types
macro_rules! declare_and_dispatch_unary {
    (
        all_types: [$($all_op:ident),* $(,)?],
        signed_types: [$($signed_op:ident),* $(,)?],
        float_types: [$($float_op:ident),* $(,)?],
        scalar_ops: [$($scalar_op:ident),* $(,)?]
    ) => {
        paste::paste! {
            // Extern C declarations
            extern "C" {
                // Operations supporting all types (including unsigned)
                $(
                    fn [<$all_op _bool>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _f8e4m3>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _f8e5m2>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _bf16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _f16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _f32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _f64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _u8>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _u16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _u32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _u64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _i8>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _i16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _i32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$all_op _i64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                )*

                // Operations supporting signed types only (bool, float, signed int)
                $(
                    fn [<$signed_op _bool>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _f8e4m3>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _f8e5m2>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _bf16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _f16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _f32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _f64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _i8>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _i16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _i32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$signed_op _i64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                )*

                // Operations supporting only float types
                $(
                    fn [<$float_op _bool>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _f8e4m3>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _f8e5m2>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _bf16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _f16>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _f32>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                    fn [<$float_op _f64>](input: *const c_void, output: *mut c_void, metadata: *const usize);
                )*

                // Scalar operations (all types, with scalar parameter)
                $(
                    fn [<$scalar_op _bool>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _f8e4m3>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _f8e5m2>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _bf16>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _f16>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _f32>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _f64>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _u8>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _u16>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _u32>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _u64>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _i8>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _i16>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _i32>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                    fn [<$scalar_op _i64>](input: *const c_void, output: *mut c_void, metadata: *const usize, scalar: *const c_void);
                )*
            }

            // Dispatch function
            unsafe fn dispatch_unary(
                name: &str,
                input: *const c_void,
                output: *mut c_void,
                metadata: *const usize,
            ) {
                match name {
                    // Operations with all types
                    $(
                        concat!(stringify!($all_op), "_bool") => [<$all_op _bool>](input, output, metadata),
                        concat!(stringify!($all_op), "_f8e4m3") => [<$all_op _f8e4m3>](input, output, metadata),
                        concat!(stringify!($all_op), "_f8e5m2") => [<$all_op _f8e5m2>](input, output, metadata),
                        concat!(stringify!($all_op), "_bf16") => [<$all_op _bf16>](input, output, metadata),
                        concat!(stringify!($all_op), "_f16") => [<$all_op _f16>](input, output, metadata),
                        concat!(stringify!($all_op), "_f32") => [<$all_op _f32>](input, output, metadata),
                        concat!(stringify!($all_op), "_f64") => [<$all_op _f64>](input, output, metadata),
                        concat!(stringify!($all_op), "_u8") => [<$all_op _u8>](input, output, metadata),
                        concat!(stringify!($all_op), "_u16") => [<$all_op _u16>](input, output, metadata),
                        concat!(stringify!($all_op), "_u32") => [<$all_op _u32>](input, output, metadata),
                        concat!(stringify!($all_op), "_u64") => [<$all_op _u64>](input, output, metadata),
                        concat!(stringify!($all_op), "_i8") => [<$all_op _i8>](input, output, metadata),
                        concat!(stringify!($all_op), "_i16") => [<$all_op _i16>](input, output, metadata),
                        concat!(stringify!($all_op), "_i32") => [<$all_op _i32>](input, output, metadata),
                        concat!(stringify!($all_op), "_i64") => [<$all_op _i64>](input, output, metadata),
                    )*

                    // Signed-only operations
                    $(
                        concat!(stringify!($signed_op), "_bool") => [<$signed_op _bool>](input, output, metadata),
                        concat!(stringify!($signed_op), "_f8e4m3") => [<$signed_op _f8e4m3>](input, output, metadata),
                        concat!(stringify!($signed_op), "_f8e5m2") => [<$signed_op _f8e5m2>](input, output, metadata),
                        concat!(stringify!($signed_op), "_bf16") => [<$signed_op _bf16>](input, output, metadata),
                        concat!(stringify!($signed_op), "_f16") => [<$signed_op _f16>](input, output, metadata),
                        concat!(stringify!($signed_op), "_f32") => [<$signed_op _f32>](input, output, metadata),
                        concat!(stringify!($signed_op), "_f64") => [<$signed_op _f64>](input, output, metadata),
                        concat!(stringify!($signed_op), "_i8") => [<$signed_op _i8>](input, output, metadata),
                        concat!(stringify!($signed_op), "_i16") => [<$signed_op _i16>](input, output, metadata),
                        concat!(stringify!($signed_op), "_i32") => [<$signed_op _i32>](input, output, metadata),
                        concat!(stringify!($signed_op), "_i64") => [<$signed_op _i64>](input, output, metadata),
                    )*

                    // Float-only operations
                    $(
                        concat!(stringify!($float_op), "_bool") => [<$float_op _bool>](input, output, metadata),
                        concat!(stringify!($float_op), "_f8e4m3") => [<$float_op _f8e4m3>](input, output, metadata),
                        concat!(stringify!($float_op), "_f8e5m2") => [<$float_op _f8e5m2>](input, output, metadata),
                        concat!(stringify!($float_op), "_bf16") => [<$float_op _bf16>](input, output, metadata),
                        concat!(stringify!($float_op), "_f16") => [<$float_op _f16>](input, output, metadata),
                        concat!(stringify!($float_op), "_f32") => [<$float_op _f32>](input, output, metadata),
                        concat!(stringify!($float_op), "_f64") => [<$float_op _f64>](input, output, metadata),
                    )*

                    _ => panic!("Unsupported unary kernel: {}", name),
                }
            }

            // Dispatch function for scalar operations
            unsafe fn dispatch_unary_scalar(
                name: &str,
                input: *const c_void,
                output: *mut c_void,
                metadata: *const usize,
                scalar: *const c_void,
            ) {
                match name {
                    // Scalar operations (all types)
                    $(
                        concat!(stringify!($scalar_op), "_bool") => [<$scalar_op _bool>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_f8e4m3") => [<$scalar_op _f8e4m3>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_f8e5m2") => [<$scalar_op _f8e5m2>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_bf16") => [<$scalar_op _bf16>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_f16") => [<$scalar_op _f16>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_f32") => [<$scalar_op _f32>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_f64") => [<$scalar_op _f64>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_u8") => [<$scalar_op _u8>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_u16") => [<$scalar_op _u16>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_u32") => [<$scalar_op _u32>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_u64") => [<$scalar_op _u64>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_i8") => [<$scalar_op _i8>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_i16") => [<$scalar_op _i16>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_i32") => [<$scalar_op _i32>](input, output, metadata, scalar),
                        concat!(stringify!($scalar_op), "_i64") => [<$scalar_op _i64>](input, output, metadata, scalar),
                    )*
                    _ => panic!("Unsupported unary scalar kernel: {}", name),
                }
            }
        }
    };
}

// Declare all operations
declare_and_dispatch_unary!(
    all_types: [
        abs, sign, square, sqrt, recip, logical_not
    ],
    signed_types: [
        neg
    ],
    float_types: [
        relu, sigmoid, tanh, gelu, softplus, silu, mish,
        sin, cos, tan,
        exp, exp2, exp10, ln, log2, log10
    ],
    scalar_ops: [
        add_scalar, sub_scalar, mul_scalar, div_scalar, pow_scalar,
        maximum_scalar, minimum_scalar,
        eq_scalar, ne_scalar, lt_scalar, le_scalar, gt_scalar, ge_scalar
    ]
);