hpt-macros 0.1.2

An internal library for generating helper functions for hpt
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

use proc_macro::TokenStream;
use quote::quote;
use syn::parse_macro_input;
use syn::Ident;

use crate::NUM_REG;

pub fn __gen_fast_reduce_simd_helper(stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(stream as Ident);

    #[cfg(target_feature = "avx2")]
    let num_registers = 16;
    #[cfg(all(
        any(target_feature = "sse", target_arch = "arm"),
        not(target_feature = "avx2")
    ))]
    let num_registers = 8;
    #[cfg(any(target_feature = "avx512f", target_arch = "aarch64"))]
    let num_registers = 32;

    let mut body = proc_macro2::TokenStream::new();
    for i in 0..num_registers as isize {
        let i = i + 1;
        let elements = 1..=i;
        let arr = quote! {
            [ #(#elements),* ]
        };
        let i_u32 = i as u32;
        body.extend(
            quote! {
            #i_u32 => {
                gen_kernel!(1, #i, inp_ptr, res_ptr, vec_size, outer_loop_size, vec_preop, vec_cumulate, inp_strides, inp_shape, prg, vec_post, #arr);
            }
        }
        );
    }
    let ret = quote! {
        match #input {
            #body
            0 => {}
            _ => unreachable!()
        }
    };
    ret.into()
}

pub fn __gen_fast_layernorm_simd_helper(stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(stream as Ident);

    #[cfg(target_feature = "avx2")]
    let num_registers = 16;
    #[cfg(all(
        any(target_feature = "sse", target_arch = "arm"),
        not(target_feature = "avx2")
    ))]
    let num_registers = 8;
    #[cfg(any(target_feature = "avx512f", target_arch = "aarch64"))]
    let num_registers = 32;

    let mut body = proc_macro2::TokenStream::new();
    for i in 0..num_registers as isize {
        let i = i + 1;
        let elements = 1..=i;
        let arr = quote! {
            [ #(#elements),* ]
        };
        let i_u32 = i as u32;
        body.extend(
            quote! {
            #i_u32 => {
                gen_kernel!(1, #i, inp_ptr, res_ptr, vec_size, outer_loop_size, inp_strides, inp_shape, prg, #arr);
            }
        }
        );
    }
    let ret = quote! {
        match #input {
            #body
            0 => {}
            _ => unreachable!()
        }
    };
    ret.into()
}

pub fn __gen_reduce_dim_not_include_simd_helper(stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(stream as Ident);

    let mut body = proc_macro2::TokenStream::new();
    for i in 0..NUM_REG as isize {
        let i = i + 1;
        let elements = 1..=i;
        let arr = quote! {
            [ #(#elements),* ]
        };
        let i_u32 = i as u32;
        body.extend(quote! {
            #i_u32 => {
                gen_kernel3!(
                    1,
                    #i,
                    outer_loop_size,
                    inp_ptr,
                    res_ptr,
                    <O as TypeCommon>::Vec::SIZE as isize,
                    intermediate_size,
                    vec_preop,
                    vec_cumulate,
                    inp_strides,
                    inp_shape,
                    prg1,
                    prg2,
                    shape_len,
                    inner_loop_size,
                    vec_post,
                    #arr
                );
            }
        });
    }
    let ret = quote! {
        match #input {
            #body,
            0 => {}
            _ => unreachable!()
        }
    };
    ret.into()
}