// CUDA error checking macro
#define CUDA_CHECK(err) \
    if (err != cudaSuccess) { \
        std::cerr << "CUDA Error: " << cudaGetErrorString(err) << " in " << __FILE__ << " at line " << __LINE__ << std::endl; \
        exit(EXIT_FAILURE); \
    }

#define OFFSET(row, col, ld) ((row) * (ld) + (col))

// transfer float4
#define FETCH_FLOAT4(pointer) (*reinterpret_cast<float4*>(&(pointer)))

// CPU convolution implementation
void conv2d_cpu(
    const float* input, const float* kernel, float* output,
    int channels, int height, int width,
    int num_kernels, int kernel_size,
    int pad, int stride)
{
    int out_height = (height + 2 * pad - kernel_size) / stride + 1;
    int out_width = (width + 2 * pad - kernel_size) / stride + 1;

    for (int n = 0; n < num_kernels; ++n) {
        for (int h = 0; h < out_height; ++h) {
            for (int w = 0; w < out_width; ++w) {
                float sum = 0.0f;
                for (int c = 0; c < channels; ++c) {
                    for (int kh = 0; kh < kernel_size; ++kh) {
                        for (int kw = 0; kw < kernel_size; ++kw) {
                            int h_in = h * stride - pad + kh;
                            int w_in = w * stride - pad + kw;
                            if (h_in >= 0 && h_in < height && w_in >= 0 && w_in < width) {
                                float val = input[(c * height + h_in) * width + w_in];
                                float wgt = kernel[((n * channels + c) * kernel_size + kh) * kernel_size + kw];
                                sum += val * wgt;
                            }
                        }
                    }
                }
                output[(n * out_height + h) * out_width + w] = sum;
            }
        }
    }
}

// im2col kernel
__global__ void im2col(
    const float* data_im,
    int channels, int height, int width,
    int ksize, int pad, int stride,
    int height_col, int width_col,
    float* data_col)
{
    int index = blockIdx.x * blockDim.x + threadIdx.x;
    int n = channels * ksize * ksize * height_col * width_col;
    if (index < n) {
        int w_out = index % width_col;
        int h_index = index / width_col;
        int h_out = h_index % height_col;
        int channel_in = h_index / height_col;
        int channel_out = channel_in * ksize * ksize;
    
        int kx = channel_out % ksize;
        int ky = (channel_out / ksize) % ksize;
        int c = channel_out / (ksize * ksize);
    
        int im_row = h_out * stride - pad + ky;
        int im_col = w_out * stride - pad + kx;
    
        float val = 0.0f;
        if (im_row >= 0 && im_row < height && im_col >= 0 && im_col < width) {
            val = data_im[(c * height + im_row) * width + im_col];
        }
        data_col[(channel_in * height_col + h_out) * width_col + w_out] = val;
    }
}

// DoubleBufferMatMul kernel
template <
    const int BLOCK_SIZE_M,
    const int BLOCK_SIZE_K,
    const int BLOCK_SIZE_N,
    const int THREAD_SIZE_Y,
    const int THREAD_SIZE_X,
    const bool ENABLE_DOUBLE_BUFFER
