yum-mirror/glslang

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

Jeff BolzImplement GL_NV_cooperative_matrix_decode_vector (#4275)3ceed991

main
branchesmainall refs →
2.9 KiB84 linesraw
1#version 450 core
2#extension GL_KHR_memory_scope_semantics : enable
3#extension GL_KHR_cooperative_matrix : enable
4#extension GL_EXT_shader_explicit_arithmetic_types : enable
5#extension GL_EXT_long_vector : enable
6#extension GL_NV_cooperative_matrix2 : enable
7#extension GL_NV_cooperative_matrix_decode_vector : enable
8#extension GL_EXT_buffer_reference : enable
9
10layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
11
12buffer BufType {
13   float16_t x[];
14} Buf;
15
16buffer BufType8 {
17   uint8_t x[];
18} Buf8;
19
20layout(buffer_reference, std430, buffer_reference_align = 4) buffer dwordBuf {
21    uint32_t d;
22};
23
24float16_t decodeF16Scalar(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
25{
26    return float16_t(b.d);
27}
28
29f16vec2 decodeF16x2(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
30{
31    return unpackFloat2x16(b.d);
32}
33
34f16vec4 decodeF16x4(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
35{
36    return f16vec4(unpackFloat2x16(b.d), unpackFloat2x16(b.d));
37}
38
39uint8_t decodeU8Scalar(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
40{
41    return uint8_t(b.d);
42}
43
44u8vec4 decodeU8x4(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
45{
46    return u8vec4(uint8_t(b.d), uint8_t(b.d >> 8), uint8_t(b.d >> 16), uint8_t(b.d >> 24));
47}
48
49vector<uint8_t, 8> decodeU8x8(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
50{
51    return vector<uint8_t, 8>(uint8_t(b.d),       uint8_t(b.d >> 8),
52                              uint8_t(b.d >> 16), uint8_t(b.d >> 24),
53                              uint8_t(b.d),       uint8_t(b.d >> 8),
54                              uint8_t(b.d >> 16), uint8_t(b.d >> 24));
55}
56
57void main()
58{
59    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA> A;
60    coopmat<uint8_t,   gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA> Au8;
61
62    tensorLayoutNV<2> t  = createTensorLayoutNV(2);
63    t = setTensorLayoutBlockSizeNV(t, 4, 8);
64    t = setTensorLayoutDimensionNV(t, 256, 512);
65
66    // Scalar-only decode (the SPV_NV_cooperative_matrix2 baseline). No
67    // capability or extension from this extension is required when the
68    // vector form is not used.
69    coopMatLoadTensorNV(A,   Buf.x,  0, t, decodeF16Scalar);
70    coopMatLoadTensorNV(Au8, Buf8.x, 0, t, decodeU8Scalar);
71
72    // Scalar + vector decode pair: the implementation may invoke either
73    // function per call site. V == 2 fp16.
74    coopMatLoadTensorNV(A,   Buf.x,  0, t, decodeF16Scalar, decodeF16x2);
75
76    // V == 4 fp16.
77    coopMatLoadTensorNV(A,   Buf.x,  0, t, decodeF16Scalar, decodeF16x4);
78
79    // V == 4 u8.
80    coopMatLoadTensorNV(Au8, Buf8.x, 0, t, decodeU8Scalar, decodeU8x4);
81
82    // V == 8 u8.
83    coopMatLoadTensorNV(Au8, Buf8.x, 0, t, decodeU8Scalar, decodeU8x8);
84}