stream : add "max_tokens" cli arg

Controls the max tokens per segment for the stream example

CMakeLists.txt

@@ -48,6 +48,8 @@ option(WHISPER_SUPPORT_SDL2            "whisper: support for libSDL2" OFF)
 
 if (APPLE)
     option(WHISPER_NO_ACCELERATE       "whisper: disable Accelerate framework" OFF)
+    option(WHISPER_NO_AVX              "whisper: disable AVX" OFF)
+    option(WHISPER_NO_AVX2             "whisper: disable AVX2" OFF)
 else()
     option(WHISPER_SUPPORT_OPENBLAS    "whisper: support for OpenBLAS" OFF)
 endif()
@@ -94,17 +96,6 @@ if (APPLE AND NOT WHISPER_NO_ACCELERATE)
     else()
         message(WARNING "Accelerate framework not found")
     endif()
-
-    find_library(FOUNDATION_LIBRARY Foundation REQUIRED)
-    find_library(METAL_FRAMEWORK    Metal      REQUIRED)
-    find_library(METALKIT_FRAMEWORK MetalKit   REQUIRED)
-    find_library(METALPERFORMANCE_FRAMEWORK MetalPerformanceShaders REQUIRED)
-
-    set(WHISPER_EXTRA_LIBS ${WHISPER_EXTRA_LIBS}
-        ${FOUNDATION_LIBRARY}
-        ${METAL_FRAMEWORK}
-        ${METALKIT_FRAMEWORK}
-        ${METALPERFORMANCE_FRAMEWORK})
 endif()
 
 if (WHISPER_SUPPORT_OPENBLAS)
@@ -154,15 +145,21 @@ if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES
 else()
     message(STATUS "x86 detected")
     if (MSVC)
-        set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /arch:AVX2")
-        set(CMAKE_C_FLAGS_RELEASE "${CMAKE_C_FLAGS_RELEASE} /arch:AVX2")
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /arch:AVX2")
+        set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /arch:AVX2")
     else()
         if (EMSCRIPTEN)
             # we require support for WASM SIMD 128-bit
             set(CMAKE_C_FLAGS   "${CMAKE_C_FLAGS}   -pthread -msimd128")
             set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -pthread")
         else()
-            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx -mavx2 -mfma -mf16c")
+            if(NOT WHISPER_NO_AVX)
+                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx")
+            endif()
+            if(NOT WHISPER_NO_AVX2)
+                set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mavx2")
+            endif()
+            set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -mfma -mf16c")
         endif()
     endif()
 endif()
@@ -179,7 +176,6 @@ set(TARGET whisper)
 
 add_library(${TARGET}
     ggml.c
-    ggml-mtl.m
     whisper.cpp
     )
 

Makefile

@@ -50,7 +50,19 @@ endif
 # TODO: probably these flags need to be tweaked on some architectures
 #       feel free to update the Makefile for your architecture and send a pull request or issue
 ifeq ($(UNAME_M),x86_64)
-	CFLAGS += -mavx -mavx2 -mfma -mf16c
+	CFLAGS += -mfma -mf16c
+	ifeq ($(UNAME_S),Darwin)
+		AVX1_M := $(shell sysctl machdep.cpu.features)
+		ifneq (,$(findstring AVX1.0,$(AVX1_M)))
+			CFLAGS += -mavx
+		endif
+		AVX2_M := $(shell sysctl machdep.cpu.leaf7_features)
+		ifneq (,$(findstring AVX2,$(AVX2_M)))
+			CFLAGS += -mavx2
+		endif
+	else
+		CFLAGS += -mavx -mavx2
+	endif
 endif
 ifeq ($(UNAME_M),amd64)
 	CFLAGS += -mavx -mavx2 -mfma -mf16c
@@ -58,8 +70,8 @@ endif
 ifndef WHISPER_NO_ACCELERATE
 	# Mac M1 - include Accelerate framework
 	ifeq ($(UNAME_S),Darwin)
-		CFLAGS  += -DGGML_USE_ACCELERATE -DGGML_PERF
-		LDFLAGS += -framework Foundation -framework Accelerate -framework Metal -framework MetalKit -framework MetalPerformanceShaders
+		CFLAGS  += -DGGML_USE_ACCELERATE
+		LDFLAGS += -framework Accelerate
 	endif
 endif
 ifneq ($(filter aarch64%,$(UNAME_M)),)
@@ -77,28 +89,26 @@ ifneq ($(filter armv8%,$(UNAME_M)),)
 	CFLAGS += -mfp16-format=ieee -mno-unaligned-access
 endif
 
-#
-# Build library + main
-#
+default: main
 
-main: examples/main/main.cpp ggml.o ggml-mtl.o whisper.o
-	$(CXX) $(CXXFLAGS) examples/main/main.cpp whisper.o ggml.o ggml-mtl.o -o main $(LDFLAGS)
-	./main -h
+#
+# Build library
+#
 
 ggml.o: ggml.c ggml.h
 	$(CC)  $(CFLAGS)   -c ggml.c -o ggml.o
 
-ggml-mtl.o: ggml-mtl.m ggml-mtl.h
-	$(CC)  $(CFLAGS)   -c ggml-mtl.m -o ggml-mtl.o
-
 whisper.o: whisper.cpp whisper.h
 	$(CXX) $(CXXFLAGS) -c whisper.cpp -o whisper.o
 
-libwhisper.a: ggml.o ggml-mtl.o whisper.o
-	$(AR) rcs libwhisper.a ggml.o ggml-mtl.o whisper.o
+libwhisper.a: ggml.o whisper.o
+	$(AR) rcs libwhisper.a ggml.o whisper.o
+
+libwhisper.so: ggml.o whisper.o
+	$(CXX) $(CXXFLAGS) -shared -o libwhisper.so ggml.o whisper.o $(LDFLAGS)
 
 clean:
-	rm -f *.o main stream bench libwhisper.a
+	rm -f *.o main stream bench libwhisper.a libwhisper.so
 
 #
 # Examples
@@ -106,6 +116,10 @@ clean:
 
 CC_SDL=`sdl2-config --cflags --libs`
 
+main: examples/main/main.cpp ggml.o whisper.o
+	$(CXX) $(CXXFLAGS) examples/main/main.cpp ggml.o whisper.o -o main $(LDFLAGS)
+	./main -h
+
 stream: examples/stream/stream.cpp ggml.o whisper.o
 	$(CXX) $(CXXFLAGS) examples/stream/stream.cpp ggml.o whisper.o -o stream $(CC_SDL) $(LDFLAGS)
 

examples/main/main.cpp

