#include <cmath>

template <typename T, int TILE_SIZE>
__global__ void mat_mul(T *A, T *B, T *C, int N, int M, int K) {
    __shared__ T sA[TILE_SIZE][TILE_SIZE];
    __shared__ T sB[TILE_SIZE][TILE_SIZE];
    
    int bx = blockIdx.x, by = blockIdx.y;
    int tx = threadIdx.x, ty = threadIdx.y;
    
    int row = by * TILE_SIZE + ty;
    int col = bx * TILE_SIZE + tx;
    
    T sum = 0;
    
    for (int tile = 0; tile < ceil((float)M/TILE_SIZE); tile++) {
        if (row < N && (tile * TILE_SIZE + tx) < M) {
            sA[ty][tx] = A[row * M + (tile * TILE_SIZE + tx)];
        } else {
            sA[ty][tx] = 0;
        }

        if ((tile * TILE_SIZE + ty) < M && col < K) {
            sB[ty][tx] = B[(tile * TILE_SIZE + ty) * K + col];
        } else {
            sB[ty][tx] = 0;
        }
        __syncthreads();
        
        for (int k = 0; k < TILE_SIZE; k++) {
            sum += sA[ty][k] * sB[k][tx];
        }
        __syncthreads();
    }
    
    if (row < N && col < K) {
        C[row * K + col] = sum;
    }
}

#define MAT_TYPE int
#define MAT_FMT "%d\t"
#define N 5
#define M 7
#define K 3
#define A_LEN (N * M)
#define B_LEN (M * K)
#define C_LEN (N * K)
#define A_SIZE (sizeof(MAT_TYPE) * N * M)
#define B_SIZE (sizeof(MAT_TYPE) * M * K)
#define C_SIZE (sizeof(MAT_TYPE) * N * K)

#include <cstdio>
#include <random>

template <typename T>
void mat_print(T *a, const char *fmt, int n, int m) {
    for (auto row = 0; row < n; row++) {
        for (auto col = 0; col < m; col++) {
            printf(fmt, a[row * m + col]);
        }
        printf("\n");
    }
}

int main() {
    std::random_device rd;
    std::mt19937 engine(rd());
    std::uniform_int_distribution<MAT_TYPE> dist(1, 10);

    MAT_TYPE buf[A_LEN + B_LEN + C_LEN];
    for (auto i = 0; i < A_LEN + B_LEN; i++) {
        buf[i] = dist(engine);
    }

    MAT_TYPE *a = buf;
    MAT_TYPE *b = a + A_LEN;
    MAT_TYPE *c = b + B_LEN;

    printf("\na\n");
    mat_print(a, MAT_FMT, N, M);
    printf("\nb\n");
    mat_print(b, MAT_FMT, M, K);

    MAT_TYPE *d_a, *d_b, *d_c;
    cudaMalloc(&d_a, A_SIZE);
    cudaMalloc(&d_b, B_SIZE);
    cudaMalloc(&d_c, C_SIZE);

    cudaMemcpy(d_a, a, A_SIZE, cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, b, B_SIZE, cudaMemcpyHostToDevice);

    dim3 blockDim(4, 4);
    dim3 threadDim(4, 4);
    mat_mul<MAT_TYPE, 4><<<blockDim, threadDim>>>(d_a, d_b, d_c, N, M, K);

    cudaMemcpy(c, d_c, C_SIZE, cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();

    cudaFree(a);
    cudaFree(b);
    cudaFree(c);

    printf("\nc\n");
    mat_print(c, MAT_FMT, N, K);
}