feat: implement mel-spectrogram with bindings

.gitignore

@@ -0,0 +1 @@
+py_example

build.toml

@@ -0,0 +1,16 @@
+[general]
+version = "0.0.1"
+
+[torch]
+name = "mel_spectrogram"
+src = [
+  "ext-torch/registration.h",
+  "ext-torch/torch_binding.cpp",
+  "ext-torch/torch_binding.h",
+]
+pysrc = ["ext-torch/__init__.py"]
+
+[kernel.mel_spectrogram]
+capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
+src = ["mel_spectrogram/mel_spectrogram.cu"]
+depends = ["torch"]

ext-torch/__init__.py

@@ -0,0 +1,15 @@
+import torch
+
+try:
+    from ._ops import ops
+except ImportError as e:
+    # Fallback for local development.
+    try:
+        import _mel_spectrogram
+        ops = torch.ops._mel_spectrogram
+    except ImportError:
+        raise e
+
+
+def mel_spectrogram(out: torch.Tensor, samples: torch.Tensor, filters: torch.Tensor, fft_size: int, fft_step: int) -> torch.Tensor:
+    return ops.mel_spectrogram(out, samples, filters, fft_size, fft_step)

ext-torch/registration.h

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <Python.h>
+
+#define _CONCAT(A, B) A##B
+#define CONCAT(A, B) _CONCAT(A, B)
+
+#define _STRINGIFY(A) #A
+#define STRINGIFY(A) _STRINGIFY(A)
+
+// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
+// could be a macro instead of a literal token.
+#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
+
+// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
+// could be a macro instead of a literal token.
+#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \
+  TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)
+
+// REGISTER_EXTENSION allows the shared library to be loaded and initialized
+// via python's import statement.
+#define REGISTER_EXTENSION(NAME)                                               \
+  PyMODINIT_FUNC CONCAT(PyInit_, NAME)() {                                     \
+    static struct PyModuleDef module = {PyModuleDef_HEAD_INIT,                 \
+                                        STRINGIFY(NAME), nullptr, 0, nullptr}; \
+    return PyModule_Create(&module);                                           \
+  }

ext-torch/torch_binding.cpp

@@ -0,0 +1,12 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("mel_spectrogram(Tensor out, Tensor samples, Tensor filters, int64_t fft_size, int64_t fft_step) -> Tensor", &mel_spectrogram);
+  ops.impl("mel_spectrogram", torch::kCUDA, &mel_spectrogram);
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

ext-torch/torch_binding.h

@@ -0,0 +1,10 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void mel_spectrogram(
+    torch::Tensor out,
+    torch::Tensor samples,
+    torch::Tensor filters,
+    int64_t fft_size,
+    int64_t fft_step);

mel_spectrogram/mel_spectrogram.cu

