compute-onnx

ONNX Runtime based compute backend for GoMLX.

It allows GoMLX models to be executed via ONNX Runtime using:

• Native Desktop/Server: CPU, CUDA (NVIDIA GPU), or MIGraphX (AMD ROCm GPU).
• WebAssembly / Browser: WebGPU, CPU (WASM SIMD), WebNN, or WebGL in browsers (see ORT Web Guide).

It supports dynamic shapes and exporting models to .onnx files.

To run the Adult dataset demo with the ONNX backend:

GOMLX_BACKEND=onnx go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo

Or targeting a specific accelerator (e.g. CUDA):

GOMLX_BACKEND=onnx:cuda go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo

Or on an AMD GPU (ROCm + MIGraphX):

GOMLX_BACKEND=onnx:migraphx go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo

Configuration can be specified in the GOMLX_BACKEND environment variable using onnx:<options> or onnxruntime:<options> (comma-separated).

• cpu: Force CPU execution.

  GOMLX_BACKEND=onnx:cpu

• cuda / gpu: Force CUDA GPU execution (uses ONNX Runtime CUDA Execution Provider via OrtIoBinding).

  GOMLX_BACKEND=onnx:cuda

• migraphx / rocm / amd (EXPERIMENTAL): Force AMD GPU execution (uses ONNX Runtime MIGraphX Execution Provider). Requires ROCm and MIGraphX to be installed (sudo apt install migraphx migraphx-dev half). If no ORT library with the MIGraphX provider is found, one is automatically installed from AMD's manylinux wheels matching the local ROCm version — it can also be installed manually with go run github.com/gomlx/compute-onnx/cmd/onnxruntime_installer -migraphx. Notes:
  • Only float32/int32/int64 graphs are supported;
  • Models with scalar (0-dimensional) inputs fails (there is a bug upstream);
  • Buffers are always tranferred back to CPU (host), which makes training slower (TODO fix).

  GOMLX_BACKEND=onnx:migraphx

• Custom Library Path: Specify an explicit path to the ONNX Runtime .so (or .dylib / .dll) shared library file. This explicitly bypasses ONNXRUNTIME_SHARED_LIBRARY_PATH.

  GOMLX_BACKEND=onnx:/path/to/libonnxruntime.so

• empty (default): Automatically detects if an NVIDIA GPU is present via nvidia-smi and defaults to CUDA if available, then checks for a discrete AMD GPU (ROCm/MIGraphX), otherwise falling back to CPU.

  GOMLX_BACKEND=onnx

Fine-tune ONNX Runtime session execution, multi-threading, and memory behavior:

• intra_op_num_threads=<int>: Sets the number of threads used to parallelize execution within an operator. Default is 0 (ORT uses all available CPU cores). When running multiple concurrent worker threads or goroutines, setting this to 1 prevents thread oversubscription.
  • Aliases: intra_threads, intraopnumthreads
• inter_op_num_threads=<int>: Sets the number of threads used to parallelize execution across independent operators in the graph (only effective when execution_mode=parallel). Default is 0.
  • Aliases: inter_threads, interopnumthreads
• execution_mode=<parallel|sequential>: Controls whether independent operators in the computation graph execute sequentially or in parallel.
  • Aliases: executionmode, or standalone parallel / sequential.
• cpu_mem_arena=<bool>: Enables or disables ONNX Runtime's CPU memory arena allocator. Default is true. Disabling (cpu_mem_arena=false or no_cpu_mem_arena) reduces memory lock contention when multiple worker threads execute concurrently.
  • Aliases: cpumemarena, or standalone cpu_mem_arena / no_cpu_mem_arena.
• mem_pattern=<bool>: Enables or disables memory pattern optimization (tracing and reusing memory allocations for static shapes). Default is true. Only valid in sequential execution mode.
  • Aliases: mempattern, or standalone mem_pattern / no_mem_pattern.
• graph_optimization_level=<level>: Controls the graph optimization level applied by ONNX Runtime:
  • 0 / disable_all / none: Disable all optimizations.
  • 1 / basic: Basic optimizations (constant folding, redundant node elimination).
  • 2 / extended: Extended optimizations (operator fusions such as LayerNorm, GELU, MatMul+Add).
  • 3 / layout: Layout optimizations.
  • 99 / all: Enable all available optimizations.
  • Aliases: graphoptimizationlevel, opt_level, optlevel.
