whisper : add CUDA-specific computation mel spectrograms (#2206)

* whisper : use polymorphic class to calculate mel spectrogram

* whisper : add cuda-specific mel spectrogram calculation

* whisper : conditionally compile cufftGetErrorString to avoid warnings

* build : add new files to makefile

* ruby : add new files to conf script

* build : fix typo in makefile

* whisper : suppress cub warning for deprecated C++ std in whisper-mel-cuda

CMakeLists.txt

@@ -364,12 +364,12 @@ if (WHISPER_CUDA)
         if (WHISPER_STATIC)
             if (WIN32)
                 # As of 12.3.1 CUDA Tookit for Windows does not offer a static cublas library
-                set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart_static CUDA::cublas CUDA::cublasLt)
+                set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart_static CUDA::cublas CUDA::cublasLt CUDA::cufft)
             else ()
-                set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart_static CUDA::cublas_static CUDA::cublasLt_static)
+                set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart_static CUDA::cublas_static CUDA::cublasLt_static CUDA::cufft_static)
             endif()
         else()
-            set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart CUDA::cublas CUDA::cublasLt)
+            set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cudart CUDA::cublas CUDA::cublasLt CUDA::cufft)
         endif()
 
         set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS} CUDA::cuda_driver)
@@ -679,6 +679,10 @@ add_library(${TARGET}
     whisper.cpp
     )
 
+if (WHISPER_CUDA)
+    target_sources(${TARGET} PRIVATE whisper-mel-cuda.cu)
+endif()
+
 include_directories (
     .
 )

Makefile

@@ -286,8 +286,8 @@ ifdef WHISPER_CUDA
 
 	CFLAGS      += -DGGML_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/$(UNAME_M)-linux/include
 	CXXFLAGS    += -DGGML_USE_CUDA -I/usr/local/cuda/include -I/opt/cuda/include -I$(CUDA_PATH)/targets/$(UNAME_M)-linux/include