@@ -0,0 +1,163 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/all.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <cufft.h>
+#include <cuda_runtime.h>
+#include <device_launch_parameters.h>
+
+const float PI = 3.14159265358979323846f;
+const float EPS = 1e-10f;
+
+// small helper function to perform atomic max on floats
+__device__ __forceinline__ float atomicMaxFloat(float* address, float val) {
+    unsigned int* address_as_uint = (unsigned int*)address;
+    unsigned int old = *address_as_uint;
+    unsigned int assumed;
+    do {
+        assumed = old;
+        old = atomicCAS(address_as_uint, assumed,
+            __float_as_uint(fmaxf(val, __uint_as_float(assumed))));
+    } while (assumed != old);
+    return __uint_as_float(old);
+}
+
+
+__global__ void apply_hanning_window_kernel(
+    float* __restrict__ windowed_samples,
+    const float* __restrict__ samples,
+    const int fft_size,
+    const int n_samples,
+    const int fft_step) {
+    
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+    
+    // max FFT size of 1024
+    __shared__ float hann[1024];
+    if (threadIdx.x < fft_size) {
+        hann[threadIdx.x] = 0.5f * (1.0f - cosf((2.0f * PI * threadIdx.x) / (float)fft_size));
+    }
+    __syncthreads();
+
+    for (int i = idx; i < n_samples; i += stride) {
+        int segment = i / fft_step;
+        int offset = i % fft_size;
+        if (offset < fft_size && (segment * fft_step + offset) < n_samples) {
+            windowed_samples[segment * fft_size + offset] = 
+                samples[segment * fft_step + offset] * hann[offset];
+        }
+    }
+}
+
+__global__ void compute_mel_energies_kernel(
+    float* __restrict__ mel_spec,
+    const float* __restrict__ fft_magnitudes,
+    const float* __restrict__ mel_filters,
+    const int n_mel,
+    const int n_fft,
+    const int n_frames) {
+    
+    int mel_idx = blockIdx.y;
+    int frame_idx = blockIdx.x * blockDim.x + threadIdx.x;
+    
+    if (mel_idx < n_mel && frame_idx < n_frames) {
+        float sum = 0.0f;
+        
+        // sum up the FFT magnitudes for each mel filter
+        for (int k = 0; k < n_fft; k++) {
+            sum += fft_magnitudes[frame_idx * n_fft + k] * 
+                   mel_filters[mel_idx * n_fft + k];
+        }
+        
+        sum = log10f(fmaxf(sum, EPS));
+        mel_spec[mel_idx * n_frames + frame_idx] = sum;
+    }
+}
+
+__global__ void normalize_mel_spec_kernel(
+    float* __restrict__ mel_spec,
+    const int n_mel,
+    const int n_frames) {
+    
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    int stride = blockDim.x * gridDim.x;
+    
+    __shared__ float block_max;
+    if (threadIdx.x == 0) {
+        block_max = -INFINITY;
+    }
+    __syncthreads();
+    
+    for (int i = idx; i < n_mel * n_frames; i += stride) {
+        atomicMaxFloat(&block_max, mel_spec[i]);
+    }
+    __syncthreads();
+    
+    float mmax = block_max - 8.0f;
+    for (int i = idx; i < n_mel * n_frames; i += stride) {
+        mel_spec[i] = fmaxf(mel_spec[i], mmax) / 4.0f + 1.0f;
+    }
+}
+
+void mel_spectrogram(
+    torch::Tensor& out,        // [n_mel, n_frames]
+    torch::Tensor& samples,    // [n_samples]
+    torch::Tensor& filters,    // [n_mel, n_fft]
+    const int fft_size,
+    const int fft_step) {
+    
+    TORCH_CHECK(samples.is_cuda(), "Input samples must be a CUDA tensor");
+    TORCH_CHECK(filters.is_cuda(), "Mel filters must be a CUDA tensor");
+    TORCH_CHECK(samples.scalar_type() == torch::kFloat32, "Input samples must be float32");
+    TORCH_CHECK(filters.scalar_type() == torch::kFloat32, "Mel filters must be float32");
+    
+    const int n_samples = samples.size(0);
+    const int n_frames = n_samples / fft_step;
+    const int n_mel = filters.size(0);
+    const int n_fft = 1 + fft_size / 2;
+    
+    // set up tensors
+    auto windowed = torch::empty({n_frames, fft_size}, 
+        torch::dtype(torch::kFloat32).device(samples.device()));
+    auto fft_out = torch::empty({n_frames, n_fft}, 
+        torch::dtype(torch::kFloat32).device(samples.device()));
+    
+    // set up kernel parameters
+    const int threads = 256;
+    dim3 block(threads);
+    dim3 grid((n_samples + threads - 1) / threads);
+    dim3 mel_grid(n_frames / threads + 1, n_mel);
+    
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(samples));
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    
+    // apply window function
+    apply_hanning_window_kernel<<<grid, block, 0, stream>>>(
+        windowed.data_ptr<float>(),
+        samples.data_ptr<float>(),
+        fft_size,
+        n_samples,
+        fft_step);
+    
+    cufftHandle plan;
+    cufftPlan1d(&plan, fft_size, CUFFT_R2C, n_frames);
+    cufftSetStream(plan, stream);
+    cufftExecR2C(plan, windowed.data_ptr<float>(), (cufftComplex*)fft_out.data_ptr<float>());
+    
+    // compute mel energies
+    compute_mel_energies_kernel<<<mel_grid, block, 0, stream>>>(
+        out.data_ptr<float>(),
+        fft_out.data_ptr<float>(),
+        filters.data_ptr<float>(),
+        n_mel,
+        n_fft,
+        n_frames);
+    
+    // normalize
+    normalize_mel_spec_kernel<<<grid, block, 0, stream>>>(
+        out.data_ptr<float>(),
+        n_mel,
+        n_frames);
+    
+    cufftDestroy(plan);
+}