#define TS 32
#define WPT 8
#define WIDTH 4
__kernel void sgemm(__global float *A, __global float *B, __global float *C, int M, int N, int K) {
    const int row = get_local_id(0);
    const int col = get_local_id(1);
    const int global_row = TS * get_group_id(0) + row;
    const int global_col = TS * get_group_id(1) + col;
    const int RTS = TS/WPT;

    __local float Asub[TS][TS];
    __local float Bsub[TS][TS];

    float intermediate_val[WPT];

    for(int w = 0; w<WPT; w++) {
        intermediate_val[w] = 0.0f;
    }
    const int num_tiles = (K+TS-1)/TS;
    for(int t=0; t<num_tiles;t++) {
        for (int w = 0; w<WPT; w++) {
            const int t_row = TS * t +row;
            const int t_col = TS * t +col;
            if (((global_row + w*RTS) < M) && (t_col < K))
                Asub[row + w*RTS][col] = A[(global_row + w*RTS) * K + t_col];
            else
                Asub[row + w*RTS][col] = 0;

            if (((t_row + w*RTS) < K) && (global_col < N)) 
                Bsub[row + w*RTS][col] = B[(t_row + w*RTS)* N + global_col];
            else
                Bsub[row + w*RTS][col] = 0;
        }
        barrier(CLK_LOCAL_MEM_FENCE);

        for(int k = 0; k < TS; k++) {
            for(int w=0; w<WPT; w++) {
                if (Asub[row + w*RTS][k] == 0 ||  Bsub[k][col] == 0)
                    continue;
                intermediate_val[w] += Asub[row + w*RTS][k] * Bsub[k][col];
            }
        }

        barrier(CLK_LOCAL_MEM_FENCE);
    }

    for(int w = 0; w<WPT; w++) {
        if ((global_row + w*RTS < M) && (global_col< N))
            C[(global_row + w*RTS) *N + global_col] = intermediate_val[w];
    }
}