Matrix multiplication SGEMM: Transposed input matrix and rectangular tiles

"WebGL2-compute" NxN matrix multiplication C = A x B (SGEMM) v.5 demo.
See Kernel 5: Transposed input matrix and rectangular tiles by Cedric Nugteren.
All A, B elements are random (0 - 1). Error "er1" is calculated as the sum of |CCPU - CGPU |/(N*N) for all matrix elements. er2 = max(|CCPU - CGPU |). See also Shader 5 benchmark.

Compute Shader 5

See also the page source
#version 310 es
#define TS 32u
#define WPT 8u                 // The amount of work-per-thread, i.e. the thread-coarsening factor
#define RTS 4u                 // The reduced tile-size in one dimension  TS/WPT
#define TSDK 16u               // The tile-size in dimension K (for kernel 5 only)
#define LPT ((TSDK*WPT)/(TS))  // The amount of loads-per-thread (assume TSN==TSM)
layout (local_size_x = TS, local_size_y = RTS, local_size_z = 1) in;
layout (std430, binding = 0) readonly buffer ssbA {
  float A[];
};
layout (std430, binding = 1) readonly buffer ssbB {
  float B[];
};
layout (std430, binding = 2) writeonly buffer ssbC {
  float C[];
};
  uniform uvec3 MNK;
  shared float Asub[TSDK][TS];    // Local memory to fit a tile of A and B
  shared float Bsub[TS][TSDK+2u];
void main() {
    uint M = MNK.x, N = MNK.y, K = MNK.z;

    // Thread identifiers
    uint row = gl_LocalInvocationID.x; // Local row ID (max: TS)
    uint col = gl_LocalInvocationID.y; // Local col ID (max: TS/WPT == RTS)
    uint globalRow = TS*gl_WorkGroupID.x + row; // Row ID of C (0..M)
    uint globalCol = TS*gl_WorkGroupID.y + col; // Col ID of C (0..N)

    // Initialise the accumulation registers
    float acc[WPT];
    for (uint w=0u; w < WPT; w++) acc[w] = 0.0;

    // Loop over all tiles
    uint numTiles = K/TSDK;
    for (uint t=0u; t < numTiles; t++) {

        // Load one tile of A and B into local memory
        for (uint l=0u; l < LPT; l++) {
            uint tiledIndex = TSDK*t + col + l*RTS;
            uint indexA = (tiledIndex)*M + TS*gl_WorkGroupID.x + row;
            uint indexB = (tiledIndex)*N + TS*gl_WorkGroupID.y + row;
            Asub[col + l*RTS][row] = A[indexA];
            Bsub[row][col + l*RTS] = B[indexB];
        }
        // Synchronise to make sure the tile is loaded
        barrier();

        // Perform the computation for a single tile
        for (uint k=0u; k < TSDK; k++)
            for (uint w=0u; w < WPT; w++)
                acc[w] += Asub[k][row] * Bsub[col + w*RTS][k];

        // Synchronise before loading the next tile
        barrier();
    }
    // Store the final result in C
    for (uint w=0u; w < WPT; w++)
        C[(globalCol + w*RTS)*M + globalRow] = acc[w];
}
Simple transpose kernel for a PxQ matrix. We choose to transpose the B matrix.
#version 310 es
#define TRANSPOSEX 32u
#define TRANSPOSEY 32u
layout (local_size_x = 32, local_size_y = 32, local_size_z = 1) in;
layout (std430, binding = 2) readonly buffer ssb2 {
  float inp[];
};
layout (std430, binding = 1) writeonly buffer ssb1 {
  float outp[];
};
uniform uvec3 MNK;
  shared float buff[TRANSPOSEX][TRANSPOSEY];
void main() {
    uint P = MNK.x, Q = MNK.y;

    // Thread identifiers
    uint tx = gl_LocalInvocationID.x;
    uint ty = gl_LocalInvocationID.y;
    uint ID0 = gl_WorkGroupID.x*TRANSPOSEX + tx; // (0..P)
    uint ID1 = gl_WorkGroupID.y*TRANSPOSEY + ty; // (0..Q)

    // Swap the x and y coordinates to perform the rotation (coalesced)
    if (ID0 < P && ID1 < Q) {
        buff[ty][tx] = inp[ID1*P + ID0];
    }

    // Synchronise all threads
    memoryBarrierShared();
    barrier();

    // We don't have to swap the x and y thread indices here,
    // because that's already done in the local memory
    uint newID0 = gl_WorkGroupID.y*TRANSPOSEY + tx;
    uint newID1 = gl_WorkGroupID.x*TRANSPOSEX + ty;

    // Store the transposed result (coalesced)
    if (newID0 < Q && newID1 < P) {
        outp[newID1*Q + newID0] = buff[tx][ty];
    }
}

SGEMM in WebGL2-compute     updated 31 Mar 2019