MusaMahmood/mcap
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

[0.0.0] OCR: an onnxruntime module and an OCR module, PP-OCRv5 mobile and server, on the GPU through CUDA when it can

Browse Source
This commit is contained in:
Musa Mahmood 2026-10-07 20:45:05 +01:00
parent 94d56f25b1
commit bb94f03b94
43 changed files with 40177 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
.gitignore vendored
View File

@ -1,6 +1,7 @@
# Build output: binaries, copied assets and logs go to bin/, intermediates to .build/ (the
# binding generators in modules/ make their own).
/bin/
/bin-headless/
.build/
# Debuggers and editors
@ -15,3 +16,10 @@
.DS_Store
Thumbs.db
.stfolder
# OCR: downloaded, not committed (modules/onnxruntime/fetch.jai, modules/OCR/fetch_models.jai)
/modules/onnxruntime/windows/
/modules/onnxruntime/linux/
/modules/onnxruntime/macos/
/modules/onnxruntime/download/
/modules/OCR/models/

13
README.md
View File

@ -54,6 +54,19 @@ The macOS check only works on a Mac, since the Objective-C module runs code at c
Run any build with `run_tests` to run the test suites (debug builds include them),
`-allow-multiple-instances` to skip the single-instance check, or `-version`.
For unattended runs (scripts, CI, agents), add `-headless`: the app exits once the tests finish,
with exit code 0 if every suite passed and 1 otherwise, and a failed assertion or runtime check
is written to stderr, the log and the crash report and exits with 1 instead of opening a message
box (`runtime_support_headless` in `modules/Runtime_Support.jai`):
```
bin/mcap-debug.exe run_tests -headless -allow-multiple-instances > test-output.txt
```
`jai first.jai - no_crash_dialog [...]` makes builds for automation: `NO_CRASH_DIALOG` makes them
headless whatever their arguments, and everything they build goes to `bin-headless/` instead of
`bin/`, so an automated build never writes over the one you run by hand.
## Starting a new app
```

1
check_targets.jai
View File

@ -34,6 +34,7 @@ BUILD_DATE :: "-";
DEBUG :: true;
COMPILE_TESTS :: true;
COMPILE_ARGS :: "";
NO_CRASH_DIALOG :: false;
S, w);
while true {
message := compiler_wait_for_message();

96
first.jai
View File

@ -9,7 +9,9 @@ APP_FOLDER_NAME :: "mcap";
APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
VERSION :: "0.0.0";
BUILD_DIRECTORY :: "bin";
// 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";
base_options: Build_Options;
@ -19,6 +21,7 @@ Pending_Build :: struct {
executable_name: string;
optimized: bool;
failed: bool;
tool: bool; // a command-line tool, not the app: no icon, manifest or bundle
}
#run,stallable build_project();
@ -40,10 +43,17 @@ build_project :: () {
else if build_release log("Building in release mode");
else log("Building in debug mode");
if build_release && array_find(args, "compile_tests") log("Compiling tests into the release build");
// jai first.jai - no_crash_dialog [...]: for automated runs. NO_CRASH_DIALOG makes the app
// headless whatever its arguments (no message boxes; see runtime_support_headless), and the
// build goes to bin-headless/ instead of bin/.
if array_find(args, "no_crash_dialog") {
build_directory = "bin-headless";
log("Building with NO_CRASH_DIALOG into %/", build_directory);
}
set_build_options_dc(.{do_output = false});
set_working_directory(#filepath); // so that we load fonts correctly.
make_directory_if_it_does_not_exist(BUILD_DIRECTORY);
make_directory_if_it_does_not_exist(build_directory);
// Workspaces compile concurrently once source is added, so every build is
// started before any is waited on.
@ -51,10 +61,15 @@ build_project :: () {
if build_debug array_add(*pending, start_main_build(optimized = false));
if build_release array_add(*pending, start_main_build(optimized = true));
// jai first.jai - ocr_bench: also bin/ocr-bench, OCR timing and output over a folder of .mcap files.
if array_find(args, "ocr_bench") array_add(*pending, start_tool_build("ocr-bench", "tools/ocr_bench.jai", optimized = build_release));
wait_for_builds(pending);
for pending if !it.failed finish_main_build(it);
for pending if !it.failed {
if it.tool log("Built %/%", build_directory, it.executable_name);
else finish_main_build(it);
}
// The assets are the same for every configuration, so they are copied once.
for pending if !it.failed { copy_assets(); break; }
@ -106,12 +121,12 @@ start_main_build :: (optimized: bool) -> Pending_Build {
}
#if OS == .WINDOWS {
file_delete(tprint("%/%.pdb", BUILD_DIRECTORY, options.output_executable_name));
file_delete(tprint("%/%.rdi", BUILD_DIRECTORY, options.output_executable_name));
file_delete(tprint("%/%.pdb", build_directory, options.output_executable_name));
file_delete(tprint("%/%.rdi", build_directory, options.output_executable_name));
}
options.temporary_storage_size = 64 * 4096; // 256kB (could be more i guess?)
options.output_path = BUILD_DIRECTORY;
options.output_path = build_directory;
w := compiler_create_workspace(options.output_executable_name);
set_build_options(options, w);
@ -133,9 +148,31 @@ start_main_build :: (optimized: bool) -> Pending_Build {
return .{ workspace = w, executable_name = options.output_executable_name, optimized = optimized };
}
start_tool_build :: (name: string, file: string, optimized: bool) -> Pending_Build {
log("Building % (%)...", name, ifx optimized then "release" else "debug");
options := base_options;
if optimized {
set_optimization(*options, .VERY_OPTIMIZED, preserve_debug_info=true);
options.backend = .LLVM;
options.output_executable_name = name;
} else {
#if CPU == .X64 options.backend = .X64;
options.output_executable_name = tprint("%-debug", name);
}
options.output_path = build_directory;
w := compiler_create_workspace(options.output_executable_name);
set_build_options(options, w);
#if OS == .WINDOWS set_build_options_dc(.{append_linker_arguments=.["/IGNORE:4099"]}, w);
compiler_begin_intercept(w);
add_build_file(file, w);
add_build_constants(w, optimized);
return .{ workspace = w, executable_name = options.output_executable_name, optimized = optimized, tool = true };
}
finish_main_build :: (build: Pending_Build) {
#if OS == .WINDOWS {
exe_path := tprint("%/%.exe", BUILD_DIRECTORY, build.executable_name);
exe_path := tprint("%/%.exe", build_directory, build.executable_name);
if build.optimized {
ico_data := create_ico_file_from_bitmap_filename("assets/images/icon.png");
@ -154,7 +191,7 @@ finish_main_build :: (build: Pending_Build) {
if build.optimized create_macos_app_bundle(build.executable_name);
}
log("Built %/%", BUILD_DIRECTORY, build.executable_name);
log("Built %/%", build_directory, build.executable_name);
}
#if OS == .MACOS {
@ -165,14 +202,14 @@ finish_main_build :: (build: Pending_Build) {
// Info.plist. The app finds data/, logs/ and images/ in the folder holding the bundle; see
// app_install_directory.
create_macos_app_bundle :: (executable_name: string) {
app := tprint("%/%", BUILD_DIRECTORY, MACOS_APP_BUNDLE);
app := tprint("%/%", build_directory, MACOS_APP_BUNDLE);
run_command("rm", "-rf", app);
executable_folder := tprint("%/Contents/MacOS", app);
resources_folder := tprint("%/Contents/Resources", app);
make_directory_if_it_does_not_exist(executable_folder, recursive = true);
make_directory_if_it_does_not_exist(resources_folder, recursive = true);
copied := copy_file(tprint("%/%", BUILD_DIRECTORY, executable_name), tprint("%/%", executable_folder, executable_name));
copied := copy_file(tprint("%/%", build_directory, executable_name), tprint("%/%", executable_folder, executable_name));
for MACOS_BUNDLE_DYLIBS copied &&= copy_file(tprint("assets/%", it), tprint("%/%", executable_folder, it));
if !copied {
log_error("ERROR: Couldn't copy the executable and its dylibs into %\n", app);
@ -181,7 +218,7 @@ finish_main_build :: (build: Pending_Build) {
run_command("chmod", "+x", tprint("%/%", executable_folder, executable_name));
// sips makes an .icns only from a size the format holds, and icon.png may be any square size.
resized := tprint("%/icon_1024.png", BUILD_DIRECTORY);
resized := tprint("%/icon_1024.png", build_directory);
result := run_command("sips", "-z", "1024", "1024", "assets/images/icon.png", "--out", resized);
if result.exit_code == 0 result = run_command("sips", "-s", "format", "icns", resized, "--out", tprint("%/%.icns", resources_folder, APP_FOLDER_NAME));
file_delete(resized);
@ -221,9 +258,38 @@ PLIST
}
copy_assets :: () {
copy_directory("assets", "bin"); // This puts the CONTENTS of assets into bin. Not sure if that's what we want, but I'll roll with it.
copy_directory("assets", build_directory); // This puts the CONTENTS of assets into the build directory. Not sure if that's what we want, but I'll roll with it.
#if OS == .MACOS copy_file("assets/libSDL3.0.dylib", "bin/libSDL3.0.dylib");
#if OS == .MACOS copy_file("assets/libSDL3.0.dylib", tprint("%/libSDL3.0.dylib", build_directory));
// OCR: onnxruntime beside the executable, where modules/onnxruntime looks for it first, and the
// models in bin/ocr-models. Both are fetched, not in git (see their fetch scripts), and large,
// so a file is only copied when it is missing or differs in size or time.
#if OS == .WINDOWS ORT_FOLDER :: "modules/onnxruntime/windows";
else #if OS == .LINUX ORT_FOLDER :: "modules/onnxruntime/linux";
else ORT_FOLDER :: "modules/onnxruntime/macos";
if file_exists(ORT_FOLDER) {
for file_list(ORT_FOLDER) {
name := path_filename(it);
if contains(name, ".dll") || contains(name, ".so") || contains(name, ".dylib") copy_if_changed(it, tprint("%/%", build_directory, name));
}
} else {
log("OCR: onnxruntime has not been fetched; run jai fetch.jai in modules/onnxruntime for OCR.");
}
make_directory_if_it_does_not_exist(tprint("%/ocr-models", build_directory));
copy_if_changed("modules/OCR/models.config", tprint("%/ocr-models/models.config", build_directory));
if file_exists("modules/OCR/models") {
for file_list("modules/OCR/models") copy_if_changed(it, tprint("%/ocr-models/%", build_directory, path_filename(it)));
} else {
log("OCR: the models have not been fetched; run jai fetch_models.jai in modules/OCR for OCR.");
}
}
copy_if_changed :: (from: string, to: string) {
from_time, from_size, from_ok := file_modtime_and_size(from);
to_time, to_size, to_ok := file_modtime_and_size(to);
if to_ok && from_ok && to_size == from_size && to_time >= from_time return;
if !copy_file(from, to) log_error("Could not copy % to %.", from, to);
}
add_build_constants :: (w: Workspace, optimized: bool) {
@ -253,6 +319,7 @@ BUILD_DATE :: "%";
DEBUG :: %;
COMPILE_TESTS :: %;
COMPILE_ARGS :: "%";
NO_CRASH_DIALOG :: %;
STRING,
escape_for_string_literal(APP_NAME),
escape_for_string_literal(APP_FOLDER_NAME),
@ -263,7 +330,8 @@ STRING,
ifx optimized then "false" else "true",
ifx !optimized || array_find(base_options.compile_time_command_line, "compile_tests") then "true" else "false",
escape_for_string_literal(join(..base_options.compile_time_command_line, " "))
escape_for_string_literal(join(..base_options.compile_time_command_line, " ")),
ifx array_find(base_options.compile_time_command_line, "no_crash_dialog") then "true" else "false"
);
add_build_string(build_constants, w);

70
modules/OCR/README.md Normal file
View File

@ -0,0 +1,70 @@
# OCR
Text in images, with where each line is: a batch of RGBA images in, each image's text lines out,
each with a box and a confidence. It runs ONNX models through onnxruntime
([`modules/onnxruntime`](../onnxruntime)): on an NVIDIA GPU through CUDA when every library that
needs is present, and on the CPU otherwise.
## Setting up
The binaries and the models are downloaded, not committed:
```
cd modules/onnxruntime && jai fetch.jai # onnxruntime 1.30.0, GPU build (CUDA 13); `- cpu` for the CPU-only build
cd modules/OCR && jai fetch_models.jai # both models (about 195 MB); `- mobile` for the light one only
jai generate.jai # in modules/onnxruntime, only when changing onnxruntime versions
```
Each checks the SHA-256 of what it downloads. Building the app (`jai first.jai`) copies onnxruntime
beside the executable and the models into `bin/ocr-models`.
## Using it
```jai
#import "OCR";
engine, ok, error := ocr_create(models_folder, "ppocr-v5-mobile"); // prefer_gpu = true
in_use, explanation := ocr_gpu_status(engine); // why not the GPU, and what to install
results := ocr_process(engine, images); // [] Ocr_Image (RGBA) -> [] Ocr_Result
for results[0].lines print("% at %,%: %\n", it.text, it.x, it.y, it.confidence);
ocr_destroy(engine);
```
`ocr_process` shares a batch out over worker threads: detection per image, then recognition of
every line of every image in batches of similar width. The results come from the context
allocator, or the one passed. `ocr_stats` gives the time spent in each stage.
## Models
`models.config` lists them; `ocr_create` takes a name. Both are PaddleOCR's PP-OCRv5 (Apache 2.0),
as converted to ONNX by RapidOCR, and read Chinese, English and Japanese:
| Name | Files | On the 20 sample captures, CPU |
| --- | --- | --- |
| `ppocr-v5-mobile` | detector 4.8 MB, recognizer 16.6 MB | about 0.5 s per capture |
| `ppocr-v5-server` | detector 88 MB, recognizer 85 MB | about 2.5 s per capture |
Another model with the same two-stage design is a new `[model NAME]` section; a different design
needs its own detection or recognition code beside `detect.jai` and `recognize.jai`.
## The GPU
onnxruntime's CUDA build needs, on PATH or beside the executable:
- the NVIDIA driver;
- the CUDA Toolkit 13.x (`cudart64_13`, `cublas64_13`, `cublasLt64_13`, `cufft64_12`), from
[developer.nvidia.com/cuda-downloads](https://developer.nvidia.com/cuda-downloads), whose installer puts its bin folder on PATH;
- cuDNN 9.x for CUDA 13 (`cudnn64_9` and its `cudnn_*64_9` companions), from
[developer.nvidia.com/cudnn-downloads](https://developer.nvidia.com/cudnn-downloads), whose bin folder for CUDA 13 you add to PATH yourself.
`ocr_create` looks for each of these first (`cuda_check.jai`) and says which are missing, rather
than letting onnxruntime fail later and obscurely. If they are all there but CUDA still cannot
start, it runs on the CPU and keeps onnxruntime's message.
## Measuring
`jai first.jai - release ocr_bench`, then `bin/ocr-bench MCAP_FOLDER`: each model on the GPU (when
available) and the CPU, the time per capture and per stage, and each capture's lines in
`bin/ocr-out/<model>-<provider>/`. Like any command-line tool here it sets
`runtime_support_headless` (`modules/Runtime_Support.jai`): a failure is printed and exits with
code 1, with no message box to wait on.

117
modules/OCR/cuda_check.jai Normal file
View File

@ -0,0 +1,117 @@
// Whether onnxruntime's CUDA provider can run here, and if not, what is missing and how to fix it.
// The provider needs the NVIDIA driver, parts of the CUDA 13 toolkit and cuDNN 9 for CUDA 13; it
// fails late and obscurely without them, so each library is looked for first, by the name the
// provider loads it by (from onnxruntime_providers_cuda.dll's imports and the names it loads).
Cuda_Library :: struct {
name: string;
part: Cuda_Part;
}
Cuda_Part :: enum { DRIVER; TOOLKIT; CUDNN; }
#if OS == .WINDOWS {
CUDA_LIBRARIES :: Cuda_Library.[
.{"nvcuda.dll", .DRIVER},
.{"cudart64_13.dll", .TOOLKIT},
.{"cublas64_13.dll", .TOOLKIT},
.{"cublasLt64_13.dll", .TOOLKIT},
.{"cufft64_12.dll", .TOOLKIT},
.{"cudnn64_9.dll", .CUDNN},
.{"cudnn_graph64_9.dll", .CUDNN},
.{"cudnn_ops64_9.dll", .CUDNN},
.{"cudnn_cnn64_9.dll", .CUDNN},
.{"cudnn_heuristic64_9.dll", .CUDNN},
.{"cudnn_engines_precompiled64_9.dll", .CUDNN},
.{"cudnn_engines_runtime_compiled64_9.dll", .CUDNN},
];
CUDA_PROVIDER_LIBRARY :: "onnxruntime_providers_cuda.dll";
} else #if OS == .LINUX {
CUDA_LIBRARIES :: Cuda_Library.[
.{"libcuda.so.1", .DRIVER},
.{"libcudart.so.13", .TOOLKIT},
.{"libcublas.so.13", .TOOLKIT},
.{"libcublasLt.so.13", .TOOLKIT},
.{"libcufft.so.12", .TOOLKIT},
.{"libcudnn.so.9", .CUDNN},
];
CUDA_PROVIDER_LIBRARY :: "libonnxruntime_providers_cuda.so";
} else {
CUDA_LIBRARIES :: Cuda_Library.[];
CUDA_PROVIDER_LIBRARY :: "";
}
// Whether the CUDA provider should be tried, and when not, an explanation for people: what is
// missing, and the steps that fix it. Temp.
ocr_cuda_check :: () -> available: bool, explanation: string {
#if OS == .MACOS {
return false, "OCR runs on the CPU: onnxruntime has no CUDA on macOS.";
}
provider_path := tprint("%/%", path_strip_filename(ort.library_path), CUDA_PROVIDER_LIBRARY);
if !file_exists(provider_path) {
return false, tprint("OCR runs on the CPU: the onnxruntime beside the program is the CPU-only build (% is not there).\nTo use an NVIDIA GPU, run `jai fetch.jai` in modules/onnxruntime (without `cpu`) and rebuild.", CUDA_PROVIDER_LIBRARY);
}
missing_driver, missing_toolkit, missing_cudnn: [..] string;
missing_driver.allocator, missing_toolkit.allocator, missing_cudnn.allocator = temp, temp, temp;
for CUDA_LIBRARIES {
if ocr_library_loads(it.name) continue;
if it.part == {
case .DRIVER; array_add(*missing_driver, it.name);
case .TOOLKIT; array_add(*missing_toolkit, it.name);
case .CUDNN; array_add(*missing_cudnn, it.name);
}
}
if !missing_driver.count && !missing_toolkit.count && !missing_cudnn.count return true, "";
builder: String_Builder;
builder.allocator = temp;
if missing_driver.count {
append(*builder, "OCR runs on the CPU: no NVIDIA driver was found (");
append(*builder, join(..missing_driver, ", ",, temp));
append(*builder, "), so there is no NVIDIA GPU to use.");
return false, builder_to_string(*builder,, temp);
}
append(*builder, "OCR runs on the CPU. The NVIDIA GPU can run it once these are installed:\n");
step := 1;
if missing_toolkit.count {
print_to_builder(*builder, "\n%. The CUDA Toolkit 13.x, which has % that are missing.\n", step, join(..missing_toolkit, ", ",, temp));
append(*builder, " https://developer.nvidia.com/cuda-downloads\n");
#if OS == .WINDOWS append(*builder, " Its installer puts its bin folder on PATH.\n");
else append(*builder, " Then make sure its lib64 folder is on LD_LIBRARY_PATH.\n");
step += 1;
}
if missing_cudnn.count {
print_to_builder(*builder, "\n%. cuDNN 9.x for CUDA 13, which has % that are missing.\n", step, join(..missing_cudnn, ", ",, temp));
append(*builder, " https://developer.nvidia.com/cudnn-downloads\n");
#if OS == .WINDOWS {
append(*builder, " Its installer does not change PATH: add the folder with these DLLs to it yourself,\n");
append(*builder, " such as C:\\Program Files\\NVIDIA\\CUDNN\\v9.x\\bin\\13.x\\x64.\n");
} else {
append(*builder, " Then make sure its lib folder is on LD_LIBRARY_PATH.\n");
}
step += 1;
}
append(*builder, "\nCopying the DLLs beside the program also works. Restart the program afterwards; ");
append(*builder, "PATH is read when it starts.");
return false, builder_to_string(*builder,, temp);
}
#scope_module
ocr_library_loads :: (name: string) -> bool {
#if OS == .WINDOWS {
// The search order the CUDA provider's own loads use: beside the program, the system folders, PATH.
module := LoadLibraryW(utf8_to_wide(name,, temp));
if !module return false;
FreeLibrary(module);
return true;
} else {
handle := dlopen(temp_c_string(name), RTLD_LAZY);
if !handle return false;
dlclose(handle);
return true;
}
}

199
modules/OCR/detect.jai Normal file
View File

@ -0,0 +1,199 @@
// Detection: where the text lines are. The detector (a DB network) gives each pixel a probability
// of being text. Pixels above det_threshold are grouped into connected regions, each kept region
// becomes a box, and the box is grown by the DB unclip distance, since the network is trained to
// mark a shrunken core of each line.
//
// Images wider than det_max_side are scaled down for detection. Images taller than det_tile_height
// (after scaling) are detected in overlapping strips, so a long scrolling capture keeps its text at
// full size; each box is kept by the strip whose own band holds its centre.
//
// Version 1 finds horizontal text only, as screenshots have: boxes are axis-aligned.
Detected_Box :: struct {
x0, y0, x1, y1: s32; // in the image, x1 and y1 exclusive
score: float;
}
DET_TILE_OVERLAP : s32 : 128; // px of each strip shared with the next, after scaling
// The image's text boxes, from the default allocator.
ocr_detect :: (engine: *Ocr_Engine, image: Ocr_Image) -> [..] Detected_Box, ok: bool, error: string {
boxes: [..] Detected_Box;
boxes.allocator = context.default_allocator;
model := *engine.model;
if image.width <= 0 || image.height <= 0 return boxes, true, "";
scale := min(1.0, cast(float) model.det_max_side / image.width);
scaled_width := max(cast(s32) 1, cast(s32) (image.width * scale + 0.5));
scaled_height := max(cast(s32) 1, cast(s32) (image.height * scale + 0.5));
tile_height := max(model.det_tile_height, cast(s32) DET_TILE_OVERLAP * 2);
tile_top: s32 = 0;
while tile_top < scaled_height {
tile_bottom := min(tile_top + tile_height, scaled_height);
// This strip's own band: what is not left to the strips above and below it.
band_top := ifx tile_top == 0 then 0 else tile_top + DET_TILE_OVERLAP / 2;
band_bottom := ifx tile_bottom == scaled_height then scaled_height else tile_bottom - DET_TILE_OVERLAP / 2;
ok, error := ocr_detect_tile(engine, image, scale, scaled_width, tile_top, tile_bottom, band_top, band_bottom, *boxes);
if !ok return boxes, false, error;
if tile_bottom == scaled_height break;
tile_top = tile_bottom - DET_TILE_OVERLAP;
}
return boxes, true, "";
}
#scope_module
// Normalization the PP-OCR detectors were trained with: ImageNet mean and deviation, applied to the
// channels in BGR order, as PaddleOCR reads images with OpenCV.
DET_MEAN :: float.[0.485, 0.456, 0.406];
DET_STD :: float.[0.229, 0.224, 0.225];
ocr_detect_tile :: (engine: *Ocr_Engine, image: Ocr_Image, scale: float, scaled_width: s32, tile_top: s32, tile_bottom: s32, band_top: s32, band_bottom: s32, boxes: *[..] Detected_Box) -> ok: bool, error: string {
model := *engine.model;
valid_w := scaled_width;
valid_h := tile_bottom - tile_top;
w := align_up_32(valid_w);
h := align_up_32(valid_h);
// Planar BGR, normalized; the padding to a multiple of 32 is the normalized zero (the mean colour).
plane := cast(s64) w * h;
input := NewArray(plane * 3, float, initialized = true,, context.default_allocator);
defer array_free(input,, context.default_allocator);
inverse := 1.0 / scale;
for y: 0..valid_h - 1 {
source_y := min(cast(s32) ((y + tile_top + 0.5) * inverse), image.height - 1);
row := image.pixels + cast(s64) source_y * image.stride;
for x: 0..valid_w - 1 {
source_x := min(cast(s32) ((x + 0.5) * inverse), image.width - 1);
p := row + source_x * 4;
i := cast(s64) y * w + x;
input[i] = (p[2] / 255.0 - DET_MEAN[0]) / DET_STD[0]; // B
input[i + plane] = (p[1] / 255.0 - DET_MEAN[1]) / DET_STD[1]; // G
input[i + plane * 2] = (p[0] / 255.0 - DET_MEAN[2]) / DET_STD[2]; // R
}
}
shape := s64.[1, 3, h, w];
out, ok, error := ocr_run(*engine.det, engine.memory_info, input, shape);
if !ok return false, error;
defer ocr_release_output(*out);
if out.dim_count != 4 || out.dims[2] != h || out.dims[3] != w {
return false, tprint("The detector gave an output of shape %, not [1, 1, %, %].", array_view(out.dims, 0, out.dim_count), h, w);
}
regions := ocr_find_regions(out.data, w, valid_w, valid_h, model.det_threshold);
defer array_free(regions,, context.default_allocator);
for regions {
if it.pixels == 0 continue;
score := cast(float) (it.probability_sum / it.pixels);
if score < model.box_threshold continue;
bw := it.x1 - it.x0;
bh := it.y1 - it.y0;
if min(bw, bh) < 3 continue; // specks
// Grown as DB's unclip does: by area * ratio / perimeter on every side.
distance := (cast(float) bw * bh * model.unclip_ratio) / (2.0 * (bw + bh));
gx0 := it.x0 - distance;
gy0 := it.y0 - distance;
gx1 := it.x1 + distance;
gy1 := it.y1 + distance;
centre_y := tile_top + (it.y0 + it.y1) / 2.0;
if centre_y < cast(float) band_top || centre_y >= cast(float) band_bottom continue; // the neighbouring strip's
box: Detected_Box;
box.x0 = clamp(cast(s32) (gx0 * inverse), 0, image.width);
box.x1 = clamp(cast(s32) (gx1 * inverse + 0.999), 0, image.width);
box.y0 = clamp(cast(s32) ((gy0 + tile_top) * inverse), 0, image.height);
box.y1 = clamp(cast(s32) ((gy1 + tile_top) * inverse + 0.999), 0, image.height);
box.score = score;
if box.x1 - box.x0 >= 2 && box.y1 - box.y0 >= 2 array_add(boxes, box);
}
return true, "";
}
align_up_32 :: (n: s32) -> s32 {
return (n + 31) & cast(s32) ~31;
}
Region :: struct {
x0, y0, x1, y1: s32; // x1 and y1 exclusive
pixels: s64;
probability_sum: float64;
}
// Groups the pixels above threshold into 8-connected regions: one pass assigning provisional labels
// and recording which touch (union-find), then one summing each region's extent and probability.
ocr_find_regions :: (probability: *float, stride: s32, width: s32, height: s32, threshold: float) -> [..] Region {
regions: [..] Region;
regions.allocator = context.default_allocator;
labels := NewArray(cast(s64) width * height, s32, initialized = true,, context.default_allocator);
defer array_free(labels,, context.default_allocator);
parent: [..] s32;
parent.allocator = context.default_allocator;
defer array_free(parent);
array_add(*parent, 0); // label 0 is background
find :: (parent: [] s32, label: s32) -> s32 {
root := label;
while parent[root] != root root = parent[root];
while parent[label] != root { next := parent[label]; parent[label] = root; label = next; } // path compression
return root;
}
unite :: (parent: [] s32, a: s32, b: s32) {
ra, rb := find(parent, a), find(parent, b);
if ra < rb parent[rb] = ra;
else if rb < ra parent[ra] = rb;
}
for y: 0..height - 1 {
for x: 0..width - 1 {
if probability[cast(s64) y * stride + x] <= threshold continue;
i := cast(s64) y * width + x;
label: s32 = 0;
// Neighbours already visited: left, and the three above.
neighbours: [4] s32;
neighbours[0] = ifx x > 0 then labels[i - 1] else 0;
neighbours[1] = ifx y > 0 && x > 0 then labels[i - width - 1] else 0;
neighbours[2] = ifx y > 0 then labels[i - width] else 0;
neighbours[3] = ifx y > 0 && x < width - 1 then labels[i - width + 1] else 0;
for neighbours {
if !it continue;
if !label label = it;
else if it != label unite(parent, label, it);
}
if !label {
label = cast(s32) parent.count;
array_add(*parent, label);
}
labels[i] = label;
}
}
// Each root gets a region.
region_of := NewArray(parent.count, s32, initialized = true,, context.default_allocator);
defer array_free(region_of,, context.default_allocator);
for y: 0..height - 1 {
for x: 0..width - 1 {
i := cast(s64) y * width + x;
if !labels[i] continue;
root := find(parent, labels[i]);
index := region_of[root] - 1;
if index < 0 {
region := array_add(*regions);
region.* = .{x0 = x, y0 = y, x1 = x + 1, y1 = y + 1};
index = cast(s32) regions.count - 1;
region_of[root] = index + 1;
}
r := *regions[index];
r.x0 = min(r.x0, x); r.x1 = max(r.x1, x + 1);
r.y0 = min(r.y0, y); r.y1 = max(r.y1, y + 1);
r.pixels += 1;
r.probability_sum += probability[cast(s64) y * stride + x];
}
}
return regions;
}

73
modules/OCR/fetch_models.jai Normal file
View File

@ -0,0 +1,73 @@
// Downloads the files models.config names into models/, checking each one's SHA-256. The models
// are not in git; run this after cloning. Files already there with the right hash are skipped.
//
// jai fetch_models.jai every model
// jai fetch_models.jai - mobile only the files ppocr-v5-mobile needs (about 21 MB)
//
// From RapidOCR's conversions of PaddleOCR's PP-OCRv5 (both Apache 2.0):
// https://www.modelscope.cn/models/RapidAI/RapidOCR
SOURCE :: "https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master";
Model_File :: struct { name: string; source_path: string; sha256: string; server: bool; }
FILES :: Model_File.[
.{"ch_PP-OCRv5_det_mobile.onnx", "onnx/PP-OCRv5/det/ch_PP-OCRv5_det_mobile.onnx", "4d97c44a20d30a81aad087d6a396b08f786c4635742afc391f6621f5c6ae78ae", false},
.{"ch_PP-OCRv5_rec_mobile.onnx", "onnx/PP-OCRv5/rec/ch_PP-OCRv5_rec_mobile.onnx", "5825fc7ebf84ae7a412be049820b4d86d77620f204a041697b0494669b1742c5", false},
.{"ppocrv5_dict.txt", "paddle/PP-OCRv5/rec/ch_PP-OCRv5_rec_mobile/ppocrv5_dict.txt", "d1979e9f794c464c0d2e0b70a7fe14dd978e9dc644c0e71f14158cdf8342af1b", false},
.{"ch_PP-OCRv5_det_server.onnx", "onnx/PP-OCRv5/det/ch_PP-OCRv5_det_server.onnx", "0f8846b1d4bba223a2a2f9d9b44022fbc22cc019051a602b41a7fda9667e4cad", true},
.{"ch_PP-OCRv5_rec_server.onnx", "onnx/PP-OCRv5/rec/ch_PP-OCRv5_rec_server.onnx", "e09385400eaaaef34ceff54aeb7c4f0f1fe014c27fa8b9905d4709b65746562a", true},
];
#run,stallable { if !fetch() exit(1); }
fetch :: () -> bool {
set_build_options_dc(.{do_output = false});
mobile_only := array_find(get_build_options().compile_time_command_line, "mobile");
here := trim_right(path_strip_filename(#file), "/\\");
folder := tprint("%/models", here);
make_directory_if_it_does_not_exist(folder);
for FILES {
if mobile_only && it.server continue;
path := tprint("%/%", folder, it.name);
if file_exists(path) && sha256_of(path) == it.sha256 { log("% is already here.", it.name); continue; }
url := tprint("%/%", SOURCE, it.source_path);
log("Downloading %...", url);
result := run_command("curl", "-L", "--fail", "--progress-bar", "-o", path, url);
if result.exit_code != 0 { log_error("curl failed for % (exit code %).", it.name, result.exit_code); return false; }
hash := sha256_of(path);
if hash != it.sha256 {
log_error("% has SHA-256 %, expected %; it was deleted.", it.name, hash, it.sha256);
file_delete(path);
return false;
}
}
log("The models are in %.", folder);
return true;
}
// Lowercase hex, or "" when it could not be worked out.
sha256_of :: (path: string) -> string {
#if OS == .WINDOWS {
result, output := run_command("certutil", "-hashfile", path, "SHA256", capture_and_return_output = true);
if result.exit_code != 0 return "";
lines := split(output, "\n");
if lines.count < 2 return "";
return to_lower_copy(replace(trim(lines[1]), " ", "")); // line 2 is the hash
} else {
command := ifx OS == .MACOS then string.["shasum", "-a", "256", path] else string.["sha256sum", path];
result, output := run_command(..command, capture_and_return_output = true);
if result.exit_code != 0 return "";
return to_lower_copy(trim(split(output, " ")[0]));
}
}
#import "Basic";
#import "Compiler";
#import "File";
#import "File_Utilities";
#import "Process";
#import "String";

40
modules/OCR/models.config Normal file
View File

@ -0,0 +1,40 @@
@version 1
# OCR models. Each [model NAME] names the files it needs from this folder (fetch_models.jai
# downloads them) and how it is run; ocr_create takes the NAME. Changes apply to engines created
# after the file is saved.
#
# title shown to people
# detector, recognizer, dictionary file names, quoted
# det_max_side wider images are scaled down to this width for detection (px)
# det_tile_height taller images are detected in overlapping strips this high (px, after scaling)
# det_threshold a pixel is text when its probability is above this
# box_threshold a text region is kept when its mean probability is at least this
# unclip_ratio how far each region is grown, as DB detection networks expect
# rec_height each text line is scaled to this height for recognition (px)
# rec_batch text lines recognized together in one run
[model ppocr-v5-mobile]
title: "PP-OCRv5 mobile (light)"
detector: "ch_PP-OCRv5_det_mobile.onnx"
recognizer: "ch_PP-OCRv5_rec_mobile.onnx"
dictionary: "ppocrv5_dict.txt"
det_max_side: 2560
det_tile_height: 1536
det_threshold: 0.3
box_threshold: 0.6
unclip_ratio: 1.5
rec_height: 48
rec_batch: 16
[model ppocr-v5-server]
title: "PP-OCRv5 server (heavy)"
detector: "ch_PP-OCRv5_det_server.onnx"
recognizer: "ch_PP-OCRv5_rec_server.onnx"
dictionary: "ppocrv5_dict.txt"
det_max_side: 2560
det_tile_height: 1536
det_threshold: 0.3
box_threshold: 0.6
unclip_ratio: 1.5
rec_height: 48
rec_batch: 16

173
modules/OCR/models.jai Normal file
View File

@ -0,0 +1,173 @@
// models.config, and running a model: creating its onnxruntime session on the GPU or the CPU, and
// running it on a float tensor.
Ocr_Model_Info :: struct {
name: string;
title: string;
detector, recognizer, dictionary: string; // file names in the models folder
det_max_side: s32 = 2560;
det_tile_height: s32 = 1536;
det_threshold: float = 0.3;
box_threshold: float = 0.6;
unclip_ratio: float = 1.5;
rec_height: s32 = 48;
rec_batch: s32 = 16;
}
OCR_MODELS_CONFIG :: "models.config";
// The models models.config lists, in its order; allocated with the context allocator. A line that
// cannot be read is logged and skipped, as in the Ex4 formats.
ocr_list_models :: (models_folder: string) -> [] Ocr_Model_Info, error: string {
models: [..] Ocr_Model_Info;
path := tprint("%/%", models_folder, OCR_MODELS_CONFIG);
text, ok := read_entire_file(path, log_errors = false,, temp);
if !ok return models, tprint("Cannot read %.", path);
model: *Ocr_Model_Info;
version_seen := false;
line_number := 0;
for raw_line: split(text, "\n",, temp) {
line_number += 1;
line := trim(ocr_strip_comment(raw_line));
if !line.count continue;
if begins_with(line, "@version") { version_seen = true; continue; }
if !version_seen return models, tprint("% does not start with @version 1.", path);
if begins_with(line, "[model ") && ends_with(line, "]") {
model = array_add(*models);
model.* = .{};
model.name = copy_string(trim(slice(line, 7, line.count - 8)));
continue;
}
colon := find_index_from_left(line, #char ":");
if colon < 0 || !model { log_error("[OCR] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; }
key := trim(slice(line, 0, colon));
value := trim(advance(line, colon + 1));
if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\"" value = slice(value, 1, value.count - 2);
number_text := value;
number, number_ok := string_to_float(number_text);
if key == {
case "title"; model.title = copy_string(value);
case "detector"; model.detector = copy_string(value);
case "recognizer"; model.recognizer = copy_string(value);
case "dictionary"; model.dictionary = copy_string(value);
case "det_max_side"; if number_ok model.det_max_side = cast(s32) number;
case "det_tile_height"; if number_ok model.det_tile_height = cast(s32) number;
case "det_threshold"; if number_ok model.det_threshold = number;
case "box_threshold"; if number_ok model.box_threshold = number;
case "unclip_ratio"; if number_ok model.unclip_ratio = number;
case "rec_height"; if number_ok model.rec_height = cast(s32) number;
case "rec_batch"; if number_ok model.rec_batch = cast(s32) number;
case; log_error("[OCR] %:%: unknown key %. Ignored.", path, line_number, key);
}
}
return models, "";
}
#scope_module
// "#" opens a comment at the start of a line or after whitespace, outside quotes.
ocr_strip_comment :: (line: string) -> string {
in_quotes := false;
for i: 0..line.count - 1 {
c := line[i];
if c == #char "\"" in_quotes = !in_quotes;
else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t") return slice(line, 0, i);
}
return line;
}
Ocr_Session :: struct {
session: *OrtSession;
input_name, output_name: *u8; // heap C strings
}
// A session for the model file; with CUDA when gpu is set. On failure, the reason.
ocr_create_session :: (env: *OrtEnv, path: string, gpu: bool, intra_op_threads: s32) -> Ocr_Session, ok: bool, error: string {
result: Ocr_Session;
api := ort.api;
options: *OrtSessionOptions;
ok, error := ort_check(api.CreateSessionOptions(*options));
if !ok return result, false, error;
defer api.ReleaseSessionOptions(options);
ort_check(api.SetSessionGraphOptimizationLevel(options, .ORT_ENABLE_ALL));
ort_check(api.SetIntraOpNumThreads(options, intra_op_threads));
ort_check(api.SetInterOpNumThreads(options, 1));
if gpu {
cuda: *OrtCUDAProviderOptionsV2;
ok, error = ort_check(api.CreateCUDAProviderOptions(*cuda));
if !ok return result, false, error;
defer api.ReleaseCUDAProviderOptions(cuda);
// Every image is a different size: an exhaustive cuDNN search would be redone for each new shape.
keys := (*u8).["device_id", "cudnn_conv_algo_search", "arena_extend_strategy"];
values := (*u8).["0", "HEURISTIC", "kSameAsRequested"];
ok, error = ort_check(api.UpdateCUDAProviderOptions(cuda, keys.data, values.data, xx keys.count));
if !ok return result, false, error;
ok, error = ort_check(api.SessionOptionsAppendExecutionProvider_CUDA_V2(options, cuda));
if !ok return result, false, error;
}
if !file_exists(path) return result, false, tprint("% is not there; run `jai fetch_models.jai` in modules/OCR and rebuild.", path);
ok, error = ort_check(api.CreateSession(env, ort_path(path), options, *result.session));
if !ok return result, false, error;
allocator: *OrtAllocator;
ort_check(api.GetAllocatorWithDefaultOptions(*allocator));
name: *u8;
ort_check(api.SessionGetInputName(result.session, 0, allocator, *name));
result.input_name = to_c_string(to_string(name),, context.default_allocator);
ort_check(api.AllocatorFree(allocator, name));
ort_check(api.SessionGetOutputName(result.session, 0, allocator, *name));
result.output_name = to_c_string(to_string(name),, context.default_allocator);
ort_check(api.AllocatorFree(allocator, name));
return result, true, "";
}
ocr_release_session :: (s: *Ocr_Session) {
if s.session ort.api.ReleaseSession(s.session);
free(s.input_name,, context.default_allocator);
free(s.output_name,, context.default_allocator);
s.* = .{};
}
// A run's output: its values stay valid until ocr_release_output.
Ocr_Output :: struct {
value: *OrtValue;
data: *float;
dims: [4] s64;
dim_count: s64;
}
// Runs the model on a float tensor of shape (the session may be shared between threads).
ocr_run :: (s: *Ocr_Session, memory_info: *OrtMemoryInfo, data: [] float, shape: [] s64) -> Ocr_Output, ok: bool, error: string {
out: Ocr_Output;
api := ort.api;
input: *OrtValue;
ok, error := ort_check(api.CreateTensorWithDataAsOrtValue(memory_info, data.data, xx (data.count * size_of(float)), shape.data, xx shape.count, .ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, *input));
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));
if !ok return out, false, error;
raw: *void;
ort_check(api.GetTensorMutableData(out.value, *raw));
out.data = cast(*float) raw;
info: *OrtTensorTypeAndShapeInfo;
ort_check(api.GetTensorTypeAndShape(out.value, *info));
count: u64;
ort_check(api.GetDimensionsCount(info, *count));
out.dim_count = min(cast(s64) count, out.dims.count);
ort_check(api.GetDimensions(info, out.dims.data, xx out.dim_count));
api.ReleaseTensorTypeAndShapeInfo(info);
return out, true, "";
}
ocr_release_output :: (out: *Ocr_Output) {
if out.value ort.api.ReleaseValue(out.value);
out.* = .{};
}

310
modules/OCR/module.jai Normal file
View File

@ -0,0 +1,310 @@
// OCR: the text in images, with where each line is. A batch of RGBA images goes in; each comes back
// as its text lines, with a box and a confidence for each.
//
// engine, ok, error := ocr_create(models_folder, "ppocr-v5-mobile");
// results := ocr_process(engine, images);
//
// The models are listed in models_folder/models.config (fetch_models.jai downloads them) and run
// with onnxruntime (modules/onnxruntime), on an NVIDIA GPU through CUDA when it can, else on the
// CPU. ocr_gpu_status says why the GPU is not in use, and what to install so it can be.
//
// Version 1 runs PP-OCRv5's two-stage pipeline: detection (detect.jai), then recognition
// (recognize.jai). Text is expected horizontal, as in screenshots.
Ocr_Provider :: enum { CPU; CUDA; }
// Tightly packed or not: stride is the bytes from one row to the next.
Ocr_Image :: struct {
pixels: *u8; // RGBA
width, height, stride: s32;
}
Ocr_Line :: struct {
x, y, width, height: s32; // in the image
text: string;
confidence: float; // 0 to 1: the mean probability of its characters
}
Ocr_Result :: struct {
lines: [] Ocr_Line; // top to bottom, then left to right
failed: bool;
error: string;
}
Ocr_Engine :: struct {
model: Ocr_Model_Info;
provider: Ocr_Provider;
gpu_explanation: string; // why the GPU is not in use; "" when it is
env: *OrtEnv;
memory_info: *OrtMemoryInfo;
det, rec: Ocr_Session;
dictionary: [] string; // class n is dictionary[n - 1]
workers: s32; // threads ocr_process uses
stats: Ocr_Stats; // since creation, or the last ocr_reset_stats
}
// Time spent in each stage, summed over the worker threads (so it can exceed the elapsed time).
Ocr_Stats :: struct {
images: s64;
lines: s64;
detect_us: s64; // preparing images, the detector, and finding the boxes
recognize_us: s64; // cropping lines, the recognizer, and decoding the text
}
ocr_stats :: (engine: *Ocr_Engine) -> Ocr_Stats { return engine.stats; }
ocr_reset_stats :: (engine: *Ocr_Engine) { engine.stats = .{}; }
// An engine for the model named in models_folder/models.config. With prefer_gpu it runs on CUDA when
// every library that needs is present; otherwise, or when CUDA fails to start, on the CPU.
ocr_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Ocr_Engine, ok: bool, error: string {
push_allocator(context.default_allocator);
loaded, load_error := ort_load();
if !loaded return null, false, load_error;
models, models_error := ocr_list_models(models_folder);
if models_error return null, false, models_error;
model: *Ocr_Model_Info;
for * models if it.name == model_name { model = it; break; }
if !model return null, false, tprint("% has no [model %].", OCR_MODELS_CONFIG, model_name);
engine := New(Ocr_Engine);
engine.model = model.*;
dictionary_text, dictionary_ok := read_entire_file(tprint("%/%", models_folder, model.dictionary), log_errors = false);
if !dictionary_ok { ocr_destroy(engine); return null, false, tprint("Cannot read the dictionary %/%.", models_folder, model.dictionary); }
entries := split(dictionary_text, "\n");
for * entries if it.count && it.*[it.count - 1] == #char "\r" it.count -= 1;
while entries.count && !entries[entries.count - 1].count entries.count -= 1; // the final newline
engine.dictionary = entries;
api := ort.api;
ok, error := ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "ocr", *engine.env));
if !ok { ocr_destroy(engine); return null, false, error; }
ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));
cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1); // efficiency cores included
detector_path := tprint("%/%", models_folder, model.detector);
recognizer_path := tprint("%/%", models_folder, model.recognizer);
use_gpu := false;
if prefer_gpu {
available, explanation := ocr_cuda_check();
use_gpu = available;
if !available engine.gpu_explanation = copy_string(explanation);
} else {
engine.gpu_explanation = copy_string("OCR runs on the CPU, as it was asked to.");
}
if use_gpu {
// A few CPU threads prepare images and decode results while the GPU runs the networks.
engine.workers = 4;
det_ok, rec_ok: bool;
det_error, rec_error: string;
engine.det, det_ok, det_error = ocr_create_session(engine.env, detector_path, true, 2);
if det_ok engine.rec, rec_ok, rec_error = ocr_create_session(engine.env, recognizer_path, true, 2);
if det_ok && rec_ok {
engine.provider = .CUDA;
} else {
ocr_release_session(*engine.det);
ocr_release_session(*engine.rec);
engine.gpu_explanation = sprint("OCR runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", ifx det_ok then rec_error else det_error);
use_gpu = false;
}
}
if !use_gpu {
// Several images at once, each network run spread over the rest of the cores.
engine.workers = clamp(cores / 4, 1, 8);
threads := max(cores / engine.workers, 1);
det_ok, rec_ok: bool;
det_error, rec_error: string;
engine.det, det_ok, det_error = ocr_create_session(engine.env, detector_path, false, threads);
if !det_ok { ocr_destroy(engine); return null, false, det_error; }
engine.rec, rec_ok, rec_error = ocr_create_session(engine.env, recognizer_path, false, threads);
if !rec_ok { ocr_destroy(engine); return null, false, rec_error; }
engine.provider = .CPU;
}
return engine, true, "";
}
ocr_provider :: (engine: *Ocr_Engine) -> Ocr_Provider {
return engine.provider;
}
// Whether the GPU is in use, and when not, why, and what to install so it can be.
ocr_gpu_status :: (engine: *Ocr_Engine) -> in_use: bool, explanation: string {
return engine.provider == .CUDA, engine.gpu_explanation;
}
ocr_destroy :: (engine: *Ocr_Engine) {
if !engine return;
ocr_release_session(*engine.det);
ocr_release_session(*engine.rec);
if engine.memory_info ort.api.ReleaseMemoryInfo(engine.memory_info);
if engine.env ort.api.ReleaseEnv(engine.env);
free(engine.gpu_explanation,, context.default_allocator);
free(engine,, context.default_allocator); // the dictionary and model strings are kept: small, and shared with ocr_list_models' caller
}
// The text in each image. Results (lines and their text) come from allocator. An image that fails
// is marked failed with the reason; the others are unaffected.
ocr_process :: (engine: *Ocr_Engine, images: [] Ocr_Image, allocator := context.allocator) -> [] Ocr_Result {
results := NewArray(images.count, Ocr_Result,, allocator);
if !images.count return results;
// Detection: images shared out among the workers.
per_image := NewArray(images.count, Image_Work,, context.default_allocator);
defer {
for * per_image array_free(it.boxes);
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;
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);
// Recognition: every line of every image, batched by width so little of each batch is padding.
jobs: [..] Line_Job;
jobs.allocator = context.default_allocator;
defer {
for jobs free(it.text,, context.default_allocator);
array_free(jobs);
}
for work: per_image {
if !work.ok continue;
for work.boxes array_add(*jobs, .{image = work.image, box = it, width = ocr_line_width(it, engine.model.rec_height)});
}
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; });
batch_size := max(engine.model.rec_batch, 1);
batch_count := (order.count + batch_size - 1) / batch_size;
Batch_Context :: struct { order: [] *Line_Job; size: s64; failed: bool; error: string; }
batches := Batch_Context.{order = order, size = batch_size};
run_on_workers(engine, batch_count, (engine: *Ocr_Engine, index: s64, data: *void) {
b := cast(*Batch_Context) data;
first := index * b.size;
count := min(b.size, b.order.count - first);
start := current_time_monotonic();
ok, error := ocr_recognize_batch(engine, array_view(b.order, first, count));
atomic_add(*engine.stats.recognize_us, to_microseconds(current_time_monotonic() - start));
if !ok { b.failed = true; b.error = copy_string(error,, context.default_allocator); } // the last error wins; any is enough to say why
}, *batches);
// Each image's lines, in reading order.
for work, image_index: per_image {
result := *results[image_index];
if !work.ok {
result.failed = true;
result.error = copy_string(work.error,, allocator);
free(work.error,, context.default_allocator);
continue;
}
if batches.failed {
result.failed = true;
result.error = copy_string(batches.error,, allocator);
continue;
}
lines: [..] Ocr_Line;
lines.allocator = allocator;
for job: jobs {
if job.image != work.image || !job.text.count continue;
array_add(*lines, .{x = job.box.x0, y = job.box.y0, width = job.box.x1 - job.box.x0, height = job.box.y1 - job.box.y0,
text = copy_string(job.text,, allocator), confidence = job.confidence});
}
// Rows by the middle of each line, then left to right within a row.
quick_sort(lines, (a: Ocr_Line, b: Ocr_Line) -> s64 {
ay, by := a.y + a.height / 2, b.y + b.height / 2;
if abs(ay - by) > min(a.height, b.height) / 2 return ifx ay < by then -1 else 1;
return a.x - b.x;
});
result.lines = lines;
engine.stats.lines += lines.count;
}
engine.stats.images += images.count;
free(batches.error,, context.default_allocator);
return results;
}
#scope_module
Image_Work :: struct {
image: *Ocr_Image;
boxes: [..] Detected_Box;
ok: bool;
error: string;
}
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};
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.
context.allocator = context.default_allocator;
shared := cast(*Shared) thread.data;
while true {
index := atomic_add(*shared.next, 1);
if index >= shared.count break;
shared.proc(shared.engine, index, shared.data);
reset_temporary_storage();
}
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;
}
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]);
}
}
#load "models.jai";
#load "detect.jai";
#load "recognize.jai";
#load "cuda_check.jai";
#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Math";
#import "Sort";
#import "System";
#import "Thread";
#import "Atomics";
#import "onnxruntime";
#if OS == .WINDOWS {
#import "Windows";
#import "Windows_Utf8";
} else {
#import "POSIX";
}

105
modules/OCR/recognize.jai Normal file
View File

@ -0,0 +1,105 @@
// Recognition: the text of each detected line. Each box is cropped from the image at full size,
// scaled to rec_height keeping its shape, and normalized to -1..1 in BGR order, as PaddleOCR's
// recognizers expect. Lines of similar width are batched, padded to the widest in the batch.
// The output is a probability per character class per step; greedy CTC decoding takes the best
// class at each step, merges repeats and drops blanks.
Line_Job :: struct {
image: *Ocr_Image;
box: Detected_Box;
width: s32; // after scaling to rec_height
// Filled in by recognition:
text: string; // default allocator
confidence: float;
}
// Recognizes jobs (one batch) in place.
ocr_recognize_batch :: (engine: *Ocr_Engine, jobs: [] *Line_Job) -> ok: bool, error: string {
if !jobs.count return true, "";
height := engine.model.rec_height;
width: s32 = 0;
for jobs width = max(width, it.width);
width = max(width, height); // the network needs a few steps even for one character
plane := cast(s64) height * width;
input := NewArray(plane * 3 * jobs.count, float, initialized = true,, context.default_allocator); // 0 is the padding
defer array_free(input,, context.default_allocator);
for job, n: jobs {
base := input.data + plane * 3 * n;
box := job.box;
bw := box.x1 - box.x0;
bh := box.y1 - box.y0;
image := job.image;
// Bilinear: the crops are mostly enlarged (small UI text up to 48 px), where nearest would be blocky.
x_scale := cast(float) bw / job.width;
y_scale := cast(float) bh / height;
for y: 0..height - 1 {
sy := (y + 0.5) * y_scale - 0.5 + box.y0;
y0 := clamp(cast(s32) floor(sy), 0, image.height - 1);
y1 := min(y0 + 1, image.height - 1);
fy := clamp(sy - y0, 0, 1);
row0 := image.pixels + cast(s64) y0 * image.stride;
row1 := image.pixels + cast(s64) y1 * image.stride;
for x: 0..job.width - 1 {
sx := (x + 0.5) * x_scale - 0.5 + box.x0;
x0 := clamp(cast(s32) floor(sx), 0, image.width - 1);
x1 := min(x0 + 1, image.width - 1);
fx := clamp(sx - x0, 0, 1);
i := cast(s64) y * width + x;
for c: 0..2 {
channel := 2 - c; // B, G, R from RGBA
top := row0[x0 * 4 + channel] * (1 - fx) + row0[x1 * 4 + channel] * fx;
bottom := row1[x0 * 4 + channel] * (1 - fx) + row1[x1 * 4 + channel] * fx;
value := top * (1 - fy) + bottom * fy;
base[i + plane * c] = (value / 255.0 - 0.5) / 0.5;
}
}
}
}
shape := s64.[jobs.count, 3, height, width];
out, ok, error := ocr_run(*engine.rec, engine.memory_info, input, shape);
if !ok return false, error;
defer ocr_release_output(*out);
if out.dim_count != 3 || out.dims[0] != jobs.count return false, tprint("The recognizer gave an output of shape %.", array_view(out.dims, 0, out.dim_count));
steps := out.dims[1];
classes := out.dims[2];
if classes != engine.dictionary.count + 2 {
return false, tprint("The recognizer has % classes, but its dictionary % characters (plus blank and space, % expected).", classes, engine.dictionary.count, engine.dictionary.count + 2);
}
for job, n: jobs {
builder: String_Builder;
builder.allocator = temp;
previous := -1;
probability_sum: float64;
characters := 0;
for t: 0..steps - 1 {
row := out.data + (n * steps + t) * classes;
best := 0;
best_p := row[0];
for c: 1..classes - 1 if row[c] > best_p { best = c; best_p = row[c]; }
if best != 0 && best != previous {
// Class 0 is the CTC blank, 1..count the dictionary, and the last one a space.
if best <= engine.dictionary.count append(*builder, engine.dictionary[best - 1]);
else append(*builder, " ");
probability_sum += best_p;
characters += 1;
}
previous = best;
}
job.text = copy_string(trim(builder_to_string(*builder,, temp)),, context.default_allocator);
job.confidence = ifx characters then cast(float) (probability_sum / characters) else 0;
}
return true, "";
}
// The width a box is recognized at: rec_height high, keeping its shape.
ocr_line_width :: (box: Detected_Box, rec_height: s32) -> s32 {
bw := box.x1 - box.x0;
bh := max(box.y1 - box.y0, 1);
return clamp(cast(s32) (cast(float) bw * rec_height / bh + 0.5), 1, REC_MAX_WIDTH);
}
REC_MAX_WIDTH :: 4096; // a line longer than this, at rec_height, is squeezed to fit

8
modules/Runtime_Support.jai
View File

@ -773,6 +773,12 @@ crash_log_hook: #type (crash_report: string);
// Set by the program: the crash report is <folder>-crash-report.txt, beside the executable on
// Windows and in the per-user data folder of this name on Linux and macOS.
crash_report_folder_name := "app";
// Set by the program when nobody is there to click a message box (the -headless argument, builds
// made with no_crash_dialog, command-line tools): a failure is still written to stderr, the crash
// report and the session log, but no message box is shown, and the process exits with code 1
// instead of going on to break into a debugger (which, with none attached, can raise another
// dialog). As in syncthing-imgui.
runtime_support_headless: bool;
#scope_file
report_error_in_yer_face_no_ctx :: (message: string, args: .. Any) #no_context {
@ -830,6 +836,8 @@ report_error_in_yer_face :: (message: string, args: .. Any, from_no_context := f
log_error(crash_report, flags = .TO_FILE_ONLY);
}
if runtime_support_headless exit(1);
// Show a Windows message box
#if OS == .WINDOWS {
W :: #import "Windows";

13307
modules/onnxruntime/bindings.jai Normal file
View File

File diff suppressed because it is too large Load Diff

128
modules/onnxruntime/fetch.jai Normal file
View File

@ -0,0 +1,128 @@
// Downloads the pinned onnxruntime release for this OS, checks its SHA-256, and puts the libraries
// in <os>/ and the C API header in include/. The binaries are not in git; run this after cloning.
//
// jai fetch.jai the GPU build (CUDA 13) on Windows and Linux; the CPU build on macOS
// jai fetch.jai - cpu the CPU-only build: smaller, and enough when there is no NVIDIA GPU
//
// Uses curl and tar, which Windows 10 and later ship with, and certutil / sha256sum / shasum.
ORT_VERSION :: "1.30.0";
Release :: struct { gpu: bool; archive: string; sha256: string; }
#if OS == .WINDOWS {
OS_FOLDER :: "windows";
RELEASES :: Release.[
.{true, "onnxruntime-win-x64-gpu_cuda13-1.30.0.zip", "8fa4b08359af682cd605892cb59077049700b640128bb93fd2c7776cf9f55bdc"},
.{false, "onnxruntime-win-x64-1.30.0.zip", "c6ba983baf5681af108599675d2a89c2d145512d02de28aed0bff177cd0ba949"},
];
LIBRARY_SUFFIX :: ".dll";
} else #if OS == .LINUX {
OS_FOLDER :: "linux";
RELEASES :: Release.[
.{true, "onnxruntime-linux-x64-gpu_cuda13-1.30.0.tgz", "382d79133112388cf94ce5855789b7c9bef12bef76a08b6b277e5a317213adcd"},
.{false, "onnxruntime-linux-x64-1.30.0.tgz", "a5ed5a3cac51fbb2e90da632ae43d19212faaa20e76484e62bcb7c23ddb3b3fd"},
];
LIBRARY_SUFFIX :: ".so";
} else #if OS == .MACOS {
OS_FOLDER :: "macos";
RELEASES :: Release.[
.{false, "onnxruntime-osx-arm64-1.30.0.tgz", "6ebb5062a934537c352937821f9fe9718e7de1a2db1122a93dd363ffd53a7012"},
];
LIBRARY_SUFFIX :: ".dylib";
}
#run,stallable { if !fetch() exit(1); }
fetch :: () -> bool {
set_build_options_dc(.{do_output = false});
args := get_build_options().compile_time_command_line;
want_gpu := !array_find(args, "cpu");
release: Release;
for RELEASES if it.gpu == want_gpu || RELEASES.count == 1 { release = it; break; }
if !release.archive { log_error("No % build of onnxruntime for this OS.", ifx want_gpu then "GPU" else "CPU"); return false; }
here := trim_right(path_strip_filename(#file), "/\\");
set_working_directory(here);
download_folder := tprint("%/download", here); // absolute paths throughout
make_directory_if_it_does_not_exist(download_folder);
archive_path := tprint("%/%", download_folder, release.archive);
if file_exists(archive_path) && sha256_of(archive_path) == release.sha256 {
log("% is already downloaded.", release.archive);
} else {
url := tprint("https://github.com/microsoft/onnxruntime/releases/download/v%/%", ORT_VERSION, release.archive);
log("Downloading %...", url);
result := run_command("curl", "-L", "--fail", "--progress-bar", "-o", archive_path, url);
if result.exit_code != 0 { log_error("curl failed (exit code %).", result.exit_code); return false; }
hash := sha256_of(archive_path);
if hash != release.sha256 {
log_error("% has SHA-256 %, expected %; it was not used.", archive_path, hash, release.sha256);
file_delete(archive_path);
return false;
}
}
extract_folder := tprint("%/extract", download_folder);
remove_folder(extract_folder);
make_directory_if_it_does_not_exist(extract_folder);
// Windows' own tar reads .zip; another tar earlier on PATH (Git's) does not.
tar := ifx OS == .WINDOWS then "C:/Windows/System32/tar.exe" else "tar";
result := run_command(tar, "-xf", archive_path, "-C", extract_folder);
if result.exit_code != 0 { log_error("tar failed (exit code %).", result.exit_code); return false; }
defer remove_folder(extract_folder);
root := tprint("%/%", extract_folder, path_strip_extension(path_strip_extension(release.archive))); // .tgz, or .zip
if !file_exists(tprint("%/lib", root)) root = tprint("%/%", extract_folder, path_strip_extension(release.archive));
// Replace whatever build was here before, so a GPU and a CPU build never mix.
library_folder := tprint("%/%", here, OS_FOLDER);
remove_folder(library_folder);
make_directory_if_it_does_not_exist(library_folder);
copied := 0;
for file_list(tprint("%/lib", root)) {
name := path_filename(it);
if !contains(name, LIBRARY_SUFFIX) continue; // .so.1.30.0 and the like included
if !copy_file(it, tprint("%/%", library_folder, name)) { log_error("Could not copy %.", name); return false; }
copied += 1;
}
make_directory_if_it_does_not_exist(tprint("%/include", here));
for file_list(tprint("%/include", root)) {
if !copy_file(it, tprint("%/include/%", here, path_filename(it))) { log_error("Could not copy %.", it); return false; }
}
write_entire_file(tprint("%/VERSION.txt", library_folder), tprint("%\n%\n", ORT_VERSION, release.archive));
log("Put % libraries in %/% and the headers in %/include (onnxruntime %, %).", copied, here, OS_FOLDER, here, ORT_VERSION, ifx release.gpu then "GPU" else "CPU");
return true;
}
// With the OS's own tools: File's delete_directory hands Windows a relative path it cannot resolve,
// and its deletion was seen racing the copies that followed.
remove_folder :: (path: string) {
if !file_exists(path) return;
#if OS == .WINDOWS run_command("cmd", "/c", "rmdir", "/s", "/q", replace(path, "/", "\\"));
else run_command("rm", "-rf", path);
}
// Lowercase hex, or "" when it could not be worked out.
sha256_of :: (path: string) -> string {
#if OS == .WINDOWS {
result, output := run_command("certutil", "-hashfile", path, "SHA256", capture_and_return_output = true);
if result.exit_code != 0 return "";
lines := split(output, "\n");
if lines.count < 2 return "";
return to_lower_copy(replace(trim(lines[1]), " ", "")); // line 2 is the hash
} else {
command := ifx OS == .MACOS then string.["shasum", "-a", "256", path] else string.["sha256sum", path];
result, output := run_command(..command, capture_and_return_output = true);
if result.exit_code != 0 return "";
return to_lower_copy(trim(split(output, " ")[0]));
}
}
#import "Basic";
#import "Compiler";
#import "File";
#import "File_Utilities";
#import "Process";
#import "String";

38
modules/onnxruntime/generate.jai Normal file
View File

@ -0,0 +1,38 @@
// Makes bindings.jai from include/onnxruntime_c_api.h: the OrtApi function table and every type it
// uses. No functions are bound to a library: module.jai loads onnxruntime at run time and finds its
// one export, OrtGetApiBase, so a program runs (without OCR) when onnxruntime is not there.
//
// jai generate.jai (after jai fetch.jai, which puts the header in include/)
#run,stallable {
set_build_options_dc(.{do_output = false});
if !generate() compiler_set_workspace_status(.FAILED);
}
generate :: () -> bool {
here := trim_right(path_strip_filename(#file), "/\\");
set_working_directory(here);
options: Generate_Bindings_Options;
{
using options;
array_add(*source_files, "include/onnxruntime_c_api.h");
array_add(*include_paths, "include");
array_add(*extra_clang_arguments, "-x", "c");
generate_library_declarations = false; // nothing is linked; see module.jai
// No libraries are named, so the default SYMBOLS_WITH_UNKNOWN_FOREIGN_LIBS strip drops OrtGetApiBase
// and the other exports: module.jai finds OrtGetApiBase at run time.
// The table and its types are all prefixed Ort; keep the names as the C API documents them.
strip_enum_prefixes = false;
auto_detect_enum_prefixes = false;
// Every declaration in onnxruntime's own headers; the system ones they include are skipped.
array_add(*path_fragments_to_treat_as_non_system_paths, "onnxruntime");
}
return generate_bindings(options, "bindings.jai");
}
#import "Basic";
#import "Compiler";
#import "File";
#import "String";
#import "Bindings_Generator";

19
modules/onnxruntime/include/cpu_provider_factory.h Normal file
View File

@ -0,0 +1,19 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "onnxruntime_c_api.h"
#ifdef __cplusplus
extern "C" {
#endif
/**
* \param use_arena zero: false. non-zero: true.
*/
ORT_EXPORT
ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_CPU, _In_ OrtSessionOptions* options, int use_arena)
ORT_ALL_ARGS_NONNULL;
#ifdef __cplusplus
}
#endif

8782
modules/onnxruntime/include/onnxruntime_c_api.h Normal file
View File

File diff suppressed because it is too large Load Diff

3820
modules/onnxruntime/include/onnxruntime_cxx_api.h Normal file
View File

File diff suppressed because it is too large Load Diff

4270
modules/onnxruntime/include/onnxruntime_cxx_inline.h Normal file
View File

File diff suppressed because it is too large Load Diff

24
modules/onnxruntime/include/onnxruntime_env_config_keys.h Normal file
View File

@ -0,0 +1,24 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
// This file contains well-known keys for OrtEnv configuration entries, which may be used to configure EPs or
// other global settings.
// Refer to OrtEnvCreationOptions::config_entries and OrtApi::CreateEnvWithOptions.
// This file does NOT specify all available keys as EPs may accept custom entries with the prefix "ep.<ep_name>.".
// Key for a boolean option that, when enabled, allows EP factories to create virtual OrtHardwareDevice
// instances via OrtEpApi::CreateHardwareDevice().
//
// This config entry is automatically set to "1" by ORT if an application registers an EP library with a registration
// name that ends in the suffix ".virtual". See OrtApi::RegisterExecutionProviderLibrary().
//
// Note: A virtual OrtHardwareDevice does not represent actual hardware on the device, and is identified via the
// metadata entry "is_virtual" with a value of "1".
//
// Allowed values:
// - "0": Default. Creation of virtual devices is not allowed.
// This is the assumed default value if this key is not present in the environment's configuration entries.
// - "1": Creation of virtual devices is allowed.
static const char* const kOrtEnvAllowVirtualDevices = "allow_virtual_devices";

3180
modules/onnxruntime/include/onnxruntime_ep_c_api.h Normal file
View File

File diff suppressed because it is too large Load Diff

59
modules/onnxruntime/include/onnxruntime_ep_device_ep_metadata_keys.h Normal file
View File

@ -0,0 +1,59 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
// This file contains well-known keys for OrtEpDevice and OrtHardwareDevice metadata entries.
// It does NOT specify all available metadata keys.
// Key for the execution provider version string. This should be available for all plugin EPs.
static const char* const kOrtEpDevice_EpMetadataKey_Version = "version";
// Key for the execution provider OS driver version.
// Value should be a 4-part dot-separated version string in the format "a.b.c.d" (e.g., "31.0.101.4502").
// This maps to the Windows DXCore adapter property DXCoreAdapterProperty::DriverVersion
// (https://learn.microsoft.com/en-us/windows/win32/api/dxcore_interface/ne-dxcore_interface-dxcoreadapterproperty).
// On non-Windows platforms, the EP should provide an equivalent OS-level driver version if available.
static const char* const kOrtEpDevice_EpMetadataKey_OSDriverVersion = "os_driver_version";
// Prefix for execution provider compatibility information stored in model metadata.
// Used when generating EP context models to store compatibility strings for each EP.
// Full key format: "ep_compatibility_info.<EP_TYPE>"
static const char* const kOrtModelMetadata_EpCompatibilityInfoPrefix = "ep_compatibility_info.";
// Key for the execution provider library path (for dynamically loaded EPs)
static const char* const kOrtEpDevice_EpMetadataKey_LibraryPath = "library_path";
// Optional metadata key for the execution provider's preferred layout of the Value KV-cache tensors
// (the past_value input and present_value output) of com.microsoft.GroupQueryAttention.
// Possible values:
// - "BNSH": (batch_size, num_heads, sequence_length, head_size). This is the assumed default value
// if this metadata key is not present, and matches the operator schema.
// - "BNHS": (batch_size, num_heads, head_size, sequence_length).
// An EP that reports "BNHS" is expected to fuse the Transpose -> GroupQueryAttention -> Transpose
// sequence that ORT inserts when the application selects that layout.
// The application passes the layout it has chosen to the session via the
// kOrtSessionOptionsGqaValueLayout session option (see onnxruntime_session_options_config_keys.h).
static const char* const kOrtEpDevice_EpMetadataKey_GqaPreferredValueLayout = "gqa_preferred_value_layout";
// Optional metadata key to determine if a OrtHardwareDevice represents a virtual (non-hardware) device.
// Possible values:
// - "0": OrtHardwareDevice is not virtual (i.e., actual hardware device). This is the assumed default value
// if this metadata key is not present.
// - "1": OrtHardwareDevice is virtual.
static const char* const kOrtHardwareDevice_MetadataKey_IsVirtual = "is_virtual";
// Key for the execution provider's weightless mode support on a specific device.
// Set by the EP during GetSupportedDevices() via CreateEpDevice() metadata.
// The app can read it via EpDevice_EpMetadata() to check device-specific weightless capability
// before calling ModelCompilationOptions_SetWeightlessEnabled().
//
// Possible values:
// - "none": EP does not support weightless mode on this device. This is the assumed default value
// if this metadata key is not present.
// - "external_only": EP supports weightless mode for external initializers only (e.g., older
// hardware/driver that must transform internal constants).
// - "all": EP supports weightless mode for all initializers (internal and external).
//
// \since Version 1.29.
static const char* const kOrtEpDevice_EpMetadataKey_WeightlessSupport = "weightless_support";

98
modules/onnxruntime/include/onnxruntime_error_code.h Normal file
View File

@ -0,0 +1,98 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
/** \addtogroup Global
* ONNX Runtime C API
* @{
*/
#ifdef __cplusplus
extern "C" {
#endif
/** \brief Error codes reported by ONNX Runtime.
*
* The error code associated with an ::OrtStatus.
*/
typedef enum OrtErrorCode {
/**
* Success. No error occurred.
*/
ORT_OK,
/**
* Generic failure that does not map to a more specific error code. Consult the error message for details.
*/
ORT_FAIL,
/**
* A caller-supplied argument was invalid (e.g. NULL pointer, out-of-range value, mismatched shape/rank, or bad
* configuration).
*/
ORT_INVALID_ARGUMENT,
/**
* A required file (such as a model file) does not exist.
*/
ORT_NO_SUCHFILE,
/**
* Legacy/unused but retained for ABI compatibility. Historically returned when a model could not be found by name in
* the ONNX Runtime Server (removed in 2022).
*/
ORT_NO_MODEL,
/**
* A hardware accelerator or backend engine reported a failure (e.g. a device crash or other device-level error).
*/
ORT_ENGINE_ERROR,
/**
* A generic runtime exception was caught. The error message is the primary source of detail.
*/
ORT_RUNTIME_EXCEPTION,
/**
* Protobuf parsing or serialization failed.
*/
ORT_INVALID_PROTOBUF,
/**
* Invalid session state for the requested operation. Despite the name, this code does not mean "success, model
* loaded"; it is returned when the session is in the wrong state for the requested call (e.g. a model is already
* loaded, the session is already initialized, or no model has been loaded yet). The name is historical and is
* retained for ABI compatibility; consult the error message for the specific condition.
*/
ORT_MODEL_LOADED,
/**
* The requested functionality is not implemented in this build.
*/
ORT_NOT_IMPLEMENTED,
/**
* The model graph is structurally invalid (e.g. recursive function definitions, invalid tensor dimensions, or
* malformed nodes).
*/
ORT_INVALID_GRAPH,
/**
* An execution provider reported a generic failure.
*/
ORT_EP_FAIL,
/**
* Model loading or session initialization was canceled at the caller's request.
*/
ORT_MODEL_LOAD_CANCELED,
/**
* The model requires compilation by an execution provider, but compilation was disabled via session options.
*/
ORT_MODEL_REQUIRES_COMPILATION,
/**
* A requested resource could not be found.
*/
ORT_NOT_FOUND,
/**
* The execution provider's hardware device has been reset.
*
* The caller should stop using the existing session (and release it) and create a new session.
*/
ORT_DEVICE_RESET,
} OrtErrorCode;
#ifdef __cplusplus
}
#endif
/// @}

134
modules/onnxruntime/include/onnxruntime_experimental_c_api.h Normal file
View File

@ -0,0 +1,134 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// Experimental C API consumer header.
//
// This header provides C function pointer typedefs and name constants for experimental ORT API functions, as well as
// any auxiliary declarations required by the experimental API functions.
//
// The function pointer typedefs and name constants should be used together with the experimental API lookup function
// `OrtApi::GetExperimentalFunction()`. A function's availability should always be checked at runtime (the lookup
// returns nullptr if the function is not present).
//
// C++ API consumers should use the companion header, onnxruntime_experimental_cxx_api.h, which provides typed
// accessors for the experimental API functions.
//
// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added, changed, or removed between
// releases without notice. Anything in this file should be treated as experimental and unstable.
//
// C usage:
// OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn fn =
// (OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn)api->GetExperimentalFunction(
// kOrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_FnName);
// if (fn) {
// OrtStatusPtr status = fn(&result);
// }
//
// C++ usage (see onnxruntime_experimental_cxx_api.h):
// if (auto* fn = Ort::Experimental::Get_OrtApi_ExperimentalApiTest_SinceV28_Fn(api)) {
// Ort::Status status(fn(&result));
// }
#pragma once
#include "onnxruntime_c_api.h"
//
// Auxiliary declarations
//
// Declarations of auxiliary types or typedefs required by experimental APIs go here.
//
// For example, if an experimental API uses a new type, OrtExperimentalType, we would declare it in this file:
// ORT_RUNTIME_CLASS(ExperimentalType);
//
ORT_RUNTIME_CLASS(ModelPackageOptions);
ORT_RUNTIME_CLASS(ModelPackageContext);
ORT_RUNTIME_CLASS(ModelPackageComponentContext);
// Opaque handle holding the EPContext callbacks and opaque state extracted from an OrtSessionOptions instance. Used by
// the experimental OrtEpApi_* EPContext data functions. Create via OrtEpApi_SessionOptions_GetEpContextConfig and
// release with OrtEpApi_ReleaseEpContextConfig.
ORT_RUNTIME_CLASS(EpContextConfig);
/** \brief Function called to write named binary data.
*
* This callback is currently used for EPContext binary data, but its contract is intentionally generic so future APIs
* can reuse it for other named data payloads. The callback is called synchronously by the component that receives it.
* ORT does not own or retain buffer after the callback returns. ORT does not serialize invocations made by different
* EP instances or worker threads.
*
* Each callback invocation represents one complete write operation for name. The callback signature does not
* provide an offset, sequence number, or final-chunk marker, so the component invoking the callback must define any
* chunked ordering and completion contract with the application. Current EPContext use should prefer a single callback
* invocation per EPContext binary unless chunking semantics are documented by the EP.
*
* The application's implementation can process the data in any way (e.g., encrypt and store, upload to cloud storage,
* or compress) before persisting it.
*
* \param[in] state Opaque pointer holding the user's state. ORT does not own or manage this pointer. The application
* must keep it valid for the duration required by the API that accepted the callback and must provide
* any synchronization required if it can be used concurrently.
* \param[in] name The file name or logical data identifier as a null-terminated UTF-8 string.
* \param[in] buffer The buffer containing data to write.
* \param[in] buffer_num_bytes The size of the buffer in bytes.
*
* \return OrtStatus* Write status. Return nullptr on success.
* On failure, use CreateStatus to provide error info with an appropriate OrtErrorCode
* (e.g., ORT_FAIL); ORT propagates the returned code. ORT will release the OrtStatus* if not null.
*/
typedef OrtStatus*(ORT_API_CALL* OrtWriteNamedBufferFunc)(_In_ void* state,
_In_ const char* name,
_In_ const void* buffer,
_In_ size_t buffer_num_bytes);
/** \brief Function called to read named binary data.
*
* This callback is currently used for EPContext binary data, but its contract is intentionally generic so future APIs
* can reuse it for other named data payloads. The application reads, processes (e.g., decrypts, decompresses,
* downloads), and returns the requested data. ORT provides an allocator so the application can allocate the output
* buffer directly. The callback is called synchronously by the component that receives it. ORT does not serialize
* invocations made by different EP instances or worker threads.
*
* \param[in] state Opaque pointer holding the user's state. ORT does not own or manage this pointer. The application
* must keep it valid for the duration required by the API that accepted the callback and must provide
* any synchronization required if it can be used concurrently.
* \param[in] name The file name or logical data identifier to read as a null-terminated UTF-8 string.
* \param[in] allocator ORT-provided allocator. The application must use this to allocate the output buffer.
* \param[out] buffer Set by the implementation to the allocated buffer containing the output data.
* \param[out] data_size Set by the implementation to the size of the output data in bytes.
*
* \return OrtStatus* Read status. Return nullptr on success.
* On failure, use CreateStatus to provide error info with an appropriate OrtErrorCode
* (e.g., ORT_FAIL); ORT propagates the returned code. ORT will release the OrtStatus* if not null.
*/
typedef OrtStatus*(ORT_API_CALL* OrtReadNamedBufferFunc)(_In_ void* state,
_In_ const char* name,
_In_ OrtAllocator* allocator,
_Outptr_ void** buffer,
_Out_ size_t* data_size);
//
// C function pointer typedefs and name constants
//
// For each ORT_EXPERIMENTAL_API(VER, RET, NAME, ...) entry in the .inc file, this generates:
//
// // Function pointer typedef:
// typedef RET(ORT_API_CALL* OrtExperimental_<NAME>_SinceV<VER>_Fn)(...) NO_EXCEPTION;
//
// // Name constant for lookup:
// static const char* const kOrtExperimental_<NAME>_SinceV<VER>_FnName = "<NAME>_SinceV<VER>";
//
// Example: ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_ExperimentalApiTest, _Out_ int64_t* out) produces:
// typedef OrtStatusPtr(ORT_API_CALL* OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn)(
// _Out_ int64_t* out) NO_EXCEPTION;
// static const char* const kOrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_FnName =
// "OrtApi_ExperimentalApiTest_SinceV28";
#define ORT_EXPERIMENTAL_API(VER, RET, NAME, ...) \
typedef RET(ORT_API_CALL* OrtExperimental_##NAME##_SinceV##VER##_Fn)(__VA_ARGS__) NO_EXCEPTION; \
static const char* const kOrtExperimental_##NAME##_SinceV##VER##_FnName = #NAME "_SinceV" #VER;
#include "onnxruntime_experimental_c_api.inc"
#undef ORT_EXPERIMENTAL_API

426
modules/onnxruntime/include/onnxruntime_experimental_c_api.inc Normal file
View File

@ -0,0 +1,426 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// Central declaration list for experimental C API functions.
//
// This file is included multiple times with different definitions of ORT_EXPERIMENTAL_API to generate constructs like
// typedefs, name constants, registration table entries, and C++ accessors.
//
// The X-macro (https://en.wikipedia.org/wiki/X_macro) ORT_EXPERIMENTAL_API has the following signature:
// ORT_EXPERIMENTAL_API(SinceVersion, ReturnType, Name, Params...)
//
// SinceVersion: The ORT API version in which the function was first introduced.
// Must be a numeric literal (e.g., 28), not a macro like ORT_API_VERSION.
// ReturnType: The return type of the function (typically OrtStatusPtr).
// Name: The base name (without version suffix). By convention, prefixed with the target stable struct
// (e.g., OrtApi_Foo for functions destined for OrtApi).
// Params: The parameter list, using SAL annotations as in other ORT C API declarations.
//
// The function lookup name is mechanically generated as Name_SinceV<SinceVersion>.
//
// Experimental function lifecycle rules:
// - Adding: Add one ORT_EXPERIMENTAL_API line with the current API version and implement it.
// - Removing: Delete the line. The versioned name is inherently unique and cannot be reused.
// - Changing a signature: Add a new entry at the current API version (producing a new name) and optionally delete
// the old entry. Both can coexist during a transition period.
// - Promoting to stable: Move the function to the stable API struct (append-only, drop the _SinceV<SinceVersion>
// suffix). Delete the experimental entry.
/** \brief Test-only experimental function that exercises the lookup mechanism end-to-end.
*
* This entry is kept permanently for regression coverage of the experimental API mechanism.
*
* \param[out] out The output of this function, the value 12345.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_ExperimentalApiTest, _Out_ int64_t* out)
// ---------------------------------------------------------------------------
// OrtModelPackageApi
//
// A model package is a directory containing one or more *components* (logical models).
// Each component has one or more *variants*, where each variant targets a single
// execution provider (EP). The package manifest declares the EP name, device type,
// and an optional compatibility string for every variant so that the runtime can
// automatically select the best variant for the hardware and EPs available in the
// caller's session options.
//
// The functions below support:
// - creating model package options that capture EP configuration from OrtSessionOptions,
// - loading a package context (manifest + metadata) from a package root path,
// - querying component/variant metadata including per-variant EP information,
// - selecting a component (which also resolves the best-matching variant),
// - querying the selected variant's name and folder path,
// - creating an OrtSession from the selected component context.
//
// Typical flow:
// 1) Create model package options:
// - OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions
// 2) Load package metadata:
// - OrtModelPackageApi_CreateModelPackageContext
// 3) Query metadata (optional):
// - OrtModelPackageApi_ModelPackage_GetSchemaVersion
// - OrtModelPackageApi_ModelPackage_GetComponentCount
// - OrtModelPackageApi_ModelPackage_GetComponentNames
// - OrtModelPackageApi_ModelPackage_GetVariantCount
// - OrtModelPackageApi_ModelPackage_GetVariantNames
// - OrtModelPackageApi_ModelPackage_GetVariantEpName
// - OrtModelPackageApi_ModelPackage_ResolveStringRef
// 4) Select a component and resolve variant:
// - OrtModelPackageApi_SelectComponent
// 5) Query selected variant info (optional):
// - OrtModelPackageApi_ModelPackageComponent_GetSelectedVariantName
// - OrtModelPackageApi_ModelPackageComponent_GetSelectedVariantFolderPath
// 6) Create session:
// - OrtModelPackageApi_CreateSession
//
// Ownership:
// - Release objects created by this API with the corresponding release functions:
// OrtModelPackageApi_ReleaseModelPackageOptions,
// OrtModelPackageApi_ReleaseModelPackageContext,
// OrtModelPackageApi_ReleaseModelPackageComponentContext.
//
// Opaque handles (OrtModelPackageOptions, OrtModelPackageContext, OrtModelPackageComponentContext)
// are declared in onnxruntime_experimental_c_api.h.
// ---------------------------------------------------------------------------
/** \brief Create model package options from an existing OrtSessionOptions.
*
* Captures EP configuration (registered execution providers and their devices) from
* the session options for use during variant selection. The resulting OrtModelPackageOptions
* is passed to OrtModelPackageApi_SelectComponent to resolve the best variant for the
* available EPs.
*
* \param[in] env The ORT environment.
* \param[in] session_options Session options containing registered EPs.
* \param[out] out Receives the newly created OrtModelPackageOptions. Must be released
* with OrtModelPackageApi_ReleaseModelPackageOptions.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions,
_In_ const OrtEnv* env,
_In_ const OrtSessionOptions* session_options,
_Outptr_ OrtModelPackageOptions** out)
/** \brief Release an OrtModelPackageOptions created by
* OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions.
*/
ORT_EXPERIMENTAL_API(28, void, OrtModelPackageApi_ReleaseModelPackageOptions,
_Frees_ptr_opt_ OrtModelPackageOptions* options)
/** \brief Create a model package context by parsing the package at the given root path.
*
* Parses the manifest.json and component metadata from the specified directory.
* The returned context provides read-only access to the package structure (components,
* variants, EP declarations).
*
* \param[in] package_root Path to the model package root directory (containing manifest.json).
* \param[out] out Receives the newly created OrtModelPackageContext. Must be released
* with OrtModelPackageApi_ReleaseModelPackageContext.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_CreateModelPackageContext,
_In_ const ORTCHAR_T* package_root,
_Outptr_ OrtModelPackageContext** out)
/** \brief Release an OrtModelPackageContext created by OrtModelPackageApi_CreateModelPackageContext. */
ORT_EXPERIMENTAL_API(28, void, OrtModelPackageApi_ReleaseModelPackageContext,
_Frees_ptr_opt_ OrtModelPackageContext* ctx)
/** \brief Get the schema version declared in the model package manifest.
*
* \param[in] ctx The model package context.
* \param[out] out_version Receives the schema version number.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetSchemaVersion,
_In_ const OrtModelPackageContext* ctx,
_Out_ int64_t* out_version)
/** \brief Get the number of components in the model package.
*
* \param[in] ctx The model package context.
* \param[out] out_count Receives the component count.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetComponentCount,
_In_ const OrtModelPackageContext* ctx,
_Out_ size_t* out_count)
/** \brief Get the names of all components in the model package.
*
* Returns a pointer to an array of UTF-8 component name strings. The array and its
* strings are owned by `ctx` and remain valid until the context is released.
*
* \param[in] ctx The model package context.
* \param[out] out_names Receives a pointer to an array of component name strings.
* \param[out] out_count Receives the number of elements in the array.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetComponentNames,
_In_ const OrtModelPackageContext* ctx,
_Outptr_result_buffer_maybenull_(*out_count) const char* const** out_names,
_Out_ size_t* out_count)
/** \brief Get the number of variants for a given component.
*
* \param[in] ctx The model package context.
* \param[in] component_name Name of the component to query.
* \param[out] out_count Receives the variant count.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetVariantCount,
_In_ const OrtModelPackageContext* ctx,
_In_ const char* component_name,
_Out_ size_t* out_count)
/** \brief Get the names of all variants for a given component.
*
* Returns a pointer to an array of UTF-8 variant name strings. The array and its
* strings are owned by `ctx` and remain valid until the context is released.
*
* \param[in] ctx The model package context.
* \param[in] component_name Name of the component to query.
* \param[out] out_variant_names Receives a pointer to an array of variant name strings.
* \param[out] out_count Receives the number of elements in the array.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetVariantNames,
_In_ const OrtModelPackageContext* ctx,
_In_ const char* component_name,
_Outptr_result_buffer_maybenull_(*out_count) const char* const** out_variant_names,
_Out_ size_t* out_count)
/** \brief Get the EP name declared for a (component, variant) pair.
*
* Each variant targets a single EP. `out_ep` receives the EP name string.
* When the variant does not declare an EP, the returned pointer is NULL.
* String memory is owned by `ctx` and remains valid until the context is released.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_GetVariantEpName,
_In_ const OrtModelPackageContext* ctx,
_In_ const char* component_name,
_In_ const char* variant_name,
_Outptr_result_maybenull_ const char** out_ep)
/** \brief Resolve a path reference declared inside the package against the model_package rules.
*
* Handles the path forms a package may use:
* - "sha256:<hex>" or "sha256:<hex>/<tail>": a content-addressed shared asset. Resolves to
* the asset's on-disk directory (honoring manifest shared_assets overrides), optionally
* joined with the confined tail. Errors if the asset is not declared/discovered.
* - any other value: a relative path resolved against `base_dir` (or the package root when
* `base_dir` is NULL), with portable-layout confinement (no absolute paths, no "..").
*
* When `must_exist` is non-zero the resolved path must exist on disk. `out_path` is owned by
* `ctx` and remains valid until the next call to this function on the same context.
*
* \param[in] ctx The package context returned by OrtModelPackageApi_CreateModelPackageContext.
* \param[in] base_dir Base directory for relative inputs. May be NULL to use the package root.
* \param[in] input The path reference to resolve.
* \param[in] must_exist Non-zero to require that the resolved path exists on disk.
* \param[out] out_path Receives the resolved path string.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackage_ResolveStringRef,
_In_ const OrtModelPackageContext* ctx,
_In_opt_ const char* base_dir,
_In_ const char* input,
_In_ int must_exist,
_Outptr_ const char** out_path)
/** \brief Select a component model and return an opaque component instance.
*
* The variant selection is also performed during this call based on the component
* metadata and the provided options. The returned `OrtModelPackageComponentContext*` is
* independent of `context` lifetime and must be released via
* OrtModelPackageApi_ReleaseModelPackageComponentContext.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_SelectComponent,
_In_ const OrtModelPackageContext* context,
_In_ const char* component_name,
_In_ const OrtModelPackageOptions* options,
_Outptr_ OrtModelPackageComponentContext** out)
/** \brief Release an OrtModelPackageComponentContext created by OrtModelPackageApi_SelectComponent. */
ORT_EXPERIMENTAL_API(28, void, OrtModelPackageApi_ReleaseModelPackageComponentContext,
_Frees_ptr_opt_ OrtModelPackageComponentContext* ctx)
/** \brief Get the name of the selected variant after OrtModelPackageApi_SelectComponent has been called.
*
* String memory is owned by `ctx` and remains valid until the context is released.
*
* \param[in] ctx The component context returned by OrtModelPackageApi_SelectComponent.
* \param[out] out_name Receives the selected variant's name string.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackageComponent_GetSelectedVariantName,
_In_ const OrtModelPackageComponentContext* ctx,
_Outptr_ const char** out_name)
/** \brief Get the folder path of the selected variant.
*
* Returns the resolved absolute path to the variant's directory on disk.
* The string is owned by `ctx` and remains valid until the context is released.
*
* \param[in] ctx The component context returned by OrtModelPackageApi_SelectComponent.
* \param[out] folder_path Receives the variant folder path string.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_ModelPackageComponent_GetSelectedVariantFolderPath,
_In_ const OrtModelPackageComponentContext* ctx,
_Outptr_ const ORTCHAR_T** folder_path)
/** \brief Create an OrtSession for a selected file within a component model variant.
*
* The chosen variant (and thus its EP selection) is determined by `context`, which
* was built from an OrtSessionOptions via OrtModelPackageApi_CreateModelPackageOptionsFromSessionOptions.
*
* Session options precedence:
* 1. session_options == NULL (default path):
* ORT uses the OrtSessionOptions that was captured when `context` was created.
* Any variant-specific session and provider options declared in the package
* metadata are merged on top.
*
* 2. session_options != NULL (advanced path):
* ORT uses the caller-provided OrtSessionOptions as-is. Variant-specific
* session and provider options from the package metadata are NOT applied.
* Use this when custom EP setup is required (e.g., shared CUDA streams,
* shared QNN EP contexts, custom allocators).
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_CreateSession,
_In_ const OrtEnv* env,
_In_ OrtModelPackageComponentContext* context,
_In_opt_ const OrtSessionOptions* session_options,
_Outptr_ OrtSession** session)
/** \brief Registers a callback to provide EPContext binary data during session load.
*
* When loading a compiled model with external (non-embedded) EPContext binary data, an execution provider can
* retrieve this callback from OrtEpContextConfig and call it instead of reading the binary data from disk.
*
* Reading happens at session load, so this callback is configured on OrtSessionOptions. The corresponding write
* callback runs only at compile time and is configured on OrtModelCompilationOptions via
* OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc.
*
* The state pointer is stored as-is and is not owned by ORT. It must remain valid while any session or EP created
* from these options may call the callback. If the same state may be used by multiple EPs or threads, the application
* is responsible for synchronization.
*
* \param[in] options The OrtSessionOptions instance.
* \param[in] read_func The OrtReadNamedBufferFunc callback. Pass NULL to clear a previously set callback (any
* previously set state is cleared as well).
* \param[in] state Opaque state passed to read_func. Can be NULL. Ignored when read_func is NULL.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_SessionOptions_SetEpContextDataReadFunc,
_Inout_ OrtSessionOptions* options, _In_opt_ OrtReadNamedBufferFunc read_func, _In_opt_ void* state)
/** \brief Sets a callback for writing EPContext binary data during compilation.
*
* When EPContext embed mode is disabled, execution providers can retrieve this callback from OrtEpContextConfig and
* call it instead of writing EPContext binary data directly to disk.
*
* This callback may be used together with OrtCompileApi::ModelCompilationOptions_SetEpContextBinaryInformation. The
* binary information still describes the compiled model/output location that EPs may use to generate stable logical
* EPContext data names or as a file-fallback location. If this callback is configured, EPs should call it for
* EPContext binary data instead of writing that data to the fallback file path.
*
* Writing happens only at compile time, so this callback is configured on OrtModelCompilationOptions. The
* corresponding read callback runs at session load and is configured on OrtSessionOptions via
* OrtApi_SessionOptions_SetEpContextDataReadFunc.
*
* The state pointer is stored as-is and is not owned by ORT. It must remain valid for the duration of the compile
* operation that may call the callback. If the same state may be used by multiple EPs or threads, the application is
* responsible for synchronization.
*
* Like OrtApi_SessionOptions_SetEpContextDataReadFunc, passing a NULL write_func clears any previously set callback
* (any previously set state is cleared as well). Calling this multiple times overwrites the previously configured
* callback.
*
* \param[in] model_compile_options The OrtModelCompilationOptions instance.
* \param[in] write_func The OrtWriteNamedBufferFunc callback used to write EPContext bytes. Pass NULL to clear a
* previously set callback (any previously set state is cleared as well).
* \param[in] state Opaque state passed to write_func. Can be NULL. Ignored when write_func is NULL.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc,
_In_ OrtModelCompilationOptions* model_compile_options,
_In_opt_ OrtWriteNamedBufferFunc write_func, _In_opt_ void* state)
/** \brief Extracts the EPContext configuration (callbacks and state) from an OrtSessionOptions instance.
*
* The EP should call this during CreateEp() while session_options is still valid, and store the returned handle for
* use during Compile(). On success, `*config` is set to a non-NULL handle that must be released with
* OrtEpApi_ReleaseEpContextConfig. On failure, an error status is returned and `*config` is not modified.
*
* The returned handle owns only ORT's copy of callback function pointers and opaque state pointer values. It does not
* own the application-provided state. The application is responsible for keeping callback state valid and
* synchronized while an EP may call callbacks retrieved from this config.
*
* \param[in] session_options The OrtSessionOptions instance.
* \param[out] config The extracted OrtEpContextConfig.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_SessionOptions_GetEpContextConfig,
_In_ const OrtSessionOptions* session_options, _Outptr_ OrtEpContextConfig** config)
/** \brief Release an OrtEpContextConfig instance.
*
* \param[in] config The OrtEpContextConfig instance to release. May be NULL.
*/
ORT_EXPERIMENTAL_API(28, void, OrtEpApi_ReleaseEpContextConfig, _Frees_ptr_opt_ OrtEpContextConfig* config)
/** \brief Get the application-provided EPContext data read callback.
*
* Returns the OrtReadNamedBufferFunc and opaque state pointer registered via
* OrtApi_SessionOptions_SetEpContextDataReadFunc. If no callback was registered, *read_func and *state are set to
* NULL. The EP is responsible for calling the callback when present and for using its own normal read path when no
* callback is present.
*
* \param[in] config The OrtEpContextConfig from OrtEpApi_SessionOptions_GetEpContextConfig.
* \param[out] read_func The registered read callback, or NULL if none was registered.
* \param[out] state Opaque state pointer passed to read_func, or NULL if none was registered.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_EpContextConfig_GetEpContextDataReadFunc,
_In_ const OrtEpContextConfig* config, _Out_ OrtReadNamedBufferFunc* read_func,
_Out_ void** state)
/** \brief Get the application-provided EPContext data write callback.
*
* Returns the OrtWriteNamedBufferFunc and opaque state pointer registered via
* OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc. If no callback was registered, *write_func and
* *state are set to NULL. The EP is responsible for calling the callback when present and for using its own normal
* write path when no callback is present.
*
* \param[in] config The OrtEpContextConfig from OrtEpApi_SessionOptions_GetEpContextConfig.
* \param[out] write_func The registered write callback, or NULL if none was registered.
* \param[out] state Opaque state pointer passed to write_func, or NULL if none was registered.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*/
ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc,
_In_ const OrtEpContextConfig* config, _Out_ OrtWriteNamedBufferFunc* write_func,
_Out_ void** state)

186
modules/onnxruntime/include/onnxruntime_experimental_cxx_api.h Normal file
View File

@ -0,0 +1,186 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// Experimental C++ API consumer header.
//
// This header provides typed inline accessors in the Ort::Experimental namespace for experimental ORT API functions,
// and any C++ wrapper types associated with the experimental API functions.
//
// It is the C++ companion to onnxruntime_experimental_c_api.h.
//
// IMPORTANT: Experimental functions are NOT part of the stable ABI. They may be added, changed, or removed between
// releases without notice. Anything in this file should be treated as experimental and unstable.
//
// Two accessor flavors are generated for each experimental function:
//
// 1. Get_<NAME>_SinceV<VER>_Fn(api)
// Returns the typed function pointer, or nullptr if the function is not available in this build.
// Use this to check availability at runtime.
//
// 2. Get_<NAME>_SinceV<VER>_FnOrThrow(api)
// Returns the typed function pointer, or throws Ort::Exception (ORT_NOT_IMPLEMENTED) if the function is not
// available in this build.
// Use this when the function is required.
//
// C++ usage (nullable):
// if (auto* fn = Ort::Experimental::Get_OrtApi_ExperimentalApiTest_SinceV28_Fn(api)) {
// Ort::Status status(fn(&result));
// }
//
// C++ usage (throwing):
// auto* fn = Ort::Experimental::Get_OrtApi_ExperimentalApiTest_SinceV28_FnOrThrow(api);
// Ort::Status status(fn(&result));
#pragma once
#include "onnxruntime_experimental_c_api.h"
#include "onnxruntime_cxx_api.h"
namespace Ort {
namespace Experimental {
//
// Nullable typed accessors
//
// For each .inc entry, this generates a typed accessor in Ort::Experimental:
//
// inline OrtExperimental_<NAME>_SinceV<VER>_Fn Get_<NAME>_SinceV<VER>_Fn(const OrtApi* api);
//
// Example: ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_ExperimentalApiTest, ...) produces:
// inline OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn
// Get_OrtApi_ExperimentalApiTest_SinceV28_Fn(const OrtApi* api) {
// return reinterpret_cast<OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn>(
// api->GetExperimentalFunction(kOrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_FnName));
// }
#define ORT_EXPERIMENTAL_API(VER, RET, NAME, ...) \
inline OrtExperimental_##NAME##_SinceV##VER##_Fn Get_##NAME##_SinceV##VER##_Fn( \
const OrtApi* api) { \
return reinterpret_cast<OrtExperimental_##NAME##_SinceV##VER##_Fn>( \
api->GetExperimentalFunction(kOrtExperimental_##NAME##_SinceV##VER##_FnName)); \
}
#include "onnxruntime_experimental_c_api.inc"
#undef ORT_EXPERIMENTAL_API
//
// Throwing typed accessors
//
// For each .inc entry, this generates a throwing accessor in Ort::Experimental:
//
// inline OrtExperimental_<NAME>_SinceV<VER>_Fn Get_<NAME>_SinceV<VER>_FnOrThrow(const OrtApi* api);
//
// It returns the non-null typed function pointer, or throws Ort::Exception (ORT_NOT_IMPLEMENTED) if the function is
// not available in this build.
//
// Example: ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_ExperimentalApiTest, ...) produces:
// inline OrtExperimental_OrtApi_ExperimentalApiTest_SinceV28_Fn
// Get_OrtApi_ExperimentalApiTest_SinceV28_FnOrThrow(const OrtApi* api) {
// auto* fn = Get_OrtApi_ExperimentalApiTest_SinceV28_Fn(api);
// if (fn == nullptr) {
// ORT_CXX_API_THROW(
// "Experimental function OrtApi_ExperimentalApiTest_SinceV28 is not available in this build",
// ORT_NOT_IMPLEMENTED);
// }
// return fn;
// }
#define ORT_EXPERIMENTAL_API(VER, RET, NAME, ...) \
inline OrtExperimental_##NAME##_SinceV##VER##_Fn Get_##NAME##_SinceV##VER##_FnOrThrow( \
const OrtApi* api) { \
auto* fn = Get_##NAME##_SinceV##VER##_Fn(api); \
if (fn == nullptr) { \
ORT_CXX_API_THROW( \
"Experimental function " #NAME "_SinceV" #VER " is not available in this build", \
ORT_NOT_IMPLEMENTED); \
} \
return fn; \
}
#include "onnxruntime_experimental_c_api.inc"
#undef ORT_EXPERIMENTAL_API
//
// Auxiliary types and helpers
//
// C++ wrapper types or helpers go here in the `Ort::Experimental` namespace.
//
// Move-only RAII owner for an OrtEpContextConfig handle, which carries the EPContext read/write callbacks and opaque
// state extracted from an OrtSessionOptions instance. The handle is released via OrtEpApi_ReleaseEpContextConfig when
// the wrapper is destroyed.
//
// Typical EP usage: construct from the session options during CreateEp(), keep the wrapper for the EP's lifetime, and
// query the callbacks via GetReadFunc() / GetWriteFunc().
class EpContextConfig {
public:
explicit EpContextConfig(std::nullptr_t) noexcept {}
explicit EpContextConfig(const SessionOptions& session_options) : EpContextConfig{session_options.GetConst()} {}
// Extracts the EPContext config from `session_options`. Throws Ort::Exception (ORT_NOT_IMPLEMENTED) if the
// experimental functions are not available in this build, or propagates any error from the extraction.
explicit EpContextConfig(ConstSessionOptions session_options) {
const OrtApi* api = &GetApi();
// Ensure the release function is available before creating a handle, so the handle can always be freed.
Get_OrtEpApi_ReleaseEpContextConfig_SinceV28_FnOrThrow(api);
auto* get_config = Get_OrtEpApi_SessionOptions_GetEpContextConfig_SinceV28_FnOrThrow(api);
ThrowOnError(get_config(static_cast<const OrtSessionOptions*>(session_options), &config_));
}
EpContextConfig(EpContextConfig&& other) noexcept : config_{other.config_} { other.config_ = nullptr; }
EpContextConfig& operator=(EpContextConfig&& other) noexcept {
if (this != &other) {
reset();
config_ = other.config_;
other.config_ = nullptr;
}
return *this;
}
EpContextConfig(const EpContextConfig&) = delete;
EpContextConfig& operator=(const EpContextConfig&) = delete;
~EpContextConfig() { reset(); }
OrtEpContextConfig* get() const noexcept { return config_; }
explicit operator bool() const noexcept { return config_ != nullptr; }
// Relinquishes ownership of the handle without releasing it.
OrtEpContextConfig* release() noexcept {
OrtEpContextConfig* released = config_;
config_ = nullptr;
return released;
}
// Releases any owned handle and resets to empty.
void reset() noexcept {
if (config_ != nullptr) {
if (auto* release_fn = Get_OrtEpApi_ReleaseEpContextConfig_SinceV28_Fn(&GetApi())) {
release_fn(config_);
}
config_ = nullptr;
}
}
// Returns the configured read callback and opaque state (both nullptr if none was set). Throws on failure.
void GetReadFunc(OrtReadNamedBufferFunc& read_func, void*& state) const {
auto* get_read_func = Get_OrtEpApi_EpContextConfig_GetEpContextDataReadFunc_SinceV28_FnOrThrow(&GetApi());
ThrowOnError(get_read_func(config_, &read_func, &state));
}
// Returns the configured write callback and opaque state (both nullptr if none was set). Throws on failure.
void GetWriteFunc(OrtWriteNamedBufferFunc& write_func, void*& state) const {
auto* get_write_func = Get_OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc_SinceV28_FnOrThrow(&GetApi());
ThrowOnError(get_write_func(config_, &write_func, &state));
}
private:
OrtEpContextConfig* config_ = nullptr;
};
} // namespace Experimental
} // namespace Ort

537
modules/onnxruntime/include/onnxruntime_float16.h Normal file
View File

@ -0,0 +1,537 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <stdint.h>
#include <cmath>
#include <cstring>
#include <limits>
namespace onnxruntime_float16 {
namespace detail {
enum class endian {
#if defined(_WIN32)
little = 0,
big = 1,
native = little,
#elif defined(__GNUC__) || defined(__clang__)
little = __ORDER_LITTLE_ENDIAN__,
big = __ORDER_BIG_ENDIAN__,
native = __BYTE_ORDER__,
#else
#error onnxruntime_float16::detail::endian is not implemented in this environment.
#endif
};
static_assert(
endian::native == endian::little || endian::native == endian::big,
"Only little-endian or big-endian native byte orders are supported.");
} // namespace detail
/// <summary>
/// Shared implementation between public and internal classes. CRTP pattern.
/// </summary>
template <class Derived>
struct Float16Impl {
protected:
/// <summary>
/// Converts from float to uint16_t float16 representation
/// </summary>
/// <param name="v"></param>
/// <returns></returns>
constexpr static uint16_t ToUint16Impl(float v) noexcept;
/// <summary>
/// Converts float16 to float
/// </summary>
/// <returns>float representation of float16 value</returns>
float ToFloatImpl() const noexcept;
/// <summary>
/// Creates an instance that represents absolute value.
/// </summary>
/// <returns>Absolute value</returns>
uint16_t AbsImpl() const noexcept {
return static_cast<uint16_t>(val & ~kSignMask);
}
/// <summary>
/// Creates a new instance with the sign flipped.
/// </summary>
/// <returns>Flipped sign instance</returns>
uint16_t NegateImpl() const noexcept {
return IsNaN() ? val : static_cast<uint16_t>(val ^ kSignMask);
}
public:
// uint16_t special values
static constexpr uint16_t kSignMask = 0x8000U;
static constexpr uint16_t kBiasedExponentMask = 0x7C00U;
static constexpr uint16_t kPositiveInfinityBits = 0x7C00U;
static constexpr uint16_t kNegativeInfinityBits = 0xFC00U;
static constexpr uint16_t kPositiveQNaNBits = 0x7E00U;
static constexpr uint16_t kNegativeQNaNBits = 0xFE00U;
static constexpr uint16_t kMaxValueBits = 0x7BFFU; // Largest normal number
static constexpr uint16_t kMinValueBits = 0xFBFFU; // Lowest normal number
static constexpr uint16_t kOneBits = 0x3C00U;
static constexpr uint16_t kMinusOneBits = 0xBC00U;
uint16_t val{0};
Float16Impl() = default;
/// <summary>
/// Checks if the value is negative
/// </summary>
/// <returns>true if negative</returns>
bool IsNegative() const noexcept {
return static_cast<int16_t>(val) < 0;
}
/// <summary>
/// Tests if the value is NaN
/// </summary>
/// <returns>true if NaN</returns>
bool IsNaN() const noexcept {
return AbsImpl() > kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value is finite
/// </summary>
/// <returns>true if finite</returns>
bool IsFinite() const noexcept {
return AbsImpl() < kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value represents positive infinity.
/// </summary>
/// <returns>true if positive infinity</returns>
bool IsPositiveInfinity() const noexcept {
return val == kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value represents negative infinity
/// </summary>
/// <returns>true if negative infinity</returns>
bool IsNegativeInfinity() const noexcept {
return val == kNegativeInfinityBits;
}
/// <summary>
/// Tests if the value is either positive or negative infinity.
/// </summary>
/// <returns>True if absolute value is infinity</returns>
bool IsInfinity() const noexcept {
return AbsImpl() == kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value is NaN or zero. Useful for comparisons.
/// </summary>
/// <returns>True if NaN or zero.</returns>
bool IsNaNOrZero() const noexcept {
auto abs = AbsImpl();
return (abs == 0 || abs > kPositiveInfinityBits);
}
/// <summary>
/// Tests if the value is normal (not zero, subnormal, infinite, or NaN).
/// </summary>
/// <returns>True if so</returns>
bool IsNormal() const noexcept {
auto abs = AbsImpl();
return (abs < kPositiveInfinityBits) // is finite
&& (abs != 0) // is not zero
&& ((abs & kBiasedExponentMask) != 0); // is not subnormal (has a non-zero exponent)
}
/// <summary>
/// Tests if the value is subnormal (denormal).
/// </summary>
/// <returns>True if so</returns>
bool IsSubnormal() const noexcept {
auto abs = AbsImpl();
return (abs < kPositiveInfinityBits) // is finite
&& (abs != 0) // is not zero
&& ((abs & kBiasedExponentMask) == 0); // is subnormal (has a zero exponent)
}
/// <summary>
/// Creates an instance that represents absolute value.
/// </summary>
/// <returns>Absolute value</returns>
Derived Abs() const noexcept { return Derived::FromBits(AbsImpl()); }
/// <summary>
/// Creates a new instance with the sign flipped.
/// </summary>
/// <returns>Flipped sign instance</returns>
Derived Negate() const noexcept { return Derived::FromBits(NegateImpl()); }
/// <summary>
/// IEEE defines that positive and negative zero are equal, this gives us a quick equality check
/// for two values by or'ing the private bits together and stripping the sign. They are both zero,
/// and therefore equivalent, if the resulting value is still zero.
/// </summary>
/// <param name="lhs">first value</param>
/// <param name="rhs">second value</param>
/// <returns>True if both arguments represent zero</returns>
static bool AreZero(const Float16Impl& lhs, const Float16Impl& rhs) noexcept {
return static_cast<uint16_t>((lhs.val | rhs.val) & ~kSignMask) == 0;
}
bool operator==(const Float16Impl& rhs) const noexcept {
if (IsNaN() || rhs.IsNaN()) {
// IEEE defines that NaN is not equal to anything, including itself.
return false;
}
return val == rhs.val;
}
bool operator!=(const Float16Impl& rhs) const noexcept { return !(*this == rhs); }
bool operator<(const Float16Impl& rhs) const noexcept {
if (IsNaN() || rhs.IsNaN()) {
// IEEE defines that NaN is unordered with respect to everything, including itself.
return false;
}
const bool left_is_negative = IsNegative();
if (left_is_negative != rhs.IsNegative()) {
// When the signs of left and right differ, we know that left is less than right if it is
// the negative value. The exception to this is if both values are zero, in which case IEEE
// says they should be equal, even if the signs differ.
return left_is_negative && !AreZero(*this, rhs);
}
return (val != rhs.val) && ((val < rhs.val) ^ left_is_negative);
}
};
// The following Float16_t conversions are based on the code from
// Eigen library.
// The conversion routines are Copyright (c) Fabian Giesen, 2016.
// The original license follows:
//
// Copyright (c) Fabian Giesen, 2016
// All rights reserved.
// Redistribution and use in source and binary forms, with or without
// modification, are permitted.
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
namespace detail {
union float32_bits {
unsigned int u;
float f;
};
} // namespace detail
template <class Derived>
inline constexpr uint16_t Float16Impl<Derived>::ToUint16Impl(float v) noexcept {
detail::float32_bits f{};
f.f = v;
constexpr detail::float32_bits f32infty = {255 << 23};
constexpr detail::float32_bits f16max = {(127 + 16) << 23};
constexpr detail::float32_bits denorm_magic = {((127 - 15) + (23 - 10) + 1) << 23};
constexpr unsigned int sign_mask = 0x80000000u;
uint16_t val = static_cast<uint16_t>(0x0u);
unsigned int sign = f.u & sign_mask;
f.u ^= sign;
// NOTE all the integer compares in this function can be safely
// compiled into signed compares since all operands are below
// 0x80000000. Important if you want fast straight SSE2 code
// (since there's no unsigned PCMPGTD).
if (f.u >= f16max.u) { // result is Inf or NaN (all exponent bits set)
val = (f.u > f32infty.u) ? 0x7e00 : 0x7c00; // NaN->qNaN and Inf->Inf
} else { // (De)normalized number or zero
if (f.u < (113 << 23)) { // resulting FP16 is subnormal or zero
// use a magic value to align our 10 mantissa bits at the bottom of
// the float. as long as FP addition is round-to-nearest-even this
// just works.
f.f += denorm_magic.f;
// and one integer subtract of the bias later, we have our final float!
val = static_cast<uint16_t>(f.u - denorm_magic.u);
} else {
unsigned int mant_odd = (f.u >> 13) & 1; // resulting mantissa is odd
// update exponent, rounding bias part 1
// Equivalent to `f.u += ((unsigned int)(15 - 127) << 23) + 0xfff`, but
// without arithmetic overflow.
f.u += 0xc8000fffU;
// rounding bias part 2
f.u += mant_odd;
// take the bits!
val = static_cast<uint16_t>(f.u >> 13);
}
}
val |= static_cast<uint16_t>(sign >> 16);
return val;
}
template <class Derived>
inline float Float16Impl<Derived>::ToFloatImpl() const noexcept {
constexpr detail::float32_bits magic = {113 << 23};
constexpr unsigned int shifted_exp = 0x7c00 << 13; // exponent mask after shift
detail::float32_bits o{};
o.u = (val & 0x7fff) << 13; // exponent/mantissa bits
unsigned int exp = shifted_exp & o.u; // just the exponent
o.u += (127 - 15) << 23; // exponent adjust
// handle exponent special cases
if (exp == shifted_exp) { // Inf/NaN?
o.u += (128 - 16) << 23; // extra exp adjust
} else if (exp == 0) { // Zero/Denormal?
o.u += 1 << 23; // extra exp adjust
o.f -= magic.f; // re-normalize
}
// Attempt to workaround the Internal Compiler Error on ARM64
// for bitwise | operator, including std::bitset
#if (defined _MSC_VER) && (defined _M_ARM || defined _M_ARM64 || defined _M_ARM64EC)
if (IsNegative()) {
return -o.f;
}
#else
// original code:
o.u |= (val & 0x8000U) << 16U; // sign bit
#endif
return o.f;
}
/// Shared implementation between public and internal classes. CRTP pattern.
template <class Derived>
struct BFloat16Impl {
protected:
/// <summary>
/// Converts from float to uint16_t float16 representation
/// </summary>
/// <param name="v"></param>
/// <returns></returns>
static uint16_t ToUint16Impl(float v) noexcept;
/// <summary>
/// Converts bfloat16 to float
/// </summary>
/// <returns>float representation of bfloat16 value</returns>
float ToFloatImpl() const noexcept;
/// <summary>
/// Creates an instance that represents absolute value.
/// </summary>
/// <returns>Absolute value</returns>
uint16_t AbsImpl() const noexcept {
return static_cast<uint16_t>(val & ~kSignMask);
}
/// <summary>
/// Creates a new instance with the sign flipped.
/// </summary>
/// <returns>Flipped sign instance</returns>
uint16_t NegateImpl() const noexcept {
return IsNaN() ? val : static_cast<uint16_t>(val ^ kSignMask);
}
public:
// uint16_t special values
static constexpr uint16_t kSignMask = 0x8000U;
static constexpr uint16_t kBiasedExponentMask = 0x7F80U;
static constexpr uint16_t kPositiveInfinityBits = 0x7F80U;
static constexpr uint16_t kNegativeInfinityBits = 0xFF80U;
static constexpr uint16_t kPositiveQNaNBits = 0x7FC1U;
static constexpr uint16_t kNegativeQNaNBits = 0xFFC1U;
static constexpr uint16_t kMaxValueBits = 0x7F7FU;
static constexpr uint16_t kMinValueBits = 0xFF7FU;
static constexpr uint16_t kRoundToNearest = 0x7FFFU;
static constexpr uint16_t kOneBits = 0x3F80U;
static constexpr uint16_t kMinusOneBits = 0xBF80U;
uint16_t val{0};
BFloat16Impl() = default;
/// <summary>
/// Checks if the value is negative
/// </summary>
/// <returns>true if negative</returns>
bool IsNegative() const noexcept {
return static_cast<int16_t>(val) < 0;
}
/// <summary>
/// Tests if the value is NaN
/// </summary>
/// <returns>true if NaN</returns>
bool IsNaN() const noexcept {
return AbsImpl() > kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value is finite
/// </summary>
/// <returns>true if finite</returns>
bool IsFinite() const noexcept {
return AbsImpl() < kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value represents positive infinity.
/// </summary>
/// <returns>true if positive infinity</returns>
bool IsPositiveInfinity() const noexcept {
return val == kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value represents negative infinity
/// </summary>
/// <returns>true if negative infinity</returns>
bool IsNegativeInfinity() const noexcept {
return val == kNegativeInfinityBits;
}
/// <summary>
/// Tests if the value is either positive or negative infinity.
/// </summary>
/// <returns>True if absolute value is infinity</returns>
bool IsInfinity() const noexcept {
return AbsImpl() == kPositiveInfinityBits;
}
/// <summary>
/// Tests if the value is NaN or zero. Useful for comparisons.
/// </summary>
/// <returns>True if NaN or zero.</returns>
bool IsNaNOrZero() const noexcept {
auto abs = AbsImpl();
return (abs == 0 || abs > kPositiveInfinityBits);
}
/// <summary>
/// Tests if the value is normal (not zero, subnormal, infinite, or NaN).
/// </summary>
/// <returns>True if so</returns>
bool IsNormal() const noexcept {
auto abs = AbsImpl();
return (abs < kPositiveInfinityBits) // is finite
&& (abs != 0) // is not zero
&& ((abs & kBiasedExponentMask) != 0); // is not subnormal (has a non-zero exponent)
}
/// <summary>
/// Tests if the value is subnormal (denormal).
/// </summary>
/// <returns>True if so</returns>
bool IsSubnormal() const noexcept {
auto abs = AbsImpl();
return (abs < kPositiveInfinityBits) // is finite
&& (abs != 0) // is not zero
&& ((abs & kBiasedExponentMask) == 0); // is subnormal (has a zero exponent)
}
/// <summary>
/// Creates an instance that represents absolute value.
/// </summary>
/// <returns>Absolute value</returns>
Derived Abs() const noexcept { return Derived::FromBits(AbsImpl()); }
/// <summary>
/// Creates a new instance with the sign flipped.
/// </summary>
/// <returns>Flipped sign instance</returns>
Derived Negate() const noexcept { return Derived::FromBits(NegateImpl()); }
/// <summary>
/// IEEE defines that positive and negative zero are equal, this gives us a quick equality check
/// for two values by or'ing the private bits together and stripping the sign. They are both zero,
/// and therefore equivalent, if the resulting value is still zero.
/// </summary>
/// <param name="lhs">first value</param>
/// <param name="rhs">second value</param>
/// <returns>True if both arguments represent zero</returns>
static bool AreZero(const BFloat16Impl& lhs, const BFloat16Impl& rhs) noexcept {
// IEEE defines that positive and negative zero are equal, this gives us a quick equality check
// for two values by or'ing the private bits together and stripping the sign. They are both zero,
// and therefore equivalent, if the resulting value is still zero.
return static_cast<uint16_t>((lhs.val | rhs.val) & ~kSignMask) == 0;
}
};
template <class Derived>
inline uint16_t BFloat16Impl<Derived>::ToUint16Impl(float v) noexcept {
uint16_t result;
if (std::isnan(v)) {
result = kPositiveQNaNBits;
} else {
auto get_msb_half = [](float fl) {
uint16_t result;
#ifdef __cpp_if_constexpr
if constexpr (detail::endian::native == detail::endian::little) {
#else
if (detail::endian::native == detail::endian::little) {
#endif
std::memcpy(&result, reinterpret_cast<char*>(&fl) + sizeof(uint16_t), sizeof(uint16_t));
} else {
std::memcpy(&result, &fl, sizeof(uint16_t));
}
return result;
};
uint16_t upper_bits = get_msb_half(v);
union {
uint32_t U32;
float F32;
};
F32 = v;
U32 += (upper_bits & 1) + kRoundToNearest;
result = get_msb_half(F32);
}
return result;
}
template <class Derived>
inline float BFloat16Impl<Derived>::ToFloatImpl() const noexcept {
if (IsNaN()) {
return std::numeric_limits<float>::quiet_NaN();
}
float result;
char* const first = reinterpret_cast<char*>(&result);
char* const second = first + sizeof(uint16_t);
#ifdef __cpp_if_constexpr
if constexpr (detail::endian::native == detail::endian::little) {
#else
if (detail::endian::native == detail::endian::little) {
#endif
std::memset(first, 0, sizeof(uint16_t));
std::memcpy(second, &val, sizeof(uint16_t));
} else {
std::memcpy(first, &val, sizeof(uint16_t));
std::memset(second, 0, sizeof(uint16_t));
}
return result;
}
} // namespace onnxruntime_float16

1117
modules/onnxruntime/include/onnxruntime_lite_custom_op.h Normal file
View File

File diff suppressed because it is too large Load Diff

54
modules/onnxruntime/include/onnxruntime_run_options_config_keys.h Normal file
View File

@ -0,0 +1,54 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
/*
* This file defines RunOptions Config Keys and format of the Config Values.
*
* The Naming Convention for a RunOptions Config Key,
* "[Area][.[SubArea1].[SubArea2]...].[Keyname]"
* Such as "ep.cuda.use_arena"
* The Config Key cannot be empty
* The maximum length of the Config Key is 128
*
* The string format of a RunOptions Config Value is defined individually for each Config.
* The maximum length of the Config Value is 1024
*/
// Key for enabling shrinkages of user listed device memory arenas.
// Expects a list of semi-colon separated key value pairs separated by colon in the following format:
// "device_0:device_id_0;device_1:device_id_1"
// No white-spaces allowed in the provided list string.
// Currently, the only supported devices are : "cpu", "gpu" (case sensitive).
// If "cpu" is included in the list, DisableCpuMemArena() API must not be called (i.e.) arena for cpu should be enabled.
// Example usage: "cpu:0;gpu:0" (or) "gpu:0"
// By default, the value for this key is empty (i.e.) no memory arenas are shrunk
static const char* const kOrtRunOptionsConfigEnableMemoryArenaShrinkage = "memory.enable_memory_arena_shrinkage";
// Set to '1' to not synchronize execution providers with CPU at the end of session run.
// Per default it will be set to '0'
// Taking CUDA EP as an example, it omit triggering cudaStreamSynchronize on the compute stream.
static const char* const kOrtRunOptionsConfigDisableSynchronizeExecutionProviders = "disable_synchronize_execution_providers";
// Set HTP performance mode for QNN HTP backend before session run.
// options for HTP performance mode: "burst", "balanced", "default", "high_performance",
// "high_power_saver", "low_balanced", "extreme_power_saver", "low_power_saver", "power_saver",
// "sustained_high_performance". Default to "default".
static const char* const kOrtRunOptionsConfigQnnPerfMode = "qnn.htp_perf_mode";
// Set HTP performance mode for QNN HTP backend post session run.
static const char* const kOrtRunOptionsConfigQnnPerfModePostRun = "qnn.htp_perf_mode_post_run";
// Set RPC control latency for QNN HTP backend
static const char* const kOrtRunOptionsConfigQnnRpcControlLatency = "qnn.rpc_control_latency";
// Set QNN Lora Config File for apply Lora in QNN context binary
static const char* const kOrtRunOptionsConfigQnnLoraConfig = "qnn.lora_config";
// Set graph annotation id for CUDA EP. Use with enable_cuda_graph=true.
// The value should be an integer. If the value is not set, the default value is 0 and
// ORT session only captures one cuda graph before another capture is requested.
// If the value is set to -1, cuda graph capture/replay is disabled in that run.
// User are not expected to set the value to 0 as it is reserved for internal use.
static const char* const kOrtRunOptionsConfigCudaGraphAnnotation = "gpu_graph_id";

800
modules/onnxruntime/include/onnxruntime_session_options_config_keys.h Normal file
View File

@ -0,0 +1,800 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
/*
* This file defines SessionOptions Config Keys and format of the Config Values.
*
* The Naming Convention for a SessionOptions Config Key,
* "[Area][.[SubArea1].[SubArea2]...].[Keyname]"
* Such as "ep.cuda.use_arena"
* The Config Key cannot be empty
* The maximum length of the Config Key is 1024
*
* The string format of a SessionOptions Config Value is defined individually for each Config.
* The maximum length of the Config Value is 8192
*/
// Key for disable PrePacking,
// If the config value is set to "1" then the prepacking is disabled, otherwise prepacking is enabled (default value)
static const char* const kOrtSessionOptionsConfigDisablePrepacking = "session.disable_prepacking";
// A value of "1" means allocators registered in the env will be used. "0" means the allocators created in the session
// will be used. Use this to override the usage of env allocators on a per session level.
static const char* const kOrtSessionOptionsConfigUseEnvAllocators = "session.use_env_allocators";
// Set to 'ORT' (case sensitive) to load an ORT format model.
// If unset, model type will default to ONNX unless inferred from filename ('.ort' == ORT format) or bytes to be ORT
static const char* const kOrtSessionOptionsConfigLoadModelFormat = "session.load_model_format";
// Set to 'ORT' (case sensitive) to save optimized model in ORT format when SessionOptions.optimized_model_path is set.
// If unset, format will default to ONNX unless optimized_model_filepath ends in '.ort'.
static const char* const kOrtSessionOptionsConfigSaveModelFormat = "session.save_model_format";
// If a value is "1", flush-to-zero and denormal-as-zero are applied. The default is "0".
// When multiple sessions are created, a main thread doesn't override changes from succeeding session options,
// but threads in session thread pools follow option changes.
// When ORT runs with OpenMP, the same rule is applied, i.e. the first session option to flush-to-zero and
// denormal-as-zero is only applied to global OpenMP thread pool, which doesn't support per-session thread pool.
// Note that an alternative way not using this option at runtime is to train and export a model without denormals
// and that's recommended because turning this option on may hurt model accuracy.
static const char* const kOrtSessionOptionsConfigSetDenormalAsZero = "session.set_denormal_as_zero";
// It controls to run quantization model in QDQ (QuantizelinearDeQuantizelinear) format or not.
// "0": enable. ORT does fusion logic for QDQ format.
// "1": disable. ORT doesn't do fusion logic for QDQ format.
// Its default value is "0" unless the DirectML execution provider is registered, in which case it defaults to "1".
static const char* const kOrtSessionOptionsDisableQuantQDQ = "session.disable_quant_qdq";
// This controls whether to prevent constant folding from folding DequantizeLinear nodes:
// "0": (default) DequantizeLinear constant folding is determined solely by session.disable_quant_qdq.
// "1": DequantizeLinear nodes are never individually constant folded.
// When session.disable_quant_qdq is "0" (default), DequantizeLinear nodes are already protected from
// constant folding to preserve QDQ node units for downstream QDQ fusion optimizers.
// When session.disable_quant_qdq is "1", then DequantizeLinear nodes are normally allowed to be
// constant folded, but setting this option to "1" still preserves DequantizeLinear nodes.
// This is useful for execution providers like WebNN that disable QDQ fusion, but which
// still need the original DQ/Q nodes to be preserved for their own quantization handling.
static const char* const kOrtSessionOptionsDisableQDQConstantFolding =
"session.disable_qdq_constant_folding";
// It controls whether to enable Double QDQ remover and Identical Children Consolidation
// "0": not to disable. ORT does remove the middle 2 Nodes from a Q->(QD->Q)->QD pairs
// "1": disable. ORT doesn't remove the middle 2 Nodes from a Q->(QD->Q)->QD pairs
// Its default value is "0"
static const char* const kOrtSessionOptionsDisableDoubleQDQRemover = "session.disable_double_qdq_remover";
// If set to "1", enables the removal of QuantizeLinear/DequantizeLinear node pairs once all QDQ handling has been
// completed. e.g. If after all QDQ handling has completed and we have -> FloatOp -> Q -> DQ -> FloatOp -> the
// Q -> DQ could potentially be removed. This will provide a performance benefit by avoiding going from float to
// 8-bit and back to float, but could impact accuracy. The impact on accuracy will be model specific and depend on
// other factors like whether the model was created using Quantization Aware Training or Post Training Quantization.
// As such, it's best to test to determine if enabling this works well for your scenario.
// The default value is "0"
// Available since version 1.11.
static const char* const kOrtSessionOptionsEnableQuantQDQCleanup = "session.enable_quant_qdq_cleanup";
// Enable or disable gelu approximation in graph optimization. "0": disable; "1": enable. The default is "0".
// GeluApproximation has side effects which may change the inference results. It is disabled by default due to this.
static const char* const kOrtSessionOptionsEnableGeluApproximation = "optimization.enable_gelu_approximation";
// Enable or disable Cast chain elimination in graph optimization. "0": disable; "1": enable. The default is "0".
// CastElimination with chain elimination has side effects which may change the inference results. It is disabled by default due to this.
static const char* const kOrtSessionOptionsEnableCastChainElimination = "optimization.enable_cast_chain_elimination";
// This setting controls whether to enable AheadOfTime function inlining.
// AOT function inlining examines the graph and attempts to inline as many locally defined functions in the model
// as possible with the help of enabled execution providers.
// This can reduce the number of function calls and improve performance because it is done before
// Level1 optimizers and constant folding. However, under some circumstances, when the EPs are not available,
// one can disable the AOT inlining, produce an optimized model and postpone AOT until run time.
// "0": enable; "1": disable.
// Its default value is "0".
static const char* const kOrtSessionOptionsDisableAheadOfTimeFunctionInlining = "session.disable_aot_function_inlining";
#ifdef ENABLE_TRAINING
// Specifies a path of the file containing a list of memory optimization configurations.
// The value should be a string indicating the file path of the config file.
// The content of the config file is a JSON struct like this:
// [
// "Gelu+Cast+:1:0",
// "Dropout+:1:1"
// ]
// Taking the example of "Gelu+Cast+:1:0",
// > "Gelu+Cast+" is the subgraph string, a valid "subgraph string" should be one subgraph representation
// output by ORT graph transformations.
// > "1" is "optimization strategy", valid values: 0 - disabled, 1 - recompute.
// > "0" is "number of subgraph to apply" which is used to control how many subgraphs to apply optimization,
// to avoid "oversaving" the memory.
static const char* const kOrtSessionOptionsMemoryOptimizerApplyConfig = "optimization.memory_optimizer_config";
// Specifies the config for detecting subgraphs for memory footprint reduction.
// The value should be a string contains int separated using commas. The default value is "0:0".
static const char* const kOrtSessionOptionsMemoryOptimizerProbeConfig = "optimization.enable_memory_probe_recompute_config";
#endif
// This setting if set should contain a comma separated list of optimizers names that should be disabled.
// Optimizers may take time to execute and affect model loading time. If you feel that a specific optimizer
// does not provider runtime benefits, but affects your model loading time you may disable it using this config
// entry. This option is not enabled in ORT_MINIMAL_BUILD build.
// A list of optimizes is available in onnxruntime/core/optimizer/graph_transformer_utils.cc
//
// Default is an empty string which means no optimizers are disabled.
static const char* const kOrtSessionOptionsDisableSpecifiedOptimizers = "optimization.disable_specified_optimizers";
// Maximum total output size in bytes that the constant folding optimizer is allowed to produce per node.
// Prevents malicious models from causing excessive memory allocation during optimization.
// If the estimated or actual output size of a constant-foldable node exceeds this limit, the node will
// not be constant folded and will instead be executed at runtime.
//
// Option values:
// - A positive integer (as string): Maximum allowed output size in bytes per constant-folded node.
// Default is "1073741824" (1 GB).
// - "0": Disable the size limit (not recommended for untrusted models).
static const char* const kOrtSessionOptionsConstantFoldingMaxOutputSizeInBytes =
"optimization.constant_folding_max_output_size_in_bytes";
// It controls whether to run graph optimizations in loop or not.
//
// "0": disable. Graph Optimization Loop is disabled.
// ```
// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
// ^ |
// | "No Loop" |
// | |
// X xxxxxxxxxxx X
// ```
// "1": enable. Graph Optimization Loop is enabled, such that, if optimizations at Level 4 are applied then
// the loop will check for any other valid optimization that can happen.
// ```
// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
// ^ |
// | "Loop only depending on Level 4" |
// | |
// ---------------------------------------------------
// ```
// "2": enable. Graph Optimization Loop is enabled, such that, if optimizations at Level 2 or above are applied then
// The loop will check for any other valid optimization that can happen.
// ```
// Level 2 --> Level 3 --> InsertCastTransforms --> Level 4
// ^ |
// | "Loop" |
// | |
// ---------------------------------------------------
// ```
// Default value is set to "1".
static const char* const kOrtSessionOptionsGraphOptimizationsLoopLevel = "session.graph_optimizations_loop_level";
// Enable or disable using device allocator for allocating initialized tensor memory. "1": enable; "0": disable. The default is "0".
// Using device allocators means the memory allocation is made using malloc/new.
static const char* const kOrtSessionOptionsUseDeviceAllocatorForInitializers = "session.use_device_allocator_for_initializers";
// Configure whether to allow the inter_op/intra_op threads spinning a number of times before blocking
// "0": thread will block if found no job to run
// "1": thread will spin a number of times before blocking
// The default is "0" when ORT is built with "ORT_CLIENT_PACKAGE_BUILD" and "1" otherwise.
// Thread spinning is disabled by default for client/on-device workloads to reduce cpu utilization and improve power efficiency.
static const char* const kOrtSessionOptionsConfigAllowInterOpSpinning = "session.inter_op.allow_spinning";
static const char* const kOrtSessionOptionsConfigAllowIntraOpSpinning = "session.intra_op.allow_spinning";
// Configure the duration in microseconds that threads spin waiting for work before blocking.
// This setting is subordinate to the allow_spinning flags (session.intra_op.allow_spinning /
// session.inter_op.allow_spinning). When allow_spinning is "0", spinning is disabled and
// the spin duration is forced to 0 regardless of this setting.
// By default (when this option is not set), the thread pool uses an iteration-count-based spin loop
// whose wall-clock duration varies by CPU architecture and pause instruction latency. This provides
// the best throughput but may result in high CPU utilization.
// Setting a positive value switches to calibrated iteration-based spinning that targets
// the specified duration. The actual spin time is a best-effort approximation based on a
// one-time measurement of the pause instruction latency; it may vary with CPU frequency
// changes. Recommended for power-sensitive or client/on-device workloads.
// Common values: 500-2000 (0.5-2ms).
// Setting to "0" with spinning enabled effectively disables spinning (equivalent to allow_spinning = false).
static const char* const kOrtSessionOptionsConfigIntraOpSpinDurationUs = "session.intra_op.spin_duration_us";
static const char* const kOrtSessionOptionsConfigInterOpSpinDurationUs = "session.inter_op.spin_duration_us";
// Configure the maximum exponential-backoff cap for the thread pool spin loop.
// When > 1, each idle spin iteration emits a growing number of SpinPause() calls
// (1, 2, 4, ..., up to this cap), which reduces the density of pause instructions
// during the spin window and lowers CPU/power usage compared to emitting one
// SpinPause() per iteration. The total wall-clock spin duration targeted by
// session.{intra,inter}_op.spin_duration_us is preserved by scaling the iteration
// count against the backoff cap.
// "1" (default) = no backoff, one SpinPause() per iteration (original behavior).
// ">= 2" = enable exponential backoff capped at this value. Typical
// values: 4 (hybrid/E-core friendly) or 8 (desktop/server).
// Values above 64 are clamped to 64.
// This setting is subordinate to allow_spinning and spin_duration_us: when
// spinning is disabled or spin_duration_us forces zero iterations, this value
// has no effect.
static const char* const kOrtSessionOptionsConfigIntraOpSpinBackoffMax = "session.intra_op.spin_backoff_max";
static const char* const kOrtSessionOptionsConfigInterOpSpinBackoffMax = "session.inter_op.spin_backoff_max";
// Key for using model bytes directly for ORT format
// If a session is created using an input byte array contains the ORT format model data,
// By default we will copy the model bytes at the time of session creation to ensure the model bytes
// buffer is valid.
// Setting this option to "1" will disable copy the model bytes, and use the model bytes directly. The caller
// has to guarantee that the model bytes are valid until the ORT session using the model bytes is destroyed.
static const char* const kOrtSessionOptionsConfigUseORTModelBytesDirectly = "session.use_ort_model_bytes_directly";
/// <summary>
/// Key for using the ORT format model flatbuffer bytes directly for initializers.
/// This avoids copying the bytes and reduces peak memory usage during model loading and initialization.
/// Requires `session.use_ort_model_bytes_directly` or `session.use_memory_mapped_ort_model` to be true.
/// If set, the flatbuffer bytes provided when creating the InferenceSession MUST remain valid for the entire
/// duration of the InferenceSession.
/// </summary>
static const char* const kOrtSessionOptionsConfigUseORTModelBytesForInitializers =
"session.use_ort_model_bytes_for_initializers";
/// <summary>
/// Key for using memory-mapped I/O to load ORT format model files.
/// When set to "1" and the session is created from a file path, ORT will use memory-mapped I/O
/// to load the .ort model file instead of reading it into a heap-allocated buffer.
/// Usage with session.use_ort_model_bytes_for_initializers will ensure Tensors point directly to the mapped bytes,
/// although the mapping must remain valid and model weights will be immutable.
/// The model load will fail if the mapping fails; fallbacks should be caller-handled.
/// </summary>
static const char* const kOrtSessionOptionsConfigUseMemoryMappedOrtModel = "session.use_memory_mapped_ort_model";
// This should only be specified when exporting an ORT format model for use on a different platform.
// If the ORT format model will be used on ARM platforms set to "1". For other platforms set to "0"
// Available since version 1.11.
static const char* const kOrtSessionOptionsQDQIsInt8Allowed = "session.qdqisint8allowed";
// x64 SSE4.1/AVX2/AVX512(with no VNNI) has overflow problem with quantizied matrix multiplication with U8S8.
// To avoid this we need to use slower U8U8 matrix multiplication instead. This option, if
// turned on, use slower U8U8 matrix multiplications. Only effective with AVX2 or AVX512
// platforms.
static const char* const kOrtSessionOptionsAvx2PrecisionMode = "session.x64quantprecision";
// Specifies how minimal build graph optimizations are handled in a full build.
// These optimizations are at the extended level or higher.
// Possible values and their effects are:
// "save": Save runtime optimizations when saving an ORT format model.
// "apply": Only apply optimizations available in a minimal build.
// ""/<unspecified>: Apply optimizations available in a full build.
// Available since version 1.11.
static const char* const kOrtSessionOptionsConfigMinimalBuildOptimizations =
"optimization.minimal_build_optimizations";
/// <summary>
/// Set to "1" to enable replay of saved runtime optimizations from an ORT format model.
/// Saved runtime optimizations contain graph rewrite instructions and must only be enabled for trusted models.
/// The default is "0".
/// </summary>
static const char* const kOrtSessionOptionsConfigEnableSavedRuntimeOptimizations =
"session.enable_saved_runtime_optimizations";
// Note: The options specific to an EP should be specified prior to appending that EP to the session options object in
// order for them to take effect.
// Specifies a list of stop op types. Nodes of a type in the stop op types and nodes downstream from them will not be
// run by the NNAPI EP.
// The value should be a ","-delimited list of op types. For example, "Add,Sub".
// If not specified, the default set of stop ops is used. To specify an empty stop ops types list and disable stop op
// exclusion, set the value to "".
static const char* const kOrtSessionOptionsConfigNnapiEpPartitioningStopOps = "ep.nnapi.partitioning_stop_ops";
// Enabling dynamic block-sizing for multithreading.
// With a positive value, thread pool will split a task of N iterations to blocks of size starting from:
// N / (num_of_threads * dynamic_block_base)
// As execution progresses, the size will decrease according to the diminishing residual of N,
// meaning the task will be distributed in smaller granularity for better parallelism.
// For some models, it helps to reduce the variance of E2E inference latency and boost performance.
// The feature will not function by default, specify any positive integer, e.g. "4", to enable it.
// Available since version 1.11.
static const char* const kOrtSessionOptionsConfigDynamicBlockBase = "session.dynamic_block_base";
// This option allows to decrease CPU usage between infrequent
// requests and forces any TP threads spinning stop immediately when the last of
// concurrent Run() call returns.
// Spinning is restarted on the next Run() call.
// Applies only to internal thread-pools
static const char* const kOrtSessionOptionsConfigForceSpinningStop = "session.force_spinning_stop";
// "1": all inconsistencies encountered during shape and type inference
// will result in failures.
// "0": in some cases warnings will be logged but processing will continue. The default.
// May be useful to expose bugs in models.
static const char* const kOrtSessionOptionsConfigStrictShapeTypeInference = "session.strict_shape_type_inference";
// "1": every model using a more recent opset than the latest released one will fail
// "0": the model may or may not work if onnxruntime cannot find an implementation, this option
// is used for development purpose.
static const char* const kOrtSessionOptionsConfigStrictAllowReleasedOpsetsOnly = "session.allow_released_opsets_only";
// The file saves configuration for partitioning node among logic streams
static const char* const kNodePartitionConfigFile = "session.node_partition_config_file";
// This Option allows setting affinities for intra op threads.
// Affinity string follows format:
// logical_processor_id,logical_processor_id;logical_processor_id,logical_processor_id
// Semicolon isolates configurations among threads, while comma split processors where ith thread expected to attach to.
// e.g.1,2,3;4,5
// specifies affinities for two threads, with the 1st thread attach to the 1st, 2nd, and 3rd processor, and 2nd thread to the 4th and 5th.
// To ease the configuration, an "interval" is also allowed:
// e.g. 1-8;8-16;17-24
// orders that the 1st thread runs on first eight processors, 2nd thread runs on next eight processors, and so forth.
// Note:
// 1. Once set, the number of thread affinities must equal to intra_op_num_threads - 1, since ort does not set affinity on the main thread which
// is started and managed by the calling app;
// 2. For windows, ort will infer the group id from a logical processor id, for example, assuming there are two groups with each has 64 logical processors,
// an id of 64 will be inferred as the last processor of the 1st group, while 65 will be interpreted as the 1st processor of the second group.
// Hence 64-65 is an invalid configuration, because a windows thread cannot be attached to processors across group boundary.
static const char* const kOrtSessionOptionsConfigIntraOpThreadAffinities = "session.intra_op_thread_affinities";
// This option will dump out the model to assist debugging any issues with layout transformation,
// and is primarily intended for developer usage. It is only relevant if an execution provider that requests
// NHWC layout is enabled such as NNAPI, XNNPACK or QNN.
//
// Default is off. Set to "1" to enable.
//
// If modified by layout transformation the model will be dumped after these steps:
// 1) insertion of the layout transformation Transpose nodes
// 2) after those are optimized using the transpose optimizer,
// 3) after the L1 transformers are applied to the updated graph.
// The model will be saved to filename post_layout_transform_step_<step_number>.onnx.
static const char* const kDebugLayoutTransformation = "session.debug_layout_transformation";
// Graph nodes that are not supported by the execution providers (EPs) explicitly added to the session are
// assigned (i.e., "fallback") to the CPU EP by default.
//
// This option allows the user to disable the fallback of unsupported graph nodes to the CPU EP.
// If this option is set to "1", session creation will fail if the execution providers other than the CPU EP cannot
// fully support all of the nodes in the graph.
//
// It is invalid to set this option and explicitly add the CPU EP to the session. In this case, session creation
// will also fail with an error.
//
// Option values:
// - "0": CPU EP fallback is not disabled. [DEFAULT]
// - "1": CPU EP fallback is disabled.
static const char* const kOrtSessionOptionsDisableCPUEPFallback = "session.disable_cpu_ep_fallback";
// Use this config when serializing a large model after optimization to specify an external initializers file
static const char* const kOrtSessionOptionsOptimizedModelExternalInitializersFileName =
"session.optimized_model_external_initializers_file_name";
// Use this config to control the minimum size of the initializer when externalizing it during serialization
static const char* const kOrtSessionOptionsOptimizedModelExternalInitializersMinSizeInBytes =
"session.optimized_model_external_initializers_min_size_in_bytes";
// Specifies the folder path used to resolve a model's external initializer data files.
// When set, external initializers are loaded from this folder instead of the model's own
// directory, overriding the model directory. This applies whether the model is loaded from a
// file path or from a memory buffer/stream. All external data files must be in the same folder.
// Typical uses include loading models with external data from memory, sharing a weights file
// across models, and weightless/cache models whose weights live outside the model directory.
// For EPContext workflows, also set kOrtSessionOptionEpContextFilePath so the EPContext
// model location remains available for resolving an external EP context binary.
static const char* const kOrtSessionOptionsModelExternalInitializersFileFolderPath =
"session.model_external_initializers_file_folder_path";
// Use this config when saving pre-packed constant initializers to an external data file.
// This allows you to memory map pre-packed initializers on model load and leave it to
// to the OS the amount of memory consumed by the pre-packed initializers. Otherwise,
// pre-packed data resides on the heap.
//
// - "0": Default is not save pre-packed initializers to a data file.
// - "1": Save pre-packed constant initializers to an external data file.
// Sample usage: sess_options.add_session_config_entry(kOrtSessionOptionsSavePrePackedConstantInitializers, "1")
static const char* const kOrtSessionOptionsSavePrePackedConstantInitializers =
"session.save_external_prepacked_constant_initializers";
// Use this config when you want to collect memory stats for each node in the graph.
// The file format is a CSV file with the following columns:
// The file will be created if it does not exist, and will be overwritten if it does.
//
// The content of the file can be used to estimate memory requirements at run time including
// the temporary allocations. This operation is preferably done on a CPU device, as the model may exceed
// device memory limits in constrained environments. When enabling this option, it is important to disable
// memory patterns, as they tend to allocate large blocks to avoid fragmentation and accommodate needs of multiple
// kernels. Memory patterns may make it difficult to allocate on a device with limited memory.
//
// The collected stats then can be used to partition the graph among the devices in a way that only the
// required memory is allocated on each device.
//
// node_name, initializers_memory, dynamic_outputs_sizes, temp_allocations_size
//
// - "full path to file": there is not a default for this option. If the file can not be opened for writing, an error will be returned.
static const char* const kOrtSessionOptionsCollectNodeMemoryStatsToFile = "session.collect_node_memory_stats_to_file";
/// This is a composite CSV setting formatted as "memory limit in kb,file name for collected stats"
/// "limit > 0": enables Capacity Aware Partitioning for Cuda EP. `limit` is optional and when absent
/// the provider may attempt to figure out the memory available automatically.
/// The setting with no pre-recorded stats is expected to look like: "limit > 0,".
/// In this case, the EP will calculate memory using the initializers referenced by the node.
/// This enables an ad-hoc and flexible scenarios with no pre-recorded stats, but may be less accurate.
/// The setting with no limit is expected to look like: ",file name for collected stats"
/// Finally a setting with both limit and pre-recorded stats absent can contain a single comma: ",".
/// The EP will attempt to place nodes on device (currently only CUDA is supported) :
/// this file is expected to be found at the same folder with the model. The file contains
/// pre-recorded stats collected when running with kOrtSessionOptionsCollectNodeMemoryStatsToFile enforce (see above)
static const char* const kOrtSessionOptionsResourceCudaPartitioningSettings =
"session.resource_cuda_partitioning_settings";
/// <summary>
/// This is a setting that contains string annotations or annotation prefixes to be matched
/// against individual nodes metadata entry 'layer_ann' to guide layer assignment during partitioning.
/// The value is a semicolon separated list of strings or string prefixes per device.
/// Format: device1(annotation1, annotation2, ...); device2(annotation1, =annotation3, ...);...
/// Where:
/// - device1, device2, ... are the recognized device names to be matched against EPs configured in
/// the given session.
/// - annotation1, annotation2, ... are annotation prefixes to be matched against node annotations. Any
/// node annotation that starts with one of these prefixes will be matched.
/// - =annotation3 indicates an exact match for annotation3. Only node annotations that are exactly
/// equal to 'annotation3' will be matched.
/// TODO: add a list of recognized devices here.
/// </summary>
static const char* const kOrtSessionOptionsLayerAssignmentSettings = "session.layer_assignment_settings";
/// <summary>
/// Name-based layer assignment. Uses the same device(pattern1, pattern2, ...); ... grammar
/// as kOrtSessionOptionsLayerAssignmentSettings but performs SUBSTRING matching against
/// Node::Name() instead of prefix/exact matching against node metadata annotations.
/// The '=' prefix (exact match) from the annotation-based grammar is rejected with an error
/// — all patterns are treated as substrings.
/// Longest matching pattern wins when multiple patterns match the same node name.
/// No subgraph inheritance is applied — each node is matched independently by its name.
///
/// MUTUALLY EXCLUSIVE with kOrtSessionOptionsLayerAssignmentSettings. Setting both returns
/// INVALID_ARGUMENT. Use annotation-based matching for models with explicit layer annotations,
/// or name-based matching for models with structured node names (HuggingFace, PyTorch exports).
/// </summary>
static const char* const kOrtSessionOptionsNameBasedLayerAssignment = "session.name_based_layer_assignment";
/// <summary>
/// Provides input shape overrides for workspace estimation in dynamic-shape models.
/// When set, the framework propagates these shapes for workspace pre-computation
/// (Level-1 budget estimation and Level-2 DeclareWorkspaceRequirements).
///
/// Format: "input_name:[d0,d1,...];input_name2:[d0,d1,...]"
/// Example: "input_ids:[8,4096];attention_mask:[8,4096]"
///
/// Each input_name must match a model graph input. Dimensions must be positive integers.
/// Unknown/symbolic dimensions in the model will be replaced by the corresponding override
/// value and propagated through a separate shape-inference graph for estimation purposes only.
/// Although an override may describe a maximum input shape, downstream inferred shapes are
/// estimation hints, not guaranteed upper bounds: operator shape transformations are not
/// necessarily monotonic. Runtime shapes are not constrained by this setting, and consumers
/// must retain runtime bounds checks and allocation fallbacks.
/// </summary>
static const char* const kOrtSessionOptionsMaxShapeOverride = "session.max_shape_override";
// Enable EP context feature to dump the partitioned graph which includes the EP context into Onnx file.
// The dumped Onnx model with EP context can be used for future inference to avoid the EP graph partitioning/compile overhead.
// "0": disable. (default)
// "1": enable.
static const char* const kOrtSessionOptionEpContextEnable = "ep.context_enable";
// Specify the file path for the ONNX model containing EP context.
// For EP context generation, defaults to original_file_name_ctx.onnx if not specified.
// During inference, EPs use this path to resolve an external EP context binary whose
// relative path is stored in an EPContext node's ep_cache_context attribute.
// To resolve an external EP context binary, set this option when the model path is
// unavailable or when kOrtSessionOptionsModelExternalInitializersFileFolderPath overrides
// it with a different directory. Specifying both paths is recommended for EPContext workflows.
// A folder is not a valid value.
static const char* const kOrtSessionOptionEpContextFilePath = "ep.context_file_path";
// Flag to specify whether to dump the EP context into the Onnx model.
// "0": dump the EP context into separate file, keep the file name in the Onnx model. (default).
// "1": dump the EP context into the Onnx model.
static const char* const kOrtSessionOptionEpContextEmbedMode = "ep.context_embed_mode";
// Specify the EPContext node name prefix to make it unique
// in case user need to merge/connect multiple EPContext nodes in one model
static const char* const kOrtSessionOptionEpContextNodeNamePrefix = "ep.context_node_name_prefix";
// Share EP related resources across sessions
static const char* const kOrtSessionOptionShareEpContexts = "ep.share_ep_contexts";
// Stop to share EP related resources across sessions from then on
static const char* const kOrtSessionOptionStopShareEpContexts = "ep.stop_share_ep_contexts";
// Used only for context model generation.
// This configuration is used when some nodes are partitioned on the CPU EP and those nodes have external initializers.
// When generating the EP context model, the new model should not rely on the old external data file used by the source ONNX model.
// Use this setting when dumping the EP context model with an external initializers file.
// If specified, all initializers will be placed inside the external data file.
// Otherwise, all initializers will be embedded inside the generated ONNX file.
// By default, this option is not set, meaning all initializers will be included within the ONNX file.
static const char* const kOrtSessionOptionsEpContextModelExternalInitializersFileName =
"ep.context_model_external_initializers_file_name";
// Internal-only flag set by OrtCompileAPI::CompileModel() to signal EPs that this session
// is being used for compilation only and will never be used for inference.
// EPs can use this to skip GPU deserialization and execution context creation, which would
// otherwise be wasteful since the session is destroyed immediately after compilation.
// This is NOT a user-facing option and must not be set directly by application code.
// "0": normal session (default)
// "1": compile-only session (set internally by OrtCompileAPI::CompileModel)
static const char* const kOrtSessionOptionCompileOnly = "session.compile_only";
// Gemm fastmath mode provides fp32 gemm acceleration with bfloat16 based matmul.
// Option values:
// - "0": Gemm FastMath mode is not enabled. [DEFAULT]
// - "1": Gemm FastMath mode is enabled.
static const char* const kOrtSessionOptionsMlasGemmFastMathArm64Bfloat16 = "mlas.enable_gemm_fastmath_arm64_bfloat16";
// Use LUT (Lookup Table) based GEMM for quantized models when available.
// Option values:
// - "0": Do not use LUT based GEMM. [DEFAULT]
// - "1": Use LUT based GEMM when available.
static const char* const kOrtSessionOptionsMlasLutGemm = "mlas.use_lut_gemm";
// Use KleidiAI kernels in MLAS if available.
// Option values:
// - "0": Use KleidiAI kernels when available. [DEFAULT]
// - "1": Disable KleidiAI kernels even if available.
static const char* const kOrtSessionOptionsMlasDisableKleidiAi = "mlas.disable_kleidiai";
// Power-user tuning option for the Arm® KleidiAI™ SME IGEMM convolution route on Arm64.
// For 2D convolutions where both SME IGEMM and the MlasGemm SGEMM fallback are valid routes, work is estimated as:
// output_h * output_w * input_channels * dilated_kernel_h * dilated_kernel_w * filter_count.
// Work above this threshold routes through the SGEMM fallback; work at or below it stays on IGEMM.
// "0" or unset uses the MLAS default heuristic, intended for typical workloads.
// This option exists for perf experimentation; the default may be retuned in future ORT releases.
static const char* const kOrtSessionOptionsMlasKleidiAiConvIgemmMaxWork = "mlas.kleidiai.conv_igemm_max_work";
// Power-user tuning option for the MLAS NCHWc pointwise (1x1) convolution algorithm.
// Controls the maximum number of input channels accumulated per kernel invocation before
// intermediate results are flushed to the output tensor. Values are rounded up to a multiple
// of the NCHWc block size. Larger values reduce output read/write round trips for deep-input
// convolutions at the cost of a larger cache working set.
// "0" or unset uses the MLAS default (128).
// This option exists for perf experimentation; the default may be retuned in future releases.
static const char* const kOrtSessionOptionsMlasNchwcPointwiseConvMaxInputChannelBatch = "mlas.nchwc_pointwise_conv_max_input_channel_batch";
// When converting DQ + MatMul -> MatMulNBits, the accuracy level of the MatMulNBits is controlled by this option.
// Refer to MatMulNBits op schema for more details.
// If not provided, default is 4.
static const char* const kOrtSessionOptionsQDQMatMulNBitsAccuracyLevel = "session.qdq_matmulnbits_accuracy_level";
// Block size used when converting per-tensor or per-axis DQ + MatMul to MatMulNBits.
// Only applies to DQ nodes without an existing block_size attribute (i.e., per-tensor or per-axis quantization).
// Positive value: explicit block_size (must be power-of-2 and >= 16, e.g., 16, 32, 64, 128).
// "0" or not provided: use default block_size of 32.
// "-1": heuristic - largest power-of-2 <= min(K, 256) that minimizes padding.
static const char* const kOrtSessionOptionsQDQMatMulNBitsBlockSize = "session.qdq_matmulnbits_block_size";
// Enable the DQ->MatMulNBits fusion graph transformer.
// "0": disabled (default). "1": enabled.
// This is typically set automatically by InferenceSession when the NvTensorRTRTX EP is registered.
static const char* const kOrtSessionOptionsEnableDQMatMulNBitsFusion = "session.enable_dq_matmulnbits_fusion";
// THIS OPTION IS NOT A REGULAR SESSION OPTION SINCE IT CAN BE MODIFIED AT ANY TIME
// Meant to be used with SetEpDynamicOptions
// Specify the type of workload for this session.
// "Default": OS determines the scheduling priority and processor performance to service this workload. [Default]
// "Efficient": OS treats this workload is efficiency oriented with low scheduling priority and efficient processor performance.
static const char* const kOrtEpDynamicOptionsWorkloadType = "ep.dynamic.workload_type";
// Disables model compilation during session initialization.
//
// If this option is set to "1", inference session creation will fail with error code ORT_MODEL_REQUIRES_COMPILATION
// if compilation is required to run the model on any Execution Provider added to the session.
// Only the following kinds of models are valid when this option is set to "1":
// - Pre-compiled models that have EPContext nodes for the compiling Execution Providers in the session.
// - Non-compiled models that run only on non-compiling Execution Providers, like CPU EP.
//
// See \href https://onnxruntime.ai/docs/execution-providers/EP-Context-Design.html for details about
// compiled models with EPContext nodes.
//
// Option values:
// - "0": EP compile is not disabled. [DEFAULT]
// - "1": EP compile is disabled.
static const char* const kOrtSessionOptionsDisableModelCompile = "session.disable_model_compile";
// Controls behavior when compiled model compatibility is SUPPORTED_PREFER_RECOMPILATION.
// "0": Allow execution with suboptimal performance. [DEFAULT]
// "1": Fail session creation to require recompilation for optimal performance.
// Note: UNSUPPORTED models always fail regardless of this setting.
static const char* const kOrtSessionOptionsFailOnSuboptimalCompiledModel =
"session.fail_on_suboptimal_compiled_model";
// THIS OPTION IS NOT A REGULAR SESSION OPTION SINCE IT CAN BE MODIFIED AT ANY TIME
// Meant to be used with SetEpDynamicOptions
// options for HTP performance mode: "burst", "balanced", "default", "high_performance",
// "high_power_saver", "low_balanced", "extreme_power_saver", "low_power_saver", "power_saver",
// "sustained_high_performance". Default to "default".
static const char* const kOrtEpDynamicOptionsQnnHtpPerformanceMode = "ep.dynamic.qnn_htp_performance_mode";
// Enables the session to record information about the subgraphs/nodes assigned to execution providers.
// When enabled, an application may call Session_GetEpGraphAssignmentInfo() to retrieve the information.
//
// Option values:
// - "0": Recording of EP graph assignment information is disabled. [DEFAULT]
// - "1": Recording of EP graph assignment information is enabled.
static const char* const kOrtSessionOptionsRecordEpGraphAssignmentInfo = "session.record_ep_graph_assignment_info";
// An application enables this option to request that EPs create compiled models (i.e., EPContext models) with EPContext
// nodes that do not store model weights internally. Instead, the weights should be provided by ONNX Runtime as
// explicit inputs to the EPContext nodes.
//
// This option is ignored by an EP if "ep.context_enable" is not set to "1".
//
// If the weights are originally stored in an external file, this allows multiple models to share the same
// external weights file.
//
// Option values:
// - "0": disable. (default)
// - "1": enable.
//
// \deprecated Since version 1.29. Use "ep.enable_weightless" instead, which covers all initializers
// (internal and external) and works in both JIT and AOT flows.
static const char* const kOrtSessionOptionEpEnableWeightlessEpContextNodes = "ep.enable_weightless_ep_context_nodes";
// Layout of the Value KV-cache tensors that the application binds to the past_value input and
// present_value output of com.microsoft.GroupQueryAttention. Applies to every GQA node in the
// model. The Key cache (past_key/present_key) is not affected.
//
// Requires onnxruntime_ENABLE_GQA_VALUE_LAYOUT, enabled by default in normal builds and automatically
// disabled in minimal, extended-minimal, and contrib-disabled builds. When disabled, setting this
// option to any value fails session initialization with ORT_INVALID_ARGUMENT. Leave it unset to load
// a model with a preconverted BNHS boundary; disabled builds do not validate or warn about its layout.
//
// Option values:
// - "BNSH": (batch_size, num_heads, sequence_length, head_size). Matches the operator schema. [DEFAULT]
// - "BNHS": (batch_size, num_heads, head_size, sequence_length).
//
// When "BNHS" is selected, ORT keeps the GQA node itself in BNSH and inserts a
// Transpose(perm=[0,1,3,2]) between the past_value graph input and the node, and another between
// the node and the present_value graph output. An EP that prefers BNHS is expected to fuse that
// Transpose -> GroupQueryAttention -> Transpose sequence into a single operation; an EP that does
// not will execute the transposes, which is correct but costs a full copy of the Value cache in
// each direction per step. The application may still bind one buffer to both past_value and
// present_value; what it loses is the GQA kernel's in-place update of that buffer, because the
// kernel now reads and writes ORT-allocated BNSH intermediates instead.
// Key buffers may remain aliased. CPU handles each cache's aliasing independently; CUDA stages the
// aliased cache when only one pair is shared, adding a cache-sized copy and scratch allocation.
// CUDA sliding-window caches still require both operator cache pairs to be shared, so they cannot
// use this unfused conversion.
//
// Query an EP's preference via the "gqa_preferred_value_layout" OrtEpDevice metadata key
// (kOrtEpDevice_EpMetadataKey_GqaPreferredValueLayout in onnxruntime_ep_device_ep_metadata_keys.h).
//
// Setting "BNSH" explicitly is a claim that the model's Value cache boundary is BNSH, and session
// initialization fails if the model already carries the BNHS conversion (as one saved from a BNHS
// session via "session.optimized_model_filepath" does). Leaving the option unset makes no claim: such
// a model loads unchanged, with a warning, exactly as it did before this option existed.
//
// Scope: this option only describes Value caches that the application itself binds, that is, a
// past_value that is a graph input and a present_value that is a graph output. A Value cache that
// stays inside the graph keeps the BNSH layout, because the application never sees it; ORT logs a
// warning naming the node in that case.
//
// Requesting "BNHS" fails session initialization when a cache is application visible but cannot be
// converted, rather than silently leaving it BNSH and letting the application bind buffers in the
// wrong layout. That happens when:
// - a past_value graph input is read by more than one node, or a present_value graph output is also
// consumed inside the graph (the layout of a shared cache cannot be changed for one reader only);
// - a node already has the layout applied to only one of past_value / present_value;
// - the Value cache is 4-bit quantized (two values are packed per byte along head_size);
// - a Value cache tensor is not rank 4;
// - a Value cache tensor reaches the boundary through a device copy node, which the conversion cannot
// be inserted across;
// - a GroupQueryAttention node is inside a subgraph (a Loop body or BeamSearch decoder), where the
// operator and its boundary are in different graphs and cannot be converted together;
// - the model is in ORT format, which does not run the graph transform that applies this option.
// Note only "BNHS" is refused there; an explicit "BNSH" is still accepted and still checked.
//
// This option takes effect at all graph optimization levels, including ORT_DISABLE_ALL, because it
// changes the layout the session expects at its inputs and outputs rather than optimizing the graph.
static const char* const kOrtSessionOptionsGqaValueLayout = "session.gqa_value_layout";
// Enable weightless mode for all initializers (internal and external).
//
// When enabled, ONNX Runtime requests that the execution provider operate without embedding or copying
// constant initializers.
//
// This option works in both JIT (non-cached) and AOT (EPContext model) flows:
// - JIT: The EP should set drop_constant_initializers to false in OrtNodeFusionOptions so that ORT
// provides the initializer data as inputs to the compiled/fused node. The EP can then access these
// initializers at Compute() time via KernelContext_GetInput().
// NOTE: Extending the lifetime of initializer data obtained via ValueInfo_GetInitializerValue() during
// Compile() so that the EP can cache and reuse data pointers directly is planned but not yet implemented.
// - AOT: ORT generates EPContext models with weightless EPContext nodes. The EP should use the
// "onnx_model_filename" EPContext node attribute or the "ep.context_source_model_path" session option
// to locate the source model's initializer data when creating a session from the compiled model.
//
// ORT checks that the EP supports weightless mode by calling OrtEpApi::GetWeightlessSupport().
// If the EP does not support it, ORT returns an error.
//
// Option values:
// - "0": disable. (default)
// - "1": enable.
//
// \since Version 1.29.
static const char* const kOrtSessionOptionEpEnableWeightless = "ep.enable_weightless";
// Specifies the file path to the original (source) ONNX model when creating a session with a weightless
// EPContext model.
//
// When an EPContext model is generated with weightless mode ("ep.enable_weightless" = "1"), the compiled
// model may not contain the original initializer data. When creating a session from the compiled model,
// the EP needs to load the initializer data from the source model. This session option provides the
// runtime location of the source model, which may differ from the path used at compile time (stored in
// the EPContext node's "onnx_model_filename" attribute).
//
// If not set, the EP falls back to the "onnx_model_filename" attribute in the EPContext node.
//
// If the source model is available as a byte buffer rather than a file path, use
// OrtApi::SessionOptionsSetWeightlessSourceModelBuffer() instead.
//
// \since Version 1.29.
static const char* const kOrtSessionOptionEpContextSourceModelPath = "ep.context_source_model_path";
// Controls the intra-op thread pool size for a session.
// Value should be a base-10 int32 string.
// Equivalent to OrtApi::SetIntraOpNumThreads.
static const char* const kOrtSessionOptionsConfigIntraOpNumThreads = "session.intra_op_num_threads";
// Controls the inter-op thread pool size for a session.
// Value should be a base-10 int32 string.
// Equivalent to OrtApi::SetInterOpNumThreads.
static const char* const kOrtSessionOptionsConfigInterOpNumThreads = "session.inter_op_num_threads";
// Enable or disable the CPU memory arena for a session.
// "0": disable; "1": enable.
// Equivalent to OrtApi::DisableCpuMemArena / OrtApi::EnableCpuMemArena.
static const char* const kOrtSessionOptionsConfigEnableCpuMemArena = "session.enable_cpu_mem_arena";
// Enable or disable memory pattern optimization for a session.
// "0": disable; "1": enable.
// Equivalent to OrtApi::DisableMemPattern / OrtApi::EnableMemPattern.
static const char* const kOrtSessionOptionsConfigEnableMemPattern = "session.enable_mem_pattern";
// Session log identifier.
// Value should be a UTF-8 string.
// Equivalent to OrtApi::SetSessionLogId.
static const char* const kOrtSessionOptionsConfigLogId = "session.log_id";
// Session log severity level.
// Value should be a base-10 int32 string (refer to OrtLoggingLevel values).
// Equivalent to OrtApi::SetSessionLogSeverityLevel.
static const char* const kOrtSessionOptionsConfigLogSeverityLevel = "session.log_severity_level";
// Session log verbosity level.
// Value should be a base-10 int32 string.
// Equivalent to OrtApi::SetSessionLogVerbosityLevel.
static const char* const kOrtSessionOptionsConfigLogVerbosityLevel = "session.log_verbosity_level";
// Enable or disable profiling for a session.
// Empty string: disable profiling.
// Non-empty string: enable profiling and use value as profile file prefix.
// Equivalent to OrtApi::DisableProfiling / OrtApi::EnableProfiling.
static const char* const kOrtSessionOptionsConfigEnableProfiling = "session.enable_profiling";
// Graph optimization level for a session.
// Value should be one of:
// "disable_all", "enable_basic", "enable_extended", "enable_layout", "enable_all".
// Equivalent to OrtApi::SetSessionGraphOptimizationLevel.
static const char* const kOrtSessionOptionsConfigGraphOptimizationLevel = "session.graph_optimization_level";
// File path for saving the optimized model.
// Value should be a path string.
// Equivalent to OrtApi::SetOptimizedModelFilePath.
static const char* const kOrtSessionOptionsConfigOptimizedModelFilePath = "session.optimized_model_filepath";
// Session execution mode.
// Value should be one of:
// "sequential", "parallel", "ort_sequential", "ort_parallel".
// Equivalent to OrtApi::SetSessionExecutionMode.
static const char* const kOrtSessionOptionsConfigExecutionMode = "session.execution_mode";
// Controls whether to use per-session thread pools.
// "0": disable per-session threads (use global thread pools).
// "1": keep per-session threads enabled (default behavior before disable call).
// Equivalent to OrtApi::DisablePerSessionThreads (one-way via API).
static const char* const kOrtSessionOptionsConfigUsePerSessionThreads = "session.use_per_session_threads";
// Enable or disable deterministic compute for a session.
// "0": disable; "1": enable.
// Equivalent to OrtApi::SetDeterministicCompute.
static const char* const kOrtSessionOptionsConfigUseDeterministicCompute = "session.use_deterministic_compute";

