"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.
#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]; } }