ds4/ds4_metal_sources.inc at nm-shared · NimbleMarkets/ds4 · GitHub

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
/* Metal kernels embedded at compile time with C23 #embed. */

typedef struct {
    const char *path;
    const char *source;
} ds4_metal_embedded_source;

static const char ds4_metal_flash_attn_source[] = {
#embed "metal/flash_attn.metal"
    , 0
};
static const char ds4_metal_dense_source[] = {
#embed "metal/dense.metal"
    , 0
};
static const char ds4_metal_moe_source[] = {
#embed "metal/moe.metal"
    , 0
};
static const char ds4_metal_dsv4_hc_source[] = {
#embed "metal/dsv4_hc.metal"
    , 0
};
static const char ds4_metal_unary_source[] = {
#embed "metal/unary.metal"
    , 0
};
static const char ds4_metal_dsv4_kv_source[] = {
#embed "metal/dsv4_kv.metal"
    , 0
};
static const char ds4_metal_dsv4_rope_source[] = {
#embed "metal/dsv4_rope.metal"
    , 0
};
static const char ds4_metal_dsv4_misc_source[] = {
#embed "metal/dsv4_misc.metal"
    , 0
};
static const char ds4_metal_argsort_source[] = {
#embed "metal/argsort.metal"
    , 0
};
static const char ds4_metal_cpy_source[] = {
#embed "metal/cpy.metal"
    , 0
};
static const char ds4_metal_concat_source[] = {
#embed "metal/concat.metal"
    , 0
};
static const char ds4_metal_get_rows_source[] = {
#embed "metal/get_rows.metal"
    , 0
};
static const char ds4_metal_sum_rows_source[] = {
#embed "metal/sum_rows.metal"
    , 0
};
static const char ds4_metal_softmax_source[] = {
#embed "metal/softmax.metal"
    , 0
};
static const char ds4_metal_repeat_source[] = {
#embed "metal/repeat.metal"
    , 0
};
static const char ds4_metal_glu_source[] = {
#embed "metal/glu.metal"
    , 0
};
static const char ds4_metal_norm_source[] = {
#embed "metal/norm.metal"
    , 0
};
static const char ds4_metal_bin_source[] = {
#embed "metal/bin.metal"
    , 0
};
static const char ds4_metal_set_rows_source[] = {
#embed "metal/set_rows.metal"
    , 0
};

static const ds4_metal_embedded_source ds4_metal_embedded_sources[] = {
    {"metal/flash_attn.metal", ds4_metal_flash_attn_source},
    {"metal/dense.metal",      ds4_metal_dense_source},
    {"metal/moe.metal",        ds4_metal_moe_source},
    {"metal/dsv4_hc.metal",    ds4_metal_dsv4_hc_source},
    {"metal/unary.metal",      ds4_metal_unary_source},
    {"metal/dsv4_kv.metal",    ds4_metal_dsv4_kv_source},
    {"metal/dsv4_rope.metal",  ds4_metal_dsv4_rope_source},
    {"metal/dsv4_misc.metal",  ds4_metal_dsv4_misc_source},
    {"metal/argsort.metal",    ds4_metal_argsort_source},
    {"metal/cpy.metal",        ds4_metal_cpy_source},
    {"metal/concat.metal",     ds4_metal_concat_source},
    {"metal/get_rows.metal",   ds4_metal_get_rows_source},
    {"metal/sum_rows.metal",   ds4_metal_sum_rows_source},
    {"metal/softmax.metal",    ds4_metal_softmax_source},
    {"metal/repeat.metal",     ds4_metal_repeat_source},
    {"metal/glu.metal",        ds4_metal_glu_source},
    {"metal/norm.metal",       ds4_metal_norm_source},
    {"metal/bin.metal",        ds4_metal_bin_source},
    {"metal/set_rows.metal",   ds4_metal_set_rows_source},
};

static const size_t ds4_metal_embedded_source_count =
    sizeof(ds4_metal_embedded_sources) / sizeof(ds4_metal_embedded_sources[0]);