yum-mirror/glslang

git clone https://git.yummers.dev/yum-mirror/glslang

Einar HovAdd support for bfloat16 and float8 tensors5ed4003a

main
branchesmainall refs →
3.9 KiB135 linesraw
1#version 460 core
2
3#extension GL_ARM_tensors : enable
4#extension GL_EXT_shader_explicit_arithmetic_types : enable
5#extension GL_EXT_bfloat16 : enable
6#extension GL_EXT_float_e5m2 : enable
7#extension GL_EXT_float_e4m3 : enable
8#extension GL_ARM_tensors_bfloat16 : enable
9#extension GL_ARM_tensors_float_e5m2 : enable
10#extension GL_ARM_tensors_float_e4m3 : enable
11
12// Int types
13uniform tensorARM<int8_t, 1> it8;
14uniform tensorARM<int16_t, 1> it16;
15uniform tensorARM<int, 1> it32;
16uniform tensorARM<int64_t, 1> it64;
17
18// Uint types
19uniform tensorARM<uint8_t, 1> ut8;
20uniform tensorARM<uint16_t, 1> ut16;
21uniform tensorARM<uint, 1> ut32;
22uniform tensorARM<uint64_t, 1> ut64;
23
24// Float types
25uniform tensorARM<float16_t, 1> ft16;
26uniform tensorARM<float, 1> ft32;
27uniform tensorARM<float64_t, 1> ft64;
28
29// bfloat16, e5m2 and e4m3 types
30uniform tensorARM<bfloat16_t, 1> bft16;
31uniform tensorARM<floate5m2_t, 1> fte5m2;
32uniform tensorARM<floate4m3_t, 1> fte4m3;
33
34void main() {
35    uint coords[] = {0};
36
37    // Int types
38    int8_t iw8 = int8_t(1);
39    int16_t iw16 = 1s;
40    int iw = 1;
41    int64_t iw64 = 1l;
42    tensorWriteARM(it8, coords, iw8);
43    tensorWriteARM(it16, coords, iw16);
44    tensorWriteARM(it32, coords, iw);
45    tensorWriteARM(it64, coords, iw64);
46    int iwv[4] = int[](1, 1, 1, 1);
47    tensorWriteARM(it32, coords, iwv);
48
49    int8_t ir8;
50    int16_t ir16;
51    int ir;
52    int64_t ir64;
53    tensorReadARM(it8, coords, ir8);
54    tensorReadARM(it16, coords, ir16);
55    tensorReadARM(it32, coords, ir);
56    tensorReadARM(it64, coords, ir64);
57    int irv[4];
58    tensorReadARM(it32, coords, irv);
59
60    // Uint types
61    uint8_t uw8 = uint8_t(1);
62    uint16_t uw16 = 1s;
63    uint uw = 1;
64    uint64_t uw64 = 1l;
65    tensorWriteARM(ut8, coords, uw8);
66    tensorWriteARM(ut16, coords, uw16);
67    tensorWriteARM(ut32, coords, uw);
68    tensorWriteARM(ut64, coords, uw64);
69    uint uwv[4] = uint[](1, 1, 1, 1);
70    tensorWriteARM(ut32, coords, uwv);
71
72    uint8_t ur8;
73    uint16_t ur16;
74    uint ur;
75    uint64_t ur64;
76    tensorReadARM(ut8, coords, ur8);
77    tensorReadARM(ut16, coords, ur16);
78    tensorReadARM(ut32, coords, ur);
79    tensorReadARM(ut64, coords, ur64);
80    uint urv[4];
81    tensorReadARM(ut32, coords, urv);
82
83    // Float types
84    float fw = 1.0;
85    float16_t fw16 = 1.0hf;
86    float64_t fw64 = 1.0;
87    tensorWriteARM(ft16, coords, fw16);
88    tensorWriteARM(ft32, coords, fw);
89    tensorWriteARM(ft64, coords, fw64);
90    float fwv[4] = float[](1.0, 1.0, 1.0, 1.0);
91    tensorWriteARM(ft32, coords, fwv);
92
93    float fr;
94    float16_t fr16;
95    float64_t fr64;
96    tensorReadARM(ft16, coords, fr16);
97    tensorReadARM(ft32, coords, fr);
98    tensorReadARM(ft64, coords, fr64);
99    float frv[4];
100    tensorReadARM(ft32, coords, frv);
101
102    // bfloat16, e5m2 and e4m3 types
103    bfloat16_t bf16w = bfloat16_t(1.0);
104    floate5m2_t fe5m2w = floate5m2_t(1.0);
105    floate4m3_t fe4m3w = floate4m3_t(1.0);
106    tensorWriteARM(bft16, coords, bf16w);
107    tensorWriteARM(fte5m2, coords, fe5m2w);
108    tensorWriteARM(fte4m3, coords, fe4m3w);
109
110    const bfloat16_t bf16one = bfloat16_t(1.0);
111    bfloat16_t bf16wv[4] = bfloat16_t[](bf16one, bf16one, bf16one, bf16one);
112    tensorWriteARM(bft16, coords, bf16wv);
113
114    const floate5m2_t fe5m2one = floate5m2_t(1.0);
115    floate5m2_t fe5m2wv[4] = floate5m2_t[](fe5m2one, fe5m2one, fe5m2one, fe5m2one);
116    tensorWriteARM(fte5m2, coords, fe5m2wv);
117
118    floate4m3_t fe4m3one = floate4m3_t(1.0);
119    floate4m3_t fe4m3wv[4] = floate4m3_t[](fe4m3one, fe4m3one, fe4m3one, fe4m3one);
120    tensorWriteARM(fte4m3, coords, fe4m3wv);
121
122    bfloat16_t bf16r;
123    floate5m2_t fe5m2r;
124    floate4m3_t fe4m3r;
125    tensorReadARM(bft16, coords, bf16r);
126    tensorReadARM(fte5m2, coords, fe5m2r);
127    tensorReadARM(fte4m3, coords, fe4m3r);
128
129    bfloat16_t bf16rv[4];
130    floate5m2_t fe5m2rv[4];
131    floate4m3_t fe4m3rv[4];
132    tensorReadARM(bft16, coords, bf16rv);
133    tensorReadARM(fte5m2, coords, fe5m2rv);
134    tensorReadARM(fte4m3, coords, fe4m3rv);
135}