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 →
3.7 KiB93 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_NV_cooperative_matrix2 : enable
6#extension GL_NV_cooperative_matrix_decode_vector : enable
7#extension GL_EXT_buffer_reference : enable
8
9layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
10
11buffer BufType {
12   float16_t x[];
13} Buf;
14
15buffer BufType8 {
16   uint8_t x[];
17} Buf8;
18
19layout(buffer_reference, std430, buffer_reference_align = 4) buffer dwordBuf {
20    uint32_t d;
21};
22
23float16_t decodeF16Scalar(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
24{
25    return float16_t(b.d);
26}
27
28f16vec2 decodeF16x2(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
29{
30    return unpackFloat2x16(b.d);
31}
32
33uint8_t decodeU8Scalar(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
34{
35    return uint8_t(b.d);
36}
37
38u8vec4 decodeU8x4(const in dwordBuf b, const in uint32_t blockCoords[2], const in uint32_t coordInBlock[2])
39{
40    return u8vec4(uint8_t(b.d), uint8_t(b.d >> 8), uint8_t(b.d >> 16), uint8_t(b.d >> 24));
41}
42
43float16_t decodeF16Scalar3d(const in dwordBuf b, const in uint32_t blockCoords[3], const in uint32_t coordInBlock[3])
44{
45    return float16_t(b.d);
46}
47
48f16vec2 decodeF16x2_3d(const in dwordBuf b, const in uint32_t blockCoords[3], const in uint32_t coordInBlock[3])
49{
50    return unpackFloat2x16(b.d);
51}
52
53void main()
54{
55    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB>           B;
56    coopmat<uint8_t,   gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB>           Bu8;
57    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Acc;
58
59    tensorLayoutNV<2> t = createTensorLayoutNV(2);
60
61    // No restrictions on Use or TensorView with v2: every combination is a
62    // valid load and the implementation chooses per call site whether to
63    // invoke the vector function. Below: gl_MatrixUseB without a view, with
64    // a transposing view, and with an identity view; gl_MatrixUseAccumulator;
65    // and a 3D layout/view variant.
66
67    // gl_MatrixUseB, no view.
68    coopMatLoadTensorNV(B,   Buf.x,  0, t, decodeF16Scalar, decodeF16x2);
69
70    // gl_MatrixUseB, transposing view (swap of innermost two).
71    tensorViewNV<2, false, 1, 0> vT2 = createTensorViewNV(2, false, 1, 0);
72    coopMatLoadTensorNV(B,   Buf.x,  0, t, vT2, decodeF16Scalar, decodeF16x2);
73    coopMatLoadTensorNV(Bu8, Buf8.x, 0, t, vT2, decodeU8Scalar,  decodeU8x4);
74
75    // gl_MatrixUseB, identity view -- previously rejected, now accepted.
76    tensorViewNV<2, false, 0, 1> vI2 = createTensorViewNV(2, false, 0, 1);
77    coopMatLoadTensorNV(B,   Buf.x,  0, t, vI2, decodeF16Scalar, decodeF16x2);
78
79    // gl_MatrixUseAccumulator -- previously rejected for vector decode, now
80    // accepted.
81    coopMatLoadTensorNV(Acc, Buf.x,  0, t, decodeF16Scalar, decodeF16x2);
82
83    // 3D layout + identity-outer/swap-innermost view (still a valid shape).
84    tensorLayoutNV<3>            t3  = createTensorLayoutNV(3);
85    tensorViewNV<3, false, 0, 2, 1> vT3 = createTensorViewNV(3, false, 0, 2, 1);
86    coopMatLoadTensorNV(B, Buf.x, 0, t3, vT3, decodeF16Scalar3d, decodeF16x2_3d);
87
88    // 3D layout + arbitrary-permutation view -- previously rejected, now
89    // accepted (the implementation simply will not invoke the vector
90    // function at sites where it would not be profitable).
91    tensorViewNV<3, false, 1, 2, 0> vW3 = createTensorViewNV(3, false, 1, 2, 0);
92    coopMatLoadTensorNV(B, Buf.x, 0, t3, vW3, decodeF16Scalar3d, decodeF16x2_3d);
93}