@@ -59,6 +59,7 @@ struct whisper_params {
 
     float word_thold = 0.01f;
 
+    bool speed_up             = false;
     bool verbose              = false;
     bool translate            = false;
     bool output_txt           = false;
@@ -104,6 +105,8 @@ bool whisper_params_parse(int argc, char ** argv, whisper_params & params) {
             params.max_len = std::stoi(argv[++i]);
         } else if (arg == "-wt" || arg == "--word-thold") {
             params.word_thold = std::stof(argv[++i]);
+        } else if (arg == "-su" || arg == "--speed-up") {
+            params.speed_up = true;
         } else if (arg == "-v" || arg == "--verbose") {
             params.verbose = true;
         } else if (arg == "--translate") {
@@ -161,6 +164,7 @@ void whisper_print_usage(int argc, char ** argv, const whisper_params & params)
     fprintf(stderr, "  -mc N,    --max-context N  maximum number of text context tokens to store (default: max)\n");
     fprintf(stderr, "  -ml N,    --max-len N      maximum segment length in characters (default: %d)\n", params.max_len);
     fprintf(stderr, "  -wt N,    --word-thold N   word timestamp probability threshold (default: %f)\n", params.word_thold);
+    fprintf(stderr, "  -su,      --speed-up       speed up audio by factor of 2 (faster processing, reduced accuracy, default: %s)\n", params.speed_up ? "true" : "false");
     fprintf(stderr, "  -v,       --verbose        verbose output\n");
     fprintf(stderr, "            --translate      translate from source language to english\n");
     fprintf(stderr, "  -otxt,    --output-txt     output result in a text file\n");
@@ -454,9 +458,30 @@ int main(int argc, char ** argv) {
         std::vector<float> pcmf32;
         {
             drwav wav;
-            if (!drwav_init_file(&wav, fname_inp.c_str(), NULL)) {
-                fprintf(stderr, "%s: failed to open WAV file '%s' - check your input\n", argv[0], fname_inp.c_str());
-                whisper_print_usage(argc, argv, {});
+
+            if (fname_inp == "-") {
+                std::vector<uint8_t> wav_data;
+                {
+                    uint8_t buf[1024];
+                    while (true)
+                    {
+                        const size_t n = fread(buf, 1, sizeof(buf), stdin);
+                        if (n == 0)
+                        {
+                            break;
+                        }
+                        wav_data.insert(wav_data.end(), buf, buf + n);
+                    }
+                }
+
+                if (drwav_init_memory(&wav, wav_data.data(), wav_data.size(), NULL) == false)
+                {
+                    fprintf(stderr, "error: failed to open WAV file from stdin\n");
+                    return 4;
+                }
+            }
+            else if (drwav_init_file(&wav, fname_inp.c_str(), NULL) == false) {
+                fprintf(stderr, "error: failed to open '%s' as WAV file\n", fname_inp.c_str());
                 return 4;
             }
 
@@ -542,6 +567,8 @@ int main(int argc, char ** argv) {
             wparams.thold_pt             = params.word_thold;
             wparams.max_len              = params.output_wts && params.max_len == 0 ? 60 : params.max_len;
 
+            wparams.speed_up             = params.speed_up;
+
             // this callback is called on each new segment
             if (!wparams.print_realtime) {
                 wparams.new_segment_callback           = whisper_print_segment_callback;

examples/stream/stream.cpp

@@ -40,7 +40,10 @@ struct whisper_params {
     int32_t step_ms    = 3000;
     int32_t length_ms  = 10000;
     int32_t capture_id = -1;
+    int32_t max_tokens = 32;
+    int32_t audio_ctx  = 0;
 
+    bool speed_up             = false;
     bool verbose              = false;
     bool translate            = false;
     bool no_context           = true;
@@ -68,6 +71,12 @@ bool whisper_params_parse(int argc, char ** argv, whisper_params & params) {
             params.length_ms = std::stoi(argv[++i]);
         } else if (arg == "-c" || arg == "--capture") {
             params.capture_id = std::stoi(argv[++i]);
+        } else if (arg == "-mt" || arg == "--max_tokens") {
+            params.max_tokens = std::stoi(argv[++i]);
+        } else if (arg == "-ac" || arg == "--audio_ctx") {
+            params.audio_ctx = std::stoi(argv[++i]);
+        } else if (arg == "-su" || arg == "--speed-up") {
+            params.speed_up = true;
         } else if (arg == "-v" || arg == "--verbose") {
             params.verbose = true;
         } else if (arg == "--translate") {
@@ -113,6 +122,9 @@ void whisper_print_usage(int argc, char ** argv, const whisper_params & params)
     fprintf(stderr, "            --step N         audio step size in milliseconds (default: %d)\n", params.step_ms);
     fprintf(stderr, "            --length N       audio length in milliseconds (default: %d)\n", params.length_ms);
     fprintf(stderr, "  -c ID,    --capture ID     capture device ID (default: -1)\n");
+    fprintf(stderr, "  -mt N,    --max_tokens N   maximum number of tokens per audio chunk (default: %d)\n", params.max_tokens);
+    fprintf(stderr, "  -ac N,    --audio_ctx N    audio context size (default: %d, 0 - all)\n", params.audio_ctx);
+    fprintf(stderr, "  -su,      --speed-up       speed up audio by factor of 2 (faster processing, reduced accuracy, default: %s)\n", params.speed_up ? "true" : "false");
     fprintf(stderr, "  -v,       --verbose        verbose output\n");
     fprintf(stderr, "            --translate      translate from source language to english\n");
     fprintf(stderr, "  -kc,      --keep-context   keep text context from earlier audio (default: false)\n");
@@ -217,6 +229,7 @@ int main(int argc, char ** argv) {
     const int n_samples = (params.step_ms/1000.0)*WHISPER_SAMPLE_RATE;
     const int n_samples_len = (params.length_ms/1000.0)*WHISPER_SAMPLE_RATE;
     const int n_samples_30s = 30*WHISPER_SAMPLE_RATE;
+    const int n_samples_keep = 0.2*WHISPER_SAMPLE_RATE;
 
     std::vector<float> pcmf32(n_samples_30s, 0.0f);
     std::vector<float> pcmf32_old;
@@ -299,7 +312,7 @@ int main(int argc, char ** argv) {
         //const int n_samples_take = std::min((int) pcmf32_old.size(), std::max(0, n_samples_30s/30 - n_samples_new));
 
         // take up to params.length_ms audio from previous iteration
-        const int n_samples_take = std::min((int) pcmf32_old.size(), std::max(0, n_samples_len - n_samples_new));
+        const int n_samples_take = std::min((int) pcmf32_old.size(), std::max(0, n_samples_keep + n_samples_len - n_samples_new));
 
         //printf("processing: take = %d, new = %d, old = %d\n", n_samples_take, n_samples_new, (int) pcmf32_old.size());
 
@@ -323,9 +336,14 @@ int main(int argc, char ** argv) {
             wparams.print_timestamps     = !params.no_timestamps;
             wparams.translate            = params.translate;
             wparams.no_context           = params.no_context;
+            wparams.single_segment       = true;
+            wparams.max_tokens           = params.max_tokens;
             wparams.language             = params.language.c_str();
             wparams.n_threads            = params.n_threads;
 
+            wparams.audio_ctx            = params.audio_ctx;
+            wparams.speed_up             = params.speed_up;
+
             if (whisper_full(ctx, wparams, pcmf32.data(), pcmf32.size()) != 0) {
                 fprintf(stderr, "%s: failed to process audio\n", argv[0]);
                 return 6;
@@ -373,7 +391,8 @@ int main(int argc, char ** argv) {
             if ((n_iter % n_new_line) == 0) {
                 printf("\n");
 
-                pcmf32_old.clear();
+                // keep part of the audio for next iteration to try to mitigate word boundary issues
+                pcmf32_old = std::vector<float>(pcmf32.end() - n_samples_keep, pcmf32.end());
             }
         }
     }

ggml-mtl.h

@@ -1,38 +0,0 @@
-#pragma once
-
-#include <stdint.h>
-#include <stddef.h>
-
-// TODO: this will hold dynamic context data in the future
-//       currently unused
-struct ggml_mtl_context {
-    void * dummy;
-};
-
-struct ggml_mtl_object {
-    int32_t id;
-    void * data;
-};
-
-struct ggml_mtl_context * ggml_mtl_init(void);
-
-struct ggml_mtl_object ggml_mtl_alloc(size_t size);
-
-// multiply matrix by vector
-void ggml_mtl_mul_mat_vec_f16(
-    struct ggml_mtl_context * ctx,
-    struct ggml_mtl_object    src0,  // matrix f16
-    const __fp16            * src1,  // vector f16
-    float                   * dst,   // vector f32
-    int                       nrows,
-    int                       ncols);
-
-// multiply matrix by matrix
-void ggml_mtl_mul_mat_f16(
-    struct ggml_mtl_context * ctx,
-    struct ggml_mtl_object    src0,  // matrix f16
-    const __fp16            * src1,  // matrix f16
-    float                   * dst,   // matrix f32
-    int                       nrows0,
-    int                       nrows1,
-    int                       ncols);

ggml-mtl.m

@@ -1,162 +0,0 @@
-#import "ggml-mtl.h"
-
-#import <Foundation/Foundation.h>
-#import <Metal/Metal.h>
-#import <MetalPerformanceShaders/MetalPerformanceShaders.h>
-
-#define GGML_MTL_MAX_BUFFERS 256
-
-// global static storage for Metal buffers
-// TODO: move this into a dynamic context
-static id<MTLBuffer> g_buffers[GGML_MTL_MAX_BUFFERS];
-
-// global MTL context
-// TODO: move this into a dynamic context
-static id<MTLDevice>       g_device;
-static id<MTLCommandQueue> g_command_queue;
-
-struct ggml_mtl_context * ggml_mtl_init() {
-    // TODO: implement properly
-    //       for now, init the global MTL context and MTL buffers
-    g_device = MTLCreateSystemDefaultDevice();
-
-    g_command_queue = [g_device newCommandQueue];
-    if (g_command_queue == nil)
-    {
-        NSLog(@"Failed to find the command queue.");
-        return nil;
-    }
-
-    return nil;
-}
-
-// search for unallocated buffer slot and use it
-struct ggml_mtl_object ggml_mtl_alloc(size_t size) {
-    // TODO: temporarily making sure that the buffers are nil at the start
-    static bool first = true;
-    if (first) {
-        for (int i = 0; i < GGML_MTL_MAX_BUFFERS; ++i) {
-            assert(g_buffers[i] == nil);
-        }
-        first = false;
-    }
-
-    struct ggml_mtl_object obj = { -1, nil };
-
-    for (int i = 0; i < GGML_MTL_MAX_BUFFERS; i++) {
-        if (g_buffers[i] == nil) {
-            g_buffers[i] = [g_device newBufferWithLength:size options:MTLResourceStorageModeManaged];
-
-            // lunk the MTL buffer to the ggml object
-            obj.id = i;
-            obj.data = [g_buffers[i] contents];
-
-            break;
-        }
-    }
-
-    return obj;
-}
-
-struct params_mul_mat_vec {
-    int  N; // rows
-    int  M; // cols
-};
-
-// multiply matrix with a vector using MPSMatrixVectorMultiplication
-void ggml_mtl_mul_mat_vec_f16(
-        struct ggml_mtl_context * ctx,
-        struct ggml_mtl_object    src0,
-        const __fp16            * src1,
-        float                   * dst,
-        int                       nrows,
-        int                       ncols) {
-    (void) ctx; // unused
-
-    // Create a command buffer to hold commands.
-    id<MTLCommandBuffer> commandBuffer = [g_command_queue commandBuffer];
-    assert(commandBuffer != nil);
-
-    // make managed device buffer to store src1
-    id<MTLBuffer> src1_buffer = [g_device newBufferWithBytes:src1 length:ncols*sizeof(__fp16) options:MTLResourceStorageModeManaged];
-    id<MTLBuffer> dst_buffer  = [g_device newBufferWithLength:nrows*sizeof(float) options:MTLResourceStorageModeManaged];
-
-    // MPSMatrixDescriptor
-    MPSMatrixDescriptor *src0_desc = [MPSMatrixDescriptor matrixDescriptorWithRows:nrows columns:ncols rowBytes:ncols*sizeof(__fp16) dataType:MPSDataTypeFloat16];
-    MPSVectorDescriptor *src1_desc = [MPSVectorDescriptor vectorDescriptorWithLength:ncols dataType:MPSDataTypeFloat16];
-    MPSVectorDescriptor *dst_desc  = [MPSVectorDescriptor vectorDescriptorWithLength:nrows dataType:MPSDataTypeFloat32];
-
-    // MPSMatrix
-    MPSMatrix *src0_mat = [[MPSMatrix alloc] initWithBuffer:g_buffers[src0.id] descriptor:src0_desc];
-    MPSVector *src1_vec = [[MPSVector alloc] initWithBuffer:src1_buffer        descriptor:src1_desc];
-    MPSVector *dst_vec  = [[MPSVector alloc] initWithBuffer:dst_buffer         descriptor:dst_desc];
-
-    // MPSMatrixVectorMultiplication
-    MPSMatrixVectorMultiplication *mul_mat_vec = [[MPSMatrixVectorMultiplication alloc] initWithDevice:g_device transpose:NO rows:nrows columns:ncols alpha:1.0 beta:0.0];
-
-    // encode
-    [mul_mat_vec encodeToCommandBuffer:commandBuffer
-                           inputMatrix:src0_mat
-                           inputVector:src1_vec
-                          resultVector:dst_vec];
-
-    [commandBuffer commit];
-    [commandBuffer waitUntilCompleted];
-
-    // copy GPU result to CPU
-    memcpy(dst, [dst_buffer contents], nrows*sizeof(float));
-}
-
-// multiply matrix with a matrix using MPSMatrixMultiplication
-void ggml_mtl_mul_mat_f16(
-        struct ggml_mtl_context * ctx,
-        struct ggml_mtl_object    src0,
-        const __fp16            * src1,
-        float                   * dst,
-        int                       nrows0,
-        int                       nrows1,
-        int                       ncols) {
-    (void) ctx; // unused
-
-    // Create a command buffer to hold commands.
-    id<MTLCommandBuffer> commandBuffer = [g_command_queue commandBuffer];
-    assert(commandBuffer != nil);
-
-    // make managed device buffer to store src1
-    id<MTLBuffer> src1_buffer = [g_device newBufferWithBytes:src1 length:ncols*nrows1*sizeof(__fp16) options:MTLResourceStorageModeManaged];
-    id<MTLBuffer> dst_buffer  = [g_device newBufferWithLength:nrows0*nrows1*sizeof(float) options:MTLResourceStorageModeManaged];
-
-    // MPSMatrixDescriptor
-    MPSMatrixDescriptor *src0_desc = [MPSMatrixDescriptor matrixDescriptorWithRows:nrows0 columns:ncols  rowBytes:ncols*sizeof(__fp16) dataType:MPSDataTypeFloat16];
-    MPSMatrixDescriptor *src1_desc = [MPSMatrixDescriptor matrixDescriptorWithRows:nrows1 columns:ncols  rowBytes:ncols*sizeof(__fp16) dataType:MPSDataTypeFloat16];
-    MPSMatrixDescriptor *dst_desc  = [MPSMatrixDescriptor matrixDescriptorWithRows:nrows1 columns:nrows0 rowBytes:nrows0*sizeof(float) dataType:MPSDataTypeFloat32];
-
-    // MPSMatrix
-    MPSMatrix *src0_mat = [[MPSMatrix alloc] initWithBuffer:g_buffers[src0.id] descriptor:src0_desc];
-    MPSMatrix *src1_mat = [[MPSMatrix alloc] initWithBuffer:src1_buffer        descriptor:src1_desc];
-    MPSMatrix *dst_mat  = [[MPSMatrix alloc] initWithBuffer:dst_buffer         descriptor:dst_desc];
-
-    //// MPSMatrixMultiplication z = x * yT
-    //MPSMatrixMultiplication *mul_mat = [[MPSMatrixMultiplication alloc] initWithDevice:g_device transposeLeft:NO transposeRight:YES resultRows:nrows resultColumns:nrows interiorColumns:ncols alpha:1.0 beta:0.0];
-
-    //// encode
-    //[mul_mat encodeToCommandBuffer:commandBuffer
-    //                   leftMatrix:src0_mat
-    //                  rightMatrix:src1_mat
-    //                 resultMatrix:dst_mat];
-
-    // MPSMatrixMultiplication zT = xT * y
-    MPSMatrixMultiplication *mul_mat = [[MPSMatrixMultiplication alloc] initWithDevice:g_device transposeLeft:NO transposeRight:YES resultRows:nrows1 resultColumns:nrows0 interiorColumns:ncols alpha:1.0 beta:0.0];
-
-    // encode
-    [mul_mat encodeToCommandBuffer:commandBuffer
-                       leftMatrix:src1_mat
-                      rightMatrix:src0_mat
-                     resultMatrix:dst_mat];
-
-    [commandBuffer commit];
-    [commandBuffer waitUntilCompleted];
-
-    // copy GPU result to CPU
-    memcpy(dst, [dst_buffer contents], nrows0*nrows1*sizeof(float));
-}

ggml.c

@@ -1,7 +1,5 @@
 #include "ggml.h"
 
-#include "ggml-mtl.h"
-
 #if defined(_MSC_VER) || defined(__MINGW32__)
 #include <malloc.h> // using malloc.h with MSC/MINGW
 #elif !defined(__FreeBSD__)
@@ -39,8 +37,14 @@ typedef HANDLE pthread_t;
 
 typedef DWORD thread_ret_t;
 static int pthread_create(pthread_t* out, void* unused, thread_ret_t(*func)(void*), void* arg) {
-    out = CreateThread(NULL, 0, func, arg, 0, NULL);
-    return out != NULL;
+    HANDLE handle = CreateThread(NULL, 0, func, arg, 0, NULL);
+    if (handle == NULL)
+    {
+        return EAGAIN;
+    }
+
+    *out = handle;
+    return 0;
 }
 
 static int pthread_join(pthread_t thread, void* unused) {
@@ -1309,8 +1313,6 @@ struct ggml_context * ggml_init(struct ggml_init_params params) {
 
     static bool first_time = true;
     if (first_time) {
-        ggml_mtl_init(); // TODO: fix this
-
         for (int i = 0; i < GGML_MAX_CONTEXTS; i++) {
             g_state.contexts[i].used = false;
         }
@@ -1466,104 +1468,6 @@ struct ggml_tensor * ggml_new_tensor_impl(
         /*.perf_cycles  =*/ 0,
         /*.perf_time_us =*/ 0,
         /*.data         =*/ data == NULL ? (void *)(result + 1) : data,
-        /*.id           =*/ -1,
-        /*.pad          =*/ { 0 },
-    };
-
-    ggml_assert_aligned(result->data);
-
-    for (int i = 0; i < n_dims; i++) {
-        result->ne[i] = ne[i];
-    }
-
-    result->nb[0] = GGML_TYPE_SIZE[type];
-    for (int i = 1; i < GGML_MAX_DIMS; i++) {
-        result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
-    }
-
-    ctx->n_objects++;
-
-    return result;
-}
-
-struct ggml_tensor * ggml_new_tensor_mtl_impl(
-        struct ggml_context * ctx,
-        enum   ggml_type type,
-        int    n_dims,
-        const int* ne,
-        void*  data) {
-    // always insert objects at the end of the context's memory pool
-    struct ggml_object * obj_cur = ctx->objects_end;
-
-    const size_t cur_offset = obj_cur == NULL ? 0 : obj_cur->offset;
-    const size_t cur_size   = obj_cur == NULL ? 0 : obj_cur->size;
-    const size_t cur_end    = cur_offset + cur_size;
-
-    struct ggml_mtl_object obj_mtl;
-    {
-        assert(data == NULL); // TODO: in-place metal buffer, need page aligned memory
-        size_t size_needed_mtl = 0;
-        if (data == NULL) {
-            size_needed_mtl += GGML_TYPE_SIZE[type];
-            for (int i = 0; i < n_dims; i++) {
-                size_needed_mtl *= ne[i];
-            }
-        }
-
-        obj_mtl = ggml_mtl_alloc(size_needed_mtl);
-    }
-
-    size_t size_needed = 0;
-    size_needed += sizeof(struct ggml_tensor);
-
-    if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) {
-        GGML_PRINT("%s: not enough space in the context's memory pool\n", __func__);
-        assert(false);
-        return NULL;
-    }
-
-    char * const mem_buffer = ctx->mem_buffer;
-
-    struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end);
-
-    *obj_new = (struct ggml_object) {
-        .offset = cur_end + GGML_OBJECT_SIZE,
-        .size   = size_needed,
-        .next   = NULL,
-    };
-
-    if (obj_cur != NULL) {
-        obj_cur->next = obj_new;
-    } else {
-        // this is the first object in this context
-        ctx->objects_begin = obj_new;
-    }
-
-    ctx->objects_end = obj_new;
-
-    //GGML_PRINT_DEBUG("%s: inserted new object at %zu\n", __func__, cur_end);
-
-    struct ggml_tensor * const result = (struct ggml_tensor *)(mem_buffer + obj_new->offset);
-
-    ggml_assert_aligned(result);
-
-    *result = (struct ggml_tensor) {
-        /*.type         =*/ type,
-        /*.n_dims       =*/ n_dims,
-        /*.ne           =*/ { 1, 1, 1, 1 },
-        /*.nb           =*/ { 0, 0, 0, 0 },
-        /*.op           =*/ GGML_OP_NONE,
-        /*.is_param     =*/ false,
-        /*.grad         =*/ NULL,
-        /*.src0         =*/ NULL,
-        /*.src1         =*/ NULL,
-        /*.opt          =*/ { NULL },
-        /*.n_tasks      =*/ 0,
-        /*.perf_runs    =*/ 0,
-        /*.perf_cycles  =*/ 0,
-        /*.perf_time_us =*/ 0,
-        /*.data         =*/ obj_mtl.data,
-        /*.id           =*/ obj_mtl.id,
         /*.pad          =*/ { 0 },
     };
 
@@ -1591,14 +1495,6 @@ struct ggml_tensor * ggml_new_tensor(
     return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL);
 }
 
-struct ggml_tensor * ggml_new_tensor_mtl(
-        struct ggml_context * ctx,
-        enum   ggml_type type,
-        int    n_dims,
-        const int* ne) {
-    return ggml_new_tensor_mtl_impl(ctx, type, n_dims, ne, NULL);
-}
-
 struct ggml_tensor * ggml_new_tensor_1d(
         struct ggml_context * ctx,
         enum   ggml_type type,
@@ -1615,15 +1511,6 @@ struct ggml_tensor * ggml_new_tensor_2d(
     return ggml_new_tensor(ctx, type, 2, ne);
 }
 
-struct ggml_tensor * ggml_new_tensor_2d_mtl(
-        struct ggml_context * ctx,
-        enum   ggml_type type,
-        int    ne0,
-        int    ne1) {
-    const int ne[2] = { ne0, ne1 };
-    return ggml_new_tensor_mtl(ctx, type, 2, ne);
-}
-
 struct ggml_tensor * ggml_new_tensor_3d(
         struct ggml_context * ctx,
         enum   ggml_type type,
@@ -4462,11 +4349,8 @@ void ggml_compute_forward_mul_mat_f16_f32(
     // nb00 <  nb01 - src0 is transposed
     //   compute by src0 columns
 
-    // are we using Metal?
-    const bool is_mtl = src0->id >= 0;
-
 #if defined(GGML_USE_ACCELERATE) || defined(GGML_USE_OPENBLAS)
-    if (ggml_compute_forward_mul_mat_use_blas(src0, src1, dst) && !is_mtl) {
+    if (ggml_compute_forward_mul_mat_use_blas(src0, src1, dst)) {
         GGML_ASSERT(nb10 == sizeof(float));
 
         if (params->ith != 0) return;
@@ -4594,20 +4478,6 @@ void ggml_compute_forward_mul_mat_f16_f32(
 
         // parallelize by src0 rows using ggml_vec_dot_f32
 
-        if (is_mtl) {
-            assert(ne02 == 1);
-            assert(ne03 == 1);
-
-            if (params->ith == 0) {
-                printf("XXXXXXXXXXX src0->ne[0] = %d, src0->ne[1] = %d\n", src0->ne[0], src0->ne[1]);
-                printf("XXXXXXXXXXX src1->ne[0] = %d, src1->ne[1] = %d\n", src1->ne[0], src1->ne[1]);
-                struct ggml_mtl_object src0_mtl = { src0->id, src0->data };
-                ggml_fp16_t * src1_fp16 = params->wdata;
-                ggml_mtl_mul_mat_f16(NULL, src0_mtl, src1_fp16, dst->data, ne01, ne11, ne00);
-            }
-            return;
-        }
-
         // total rows in src0
         const int nr = ne01*ne02*ne03;
 

ggml.h

@@ -1,5 +1,174 @@
 #pragma once
 
+//
+// GGML Tensor Library
+//
+// This documentation is still a work in progress.
+// If you wish some specific topics to be covered, feel free to drop a comment:
+//
+//   https://github.com/ggerganov/whisper.cpp/issues/40
+//
+// ## Overview
+//
+// This library implements:
+//
+//  - a set of tensor operations
+//  - automatic differentiation
+//  - basic optimization algorithms
+//
+// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes,
+// but is not limited to, the following:
+//
+//  - linear regression
+//  - support vector machines
+//  - neural networks
+//
+// The library allows the user to define a certain function using the available tensor operations. This function
+// definition is represented internally via a computation graph. Each tensor operation in the function definition
+// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the
+// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized
+// using one of the available optimization algorithms.
+//
+// For example, here we define the function: f(x) = a*x^2 + b
+//
+//   {
+//       struct ggml_init_params params = {
+//           .mem_size   = 16*1024*1024,
+//           .mem_buffer = NULL,
+//       };
+//
+//       // memory allocation happens here
+//       struct ggml_context * ctx = ggml_init(params);
+//
+//       struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//
+//       ggml_set_param(ctx, x); // x is an input variable
+//
+//       struct ggml_tensor * a  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//       struct ggml_tensor * b  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+//       struct ggml_tensor * x2 = ggml_mul(ctx, x, x);
+//       struct ggml_tensor * f  = ggml_add(ctx, ggml_mul(ctx, a, x2), b);
+//
+//       ...
+//   }
+//
+// Notice that the function definition above does not involve any actual computation. The computation is performed only
+// when the user explicitly requests it. For example, to compute the function's value at x = 2.0:
+//
+//   {
+//       ...
+//
+//       struct ggml_cgraph gf = ggml_build_forward(f);
+//
+//       // set the input variable and parameter values
+//       ggml_set_f32(x, 2.0f);
+//       ggml_set_f32(a, 3.0f);
+//       ggml_set_f32(b, 4.0f);
+//
+//       ggml_graph_compute(ctx0, &gf);
+//
+//       printf("f = %f\n", ggml_get_f32_1d(f, 0));
+//
+//       ...
+//   }
+//
+// The actual computation is performed in the ggml_graph_compute() function.
+//
+// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the
+// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know
+// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory
+// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was
+// actually needed.
+//
+// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic
+// differentiation and optimization algorithms.
+//
+// The described approach allows to define the function graph once and then compute its forward or backward graphs
+// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way
+// the user can avoid the memory allocation overhead at runtime.
+//
+// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class
+// citizens, but in theory the library can be extended to support FP8 and integer data types.
+//
+// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary
+// and binary operations. Most of the available operations fall into one of these two categories. With time, it became
+// clear that the library needs to support more complex operations. The way to support these operations is not clear
+// yet, but a few examples are demonstrated in the following operations:
+//
+//   - ggml_permute()
+//   - ggml_conv_1d_1s()
+//   - ggml_conv_1d_2s()
+//
+// For each tensor operator, the library implements a forward and backward computation function. The forward function
+// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the
+// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a
+// calculus class, or watch the following video:
+//
+//   What is Automatic Differentiation?
+//   https://www.youtube.com/watch?v=wG_nF1awSSY
+//
+//
+// ## Tensor data (struct ggml_tensor)
+//
+// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of
+// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains
+// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example:
+//
+//   {
+//       struct ggml_tensor * c = ggml_add(ctx, a, b);
+//
+//       assert(c->src[0] == a);
+//       assert(c->src[1] == b);
+//   }
+//
+// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the
+// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows
+// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and
+// permutation. All tensor operations have to take the stride into account and not assume that the tensor is
+// contiguous in memory.
+//
+// The data of the tensor is accessed via the "data" pointer. For example:
+//
+//   {
+//       struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2, 3);
+//
+//       // a[1, 2] = 1.0f;
+//       *(float *) ((char *) a->data + 2*a->nb[1] + 1*a->nb[0]) = 1.0f;
+//
+//       // a[2, 0] = 2.0f;
+//       *(float *) ((char *) a->data + 0*a->nb[1] + 2*a->nb[0]) = 2.0f;
+//
+//       ...
+//   }
+//
+// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used.
+//
+// ## The matrix multiplication operator (ggml_mul_mat)
+//
+// TODO
+//
+//
+// ## Multi-threading
+//
+// TODO
+//
+//
+// ## Overview of ggml.c
+//
+// TODO
+//
+//
+// ## SIMD optimizations
+//
+// TODO
+//
+//
+// ## Debugging ggml
+//
+// TODO
+//
+//
+
 #ifdef  __cplusplus
 extern "C" {
 #endif
@@ -21,7 +190,8 @@ typedef __fp16 ggml_fp16_t;
 typedef uint16_t ggml_fp16_t;
 #endif
 
-float ggml_fp16_to_fp32(ggml_fp16_t x);
+// convert FP16 <-> FP32
+float       ggml_fp16_to_fp32(ggml_fp16_t x);
 ggml_fp16_t ggml_fp32_to_fp16(float x);
 
 struct ggml_object;
@@ -36,6 +206,7 @@ enum ggml_type {
     GGML_TYPE_COUNT,
 };
 
+// available tensor operations:
 enum ggml_op {
     GGML_OP_NONE = 0,
 
@@ -108,8 +279,7 @@ struct ggml_tensor {
     int64_t perf_time_us;
 
     void * data;
-    int32_t id; // TODO: mtl buffer id
-    char pad[4];
+    char padding[8];
 };
 
 // computation graph
@@ -137,7 +307,7 @@ struct ggml_init_params {
     void * mem_buffer; // if NULL, memory will be allocated internally
 };
 
-void ggml_time_init(void);
+void    ggml_time_init(void); // call this once at the beginning of the program
 int64_t ggml_time_ms(void);
 int64_t ggml_time_us(void);
 int64_t ggml_cycles(void);
@@ -174,12 +344,6 @@ struct ggml_tensor * ggml_new_tensor_2d(
         int    ne0,
         int    ne1);
 
-struct ggml_tensor * ggml_new_tensor_2d_mtl(
-        struct ggml_context * ctx,
-        enum   ggml_type type,
-        int    ne0,
-        int    ne1);
-
 struct ggml_tensor * ggml_new_tensor_3d(
         struct ggml_context * ctx,
         enum   ggml_type type,

models/convert-pt-to-ggml.py

@@ -297,8 +297,6 @@ for name in list_vars.keys():
                 name == "encoder.conv2.bias"   or \
                 name == "encoder.positional_embedding" or \
                 name == "decoder.positional_embedding":
-            ftype = 0
-            data = data.astype(np.float32)
             print("  Converting to float32")
             data = data.astype(np.float32)
             ftype = 0

whisper.cpp

@@ -424,6 +424,9 @@ struct whisper_context {
     int64_t t_last;
     whisper_token tid_last;
     std::vector<float> energy; // PCM signal energy
+
+    // [EXPERIMENTAL] speed-up techniques
+    int32_t exp_n_audio_ctx; // 0 - use default
 };
 
 // load the model from a ggml file
@@ -613,7 +616,7 @@ static bool whisper_model_load(const std::string & fname, whisper_context & wctx
         const int n_audio_state = hparams.n_audio_state;
         const int n_audio_layer = hparams.n_audio_layer;
 
-        const int n_text_ctx = hparams.n_text_ctx;
+        const int n_text_ctx   = hparams.n_text_ctx;
         const int n_text_state = hparams.n_text_state;
         const int n_text_layer = hparams.n_text_layer;
 
@@ -748,7 +751,7 @@ static bool whisper_model_load(const std::string & fname, whisper_context & wctx
         const int n_audio_state = hparams.n_audio_state;
         const int n_audio_layer = hparams.n_audio_layer;
 
-        const int n_text_ctx = hparams.n_text_ctx;
+        const int n_text_ctx   = hparams.n_text_ctx;
         const int n_text_state = hparams.n_text_state;
         const int n_text_layer = hparams.n_text_layer;
 
@@ -788,10 +791,10 @@ static bool whisper_model_load(const std::string & fname, whisper_context & wctx
                 layer.mlp_ln_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
                 layer.mlp_ln_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
 
-                layer.mlp_0_w = ggml_new_tensor_2d_mtl(ctx, wtype,       n_audio_state, 4*n_audio_state); // offload to GPU
+                layer.mlp_0_w = ggml_new_tensor_2d(ctx, wtype,           n_audio_state, 4*n_audio_state);
                 layer.mlp_0_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 4*n_audio_state);
 
-                layer.mlp_1_w = ggml_new_tensor_2d_mtl(ctx, wtype,     4*n_audio_state, n_audio_state);   // offload to GPU
+                layer.mlp_1_w = ggml_new_tensor_2d(ctx, wtype,         4*n_audio_state, n_audio_state);
                 layer.mlp_1_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32,   n_audio_state);
 
                 layer.attn_ln_0_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state);
@@ -967,7 +970,7 @@ static bool whisper_model_load(const std::string & fname, whisper_context & wctx
 
         // key/value memory for the cross-attention layer
         {
-            const int n_audio_ctx   = hparams.n_audio_ctx;
+            const int n_audio_ctx = hparams.n_audio_ctx;
 
             const int n_mem      = n_text_layer*n_audio_ctx;
             const int n_elements = n_text_state*n_mem;
@@ -1076,13 +1079,11 @@ static bool whisper_encode(
     const auto & mel_inp = wctx.mel;
     const auto & hparams = model.hparams;
 
-    const int n_ctx   = hparams.n_audio_ctx;
+    const int n_ctx   = wctx.exp_n_audio_ctx > 0 ? wctx.exp_n_audio_ctx : hparams.n_audio_ctx;
     const int n_state = hparams.n_audio_state;
     const int n_head  = hparams.n_audio_head;
     const int n_layer = hparams.n_audio_layer;
 
-    const int N = n_ctx;
-
     const int n_mels = hparams.n_mels;
     assert(mel_inp.n_mel == n_mels);
 
@@ -1132,7 +1133,30 @@ static bool whisper_encode(
         cur = ggml_gelu(ctx0, cur);
     }
 
-    cur = ggml_add(ctx0, model.e_pe, ggml_transpose(ctx0, cur));
+    // ===================================================================
+    // NOTE: experimenting with partial evaluation of the encoder (ignore)
+    //static int iter = -1;
+    //const int n_iter = 1500/n_ctx;
+
+    //iter = (iter + 1) % n_iter;
+
+    //if (iter == 0) {
+    //    memset(model.memory_cross_k->data, 0, ggml_nbytes(model.memory_cross_k));
+    //    memset(model.memory_cross_v->data, 0, ggml_nbytes(model.memory_cross_v));
+    //}
+
+    static int iter = 0;
+
+    const size_t e_pe_stride = model.e_pe->ne[0]*ggml_element_size(model.e_pe);
+    const size_t e_pe_offset = model.e_pe->ne[0]*ggml_element_size(model.e_pe)*n_ctx*iter;
+
+    struct ggml_tensor * e_pe = ggml_view_2d(ctx0, model.e_pe, model.e_pe->ne[0], n_ctx, e_pe_stride, e_pe_offset);
+
+    cur = ggml_add(ctx0, e_pe, ggml_transpose(ctx0, cur));
+    // ===================================================================
+
+    // original:
+    //cur = ggml_add(ctx0, model.e_pe, ggml_transpose(ctx0, cur));
 
     struct ggml_tensor * inpL = cur;
 
@@ -1198,14 +1222,14 @@ static bool whisper_encode(
                 ggml_permute(ctxL,
                         ggml_cpy(ctxL,
                             Qcur,
-                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, N)),
+                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
                         0, 2, 1, 3);
 
             struct ggml_tensor * K =
                 ggml_permute(ctxL,
                         ggml_cpy(ctxL,
                             Kcur,
-                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, N)),
+                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
                         0, 2, 1, 3);
 
             struct ggml_tensor * V =
@@ -1213,9 +1237,9 @@ static bool whisper_encode(
                         ggml_permute(ctxL,
                             ggml_reshape_3d(ctxL,
                                 Vcur,
-                                n_state/n_head, n_head, N),
+                                n_state/n_head, n_head, n_ctx),
                             1, 2, 0, 3),
-                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, N, n_state/n_head, n_head)
+                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_ctx, n_state/n_head, n_head)
                         );
 
             struct ggml_tensor * KQV = ggml_flash_attn(ctxL, Q, K, V, false);
@@ -1224,14 +1248,14 @@ static bool whisper_encode(
                 ggml_permute(ctxL,
                         ggml_cpy(ctxL,
                             Qcur,
-                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F32, n_state/n_head, n_head, N)),
+                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F32, n_state/n_head, n_head, n_ctx)),
                         0, 2, 1, 3);
 
             struct ggml_tensor * K =
                 ggml_permute(ctxL,
                         ggml_cpy(ctxL,
                             Kcur,
-                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, N)),
+                            ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
                         0, 2, 1, 3);
 
             // K * Q
@@ -1249,7 +1273,7 @@ static bool whisper_encode(
             //    ggml_permute(ctxL,
             //            ggml_cpy(ctxL,
             //                Vcur,
-            //                ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, N)),
+            //                ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_head, n_ctx)),
             //            1, 2, 0, 3);
 
             //struct ggml_tensor * KQV = ggml_mul_mat(ctxL, V_trans, KQ_soft_max);
@@ -1259,9 +1283,9 @@ static bool whisper_encode(
                         ggml_permute(ctxL,
                             ggml_reshape_3d(ctxL,
                                 Vcur,
-                                n_state/n_head, n_head, N),
+                                n_state/n_head, n_head, n_ctx),
                             0, 2, 1, 3),
-                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, N, n_head)
+                        ggml_new_tensor_3d(ctxL, GGML_TYPE_F16, n_state/n_head, n_ctx, n_head)
                         );
 
             struct ggml_tensor * KQV = ggml_mul_mat(ctxL, ggml_transpose(ctxL, V), KQ_soft_max);
@@ -1271,7 +1295,7 @@ static bool whisper_encode(
 
             cur = ggml_cpy(ctxL,
                     KQV_merged,
-                    ggml_new_tensor_2d(ctxL, GGML_TYPE_F32, n_state, N));
+                    ggml_new_tensor_2d(ctxL, GGML_TYPE_F32, n_state, n_ctx));
         }
 
         // projection
@@ -1342,7 +1366,7 @@ static bool whisper_encode(
             ggml_build_forward_expand(&gf, inpO);
             ggml_graph_compute       (ctxL, &gf);
 
-            ggml_graph_print(&gf);
+            //ggml_graph_print(&gf);
         }
 
         // TODO: this is a hack to have per-layer computation graphs - need to come up with something better
@@ -1425,6 +1449,8 @@ static bool whisper_encode(
                         Vcross),
                     Vcross);
 
+            //struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_cross_k, n_state*n_ctx, (ggml_element_size(model.memory_cross_k)*n_state)*(il*hparams.n_audio_ctx + iter*n_ctx));
+            //struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_cross_v, n_state*n_ctx, (ggml_element_size(model.memory_cross_v)*n_state)*(il*hparams.n_audio_ctx + iter*n_ctx));
             struct ggml_tensor * k = ggml_view_1d(ctx0, model.memory_cross_k, n_state*n_ctx, (ggml_element_size(model.memory_cross_k)*n_state)*(il*n_ctx));
             struct ggml_tensor * v = ggml_view_1d(ctx0, model.memory_cross_v, n_state*n_ctx, (ggml_element_size(model.memory_cross_v)*n_state)*(il*n_ctx));
 
@@ -1474,7 +1500,7 @@ static bool whisper_decode(
     const int n_layer = hparams.n_text_layer;
 
     const int N = n_tokens;
-    const int M = hparams.n_audio_ctx;
+    const int M = wctx.exp_n_audio_ctx > 0 ? wctx.exp_n_audio_ctx : hparams.n_audio_ctx;
 
     struct ggml_init_params params = {
             .mem_size   = wctx.buf_compute.size(),
@@ -2031,6 +2057,7 @@ static bool log_mel_spectrogram(
     const int n_mel,
     const int n_threads,
     const whisper_filters & filters,
+    const bool speed_up,
     whisper_mel & mel) {
 
     // Hanning window
@@ -2044,7 +2071,7 @@ static bool log_mel_spectrogram(
     mel.n_len = (n_samples)/fft_step;
     mel.data.resize(mel.n_mel*mel.n_len);
 
-    const int n_fft = 1 + fft_size/2;
+    const int n_fft = 1 + (speed_up ? fft_size/4 : fft_size/2);
 
     //printf("%s: n_samples = %d, n_len = %d\n", __func__, n_samples, mel.n_len);
     //printf("%s: recording length: %f s\n", __func__, (float) n_samples/sample_rate);
@@ -2091,6 +2118,13 @@ static bool log_mel_spectrogram(
                     //}
                 }
 
+                if (speed_up) {
+                    // scale down in the frequency domain results in a speed up in the time domain
+                    for (int j = 0; j < n_fft; j++) {
+                        fft_out[j] = 0.5*(fft_out[2*j] + fft_out[2*j + 1]);
+                    }
+                }
+
                 // mel spectrogram
                 for (int j = 0; j < mel.n_mel; j++) {
                     double sum = 0.0;
@@ -2171,7 +2205,21 @@ void whisper_free(struct whisper_context * ctx) {
 int whisper_pcm_to_mel(struct whisper_context * ctx, const float * samples, int n_samples, int n_threads) {
     const int64_t t_start_us = ggml_time_us();
 
-    if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, WHISPER_N_FFT, WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, ctx->mel)) {
+    if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, WHISPER_N_FFT, WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, false, ctx->mel)) {
+        fprintf(stderr, "%s: failed to compute mel spectrogram\n", __func__);
+        return -1;
+    }
+
+    ctx->t_mel_us = ggml_time_us() - t_start_us;
+
+    return 0;
+}
+
+// same as whisper_pcm_to_mel, but applies a Phase Vocoder to speed up the audio x2
+int whisper_pcm_to_mel_phase_vocoder(struct whisper_context * ctx, const float * samples, int n_samples, int n_threads) {
+    const int64_t t_start_us = ggml_time_us();
+
+    if (!log_mel_spectrogram(samples, n_samples, WHISPER_SAMPLE_RATE, 2*WHISPER_N_FFT, 2*WHISPER_HOP_LENGTH, WHISPER_N_MEL, n_threads, ctx->model.filters, true, ctx->mel)) {
         fprintf(stderr, "%s: failed to compute mel spectrogram\n", __func__);
         return -1;
     }
@@ -2343,6 +2391,7 @@ struct whisper_full_params whisper_full_default_params(enum whisper_sampling_str
 
                     /*.translate            =*/ false,
                     /*.no_context           =*/ false,
+                    /*.single_segment       =*/ false,
                     /*.print_special_tokens =*/ false,
                     /*.print_progress       =*/ true,
                     /*.print_realtime       =*/ false,
@@ -2352,6 +2401,10 @@ struct whisper_full_params whisper_full_default_params(enum whisper_sampling_str
                     /*.thold_pt             =*/ 0.01f,
                     /*.thold_ptsum          =*/ 0.01f,
                     /*.max_len              =*/ 0,
+                    /*.max_tokens           =*/ 0,
+
+                    /*.speed_up             =*/ false,
+                    /*.audio_ctx            =*/ 0,
 
                     /*.language             =*/ "en",
 
@@ -2381,6 +2434,7 @@ struct whisper_full_params whisper_full_default_params(enum whisper_sampling_str
 
                     /*.translate            =*/ false,
                     /*.no_context           =*/ false,
+                    /*.single_segment       =*/ false,
                     /*.print_special_tokens =*/ false,
                     /*.print_progress       =*/ true,
                     /*.print_realtime       =*/ false,
@@ -2390,6 +2444,10 @@ struct whisper_full_params whisper_full_default_params(enum whisper_sampling_str
                     /*.thold_pt             =*/ 0.01f,
                     /*.thold_ptsum          =*/ 0.01f,
                     /*.max_len              =*/ 0,
+                    /*.max_tokens           =*/ 0,
+
+                    /*.speed_up             =*/ false,
+                    /*.audio_ctx            =*/ 0,
 
                     /*.language             =*/ "en",
 
@@ -2485,9 +2543,16 @@ int whisper_full(
     result_all.clear();
 
     // compute log mel spectrogram
-    if (whisper_pcm_to_mel(ctx, samples, n_samples, params.n_threads) != 0) {
-        fprintf(stderr, "%s: failed to compute log mel spectrogram\n", __func__);
-        return -1;
+    if (params.speed_up) {
+        if (whisper_pcm_to_mel_phase_vocoder(ctx, samples, n_samples, params.n_threads) != 0) {
+            fprintf(stderr, "%s: failed to compute log mel spectrogram\n", __func__);
+            return -1;
+        }
+    } else {
+        if (whisper_pcm_to_mel(ctx, samples, n_samples, params.n_threads) != 0) {
+            fprintf(stderr, "%s: failed to compute log mel spectrogram\n", __func__);
+            return -1;
+        }
     }
 
     if (params.token_timestamps) {
@@ -2513,6 +2578,9 @@ int whisper_full(
         prompt_past.clear();
     }
 
+    // overwrite audio_ctx
+    ctx->exp_n_audio_ctx = params.audio_ctx;
+
     // these tokens determine the task that will be performed
     std::vector<whisper_token> prompt_init = { whisper_token_sot(ctx) };
     if (whisper_is_multilingual(ctx)) {
@@ -2623,7 +2691,7 @@ int whisper_full(
                 //}
 
                 // end of text token
-                if (token.id == whisper_token_eot(ctx)) {
+                if (token.id == whisper_token_eot(ctx) || (params.max_tokens > 0 && i > params.max_tokens)) {
                     if (result_len == 0) {
                         if (seek + seek_delta + 100 >= seek_end) {
                             result_len = i + 1;
@@ -2632,6 +2700,12 @@ int whisper_full(
                             fprintf(stderr, "\n%s: failed to generate timestamp token - this should not happen\n\n", __func__);
                         }
                     }
+
+                    if (params.single_segment) {
+                        result_len = i + 1;
+                        seek_delta = 100*WHISPER_CHUNK_SIZE;
+                    }
+
                     break;
                 }
 
@@ -2673,16 +2747,19 @@ int whisper_full(
                 if (tokens_cur[i].id > whisper_token_beg(ctx)) {
                     const auto t1 = seek + 2*(tokens_cur[i].tid - whisper_token_beg(ctx));
                     if (!text.empty()) {
+                        const auto tt0 = params.speed_up ? 2*t0 : t0;
+                        const auto tt1 = params.speed_up ? 2*t1 : t1;
+
                         if (params.print_realtime) {
                             if (params.print_timestamps) {
-                                printf("[%s --> %s]  %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text.c_str());
+                                printf("[%s --> %s]  %s\n", to_timestamp(tt0).c_str(), to_timestamp(tt1).c_str(), text.c_str());
                             } else {
                                 printf("%s", text.c_str());
                                 fflush(stdout);
                             }
                         }
 
-                        result_all.push_back({ t0, t1, text, {} });
+                        result_all.push_back({ tt0, tt1, text, {} });
                         for (int j = i0; j <= i; j++) {
                             result_all.back().tokens.push_back(tokens_cur[j]);
                         }
@@ -2714,16 +2791,19 @@ int whisper_full(
             if (!text.empty()) {
                 const auto t1 = seek + seek_delta;
 
+                const auto tt0 = params.speed_up ? 2*t0 : t0;
+                const auto tt1 = params.speed_up ? 2*t1 : t1;
+
                 if (params.print_realtime) {
                     if (params.print_timestamps) {
-                        printf("[%s --> %s]  %s\n", to_timestamp(t0).c_str(), to_timestamp(t1).c_str(), text.c_str());
+                        printf("[%s --> %s]  %s\n", to_timestamp(tt0).c_str(), to_timestamp(tt1).c_str(), text.c_str());
                     } else {
                         printf("%s", text.c_str());
                         fflush(stdout);
                     }
                 }
 
-                result_all.push_back({ t0, t1, text, {} });
+                result_all.push_back({ tt0, tt1, text, {} });
                 for (int j = i0; j < (int) tokens_cur.size(); j++) {
                     result_all.back().tokens.push_back(tokens_cur[j]);
                 }
@@ -2805,7 +2885,7 @@ int whisper_full_parallel(
 
             // key/value memory for the cross-attention layer
             {
-                const int n_audio_ctx   = hparams.n_audio_ctx;
+                const int n_audio_ctx = hparams.n_audio_ctx;
 
                 const int n_mem      = n_text_layer*n_audio_ctx;
                 const int n_elements = n_text_state*n_mem;

whisper.h

@@ -191,6 +191,7 @@ extern "C" {
 
         bool translate;
         bool no_context;
+        bool single_segment; // force single segment output (useful for streaming)
         bool print_special_tokens;
         bool print_progress;
         bool print_realtime;
@@ -201,6 +202,11 @@ extern "C" {
         float thold_pt;         // timestamp token probability threshold (~0.01)
         float thold_ptsum;      // timestamp token sum probability threshold (~0.01)
         int   max_len;          // max segment length in characters
+        int   max_tokens;       // max tokens per segment (0 = no limit)
+
+        // [EXPERIMENTAL] speed-up techniques
+        bool speed_up;  // speed-up the audio by 2x using Phase Vocoder
+        int  audio_ctx; // overwrite the audio context size (0 = use default)
 
         const char * language;