tensor_cores_doesn't_work

#include <cmath>
#include <iostream>
#include <cuda_fp16.h>  // Required for half precision operations
#include <mma.h>        // Required for tensor core operations
#include "gpu-new-forward.h"

using namespace nvcuda;  // Namespace for tensor core operations

#define TILE_WIDTH 16   // Tile width compatible with tensor cores
#define BLOCK_SIZE 512  // Block size for other kernels

// Error checking macro
#define CUDA_CHECK_ERROR(call)                                                        \
    do {                                                                              \
        cudaError_t err = call;                                                       \
        if (err != cudaSuccess) {                                                     \
            std::cerr << "CUDA error in " << __FILE__ << ":" << __LINE__ << " - "     \
                      << cudaGetErrorString(err) << std::endl;                        \
            exit(EXIT_FAILURE);                                                       \
        }                                                                             \
    } while (0)

// Original matrix unrolling kernel (unchanged)
__global__ void matrix_unrolling_kernel(const float *input, float *output,
                                        const int Batch, const int Channel,
                                        const int Height, const int Width,
                                        const int K) {
    #define in_4d(i3, i2, i1, i0) input[(i3) * (Channel * Height * Width) + (i2) * (Height * Width) + (i1) * (Width) + i0]
    #define out_3d(i1, i0) output[(i1) * (Batch * W_unroll) + i0]

    const size_t Height_out = Height - K + 1;
    const size_t Width_out = Width - K + 1;
    const size_t W_unroll = Height_out * Width_out;
    const size_t H_unroll = Channel * K * K;
    const size_t W_total_unroll = Batch * W_unroll;

    const size_t c = blockIdx.x * blockDim.x + threadIdx.x;
    const size_t hw_pos = blockIdx.y * blockDim.y + threadIdx.y;
    const size_t batch_idx = blockIdx.z * blockDim.z + threadIdx.z;

    const size_t h_out = hw_pos / Width_out;
    const size_t w_out = hw_pos % Width_out;

    if (c >= Channel || h_out >= Height_out || w_out >= Width_out || batch_idx >= Batch) {
        return;
    }

    const size_t w_unroll = h_out * Width_out + w_out;
    const size_t w_total_unroll = batch_idx * W_unroll + w_unroll;
    const size_t w_base = c * K * K;

    for (int p = 0; p < K; p++) {
        for (int q = 0; q < K; q++) {
            int h_unroll = w_base + p * K + q;
            out_3d(h_unroll, w_total_unroll) = in_4d(batch_idx, c, h_out + p, w_out + q);
        }
    }

    #undef in_4d
    #undef out_3d
}

// Helper kernel to convert float to half precision
__global__ void convertToHalf(half *out, const float *in, int size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < size) {
        out[idx] = __float2half(in[idx]);
    }
}

// Optimized matrix multiplication kernel using tensor cores
__global__ void matrixMultiplyTensorCores(const half *A, const half *B, float *C,
                                         int numARows, int numAColumns,
                                         int numBRows, int numBColumns,
                                         int numCRows, int numCColumns)
{
    // Define tensor core operation dimensions
    const int WMMA_M = 16;
    const int WMMA_N = 16;
    const int WMMA_K = 16;

    // Calculate the warp and lane indices
    int warpM = (blockIdx.y * blockDim.y + threadIdx.y) / 32;
    int warpN = (blockIdx.x * blockDim.x + threadIdx.x) / 32;

    // Declare fragments
    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, half, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> acc_frag;

    // Initialize the output to zero
    wmma::fill_fragment(acc_frag, 0.0f);

    // Calculate the starting row and column for A and B
    int aRow = warpM * WMMA_M;
    int bCol = warpN * WMMA_N;

    // Loop over K
    for (int k = 0; k < numAColumns; k += WMMA_K) {
        if (aRow < numARows && bCol < numCColumns && (k + WMMA_K) <= numAColumns && (k + WMMA_K) <= numBRows) {
            // Load the inputs
            wmma::load_matrix_sync(a_frag, A + aRow * numAColumns + k, numAColumns);
            wmma::load_matrix_sync(b_frag, B + k * numBColumns + bCol, numBColumns);

            // Perform the matrix multiplication
            wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);
        }
    }

    // Store the output
    if (aRow < numCRows && bCol < numCColumns) {
        wmma::store_matrix_sync(C + aRow * numCColumns + bCol, acc_frag, numCColumns, wmma::mem_row_major);
    }
}