>
__global__ void DoubleBufferMatMul(
    float* __restrict__ A,
    float* __restrict__ B,
    float* __restrict__ C,
    const int M, const int K, const int N)
{
    // Block index
    int bx = blockIdx.x;
    int by = blockIdx.y;

    // Thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    // size of thread block
    const int bszx = BLOCK_SIZE_N / THREAD_SIZE_X;
    const int bszy = BLOCK_SIZE_M / THREAD_SIZE_Y;
    const int THREAD_NUM_PER_BLOCK = bszy * bszx;

    // thread id
    const int tid = ty * bszx + tx;

    // shared memory
    __shared__ float As[BLOCK_SIZE_M * 2][BLOCK_SIZE_K];
    __shared__ float Bs[BLOCK_SIZE_K * 2][BLOCK_SIZE_N];

    // registers for C
    float accum[THREAD_SIZE_Y][THREAD_SIZE_X] = { 0 };
    // registers for A and B
    float frag_a[THREAD_SIZE_Y];
    float frag_b[THREAD_SIZE_X];

    // threads needed to load one row of tile
    // / 4 is because float4 is used
    const int A_TILE_THREAD_PER_ROW = BLOCK_SIZE_K / 4;
    const int B_TILE_THREAD_PER_ROW = BLOCK_SIZE_N / 4;

    // row number and col number that needs to be loaded by this thread
    const int A_TILE_ROW_START = tid / A_TILE_THREAD_PER_ROW;
    const int B_TILE_ROW_START = tid / B_TILE_THREAD_PER_ROW;

    const int A_TILE_COL = tid % A_TILE_THREAD_PER_ROW * 4;
    const int B_TILE_COL = tid % B_TILE_THREAD_PER_ROW * 4;

    // row stride that thread uses to load multiple rows of a tile
    const int A_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / A_TILE_THREAD_PER_ROW;
    const int B_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / B_TILE_THREAD_PER_ROW;

    int A_SM_OFFSET = 0;
    int B_SM_OFFSET = 0;

    for (int tile_idx = 0; tile_idx < K; tile_idx += BLOCK_SIZE_K) {
        // load A from global memory to shared memory
#pragma unroll
        for (int i = 0; i < BLOCK_SIZE_M; i += A_TILE_ROW_STRIDE) {
            int row = BLOCK_SIZE_M * by + A_TILE_ROW_START + i;
            int col = A_TILE_COL + tile_idx;
            if (row < M && col < K) {
                FETCH_FLOAT4(As[A_TILE_ROW_START + i + A_SM_OFFSET][A_TILE_COL]) = FETCH_FLOAT4(A[OFFSET(
                    row, // row
                    col, // col
                    K)]);
            } else {
                As[A_TILE_ROW_START + i + A_SM_OFFSET][A_TILE_COL] = 0.0f;
            }
        }

        // load B from global memory to shared memory
#pragma unroll
        for (int i = 0; i < BLOCK_SIZE_K; i += B_TILE_ROW_STRIDE) {
            int row = tile_idx + B_TILE_ROW_START + i;
            int col = B_TILE_COL + BLOCK_SIZE_N * bx;
            if (row < K && col < N) {
                FETCH_FLOAT4(Bs[B_TILE_ROW_START + i + B_SM_OFFSET][B_TILE_COL]) = FETCH_FLOAT4(B[OFFSET(
                    row, // row
                    col, // col
                    N)]);
            } else {
                Bs[B_TILE_ROW_START + i + B_SM_OFFSET][B_TILE_COL] = 0.0f;
            }
        }

        __syncthreads();

        // compute c
#pragma unroll
        for (int k = 0; k < BLOCK_SIZE_K; ++k) {
            // load A from shared memory to register
#pragma unroll
            for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
                frag_a[thread_y] = As[ty * THREAD_SIZE_Y + thread_y + A_SM_OFFSET][k];
            }

            // load B from shared memory to register
#pragma unroll
            for (int thread_x = 0; thread_x < THREAD_SIZE_X; thread_x += 4) {
                FETCH_FLOAT4(frag_b[thread_x]) = FETCH_FLOAT4(Bs[k + B_SM_OFFSET][THREAD_SIZE_X * tx + thread_x]);
            }

#pragma unroll
            for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
                for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
                    accum[thread_y][thread_x] += frag_a[thread_y] * frag_b[thread_x];
                }
            }

        }
        // change offset so new block of A and B will not override the previous data
        // therefore __syncthreads() is not needed;
        A_SM_OFFSET ^= BLOCK_SIZE_M;
        B_SM_OFFSET ^= BLOCK_SIZE_K;
    }

    // store back to C
#pragma unroll
    for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
        for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
            int row = BLOCK_SIZE_M * by + ty * THREAD_SIZE_Y + thread_y;
            int col = BLOCK_SIZE_N * bx + tx * THREAD_SIZE_X + thread_x;
            if (row < M && col < N) {
                C[OFFSET(
                    row,
                    col,
                    N)] = accum[thread_y][thread_x];
            }
        }
    }
}

// FusedIm2ColDoubleBufferMatMul kernel
template <
    const int BLOCK_SIZE_M,
    const int BLOCK_SIZE_K,
    const int BLOCK_SIZE_N,
    const int THREAD_SIZE_Y,
    const int THREAD_SIZE_X,
    const bool ENABLE_DOUBLE_BUFFER
