

@cuda.jit(argtypes=[f4[:, :], f4[:, :], f4[:, :]])
def cu_square_matrix_mul(A, B, C):
    sA = cuda.shared.array(shape=(tpb, tpb), dtype=f4)
    sB = cuda.shared.array(shape=(tpb, tpb), dtype=f4)
    tx = cuda.threadIdx.x
    ty = cuda.threadIdx.y
    bx = cuda.blockIdx.x
    by = cuda.blockIdx.y
    bw = cuda.blockDim.x
    bh = cuda.blockDim.y
    x = (tx + (bx * bw))
    y = (ty + (by * bh))
    acc = 0.0
    for i in range(bpg):
        if ((x < n) and (y < n)):
            sA[(ty, tx)] = A[(y, (tx + (i * tpb)))]
            sB[(ty, tx)] = B[((ty + (i * tpb)), x)]
        cuda.syncthreads()
        if ((x < n) and (y < n)):
            for j in range(tpb):
                acc += (sA[(ty, j)] * sB[(j, tx)])
        cuda.syncthreads()
    if ((x < n) and (y < n)):
        C[(y, x)] = acc

dataset/ETHPy150Open numba/numba/examples/cudajit/matmul_smem.py cu_square_matrix_mul/original