731
modules/onnxruntime/include/onnxruntime_training_c_api.h Normal file
View File

@ -0,0 +1,731 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// This file contains the training c apis.
#pragma once
#include <stdbool.h>
#include "onnxruntime_c_api.h"
/** \page training_c_cpp_api Training C & C++ APIs
*
* Training C and C++ APIs are an extension of the \ref c_cpp_api "onnxruntime core C and C++ APIs" and should be used in conjunction with them.
*
* In order to train a model with onnxruntime, the following training artifacts must be generated:
* - The training onnx model
* - The checkpoint file
* - The optimizer onnx model
* - The eval onnx model model (optional)
*
* These training artifacts can be generated as part of an offline step using the python [utilities](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md) made available in the `onnxruntime-training` python package.
*
* After these artifacts have been generated, the C and C++ utilities listed in this documentation can be leveraged to perform training.
*
* If any problem is encountered, please create an [issue](https://github.com/microsoft/onnxruntime/issues/new) with your scenario and requirements, and we will be sure to respond and follow up on the request.
*
* <h1>Training C API</h1>
*
* ::OrtTrainingApi - Training C API functions.
*
* This C structure contains functions that enable users to perform training with onnxruntime.
*
* _Sample Code_:
*
* ```c
* #include <onnxruntime_training_api.h>
*
* OrtApi* g_ort_api = OrtGetApiBase()->GetApi(ORT_API_VERSION);
* OrtTrainingApi* g_ort_training_api = g_ort_api->GetTrainingApi(ORT_API_VERSION);
*
* OrtEnv* env = NULL;
* g_ort_api->CreateEnv(logging_level, logid, &env);
* OrtSessionOptions* session_options = NULL;
* g_ort_api->CreateSessionOptions(&session_options);
*
* OrtCheckpointState* state = NULL;
* g_ort_training_api->LoadCheckpoint(path_to_checkpoint, &state);
*
* OrtTrainingSession* training_session = NULL;
* g_ort_training_api->CreateTrainingSession(env, session_options, training_model_path,
* state, eval_model_path, optimizer_model_path,
* &training_session);
* // Training loop
* {
* g_ort_training_api->TrainStep(...);
* g_ort_training_api->OptimizerStep(...);
* g_ort_training_api->LazyResetGrad(...);
* }
*
* g_ort_training_api->ExportModelForInferencing(training_session, inference_model_path, ...);
* g_ort_training_api->SaveCheckpoint(state, path_to_checkpoint, false);
*
* g_ort_training_api->ReleaseTrainingSession(training_session);
* g_ort_training_api->ReleaseCheckpointState(state);
* ```
*
* > **Note**
* > The ::OrtCheckpointState contains the entire training state that the ::OrtTrainingSession uses. As a result, the training session must always have access to the state. That is to say, the ::OrtCheckpointState instance must outlive the lifetime of the ::OrtTrainingSession instance.
*
* <h1>Training C++ API</h1>
*
* @ref TrainingCpp - Training C++ API classes and functions.
*
* These C++ classes and functions enable users to perform training with onnxruntime.
*
* _Sample Code_:
*
* ```cc
* #include <onnxruntime_training_cxx_api.h>
*
* Ort::Env env;
* Ort::SessionOptions session_options;
*
* auto state = Ort::CheckpointState::LoadCheckpoint(path_to_checkpoint);
* auto training_session = Ort::TrainingSession(env, session_options, state, training_model_path,
* eval_model_path, optimizer_model_path);
*
* // Training Loop
* {
* training_session.TrainStep(...);
* training_session.OptimizerStep(...);
* training_session.LazyResetGrad(...);
* }
*
* training_session->ExportModelForInferencing(inference_model_path, ...);
* Ort::CheckpointState::SaveCheckpoint(state, path_to_checkpoint, false);
* ```
* > **Note**
* > The ::Ort::CheckpointState contains the entire training state that the ::Ort::TrainingSession uses. As a result, the training session must always have access to the state. That is to say, the ::Ort::CheckpointState instance must outlive the lifetime of the ::Ort::TrainingSession instance.
*/
/** @defgroup TrainingC Ort Training C API
* @{
*/
ORT_RUNTIME_CLASS(TrainingSession); // Type that enables performing training for the given user models.
ORT_RUNTIME_CLASS(CheckpointState); // Type that holds the training states for the training session.
/** \brief Type of property to be added to or returned from the ::OrtCheckpointState.
*/
typedef enum OrtPropertyType {
OrtIntProperty = 0,
OrtFloatProperty = 1,
OrtStringProperty = 2,
} OrtPropertyType;
/** \brief The Training C API that holds onnxruntime training function pointers
*
* All the Training C API functions are defined inside this structure as pointers to functions.
* Call OrtApi::GetTrainingApi to get a pointer to this struct.
*
* \nosubgrouping
*/
struct OrtTrainingApi {
/// \name Accessing The Training Session State
/// @{
/** \brief Load a checkpoint state from a file on disk into checkpoint_state.
*
* This function will parse a checkpoint file, pull relevant data and load the training
* state into the checkpoint_state. This checkpoint state can then be used to create the
* training session by invoking OrtTrainingApi::CreateTrainingSession. By doing so, the training
* session will resume training from the given checkpoint state.
* \note Note that the training session created with a checkpoint state uses this state to store the entire
* training state (including model parameters, its gradients, the optimizer states and the properties).
* As a result, it is required that the checkpoint state outlive the lifetime of the training session.
* \note Note that the checkpoint file can be either the complete checkpoint or the nominal checkpoint.
*
* \param[in] checkpoint_path Path to the checkpoint file
* \param[out] checkpoint_state Checkpoint state that contains the states of the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(LoadCheckpoint, _In_ const ORTCHAR_T* checkpoint_path,
_Outptr_ OrtCheckpointState** checkpoint_state);
/** \brief Save the given state to a checkpoint file on disk.
*
* This function serializes the provided checkpoint state to a file on disk.
* This checkpoint can later be loaded by invoking OrtTrainingApi::LoadCheckpoint to resume
* training from this snapshot of the state.
*
* \param[in] checkpoint_state The checkpoint state to save.
* \param[in] checkpoint_path Path to the checkpoint file.
* \param[in] include_optimizer_state Flag to indicate whether to save the optimizer state or not.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(SaveCheckpoint, _In_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* checkpoint_path,
const bool include_optimizer_state);
/// @}
/// \name Implementing The Training Loop
/// @{
/** \brief Create a training session that can be used to begin or resume training.
*
* This function creates a training session based on the env and session options provided that can
* begin or resume training from a given checkpoint state for the given onnx models.
* The checkpoint state represents the parameters of the training session which will be moved
* to the device specified by the user through the session options (if necessary).
* The training session requires four training artifacts
* - The training onnx model
* - The evaluation onnx model (optional)
* - The optimizer onnx model
* - The checkpoint file
*
* These artifacts can be generated using the `onnxruntime-training` python [utility](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md).
*
* \param[in] env Environment to be used for the training session.
* \param[in] options Session options that the user can customize for this training session.
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
* \param[in] train_model_path Model to be used to perform training.
* \param[in] eval_model_path Model to be used to perform evaluation.
* \param[in] optimizer_model_path Model to be used to perform gradient descent.
* \param[out] out Created training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(CreateTrainingSession, _In_ const OrtEnv* env, _In_ const OrtSessionOptions* options,
_Inout_ OrtCheckpointState* checkpoint_state, _In_ const ORTCHAR_T* train_model_path,
_In_ const ORTCHAR_T* eval_model_path, _In_ const ORTCHAR_T* optimizer_model_path,
_Outptr_result_maybenull_ OrtTrainingSession** out);
/** \brief Create a training session that can be used to begin or resume training.
* This api provides a way to load all the training artifacts from buffers instead of files.
*
* \param[in] env Environment to be used for the training session.
* \param[in] options Session options that the user can customize for this training session.
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
* \param[in] train_model_data Buffer containing the model data to be used to perform training
* \param[in] train_data_length Length of the buffer containing train_model_data
* \param[in] eval_model_data Buffer containing the model data to be used to perform evaluation
* \param[in] eval_data_length Length of the buffer containing eval_model_data
* \param[in] optim_model_data Buffer containing the model data to be used to perform weight update
* \param[in] optim_data_length Length of the buffer containing optim_model_data
* \param[out] out Created training session.
*
*/
ORT_API2_STATUS(CreateTrainingSessionFromBuffer, _In_ const OrtEnv* env,
_In_ const OrtSessionOptions* options, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const void* train_model_data, size_t train_data_length,
_In_ const void* eval_model_data, size_t eval_data_length,
_In_ const void* optim_model_data, size_t optim_data_length,
_Outptr_result_maybenull_ OrtTrainingSession** out);
/// @}
/// \name Model IO Information
/// @{
/** \brief Retrieves the number of user outputs in the training model.
*
* This function returns the number of outputs of the training model so that the user can
* allocate space for the number of outputs when OrtTrainingApi::TrainStep is invoked.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] out Number of user outputs in the training model.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
/** \brief Retrieves the number of user outputs in the eval model.
*
* This function returns the number of outputs of the eval model so that the user can
* allocate space for the number of outputs when OrtTrainingApi::EvalStep is invoked.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] out Number of user outputs in the eval model.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetEvalModelOutputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
/** \brief Retrieves the names of user outputs in the training model.
*
* This function returns the names of outputs of the training model that can be associated with the OrtValue(s)
* returned by the OrtTrainingApi::TrainStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] index Index of the output name requested.
* \param[in] allocator Allocator to use to allocate the memory for the name.
* \param[out] output Name of the training model output at the given index.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetTrainingModelOutputName, _In_ const OrtTrainingSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output);
/** \brief Retrieves the names of user outputs in the eval model.
*
* This function returns the names of outputs of the eval model that can be associated with the OrtValue(s) returned
* by the OrtTrainingApi::EvalStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] index Index of the output name requested.
* \param[in] allocator Allocator to use to allocate the memory for the name.
* \param[out] output Name of the eval model output at the given index.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetEvalModelOutputName, _In_ const OrtTrainingSession* sess, size_t index, _Inout_ OrtAllocator* allocator, _Outptr_ char** output);
/// @}
/// \name Implementing The Training Loop
/// @{
/** \brief Reset the gradients of all trainable parameters to zero lazily.
*
* This function sets the internal state of the training session such that the gradients of the trainable
* parameters in the OrtCheckpointState will be scheduled to be reset just before the new gradients are
* computed on the next invocation of the next OrtTrainingApi::TrainStep.
*
* \param[in] session The `this` pointer to the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(LazyResetGrad, _Inout_ OrtTrainingSession* session);
/** \brief Computes the outputs of the training model and the gradients of the trainable parameters for the given inputs
*
* This function performs a training step that computes the outputs of the training model and the gradients
* of the trainable parameters for the given inputs. The train step is performed based on the training model
* that was provided to the training session.
* The OrtTrainingApi::TrainStep is equivalent of running forward propagation and backward propagation in a single
* step.
* The gradients computed are stored inside the training session state so they can be later consumed
* by the OrtTrainingApi::OptimizerStep function.
* The gradients can be lazily reset by invoking the OrtTrainingApi::LazyResetGrad function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] run_options Run options for this training step.
* \param[in] inputs_len Number of user inputs to the training model.
* \param[in] inputs The user inputs to the training model.
* \param[in] outputs_len Number of user outputs expected from this training step.
* \param[out] outputs User outputs computed by train step.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainStep, _Inout_ OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
_In_ size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
_In_ size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
/** \brief Computes the outputs for the eval model for the given inputs
*
* This function performs an eval step that computes the outputs of the eval model for the given inputs.
* The eval step is performed based on the eval model that was provided to the training session.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] run_options Run options for this eval step.
* \param[in] inputs_len Number of user inputs to the eval model.
* \param[in] inputs The user inputs to the eval model.
* \param[in] outputs_len Number of user outputs expected from this eval step.
* \param[out] outputs User outputs computed by eval step.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(EvalStep, _In_ const OrtTrainingSession* sess, _In_opt_ const OrtRunOptions* run_options,
_In_ size_t inputs_len, _In_reads_(inputs_len) const OrtValue* const* inputs,
_In_ size_t outputs_len, _Inout_updates_all_(outputs_len) OrtValue** outputs);
/** \brief Sets the learning rate for this training session.
*
* This function allows users to set the learning rate for the training session. The current
* learning rate is maintained by the training session and can be overwritten by invoking
* this function with the desired learning rate. This function should not be used when a valid
* learning rate scheduler is registered. It should be used either to set the learning rate
* derived from a custom learning rate scheduler or to set a constant learning rate to be used
* throughout the training session.
* \note Please note that this function does not set the initial learning rate that may be needed
* by the predefined learning rate schedulers. To set the initial learning rate for learning
* rate schedulers, please look at the function OrtTrainingApi::RegisterLinearLRScheduler.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] learning_rate Desired learning rate to be set.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(SetLearningRate, _Inout_ OrtTrainingSession* sess, _In_ float learning_rate);
/** \brief Gets the current learning rate for this training session.
*
* This function allows users to get the learning rate for the training session. The current
* learning rate is maintained by the training session, and users can query it for the purpose
* of implementing their own learning rate schedulers.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] learning_rate Learning rate currently in use by the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetLearningRate, _Inout_ OrtTrainingSession* sess, _Out_ float* learning_rate);
/** \brief Performs the weight updates for the trainable parameters using the optimizer model.
*
* This function performs the weight update step that updates the trainable parameters such that they
* take a step in the direction of their gradients (gradient descent). The optimizer step is performed
* based on the optimizer model that was provided to the training session.
* The updated parameters are stored inside the training state so that they can be used by the next
* OrtTrainingApi::TrainStep function call.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] run_options Run options for this optimizer step.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(OptimizerStep, _Inout_ OrtTrainingSession* sess,
_In_opt_ const OrtRunOptions* run_options);
/** \brief Registers a linear learning rate scheduler for the training session.
*
* Register a linear learning rate scheduler that decays the learning rate by linearly updated
* multiplicative factor from the initial learning rate set on the training session to 0. The decay
* is performed after the initial warm up phase where the learning rate is linearly incremented
* from 0 to the initial learning rate provided.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] warmup_step_count Warmup steps for LR warmup.
* \param[in] total_step_count Total step count.
* \param[in] initial_lr The initial learning rate to be used by the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(RegisterLinearLRScheduler, _Inout_ OrtTrainingSession* sess, _In_ const int64_t warmup_step_count,
_In_ const int64_t total_step_count, _In_ const float initial_lr);
/** \brief Update the learning rate based on the registered learing rate scheduler.
*
* Takes a scheduler step that updates the learning rate that is being used by the training session.
* This function should typically be called before invoking the optimizer step for each round,
* or as determined necessary to update the learning rate being used by the training session.
* \note Please note that a valid predefined learning rate scheduler must be first registered to invoke this
* function.
*
* \param[in] sess The `this` pointer to the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(SchedulerStep, _Inout_ OrtTrainingSession* sess);
/// @}
/// \name Accessing The Training Session State
/// @{
/** \brief Retrieves the size of all the parameters.
*
* Calculates the total number of primitive (datatype of the parameters) elements of all the parameters in the
* training state.
* When trainable_only argument is true, the size is calculated for trainable params only.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] out Size of all parameter elements.
* \param[in] trainable_only Whether to skip non-trainable parameters
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetParametersSize, _Inout_ OrtTrainingSession* sess, _Out_ size_t* out, bool trainable_only);
/** \brief Copy all parameters to a contiguous buffer held by the argument parameters_buffer
*
* The parameters_buffer has to be of the size given by GetParametersSize api call,
* with matching setting for the argument trainable_only. All the target parameters must be of the same
* datatype. The OrtValue must be pre-allocated onto
* the desired device. This is a complementary function to OrtTrainingApi::CopyBufferToParameters.
* Parameter ordering is preserved.
* User is responsible for allocating and freeing the resources used by the parameters_buffer.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] trainable_only Whether to skip non-trainable parameters
* \param[out] parameters_buffer The pre-allocated OrtValue buffer to copy onto.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(CopyParametersToBuffer, _Inout_ OrtTrainingSession* sess,
_Inout_ OrtValue* parameters_buffer, bool trainable_only);
/** \brief Copy parameter values from the given contiguous buffer held by parameters_buffer to the training state
*
* The parameters_buffer argument has to be of the size given by OrtTrainingApi::GetParametersSize api call,
* with matching setting for trainable_only argument. All the target parameters must be of the same
* datatype. This is a complementary function to OrtTrainingApi::CopyParametersToBuffer
* and can be used to load updated buffer values onto the training state.
* Parameter ordering is preserved.
* User is responsible for allocating and freeing the resources used by the parameters_buffer.
* In case the training session was created with a nominal checkpoint, invoking this function is required
* to load the updated parameters onto the checkpoint to complete it.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] trainable_only Whether to skip non-trainable parameters
* \param[out] parameters_buffer The pre-allocated OrtValue buffer to copy from.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(CopyBufferToParameters, _Inout_ OrtTrainingSession* sess,
_Inout_ OrtValue* parameters_buffer, bool trainable_only);
/// @}
/// \name Release Training Resources
/// @{
/** \brief Frees up the memory used up by the training session.
*
* This function frees up any memory that was allocated in the training session. The training
* session can no longer be used after this call.
*
*/
ORT_CLASS_RELEASE(TrainingSession);
/** \brief Frees up the memory used up by the checkpoint state.
*
* This function frees up any memory that was allocated in the checkpoint state. The checkpoint
* state can no longer be used after this call.
* \note Note that the checkpoint state must be released only after the training session has been released.
*
*/
ORT_CLASS_RELEASE(CheckpointState);
/// @}
/// \name Prepare For Inferencing
/// @{
/** \brief Export a model that can be used for inferencing.
*
* If the training session was provided with an eval model, the training session can generate
* an inference model if it knows the inference graph outputs. The input inference graph outputs
* are used to prune the eval model so that the inference model's outputs align with the provided outputs.
* The exported model is saved at the path provided and can be used for inferencing with InferenceSession.
* \note Note that the function re-loads the eval model from the path provided to OrtTrainingApi::CreateTrainingSession
* and expects that this path still be valid.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] inference_model_path Path where the inference model should be serialized to.
* \param[in] graph_outputs_len Size of the graph output names array.
* \param[in] graph_output_names Names of the outputs that are needed in the inference model.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(ExportModelForInferencing, _Inout_ OrtTrainingSession* sess,
_In_ const ORTCHAR_T* inference_model_path, size_t graph_outputs_len,
_In_reads_(graph_outputs_len) const char* const* graph_output_names);
/// @}
/// \name Training Utilities
/// @{
/** \brief Sets the seed used for random number generation in Onnxruntime.
*
* Use this function to generate reproducible results. It should be noted that completely reproducible
* results are not guaranteed.
*
* \param[in] seed The seed to be set.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(SetSeed, _In_ const int64_t seed);
/// @}
/// \name Model IO Information
/// @{
/** \brief Retrieves the number of user inputs in the training model.
*
* This function returns the number of inputs of the training model so that the user can accordingly
* allocate the OrtValue(s) provided to the OrtTrainingApi::TrainStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] out Number of user inputs in the training model.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetTrainingModelInputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
/** \brief Retrieves the number of user inputs in the eval model.
*
* This function returns the number of inputs of the eval model so that the user can accordingly
* allocate the OrtValue(s) provided to the OrtTrainingApi::EvalStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[out] out Number of user inputs in the eval model.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetEvalModelInputCount, _In_ const OrtTrainingSession* sess, _Out_ size_t* out);
/** \brief Retrieves the name of the user input at given index in the training model.
*
* This function returns the names of inputs of the training model that can be associated with the
* OrtValue(s) provided to the OrtTrainingApi::TrainStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] index The index of the training model input name requested.
* \param[in] allocator The allocator to use to allocate the memory for the requested name.
* \param[out] output Name of the user input for the training model at the given index.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetTrainingModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
_In_ OrtAllocator* allocator, _Outptr_ char** output);
/** \brief Retrieves the name of the user input at given index in the eval model.
*
* This function returns the names of inputs of the eval model that can be associated with the OrtValue(s) provided
* to the OrtTrainingApi::EvalStep function.
*
* \param[in] sess The `this` pointer to the training session.
* \param[in] index The index of the eval model input name requested.
* \param[in] allocator The allocator to use to allocate the memory for the requested name.
* \param[out] output Name of the user input for the eval model at the given index.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(TrainingSessionGetEvalModelInputName, _In_ const OrtTrainingSession* sess, size_t index,
_In_ OrtAllocator* allocator, _Outptr_ char** output);
/// @}
/// \name Accessing The Training Session State
/// @{
/** \brief Adds or updates the given property to/in the checkpoint state.
*
* Runtime properties such as epoch, training step, best score, and others can be added to the checkpoint
* state by the user by calling this function with the corresponding property name and value.
* The given property name must be unique to be able to successfully add the property.
*
* \param[in] checkpoint_state The checkpoint state which should hold the property.
* \param[in] property_name Name of the property being added or updated.
* \param[in] property_type Type of the property associated with the given name.
* \param[in] property_value Property value associated with the given name.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(AddProperty, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _In_ enum OrtPropertyType property_type,
_In_ void* property_value);
/** \brief Gets the property value associated with the given name from the checkpoint state.
*
* Gets the property value from an existing entry in the checkpoint state. The property must
* exist in the checkpoint state to be able to retrieve it successfully.
*
* \param[in] checkpoint_state The checkpoint state that is currently holding the property.
* \param[in] property_name Name of the property being retrieved.
* \param[in] allocator Allocator used to allocate the memory for the property_value.
* \param[out] property_type Type of the property associated with the given name.
* \param[out] property_value Property value associated with the given name.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetProperty, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* property_name, _Inout_ OrtAllocator* allocator,
_Out_ enum OrtPropertyType* property_type, _Outptr_ void** property_value);
/// @}
/// \name Accessing The Training Session State
/// @{
/** \brief Load a checkpoint state from a buffer into checkpoint_state.
*
* This function will parse a checkpoint bytes buffer, pull relevant data and load the training
* state into the checkpoint_state. This checkpoint state can then be used to create the
* training session by invoking OrtTrainingApi::CreateTrainingSession. By doing so, the training
* session will resume training from the given checkpoint state.
* \note Note that the training session created with a checkpoint state uses this state to store the entire
* training state (including model parameters, its gradients, the optimizer states and the properties).
* As a result, it is required that the checkpoint state outlive the lifetime of the training session.
*
* \param[in] checkpoint_buffer Path to the checkpoint bytes buffer.
* \param[in] num_bytes Number of bytes in the checkpoint buffer.
* \param[out] checkpoint_state Checkpoint state that contains the states of the training session.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(LoadCheckpointFromBuffer, _In_ const void* checkpoint_buffer,
_In_ const size_t num_bytes, _Outptr_ OrtCheckpointState** checkpoint_state);
/** \brief Retrieves the type and shape information of the parameter associated with the given parameter name.
*
* This function retrieves the type and shape of the parameter associated with the given parameter name.
* The parameter must exist in the checkpoint state to be able to retrieve its type and shape information successfully.
*
* \param[in] checkpoint_state The checkpoint state.
* \param[in] parameter_name Name of the parameter being retrieved.
* \param[out] parameter_type_and_shape The type and shape of the parameter being retrieved.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetParameterTypeAndShape, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* parameter_name, _Outptr_ OrtTensorTypeAndShapeInfo** parameter_type_and_shape);
/** \brief Updates the data associated with the model parameter in the checkpoint state for the given parameter name.
*
* This function updates a model parameter in the checkpoint state with the given parameter data.
* The training session must be already created with the checkpoint state that contains the parameter
* being updated. The given parameter is copied over to the registered device for the training session.
* The parameter must exist in the checkpoint state to be able to update it successfully.
*
* \param[in] checkpoint_state The checkpoint state.
* \param[in] parameter_name Name of the parameter being updated.
* \param[in] parameter The parameter data that should replace the existing parameter data.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(UpdateParameter, _Inout_ OrtCheckpointState* checkpoint_state,
_In_ const char* parameter_name, _In_ OrtValue* parameter);
/** \brief Gets the data associated with the model parameter from the checkpoint state for the given parameter name.
*
* This function retrieves the model parameter data from the checkpoint state for the given parameter name.
* The parameter is copied over and returned as an OrtValue. The training session must be already created
* with the checkpoint state that contains the parameter being retrieved.
* The parameter must exist in the checkpoint state to be able to retrieve it successfully.
*
* \param[in] checkpoint_state The checkpoint state.
* \param[in] parameter_name Name of the parameter being retrieved.
* \param[in] allocator Allocator used to allocate the memory for the parameter.
* \param[out] parameter The parameter data that is retrieved from the checkpoint state.
*
* \snippet{doc} snippets.dox OrtStatus Return Value
*
*/
ORT_API2_STATUS(GetParameter, _In_ const OrtCheckpointState* checkpoint_state,
_In_ const char* parameter_name, _Inout_ OrtAllocator* allocator,
_Outptr_ OrtValue** parameter);
/// @}
};
typedef struct OrtTrainingApi OrtTrainingApi;
/// @}

