<commit_before>/*
Original works by:
--------------------------------------------------------
MAGMA
Copyright (c) 2017 The University of Tennessee. All rights reserved.
Licensed under modified BSD license
*/


// These parameters will be determined by utils.read_code
//#define DIM_X  ${DIM_X}
//#define DIM_Y  ${DIM_Y}
//#define BLK_M  ${BLK_M}
//#define BLK_N  ${BLK_N}
//#define BLK_K  ${BLK_K}
//#define DIM_XA  ${DIM_XA}
//#define DIM_YA  ${DIM_YA}
//#define DIM_XB  ${DIM_XB}
//#define DIM_YB  ${DIM_YB}
//#define THR_N  ${THR_N}
//#define THR_M  ${THR_M}


extern "C" __global__
void sgemm(
        int M, int N, int K,
        const float* A,
        const float* B,
        float * C)
{
    int idx = threadIdx.x;
    int idy = threadIdx.y;

    int idt = DIM_X * idy + idx;

    int idxA = idt % DIM_XA;
    int idyA = idt / DIM_XA;

    int idxB = idt % DIM_XB;
    int idyB = idt / DIM_XB;

    int blx = blockIdx.x;
    int bly = blockIdx.y;

    __shared__ float sA[BLK_K][BLK_M + 1];
    __shared__ float sB[BLK_N][BLK_K + 1];

    // registers for the innermost loop
    float rC[THR_N][THR_M];
    float rA[THR_M];
    float rB[THR_N];

    float ra[BLK_K / DIM_YA][BLK_M / DIM_XA];
    float rb[BLK_N / DIM_YB][BLK_K / DIM_XB];

    const float* offs_dA = A + blx * BLK_M       + idyA * M + idxA;
    int boundA = (M * (K - 1) + M) - (blx * BLK_M + idyA * M + idxA) - 1;
    const float* offs_dB = B + bly * BLK_N * K + idyB * K + idxB;
    int boundB = (K * (N - 1) + K) - (bly * BLK_N * K + idyB * K + idxB) - 1;

    int m, n, k, kk;
    
    #pragma unroll
    for (n = 0; n < THR_N; n++)
        #pragma unroll
        for (m = 0 ; m < THR_M; m++)
            rC[n][m] = 0;

    // blockwise transpose to transpose load
    #pragma unroll
    for (n = 0; n < BLK_K; n += DIM_YA)
        #pragma unroll
        for (m = 0; m < BLK_M; m += DIM_XA)
            sA[n + idyA][m + idxA] = offs_dA[n * M + m];
    // blockwise transpose to transpose load
    #pragma unroll
    for (n = 0; n < BLK_N; n += DIM_YB)
        #pragma unroll
        for (m = 0; m < BLK_K; m += DIM_XB)
            sB[n + idyB][m + idxB] = offs_dB[n * K + m];
    __syncthreads();

    for (kk = 0; kk < K - BLK_K; kk += BLK_K)
    {
        offs_dA += BLK_K * M;
        boundA -= BLK_K * M;
        offs_dB += BLK_K;
        boundB -= BLK_K;
        
        #pragma unroll
        for (n = 0; n < BLK_K / DIM_YA; n++)
            #pragma unroll
            for (m = 0; m < BLK_M / DIM_XA; m++)
                ra[n][m] = offs_dA[n * DIM_YA * M + m * DIM_XA];

        #pragma unroll
        for (n = 0; n < BLK_N / DIM_YB; n++)
            #pragma unroll
            for (m = 0; m < BLK_K / DIM_XB; m++)
                rb[n][m] = offs_dB[n * DIM_YB * K + m * DIM_XB];

        // multiply
        #pragma unroll
        for (k = 0; k < BLK_K; k++)
        {
            #pragma unroll
            for (m = 0; m < THR_M; m++)
                rA[m] = sA[k][m * DIM_X + idx];
            
            #pragma unroll
            for (n = 0; n < THR_N; n++)
                rB[n] = sB[n * DIM_Y + idy][k];

            #pragma unroll
            for (n = 0; n < THR_N; n++) {
                #pragma unroll
                for (m = 0; m < THR_M; m++) {
                    rC[n][m] += rA[m] * rB[n];
                }
            }
        }
        __syncthreads();

        // store A regs->smem
        #pragma unroll
        for (n = 0; n < BLK_K / DIM_YA; n++)
            #pragma unroll
            for (m = 0; m < BLK_M / DIM_XA; m++)
                sA[n * DIM_YA + idyA][m * DIM_XA + idxA] = ra[n][m];

        #pragma unroll
        for (n = 0; n < BLK_N / DIM_YB; n++)
            #pragma unroll
            for (m = 0; m < BLK_K / DIM_XB; m++)
                sB[n * DIM_YB + idyB][m * DIM_XB + idxB] = rb[n][m];
        __syncthreads();
    }

    // Multiply last full (BLK_K) or partial block of columns of A and
    // rows of B.
    // It's okay that m,n exceed matrix bounds as all work is in registers
    // or shared memory, and out-of-bounds rC[n][m] will not be saved later.

    kk = K - kk;
    #pragma unroll
    for (k = 0; k < kk; k++)
    {
        #pragma unroll
        for (m = 0; m < THR_M; m++)
            rA[m] = sA[k][m * DIM_X + idx];

        #pragma unroll
        for (n = 0; n < THR_N; n++)
            rB[n] = sB[n * DIM_Y + idy][k];
        
        #pragma unroll
        for (n = 0; n < THR_N; n++) {
            #pragma unroll
            for (m = 0; m < THR_M; m++){
                rC[n][m] += rA[m] * rB[n];
            }
        }
    }

    #pragma unroll
    for (n = 0; n < THR_N; n++) {
        int coord_dCn = bly * BLK_N + n * DIM_Y + idy;
        #pragma unroll
        for (m = 0; m < THR_M; m++) {
            int coord_dCm = blx * BLK_M + m * DIM_X + idx;
            if (coord_dCm < M && coord_dCn < N) {
                C[coord_dCn * M + coord_dCm] = rC[n][m];
            }
        }
    }
}
<commit_msg>use braces correctly<commit_after>/*
Original works by:
--------------------------------------------------------
MAGMA
Copyright (c) 2017 The University of Tennessee. All rights reserved.
Licensed under modified BSD license
*/