>
__global__ void FusedIm2ColDoubleBufferMatMul(
    float* __restrict__ input,
    float* __restrict__ kernel,
    float* __restrict__ output,
    const int C, const int H, const int W,
    const int K, const int P, const int S,
    const int OH, const int OW)
{
    // Block index
    int bx = blockIdx.x;
    int by = blockIdx.y;

    // Thread index
    int tx = threadIdx.x;
    int ty = threadIdx.y;

    // size of thread block
    const int bszx = BLOCK_SIZE_N / THREAD_SIZE_X;
    const int bszy = BLOCK_SIZE_M / THREAD_SIZE_Y;
    const int THREAD_NUM_PER_BLOCK = bszy * bszx;

    // thread id
    const int tid = ty * bszx + tx;

    // shared memory
    __shared__ float As[BLOCK_SIZE_M * 2][BLOCK_SIZE_K];
    __shared__ float Bs[BLOCK_SIZE_K * 2][BLOCK_SIZE_N];

    // registers for C
    float accum[THREAD_SIZE_Y][THREAD_SIZE_X] = { 0 };
    // registers for A and B
    float frag_a[THREAD_SIZE_Y];
    float frag_b[THREAD_SIZE_X];

    // threads needed to load one row of tile
    // / 4 is because float4 is used
    const int A_TILE_THREAD_PER_ROW = BLOCK_SIZE_K / 4;
    const int B_TILE_THREAD_PER_ROW = BLOCK_SIZE_N / 4;

    // row number and col number that needs to be loaded by this thread
    const int A_TILE_ROW_START = tid / A_TILE_THREAD_PER_ROW;
    const int B_TILE_ROW_START = tid / B_TILE_THREAD_PER_ROW;

    const int A_TILE_COL = tid % A_TILE_THREAD_PER_ROW * 4;
    const int B_TILE_COL = tid % B_TILE_THREAD_PER_ROW * 4;

    // row stride that thread uses to load multiple rows of a tile
    const int A_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / A_TILE_THREAD_PER_ROW;
    const int B_TILE_ROW_STRIDE = THREAD_NUM_PER_BLOCK / B_TILE_THREAD_PER_ROW;

    int A_SM_OFFSET = 0;
    int B_SM_OFFSET = 0;

    const int tile_size = C * P * P;

    for (int tile_idx = 0; tile_idx < tile_size; tile_idx += BLOCK_SIZE_K) {
        // load A (kernel) from global memory to shared memory
#pragma unroll
        for (int i = 0; i < BLOCK_SIZE_M; i += A_TILE_ROW_STRIDE) {
            int row = BLOCK_SIZE_M * by + A_TILE_ROW_START + i;
            int col = A_TILE_COL + tile_idx;
            if (row < K && col < tile_size) {
                FETCH_FLOAT4(As[A_TILE_ROW_START + i + A_SM_OFFSET][A_TILE_COL]) = FETCH_FLOAT4(kernel[OFFSET(
                    row, // row
                    col, // col
                    tile_size)]);
            } else {
                As[A_TILE_ROW_START + i + A_SM_OFFSET][A_TILE_COL] = 0.0f;
            }
        }

        // load B (input) into shared memory on-the-fly from the input image
#pragma unroll
        for (int i = 0; i < BLOCK_SIZE_K; i += B_TILE_ROW_STRIDE) {
            int k_row = tile_idx + B_TILE_ROW_START + i;
            int n_col = BLOCK_SIZE_N * bx + B_TILE_COL;
            if (k_row < tile_size) {
#pragma unroll
                for (int ii = 0; ii < 4; ++ii) { // Handle float4
                    int k = k_row + ii;
                    if (k < tile_size) {
                        int c = k / (P * P);
                        int kh = (k / P) % P;
                        int kw = k % P;

#pragma unroll
                        for (int jj = 0; jj < 4; ++jj) {
                            int n = n_col + jj;
                            if (n < OH * OW) {
                                int h_out = n / OW;
                                int w_out = n % OW;

                                int h_in = h_out * S - P / 2 + kh;
                                int w_in = w_out * S - P / 2 + kw;

                                if (h_in >= 0 && h_in < H && w_in >= 0 && w_in < W) {
                                    Bs[B_TILE_ROW_START + i + B_SM_OFFSET + ii][n_col + jj - BLOCK_SIZE_N * bx] =
                                        input[(c * H + h_in) * W + w_in];
                                } else {
                                    Bs[B_TILE_ROW_START + i + B_SM_OFFSET + ii][n_col + jj - BLOCK_SIZE_N * bx] = 0.0f;
                                }
                            }
                        }
                    }
                }
            }
        }

        __syncthreads();

        // compute c
#pragma unroll
        for (int k = 0; k < BLOCK_SIZE_K; ++k) {
            // load A from shared memory to register
#pragma unroll
            for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
                frag_a[thread_y] = As[ty * THREAD_SIZE_Y + thread_y + A_SM_OFFSET][k];
            }

            // load B from shared memory to register
#pragma unroll
            for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
                frag_b[thread_x] = Bs[k + B_SM_OFFSET][THREAD_SIZE_X * tx + thread_x];
            }

#pragma unroll
            for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
                for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
                    accum[thread_y][thread_x] += frag_a[thread_y] * frag_b[thread_x];
                }
            }

        }
        // change offset so new block of A and B will not override the previous data
        // therefore __syncthreads() is not needed;
        A_SM_OFFSET ^= BLOCK_SIZE_M;
        B_SM_OFFSET ^= BLOCK_SIZE_K;
    }

    // store back to C (output)