418
modules/onnxruntime/include/onnxruntime_training_cxx_api.h Normal file
View File

@ -0,0 +1,418 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "onnxruntime_training_c_api.h"
#include <optional>
#include <variant>
namespace Ort::detail {
#define ORT_DECLARE_TRAINING_RELEASE(NAME) \
void OrtRelease(Ort##NAME* ptr);
// These release methods must be forward declared before including onnxruntime_cxx_api.h
// otherwise class Base won't be aware of them
ORT_DECLARE_TRAINING_RELEASE(CheckpointState);
ORT_DECLARE_TRAINING_RELEASE(TrainingSession);
} // namespace Ort::detail
#include "onnxruntime_cxx_api.h"
namespace Ort {
/// <summary>
/// This function returns the C training api struct with the pointers to the ort training C functions.
/// If using C++, please use the class instances instead of invoking the C functions directly.
/// </summary>
/// <returns>OrtTrainingApi struct with ort training C function pointers.</returns>
inline const OrtTrainingApi& GetTrainingApi() { return *GetApi().GetTrainingApi(ORT_API_VERSION); }
namespace detail {
#define ORT_DEFINE_TRAINING_RELEASE(NAME) \
inline void OrtRelease(Ort##NAME* ptr) { GetTrainingApi().Release##NAME(ptr); }
ORT_DEFINE_TRAINING_RELEASE(CheckpointState);
ORT_DEFINE_TRAINING_RELEASE(TrainingSession);
#undef ORT_DECLARE_TRAINING_RELEASE
#undef ORT_DEFINE_TRAINING_RELEASE
} // namespace detail
using Property = std::variant<int64_t, float, std::string>;
/**
* \defgroup TrainingCpp Ort Training C++ API
* @{
*/
/** \brief Holds the state of the training session.
*
* This class holds the entire training session state that includes model parameters, their gradients,
* optimizer parameters, and user properties. The Ort::TrainingSession leverages the Ort::CheckpointState
* by accessing and updating the contained training state.
* \note Note that the training session created with a checkpoint state uses this state to store the entire
* training state (including model parameters, its gradients, the optimizer states and the properties).
* The Ort::TrainingSession does not hold a copy of the Ort::CheckpointState and as a result, it is required
* that the checkpoint state outlive the lifetime of the training session.
* \note Note that the checkpoint state can be either the complete checkpoint state or the nominal checkpoint
* state depending on the version provided while loading the checkpoint.
*
*/
class CheckpointState : public detail::Base<OrtCheckpointState> {
private:
CheckpointState(OrtCheckpointState* checkpoint_state) { p_ = checkpoint_state; }
public:
// Construct the checkpoint state by loading the checkpoint by calling LoadCheckpoint
CheckpointState() = delete;
/// \name Accessing The Training Session State
/// @{
/** \brief Load a checkpoint state from a file on disk into checkpoint_state.
*
* This function will parse a checkpoint file, pull relevant data and load the training
* state and return an instance of Ort::CheckpointState. This checkpoint state can then be used to create the
* training session by instantiating Ort::TrainingSession. By doing so, the training session will resume
* training from the given checkpoint state.
*
* \param[in] path_to_checkpoint Path to the checkpoint file
* \return Ort::CheckpointState object which holds the state of the training session parameters.
*
*/
static CheckpointState LoadCheckpoint(const std::basic_string<ORTCHAR_T>& path_to_checkpoint);
/** \brief Load a checkpoint state from a buffer.
*
* This function will parse a checkpoint buffer, pull relevant data and load the training
* state and return an instance of Ort::CheckpointState. This checkpoint state can then be used to create the
* training session by instantiating Ort::TrainingSession. By doing so, the training session will resume
* training from the given checkpoint state.
*
* \param[in] buffer Buffer containing the checkpoint data.
* \return Ort::CheckpointState object which holds the state of the training session parameters.
*
*/
static CheckpointState LoadCheckpointFromBuffer(const std::vector<uint8_t>& buffer);
/** \brief Save the given state to a checkpoint file on disk.
*
* This function serializes the provided checkpoint state to a file on disk.
* This checkpoint can later be loaded by invoking Ort::CheckpointState::LoadCheckpoint to resume
* training from this snapshot of the state.
*
* \param[in] checkpoint_state The checkpoint state to save.
* \param[in] path_to_checkpoint Path to the checkpoint file.
* \param[in] include_optimizer_state Flag to indicate whether to save the optimizer state or not.
*
*/
static void SaveCheckpoint(const CheckpointState& checkpoint_state,
const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
const bool include_optimizer_state = false);
/** \brief Adds or updates the given property to/in the checkpoint state.
*
* Runtime properties such as epoch, training step, best score, and others can be added to the checkpoint
* state by the user by calling this function with the corresponding property name and value.
* The given property name must be unique to be able to successfully add the property.
*
* \param[in] property_name Name of the property being added or updated.
* \param[in] property_value Property value associated with the given name.
*
*/
void AddProperty(const std::string& property_name, const Property& property_value);
/** \brief Gets the property value associated with the given name from the checkpoint state.
*
* Gets the property value from an existing entry in the checkpoint state. The property must
* exist in the checkpoint state to be able to retrieve it successfully.
*
* \param[in] property_name Name of the property being retrieved.
* \return Property value associated with the given property name.
*
*/
Property GetProperty(const std::string& property_name);
/** \brief Updates the data associated with the model parameter in the checkpoint state for the given parameter name.
*
* This function updates a model parameter in the checkpoint state with the given parameter data.
* The training session must be already created with the checkpoint state that contains the parameter
* being updated. The given parameter is copied over to the registered device for the training session.
* The parameter must exist in the checkpoint state to be able to update it successfully.
*
* \param[in] parameter_name Name of the parameter being updated.
* \param[in] parameter The parameter data that should replace the existing parameter data.
*
*/
void UpdateParameter(const std::string& parameter_name, const Value& parameter);
/** \brief Gets the data associated with the model parameter from the checkpoint state for the given parameter name.
*
* This function retrieves the model parameter data from the checkpoint state for the given parameter name.
* The parameter is copied over to the provided OrtValue. The training session must be already created
* with the checkpoint state that contains the parameter being retrieved.
* The parameter must exist in the checkpoint state to be able to retrieve it successfully.
*
* \param[in] parameter_name Name of the parameter being retrieved.
* \return The parameter data that is retrieved from the checkpoint state.
*
*/
Value GetParameter(const std::string& parameter_name);
/// @}
};
/** \brief Trainer class that provides training, evaluation and optimizer methods for training an ONNX models.
*
* The training session requires four training artifacts
* - The training onnx model
* - The evaluation onnx model (optional)
* - The optimizer onnx model
* - The checkpoint file
*
* These artifacts can be generated using the `onnxruntime-training` python [utility](https://github.com/microsoft/onnxruntime/blob/main/orttraining/orttraining/python/training/onnxblock/README.md).
*
*/
class TrainingSession : public detail::Base<OrtTrainingSession> {
private:
size_t training_model_output_count_, eval_model_output_count_;
public:
/// \name Constructing the Training Session
/// @{
/** \brief Create a training session that can be used to begin or resume training.
*
* This constructor instantiates the training session based on the env and session options provided that can
* begin or resume training from a given checkpoint state for the given onnx models.
* The checkpoint state represents the parameters of the training session which will be moved
* to the device specified by the user through the session options (if necessary).
*
* \param[in] env Env to be used for the training session.
* \param[in] session_options SessionOptions that the user can customize for this training session.
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
* \param[in] train_model_path Model to be used to perform training.
* \param[in] eval_model_path Model to be used to perform evaluation.
* \param[in] optimizer_model_path Model to be used to perform gradient descent.
*
*/
TrainingSession(const Env& env, const SessionOptions& session_options, CheckpointState& checkpoint_state,
const std::basic_string<ORTCHAR_T>& train_model_path,
const std::optional<std::basic_string<ORTCHAR_T>>& eval_model_path = std::nullopt,
const std::optional<std::basic_string<ORTCHAR_T>>& optimizer_model_path = std::nullopt);
/** \brief Create a training session that can be used to begin or resume training.
* This constructor allows the users to load the models from buffers instead of files.
*
* \param[in] env Env to be used for the training session.
* \param[in] session_options SessionOptions that the user can customize for this training session.
* \param[in] checkpoint_state Training states that the training session uses as a starting point for training.
* \param[in] train_model_data Buffer containing training model data.
* \param[in] eval_model_data Buffer containing evaluation model data.
* \param[in] optim_model_data Buffer containing optimizer model (used for performing weight/parameter update).
*
*/
TrainingSession(const Env& env, const SessionOptions& session_options, CheckpointState& checkpoint_state,
const std::vector<uint8_t>& train_model_data, const std::vector<uint8_t>& eval_model_data = {},
const std::vector<uint8_t>& optim_model_data = {});
/// @}
/// \name Implementing The Training Loop
/// @{
/** \brief Computes the outputs of the training model and the gradients of the trainable parameters for the given inputs
*
* This function performs a training step that computes the outputs of the training model and the gradients
* of the trainable parameters for the given inputs. The train step is performed based on the training model
* that was provided to the training session.
* The Ort::TrainingSession::TrainStep is equivalent of running forward propagation and backward propagation in a single
* step.
* The gradients computed are stored inside the training session state so they can be later consumed
* by the Ort::TrainingSession::OptimizerStep function.
* The gradients can be lazily reset by invoking the Ort::TrainingSession::LazyResetGrad function.
*
* \param[in] input_values The user inputs to the training model.
* \return A std::vector of Ort::Value objects that represents the output of the forward pass of the training model.
*
*
*/
std::vector<Value> TrainStep(const std::vector<Value>& input_values);
/** \brief Reset the gradients of all trainable parameters to zero lazily.
*
* This function sets the internal state of the training session such that the gradients of the trainable
* parameters in the OrtCheckpointState will be scheduled to be reset just before the new gradients are
* computed on the next invocation of the next Ort::TrainingSession::TrainStep.
*
*/
void LazyResetGrad();
/** \brief Computes the outputs for the eval model for the given inputs
*
* This function performs an eval step that computes the outputs of the eval model for the given inputs.
* The eval step is performed based on the eval model that was provided to the training session.
*
* \param[in] input_values The user inputs to the eval model.
* \return A std::vector of Ort::Value objects that represents the output of the eval pass.
*
*/
std::vector<Value> EvalStep(const std::vector<Value>& input_values);
/** \brief Sets the learning rate for this training session.
*
* This function allows users to set the learning rate for the training session. The current
* learning rate is maintained by the training session and can be overwritten by invoking
* this function with the desired learning rate. This function should not be used when a valid
* learning rate scheduler is registered. It should be used either to set the learning rate
* derived from a custom learning rate scheduler or to set a constant learning rate to be used
* throughout the training session.
* \note Please note that this function does not set the initial learning rate that may be needed
* by the predefined learning rate schedulers. To set the initial learning rate for learning
* rate schedulers, please look at the function Ort::TrainingSession::RegisterLinearLRScheduler.
*
* \param[in] learning_rate Desired learning rate to be set.
*
*/
void SetLearningRate(float learning_rate);
/** \brief Gets the current learning rate for this training session.
*
* This function allows users to get the learning rate for the training session. The current
* learning rate is maintained by the training session, and users can query it for the purpose
* of implementing their own learning rate schedulers.
*
* \return float representing the current learning rate.
*
*/
float GetLearningRate() const;
/** \brief Registers a linear learning rate scheduler for the training session.
*
* Register a linear learning rate scheduler that decays the learning rate by linearly updated
* multiplicative factor from the initial learning rate set on the training session to 0. The decay
* is performed after the initial warm up phase where the learning rate is linearly incremented
* from 0 to the initial learning rate provided.
*
* \param[in] warmup_step_count Warmup steps for LR warmup.
* \param[in] total_step_count Total step count.
* \param[in] initial_lr The initial learning rate to be used by the training session.
*
*/
void RegisterLinearLRScheduler(int64_t warmup_step_count, int64_t total_step_count,
float initial_lr);
/** \brief Update the learning rate based on the registered learing rate scheduler.
*
* Takes a scheduler step that updates the learning rate that is being used by the training session.
* This function should typically be called before invoking the optimizer step for each round,
* or as determined necessary to update the learning rate being used by the training session.
* \note Please note that a valid predefined learning rate scheduler must be first registered to invoke this
* function.
*
*/
void SchedulerStep();
/** \brief Performs the weight updates for the trainable parameters using the optimizer model.
*
* This function performs the weight update step that updates the trainable parameters such that they
* take a step in the direction of their gradients (gradient descent). The optimizer step is performed
* based on the optimizer model that was provided to the training session.
* The updated parameters are stored inside the training state so that they can be used by the next
* Ort::TrainingSession::TrainStep function call.
*
*/
void OptimizerStep();
/// @}
/// \name Prepare For Inferencing
/// @{
/** \brief Export a model that can be used for inferencing.
*
* If the training session was provided with an eval model, the training session can generate
* an inference model if it knows the inference graph outputs. The input inference graph outputs
* are used to prune the eval model so that the inference model's outputs align with the provided outputs.
* The exported model is saved at the path provided and can be used for inferencing with Ort::Session.
* \note Note that the function re-loads the eval model from the path provided to Ort::TrainingSession
* and expects that this path still be valid.
*
* \param[in] inference_model_path Path where the inference model should be serialized to.
* \param[in] graph_output_names Names of the outputs that are needed in the inference model.
*
*/
void ExportModelForInferencing(const std::basic_string<ORTCHAR_T>& inference_model_path,
const std::vector<std::string>& graph_output_names);
/// @}
/// \name Model IO Information
/// @{
/** \brief Retrieves the names of the user inputs for the training and eval models.
*
* This function returns the names of inputs of the training or eval model that can be associated
* with the Ort::Value(s) provided to the Ort::TrainingSession::TrainStep or Ort::TrainingSession::EvalStep
* function.
*
* \param[in] training Whether the training model input names are requested or eval model input names.
* \return Graph input names for either the training model or the eval model.
*
*/
std::vector<std::string> InputNames(const bool training);
/** \brief Retrieves the names of the user outputs for the training and eval models.
*
* This function returns the names of outputs of the training or eval model that can be associated
* with the Ort::Value(s) returned by the Ort::TrainingSession::TrainStep or Ort::TrainingSession::EvalStep
* function.
*
* \param[in] training Whether the training model output names are requested or eval model output names.
* \return Graph output names for either the training model or the eval model.
*
*/
std::vector<std::string> OutputNames(const bool training);
/// @}
/// \name Accessing The Training Session State
/// @{
/** \brief Returns a contiguous buffer that holds a copy of all training state parameters
*
* \param[in] only_trainable Whether to only copy trainable parameters or to copy all parameters.
* \return Contiguous buffer to the model parameters.
*
*/
Value ToBuffer(const bool only_trainable);
/** \brief Loads the training session model parameters from a contiguous buffer
*
* In case the training session was created with a nominal checkpoint, invoking this function is required
* to load the updated parameters onto the checkpoint to complete it.
*
* \param[in] buffer Contiguous buffer to load the parameters from.
*/
void FromBuffer(Value& buffer);
/// @}
};
/// \name Training Utilities
/// @{
/** \brief This function sets the seed for generating random numbers.
*
* Use this function to generate reproducible results. It should be noted that completely
* reproducible results are not guaranteed.
*
* \param[in] seed Manual seed to use for random number generation.
*/
void SetSeed(const int64_t seed);
/// @}
/// @}
} // namespace Ort
#include "onnxruntime_training_cxx_inline.h"

295
modules/onnxruntime/include/onnxruntime_training_cxx_inline.h Normal file
View File

@ -0,0 +1,295 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "onnxruntime_training_c_api.h"
#include "onnxruntime_cxx_api.h"
namespace Ort {
inline TrainingSession::TrainingSession(const Env& env, const SessionOptions& session_options,
CheckpointState& checkpoint_state,
const std::basic_string<ORTCHAR_T>& train_model_path,
const std::optional<std::basic_string<ORTCHAR_T>>& eval_model_path,
const std::optional<std::basic_string<ORTCHAR_T>>& optimizer_model_path) {
ThrowOnError(GetTrainingApi().CreateTrainingSession(
env, session_options, checkpoint_state,
train_model_path.c_str(),
eval_model_path.has_value() ? eval_model_path.value().c_str() : nullptr,
optimizer_model_path.has_value() ? optimizer_model_path.value().c_str() : nullptr,
&p_));
ThrowOnError(GetTrainingApi().TrainingSessionGetTrainingModelOutputCount(p_, &training_model_output_count_));
ThrowOnError(GetTrainingApi().TrainingSessionGetEvalModelOutputCount(p_, &eval_model_output_count_));
}
inline TrainingSession::TrainingSession(const Env& env, const SessionOptions& session_options,
CheckpointState& checkpoint_state,
const std::vector<uint8_t>& train_model_data,
const std::vector<uint8_t>& eval_model_data,
const std::vector<uint8_t>& optim_model_data) {
ThrowOnError(GetTrainingApi().CreateTrainingSessionFromBuffer(
env, session_options, checkpoint_state,
train_model_data.data(), train_model_data.size(),
eval_model_data.data(), eval_model_data.size(),
optim_model_data.data(), optim_model_data.size(),
&p_));
ThrowOnError(GetTrainingApi().TrainingSessionGetTrainingModelOutputCount(p_, &training_model_output_count_));
ThrowOnError(GetTrainingApi().TrainingSessionGetEvalModelOutputCount(p_, &eval_model_output_count_));
}
inline std::vector<Value> TrainingSession::TrainStep(const std::vector<Value>& input_values) {
std::vector<Value> output_values;
output_values.reserve(training_model_output_count_);
for (size_t i = 0; i < training_model_output_count_; i++) output_values.emplace_back(nullptr);
auto ort_input_values = reinterpret_cast<const OrtValue* const*>(input_values.data());
auto ort_output_values = reinterpret_cast<OrtValue**>(output_values.data());
RunOptions run_options;
ThrowOnError(GetTrainingApi().TrainStep(
p_, run_options, input_values.size(), ort_input_values,
training_model_output_count_, ort_output_values));
return output_values;
}
inline void TrainingSession::LazyResetGrad() {
ThrowOnError(GetTrainingApi().LazyResetGrad(p_));
}
inline std::vector<Value> TrainingSession::EvalStep(const std::vector<Value>& input_values) {
std::vector<Value> output_values;
output_values.reserve(eval_model_output_count_);
for (size_t i = 0; i < eval_model_output_count_; i++) output_values.emplace_back(nullptr);
auto ort_input_values = reinterpret_cast<const OrtValue* const*>(input_values.data());
auto ort_output_values = reinterpret_cast<OrtValue**>(output_values.data());
RunOptions run_options;
ThrowOnError(GetTrainingApi().EvalStep(
p_, run_options, input_values.size(), ort_input_values,
eval_model_output_count_, ort_output_values));
return output_values;
}
inline void TrainingSession::SetLearningRate(float learning_rate) {
ThrowOnError(GetTrainingApi().SetLearningRate(p_, learning_rate));
}
inline float TrainingSession::GetLearningRate() const {
float learning_rate = 0;
ThrowOnError(GetTrainingApi().GetLearningRate(p_, &learning_rate));
return learning_rate;
}
inline void TrainingSession::RegisterLinearLRScheduler(int64_t warmup_step_count, int64_t total_step_count,
float initial_lr) {
ThrowOnError(GetTrainingApi().RegisterLinearLRScheduler(p_, warmup_step_count, total_step_count,
initial_lr));
}
inline void TrainingSession::SchedulerStep() {
ThrowOnError(GetTrainingApi().SchedulerStep(p_));
}
inline void TrainingSession::OptimizerStep() {
RunOptions run_options;
ThrowOnError(GetTrainingApi().OptimizerStep(p_, run_options));
}
inline std::vector<std::string> TrainingSession::InputNames(const bool training) {
auto& input_count_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelInputCount
: GetTrainingApi().TrainingSessionGetEvalModelInputCount;
auto& input_name_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelInputName
: GetTrainingApi().TrainingSessionGetEvalModelInputName;
size_t input_count = 0;
ThrowOnError(input_count_function(p_, &input_count));
std::vector<std::string> input_names(input_count);
AllocatorWithDefaultOptions allocator;
for (size_t index = 0; index < input_count; ++index) {
char* input_name;
ThrowOnError(input_name_function(p_, index, allocator, &input_name));
input_names[index] = std::string(input_name);
allocator.Free(input_name);
}
return input_names;
}
inline std::vector<std::string> TrainingSession::OutputNames(const bool training) {
auto& output_count_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelOutputCount
: GetTrainingApi().TrainingSessionGetEvalModelOutputCount;
auto& output_name_function = training ? GetTrainingApi().TrainingSessionGetTrainingModelOutputName
: GetTrainingApi().TrainingSessionGetEvalModelOutputName;
size_t output_count = 0;
ThrowOnError(output_count_function(p_, &output_count));
std::vector<std::string> output_names(output_count);
AllocatorWithDefaultOptions allocator;
for (size_t index = 0; index < output_count; ++index) {
char* output_name;
ThrowOnError(output_name_function(p_, index, allocator, &output_name));
output_names[index] = std::string(output_name);
allocator.Free(output_name);
}
return output_names;
}
inline Value TrainingSession::ToBuffer(const bool only_trainable) {
size_t buffer_size = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &buffer_size, only_trainable));
std::array<int64_t, 1> buffer_shape{static_cast<int64_t>(buffer_size)};
AllocatorWithDefaultOptions allocator;
Value buffer = Value::CreateTensor(allocator, buffer_shape.data(), 1U,
ONNXTensorElementDataType::ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);
ThrowOnError(GetTrainingApi().CopyParametersToBuffer(p_, buffer, only_trainable));
return buffer;
}
inline void TrainingSession::FromBuffer(Value& buffer) {
if (!buffer.IsTensor()) {
ThrowStatus(Status("Incorrect buffer received. Expected a tensor buffer.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
auto tensor_info = buffer.GetTensorTypeAndShapeInfo();
auto buffer_shape = tensor_info.GetShape();
if (buffer_shape.size() != 1U) {
ThrowStatus(Status("Incorrect buffer received. Expected a contiguous tensor buffer.",
OrtErrorCode::ORT_INVALID_ARGUMENT));
}
auto buffer_size = buffer_shape.front();
size_t session_buffer_size = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &session_buffer_size, false));
if (buffer_size == static_cast<int64_t>(session_buffer_size)) {
ThrowOnError(GetTrainingApi().CopyBufferToParameters(p_, buffer, false));
return;
}
size_t session_buffer_size_trainable_only = 0U;
ThrowOnError(GetTrainingApi().GetParametersSize(p_, &session_buffer_size_trainable_only, true));
if (buffer_size == static_cast<int64_t>(session_buffer_size_trainable_only)) {
ThrowOnError(GetTrainingApi().CopyBufferToParameters(p_, buffer, true));
return;
} else {
ThrowStatus(Status("Incorrect buffer size received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
}
inline CheckpointState CheckpointState::LoadCheckpoint(const std::basic_string<ORTCHAR_T>& path_to_checkpoint) {
OrtCheckpointState* checkpoint_state;
ThrowOnError(GetTrainingApi().LoadCheckpoint(path_to_checkpoint.c_str(), &checkpoint_state));
return CheckpointState(checkpoint_state);
}
inline CheckpointState CheckpointState::LoadCheckpointFromBuffer(const std::vector<uint8_t>& buffer) {
OrtCheckpointState* checkpoint_state;
ThrowOnError(GetTrainingApi().LoadCheckpointFromBuffer(buffer.data(), buffer.size(), &checkpoint_state));
return CheckpointState(checkpoint_state);
}
inline void CheckpointState::SaveCheckpoint(const CheckpointState& checkpoint_states,
const std::basic_string<ORTCHAR_T>& path_to_checkpoint,
const bool include_optimizer_state) {
ThrowOnError(GetTrainingApi().SaveCheckpoint(checkpoint_states, path_to_checkpoint.c_str(),
include_optimizer_state));
}
inline void TrainingSession::ExportModelForInferencing(const std::basic_string<ORTCHAR_T>& inference_model_path,
const std::vector<std::string>& graph_output_names) {
std::vector<const char*> output_names;
output_names.reserve(graph_output_names.size());
for (const auto& output_name : graph_output_names) {
output_names.push_back(output_name.c_str());
}
ThrowOnError(GetTrainingApi().ExportModelForInferencing(
p_, inference_model_path.c_str(), graph_output_names.size(), output_names.data()));
}
inline void SetSeed(const int64_t seed) {
ThrowOnError(GetTrainingApi().SetSeed(seed));
}
inline void CheckpointState::AddProperty(const std::string& property_name, const Property& property_value) {
if (std::holds_alternative<int64_t>(property_value)) {
int64_t value = std::get<int64_t>(property_value);
void* value_p = &value;
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtIntProperty, value_p));
} else if (std::holds_alternative<float>(property_value)) {
float value = std::get<float>(property_value);
void* value_p = &value;
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtFloatProperty, value_p));
} else if (std::holds_alternative<std::string>(property_value)) {
std::string value = std::get<std::string>(property_value);
auto buffer = std::make_unique<char[]>(value.length() + 1);
memcpy(buffer.get(), value.c_str(), value.length());
// AddProperty takes a char* and calls PropertyBag::AddProperty which takes a std::string. The data will be
// copied at that point so buffer can free the local allocation once the call is made.
ThrowOnError(GetTrainingApi().AddProperty(p_, property_name.c_str(), OrtPropertyType::OrtStringProperty,
buffer.get()));
} else {
ThrowStatus(Status("Unknown property type received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
}
}
inline Property CheckpointState::GetProperty(const std::string& property_name) {
void* property_value = nullptr;
OrtPropertyType property_type;
AllocatorWithDefaultOptions allocator;
ThrowOnError(GetTrainingApi().GetProperty(p_, property_name.c_str(), allocator, &property_type, &property_value));
Property property;
switch (property_type) {
case OrtPropertyType::OrtIntProperty: {
auto value_p = reinterpret_cast<int64_t*>(property_value);
property = *value_p;
allocator.Free(property_value);
break;
}
case OrtPropertyType::OrtFloatProperty: {
auto value_p = reinterpret_cast<float*>(property_value);
property = *value_p;
allocator.Free(property_value);
break;
}
case OrtPropertyType::OrtStringProperty: {
auto value_p = reinterpret_cast<char*>(property_value);
property = std::string(value_p);
allocator.Free(property_value);
break;
}
default: {
ThrowStatus(Status("Unknown property type received.", OrtErrorCode::ORT_INVALID_ARGUMENT));
break;
}
}
return property;
}
inline void CheckpointState::UpdateParameter(const std::string& parameter_name, const Value& parameter) {
ThrowOnError(GetTrainingApi().UpdateParameter(p_, parameter_name.c_str(), parameter));
}
inline Value CheckpointState::GetParameter(const std::string& parameter_name) {
AllocatorWithDefaultOptions allocator;
OrtValue* parameter;
ThrowOnError(GetTrainingApi().GetParameter(p_, parameter_name.c_str(), allocator, &parameter));
return Value{parameter};
}
} // namespace Ort

18
modules/onnxruntime/include/provider_options.h Normal file
View File

@ -0,0 +1,18 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>
#include <unordered_map>
#include <vector>
namespace onnxruntime {
// data types for execution provider options
using ProviderOptions = std::unordered_map<std::string, std::string>;
using ProviderOptionsVector = std::vector<ProviderOptions>;
using ProviderOptionsMap = std::unordered_map<std::string, ProviderOptions>;
} // namespace onnxruntime

108
modules/onnxruntime/module.jai Normal file
View File

@ -0,0 +1,108 @@
// onnxruntime, loaded at run time. Nothing is linked: ort_load finds the library, takes its one
// export, OrtGetApiBase, and from it the OrtApi function table (bindings.jai, from generate.jai).
// A program that imports this module starts without onnxruntime and only fails in ort_load.
//
// Version: see fetch.jai (ORT_VERSION) and <os>/VERSION.txt. The bindings are for ORT_API_VERSION;
// an older library is refused, a newer one serves this version of the table.
ort: struct {
api: *OrtApi;
library_path: string; // the library that was loaded
version: string; // onnxruntime's own version string, such as 1.30.0
}
// Loads onnxruntime from the first folder that has it: beside the running executable, then each
// of folders. A full path is always used: Windows keeps an older onnxruntime.dll in System32 for
// its own use, which a bare name could find first.
ort_load :: (folders: [] string = .[]) -> ok: bool, error: string {
if ort.api return true, "";
candidates: [..] string;
candidates.allocator = temp;
array_add(*candidates, path_strip_filename(get_path_of_running_executable()));
for folders array_add(*candidates, it);
tried: [..] string;
tried.allocator = temp;
for folder: candidates {
path := tprint("%/%", trim_right(folder, "/\\"), ORT_LIBRARY_NAME);
if !file_exists(path) { array_add(*tried, path); continue; }
get_api_base, load_error := ort_open_library(path);
if !get_api_base return false, load_error;
base := get_api_base();
if !base return false, tprint("% returned no API.", path);
version := to_string(base.GetVersionString());
api := base.GetApi(ORT_API_VERSION);
if !api return false, tprint("% is onnxruntime %, which does not offer API version % (these bindings); use the version fetch.jai pins.", path, version, ORT_API_VERSION);
ort.api = api;
ort.library_path = copy_string(path,, context.default_allocator);
ort.version = copy_string(version,, context.default_allocator);
return true, "";
}
return false, tprint("% was not found. Looked for:\n %\nRun `jai fetch.jai` in modules/onnxruntime and rebuild, so the build copies it beside the executable.",
ORT_LIBRARY_NAME, join(..tried, "\n ",, temp));
}
// A failed status as its message, and the status released. ok and "" for success (a null status).
ort_check :: (status: OrtStatusPtr) -> ok: bool, message: string {
if !status return true, "";
message := tprint("%", to_string(ort.api.GetErrorMessage(status)));
ort.api.ReleaseStatus(status);
return false, message;
}
// A path as the C API takes it: UTF-16 on Windows, UTF-8 elsewhere. The bindings were generated on
// Windows, where ORTCHAR_T is a wchar_t, so elsewhere the pointer is cast to the type they expect. Temp.
ort_path :: (path: string) -> *ORTCHAR_T {
#if OS == .WINDOWS {
return utf8_to_wide(path,, temp);
} else {
return cast(*ORTCHAR_T) temp_c_string(path);
}
}
#load "bindings.jai";
#scope_module
Get_Api_Base :: #type () -> *OrtApiBase #c_call;
#if OS == .WINDOWS {
ORT_LIBRARY_NAME :: "onnxruntime.dll";
ort_open_library :: (path: string) -> Get_Api_Base, error: string {
// Altered search path: the libraries onnxruntime loads in turn (onnxruntime_providers_shared.dll,
// onnxruntime_providers_cuda.dll) are looked for in its folder.
windows_path := replace(path, "/", "\\",, temp);
module := LoadLibraryExW(utf8_to_wide(windows_path,, temp), null, LOAD_WITH_ALTERED_SEARCH_PATH);
if !module return null, tprint("Could not load % (Windows error %).", path, GetLastError());
get_api_base := cast(Get_Api_Base) GetProcAddress(module, "OrtGetApiBase");
if !get_api_base return null, tprint("% has no OrtGetApiBase; it is not onnxruntime.", path);
return get_api_base, "";
}
#import "Windows";
#import "Windows_Utf8";
} else {
#if OS == .MACOS ORT_LIBRARY_NAME :: "libonnxruntime.dylib";
else ORT_LIBRARY_NAME :: "libonnxruntime.so";
ort_open_library :: (path: string) -> Get_Api_Base, error: string {
handle := dlopen(temp_c_string(path), RTLD_NOW);
if !handle return null, tprint("Could not load %: %", path, to_string(dlerror()));
get_api_base := cast(Get_Api_Base) dlsym(handle, "OrtGetApiBase");
if !get_api_base return null, tprint("% has no OrtGetApiBase; it is not onnxruntime.", path);
return get_api_base, "";
}
#import "POSIX";
}
#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "System";