// These parameters will be determined by utils.read_code
//#define DIM_X  ${DIM_X}
//#define DIM_Y  ${DIM_Y}
//#define BLK_M  ${BLK_M}
//#define BLK_N  ${BLK_N}
//#define BLK_K  ${BLK_K}
//#define DIM_XA  ${DIM_XA}
//#define DIM_YA  ${DIM_YA}
//#define DIM_XB  ${DIM_XB}
//#define DIM_YB  ${DIM_YB}
//#define THR_N  ${THR_N}
//#define THR_M  ${THR_M}


extern "C" __global__
void sgemm(
        int M, int N, int K,
        const float* A,
        const float* B,
        float * C)
{
    int idx = threadIdx.x;
    int idy = threadIdx.y;

    int idt = DIM_X * idy + idx;

    int idxA = idt % DIM_XA;
    int idyA = idt / DIM_XA;

    int idxB = idt % DIM_XB;
    int idyB = idt / DIM_XB;

    int blx = blockIdx.x;
    int bly = blockIdx.y;

    __shared__ float sA[BLK_K][BLK_M + 1];
    __shared__ float sB[BLK_N][BLK_K + 1];

    // registers for the innermost loop
    float rC[THR_N][THR_M];
    float rA[THR_M];
    float rB[THR_N];

    float ra[BLK_K / DIM_YA][BLK_M / DIM_XA];
    float rb[BLK_N / DIM_YB][BLK_K / DIM_XB];

    const float* offs_dA = A + blx * BLK_M       + idyA * M + idxA;
    int boundA = (M * (K - 1) + M) - (blx * BLK_M + idyA * M + idxA) - 1;
    const float* offs_dB = B + bly * BLK_N * K + idyB * K + idxB;
    int boundB = (K * (N - 1) + K) - (bly * BLK_N * K + idyB * K + idxB) - 1;

    int m, n, k, kk;
    
    #pragma unroll
    for (n = 0; n < THR_N; n++) {
        #pragma unroll
        for (m = 0 ; m < THR_M; m++) {
            rC[n][m] = 0;
        }
    }

    // blockwise transpose to transpose load
    #pragma unroll
    for (n = 0; n < BLK_K; n += DIM_YA) {
        #pragma unroll
        for (m = 0; m < BLK_M; m += DIM_XA) {
            sA[n + idyA][m + idxA] = offs_dA[n * M + m];
        }
    }
    // blockwise transpose to transpose load
    #pragma unroll
    for (n = 0; n < BLK_N; n += DIM_YB) {
        #pragma unroll
        for (m = 0; m < BLK_K; m += DIM_XB) {
            sB[n + idyB][m + idxB] = offs_dB[n * K + m];
        }
    }
    __syncthreads();

    for (kk = 0; kk < K - BLK_K; kk += BLK_K)
    {
        offs_dA += BLK_K * M;
        boundA -= BLK_K * M;
        offs_dB += BLK_K;
        boundB -= BLK_K;
        
        #pragma unroll
        for (n = 0; n < BLK_K / DIM_YA; n++) {
            #pragma unroll
            for (m = 0; m < BLK_M / DIM_XA; m++) {
                ra[n][m] = offs_dA[n * DIM_YA * M + m * DIM_XA];
            }
        }

        #pragma unroll
        for (n = 0; n < BLK_N / DIM_YB; n++) {
            #pragma unroll
            for (m = 0; m < BLK_K / DIM_XB; m++) {
                rb[n][m] = offs_dB[n * DIM_YB * K + m * DIM_XB];
            }
        }

        // multiply
        #pragma unroll
        for (k = 0; k < BLK_K; k++)
        {
            #pragma unroll
            for (m = 0; m < THR_M; m++) {
                rA[m] = sA[k][m * DIM_X + idx];
            }
            
            #pragma unroll
            for (n = 0; n < THR_N; n++) {
                rB[n] = sB[n * DIM_Y + idy][k];
            }

            #pragma unroll
            for (n = 0; n < THR_N; n++) {
                #pragma unroll
                for (m = 0; m < THR_M; m++) {
                    rC[n][m] += rA[m] * rB[n];
                }
            }
        }
        __syncthreads();

        // store A regs->smem
        #pragma unroll
        for (n = 0; n < BLK_K / DIM_YA; n++)
        {
            #pragma unroll
            for (m = 0; m < BLK_M / DIM_XA; m++)
            {
                sA[n * DIM_YA + idyA][m * DIM_XA + idxA] = ra[n][m];
            }
        }

        #pragma unroll
        for (n = 0; n < BLK_N / DIM_YB; n++)
        {
            #pragma unroll
            for (m = 0; m < BLK_K / DIM_XB; m++)
            {
                if (BLK_N <= n * DIM_YB + idyB)
                    printf("aaa %d %d\n", BLK_N, n * DIM_YB + idyB);
                if ((BLK_K + 1) <= (m * DIM_XB + idxB))
                    printf("bbb %d %d %d %d\n", BLK_K + 1, m * DIM_XB + idxB, m, BLK_K / DIM_XB);
                sB[n * DIM_YB + idyB][m * DIM_XB + idxB] = rb[n][m];
            }
        }
        __syncthreads();
    }

    // Multiply last full (BLK_K) or partial block of columns of A and
    // rows of B.
    // It's okay that m,n exceed matrix bounds as all work is in registers
    // or shared memory, and out-of-bounds rC[n][m] will not be saved later.

    kk = K - kk;
    #pragma unroll
    for (k = 0; k < kk; k++)
    {
        #pragma unroll
        for (m = 0; m < THR_M; m++) {
            rA[m] = sA[k][m * DIM_X + idx];
        }

        #pragma unroll
        for (n = 0; n < THR_N; n++) {
            rB[n] = sB[n * DIM_Y + idy][k];
        }
        
        #pragma unroll
        for (n = 0; n < THR_N; n++) {
            #pragma unroll
            for (m = 0; m < THR_M; m++) {
                rC[n][m] += rA[m] * rB[n];
            }
        }
    }

    #pragma unroll
    for (n = 0; n < THR_N; n++) {
        int coord_dCn = bly * BLK_N + n * DIM_Y + idy;
        #pragma unroll
        for (m = 0; m < THR_M; m++) {
            int coord_dCm = blx * BLK_M + m * DIM_X + idx;
            if (coord_dCm < M && coord_dCn < N) {
                C[coord_dCn * M + coord_dCm] = rC[n][m];
            }
        }
    }
}

{'fraction_non_alphanumeric': 0.15617537491641992, 'fraction_numerical': 0.007450568344636546, 'mean_word_length': 1.3714609286523216, 'pattern_counts': {'":': 0, '<': 63, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7449371', 'n_tokens_mistral': 4530, 'n_tokens_neox': 4310, 'n_words': 1169}