// Original matrix permute kernel (unchanged)
__global__ void matrix_permute_kernel(const float *input, float *output, int Map_out,
                                      int Batch, int image_size) {
    int b = blockIdx.y;
    int x = blockIdx.x * BLOCK_SIZE + threadIdx.x;
    if (x < image_size) {
        for (int m = 0; m < Map_out; m++) {
            output[b * Map_out * image_size + m * image_size + x] =
                    input[m * Batch * image_size + b * image_size + x];
        }
    }
}

__host__ void GPUInterface::conv_forward_gpu_prolog(const float *host_output,
    const float *host_input, const float *host_mask, float **device_output_ptr,
    float **device_input_ptr, float **device_mask_ptr, const int Batch,
    const int Map_out, const int Channel, const int Height, const int Width, const int K)
{
    // Calculate output dimensions
    const int Height_out = Height - K + 1;
    const int Width_out = Width - K + 1;

    // Calculate memory sizes
    const size_t input_size = Batch * Channel * Height * Width * sizeof(float);
    const size_t mask_size = Map_out * Channel * K * K * sizeof(float);
    const size_t output_size = Batch * Map_out * Height_out * Width_out * sizeof(float);

    // Allocate device memory with error checking
    CUDA_CHECK_ERROR(cudaMalloc((void**)device_input_ptr, input_size));
    CUDA_CHECK_ERROR(cudaMalloc((void**)device_mask_ptr, mask_size));
    CUDA_CHECK_ERROR(cudaMalloc((void**)device_output_ptr, output_size));

    // Copy input data to device with error checking
    CUDA_CHECK_ERROR(cudaMemcpy(*device_input_ptr, host_input, input_size, cudaMemcpyHostToDevice));
    CUDA_CHECK_ERROR(cudaMemcpy(*device_mask_ptr, host_mask, mask_size, cudaMemcpyHostToDevice));
}

__host__ void GPUInterface::conv_forward_gpu(float *device_output, const float *device_input,
    const float *device_mask, const int Batch, const int Map_out, const int Channel,
    const int Height, const int Width, const int K)
{
    const int Height_out = Height - K + 1;
    const int Width_out = Width - K + 1;
    const int Height_unrolled = Channel * K * K;
    const int Width_unrolled = Batch * Height_out * Width_out;

    // Allocate intermediate matrices
    float *unrolled_matrix;
    float *matmul_output;
    half *A_half, *B_half; // For tensor cores

    size_t unrolled_size = static_cast<size_t>(Batch) * Channel * K * K * Height_out * Width_out * sizeof(float);
    size_t matmul_size = static_cast<size_t>(Batch) * Map_out * Height_out * Width_out * sizeof(float);

    CUDA_CHECK_ERROR(cudaMalloc((void**)&unrolled_matrix, unrolled_size));
    CUDA_CHECK_ERROR(cudaMalloc((void**)&matmul_output, matmul_size));

    // Allocate half-precision memory
    size_t A_half_size = static_cast<size_t>(Map_out) * Height_unrolled * sizeof(half);
    size_t B_half_size = static_cast<size_t>(Height_unrolled) * Width_unrolled * sizeof(half);
    CUDA_CHECK_ERROR(cudaMalloc((void**)&A_half, A_half_size));
    CUDA_CHECK_ERROR(cudaMalloc((void**)&B_half, B_half_size));

    // Set dimensions for unrolling
    dim3 blockDim_unroll(16, 16, 1);  // Adjusted for better occupancy
    dim3 gridDim_unroll(
        (Channel + blockDim_unroll.x - 1) / blockDim_unroll.x,
        (Height_out * Width_out + blockDim_unroll.y - 1) / blockDim_unroll.y,
        Batch
    );

    // Perform matrix unrolling
    matrix_unrolling_kernel<<<gridDim_unroll, blockDim_unroll>>>(device_input, unrolled_matrix,
        Batch, Channel, Height, Width, K);
    CUDA_CHECK_ERROR(cudaGetLastError());
    CUDA_CHECK_ERROR(cudaDeviceSynchronize());

    // Convert device_mask and unrolled_matrix from float to half
    // Assuming device_mask is of size Map_out * Channel * K * K
    int size_A = Map_out * Channel * K * K;
    int size_B = static_cast<int>(Batch) * Height_unrolled * Width_out; // Adjust as per actual data
    // However, to match the sizes, we need to ensure that A and B are properly sized

    // Calculate total number of elements for A and B
    int total_A_elements = Map_out * Height_unrolled;
    int total_B_elements = Height_unrolled * Width_unrolled;

    // Launch convertToHalf kernel for A (mask)
    int threads = 256;
    int blocks_A = (total_A_elements + threads - 1) / threads;
    convertToHalf<<<blocks_A, threads>>>(A_half, device_mask, total_A_elements);
    CUDA_CHECK_ERROR(cudaGetLastError());

    // Launch convertToHalf kernel for B (unrolled_matrix)
    int blocks_B = (total_B_elements + threads - 1) / threads;
    convertToHalf<<<blocks_B, threads>>>(B_half, unrolled_matrix, total_B_elements);
    CUDA_CHECK_ERROR(cudaGetLastError());
    CUDA_CHECK_ERROR(cudaDeviceSynchronize());

    // Set dimensions for tensor core matrix multiplication
    dim3 dimGrid((Width_unrolled + TILE_WIDTH - 1) / TILE_WIDTH,
                 (Map_out + TILE_WIDTH - 1) / TILE_WIDTH,
                 1);
    dim3 dimBlock(TILE_WIDTH, TILE_WIDTH, 1);

    // Perform matrix multiplication using tensor cores
    matrixMultiplyTensorCores<<<dimGrid, dimBlock>>>(A_half, B_half, matmul_output,
        Map_out, Height_unrolled, Height_unrolled, Width_unrolled,
        Map_out, Width_unrolled);
    CUDA_CHECK_ERROR(cudaGetLastError());
    CUDA_CHECK_ERROR(cudaDeviceSynchronize());

    // Permute the result
    const int out_image_size = Height_out * Width_out;
    dim3 permute_kernel_grid_dim((out_image_size + BLOCK_SIZE - 1) / BLOCK_SIZE, Batch, 1);
    matrix_permute_kernel<<<permute_kernel_grid_dim, BLOCK_SIZE>>>(matmul_output,
        device_output, Map_out, Batch, out_image_size);
    CUDA_CHECK_ERROR(cudaGetLastError());
    CUDA_CHECK_ERROR(cudaDeviceSynchronize());

    // Free intermediate buffers
    CUDA_CHECK_ERROR(cudaFree(matmul_output));
    CUDA_CHECK_ERROR(cudaFree(unrolled_matrix));
    CUDA_CHECK_ERROR(cudaFree(A_half));
    CUDA_CHECK_ERROR(cudaFree(B_half));
}

