whisper.rn 0.4.0-rc.4 → 0.4.0-rc.5

package/android/src/main/jni.cpp

@@ -10,6 +10,7 @@
 #include "whisper.h"
 #include "rn-whisper.h"
 #include "ggml.h"
+#include "jni-utils.h"
 
 #define UNUSED(x) (void)(x)
 #define TAG "JNI"
@@ -96,7 +97,8 @@ static void input_stream_close(void *ctx) {
 
 static struct whisper_context *whisper_init_from_input_stream(
     JNIEnv *env,
-    jobject input_stream // PushbackInputStream
+    jobject input_stream, // PushbackInputStream
+    struct whisper_context_params cparams
 ) {
     input_stream_context *context = new input_stream_context;
     context->env = env;
@@ -108,7 +110,7 @@ static struct whisper_context *whisper_init_from_input_stream(
         .eof = &input_stream_is_eof,
         .close = &input_stream_close
     };
-    return whisper_init(&loader);
+    return whisper_init_with_params(&loader, cparams);
 }
 
 // Load model from asset
@@ -127,7 +129,8 @@ static void asset_close(void *ctx) {
 static struct whisper_context *whisper_init_from_asset(
     JNIEnv *env,
     jobject assetManager,
-    const char *asset_path
+    const char *asset_path,
+    struct whisper_context_params cparams
 ) {
     LOGI("Loading model from asset '%s'\n", asset_path);
     AAssetManager *asset_manager = AAssetManager_fromJava(env, assetManager);
@@ -142,7 +145,7 @@ static struct whisper_context *whisper_init_from_asset(
         .eof = &asset_is_eof,
         .close = &asset_close
     };
-    return whisper_init(&loader);
+    return whisper_init_with_params(&loader, cparams);
 }
 
 extern "C" {
@@ -151,9 +154,10 @@ JNIEXPORT jlong JNICALL
 Java_com_rnwhisper_WhisperContext_initContext(
         JNIEnv *env, jobject thiz, jstring model_path_str) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
     struct whisper_context *context = nullptr;
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
-    context = whisper_init_from_file(model_path_chars);
+    context = whisper_init_from_file_with_params(model_path_chars, cparams);
     env->ReleaseStringUTFChars(model_path_str, model_path_chars);
     return reinterpret_cast<jlong>(context);
 }
@@ -166,9 +170,10 @@ Java_com_rnwhisper_WhisperContext_initContextWithAsset(
     jstring model_path_str
 ) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
     struct whisper_context *context = nullptr;
     const char *model_path_chars = env->GetStringUTFChars(model_path_str, nullptr);
-    context = whisper_init_from_asset(env, asset_manager, model_path_chars);
+    context = whisper_init_from_asset(env, asset_manager, model_path_chars, cparams);
     env->ReleaseStringUTFChars(model_path_str, model_path_chars);
     return reinterpret_cast<jlong>(context);
 }
@@ -180,30 +185,65 @@ Java_com_rnwhisper_WhisperContext_initContextWithInputStream(
     jobject input_stream
 ) {
     UNUSED(thiz);
+    struct whisper_context_params cparams;
     struct whisper_context *context = nullptr;
-    context = whisper_init_from_input_stream(env, input_stream);
+    context = whisper_init_from_input_stream(env, input_stream, cparams);
     return reinterpret_cast<jlong>(context);
 }
 
-JNIEXPORT jboolean JNICALL
-Java_com_rnwhisper_WhisperContext_vadSimple(
-    JNIEnv *env,
-    jobject thiz,
-    jfloatArray audio_data,
-    jint audio_data_len,
-    jfloat vad_thold,
-    jfloat vad_freq_thold
-) {
-    UNUSED(thiz);
 
-    std::vector<float> samples(audio_data_len);
-    jfloat *audio_data_arr = env->GetFloatArrayElements(audio_data, nullptr);
-    for (int i = 0; i < audio_data_len; i++) {
-        samples[i] = audio_data_arr[i];
+struct whisper_full_params createFullParams(JNIEnv *env, jobject options) {
+    struct whisper_full_params params = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
+
+    params.print_realtime = false;
+    params.print_progress = false;
+    params.print_timestamps = false;
+    params.print_special = false;
+    
+    int max_threads = std::thread::hardware_concurrency();
+    // Use 2 threads by default on 4-core devices, 4 threads on more cores
+    int default_n_threads = max_threads == 4 ? 2 : min(4, max_threads);
+    int n_threads = readablemap::getInt(env, options, "maxThreads", default_n_threads);
+    params.n_threads = n_threads > 0 ? n_threads : default_n_threads;
+    params.translate = readablemap::getBool(env, options, "translate", false);
+    params.speed_up = readablemap::getBool(env, options, "speedUp", false);
+    params.token_timestamps = readablemap::getBool(env, options, "tokenTimestamps", false);
+    params.offset_ms = 0;
+    params.no_context = true;
+    params.single_segment = false;
+
+    int beam_size = readablemap::getInt(env, options, "beamSize", -1);
+    if (beam_size > -1) {
+        params.strategy = WHISPER_SAMPLING_BEAM_SEARCH;
+        params.beam_search.beam_size = beam_size;
     }
-    bool is_speech = rn_whisper_vad_simple(samples, WHISPER_SAMPLE_RATE, 1000, vad_thold, vad_freq_thold, false);
-    env->ReleaseFloatArrayElements(audio_data, audio_data_arr, JNI_ABORT);
-    return is_speech;
+    int best_of = readablemap::getInt(env, options, "bestOf", -1);
+    if (best_of > -1) params.greedy.best_of = best_of;
+    int max_len = readablemap::getInt(env, options, "maxLen", -1);
+    if (max_len > -1) params.max_len = max_len;
+    int max_context = readablemap::getInt(env, options, "maxContext", -1);
+    if (max_context > -1) params.n_max_text_ctx = max_context;
+    int offset = readablemap::getInt(env, options, "offset", -1);
+    if (offset > -1) params.offset_ms = offset;
+    int duration = readablemap::getInt(env, options, "duration", -1);
+    if (duration > -1) params.duration_ms = duration;
+    int word_thold = readablemap::getInt(env, options, "wordThold", -1);
+    if (word_thold > -1) params.thold_pt = word_thold;
+    float temperature = readablemap::getFloat(env, options, "temperature", -1);
+    if (temperature > -1) params.temperature = temperature;
+    float temperature_inc = readablemap::getFloat(env, options, "temperatureInc", -1);
+    if (temperature_inc > -1) params.temperature_inc = temperature_inc;
+    jstring prompt = readablemap::getString(env, options, "prompt", nullptr);
+    if (prompt != nullptr) {
+        params.initial_prompt = env->GetStringUTFChars(prompt, nullptr);
+        env->DeleteLocalRef(prompt);
+    }
+    jstring language = readablemap::getString(env, options, "language", nullptr);
+    if (language != nullptr) {
+        params.language = env->GetStringUTFChars(language, nullptr);
+        env->DeleteLocalRef(language);
+    }
+    return params;
 }
 
 struct callback_context {
@@ -212,102 +252,23 @@ struct callback_context {
 };
 
 JNIEXPORT jint JNICALL
-Java_com_rnwhisper_WhisperContext_fullTranscribe(
+Java_com_rnwhisper_WhisperContext_fullWithNewJob(
     JNIEnv *env,
     jobject thiz,
     jint job_id,
     jlong context_ptr,
     jfloatArray audio_data,
     jint audio_data_len,
-    jint n_threads,
-    jint max_context,
-    int word_thold,
-    int max_len,
-    jboolean token_timestamps,
-    jint offset,
-    jint duration,
-    jfloat temperature,
-    jfloat temperature_inc,
-    jint beam_size,
-    jint best_of,
-    jboolean speed_up,
-    jboolean translate,
-    jstring language,
-    jstring prompt,
+    jobject options,
     jobject callback_instance
 ) {
     UNUSED(thiz);
     struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
     jfloat *audio_data_arr = env->GetFloatArrayElements(audio_data, nullptr);
 
-    int max_threads = std::thread::hardware_concurrency();
-    // Use 2 threads by default on 4-core devices, 4 threads on more cores
-    int default_n_threads = max_threads == 4 ? 2 : min(4, max_threads);
-
     LOGI("About to create params");
 
-    struct whisper_full_params params = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
-
-    if (beam_size > -1) {
-        params.strategy = WHISPER_SAMPLING_BEAM_SEARCH;
-        params.beam_search.beam_size = beam_size;
-    }
-
-    params.print_realtime = false;
-    params.print_progress = false;
-    params.print_timestamps = false;
-    params.print_special = false;
-    params.translate = translate;
-    const char *language_chars = env->GetStringUTFChars(language, nullptr);
-    params.language = language_chars;
-    params.n_threads = n_threads > 0 ? n_threads : default_n_threads;
-    params.speed_up = speed_up;
-    params.offset_ms = 0;
-    params.no_context = true;
-    params.single_segment = false;
-
-    if (max_len > -1) {
-        params.max_len = max_len;
-    }
-    params.token_timestamps = token_timestamps;
-
-    if (best_of > -1) {
-        params.greedy.best_of = best_of;
-    }
-    if (max_context > -1) {
-        params.n_max_text_ctx = max_context;
-    }
-    if (offset > -1) {
-        params.offset_ms = offset;
-    }
-    if (duration > -1) {
-        params.duration_ms = duration;
-    }
-    if (word_thold > -1) {
-        params.thold_pt = word_thold;
-    }
-    if (temperature > -1) {
-        params.temperature = temperature;
-    }
-    if (temperature_inc > -1) {
-        params.temperature_inc = temperature_inc;
-    }
-    if (prompt != nullptr) {
-        params.initial_prompt = env->GetStringUTFChars(prompt, nullptr);
-    }
-
-    // abort handlers
-    bool* abort_ptr = rn_whisper_assign_abort_map(job_id);
-    params.encoder_begin_callback = [](struct whisper_context * /*ctx*/, struct whisper_state * /*state*/, void * user_data) {
-        bool is_aborted = *(bool*)user_data;
-        return !is_aborted;
-    };
-    params.encoder_begin_callback_user_data = abort_ptr;
-    params.abort_callback = [](void * user_data) {
-        bool is_aborted = *(bool*)user_data;
-        return is_aborted;
-    };
-    params.abort_callback_user_data = abort_ptr;
+    whisper_full_params params = createFullParams(env, options);
 
     if (callback_instance != nullptr) {
         callback_context *cb_ctx = new callback_context;
@@ -335,6 +296,8 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
         params.new_segment_callback_user_data = cb_ctx;
     }
 
+    rnwhisper::job* job = rnwhisper::job_new(job_id, params);
+
     LOGI("About to reset timings");
     whisper_reset_timings(context);
 
@@ -344,11 +307,122 @@ Java_com_rnwhisper_WhisperContext_fullTranscribe(
         // whisper_print_timings(context);
     }
     env->ReleaseFloatArrayElements(audio_data, audio_data_arr, JNI_ABORT);
-    env->ReleaseStringUTFChars(language, language_chars);
-    if (rn_whisper_transcribe_is_aborted(job_id)) {
-        code = -999;
+    
+    if (job->is_aborted()) code = -999;
+    rnwhisper::job_remove(job_id);
+    return code;
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_createRealtimeTranscribeJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jobject options
+) {
+    whisper_full_params params = createFullParams(env, options);
+    rnwhisper::job* job = rnwhisper::job_new(job_id, params);
+    rnwhisper::vad_params vad;
+    vad.use_vad = readablemap::getBool(env, options, "useVad", false);
+    vad.vad_ms = readablemap::getInt(env, options, "vadMs", 2000);
+    vad.vad_thold = readablemap::getFloat(env, options, "vadThold", 0.6f);
+    vad.freq_thold = readablemap::getFloat(env, options, "vadFreqThold", 100.0f);
+
+    jstring audio_output_path = readablemap::getString(env, options, "audioOutputPath", nullptr);
+    const char* audio_output_path_str = nullptr;
+    if (audio_output_path != nullptr) {
+        audio_output_path_str = env->GetStringUTFChars(audio_output_path, nullptr);
+        env->DeleteLocalRef(audio_output_path);
+    }
+    job->set_realtime_params(
+        vad,
+        readablemap::getInt(env, options, "realtimeAudioSec", 0),
+        readablemap::getInt(env, options, "realtimeAudioSliceSec", 0),
+        audio_output_path_str
+    );
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_finishRealtimeTranscribeJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jintArray slice_n_samples
+) {
+    UNUSED(env);
+    UNUSED(thiz);
+    UNUSED(context_ptr);
+
+    rnwhisper::job *job = rnwhisper::job_get(job_id);
+    if (job->audio_output_path != nullptr) {
+        RNWHISPER_LOG_INFO("job->params.language: %s\n", job->params.language);
+        std::vector<int> slice_n_samples_vec;
+        jint *slice_n_samples_arr = env->GetIntArrayElements(slice_n_samples, nullptr);
+        slice_n_samples_vec = std::vector<int>(slice_n_samples_arr, slice_n_samples_arr + env->GetArrayLength(slice_n_samples));
+        env->ReleaseIntArrayElements(slice_n_samples, slice_n_samples_arr, JNI_ABORT);
+
+        // TODO: Append in real time so we don't need to keep all slices & also reduce memory usage
+        rnaudioutils::save_wav_file(
+            rnaudioutils::concat_short_buffers(job->pcm_slices, slice_n_samples_vec),
+            job->audio_output_path
+        );
+    }
+    rnwhisper::job_remove(job_id);
+}
+
+JNIEXPORT jboolean JNICALL
+Java_com_rnwhisper_WhisperContext_vadSimple(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jint slice_index,
+    jint n_samples,
+    jint n
+) {
+    UNUSED(thiz);
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    return job->vad_simple(slice_index, n_samples, n);
+}
+
+JNIEXPORT void JNICALL
+Java_com_rnwhisper_WhisperContext_putPcmData(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jshortArray pcm,
+    jint slice_index,
+    jint n_samples,
+    jint n
+) {
+    UNUSED(thiz);
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    jshort *pcm_arr = env->GetShortArrayElements(pcm, nullptr);
+    job->put_pcm_data(pcm_arr, slice_index, n_samples, n);
+    env->ReleaseShortArrayElements(pcm, pcm_arr, JNI_ABORT);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_rnwhisper_WhisperContext_fullWithJob(
+    JNIEnv *env,
+    jobject thiz,
+    jint job_id,
+    jlong context_ptr,
+    jint slice_index,
+    jint n_samples
+) {
+    UNUSED(thiz);
+    struct whisper_context *context = reinterpret_cast<struct whisper_context *>(context_ptr);
+
+    rnwhisper::job* job = rnwhisper::job_get(job_id);
+    float* pcmf32 = job->pcm_slice_to_f32(slice_index, n_samples);
+    int code = whisper_full(context, job->params, pcmf32, n_samples);
+    free(pcmf32);
+    if (code == 0) {
+        // whisper_print_timings(context);
     }
-    rn_whisper_remove_abort_map(job_id);
+    if (job->is_aborted()) code = -999;
     return code;
 }
 
@@ -359,7 +433,8 @@ Java_com_rnwhisper_WhisperContext_abortTranscribe(
     jint job_id
 ) {
     UNUSED(thiz);
-    rn_whisper_abort_transcribe(job_id);
+    rnwhisper::job *job = rnwhisper::job_get(job_id);
+    if (job) job->abort();
 }
 
 JNIEXPORT void JNICALL
@@ -368,7 +443,7 @@ Java_com_rnwhisper_WhisperContext_abortAllTranscribe(
     jobject thiz
 ) {
     UNUSED(thiz);
-    rn_whisper_abort_all_transcribe();
+    rnwhisper::job_abort_all();
 }
 
 JNIEXPORT jint JNICALL

package/cpp/README.md

@@ -1,4 +1,4 @@
 # Note
 
-- Only `rn-whisper.h` / `rn-whisper.cpp` are the specific files for this project, others are sync from [whisper.cpp](https://github.com/ggerganov/whisper.cpp).
+- Only `rn-*` are the specific files for this project, others are sync from [whisper.cpp](https://github.com/ggerganov/whisper.cpp).
 - We can update the native source by using the [bootstrap](../scripts/bootstrap.sh) script.

package/cpp/coreml/whisper-encoder-impl.h

@@ -123,7 +123,7 @@ API_AVAILABLE(macos(12.0), ios(15.0), watchos(8.0), tvos(15.0)) __attribute__((v
 
 /**
     Make a prediction using the convenience interface
-    @param logmel_data as 1 × 80 × 3000 3-dimensional array of floats:
+    @param logmel_data as 1 × n_mel × 3000 3-dimensional array of floats:
     @param error If an error occurs, upon return contains an NSError object that describes the problem. If you are not interested in possible errors, pass in NULL.
     @return the prediction as whisper_encoder_implOutput
 */

package/cpp/coreml/whisper-encoder.h

@@ -3,6 +3,8 @@
 // Code is derived from the work of Github user @wangchou
 // ref: https://github.com/wangchou/callCoreMLFromCpp
 
+#include <stdint.h>
+
 #if __cplusplus
 extern "C" {
 #endif
@@ -14,6 +16,8 @@ void whisper_coreml_free(struct whisper_coreml_context * ctx);
 
 void whisper_coreml_encode(
         const whisper_coreml_context * ctx,
+                             int64_t   n_ctx,
+                             int64_t   n_mel,
                                float * mel,
                                float * out);
 

package/cpp/coreml/whisper-encoder.mm

@@ -48,13 +48,15 @@ void whisper_coreml_free(struct whisper_coreml_context * ctx) {
 
 void whisper_coreml_encode(
         const whisper_coreml_context * ctx,
+                             int64_t   n_ctx,
+                             int64_t   n_mel,
                                float * mel,
                                float * out) {
     MLMultiArray * inMultiArray = [
         [MLMultiArray alloc] initWithDataPointer: mel
-                                           shape: @[@1, @80, @3000]
+                                           shape: @[@1, @(n_mel), @(n_ctx)]
                                         dataType: MLMultiArrayDataTypeFloat32
-                                         strides: @[@(240000), @(3000), @1]
+                                         strides: @[@(n_ctx*n_mel), @(n_ctx), @1]
                                      deallocator: nil
                                            error: nil
     ];

package/cpp/ggml-alloc.c

@@ -137,7 +137,7 @@ void wsp_ggml_tallocr_alloc(wsp_ggml_tallocr_t alloc, struct wsp_ggml_tensor * t
 
 #ifdef WSP_GGML_ALLOCATOR_DEBUG
     add_allocated_tensor(alloc, tensor);
-    size_t cur_max = (char*)addr - (char*)alloc->data + size;
+    size_t cur_max = (char*)addr - (char*)alloc->base + size;
     if (cur_max > alloc->max_size) {
         printf("max_size = %.2f MB: tensors: ", cur_max / 1024.0 / 1024.0);
         for (int i = 0; i < 1024; i++) {
@@ -168,10 +168,6 @@ static void wsp_ggml_tallocr_free_tensor(wsp_ggml_tallocr_t alloc, struct wsp_gg
     size = aligned_offset(NULL, size, alloc->alignment);
     AT_PRINTF("%s: freeing %s at %p (%zu bytes) - n_free_blocks = %d\n", __func__, tensor->name, ptr, size, alloc->n_free_blocks);
 
-    if (!alloc->measure) {
-        wsp_ggml_backend_buffer_free_tensor(alloc->buffer, tensor);
-    }
-
 #ifdef WSP_GGML_ALLOCATOR_DEBUG
     remove_allocated_tensor(alloc, tensor);
 #endif
@@ -237,7 +233,7 @@ void wsp_ggml_tallocr_reset(wsp_ggml_tallocr_t alloc) {
 }
 
 wsp_ggml_tallocr_t wsp_ggml_tallocr_new(void * data, size_t size, size_t alignment) {
-    struct wsp_ggml_backend_buffer * buffer = wsp_ggml_backend_cpu_buffer_from_ptr(NULL, data, size);
+    struct wsp_ggml_backend_buffer * buffer = wsp_ggml_backend_cpu_buffer_from_ptr(data, size);
 
     wsp_ggml_tallocr_t alloc = (wsp_ggml_tallocr_t)malloc(sizeof(struct wsp_ggml_tallocr));
 
@@ -446,18 +442,19 @@ static wsp_ggml_tallocr_t node_tallocr(wsp_ggml_gallocr_t galloc, struct wsp_ggm
     return galloc->hash_allocs[wsp_ggml_hash_find_or_insert(galloc->hash_set, node)];
 }
 
-static void init_view(wsp_ggml_gallocr_t galloc, struct wsp_ggml_tensor * view) {
+static void init_view(wsp_ggml_gallocr_t galloc, struct wsp_ggml_tensor * view, bool update_backend) {
     wsp_ggml_tallocr_t alloc = node_tallocr(galloc, view);
 
-    //printf("init_view: %s from src %s\n", view->name, view->view_src->name);
     WSP_GGML_ASSERT(view->view_src != NULL && view->view_src->data != NULL);
-    view->backend = view->view_src->backend;
+    if (update_backend) {
+        view->backend = view->view_src->backend;
+    }
     view->buffer  = view->view_src->buffer;
     view->data    = (char *)view->view_src->data + view->view_offs;
 
     // FIXME: the view should be initialized by the owning buffer, but currently this breaks the CUDA backend
     // due to the wsp_ggml_tensor_extra_gpu ring buffer overwriting the KV cache extras
-    assert(wsp_ggml_tallocr_is_measure(alloc) || !view->buffer || view->buffer->backend == alloc->buffer->backend);
+    assert(wsp_ggml_tallocr_is_measure(alloc) || !view->buffer || view->buffer->buft == alloc->buffer->buft);
 
     if (!alloc->measure) {
         wsp_ggml_backend_buffer_init_tensor(alloc->buffer, view);
@@ -469,7 +466,7 @@ static void allocate_node(wsp_ggml_gallocr_t galloc, struct wsp_ggml_tensor * no
 
     if (node->data == NULL) {
         if (wsp_ggml_is_view(node)) {
-            init_view(galloc, node);
+            init_view(galloc, node, true);
         } else {
             // see if we can reuse a parent's buffer (inplace)
             if (wsp_ggml_op_can_inplace(node->op)) {
@@ -499,15 +496,14 @@ static void allocate_node(wsp_ggml_gallocr_t galloc, struct wsp_ggml_tensor * no
                                 AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name);
                                 node->view_src = view_src;
                                 view_src_hn->n_views += 1;
-                                init_view(galloc, node);
+                                init_view(galloc, node, false);
                                 return;
                             }
-                        }
-                        else {
+                        } else {
                             AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name);
                             node->view_src = parent;
                             p_hn->n_views += 1;
-                            init_view(galloc, node);
+                            init_view(galloc, node, false);
                             return;
                         }
                     }
@@ -537,7 +533,7 @@ static void wsp_ggml_tallocr_alloc_graph_impl(wsp_ggml_gallocr_t galloc, struct
             hash_get(galloc, view_src)->n_views += 1;
             if (node->buffer == NULL && node->data != NULL) {
                 // view of a pre-allocated tensor, didn't call init_view() yet
-                init_view(galloc, node);
+                init_view(galloc, node, true);
             }
         }
 
@@ -548,7 +544,7 @@ static void wsp_ggml_tallocr_alloc_graph_impl(wsp_ggml_gallocr_t galloc, struct
             }
             hash_get(galloc, parent)->n_children += 1;
             if (wsp_ggml_is_view(parent) && parent->buffer == NULL && parent->data != NULL) {
-                init_view(galloc, parent);
+                init_view(galloc, parent, true);
             }
         }
    }
@@ -663,7 +659,7 @@ size_t wsp_ggml_gallocr_alloc_graph(wsp_ggml_gallocr_t galloc, wsp_ggml_tallocr_
     return max_size;
 }
 
-void wsp_ggml_gallocr_alloc_graph_n(wsp_ggml_gallocr_t galloc, struct wsp_ggml_cgraph * graph, struct wsp_ggml_hash_set hash_set, wsp_ggml_tallocr_t * hash_node_alloct) {
+void wsp_ggml_gallocr_alloc_graph_n(wsp_ggml_gallocr_t galloc, struct wsp_ggml_cgraph * graph, struct wsp_ggml_hash_set hash_set, wsp_ggml_tallocr_t * hash_node_talloc) {
     const size_t hash_size = hash_set.size;
 
     WSP_GGML_ASSERT(hash_size >= (size_t)(graph->n_nodes + graph->n_leafs));
@@ -686,7 +682,7 @@ void wsp_ggml_gallocr_alloc_graph_n(wsp_ggml_gallocr_t galloc, struct wsp_ggml_c
     // reset hash values
     memset(galloc->hash_values, 0, sizeof(struct hash_node) * hash_size);
 
-    galloc->hash_allocs = hash_node_alloct;
+    galloc->hash_allocs = hash_node_talloc;
 
     wsp_ggml_tallocr_alloc_graph_impl(galloc, graph);
 
@@ -764,3 +760,43 @@ size_t wsp_ggml_allocr_max_size(wsp_ggml_allocr_t alloc) {
 size_t wsp_ggml_allocr_alloc_graph(wsp_ggml_allocr_t alloc, struct wsp_ggml_cgraph * graph) {
     return wsp_ggml_gallocr_alloc_graph(alloc->galloc, alloc->talloc, graph);
 }
+
+// utils
+wsp_ggml_backend_buffer_t wsp_ggml_backend_alloc_ctx_tensors_from_buft(struct wsp_ggml_context * ctx, wsp_ggml_backend_buffer_type_t buft) {
+    WSP_GGML_ASSERT(wsp_ggml_get_no_alloc(ctx) == true);
+
+    size_t alignment = wsp_ggml_backend_buft_get_alignment(buft);
+
+    size_t nbytes = 0;
+    for (struct wsp_ggml_tensor * t = wsp_ggml_get_first_tensor(ctx); t != NULL; t = wsp_ggml_get_next_tensor(ctx, t)) {
+        if (t->data == NULL && t->view_src == NULL) {
+            nbytes += WSP_GGML_PAD(wsp_ggml_backend_buft_get_alloc_size(buft, t), alignment);
+        }
+    }
+
+    if (nbytes == 0) {
+        fprintf(stderr, "%s: no tensors to allocate\n", __func__);
+        return NULL;
+    }
+
+    wsp_ggml_backend_buffer_t buffer = wsp_ggml_backend_buft_alloc_buffer(buft, nbytes);
+    wsp_ggml_tallocr_t tallocr = wsp_ggml_tallocr_new_from_buffer(buffer);
+
+    for (struct wsp_ggml_tensor * t = wsp_ggml_get_first_tensor(ctx); t != NULL; t = wsp_ggml_get_next_tensor(ctx, t)) {
+        if (t->data == NULL) {
+            if (t->view_src == NULL) {
+                wsp_ggml_tallocr_alloc(tallocr, t);
+            } else {
+                wsp_ggml_backend_view_init(buffer, t);
+            }
+        }
+    }
+
+    wsp_ggml_tallocr_free(tallocr);
+
+    return buffer;
+}
+
+wsp_ggml_backend_buffer_t wsp_ggml_backend_alloc_ctx_tensors(struct wsp_ggml_context * ctx, wsp_ggml_backend_t backend) {
+    return wsp_ggml_backend_alloc_ctx_tensors_from_buft(ctx, wsp_ggml_backend_get_default_buffer_type(backend));
+}

package/cpp/ggml-alloc.h

@@ -8,6 +8,7 @@ extern "C" {
 
 struct wsp_ggml_backend;
 struct wsp_ggml_backend_buffer;
+struct wsp_ggml_backend_buffer_type;
 
 //
 // Legacy API
@@ -80,6 +81,12 @@ WSP_GGML_API void   wsp_ggml_gallocr_alloc_graph_n(
                     struct wsp_ggml_hash_set hash_set,
                     wsp_ggml_tallocr_t * hash_node_talloc);
 
+
+// Utils
+// Create a buffer and allocate all the tensors in a wsp_ggml_context
+WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors_from_buft(struct wsp_ggml_context * ctx, struct wsp_ggml_backend_buffer_type * buft);
+WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors(struct wsp_ggml_context * ctx, struct wsp_ggml_backend * backend);
+
 #ifdef  __cplusplus
 }
 #endif