package onnx
import (
"fmt"
"os"
"path/filepath"
"runtime"
"strings"
"sync"
ort "github.com/yalue/onnxruntime_go"
"golang.org/x/sys/unix"
)
var (
runtimeOnce sync.Once
runtimeErr error
)
// FindORTLibrary finds the libonnxruntime shared library: the explicit path,
// then $CHIRAMI_ORT_LIB (which the nix package sets), next to the executable,
// then the usual system locations.
func FindORTLibrary(explicitPath string) (string, error) {
if explicitPath != "" {
if _, err := os.Stat(explicitPath); err == nil {
return explicitPath, nil
}
return "", fmt.Errorf("specified onnxruntime library not found: %s", explicitPath)
}
if envPath := os.Getenv("CHIRAMI_ORT_LIB"); envPath != "" {
if _, err := os.Stat(envPath); err == nil {
return envPath, nil
}
return "", fmt.Errorf("CHIRAMI_ORT_LIB points to non-existent path: %s", envPath)
}
libName := "libonnxruntime.so"
if runtime.GOOS == "darwin" {
libName = "libonnxruntime.dylib"
}
// Check next to current executable
if exe, err := os.Executable(); err == nil {
p := filepath.Join(filepath.Dir(exe), libName)
if _, err := os.Stat(p); err == nil {
return p, nil
}
}
// Standard search locations
candidates := []string{
"/opt/homebrew/lib/" + libName,
"/usr/local/lib/" + libName,
"/usr/lib/" + libName,
"/usr/lib64/" + libName,
libName,
}
for _, c := range candidates {
if _, err := os.Stat(c); err == nil {
return c, nil
}
}
return "", fmt.Errorf("could not find %s: set CHIRAMI_ORT_LIB", libName)
}
// InitRuntime initializes the ONNX Runtime environment once.
func InitRuntime(explicitLibPath string) error {
runtimeOnce.Do(func() {
if ort.IsInitialized() {
return
}
libPath, err := FindORTLibrary(explicitLibPath)
if err != nil {
runtimeErr = err
return
}
ort.SetSharedLibraryPath(libPath)
err = ort.InitializeEnvironment()
if err != nil {
runtimeErr = fmt.Errorf("initialize onnxruntime (%s): %w", libPath, err)
return
}
})
return runtimeErr
}
// newSession opens one model so that its weights stay on disk: the model is
// kept in ORT format, memory-mapped, and ORT is told to use those bytes in
// place instead of copying them (and not to repack weights or keep arenas).
// The weights then count as page cache the OS can drop, not as chirami-ocr's
// own memory, and a model the image never needs is never read in: about
// 120 MB of real memory instead of 700 MB. The ORT-format file is made from
// the .onnx on first use, and again if this onnxruntime cannot load it.
func newSession(path string, inputs, outputs []string) (*ort.DynamicAdvancedSession, error) {
ortPath := strings.TrimSuffix(path, ".onnx") + ".ort"
s, err := mmapSession(ortPath, inputs, outputs)
if err == nil {
return s, nil
}
if err := convert(path, ortPath); err != nil {
return nil, err
}
return mmapSession(ortPath, inputs, outputs)
}
func sessionOptions() (*ort.SessionOptions, error) {
opts, err := ort.NewSessionOptions()
if err != nil {
return nil, err
}
for _, err := range []error{
opts.SetCpuMemArena(false),
opts.SetMemPattern(false),
opts.AddSessionConfigEntry("session.disable_prepacking", "1"),
} {
if err != nil {
opts.Destroy()
return nil, err
}
}
return opts, nil
}
// mmapSession opens an ORT-format model in place. The mapping lives as long
// as the helper process does.
func mmapSession(path string, inputs, outputs []string) (*ort.DynamicAdvancedSession, error) {
f, err := os.Open(path)
if err != nil {
return nil, err
}
defer f.Close()
st, err := f.Stat()
if err != nil {
return nil, err
}
b, err := unix.Mmap(int(f.Fd()), 0, int(st.Size()), unix.PROT_READ, unix.MAP_SHARED)
if err != nil {
return nil, err
}
opts, err := sessionOptions()
if err != nil {
unix.Munmap(b)
return nil, err
}
defer opts.Destroy()
opts.AddSessionConfigEntry("session.use_ort_model_bytes_directly", "1")
opts.AddSessionConfigEntry("session.use_ort_model_bytes_for_initializers", "1")
s, err := ort.NewDynamicAdvancedSessionWithONNXData(b, inputs, outputs, opts)
if err != nil {
unix.Munmap(b)
return nil, err
}
return s, nil
}
// convert saves the ONNX model at src as an optimised ORT-format model at
// dst, atomically.
func convert(src, dst string) error {
opts, err := sessionOptions()
if err != nil {
return err
}
defer opts.Destroy()
tmp := dst + ".tmp"
opts.AddSessionConfigEntry("session.save_model_format", "ORT")
if err := opts.SetOptimizedModelFilePath(tmp); err != nil {
return err
}
s, err := ort.NewDynamicAdvancedSession(src, nil, nil, opts)
if err != nil {
return fmt.Errorf("convert %s: %w", filepath.Base(src), err)
}
s.Destroy()
return os.Rename(tmp, dst)
}