• session_clones=<int>: Maximum number of concurrent ONNX Runtime session clones pooled per compiled Executable. Default is 8 on CPU (1 on GPU). Clones are allocated on-demand as concurrent goroutines execute, eliminating session lock contention across workers while avoiding unnecessary memory allocation for single-threaded tasks.
  • Aliases: sessionclones, clones, session_pool, sessionpool.

Example for high-throughput multi-worker concurrent inference:

GOMLX_BACKEND="onnx:cpu,intra_op_num_threads=1,inter_op_num_threads=1,cpu_mem_arena=false,execution_mode=parallel"

This allows one to export GOMLX trained (or fine-tuned) models to ONNX.

See an example in UCI-Adult demo. If you have a pre-trained file in a directory called base:

GOMLX_BACKEND=onnx go run -tags=onnx ./examples/adult/demo/ -checkpoint "base" -save_onnx="/tmp/a.onnx" -vmodule=save_onnx=1

The backend automatically locates or manages the required ONNX Runtime shared library (libonnxruntime.so / onnxruntime.dll):

• Custom Library Path: Set the ONNXRUNTIME_SHARED_LIBRARY_PATH environment variable or pass an explicit library path in the backend configuration (e.g. GOMLX_BACKEND=onnx:/path/to/libonnxruntime.so) to point directly to the shared library binary. Passing an explicit path in the configuration bypasses ONNXRUNTIME_SHARED_LIBRARY_PATH.
• Auto-Installation: If no library path is provided, the backend automatically downloads and extracts prebuilt official ONNX Runtime binaries locally (e.g. ~/.local/lib/onnxruntime/ on Linux).
• Disabling Auto-Installation: Set the environment variable GOMLX_NO_AUTO_INSTALL=1 or call onnxbackend.EnableAutoInstall(false) programmatically to disable automatic downloads (useful for offline environments or container deployments).

The MIGraphX execution provider relies on a local ROCm installation:

• ROCM_PATH: Directory where ROCm is installed (defaults to /opt/rocm). It is used to locate rocminfo and the HIP/MIGraphX libraries when auto-detecting an AMD GPU and its ROCm version.
• GOMLX_MIGRAPHX_CACHE_DIR: Directory where the MIGraphX compiled-program (.mxr) for each model is cached, skipping the expensive graph compilation on subsequent runs. Equivalent to the migraphx_cache_dir config key (e.g. GOMLX_BACKEND=onnx:migraphx,migraphx_cache_dir=/tmp/mxr); an empty value disables caching.

If graph compilation or session creation fails in ONNX Runtime, setting the GOMLX_ONNX_SAVE_ON_FAILURE environment variable instructs the backend to automatically save the serialized ONNX model protobuf to the specified file path before returning the compilation error:

GOMLX_ONNX_SAVE_ON_FAILURE="/tmp/failed_model.onnx" go run -tags=onnx ...

This allows you to inspect the invalid graph using onnx_printer or Netron to diagnose the failure.

This repository includes a CLI tool in cmd/onnx_printer to inspect and pretty-print .onnx model files in the terminal:

go run github.com/gomlx/compute-onnx/cmd/onnx_printer path/to/model.onnx

It formats input, output, and node tensor shapes using GoMLX shapes.Shape (including named dynamic axes) and prints each graph operation on a single line. Tensor constants and initializers are truncated to 10 elements by default (controlled via -max_items / -n).

Example usage:

# Print model details with a maximum of 5 items for constant values
go run github.com/gomlx/compute-onnx/cmd/onnx_printer -max_items 5 /tmp/model.onnx

# Read from stdin
cat /tmp/model.onnx | go run github.com/gomlx/compute-onnx/cmd/onnx_printer

Tip: For interactive graphical visualization of ONNX models, you can open .onnx model files using Netron.

• Backend Log Level (log=<level>): Configures ONNX Runtime's internal logging severity level.

  • log=0: Errors only (severity level 3 / ERROR)
  • log=1: Warnings (severity level 2 / WARNING)
  • log=2: Informational (severity level 1 / INFO)
  • log=3: Verbose (severity level 0 / VERBOSE)

  Example:

  GOMLX_BACKEND="onnx:cuda,log=2"

• Execution Timing Log (-vmodule=executable=1): Enables per-step execution timing breakdown printed via klog using humanize.Duration.

  Example:

  GOMLX_BACKEND=onnx:cuda go run -tags=onnx github.com/gomlx/gomlx/examples/adult/demo -vmodule=executable=1

• Go
• ONNX
• ONNX Runtime