__host__ void GPUInterface::conv_forward_gpu_epilog(float *host_output, float *device_output, float *device_input, float *device_mask, const int Batch,
    const int Map_out, const int Channel, const int Height, const int Width, const int K)
{
    const int Height_out = Height - K + 1;
    const int Width_out = Width - K + 1;
    const int output_size = Batch * Map_out * Height_out * Width_out * sizeof(float);

    // Copy output back to host with error checking
    CUDA_CHECK_ERROR(cudaMemcpy(host_output, device_output, output_size, cudaMemcpyDeviceToHost));

    // Free device memory with error checking
    CUDA_CHECK_ERROR(cudaFree(device_output));
    CUDA_CHECK_ERROR(cudaFree(device_input));
    CUDA_CHECK_ERROR(cudaFree(device_mask));
}

__host__ void GPUInterface::get_device_properties()
{
    int deviceCount;
    CUDA_CHECK_ERROR(cudaGetDeviceCount(&deviceCount));

    for(int dev = 0; dev < deviceCount; dev++)
    {
        cudaDeviceProp deviceProp;
        CUDA_CHECK_ERROR(cudaGetDeviceProperties(&deviceProp, dev));

        std::cout<<"Device "<<dev<<" name: "<<deviceProp.name<<std::endl;
        std::cout<<"Computational capabilities: "<<deviceProp.major<<"."<<deviceProp.minor<<std::endl;
        std::cout<<"Max Global memory size: "<<deviceProp.totalGlobalMem<<std::endl;
        std::cout<<"Max Constant memory size: "<<deviceProp.totalConstMem<<std::endl;
        std::cout<<"Max Shared memory size per block: "<<deviceProp.sharedMemPerBlock<<std::endl;
        std::cout<<"Max threads per block: "<<deviceProp.maxThreadsPerBlock<<std::endl;
        std::cout<<"Max block dimensions: "<<deviceProp.maxThreadsDim[0]<<" x, "<<deviceProp.maxThreadsDim[1]<<" y, "<<deviceProp.maxThreadsDim[2]<<" z"<<std::endl;
        std::cout<<"Max grid dimensions: "<<deviceProp.maxGridSize[0]<<" x, "<<deviceProp.maxGridSize[1]<<" y, "<<deviceProp.maxGridSize[2]<<" z"<<std::endl;
        std::cout<<"Warp Size: "<<deviceProp.warpSize<<std::endl;
    }
}