take a look at this code, we will build on top of it, but before we do that, I want you check something and see if it is right. after the first 3 matrix multiplications, the output will be 12 x 128 but there is also another input coming from another source, this input is 12 x 512. I need these two to be combined, in a way that we would have 12 x 512 first then we concatenate 12 x  128 to it at the end of it. here in this concatenate we have two dimensions, row and col, but i think it is unnecessary, as the data is flat in the memory? like (12 x 512 = 6144) + (12 x 128 =  1536) so in total 7680 right. so maybe the kernel can be even simpler and faster? if so please implement that and we will proceed with the next part. #include <iostream>
#include <cuda_runtime.h>
#include <cublas_v2.h>

// Define initial matrix sizes
#define M1 12
#define K1 3
#define N1 32

#define K2 N1
#define N2 64

#define K3 N2
#define N3 128

#define D_DIM 512 // Dimension of the additional input

// Tile size for custom kernel
#define TILE_SIZE 16

// Error checking macro
#define CUDA_CHECK(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)

// Custom matrix multiplication kernel with tiling and loop unrolling
__global__ void matMulOptimizedKernel(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
    int M, int N, int K)
{
    __shared__ float As[TILE_SIZE][TILE_SIZE];
    __shared__ float Bs[TILE_SIZE][TILE_SIZE];

    int row = blockIdx.y * TILE_SIZE + threadIdx.y;
    int col = blockIdx.x * TILE_SIZE + threadIdx.x;

    float Cvalue = 0.0f;

    for (int t = 0; t < (K + TILE_SIZE - 1) / TILE_SIZE; ++t) {
        int tiledCol = t * TILE_SIZE + threadIdx.x;
        int tiledRow = t * TILE_SIZE + threadIdx.y;

        if (row < M && tiledCol < K)
            As[threadIdx.y][threadIdx.x] = A[row * K + tiledCol];
        else
            As[threadIdx.y][threadIdx.x] = 0.0f;

        if (col < N && tiledRow < K)
            Bs[threadIdx.y][threadIdx.x] = B[tiledRow * N + col];
        else
            Bs[threadIdx.y][threadIdx.x] = 0.0f;

        __syncthreads();

#pragma unroll
        for (int i = 0; i < TILE_SIZE; ++i) {
            Cvalue += As[threadIdx.y][i] * Bs[i][threadIdx.x];
        }

        __syncthreads();
    }

    if (row < M && col < N)
        C[row * N + col] = Cvalue;
}

// Kernel to concatenate two matrices horizontally (assuming row-major order)
__global__ void concatenateMatrices(const float* A, const float* B, float* C, int M, int s, int d)
{
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < (N1 + N2)) {
        if (col < N1)
            C[row * (N1 + N2) + col] = A[row * N1 + col];
        else
            C[row * (N1 + N2) + col] = B[row * N2 + (col - N1)];
    }
}

void matMulOptimized(const float* d_A, const float* d_B, float* d_C, int M, int N, int K)
{
    dim3 blockSize(TILE_SIZE, TILE_SIZE);
    dim3 gridSize((N + TILE_SIZE - 1) / TILE_SIZE, (M + TILE_SIZE - 1) / TILE_SIZE);

    matMulOptimizedKernel << <gridSize, blockSize >> > (d_A, d_B, d_C, M, N, K);
    CUDA_CHECK(cudaGetLastError());
}

