mlp_cuda.cuh

#pragma once
#include <cuda_runtime.h>
#include <vector>

// The CUDA Kernel
__global__ void forward_kernel(const double* weights, const double* input, const double* biases, double* output, int inputSize, int outputSize) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    
    if (i < outputSize) {
        double sum = 0.0;
        for (int j = 0; j < inputSize; j++) {
            sum += weights[i * inputSize + j] * input[j];
        }
        output[i] = sum + biases[i];
    }
}

// The Supplementary Wrapper
inline std::vector<double> launch_cuda_forward(const std::vector<double>& flat_weights, const std::vector<double>& input, const std::vector<double>& biases) {
    int outputSize = biases.size();
    int inputSize = input.size();

    double *d_weights, *d_input, *d_biases, *d_output;

    cudaMalloc(&d_weights, outputSize * inputSize * sizeof(double));
    cudaMalloc(&d_input, inputSize * sizeof(double));
    cudaMalloc(&d_biases, outputSize * sizeof(double));
    cudaMalloc(&d_output, outputSize * sizeof(double));

    cudaMemcpy(d_weights, flat_weights.data(), outputSize * inputSize * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(d_input, input.data(), inputSize * sizeof(double), cudaMemcpyHostToDevice);
    cudaMemcpy(d_biases, biases.data(), outputSize * sizeof(double), cudaMemcpyHostToDevice);

    int blockSize = 256;
    int gridSize = (outputSize + blockSize - 1) / blockSize;
    forward_kernel<<<gridSize, blockSize>>>(d_weights, d_input, d_biases, d_output, inputSize, outputSize);
    
    cudaDeviceSynchronize();

    std::vector<double> output(outputSize);
    cudaMemcpy(output.data(), d_output, outputSize * sizeof(double), cudaMemcpyDeviceToHost);

    cudaFree(d_weights);
    cudaFree(d_input);
    cudaFree(d_biases);
    cudaFree(d_output);

    return output;
}