spectrogram.cu

// cuda_spectrogram.cu
#include <cuda_runtime.h>
#include <iostream>
#include <fstream>
#include <cmath>
#include <vector>
#include <complex>

#define BLOCK_SIZE 256
#define WINDOW_SIZE 512
#define HOP_SIZE 256  // 50% overlap

__device__ float hann(int n, int N) {
    return 0.5f * (1.0f - cosf(2.0f * M_PI * n / (N - 1)));
}

__global__ void stftKernel(const float* input, float* output, int numFrames, int frameSize, int hopSize) {
    int frame = blockIdx.x;
    int k = threadIdx.x;

    if (frame >= numFrames || k >= frameSize)
        return;

    int start = frame * hopSize;
    __shared__ float windowed[WINDOW_SIZE];

    // Windowed input
    if (start + k < numFrames * hopSize)
        windowed[k] = input[start + k] * hann(k, frameSize);
    else
        windowed[k] = 0.0f;

    __syncthreads();

    // Compute DFT (naive, slow but OK for small window)
    float real = 0.0f, imag = 0.0f;
    for (int n = 0; n < frameSize; n++) {
        float angle = 2.0f * M_PI * k * n / frameSize;
        real += windowed[n] * cosf(angle);
        imag -= windowed[n] * sinf(angle);
    }

    float mag = real * real + imag * imag;
    output[frame * frameSize + k] = mag; // power spectrum
}

void readRawAudio(const char* filename, std::vector<float>& data) {
    std::ifstream f(filename, std::ios::binary);
    if (!f) {
        std::cerr << "Error: Cannot open input audio\n";
        exit(1);
    }
    f.seekg(0, std::ios::end);
    size_t size = f.tellg() / sizeof(float);
    f.seekg(0);
    data.resize(size);
    f.read(reinterpret_cast<char*>(data.data()), size * sizeof(float));
    f.close();
}

void writeSpectrogram(const char* filename, const std::vector<float>& spec) {
    std::ofstream f(filename, std::ios::binary);
    f.write(reinterpret_cast<const char*>(spec.data()), spec.size() * sizeof(float));
    f.close();
}

int main(int argc, char** argv) {
    if (argc != 3) {
        std::cout << "Usage: ./spectrogram input.raw output.raw\n";
        return 1;
    }

    std::vector<float> audio;
    readRawAudio(argv[1], audio);
    int frameSize = WINDOW_SIZE;
    int hopSize = HOP_SIZE;
    int numFrames = (audio.size() - frameSize) / hopSize;

    size_t inputBytes = audio.size() * sizeof(float);
    size_t outputBytes = numFrames * frameSize * sizeof(float);

    float *d_input, *d_output;
    cudaMalloc(&d_input, inputBytes);
    cudaMalloc(&d_output, outputBytes);

    cudaMemcpy(d_input, audio.data(), inputBytes, cudaMemcpyHostToDevice);

    stftKernel<<<numFrames, frameSize>>>(d_input, d_output, numFrames, frameSize, hopSize);
    cudaDeviceSynchronize();

    std::vector<float> spectrogram(numFrames * frameSize);
    cudaMemcpy(spectrogram.data(), d_output, outputBytes, cudaMemcpyDeviceToHost);

    writeSpectrogram(argv[2], spectrogram);

    cudaFree(d_input);
    cudaFree(d_output);

    std::cout << "Spectrogram written to " << argv[2] << "\n";
    return 0;
}