-	LDFLAGS     += -lcuda -lcublas -lculibos -lcudart -lcublasLt -lpthread -ldl -lrt -L/usr/local/cuda/lib64 -L/opt/cuda/lib64 -L$(CUDA_PATH)/targets/$(UNAME_M)-linux/lib -L/usr/lib/wsl/lib
-	WHISPER_OBJ += ggml-cuda.o
+	LDFLAGS     += -lcuda -lcublas -lculibos -lcudart -lcublasLt -lcufft -lpthread -ldl -lrt -L/usr/local/cuda/lib64 -L/opt/cuda/lib64 -L$(CUDA_PATH)/targets/$(UNAME_M)-linux/lib -L/usr/lib/wsl/lib
+	WHISPER_OBJ += ggml-cuda.o whisper-mel-cuda.o
 	WHISPER_OBJ += $(patsubst %.cu,%.o,$(wildcard ggml-cuda/*.cu))
 	NVCC        = nvcc
 	NVCCFLAGS   = --forward-unknown-to-host-compiler -arch=$(CUDA_ARCH_FLAG)
@@ -299,6 +299,9 @@ ggml-cuda.o: ggml-cuda.cu ggml-cuda.h ggml.h ggml-backend.h ggml-backend-impl.h
 	$(NVCC) $(NVCCFLAGS) $(CXXFLAGS) -Wno-pedantic -c $< -o $@
 endif
 
+whisper-mel-cuda.o: whisper-mel-cuda.cu whisper.h ggml.h ggml-backend.h whisper-mel.hpp whisper-mel-cuda.hpp
+	$(NVCC) $(NVCCFLAGS) $(CXXFLAGS) -Wno-pedantic -c $< -o $@
+
 ifdef WHISPER_HIPBLAS
 	ROCM_PATH   ?= /opt/rocm
 	HIPCC       ?= $(ROCM_PATH)/bin/hipcc
@@ -404,7 +407,7 @@ ggml-quants.o: ggml-quants.c ggml.h ggml-quants.h
 
 WHISPER_OBJ += ggml.o ggml-alloc.o ggml-backend.o ggml-quants.o
 
-whisper.o: whisper.cpp whisper.h ggml.h ggml-cuda.h
+whisper.o: whisper.cpp whisper.h whisper-mel.hpp ggml.h ggml-cuda.h
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 
 ifndef WHISPER_COREML

bindings/ruby/ext/extconf.rb

@@ -1,6 +1,7 @@
 require 'mkmf'
 system("cp #{File.join(File.dirname(__FILE__),'..','..','..','whisper.cpp')} .")
 system("cp #{File.join(File.dirname(__FILE__),'..','..','..','whisper.h')} .")
+system("cp #{File.join(File.dirname(__FILE__),'..','..','..','whisper-mel.hpp')} .")
 system("cp #{File.join(File.dirname(__FILE__),'..','..','..','ggml.h')} .")
 system("cp #{File.join(File.dirname(__FILE__),'..','..','..','ggml.c')} .")
 system("cp #{File.join(File.dirname(__FILE__),'..','..','..','ggml-impl.h')} .")

whisper-mel-cuda.cu

@@ -0,0 +1,342 @@
+#define CUB_IGNORE_DEPRECATED_CPP_DIALECT
+#include "whisper-mel-cuda.hpp"
+#include "whisper.h"
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cufft.h>
+#include <cublas_v2.h>
+#include <cuComplex.h>
+#include <cub/device/device_reduce.cuh>
+
+#include <algorithm>
+
+#if defined(_MSC_VER)
+#pragma warning(disable: 4324) // added padding
+#endif
+
+#ifndef NDEBUG
+#   define DO_CHECKS 1
+#else
+#   define DO_CHECKS 0
+#endif
+
+namespace {
+
+#if DO_CHECKS
+const char* cufftGetErrorString(cufftResult_t res) {
+    switch (res) {
+    case CUFFT_SUCCESS: return "The cuFFT operation was successful";
+    case CUFFT_INVALID_PLAN: return "cuFFT was passed an invalid plan handle";
+    case CUFFT_ALLOC_FAILED: return "cuFFT failed to allocate GPU or CPU memory";
+    case CUFFT_INVALID_TYPE: return "No longer used";
+    case CUFFT_INVALID_VALUE: return "User specified an invalid pointer or parameter";
+    case CUFFT_INTERNAL_ERROR: return "Driver or internal cuFFT library error";
+    case CUFFT_EXEC_FAILED: return "Failed to execute an FFT on the GPU";
+    case CUFFT_SETUP_FAILED: return "The cuFFT library failed to initialize";
+    case CUFFT_INVALID_SIZE: return "User specified an invalid transform size";
+    case CUFFT_UNALIGNED_DATA: return "No longer used";
+    case CUFFT_INCOMPLETE_PARAMETER_LIST: return "Missing parameters in call";
+    case CUFFT_INVALID_DEVICE: return "Execution of a plan was on different GPU than plan creation";
+    case CUFFT_PARSE_ERROR: return "Internal plan database error";
+    case CUFFT_NO_WORKSPACE: return "No workspace has been provided prior to plan execution";
+    case CUFFT_NOT_IMPLEMENTED: return "Function does not implement functionality for parameters given.";
+    case CUFFT_LICENSE_ERROR: return "Used in previous versions.";
+    case CUFFT_NOT_SUPPORTED: return "Operation is not supported for parameters given.";
+    default: return "Unknown error";
+    }
+}
+
+#   define CUDA_CHECK_GEN(err, success, error_fn)                                            \
+         do {                                                                                \
+            auto err_ = (err);                                                               \
+            if (err_ != (success)) {                                                         \
+                fprintf(stderr, "%s %s:%d - %s\n", #err, __FILE__, __LINE__, error_fn(err_)); \
+            }                                                                                \
+        } while (0)
+#else
+#   define CUDA_CHECK_GEN(err, success, error_fn) err
+#endif
+
+#define CUDA_CHECK(err) CUDA_CHECK_GEN(err, cudaSuccess, cudaGetErrorString)
+#define CUBLAS_CHECK(err) CUDA_CHECK_GEN(err, CUBLAS_STATUS_SUCCESS, cublasGetStatusString)
+#define CUFFT_CHECK(err) CUDA_CHECK_GEN(err, CUFFT_SUCCESS, cufftGetErrorString)
+
+__global__ void k_fill_stft_input(
+    const float * padded_samples,
+    const int n_frames,
+    const float * hann_window,
+    float * stft_in
+) {
+    auto y = blockIdx.y * blockDim.y + threadIdx.y;
+    // if (y >= n_frames) return;
+    auto x = blockIdx.x * blockDim.x + threadIdx.x;
+    // if (x >= WHISPER_N_FFT) return;
+
+    auto line = padded_samples + y * WHISPER_HOP_LENGTH;
+    auto outLine = stft_in + y * WHISPER_N_FFT;
+
+    outLine[x] = line[x] * hann_window[x];
+}
+
+__global__ void k_calc_magnitudes(
+    const cuComplex* stft_out,
+    const int n_frames,
+    float * magnitudes
+) {
+    auto y = blockIdx.y * blockDim.y + threadIdx.y;
+    // if (y >= n_frames) return;
+    auto x = blockIdx.x * blockDim.x + threadIdx.x;
+    // if (x >= WHISPER_N_FFT_HALF) return;
+
+    auto idx = y * WHISPER_N_FFT_HALF + x;
+
+    auto r = stft_out[idx].x;
+    auto i = stft_out[idx].y;
+    magnitudes[idx] = r * r + i * i;
+}
+
+__global__ void k_calc_log_mel(
+    const float * mel_data,
+    const int n_mel,
+    const float * max_val,
+    float * log_mel
+) {
+    auto x = blockIdx.x * blockDim.x + threadIdx.x;
+    if (x >= n_mel) return;
+
+    float val = mel_data[x];
+
+    constexpr float e = 1e-10f;
+    if (val < e) val = e;
+
+    val = log10(val);
+
+    const float max = log10(*max_val) - 8.f;
+    if (val < max) val = max;
+
+    log_mel[x] = (val + 4) / 4;
+}
+
+void fill_stft_input(
+    const float * padded_samples,
+    int n_frames,
+    const float * hann_window,
+    float * stft_in,
+    cudaStream_t stream
+) {
+    dim3 block(WHISPER_N_FFT, 1);
+    dim3 grid(1, n_frames);
+
+    k_fill_stft_input<<<grid, block, 0, stream>>>(padded_samples, n_frames, hann_window, stft_in);
+}
+
+void calc_magnitudes(
+    const cuComplex* stft_out,
+    int n_frames,
+    float * magnitudes,
+    cudaStream_t stream
+) {
+    dim3 block(WHISPER_N_FFT_HALF, 1);
+    dim3 grid(1, n_frames);
+    k_calc_magnitudes<<<grid, block, 0, stream>>>(stft_out, n_frames, magnitudes);
+}
+
+constexpr auto LOG_MEL_PREFIX_SIZE = 256;
+
+size_t get_log_mel_temp_storage_size() {
+    constexpr auto maxPaddedSamples = 2 * WHISPER_N_SAMPLES + WHISPER_N_FFT;
+    constexpr auto maxFrames = 1 + (maxPaddedSamples - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+    constexpr auto maxMels = 160;
+
+    size_t nbytes = 0;
+    float * temp = nullptr;
+    cub::DeviceReduce::Max(nullptr, nbytes, temp, temp, maxFrames * maxMels);
+    return nbytes + LOG_MEL_PREFIX_SIZE;
+}
+
+void calc_log_mel(
+    const float * mel_data,
+    int n_mel,
+    void * tempStorage,
+    int tempStorageSize,
+    float * log_mel,
+    cudaStream_t stream
+) {
+    float * max_val = reinterpret_cast<float *>(tempStorage);
+    void * maxTemp = reinterpret_cast<char*>(tempStorage) + LOG_MEL_PREFIX_SIZE;
+
+    size_t nbytes = size_t(tempStorageSize - LOG_MEL_PREFIX_SIZE);
+    cub::DeviceReduce::Max(maxTemp, nbytes, mel_data, max_val, n_mel, stream);
+
+    int block = 256;
+    int grid = (n_mel + block - 1) / block;
+
+    k_calc_log_mel<<<grid, block, 0, stream>>>(mel_data, n_mel, max_val, log_mel);
+}
+
+class mel_calc_cuda : public whisper_mel_calc {
+    const int m_n_mel;
+
+    ggml_backend_t m_backend = nullptr;
+
+    cudaStream_t m_stream = nullptr;
+    cublasHandle_t m_cublas_handle = nullptr;
+
+    float * m_hann_window = nullptr;
+
+    size_t m_cufft_workspace_size = 0;
+    void * m_cufft_workspace = nullptr;
+
+    float * m_filters = nullptr;
+
+    size_t m_log_mel_temp_storage_size = 0;
+    void * m_log_mel_temp_storage = nullptr;
+public:
+    mel_calc_cuda(ggml_backend_t backend, const whisper_filters& filters)
+        : m_n_mel(filters.n_mel)
+        , m_backend(backend)
+    {
+        if (filters.n_fft != WHISPER_N_FFT_HALF) {
+            throw std::invalid_argument("MelFilters n_frames must be WHISPER_N_FFT_HALF");
+        }
+        assert(filters.data.size() == filters.n_mel * WHISPER_N_FFT_HALF);
+
+        CUDA_CHECK(cudaStreamCreate(&m_stream));
+        CUBLAS_CHECK(cublasCreate(&m_cublas_handle));
+        CUBLAS_CHECK(cublasSetMathMode(m_cublas_handle, CUBLAS_TF32_TENSOR_OP_MATH));
+        CUBLAS_CHECK(cublasSetStream(m_cublas_handle, m_stream));
+
+        // create Hann window
+        {
+            auto hw = whisper_mel_calc::hann_window();
+            CUDA_CHECK(cudaMallocAsync(&m_hann_window, hw.len * sizeof(float), m_stream));
+            CUDA_CHECK(cudaMemcpyAsync(m_hann_window, hw.data, hw.len * sizeof(float), cudaMemcpyHostToDevice, m_stream));
+        }
+
+        // create working area
+        {
+            constexpr auto maxPaddedSamples = 2 * WHISPER_N_SAMPLES + WHISPER_N_FFT;
+            constexpr auto maxFrames = 1 + (maxPaddedSamples - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+            CUFFT_CHECK(cufftEstimate1d(WHISPER_N_FFT, CUFFT_R2C, maxFrames, &m_cufft_workspace_size));
+            CUDA_CHECK(cudaMallocAsync(&m_cufft_workspace, m_cufft_workspace_size, m_stream));
+        }
+
+        // fill filters
+        {
+            auto& f = filters.data;
+            CUDA_CHECK(cudaMallocAsync(&m_filters, f.size() * sizeof(float), m_stream));
+            CUDA_CHECK(cudaMemcpyAsync(m_filters, f.data(), f.size() * sizeof(float), cudaMemcpyHostToDevice, m_stream));
+        }
+
+        {
+            m_log_mel_temp_storage_size = get_log_mel_temp_storage_size();
+            CUDA_CHECK(cudaMallocAsync(&m_log_mel_temp_storage, m_log_mel_temp_storage_size, m_stream));
+        }
+    }
+
+    ~mel_calc_cuda() {
+        CUDA_CHECK(cudaStreamSynchronize(m_stream));
+        CUDA_CHECK(cudaStreamDestroy(m_stream));
+        CUDA_CHECK(cudaFree(m_hann_window));
+        CUDA_CHECK(cudaFree(m_cufft_workspace));
+        CUDA_CHECK(cudaFree(m_filters));
+        CUDA_CHECK(cudaFree(m_log_mel_temp_storage));
+    }
+
+    virtual whisper_mel calculate(whisper_span<const float> samples, int /*n_threads*/) const override {
+        const size_t mirror_pad = WHISPER_N_FFT / 2;
+        const size_t padded_size = samples.len + WHISPER_N_SAMPLES + WHISPER_N_FFT;
+
+        // pad
+        std::vector<float> padded_samples(padded_size);
+        std::reverse_copy(samples.data + 1, samples.data + 1 + mirror_pad, padded_samples.begin()); // reflect
+        std::copy(samples.data, samples.data + samples.len, padded_samples.begin() + mirror_pad); // copy
+
+        // fill the rest of the data
+        // it should canonically be mirrored at the end as well,
+        // but we just assume the last MEL_FRAME_SIZE/2 samples are zeros
+        std::fill(padded_samples.begin() + mirror_pad + samples.len, padded_samples.end(), 0.f);
+
+        const auto n_frames = 1 + (padded_samples.size() - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+
+        float * cu_padded_samples = nullptr;
+        CUDA_CHECK(cudaMallocAsync(&cu_padded_samples, padded_samples.size() * sizeof(float), m_stream));
+        CUDA_CHECK(cudaMemcpyAsync(cu_padded_samples, padded_samples.data(), padded_samples.size() * sizeof(float), cudaMemcpyHostToDevice, m_stream));
+
+        float * stft_in = nullptr; // contiguous buffer for stft input
+        CUDA_CHECK(cudaMallocAsync(&stft_in, n_frames * WHISPER_N_FFT * sizeof(float), m_stream));
+
+        fill_stft_input(cu_padded_samples, int(n_frames), m_hann_window, stft_in, m_stream);
+
+        cufftComplex* stft_out;
+        CUDA_CHECK(cudaMallocAsync(&stft_out, n_frames * WHISPER_N_FFT_HALF * sizeof(cufftComplex), m_stream));
+
+        cufftHandle plan;
+        CUFFT_CHECK(cufftCreate(&plan));
+        CUFFT_CHECK(cufftSetAutoAllocation(plan, 0));
+        {
+            size_t waSize;
+            CUFFT_CHECK(cufftMakePlan1d(plan, WHISPER_N_FFT, CUFFT_R2C, int(n_frames), &waSize));
+            assert(waSize <= m_cufft_workspace_size);
+            CUFFT_CHECK(cufftSetWorkArea(plan, m_cufft_workspace));
+            CUFFT_CHECK(cufftSetStream(plan, m_stream));
+        }
+        CUFFT_CHECK(cufftExecR2C(plan, stft_in, stft_out));
+
+        const auto n_mag_frames = n_frames - 1; // drop last frame
+        float * magnitudes;
+        CUDA_CHECK(cudaMallocAsync(&magnitudes, n_mag_frames * WHISPER_N_FFT_HALF * sizeof(float), m_stream));
+        calc_magnitudes(stft_out, int(n_mag_frames), magnitudes, m_stream);
+
+        float * mel_data = nullptr;
+        CUDA_CHECK(cudaMallocAsync(&mel_data, m_n_mel * n_mag_frames * sizeof(float), m_stream));
+
+        const float fone = 1.0f, fzero = 0.0f;
+        CUBLAS_CHECK(cublasSgemm(m_cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N,
+            int(n_mag_frames), m_n_mel, WHISPER_N_FFT_HALF,
+            &fone,
+            magnitudes, WHISPER_N_FFT_HALF,
+            m_filters, WHISPER_N_FFT_HALF,
+            &fzero,
+            mel_data, int(n_mag_frames)));
+
+        float * log_mels = nullptr;
+        CUDA_CHECK(cudaMallocAsync(&log_mels, m_n_mel * n_mag_frames * sizeof(float), m_stream));
+
+        calc_log_mel(
+            mel_data, int(m_n_mel * n_mag_frames),
+            m_log_mel_temp_storage, int(m_log_mel_temp_storage_size),
+            log_mels, m_stream);
+
+        whisper_mel ret;
+        ret.n_mel = m_n_mel;
+        ret.n_len = int(n_mag_frames);
+        // Calculate semi-padded sample length to ensure compatibility
+        ret.n_len_org = 1 + int(samples.len + mirror_pad - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+        ret.data.resize(m_n_mel * n_mag_frames);
+        CUDA_CHECK(cudaMemcpyAsync(ret.data.data(), log_mels, ret.data.size() * sizeof(float), cudaMemcpyDeviceToHost, m_stream));
+
+        CUDA_CHECK(cudaStreamSynchronize(m_stream));
+
+        // cleanup
+        CUFFT_CHECK(cufftDestroy(plan));
+        CUDA_CHECK(cudaFreeAsync(log_mels, m_stream));
+        CUDA_CHECK(cudaFreeAsync(mel_data, m_stream));
+        CUDA_CHECK(cudaFreeAsync(magnitudes, m_stream));
+        CUDA_CHECK(cudaFreeAsync(stft_out, m_stream));
+        CUDA_CHECK(cudaFreeAsync(stft_in, m_stream));
+        CUDA_CHECK(cudaFreeAsync(cu_padded_samples, m_stream));
+
+        return ret;
+    }
+};
+
+}
+
+whisper_mel_calc * whisper_mel_calc_create_cuda(ggml_backend_t backend, const whisper_filters & filters) {
+    if (filters.n_fft != WHISPER_N_FFT_HALF) {
+        return nullptr;
+    }
+    return new mel_calc_cuda(backend, filters);
+}

whisper-mel-cuda.hpp

@@ -0,0 +1,3 @@
+#include "whisper-mel.hpp"
+
+whisper_mel_calc * whisper_mel_calc_create_cuda(ggml_backend_t backend, const whisper_filters & filters);

whisper-mel.hpp

@@ -0,0 +1,33 @@
+#pragma once
+#include "ggml-backend.h"
+#include <vector>
+
+struct whisper_mel {
+    int n_len;
+    int n_len_org;
+    int n_mel;
+
+    std::vector<float> data;
+};
+
+struct whisper_filters {
+    int32_t n_mel;
+    int32_t n_fft;
+
+    std::vector<float> data;
+};
+
+template <typename T>
+struct whisper_span {
+    T * data;
+    int len;
+};
+
+struct whisper_mel_calc {
+    virtual ~whisper_mel_calc();
+    virtual whisper_mel calculate(whisper_span<const float> samples, int n_threads) const = 0;
+    static whisper_span<const float> hann_window();
+};
+
+// returns a new pointer which needs to be freed with delete
+whisper_mel_calc * whisper_mel_calc_create(ggml_backend_t backend, const whisper_filters & filters);

whisper.cpp

@@ -10,6 +10,7 @@
 
 #ifdef GGML_USE_CUDA
 #include "ggml-cuda.h"
+#include "whisper-mel-cuda.hpp"
 #endif
 
 #ifdef GGML_USE_SYCL
@@ -24,6 +25,8 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 
+#include "whisper-mel.hpp"
+
 #include <atomic>
 #include <algorithm>
 #include <cassert>
@@ -380,21 +383,6 @@ static const std::map<whisper_alignment_heads_preset, whisper_aheads> g_aheads {
 
 static std::vector<uint32_t> get_alignment_heads_by_layer(const whisper_context_params & cparams, int il, int32_t n_text_layer, int32_t n_head);
 
-struct whisper_mel {
-    int n_len;
-    int n_len_org;
-    int n_mel;
-
-    std::vector<float> data;
-};
-
-struct whisper_filters {
-    int32_t n_mel;
-    int32_t n_fft;
-
-    std::vector<float> data;
-};
-
 struct whisper_vocab {
     using id    = int32_t;
     using token = std::string;
@@ -883,6 +871,8 @@ struct whisper_context {
     whisper_model model;
     whisper_vocab vocab;
 
+    whisper_mel_calc * mel_calc = nullptr;
+
     whisper_state * state = nullptr;
 
     ggml_backend_t backend = nullptr;
@@ -2894,6 +2884,14 @@ struct whisper_global_cache {
 } global_cache;
 }
 
+// Mel spectrogram
+
+whisper_mel_calc::~whisper_mel_calc() = default; // export vtable
+
+whisper_span<const float> whisper_mel_calc::hann_window() {
+    return {global_cache.hann_window, WHISPER_N_FFT};
+}
+
 // naive Discrete Fourier Transform
 // input is real-valued
 // output is complex-valued
@@ -2976,8 +2974,10 @@ static void fft(const std::vector<float> & in, std::vector<float> & out) {
 }
 
 static void log_mel_spectrogram_worker_thread(int ith, const float * hann, const std::vector<float> & samples,
-                                              int n_samples, int frame_size, int frame_step, int n_threads,
+                                              int n_samples, int n_threads,
                                               const whisper_filters & filters, whisper_mel & mel) {
+    const auto frame_size = WHISPER_N_FFT;
+    const auto frame_step = WHISPER_HOP_LENGTH;
     std::vector<float> fft_in(frame_size, 0.0);
     std::vector<float> fft_out(2 * frame_size);
     int n_fft = filters.n_fft;
@@ -3041,99 +3041,95 @@ static void log_mel_spectrogram_worker_thread(int ith, const float * hann, const
         }
     }
 }
+namespace {
+struct mel_calc_cpu : public whisper_mel_calc {
+    const whisper_filters& m_filters;
+    mel_calc_cpu(const whisper_filters & filters) : m_filters(filters) {}
 
-// ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L110-L157
-static bool log_mel_spectrogram(
-              whisper_state & wstate,
-              const float * samples,
-              const int   n_samples,
-              const int   /*sample_rate*/,
-              const int   frame_size,
-              const int   frame_step,
-              const int   n_mel,
-              const int   n_threads,
-              const whisper_filters & filters,
-              const bool   debug,
-              whisper_mel & mel) {
-    const int64_t t_start_us = ggml_time_us();
+    // ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L110-L157
+    whisper_mel calculate(whisper_span<const float> ssamples, int n_threads) const override {
+        // Hann window
+        const float * hann = global_cache.hann_window;
 
-    // Hann window
-    WHISPER_ASSERT(frame_size == WHISPER_N_FFT && "Unsupported frame_size");
-    const float * hann = global_cache.hann_window;
+        // Calculate the length of padding
+        int64_t stage_1_pad = WHISPER_SAMPLE_RATE * 30;
+        int64_t stage_2_pad = WHISPER_N_FFT / 2;
 
-    // Calculate the length of padding
-    int64_t stage_1_pad = WHISPER_SAMPLE_RATE * 30;
-    int64_t stage_2_pad = frame_size / 2;
+        const int n_samples = int(ssamples.len);
+        const float * samples = ssamples.data;
 
-    // Initialize a vector and copy data from C array to it.
-    std::vector<float> samples_padded;
-    samples_padded.resize(n_samples + stage_1_pad + stage_2_pad * 2);
-    std::copy(samples, samples + n_samples, samples_padded.begin() + stage_2_pad);
+        // Initialize a vector and copy data from C array to it.
+        std::vector<float> samples_padded;
+        samples_padded.resize(n_samples + stage_1_pad + stage_2_pad * 2);
+        std::copy(samples, samples + n_samples, samples_padded.begin() + stage_2_pad);
 
-    // pad 30 seconds of zeros at the end of audio (480,000 samples) + reflective pad 200 samples at the end of audio
-    std::fill(samples_padded.begin() + n_samples + stage_2_pad, samples_padded.begin() + n_samples + stage_1_pad + 2 * stage_2_pad, 0);
+        // pad 30 seconds of zeros at the end of audio (480,000 samples) + reflective pad 200 samples at the end of audio
+        std::fill(samples_padded.begin() + n_samples + stage_2_pad, samples_padded.begin() + n_samples + stage_1_pad + 2 * stage_2_pad, 0);
 
-    // reflective pad 200 samples at the beginning of audio
-    std::reverse_copy(samples + 1, samples + 1 + stage_2_pad, samples_padded.begin());
+        // reflective pad 200 samples at the beginning of audio
+        std::reverse_copy(samples + 1, samples + 1 + stage_2_pad, samples_padded.begin());
 
-    mel.n_mel     = n_mel;
-    // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/SpectralOps.cpp#L936
-    // Calculate number of frames + remove the last frame
-    mel.n_len     = (samples_padded.size() - frame_size) / frame_step;
-    // Calculate semi-padded sample length to ensure compatibility
-    mel.n_len_org = 1 + (n_samples + stage_2_pad - frame_size) / frame_step;
-    mel.data.resize(mel.n_mel * mel.n_len);
+        whisper_mel mel;
+        mel.n_mel     = m_filters.n_mel;
+        // https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/SpectralOps.cpp#L936
+        // Calculate number of frames + remove the last frame
+        mel.n_len     = (samples_padded.size() - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+        // Calculate semi-padded sample length to ensure compatibility
+        mel.n_len_org = 1 + (n_samples + stage_2_pad - WHISPER_N_FFT) / WHISPER_HOP_LENGTH;
+        mel.data.resize(mel.n_mel * mel.n_len);
 
 
-    {
-        std::vector<std::thread> workers(n_threads - 1);
-        for (int iw = 0; iw < n_threads - 1; ++iw) {
-            workers[iw] = std::thread(
-                    log_mel_spectrogram_worker_thread, iw + 1, hann, samples_padded,
-                    n_samples + stage_2_pad, frame_size, frame_step, n_threads,
-                    std::cref(filters), std::ref(mel));
-        }
-
-        // main thread
-        log_mel_spectrogram_worker_thread(0, hann, samples_padded, n_samples + stage_2_pad, frame_size, frame_step, n_threads, filters, mel);
+        {
+            std::vector<std::thread> workers(n_threads - 1);
+            for (int iw = 0; iw < n_threads - 1; ++iw) {
+                workers[iw] = std::thread(
+                        log_mel_spectrogram_worker_thread, iw + 1, hann, samples_padded,
+                        n_samples + stage_2_pad, n_threads,
+                        std::cref(m_filters), std::ref(mel));
+            }
 
-        for (int iw = 0; iw < n_threads - 1; ++iw) {
-            workers[iw].join();
-        }
-    }
+            // main thread
+            log_mel_spectrogram_worker_thread(0, hann, samples_padded, n_samples + stage_2_pad, n_threads, m_filters, mel);
 
-    // clamping and normalization
-    double mmax = -1e20;
-    for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
-        if (mel.data[i] > mmax) {
-            mmax = mel.data[i];
+            for (int iw = 0; iw < n_threads - 1; ++iw) {
+                workers[iw].join();
+            }
         }
-    }
-
-    mmax -= 8.0;
 
-    for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
-        if (mel.data[i] < mmax) {
-            mel.data[i] = mmax;
+        // clamping and normalization
+        double mmax = -1e20;
+        for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
+            if (mel.data[i] > mmax) {
+                mmax = mel.data[i];
+            }
         }
 
-        mel.data[i] = (mel.data[i] + 4.0)/4.0;
-    }
+        mmax -= 8.0;
 
-    wstate.t_mel_us += ggml_time_us() - t_start_us;
+        for (int i = 0; i < mel.n_mel*mel.n_len; i++) {
+            if (mel.data[i] < mmax) {
+                mel.data[i] = mmax;
+            }
 
-    // Dump log_mel_spectrogram
-    if (debug) {
-        std::ofstream outFile("log_mel_spectrogram.json");
-        outFile << "[";
-        for (uint64_t i = 0; i < mel.data.size() - 1; i++) {
-            outFile << mel.data[i] << ", ";
+            mel.data[i] = (mel.data[i] + 4.0)/4.0;
         }
-        outFile << mel.data[mel.data.size() - 1] << "]";
-        outFile.close();
+
+        return mel;
     }
+};
+}
 
-    return true;
+whisper_mel_calc * whisper_mel_calc_create(ggml_backend_t backend, const whisper_filters & filters) {
+#if GGML_USE_CUDA
+    if (ggml_backend_is_cuda(backend)) {
+        auto ret = whisper_mel_calc_create_cuda(backend, filters);
+        // run a warmup to avoid the first kernel launch overhead (thus we get the best perf even on the first run)
+        const float warmup[256] = {0};
+        ret->calculate({warmup, 256}, 1);
+        return ret;
+    } else
+#endif
+        return new mel_calc_cpu(filters);
 }
 
 // split text into tokens
@@ -3593,6 +3589,8 @@ struct whisper_context * whisper_init_with_params_no_state(struct whisper_model_
         return nullptr;
     }
 
+    ctx->mel_calc = whisper_mel_calc_create(ctx->backend, ctx->model.filters);
+
     loader->close(loader->context);
 
     return ctx;
@@ -3713,6 +3711,8 @@ void whisper_free(struct whisper_context * ctx) {
 
         ggml_backend_free(ctx->backend);
 
+        delete ctx->mel_calc;
+        ctx->mel_calc = nullptr;
         delete ctx;
     }
 }
@@ -3730,11 +3730,21 @@ void whisper_free_params(struct whisper_full_params * params) {
 }
 
 int whisper_pcm_to_mel_with_state(struct whisper_context * ctx, struct whisper_state * state, const float * samples, int n_samples, int n_threads) {
-    if (!log_mel_spectrogram(*state, samples, n_samples, WHISPER_SAMPLE_RATE, WHISPER_N_FFT, WHISPER_HOP_LENGTH, ctx->model.filters.n_mel, n_threads, ctx->model.filters, false, state->mel)) {
-        WHISPER_LOG_ERROR("%s: failed to compute mel spectrogram\n", __func__);
-        return -1;
-    }
+    const int64_t t_start_us = ggml_time_us();
+    state->mel = ctx->mel_calc->calculate({samples, n_samples}, n_threads);
+    state->t_mel_us += ggml_time_us() - t_start_us;
 
+    // Dump log_mel_spectrogram
+    //{
+    //    auto& mel = state->mel;
+    //    std::ofstream outFile("log_mel_spectrogram.json");
+    //    outFile << "[";
+    //    for (uint64_t i = 0; i < mel.data.size() - 1; i++) {
+    //        outFile << mel.data[i] << ", ";
+    //    }
+    //    outFile << mel.data[mel.data.size() - 1] << "]";
+    //    outFile.close();
+    //}
     return 0;
 }
 

whisper.h

@@ -31,8 +31,10 @@
 
 #define WHISPER_SAMPLE_RATE 16000
 #define WHISPER_N_FFT       400
+#define WHISPER_N_FFT_HALF  (WHISPER_N_FFT / 2 + 1)
 #define WHISPER_HOP_LENGTH  160
 #define WHISPER_CHUNK_SIZE  30
+#define WHISPER_N_SAMPLES   (WHISPER_SAMPLE_RATE * WHISPER_CHUNK_SIZE)
 
 #ifdef __cplusplus
 extern "C" {