5
src/App.jai
View File

@ -14,6 +14,7 @@ USING_WINDOWS :: OS == .WINDOWS;
#load "app/main.jai"; // entry point
#load "app/app_ui.jai";
#load "app/viewer.jai";
#load "app/ocr_ui.jai";
#load "app/exit_shutdown.jai";
#load "base/circle_buffer.jai";
@ -48,6 +49,8 @@ Shared_Channel :: #import "Shared_Memory_Channel";
#import "stb_image";
#import "stb_image_resize";
#import "Capture_File"; // the .mcap format, copied from D:\SnagitIndex\SnagitConverter
#import "OCR";
#import "onnxruntime";
// Basic's arenas import Secure_Memory at its defaults; this import adds its test suite.
using Secure_Memory :: #import "Secure_Memory"()(LOAD_TESTS=COMPILE_TESTS);
@ -162,6 +165,8 @@ App :: struct {
allow_multiple_instances: bool; // Skips the single-instance check, so a dev build can run beside an installed one.
should_exit_early: bool;
run_tests: bool;
headless: bool; // -headless (or a no_crash_dialog build): no message boxes on failures, and run_tests exits when done
exit_code: s32; // the process's exit code when the main loop ends
last_window_restore_time: Apollo_Time;
imgui_context: *ImGui.ImGuiContext;

3
src/app/app_ui.jai
View File

@ -7,6 +7,7 @@ APP_MAIN_WINDOW_LABEL :: "Main";
// Once, after ImGui, the fonts, the theme and the settings are up and before the first frame.
app_ui_startup :: () {
viewer_startup();
ocr_ui_startup();
}
// Every frame, before input is handed to ImGui. Poll background work here.
@ -17,9 +18,11 @@ app_ui_frame_start :: () {
// stop background work here. Must not touch ImGui, since the main thread is still drawing frames.
app_ui_shutdown :: () {
viewer_shutdown();
ocr_ui_shutdown();
}
app_main_window_ui :: () {
push_imgui_window(APP_MAIN_WINDOW_LABEL);
viewer_ui();
ocr_ui_frame();
}

6
src/app/main.jai
View File

@ -14,6 +14,8 @@ main :: () {
args_combined := join(..app.args, " ");
app.run_tests = array_find(app.args, "run_tests");
app.headless = array_find(app.args, "-headless") || NO_CRASH_DIALOG;
runtime_support_headless = app.headless;
app.allow_multiple_instances = array_find(app.args, "-allow-multiple-instances");
defer report_memory_leaks();
@ -127,7 +129,8 @@ main :: () {
if app.run_tests {
#if COMPILE_TESTS {
testing_init(assert_on_fail=DEBUG, log_successful_tests=DEBUG);
// Headless, a failed suite is counted and reported in the exit code rather than asserted on.
testing_init(assert_on_fail = DEBUG && !app.headless, log_successful_tests=DEBUG);
run_all_tests();
} else {
log_error("run_tests was passed, but this build was compiled without tests. Rebuild with the compile_tests argument to include them.");
@ -343,6 +346,7 @@ main :: () {
if ini_filename ImGui.SaveIniSettingsToDisk(ini_filename);
// By the time the loop above exits, app_exit_shutdown_poll (called from app_frame) has already
// confirmed the background shutdown work finished.
if app.exit_code exit(app.exit_code); // e.g. a headless test run with failures
}
// Returns false when the main loop should exit.

256
src/app/ocr_ui.jai Normal file
View File

@ -0,0 +1,256 @@
// OCR in the app: an engine made on a background thread at startup (modules/OCR), an OCR window that
// says whether it runs on the GPU and, if not, what to install, and the OCR button in the Capture
// window, which reads the capture shown and draws its lines over it.
//
// Engines take a moment to make and OCR takes a moment to run, so both happen on one background
// thread; the UI only looks at what it has finished.
OCR_WINDOW :: "OCR";
OCR_DEFAULT_MODEL :: "ppocr-v5-mobile";
ocr_ui: struct {
models: [] Ocr_Model_Info;
models_error: string;
model_index: s32; // the model chosen in the window
// Owned by the main thread while thread_busy is false; by the background thread while it is true.
engine: *Ocr_Engine;
engine_model: string; // the model engine runs
engine_error: string; // why the last engine could not be made
thread: Thread;
thread_started: bool;
thread_busy: bool; // atomic: set by the main thread, cleared by the background one
task: Ocr_Task;
// The last capture read, for the Capture window.
lines: [] Ocr_Line;
lines_index: s32 = -1; // the capture they are for
lines_error: string;
lines_seconds: float64;
show_boxes := true;
window_open := false;
toast_shown := false;
}
Ocr_Task_Kind :: enum { NONE; CREATE; READ; }
Ocr_Task :: struct {
kind: Ocr_Task_Kind;
model: string;
path: string; // READ: the .mcap
index: s32; // READ: its capture index
}
ocr_models_folder :: () -> string {
return tprint("%ocr-models", path_strip_filename(get_path_of_running_executable()));
}
ocr_ui_startup :: () {
push_allocator(context.default_allocator);
models, error := ocr_list_models(ocr_models_folder());
ocr_ui.models = models;
if error ocr_ui.models_error = copy_string(error);
for ocr_ui.models if it.name == OCR_DEFAULT_MODEL ocr_ui.model_index = cast(s32) it_index;
if ocr_ui.models.count ocr_ui_start(.{kind = .CREATE, model = ocr_ui.models[ocr_ui.model_index].name});
}
ocr_ui_shutdown :: () {
if ocr_ui.thread_started thread_is_done(*ocr_ui.thread, -1); // finishes what it is doing
ocr_destroy(ocr_ui.engine);
}
// Starts task on the background thread; false while it is still busy with the last one.
ocr_ui_start :: (task: Ocr_Task) -> bool {
if ocr_ui.thread_busy return false;
if ocr_ui.thread_started thread_is_done(*ocr_ui.thread, -1); // finished: thread_busy is clear
push_allocator(context.default_allocator);
ocr_ui.task = task;
ocr_ui.task.model = copy_string(task.model);
ocr_ui.task.path = copy_string(task.path);
ocr_ui.thread_busy = true;
if ocr_ui.thread_started thread_deinit(*ocr_ui.thread);
ocr_ui.thread_started = thread_init(*ocr_ui.thread, ocr_ui_thread_proc, temporary_storage_size = 4 * 1024 * 1024, name = "ocr");
if !ocr_ui.thread_started { ocr_ui.thread_busy = false; return false; }
thread_start(*ocr_ui.thread);
return true;
}
ocr_ui_thread_proc :: (thread: *Thread) -> s64 {
context.allocator = context.default_allocator;
task := ocr_ui.task;
if task.kind == {
case .CREATE;
ocr_destroy(ocr_ui.engine);
ocr_ui.engine = null;
engine, ok, error := ocr_create(ocr_models_folder(), task.model);
ocr_ui.engine = engine;
ocr_ui.engine_model = task.model;
ocr_ui.engine_error = ifx ok then "" else copy_string(error);
if ok log("[OCR] % on %.%", task.model, engine.provider, ifx engine.gpu_explanation then tprint(" %", engine.gpu_explanation) else "");
else log_error("[OCR] Could not start %: %", task.model, error);
case .READ;
start := current_time_monotonic();
lines, error := ocr_ui_read_capture(task.path);
ocr_ui.lines = lines;
ocr_ui.lines_error = error;
ocr_ui.lines_index = task.index;
ocr_ui.lines_seconds = to_float64_seconds(current_time_monotonic() - start);
}
free(task.model);
free(task.path);
atomic_swap(*ocr_ui.thread_busy, false);
return 0;
}
// The lines in a capture's image as the Capture window shows it: the rendered one, annotations and
// all, or the original when there is none. Heap.
ocr_ui_read_capture :: (path: string) -> [] Ocr_Line, error: string {
if !ocr_ui.engine return .[], copy_string("No OCR engine is running.");
data, read_ok := read_entire_file(path, log_errors = false,, temp);
if !read_ok return .[], sprint("Cannot read %.", path);
m, ok, error := mcap_parse(cast([] u8) data);
if !ok return .[], copy_string(error);
png := mcap_rendered(m);
if !png.count png = m.original;
w, h, channels: s32;
pixels := stbi_load_from_memory(png.data, cast(s32) png.count, *w, *h, *channels, 4);
if !pixels return .[], copy_string("Its image cannot be decoded.");
defer stbi_image_free(pixels);
images := Ocr_Image.[.{pixels, w, h, w * 4}];
results := ocr_process(ocr_ui.engine, images);
defer array_free(results);
if results[0].failed return .[], results[0].error;
return results[0].lines, "";
}
ocr_ui_free_lines :: () {
for ocr_ui.lines free(it.text);
array_free(ocr_ui.lines);
free(ocr_ui.lines_error);
ocr_ui.lines = .[];
ocr_ui.lines_error = "";
ocr_ui.lines_index = -1;
}
// ---------------------------------------------------------------- drawing
// Each frame: the toast once the engine is up on the CPU though a GPU could run it, and the OCR window.
ocr_ui_frame :: () { using ImGui;
busy := ocr_ui.thread_busy;
if !busy && ocr_ui.engine && !ocr_ui.toast_shown {
ocr_ui.toast_shown = true;
in_use, explanation := ocr_gpu_status(ocr_ui.engine);
if !in_use && contains(explanation, "can run it once") {
toast_show(.Error, "OCR runs on the CPU: CUDA libraries are missing. See the OCR window for what to install.");
ocr_ui.window_open = true;
}
}
if !ocr_ui.window_open return;
SetNextWindowSize(.{640, 360}, .FirstUseEver);
if Begin(OCR_WINDOW, *ocr_ui.window_open) {
if ocr_ui.models_error {
TextColored(app_theme.palette.warning, "%", ocr_ui.models_error);
TextWrapped("Run `jai fetch_models.jai` in modules/OCR and `jai fetch.jai` in modules/onnxruntime, then rebuild.");
} else {
// The model: changing it makes a new engine on the background thread.
preview := ifx ocr_ui.models.count then ocr_ui.models[ocr_ui.model_index].title else "";
SetNextItemWidth(320);
if BeginCombo("Model", temp_c_string(preview)) {
for ocr_ui.models {
if Selectable(temp_c_string(it.title), it_index == ocr_ui.model_index) && it_index != ocr_ui.model_index && !busy {
ocr_ui.model_index = cast(s32) it_index;
ocr_ui.toast_shown = true; // said once already, or not needed
ocr_ui_free_lines();
ocr_ui_start(.{kind = .CREATE, model = it.name});
}
}
EndCombo();
}
if busy && ocr_ui.task.kind == .CREATE {
TextColored(app_theme.palette.dim, "Starting % ...", ocr_ui.task.model);
} else if ocr_ui.engine_error {
TextColored(app_theme.palette.warning, "OCR could not start:");
TextWrapped("%", ocr_ui.engine_error);
} else if ocr_ui.engine && !busy {
in_use, explanation := ocr_gpu_status(ocr_ui.engine);
Text("Running on: %", ifx in_use then "the GPU (CUDA)" else "the CPU");
SameLine();
TextColored(app_theme.palette.dim, "onnxruntime %, % worker threads", ort.version, ocr_ui.engine.workers);
if explanation { Separator(); text_with_links(explanation); }
}
// The text of the capture last read, in reading order.
if !busy && ocr_ui.lines_index >= 0 && ocr_ui.lines_index < viewer.captures.count {
Separator();
Text("Text of %:", viewer.captures[ocr_ui.lines_index].name);
BeginChild("ocr text");
ocr_ui_capture_text(ocr_ui.lines_index);
EndChild();
}
}
}
End();
}
// The OCR row of the Capture window, above its image: the button, and how the last read went.
ocr_ui_capture_controls :: (index: s32, path: string) { using ImGui;
busy := ocr_ui.thread_busy;
can_read := ocr_ui.engine && !busy;
if !can_read BeginDisabled();
if SmallButton("OCR") {
ocr_ui_free_lines();
ocr_ui_start(.{kind = .READ, path = path, index = index});
}
if !can_read EndDisabled();
SameLine();
if busy && ocr_ui.task.kind == .READ && ocr_ui.task.index == index {
TextColored(app_theme.palette.dim, "Reading...");
} else if !busy && ocr_ui.lines_index == index {
if ocr_ui.lines_error TextColored(app_theme.palette.warning, "%", ocr_ui.lines_error);
else TextColored(app_theme.palette.dim, "% lines in % s", ocr_ui.lines.count, FF(ocr_ui.lines_seconds, 2));
SameLine();
Checkbox("Boxes", *ocr_ui.show_boxes);
} else if !ocr_ui.engine {
TextColored(app_theme.palette.dim, "(OCR is not running; see the OCR window)");
}
SameLine();
if SmallButton("OCR window") ocr_ui.window_open = true;
}
// The lines read from this capture, as boxes over its image, drawn at image_min with scale; and their
// text in a list below the image area.
ocr_ui_capture_overlay :: (index: s32, image_min: ImGui.ImVec2, scale: float) { using ImGui;
if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.show_boxes return;
draw := GetWindowDrawList();
colour := GetColorU32(.PlotHistogram);
for ocr_ui.lines {
p0 := ImVec2.{image_min.x + it.x * scale, image_min.y + it.y * scale};
p1 := ImVec2.{p0.x + it.width * scale, p0.y + it.height * scale};
draw.AddRect(draw, p0, p1, colour, 0, 1.5);
if IsMouseHoveringRect(p0, p1) SetTooltip("% (%)", it.text, FF(it.confidence, 2));
}
}
// Text a line at a time, each https:// address in it drawn as a link that opens in the browser.
text_with_links :: (text: string) { using ImGui;
for line: split(text, "\n",, temp) {
start := find_index_from_left(line, "https://");
if start < 0 { TextWrapped("%", line); continue; }
end := start;
while end < line.count && line[end] != #char " " && line[end] != #char ")" && line[end] != #char "," end += 1;
url := slice(line, start, end - start);
if start > 0 { TextUnformatted(slice(line, 0, start)); SameLine(0, 0); }
TextLinkOpenURL(temp_c_string(url));
if IsItemHovered() SetTooltip("Open % in the browser", url);
if end < line.count { SameLine(0, 0); TextUnformatted(advance(line, end)); }
}
}
ocr_ui_capture_text :: (index: s32) { using ImGui;
if ocr_ui.thread_busy || ocr_ui.lines_index != index || !ocr_ui.lines.count return;
for ocr_ui.lines TextUnformatted(it.text);
}

3
src/app/viewer.jai
View File

@ -526,6 +526,7 @@ viewer_full_window :: () { using ImGui;
if SliderFloat("Zoom", *viewer.full_zoom, VIEWER_MIN_ZOOM, VIEWER_MAX_ZOOM, "%.0f%%", .Logarithmic) viewer.full_fit = .OFF;
SameLine();
if SmallButton("100%") { viewer.full_fit = .OFF; viewer.full_zoom = 100; }
ocr_ui_capture_controls(viewer.full_index, capture.path); // app/ocr_ui.jai
// Left and Right step through the captures in grid order while this window has focus.
if IsWindowFocused(.RootAndChildWindows) {
@ -558,6 +559,8 @@ viewer_full_window :: () { using ImGui;
}
zoom := clamp(viewer.full_zoom, VIEWER_MIN_ZOOM, VIEWER_MAX_ZOOM) / 100;
Image(gpu_texture_ref(viewer.full_texture), .{source_w * zoom, source_h * zoom});
// OCR boxes are in the image's own pixels, which is the captured size.
ocr_ui_capture_overlay(viewer.full_index, GetItemRectMin(), zoom);
EndChild();
// The previous image stays up while the next one loads, so stepping through does not flicker.
}

9
src/test/testing.jai
View File

@ -8,6 +8,7 @@ test_thread_group: *Thread_Group;
Test_Data :: struct {
work_added: s32;
work_completed: s32;
work_failed: s32;
log_successful_tests: bool;
assert_on_fail: bool;
@ -121,6 +122,7 @@ testing_cleanup :: () {
log("[TEST PASSED][%] suite passed", tt.module);
} else {
log_error("[TEST FAILED][%] suite reported failures", tt.module);
test_data.work_failed += 1;
if test_data.assert_on_fail assert(false, tprint("Test suite % failed.", tt.module));
}
}
@ -134,6 +136,13 @@ testing_cleanup :: () {
if test_data.work_added == test_data.work_completed {
log("All % test tasks complete.", test_data.work_completed);
// Headless (run_tests -headless), the app exits once the tests are done: 0 if every suite
// passed, 1 otherwise.
if app.headless {
if test_data.work_failed log_error("% of % test suites failed.", test_data.work_failed, test_data.work_completed);
app.exit_code = ifx test_data.work_failed then cast(s32) 1 else 0;
app.should_exit = true;
}
testing_finish();
}
}

148
tools/ocr_bench.jai Normal file
View File

@ -0,0 +1,148 @@
// ocr-bench: runs OCR over a folder of .mcap captures with each model, on the GPU when it can and
// on the CPU, and reports the time per image and per stage. Each capture's lines are written to
// OUT/<model>-<provider>/<name>.txt, to read against the images.
//
// ocr-bench MCAP_FOLDER [-models a,b] [-cpu] [-batch N] [-limit N] [-out FOLDER]
//
// -models the models to run (default: every model in ocr-models/models.config)
// -cpu the CPU only; by default the GPU too, when it is available
// -batch images per ocr_process call (default 8)
// -limit at most this many captures (default all)
// -out where the text goes (default ocr-out beside this program)
//
// The models are read from ocr-models beside this program, where the build copies them.
USAGE :: "Usage: ocr-bench MCAP_FOLDER [-models a,b] [-cpu] [-batch N] [-limit N] [-out FOLDER]\n";
main :: () {
runtime_support_headless = true; // a command-line tool: a failure is reported on the console and exits with 1, no message box
args := get_command_line_arguments();
if args.count < 2 { print(USAGE); exit(2); }
folder := args[1];
model_names: [..] string;
cpu_only := false;
batch := 8;
limit := S64_MAX_VALUE;
exe_folder := path_strip_filename(get_path_of_running_executable());
// Lasting copies: temporary storage is reset after every batch.
out_folder := sprint("%ocr-out", exe_folder);
i := 2;
while i < args.count {
arg := args[i];
value := ifx i + 1 < args.count then args[i + 1] else "";
if arg == {
case "-models"; for split(value, ",") array_add(*model_names, trim(it)); i += 2;
case "-cpu"; cpu_only = true; i += 1;
case "-batch"; batch = max(1, string_to_int(value)); i += 2;
case "-limit"; limit = max(1, string_to_int(value)); i += 2;
case "-out"; out_folder = value; i += 2;
case; print("Unknown option %.\n%", arg, USAGE); exit(2);
}
}
models_folder := sprint("%ocr-models", exe_folder);
if !model_names.count {
models, error := ocr_list_models(models_folder);
if error { print("%\n", error); exit(1); }
for models array_add(*model_names, it.name);
}
// The captures, decoded once: the original image of each, as the camera took it.
Capture :: struct { name: string; pixels: *u8; width, height: s32; }
captures: [..] Capture;
decode_start := current_time_monotonic();
for file_list(folder) {
if captures.count >= limit break;
if !ends_with_nocase(it, ".mcap") continue;
data, ok := read_entire_file(it);
if !ok continue;
m, parsed, error := mcap_parse(cast([] u8) data);
if !parsed { print("Skipping %: %\n", it, error); continue; }
w, h, channels: s32;
pixels := stbi_load_from_memory(m.original.data, cast(s32) m.original.count, *w, *h, *channels, 4);
if !pixels { print("Skipping %: its image cannot be decoded.\n", it); continue; }
array_add(*captures, .{copy_string(path_strip_extension(path_filename(it))), pixels, w, h});
free(data);
}
if !captures.count { print("No .mcap files in %.\n", folder); exit(1); }
quick_sort(captures, (a: Capture, b: Capture) -> s64 { return compare(a.name, b.name); });
megapixels: float64;
for captures megapixels += cast(float64) it.width * it.height / 1_000_000;
print("% captures (% megapixels) decoded in % s.\n\n", captures.count, FF(megapixels, 1), FF(seconds_since(decode_start), 2));
images := NewArray(captures.count, Ocr_Image);
for captures images[it_index] = .{it.pixels, it.width, it.height, it.width * 4};
for model_name: model_names {
for prefer_gpu: bool.[true, false] {
if prefer_gpu && cpu_only continue;
create_start := current_time_monotonic();
engine, ok, error := ocr_create(models_folder, model_name, prefer_gpu);
if !ok { print("% : %\n\n", model_name, error); break; }
create_seconds := seconds_since(create_start);
in_use, explanation := ocr_gpu_status(engine);
provider := ocr_provider(engine);
if prefer_gpu && !in_use {
// Said once; the CPU run that follows is the same thing.
print("%\n\n", explanation);
ocr_destroy(engine);
continue;
}
print("== % on % (sessions created in % s, % worker threads)\n", model_name, provider, FF(create_seconds, 2), engine.workers);
text_folder := sprint("%/%-%", out_folder, model_name, to_lower_copy(tprint("%", provider),, temp));
make_directory_if_it_does_not_exist(text_folder, recursive = true);
// The first image alone: the first run pays for setting up kernels and memory.
first_start := current_time_monotonic();
ocr_process(engine, array_view(images, 0, 1),, temp);
first_seconds := seconds_since(first_start);
ocr_reset_stats(engine);
start := current_time_monotonic();
lines, characters, failures: s64;
done := 0;
while done < images.count {
count := min(batch, images.count - done);
results := ocr_process(engine, array_view(images, done, count),, temp);
for result, n: results {
name := captures[done + n].name;
if result.failed { failures += 1; print(" % failed: %\n", name, result.error); continue; }
builder: String_Builder;
builder.allocator = temp;
for result.lines {
print_to_builder(*builder, "% % % % % %\n", it.x, it.y, it.width, it.height, FF(it.confidence, 3), it.text);
characters += it.text.count;
}
lines += result.lines.count;
write_entire_file(tprint("%/%.txt", text_folder, name), *builder);
}
done += count;
reset_temporary_storage();
}
seconds := seconds_since(start);
stats := ocr_stats(engine);
print(" first image alone: % s (warm-up)\n", FF(first_seconds, 3));
print(" all % images: % s, % ms per image, % megapixels/s, % failed\n", images.count, FF(seconds, 2), FF(1000 * seconds / images.count, 1), FF(megapixels / seconds, 2), failures);
print(" found: % lines, % bytes of text\n", lines, characters);
print(" stage time (summed over threads): detection % s, recognition % s\n", FF(stats.detect_us / 1_000_000.0, 2), FF(stats.recognize_us / 1_000_000.0, 2));
print(" text in %\n\n", text_folder);
ocr_destroy(engine);
}
}
}
seconds_since :: (start: Apollo_Time) -> float64 {
return to_float64_seconds(current_time_monotonic() - start);
}
S64_MAX_VALUE :: 0x7FFF_FFFF_FFFF_FFFF;
#import "Basic";
#import "String";
#import "File";
#import "File_Utilities";
#import "Sort";
#import "System";
#import "stb_image";
#import "Capture_File";
#import "OCR";