#pragma unroll
    for (int thread_y = 0; thread_y < THREAD_SIZE_Y; ++thread_y) {
#pragma unroll
        for (int thread_x = 0; thread_x < THREAD_SIZE_X; ++thread_x) {
            int row = BLOCK_SIZE_M * by + ty * THREAD_SIZE_Y + thread_y;
            int col = BLOCK_SIZE_N * bx + tx * THREAD_SIZE_X + thread_x;
            if (row < K && col < OH * OW) {
                output[row * (OH * OW) + col] = accum[thread_y][thread_x];
            }
        }
    }
}

int main() {
    // Input dimensions
    const int channels = 3;      // Input channels (C)
    const int height = 224;       // Input height (H)
    const int width = 224;        // Input width (W)

    // Kernel dimensions
    const int num_kernels = 32;  // Number of kernels/output channels (K)
    const int kernel_size = 3;   // Kernel size (P)
    const int pad = 1;           // Padding
    const int stride = 2;        // Stride (S)

    // Output dimensions
    const int out_height = (height + 2 * pad - kernel_size) / stride + 1; // (OH)
    const int out_width = (width + 2 * pad - kernel_size) / stride + 1;   // (OW)

    // Compute size of im2col matrix
    const int ksize = kernel_size;
    const int height_col = out_height;
    const int width_col = out_width;
    const int channels_col = channels * ksize * ksize;
    const int col_matrix_height = channels_col;
    const int col_matrix_width = height_col * width_col;

    // GEMM dimensions
    const int M = num_kernels;          // Number of output channels
    const int K = channels * kernel_size * kernel_size; // Shared dimension
    const int N = out_height * out_width; // Output spatial size

    // Allocate host memory
    size_t input_size = channels * height * width * sizeof(float);
    size_t kernel_size_bytes = num_kernels * channels * kernel_size * kernel_size * sizeof(float);
    size_t output_size = num_kernels * out_height * out_width * sizeof(float);
    size_t col_matrix_size = channels_col * out_height * out_width * sizeof(float);

    float* h_input = (float*)malloc(input_size);
    float* h_kernel = (float*)malloc(kernel_size_bytes);
    float* h_output_cpu = (float*)malloc(output_size);
    float* h_output_gpu = (float*)malloc(output_size);
    float* h_output_fused = (float*)malloc(output_size);

    // Initialize input and kernel with random values
    srand(time(NULL));
    for (int i = 0; i < channels * height * width; ++i) {
        h_input[i] = static_cast<float>(rand()) / RAND_MAX;
    }
    for (int i = 0; i < num_kernels * channels * kernel_size * kernel_size; ++i) {
        h_kernel[i] = static_cast<float>(rand()) / RAND_MAX;
    }

    // CPU convolution
    std::cout << "Running CPU convolution..." << std::endl;
    clock_t cpu_start = clock();
    conv2d_cpu(h_input, h_kernel, h_output_cpu,
               channels, height, width,
               num_kernels, kernel_size,
               pad, stride);
    clock_t cpu_end = clock();
    double cpu_time = 1000.0 * (cpu_end - cpu_start) / CLOCKS_PER_SEC;
    std::cout << "CPU convolution time: " << cpu_time << " ms" << std::endl;

    // GPU computation

    // Allocate device memory
    float* d_input;
    float* d_kernel;
    float* d_output;
    float* d_col;
    float* d_output_fused;

    CUDA_CHECK(cudaMalloc((void**)&d_input, input_size));
    CUDA_CHECK(cudaMalloc((void**)&d_kernel, kernel_size_bytes));
    CUDA_CHECK(cudaMalloc((void**)&d_output, output_size));
    CUDA_CHECK(cudaMalloc((void**)&d_col, col_matrix_size));
    CUDA_CHECK(cudaMalloc((void**)&d_output_fused, output_size));

    // Copy data to device
    CUDA_CHECK(cudaMemcpy(d_input, h_input, input_size, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_kernel, h_kernel, kernel_size_bytes, cudaMemcpyHostToDevice));

    // Run im2col kernel
    int im2col_threads = 256;
    int im2col_blocks = (channels_col * out_height * out_width + im2col_threads - 1) / im2col_threads;

    std::cout << "Running GPU convolution with im2col + DoubleBufferMatMul..." << std::endl;

    cudaEvent_t start, stop;
    float elapsedTime;

    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    cudaEventRecord(start, 0);

    im2col<<<im2col_blocks, im2col_threads>>>(
        d_input, channels, height, width,
        kernel_size, pad, stride,
        out_height, out_width,
        d_col);

    CUDA_CHECK(cudaGetLastError());

    // Define kernel parameters for DoubleBufferMatMul
    const int BLOCK_SIZE_M = 64;  // Adjust as per your GPU
    const int BLOCK_SIZE_K = 16;
    const int BLOCK_SIZE_N = 64;
    const int THREAD_SIZE_Y = 8;
    const int THREAD_SIZE_X = 8;
    const bool ENABLE_DOUBLE_BUFFER = true;

    // Ensure block sizes are compatible
    static_assert(BLOCK_SIZE_M % THREAD_SIZE_Y == 0, "BLOCK_SIZE_M must be divisible by THREAD_SIZE_Y");
    static_assert(BLOCK_SIZE_N % THREAD_SIZE_X == 0, "BLOCK_SIZE_N must be divisible by THREAD_SIZE_X");
    static_assert(BLOCK_SIZE_K % 4 == 0, "BLOCK_SIZE_K must be divisible by 4");

    dim3 block_dim(BLOCK_SIZE_N / THREAD_SIZE_X, BLOCK_SIZE_M / THREAD_SIZE_Y);
    dim3 grid_dim((N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N, (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M);

    // Run DoubleBufferMatMul kernel
    DoubleBufferMatMul<BLOCK_SIZE_M, BLOCK_SIZE_K, BLOCK_SIZE_N, THREAD_SIZE_Y, THREAD_SIZE_X, ENABLE_DOUBLE_BUFFER><<<grid_dim, block_dim>>>(
        d_kernel, // A matrix (weights)
        d_col,    // B matrix (im2col output)
        d_output, // C matrix (output)
        M, K, N
    );

    CUDA_CHECK(cudaGetLastError());

    cudaEventRecord(stop, 0);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&elapsedTime, start, stop);
    std::cout << "im2col + DoubleBufferMatMul execution time: " << elapsedTime << " ms" << std::endl;

    // Copy result back to host
    CUDA_CHECK(cudaMemcpy(h_output_gpu, d_output, output_size, cudaMemcpyDeviceToHost));

    // Run fused convolution kernel
    std::cout << "Running fused convolution kernel..." << std::endl;

    cudaEventRecord(start, 0);

    FusedIm2ColDoubleBufferMatMul<BLOCK_SIZE_M, BLOCK_SIZE_K, BLOCK_SIZE_N, THREAD_SIZE_Y, THREAD_SIZE_X, ENABLE_DOUBLE_BUFFER><<<grid_dim, block_dim>>>(
        d_input,       // Input image
        d_kernel,      // Kernel weights
        d_output_fused,// Output
        channels, height, width,  // Input dimensions
        num_kernels, kernel_size, stride,   // Kernel parameters
        out_height, out_width     // Output dimensions
    );

    CUDA_CHECK(cudaGetLastError());

    cudaEventRecord(stop, 0);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&elapsedTime, start, stop);
    std::cout << "Fused convolution kernel execution time: " << elapsedTime << " ms" << std::endl;

    // Copy fused output back to host
    CUDA_CHECK(cudaMemcpy(h_output_fused, d_output_fused, output_size, cudaMemcpyDeviceToHost));

    // Compare CPU and GPU results (im2col + DoubleBufferMatMul)
    float max_error = 0.0f;
    float sum_error = 0.0f;
    for (int i = 0; i < num_kernels * out_height * out_width; ++i) {
        float error = fabs(h_output_cpu[i] - h_output_gpu[i]);
        sum_error += error;
        if (error > max_error) max_error = error;
    }
    std::cout << "im2col + DoubleBufferMatMul - Max error: " << max_error << std::endl;
    std::cout << "im2col + DoubleBufferMatMul - Average error: " << sum_error / (num_kernels * out_height * out_width) << std::endl;

    // Compare CPU and Fused GPU results
    max_error = 0.0f;
    sum_error = 0.0f;
    for (int i = 0; i < num_kernels * out_height * out_width; ++i) {
        float error = fabs(h_output_cpu[i] - h_output_fused[i]);
        sum_error += error;
        if (error > max_error) max_error = error;
    }
    std::cout << "Fused Kernel - Max error: " << max_error << std::endl;
    std::cout << "Fused Kernel - Average error: " << sum_error / (num_kernels * out_height * out_width) << std::endl;

    // Free device memory
    cudaFree(d_input);
    cudaFree(d_kernel);
    cudaFree(d_output);
    cudaFree(d_col);
    cudaFree(d_output_fused);

    // Free host memory
    free(h_input);
    free(h_kernel);
    free(h_output_cpu);
    free(h_output_gpu);
    free(h_output_fused);

    return 0;
}// Define kernel parameters for DoubleBufferMatMul
    const int BLOCK_SIZE_M = 64;  // Adjust as per your GPU
    const int BLOCK_SIZE_K = 16;
    const int BLOCK_SIZE_N = 64;
    const int THREAD_SIZE_Y = 8;
    const int THREAD_SIZE_X = 8;
    const bool ENABLE_DOUBLE_BUFFER = true;

    // Ensure block sizes are compatible
    static_assert(BLOCK_SIZE_M % THREAD_SIZE_Y == 0, "BLOCK_SIZE_M must be divisible by THREAD_SIZE_Y");
    static_assert(BLOCK_SIZE_N % THREAD_SIZE_X == 0, "BLOCK_SIZE_N must be divisible by THREAD_SIZE_X");
    static_assert(BLOCK_SIZE_K % 4 == 0, "BLOCK_SIZE_K must be divisible by 4");

    dim3 block_dim(BLOCK_SIZE_N / THREAD_SIZE_X, BLOCK_SIZE_M / THREAD_SIZE_Y);
    dim3 grid_dim((N + BLOCK_SIZE_N - 1) / BLOCK_SIZE_N, (M + BLOCK_SIZE_M - 1) / BLOCK_SIZE_M);

    // Run DoubleBufferMatMul kernel
    DoubleBufferMatMul<BLOCK_SIZE_M, BLOCK_SIZE_K, BLOCK_SIZE_N, THREAD_SIZE_Y, THREAD_SIZE_X, ENABLE_DOUBLE_BUFFER><<<grid_dim, block_dim>>>(
        d_kernel, // A matrix (weights)
        d_col,    // B matrix (im2col output)
        d_output, // C matrix (output)
        M, K, N
    ); is this grid dim and block dim correct cuz, the input dim is 3 x 224 x 224 and the filter is 32 x 3 x 3 x 3 and the output becomes 32 x 112 x 112 but with this setup i get 8 x 8 for blockdim and 196 x 1 x 1 for grid dim, is this right? can you fix it if wrong