1#version 450
 2
 3#include "types.glsl"
 4#include "generic_unary_head.glsl"
 5
 6// workgroup does 32x32 tile, but uses 32x8 threads
 7#define TILE_DIM 32
 8layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in;
 9
10shared uint sh[TILE_DIM][TILE_DIM + 1];
11
12void iter(uvec3 wg_id) {
13    const uint tile_col = wg_id.x;
14    const uint tile_row = wg_id.y;
15
16    const uint tid_col = gl_LocalInvocationID.x;
17    const uint tid_row = gl_LocalInvocationID.y;
18
19    const uint i2 = wg_id.z % p.ne12;
20    const uint i3 = wg_id.z / p.ne12;
21    const uint i02 = i2;
22    const uint i03 = i3;
23
24    // The workgroup does TILE_DIM x TILE_DIM, but swaps the LSBs of the
25    // src coords to make memory accesses contiguous, dst has tid.x in i0,
26    // src has tid.x in i01
27
28    [[unroll]] for (uint y = 0; y < 4; ++y) {
29        const uint i00 = tile_col * TILE_DIM + tid_row + 8 * y;
30        const uint i01 = tile_row * TILE_DIM + tid_col;
31        if (i00 < p.ne00 && i01 < p.ne01 && i02 < p.ne02 && i03 < p.ne03) {
32            const uint src_idx = i00 * p.nb00 + i01 * p.nb01 + i02 * p.nb02 + i03 * p.nb03;
33            sh[tid_row + 8 * y][tid_col] = uint(data_a[get_aoffset() + src_idx]);
34        }
35    }
36
37    barrier();
38
39    [[unroll]] for (uint y = 0; y < 4; ++y) {
40        const uint i0 = tile_col * TILE_DIM + tid_col;
41        const uint i1 = tile_row * TILE_DIM + tid_row + 8 * y;
42        if (i0 < p.ne10 && i1 < p.ne11 && i2 < p.ne12 && i3 < p.ne13) {
43            const uint dst_idx = i0 * p.nb10 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13;
44            // load transposed
45            data_d[get_doffset() + dst_idx] = D_TYPE(sh[tid_col][tid_row + 8 * y]);
46        }
47    }
48}
49
50#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b))
51
52void main() {
53    uint z = gl_WorkGroupID.z;
54    uint y = gl_WorkGroupID.y;
55    bool need_barrier = false;
56    for (uint z = gl_WorkGroupID.z; z < p.ne12 * p.ne13; z += gl_NumWorkGroups.z) {
57        for (uint y = gl_WorkGroupID.y; y < CEIL_DIV(p.ne11, TILE_DIM); y += gl_NumWorkGroups.y) {
58            for (uint x = gl_WorkGroupID.x; x < CEIL_DIV(p.ne10, TILE_DIM); x += gl_NumWorkGroups.x) {
59                if (need_barrier) {
60                    barrier();
61                }
62                need_barrier = true;
63                iter(uvec3(x, y, z));
64            }
65        }
66    }
67}