int main()
{
    // Initialize matrices
    float* h_A1, * h_B1, * h_B2, * h_B3, * h_D;
    float* d_A1, * d_B1, * d_B2, * d_B3, * d_D, * d_E, * d_Vec;
    float* d_C1, * d_C2, * d_C3;
    float* d_W1, * d_W2, * d_W3, * d_W4;
    float* d_F1, * d_F2, * d_F3, * d_F4;

    // Size calculations
    size_t size_A1 = M1 * K1 * sizeof(float);
    size_t size_B1 = K1 * N1 * sizeof(float);
    size_t size_B2 = K2 * N2 * sizeof(float);
    size_t size_B3 = K3 * N3 * sizeof(float);
    size_t size_D = M1 * D_DIM * sizeof(float); // Size of h_D (12 x 512)
    size_t size_E = M1 * (N3 + D_DIM) * sizeof(float); // Size of concatenated matrix (12 x 640)

    size_t size_C1 = M1 * N1 * sizeof(float);
    size_t size_C2 = M1 * N2 * sizeof(float);
    size_t size_C3 = M1 * N3 * sizeof(float);

    size_t size_Vec = M1 * (N3 + D_DIM) * sizeof(float); // 12 x 640 -> 7680 x 1

    size_t size_W1 = 1024 * 7680 * sizeof(float);
    size_t size_W2 = 512 * 1024 * sizeof(float);
    size_t size_W3 = 128 * 512 * sizeof(float);
    size_t size_W4 = 12 * 128 * sizeof(float);

    size_t size_F1 = 1024 * sizeof(float);
    size_t size_F2 = 512 * sizeof(float);
    size_t size_F3 = 128 * sizeof(float);
    size_t size_F4 = 12 * sizeof(float);

    // Allocate pinned host memory (initial inputs and weights)
    CUDA_CHECK(cudaMallocHost((void**)&h_A1, size_A1));
    CUDA_CHECK(cudaMallocHost((void**)&h_B1, size_B1));
    CUDA_CHECK(cudaMallocHost((void**)&h_B2, size_B2));
    CUDA_CHECK(cudaMallocHost((void**)&h_B3, size_B3));
    CUDA_CHECK(cudaMallocHost((void**)&h_D, size_D));

    // Initialize host matrices with random values
    srand(0);
    for (int i = 0; i < M1 * K1; ++i) h_A1[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < K1 * N1; ++i) h_B1[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < K2 * N2; ++i) h_B2[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < K3 * N3; ++i) h_B3[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < M1 * D_DIM; ++i) h_D[i] = static_cast<float>(rand()) / RAND_MAX; // Initialize h_D

    // Allocate device memory
    CUDA_CHECK(cudaMalloc((void**)&d_A1, size_A1));
    CUDA_CHECK(cudaMalloc((void**)&d_B1, size_B1));
    CUDA_CHECK(cudaMalloc((void**)&d_B2, size_B2));
    CUDA_CHECK(cudaMalloc((void**)&d_B3, size_B3));
    CUDA_CHECK(cudaMalloc((void**)&d_D, size_D)); // 12 x 512 input
    CUDA_CHECK(cudaMalloc((void**)&d_C1, size_C1)); // Intermediate results
    CUDA_CHECK(cudaMalloc((void**)&d_C2, size_C2));
    CUDA_CHECK(cudaMalloc((void**)&d_C3, size_C3));

    CUDA_CHECK(cudaMalloc((void**)&d_E, size_E)); // Concatenated matrix (12 x 640)
    d_Vec = d_E; // The vector is the same memory as d_E

    // Allocate device memory for weights and results
    CUDA_CHECK(cudaMalloc((void**)&d_W1, size_W1));
    CUDA_CHECK(cudaMalloc((void**)&d_W2, size_W2));
    CUDA_CHECK(cudaMalloc((void**)&d_W3, size_W3));
    CUDA_CHECK(cudaMalloc((void**)&d_W4, size_W4));

    CUDA_CHECK(cudaMalloc((void**)&d_F1, size_F1));
    CUDA_CHECK(cudaMalloc((void**)&d_F2, size_F2));
    CUDA_CHECK(cudaMalloc((void**)&d_F3, size_F3));
    CUDA_CHECK(cudaMalloc((void**)&d_F4, size_F4));

    // Initialize weights on host and copy to device
    float* h_W1, * h_W2, * h_W3, * h_W4;
    CUDA_CHECK(cudaMallocHost((void**)&h_W1, size_W1));
    CUDA_CHECK(cudaMallocHost((void**)&h_W2, size_W2));
    CUDA_CHECK(cudaMallocHost((void**)&h_W3, size_W3));
    CUDA_CHECK(cudaMallocHost((void**)&h_W4, size_W4));

    for (int i = 0; i < 1024 * 7680; ++i) h_W1[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < 512 * 1024; ++i) h_W2[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < 128 * 512; ++i) h_W3[i] = static_cast<float>(rand()) / RAND_MAX;
    for (int i = 0; i < 12 * 128; ++i) h_W4[i] = static_cast<float>(rand()) / RAND_MAX;

    // Copy all data to device at the beginning
    CUDA_CHECK(cudaMemcpy(d_A1, h_A1, size_A1, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B1, h_B1, size_B1, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B2, h_B2, size_B2, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B3, h_B3, size_B3, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_D, h_D, size_D, cudaMemcpyHostToDevice)); // Copy h_D

    CUDA_CHECK(cudaMemcpy(d_W1, h_W1, size_W1, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_W2, h_W2, size_W2, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_W3, h_W3, size_W3, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_W4, h_W4, size_W4, cudaMemcpyHostToDevice));

    // Free host memory for weights
    CUDA_CHECK(cudaFreeHost(h_W1));
    CUDA_CHECK(cudaFreeHost(h_W2));
    CUDA_CHECK(cudaFreeHost(h_W3));
    CUDA_CHECK(cudaFreeHost(h_W4));

    // Create cuBLAS handle
    cublasHandle_t handle;
    cublasCreate(&handle);

    // Constants for cuBLAS operations
    const float alpha = 1.0f;
    const float beta = 0.0f;

    // Timing computation operations
    cudaEvent_t startEvent, stopEvent;
    float totalTime = 0;
    CUDA_CHECK(cudaEventCreate(&startEvent));
    CUDA_CHECK(cudaEventCreate(&stopEvent));
    CUDA_CHECK(cudaEventRecord(startEvent, 0));

    // Computation begins here (no cudaMemcpy or cudaMalloc calls from this point)

    // C1 = A1 * B1
    matMulOptimized(d_A1, d_B1, d_C1, M1, N1, K1);

    // C2 = C1 * B2
    matMulOptimized(d_C1, d_B2, d_C2, M1, N2, K2);

    // C3 = C2 * B3
    matMulOptimized(d_C2, d_B3, d_C3, M1, N3, K3);

    // Concatenate d_C3 (12 x 128) and d_D (12 x 512) to get d_E (12 x 640)
    dim3 blockSize_concat(16, 16);
    dim3 gridSize_concat((N3 + D_DIM + 15) / 16, (M1 + 15) / 16);
    concatenateMatrices << <gridSize_concat, blockSize_concat >> > (d_C3, d_D, d_E, M1, N3, D_DIM);
    CUDA_CHECK(cudaGetLastError());

    // Reshape d_E (12 x 640) to d_Vec (7680 x 1)
    // Since we're using row-major order, we can interpret the memory of d_E as a vector

    // Next, multiply d_W1 (1024 x 7680) with d_Vec (7680 x 1)
    // d_F1 = d_W1 * d_Vec
    cublasSgemv(handle, CUBLAS_OP_N, 1024, 7680, &alpha, d_W1, 1024, d_Vec, 1, &beta, d_F1, 1);

    // Multiply d_W2 (512 x 1024) with d_F1 (1024 x 1)
    // d_F2 = d_W2 * d_F1
    cublasSgemv(handle, CUBLAS_OP_N, 512, 1024, &alpha, d_W2, 512, d_F1, 1, &beta, d_F2, 1);

    // Multiply d_W3 (128 x 512) with d_F2 (512 x 1)
    // d_F3 = d_W3 * d_F2
    cublasSgemv(handle, CUBLAS_OP_N, 128, 512, &alpha, d_W3, 128, d_F2, 1, &beta, d_F3, 1);

    // Multiply d_W4 (12 x 128) with d_F3 (128 x 1)
    // d_F4 = d_W4 * d_F3
    cublasSgemv(handle, CUBLAS_OP_N, 12, 128, &alpha, d_W4, 12, d_F3, 1, &beta, d_F4, 1);

    CUDA_CHECK(cudaEventRecord(stopEvent, 0));
    CUDA_CHECK(cudaEventSynchronize(stopEvent));
    CUDA_CHECK(cudaEventElapsedTime(&totalTime, startEvent, stopEvent));
    std::cout << "Total computation execution time: " << totalTime << " ms" << std::endl;

    CUDA_CHECK(cudaEventDestroy(startEvent));
    CUDA_CHECK(cudaEventDestroy(stopEvent));

    // Optional: Copy final result back to host (if needed)
    // float h_F4[12];
    // CUDA_CHECK(cudaMemcpy(h_F4, d_F4, 12 * sizeof(float), cudaMemcpyDeviceToHost));

    // Clean up
    cublasDestroy(handle);
    CUDA_CHECK(cudaFree(d_A1));
    CUDA_CHECK(cudaFree(d_B1));
    CUDA_CHECK(cudaFree(d_B2));
    CUDA_CHECK(cudaFree(d_B3));
    CUDA_CHECK(cudaFree(d_D));

    CUDA_CHECK(cudaFree(d_C1));
    CUDA_CHECK(cudaFree(d_C2));
    CUDA_CHECK(cudaFree(d_C3));
    CUDA_CHECK(cudaFree(d_E));

    CUDA_CHECK(cudaFree(d_W1));
    CUDA_CHECK(cudaFree(d_W2));
    CUDA_CHECK(cudaFree(d_W3));
    CUDA_CHECK(cudaFree(d_W4));

    CUDA_CHECK(cudaFree(d_F1));
    CUDA_CHECK(cudaFree(d_F2));
    CUDA_CHECK(cudaFree(d_F3));
    CUDA_CHECK(cudaFree(d_F4));

    CUDA_CHECK(cudaFreeHost(h_A1));
    CUDA_CHECK(cudaFreeHost(h_B1));
    CUDA_CHECK(cudaFreeHost(h_B2));
    CUDA_CHECK(cudaFreeHost(h_B3));
    CUDA_CHECK(cudaFreeHost(h_D));

    return 0;
}