[0.1.24] GPU memory: OCR runs largest first, gives its arenas' blocks back after each stage, and keeps its worker threads
OCR filled the GPU: 9 GB for the mobile model on gpu-bench's 30 images, past 16 GB for the server model, which Windows then paged into shared memory (1.5 to 11 s a capture). onnxruntime's logged arena extensions showed why: the arena keeps every block it takes from CUDA and never joins or frees them, and OCR's inputs grow from run to run (recognizer batches went narrowest first; images in any order across calls), so each larger run left another block behind: 5.7 GB of recognizer blocks where the widest batch needs 0.9 GB. - Detection runs the largest images first, recognition the widest batches. - The detector's and recognizer's arenas shrink once each stage of an ocr_process call is done, by a run on a small input with memory.enable_memory_arena_shrinkage. Shrinking after every run cost 33 to 59% of the speed. - The worker threads are made with the engine and kept, not made for every stage of every call (onnxruntime keeps GPU state for each thread that runs a session, microsoft/onnxruntime#33037); 6% faster. - OCR's arenas grow by what is asked (kSameAsRequested). - ort.log_level sets what the engines' onnxruntime environments log; INFO shows every arena extension. Mobile OCR on gpu-bench (RTX 5080 laptop, 4 workers): peak 5.7 GB over the baseline, 2 GB between calls, 142 ms a capture, against 9.0 GB and 126 ms. The text read differs by a line in 2,986 (batches are grouped differently). The server model is not yet measured again.
This commit is contained in:
parent
c351ca92c0
commit
8c8b2c6e38
19
README.md
19
README.md
@ -101,16 +101,15 @@ grid (in its order, so through a search's matches).
|
||||
- **Model:** PP-OCRv5 mobile (light, about 85 ms a capture on the GPU) or server (more accurate,
|
||||
about 0.3 s when it runs well; see below). The choice is kept in the settings file (`model` under
|
||||
`[ocr]`).
|
||||
- **Known issue, the server model on the GPU:** its speed is bimodal. The same program on the same
|
||||
20 captures (RTX 5080 laptop, 16 GB) runs either at about 0.27 s a capture with a peak of about
|
||||
2.4 GB of process memory, or at about 1.5 s with about 6.6 GB, with no change between runs and
|
||||
15 GB of the GPU's memory free beforehand. Longer runs fare worse: 50 captures peaked at 16 GB, and
|
||||
100 averaged about 4.6 s a capture. The extra memory looks like GPU allocations spilling into shared
|
||||
system memory, which Windows then pages: what the GPU memory defaults (onnxruntime's
|
||||
kNextPowerOfTwo arena, no cuDNN max workspace, 4 workers) were meant to stop, and which they
|
||||
only partly did. It is the same with Jai's allocator and with rpmalloc. Until it is fixed, scan a
|
||||
large library with the mobile model. Measured on 2026-10-13 with
|
||||
`ocr-bench -models ppocr-v5-server -gpu-only -limit 20`.
|
||||
- **GPU memory** (fixed in 0.1.24): OCR used to fill the GPU, 9 GB with the mobile model on
|
||||
gpu-bench's 30 images and past 16 GB with the server model, which Windows then paged into shared
|
||||
system memory, running it at 1.5 to 11 s a capture. onnxruntime's arena keeps every block it takes
|
||||
from CUDA, and OCR's inputs change size every run, so each larger image or wider batch of lines left
|
||||
another block behind. OCR now runs the largest images and widest batches first, has the detector's
|
||||
and recognizer's arenas give back what they hold after each stage of a call, and keeps its 4 worker
|
||||
threads for the engine's life. The mobile model now peaks at 5.7 GB and drops to 2 GB between
|
||||
calls, at 142 ms a capture against 126 before (gpu-bench, 2026-10-17). The server model has not been
|
||||
measured again yet.
|
||||
- **Running on:** the GPU or the CPU, and if the CPU, what to install for the GPU.
|
||||
- **Scan the library** reads every capture the model has not read yet and, with **Labels too**,
|
||||
labels every capture not yet labelled with the current list. While a search is active the button
|
||||
|
||||
@ -47,7 +47,7 @@ APP_NAME :: "mcap";
|
||||
APP_FOLDER_NAME :: "mcap";
|
||||
GIT_BRANCH :: "check";
|
||||
GIT_HASH :: "-";
|
||||
VERSION :: "0.1.23";
|
||||
VERSION :: "0.1.24";
|
||||
BUILD_DATE :: "-";
|
||||
DEBUG :: true;
|
||||
COMPILE_TESTS :: true;
|
||||
|
||||
@ -8,7 +8,7 @@ APP_NAME :: "mcap";
|
||||
APP_FOLDER_NAME :: "mcap";
|
||||
APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
|
||||
|
||||
VERSION :: "0.1.23";
|
||||
VERSION :: "0.1.24";
|
||||
// Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated
|
||||
// builds never write over the ones you run by hand.
|
||||
build_directory := "bin";
|
||||
|
||||
@ -102,7 +102,7 @@ labels_create :: (models_folder: string, model_name: string, prefer_gpu := true)
|
||||
engine := New(Labels_Engine);
|
||||
engine.model = model.*;
|
||||
api := ort.api;
|
||||
ok, error := ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "labels", *engine.env));
|
||||
ok, error := ort_check(api.CreateEnv(ort.log_level, "labels", *engine.env));
|
||||
if !ok { labels_destroy(engine); return null, false, error; }
|
||||
ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));
|
||||
|
||||
|
||||
@ -89,9 +89,11 @@ ocr_create_sessions :: (engine: *Ocr_Engine, models_folder: string, gpu: bool, i
|
||||
model := *engine.model;
|
||||
ok: bool;
|
||||
error: string;
|
||||
engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads, engine.options);
|
||||
// The detector's and recognizer's inputs change size with nearly every run: their arenas give back
|
||||
// what no run is using once per call (ocr_shrink_arena), or they keep a block of every size.
|
||||
engine.det, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.detector), gpu, intra_op_threads, engine.options, shrink_arena = true);
|
||||
if !ok return false, error;
|
||||
engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads, engine.options);
|
||||
engine.rec, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.recognizer), gpu, intra_op_threads, engine.options, shrink_arena = true);
|
||||
if !ok return false, error;
|
||||
if model.textline_ori {
|
||||
engine.orientation, ok, error = ocr_create_session(engine.env, tprint("%/%", models_folder, model.textline_ori), gpu, intra_op_threads, engine.options);
|
||||
@ -115,6 +117,19 @@ ocr_release_sessions :: (engine: *Ocr_Engine) {
|
||||
// engine destroyed without ever having run left the next engine to crash in cublasLtCtxInit, with a
|
||||
// DEP access violation, on its first run; an engine that has run once leaves nothing broken. It also
|
||||
// moves the first run's setup into creation. On failure, the reason.
|
||||
// On CUDA, the session's arena (the detector's or the recognizer's) gives back every block no run is
|
||||
// using: called by ocr_process once a stage's runs are all done, as a run on a small input, since
|
||||
// onnxruntime shrinks an arena only at the end of a run. Once per run instead cost a third to a half
|
||||
// of OCR's speed.
|
||||
ocr_shrink_arena :: (engine: *Ocr_Engine, session: *Ocr_Session) {
|
||||
if engine.provider != .CUDA || !session.session return;
|
||||
side := ifx session == *engine.rec then cast(s64) engine.model.rec_height else 32;
|
||||
shape := s64.[1, 3, side, side];
|
||||
input := NewArray(3 * side * side, float, initialized = true,, temp);
|
||||
out, ok := ocr_run(session, engine.memory_info, input, shape, shrink = true);
|
||||
if ok ocr_release_output(*out);
|
||||
}
|
||||
|
||||
ocr_warm_up :: (engine: *Ocr_Engine) -> ok: bool, error: string {
|
||||
Warm_Up :: struct { session: *Ocr_Session; height, width: s64; }
|
||||
runs := Warm_Up.[
|
||||
@ -135,8 +150,8 @@ ocr_warm_up :: (engine: *Ocr_Engine) -> ok: bool, error: string {
|
||||
}
|
||||
|
||||
// A session for the model file, set up as Ocr_Options says (modules/onnxruntime/session.jai).
|
||||
ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, ocr_options: Ocr_Options) -> Ocr_Session, ok: bool, error: string {
|
||||
settings := Ort_Gpu_Settings.{arena_extend_strategy = ocr_options.cuda_arena_extend_strategy, cudnn_max_workspace = ocr_options.cudnn_max_workspace};
|
||||
ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32, ocr_options: Ocr_Options, shrink_arena := false) -> Ocr_Session, ok: bool, error: string {
|
||||
settings := Ort_Gpu_Settings.{arena_extend_strategy = ocr_options.cuda_arena_extend_strategy, cudnn_max_workspace = ocr_options.cudnn_max_workspace, shrink_arena = shrink_arena};
|
||||
session, ok, error := ort_create_session(env, path, gpu, intra_op_threads, settings, missing_hint = "run `jai fetch_models.jai` in modules/OCR and rebuild");
|
||||
return session, ok, error;
|
||||
}
|
||||
|
||||
@ -47,6 +47,7 @@ Ocr_Engine :: struct {
|
||||
rec_batch: s32; // text lines per recognizer run: the model's rec_batch_gpu or rec_batch_cpu
|
||||
options: Ocr_Options;
|
||||
stats: Ocr_Stats; // since creation, or the last ocr_reset_stats
|
||||
pool: Ocr_Pool; // the worker threads, for the engine's life (run_on_workers)
|
||||
|
||||
// When set, and true, ocr_process stops starting work: the batch of lines or image each worker is
|
||||
// on is finished, and then every image fails as "cancelled" (none is part read). Set by the caller.
|
||||
@ -72,13 +73,13 @@ Ocr_Stats :: struct {
|
||||
ocr_stats :: (engine: *Ocr_Engine) -> Ocr_Stats { return engine.stats; }
|
||||
ocr_reset_stats :: (engine: *Ocr_Engine) { engine.stats = .{}; }
|
||||
|
||||
// How an engine runs on the GPU, beyond its model's models.config entry. The defaults are what ran
|
||||
// fastest on the 20 sample captures (RTX 5080 laptop, onnxruntime 1.30.0): nearly every image and
|
||||
// batch is a different size, and with onnxruntime's exact-size arena and cuDNN's largest workspaces
|
||||
// GPU memory filled up and the server model ran several times slower.
|
||||
// How an engine runs on the GPU, beyond its model's models.config entry. Nearly every image and
|
||||
// batch is a different size. onnxruntime's arena keeps each block it takes from CUDA and never joins
|
||||
// or frees them, so inputs that grow one run after another leave a block of every size behind: the
|
||||
// runs go largest first (ocr_process) so later ones fit in what the first took.
|
||||
Ocr_Options :: struct {
|
||||
gpu_workers: s32 = 4; // threads preparing images and decoding results while the GPU runs the networks; more let the server model fill GPU memory
|
||||
cuda_arena_extend_strategy := "kNextPowerOfTwo"; // reuses freed blocks for other sizes; "kSameAsRequested" grows by exactly what is asked
|
||||
gpu_workers: s32 = 4; // threads preparing images and decoding results while the GPU runs the networks
|
||||
cuda_arena_extend_strategy := "kSameAsRequested"; // how much the arena takes when it runs short: what is asked; "kNextPowerOfTwo" takes twice as much each time, up to twice what the runs need
|
||||
cudnn_max_workspace := false; // let cuDNN take as much working memory as it finds best (onnxruntime's own default)
|
||||
rec_width_step: s32 = 64; // a recognizer batch's width is rounded up to a multiple of this, so shapes recur
|
||||
}
|
||||
@ -109,7 +110,7 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op
|
||||
engine.dictionary = entries;
|
||||
|
||||
api := ort.api;
|
||||
ok, error := ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "ocr", *engine.env));
|
||||
ok, error := ort_check(api.CreateEnv(ort.log_level, "ocr", *engine.env));
|
||||
if !ok { ocr_destroy(engine); return null, false, error; }
|
||||
ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));
|
||||
|
||||
@ -149,6 +150,7 @@ ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true, op
|
||||
engine.provider = .CPU;
|
||||
engine.rec_batch = max(engine.model.rec_batch_cpu, 1);
|
||||
}
|
||||
ocr_pool_start(engine);
|
||||
return engine, true, "";
|
||||
}
|
||||
|
||||
@ -163,6 +165,7 @@ ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string {
|
||||
|
||||
ocr_destroy :: (engine: *Ocr_Engine) {
|
||||
if !engine return;
|
||||
ocr_pool_stop(engine);
|
||||
if engine.provider == .CUDA ort_gpu_lock();
|
||||
ocr_release_sessions(engine);
|
||||
if engine.provider == .CUDA ort_gpu_unlock();
|
||||
@ -185,13 +188,22 @@ ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.
|
||||
array_free(per_image,, context.default_allocator);
|
||||
}
|
||||
for * per_image it.image = *images[it_index];
|
||||
run_on_workers(engine, per_image.count, (engine: *Ocr_Engine, index: s64, data: *void) {
|
||||
work := cast(*Image_Work) data + index;
|
||||
// The largest first: the detector's arena then holds a block that fits every later image
|
||||
// (Ocr_Options). Smallest first, it took a new one for each larger image and kept them all.
|
||||
detect_order := NewArray(per_image.count, *Image_Work,, context.default_allocator);
|
||||
defer array_free(detect_order,, context.default_allocator);
|
||||
for * per_image detect_order[it_index] = it;
|
||||
quick_sort(detect_order, (a: *Image_Work, b: *Image_Work) -> s64 {
|
||||
return cast(s64) b.image.width * b.image.height - cast(s64) a.image.width * a.image.height;
|
||||
});
|
||||
run_on_workers(engine, detect_order.count, (engine: *Ocr_Engine, index: s64, data: *void) {
|
||||
work := (cast(**Image_Work) data)[index];
|
||||
start := current_time_monotonic();
|
||||
work.boxes, work.ok, work.error = ocr_detect(engine, work.image.*);
|
||||
atomic_add(*engine.stats.detect_us, to_microseconds(current_time_monotonic() - start));
|
||||
if work.error work.error = copy_string(work.error,, context.default_allocator);
|
||||
}, per_image.data);
|
||||
}, detect_order.data);
|
||||
ocr_shrink_arena(engine, *engine.det); // before recognition, so the two stages' blocks are not held at once
|
||||
if ocr_cancelled(engine) return ocr_all_cancelled(results, allocator);
|
||||
|
||||
// Recognition: every line of every image, batched by width so little of each batch is padding.
|
||||
@ -247,7 +259,10 @@ ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.
|
||||
order := NewArray(jobs.count, *Line_Job,, context.default_allocator);
|
||||
defer array_free(order,, context.default_allocator);
|
||||
for * jobs order[it_index] = it;
|
||||
quick_sort(order, (a: *Line_Job, b: *Line_Job) -> s64 { return a.width - b.width; });
|
||||
// Widest first, so each batch fits in the arena block the first one took (Ocr_Options). Narrowest
|
||||
// first, every batch was wider than the last and took a new block: 5.7 GB of them for the mobile
|
||||
// model on gpu-bench's 30 images, where the widest batch needs 0.9 GB.
|
||||
quick_sort(order, (a: *Line_Job, b: *Line_Job) -> s64 { return b.width - a.width; });
|
||||
|
||||
batch_size := max(engine.rec_batch, 1);
|
||||
batch_count := (order.count + batch_size - 1) / batch_size;
|
||||
@ -281,6 +296,8 @@ ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.
|
||||
}
|
||||
}
|
||||
|
||||
ocr_shrink_arena(engine, *engine.rec); // models.jai: what this call's runs took, given back
|
||||
|
||||
// Each image's lines, in reading order.
|
||||
for work, image_index: per_image {
|
||||
result := *results[image_index];
|
||||
@ -337,48 +354,93 @@ Image_Work :: struct {
|
||||
Work_Proc :: #type (engine: *Ocr_Engine, index: s64, data: *void);
|
||||
|
||||
// Calls proc for each index below count, on engine.workers threads at once.
|
||||
run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void) {
|
||||
if count <= 0 return;
|
||||
Shared :: struct { engine: *Ocr_Engine; count: s64; next: s64; proc: Work_Proc; data: *void; }
|
||||
shared := Shared.{engine = engine, count = count, proc = proc, data = data};
|
||||
// The engine's worker threads: started with it and kept until ocr_destroy, never one set per call.
|
||||
// onnxruntime (1.26 to 1.30 at least) sets up a CUDA stream with cuBLAS and cuDNN handles for every
|
||||
// OS thread that runs a session, about 10 MB each, kept after the thread ends (microsoft/onnxruntime
|
||||
// #33037). Workers made afresh for every stage of every call, as they were, gave the OCR sessions
|
||||
// thousands of threads in a scan, and starting them cost time: 6% on gpu-bench's mobile OCR.
|
||||
Ocr_Pool :: struct {
|
||||
threads: [] Thread;
|
||||
started: s64;
|
||||
mutex: Mutex; // one run_on_workers at a time
|
||||
wake: Semaphore; // signalled once per worker for each job, and for quitting
|
||||
done: Semaphore; // each worker signals it when it has finished its part of a job
|
||||
quit: bool;
|
||||
// The job: items 0..count-1, taken in turn.
|
||||
count, next: s64;
|
||||
proc: Work_Proc;
|
||||
data: *void;
|
||||
}
|
||||
|
||||
worker_proc :: (thread: *Thread) -> s64 {
|
||||
// The context comes from the calling thread, allocator included, and that may be its temporary
|
||||
// allocator, which belongs to that thread alone. Workers allocate from the heap, and use their
|
||||
// own temporary storage.
|
||||
ocr_pool_start :: (engine: *Ocr_Engine) {
|
||||
pool := *engine.pool;
|
||||
init(*pool.mutex);
|
||||
init(*pool.wake);
|
||||
init(*pool.done);
|
||||
// thread_init allocates the thread's temporary storage from the context allocator.
|
||||
push_allocator(context.default_allocator);
|
||||
pool.threads = NewArray(max(engine.workers, 1), Thread);
|
||||
for * pool.threads {
|
||||
if !thread_init(it, ocr_pool_worker, temporary_storage_size = 1024 * 1024) break;
|
||||
thread_start(it, engine); // sets thread.data
|
||||
pool.started += 1;
|
||||
}
|
||||
}
|
||||
|
||||
ocr_pool_stop :: (engine: *Ocr_Engine) {
|
||||
pool := *engine.pool;
|
||||
if !pool.threads.count return;
|
||||
pool.quit = true;
|
||||
for 1..pool.started signal(*pool.wake);
|
||||
for 0..pool.started - 1 {
|
||||
thread_is_done(*pool.threads[it], -1);
|
||||
thread_deinit(*pool.threads[it]);
|
||||
}
|
||||
array_free(pool.threads,, context.default_allocator);
|
||||
destroy(*pool.wake);
|
||||
destroy(*pool.done);
|
||||
destroy(*pool.mutex);
|
||||
pool.* = .{};
|
||||
}
|
||||
|
||||
ocr_pool_worker :: (thread: *Thread) -> s64 {
|
||||
// Workers allocate from the heap, and use their own temporary storage.
|
||||
context.allocator = context.default_allocator;
|
||||
shared := cast(*Shared) thread.data;
|
||||
engine := cast(*Ocr_Engine) thread.data;
|
||||
pool := *engine.pool;
|
||||
while true {
|
||||
if ocr_cancelled(shared.engine) break; // what is left fails as cancelled (ocr_process)
|
||||
index := atomic_add(*shared.next, 1);
|
||||
if index >= shared.count break;
|
||||
shared.proc(shared.engine, index, shared.data);
|
||||
Default_Allocator.thread_idle(); // parked until the next job: give back what this thread's heap holds free
|
||||
wait_for(*pool.wake);
|
||||
if pool.quit break;
|
||||
while true {
|
||||
if ocr_cancelled(engine) break; // what is left fails as cancelled (ocr_process)
|
||||
index := atomic_add(*pool.next, 1);
|
||||
if index >= pool.count break;
|
||||
pool.proc(engine, index, pool.data);
|
||||
reset_temporary_storage();
|
||||
}
|
||||
signal(*pool.done);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
// thread_init allocates the thread's temporary storage from the context allocator and tells it
|
||||
// the thread is starting; the caller's may be a temporary allocator, which refuses that.
|
||||
push_allocator(context.default_allocator);
|
||||
thread_count := min(cast(s64) engine.workers, count);
|
||||
threads := NewArray(thread_count, Thread,, context.default_allocator);
|
||||
defer array_free(threads,, context.default_allocator);
|
||||
started := 0;
|
||||
for * threads {
|
||||
if !thread_init(it, worker_proc, temporary_storage_size = 1024 * 1024) break;
|
||||
thread_start(it, *shared); // sets thread.data
|
||||
started += 1;
|
||||
run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void) {
|
||||
if count <= 0 return;
|
||||
pool := *engine.pool;
|
||||
if !pool.started { // no threads: do it here
|
||||
for 0..count - 1 {
|
||||
if ocr_cancelled(engine) break;
|
||||
proc(engine, it, data);
|
||||
}
|
||||
if !started { // no threads: do it here
|
||||
here: Thread;
|
||||
here.data = *shared;
|
||||
worker_proc(*here);
|
||||
}
|
||||
for 0..started - 1 {
|
||||
thread_is_done(*threads[it], -1);
|
||||
thread_deinit(*threads[it]);
|
||||
return;
|
||||
}
|
||||
lock(*pool.mutex);
|
||||
defer unlock(*pool.mutex);
|
||||
pool.count, pool.next, pool.proc, pool.data = count, 0, proc, data;
|
||||
// Every worker is woken and reports back, even with fewer items than workers: a worker that finds
|
||||
// nothing left reports at once, so no wake is left over for the next job.
|
||||
for 1..pool.started signal(*pool.wake);
|
||||
for 1..pool.started wait_for(*pool.done);
|
||||
}
|
||||
|
||||
#load "models.jai";
|
||||
@ -395,6 +457,7 @@ run_on_workers :: (engine: *Ocr_Engine, count: s64, proc: Work_Proc, data: *void
|
||||
#import "Thread";
|
||||
#import "Atomics";
|
||||
#import "onnxruntime";
|
||||
Default_Allocator :: #import "Default_Allocator"; // rpmalloc (modules/Default_Allocator), for thread_idle where the workers wait
|
||||
#if OS == .WINDOWS {
|
||||
#import "Windows";
|
||||
#import "Windows_Utf8";
|
||||
|
||||
@ -148,7 +148,7 @@ segment_create :: (models_folder: string, model_name: string, prefer_gpu := true
|
||||
if !ok { segment_destroy(engine); return null, false, error; }
|
||||
|
||||
api := ort.api;
|
||||
ok, error = ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "segment", *engine.env));
|
||||
ok, error = ort_check(api.CreateEnv(ort.log_level, "segment", *engine.env));
|
||||
if !ok { segment_destroy(engine); return null, false, error; }
|
||||
ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));
|
||||
|
||||
|
||||
@ -9,6 +9,9 @@ ort: struct {
|
||||
api: *OrtApi;
|
||||
library_path: string; // the library that was loaded
|
||||
version: string; // onnxruntime's own version string, such as 1.30.0
|
||||
// What the engines' environments log, to the console: set before the first engine is made, as
|
||||
// onnxruntime keeps one environment per process. INFO shows each GPU arena extension and its size.
|
||||
log_level := OrtLoggingLevel.ORT_LOGGING_LEVEL_WARNING;
|
||||
}
|
||||
|
||||
// Loads onnxruntime from the first folder that has it: beside the running executable, then each
|
||||
|
||||
@ -20,17 +20,21 @@ ort_gpu_held: s32;
|
||||
|
||||
Ort_Session :: struct {
|
||||
session: *OrtSession;
|
||||
shrink_options: *OrtRunOptions; // for a run that shrinks the arena afterwards (ort_run's shrink); null when the settings did not ask for it
|
||||
input_name, output_name: *u8; // heap C strings
|
||||
input_shape: [4] s64; // as the model declares it: -1 for a size it leaves open
|
||||
input_rank: s64;
|
||||
}
|
||||
|
||||
// How a CUDA session uses GPU memory. The defaults came from measuring OCR on the sample captures
|
||||
// (RTX 5080 laptop): with nearly every input a new size, onnxruntime's exact-size arena filled GPU
|
||||
// memory, and cuDNN's largest workspaces made it worse.
|
||||
// How a CUDA session uses GPU memory. Its arena takes blocks from CUDA as runs need them, and by
|
||||
// default never joins or frees them: a model whose input size changes (OCR's) leaves a block of
|
||||
// every larger size behind, gigabytes of them, unless the arena is shrunk now and then. Shrinking
|
||||
// after every run costs a third to a half of OCR's speed (CUDA frees and allocates again, waiting for
|
||||
// the GPU), so OCR shrinks once per call (ocr_process).
|
||||
Ort_Gpu_Settings :: struct {
|
||||
arena_extend_strategy := "kNextPowerOfTwo"; // reuses freed blocks for other sizes; "kSameAsRequested" grows by exactly what is asked
|
||||
arena_extend_strategy := "kNextPowerOfTwo"; // how much the arena takes when it runs short: twice as much each time; "kSameAsRequested" takes what is asked
|
||||
cudnn_max_workspace := false; // let cuDNN take as much working memory as it finds best (onnxruntime's own default)
|
||||
shrink_arena := false; // a run with ort_run's shrink set then gives back the arena's blocks no run is using
|
||||
}
|
||||
|
||||
// A session for the model file at path; with CUDA when gpu is set. output names the output that
|
||||
@ -65,6 +69,12 @@ ort_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads:
|
||||
if !file_exists(path) return result, false, tprint("% is not there%.", path, ifx missing_hint then tprint("; %", missing_hint) else "");
|
||||
ok, error = ort_check(api.CreateSession(env, ort_path(path), options, *result.session));
|
||||
if !ok return result, false, error;
|
||||
if gpu && settings.shrink_arena {
|
||||
// Device 0's arena, as device_id above; the CPU's is left alone.
|
||||
ok, error = ort_check(api.CreateRunOptions(*result.shrink_options));
|
||||
if ok ok, error = ort_check(api.AddRunConfigEntry(result.shrink_options, "memory.enable_memory_arena_shrinkage", "gpu:0"));
|
||||
if !ok { ort_release_session(*result); return result, false, error; }
|
||||
}
|
||||
|
||||
allocator: *OrtAllocator;
|
||||
ort_check(api.GetAllocatorWithDefaultOptions(*allocator));
|
||||
@ -107,6 +117,7 @@ ort_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads:
|
||||
|
||||
ort_release_session :: (s: *Ort_Session) {
|
||||
if s.session ort.api.ReleaseSession(s.session);
|
||||
if s.shrink_options ort.api.ReleaseRunOptions(s.shrink_options);
|
||||
free(s.input_name,, context.default_allocator);
|
||||
free(s.output_name,, context.default_allocator);
|
||||
s.* = .{};
|
||||
@ -121,8 +132,9 @@ Ort_Output :: struct {
|
||||
}
|
||||
|
||||
// Runs the session on one input tensor of shape, bytes long at data, with element type element.
|
||||
// A session may be run from several threads at once.
|
||||
ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s64, element: ONNXTensorElementDataType, shape: [] s64) -> Ort_Output, ok: bool, error: string {
|
||||
// A session may be run from several threads at once. With shrink, and a session made with
|
||||
// shrink_arena, the arena gives back the blocks no run is using once this run is done.
|
||||
ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s64, element: ONNXTensorElementDataType, shape: [] s64, shrink := false) -> Ort_Output, ok: bool, error: string {
|
||||
out: Ort_Output;
|
||||
api := ort.api;
|
||||
input: *OrtValue;
|
||||
@ -130,7 +142,7 @@ ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s6
|
||||
if !ok return out, false, error;
|
||||
defer api.ReleaseValue(input);
|
||||
|
||||
ok, error = ort_check(api.Run(s.session, null, *s.input_name, *input, 1, *s.output_name, 1, *out.value));
|
||||
ok, error = ort_check(api.Run(s.session, ifx shrink then s.shrink_options else null, *s.input_name, *input, 1, *s.output_name, 1, *out.value));
|
||||
if !ok return out, false, error;
|
||||
|
||||
raw: *void;
|
||||
@ -146,8 +158,8 @@ ort_run :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: *void, bytes: s6
|
||||
return out, true, "";
|
||||
}
|
||||
|
||||
ort_run_float :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64) -> Ort_Output, ok: bool, error: string {
|
||||
out, ok, error := ort_run(s, memory_info, data.data, data.count * size_of(float), .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, shape);
|
||||
ort_run_float :: (s: *Ort_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64, shrink := false) -> Ort_Output, ok: bool, error: string {
|
||||
out, ok, error := ort_run(s, memory_info, data.data, data.count * size_of(float), .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, shape, shrink);
|
||||
return out, ok, error;
|
||||
}
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user