Add ONNX Runtime backend support (OpenVINO / DirectML / TensorRT / CPU, Windows & Linux)

Compiling.md

@@ -152,3 +152,32 @@ As also mentioned in the instructions below but repeated here for visibility, if
    * Pre-trained neural nets are available at [the main training website](https://katagotraining.org/).
    * You will probably want to edit `configs/gtp_example.cfg` (see "Tuning for Performance" above).
    * If using OpenCL, you will want to verify that KataGo is picking up the correct device when you run it (e.g. some systems may have both an Intel CPU OpenCL and GPU OpenCL, if KataGo appears to pick the wrong one, you can correct this by specifying `openclGpuToUse` in `configs/gtp_example.cfg`).
+
+## ONNX Runtime backend (optional)
+The `ONNX` backend runs inference through [ONNX Runtime](https://onnxruntime.ai/), which supports several execution providers (CPU, OpenVINO for Intel GPUs/NPUs, CUDA, TensorRT, etc.). It reuses KataGo's built-in `OnnxModelBuilder` (the same graph emitter the TensorRT backend uses), so its IO protocol and post-processing are identical to TensorRT; only the runtime differs. It is useful when you want to run KataGo on a non-NVIDIA accelerator that already has an ONNX Runtime execution provider, or for cross-vendor benchmarking.
+
+> **Note**: This backend is more involved to set up than the built-in backends above, because the official prebuilt ONNX Runtime packages do **not** ship the execution providers you may need (e.g. the OpenVINO EP). You generally have to build ONNX Runtime from source with the provider(s) you want enabled.
+
+### Requirements
+   * Everything KataGo normally needs (CMake, a C++17 compiler, zlib).
+   * ONNX Runtime, built from source with the execution provider(s) you intend to use. For the OpenVINO EP, build ONNX Runtime with `--use_openvino` against an installed OpenVINO toolkit. See https://onnxruntime.ai/docs/install/ for build instructions.
+   * Protobuf. The ONNX graph is serialized as an ONNX `ModelProto`, so `find_package(Protobuf)` must succeed. A protobuf 3.x (no abseil dependency) works; the version bundled in the ONNX Runtime source build tree is known to work.
+   * If using the OpenVINO EP, the OpenVINO runtime toolkit itself, plus its runtime DLLs at runtime (see below).
+
+### Compile
+   * Point CMake at your ONNX Runtime install tree and protobuf, and select the backend:
+     ```
+     cmake -S KataGo/cpp -B KataGo/cpp/build -DUSE_BACKEND=ONNX ^
+       -DONNXRUNTIME_ROOT=<path-to-onnxruntime-install> ^
+       -DProtobuf_PROTOC_EXECUTABLE=<protoc> ^
+       -DProtobuf_INCLUDE_DIR=<protobuf-include> ^
+       -DProtobuf_LIBRARY=<protobuf-lib>
+     cmake --build KataGo/cpp/build -j
+     ```
+   * `-DONNXRUNTIME_ROOT` should contain `include/onnxruntime/`, `lib/onnxruntime.lib` (or `.so`/`.dylib`), and the provider DLLs.
+   * As with other backends, `-DNO_GIT_REVISION=1` avoids embedding the git hash, and `-DBUILD_DISTRIBUTED=1` enables distributed-training support.
+
+### Runtime
+   * The `onnxruntime` shared library must be on your path or beside the executable.
+   * When using the OpenVINO EP, also deploy the OpenVINO runtime DLLs beside the executable (`openvino.dll`, `openvino_intel_gpu_plugin.dll`, `tbb12.dll`, `cache.json`, etc.), or put them on the system path.
+   * Configure the provider in `configs/gtp_example.cfg` via the `onnx*` keys, e.g. `onnxProvider=openvino` and `onnxOpenVINODeviceType=GPU`. See the ONNX settings block in `configs/gtp_example.cfg` for the full list.

cpp/CMakeLists.txt

@@ -44,7 +44,7 @@ endif()
 set(BUILD_DISTRIBUTED 0 CACHE BOOL "Build with http support for contributing to distributed training")
 set(USE_BACKEND CACHE STRING "Neural net backend")
 string(TOUPPER "${USE_BACKEND}" USE_BACKEND)
-set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL)
+set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL ONNX)
 
 set(USE_TCMALLOC 0 CACHE BOOL "Use TCMalloc")
 set(NO_GIT_REVISION 0 CACHE BOOL "Disable embedding the git revision into the compiled exe")
@@ -164,8 +164,13 @@ elseif(USE_BACKEND STREQUAL "EIGEN")
   set(NEURALNET_BACKEND_SOURCES
     neuralnet/eigenbackend.cpp
     )
+elseif(USE_BACKEND STREQUAL "ONNX")
+  message(STATUS "-DUSE_BACKEND=ONNX, using ONNX Runtime backend.")
+  set(NEURALNET_BACKEND_SOURCES
+    neuralnet/onnxbackend.cpp
+    )
 elseif(USE_BACKEND STREQUAL "")
-  message(WARNING "${ColorBoldRed}WARNING: Using dummy neural net backend, intended for non-neural-net testing only, will fail on any code path requiring a neural net. To use neural net, specify -DUSE_BACKEND=CUDA or -DUSE_BACKEND=TENSORRT or -DUSE_BACKEND=OPENCL or -DUSE_BACKEND=EIGEN to compile with the respective backend.${ColorReset}")
+  message(WARNING "${ColorBoldRed}WARNING: Using dummy neural net backend, intended for non-neural-net testing only, will fail on any code path requiring a neural net. To use neural net, specify -DUSE_BACKEND=CUDA or -DUSE_BACKEND=TENSORRT or -DUSE_BACKEND=OPENCL or -DUSE_BACKEND=EIGEN or -DUSE_BACKEND=ONNX to compile with the respective backend.${ColorReset}")
   set(NEURALNET_BACKEND_SOURCES neuralnet/dummybackend.cpp)
 else()
   message(FATAL_ERROR "Unrecognized backend: " ${USE_BACKEND})
@@ -535,6 +540,58 @@ elseif(USE_BACKEND STREQUAL "EIGEN")
       message(STATUS "Found Eigen3 at ${EIGEN3_INCLUDE_DIRS}")
     endif()
   endif()
+elseif(USE_BACKEND STREQUAL "ONNX")
+  target_compile_definitions(katago PRIVATE USE_ONNX_BACKEND)
+
+  # ONNX Runtime install tree (include/ lib/ bin/). The official prebuilt ORT packages do
+  # NOT ship the OpenVINO execution provider, so for Intel GPU acceleration ORT must be
+  # built from source with --use_openvino GPU (see the project's build notes).
+  set(ONNXRUNTIME_ROOT "" CACHE PATH "Path to ONNX Runtime package root (containing include/, lib/, bin/)")
+  if(NOT IS_DIRECTORY "${ONNXRUNTIME_ROOT}")
+    message(FATAL_ERROR "ONNXRUNTIME_ROOT does not exist: ${ONNXRUNTIME_ROOT}. Set -DONNXRUNTIME_ROOT=<ort install dir>.")
+  endif()
+  set(ONNXRUNTIME_INCLUDE_DIR "${ONNXRUNTIME_ROOT}/include/onnxruntime")
+  if(NOT IS_DIRECTORY "${ONNXRUNTIME_INCLUDE_DIR}")
+    message(FATAL_ERROR "ONNX Runtime include directory not found: ${ONNXRUNTIME_INCLUDE_DIR}")
+  endif()
+  target_include_directories(katago SYSTEM PRIVATE "${ONNXRUNTIME_INCLUDE_DIR}")
+  if(WIN32)
+    set(ONNXRUNTIME_LIB "${ONNXRUNTIME_ROOT}/lib/onnxruntime.lib")
+    file(GLOB ONNXRUNTIME_DLLS "${ONNXRUNTIME_ROOT}/lib/*.dll" "${ONNXRUNTIME_ROOT}/bin/*.dll")
+  else()
+    find_library(ONNXRUNTIME_LIB onnxruntime HINTS "${ONNXRUNTIME_ROOT}/lib" "${ONNXRUNTIME_ROOT}/bin" "${ONNXRUNTIME_ROOT}")
+  endif()
+  if(NOT ONNXRUNTIME_LIB OR ONNXRUNTIME_LIB STREQUAL "ONNXRUNTIME_LIB-NOTFOUND" OR NOT EXISTS "${ONNXRUNTIME_LIB}")
+    message(FATAL_ERROR "Could not find onnxruntime library under ${ONNXRUNTIME_ROOT}. Looked for: ${ONNXRUNTIME_LIB}")
+  endif()
+  target_link_libraries(katago ${ONNXRUNTIME_LIB})
+
+  # The ONNX backend emits an ONNX ModelProto via the same OnnxModelBuilder as the
+  # TensorRT backend and hands the serialized bytes to Ort::Session. Generate onnx.pb.h
+  # from the vendored external/onnx/onnx.proto and link our own protobuf; the handoff to
+  # ORT is serialized bytes, so there is no ABI contact with whatever protobuf lives
+  # inside the ORT DLL. (Protobuf and protoc must be findable by find_package(Protobuf);
+  # for an ORT built from source these live under its _deps/protobuf-build.)
+  find_package(Protobuf REQUIRED)
+  message(STATUS "Found Protobuf version: ${Protobuf_VERSION}")
+  set(ONNX_PROTO_DIR "${CMAKE_CURRENT_SOURCE_DIR}/external/onnx")
+  protobuf_generate_cpp(ONNX_PROTO_SRCS ONNX_PROTO_HDRS "${ONNX_PROTO_DIR}/onnx.proto")
+  set_source_files_properties(${ONNX_PROTO_SRCS} PROPERTIES COMPILE_OPTIONS "-w")
+  target_sources(katago PRIVATE ${ONNX_PROTO_SRCS} neuralnet/onnxmodelbuilder.cpp)
+  target_include_directories(katago SYSTEM PRIVATE ${CMAKE_CURRENT_BINARY_DIR} ${Protobuf_INCLUDE_DIRS})
+  target_link_libraries(katago ${Protobuf_LIBRARIES})
+
+  # Deploy the ORT runtime DLLs next to katago.exe so the build dir is self-contained.
+  # NOTE: OpenVINO's own runtime DLLs are not shipped by ORT and must be copied
+  # separately (see the project's build notes).
+  if(WIN32 AND ONNXRUNTIME_DLLS)
+    foreach(_onnxruntime_dll IN LISTS ONNXRUNTIME_DLLS)
+      add_custom_command(TARGET katago POST_BUILD
+        COMMAND ${CMAKE_COMMAND} -E copy_if_different
+          "${_onnxruntime_dll}"
+          $<TARGET_FILE_DIR:katago>)
+    endforeach()
+  endif()
 endif()
 
 if(USE_BIGGER_BOARDS_EXPENSIVE)

cpp/configs/gtp_example.cfg

@@ -365,6 +365,15 @@ searchFactorWhenWinningThreshold = 0.95
 # if running out of memory, or using multiple GPUs that expect to share work.
 # nnMaxBatchSize = <integer>
 
+# Per-server-thread max batch size overrides (optional, useful for
+# heterogeneous multi-device setups like mixing NPU + GPU + CPU).
+# Each entry limits the batch size for one server thread, allowing you
+# to give the fast device large batches for high utilisation while
+# capping the slow device to small batches so it doesn't starve the
+# fast one. If not set, every thread uses the global nnMaxBatchSize.
+# nnMaxBatchSizeThread0 = 8     # e.g. fast device (GPU/NPU) - large batches
+# nnMaxBatchSizeThread1 = 2     # e.g. slow device (CPU) - tiny batches
+
 # Controls the neural network cache size, which is the primary RAM/memory use.
 # KataGo will cache up to (2 ** nnCacheSizePowerOfTwo) many neural net
 # evaluations in case of transpositions in the tree.
@@ -462,6 +471,59 @@ searchFactorWhenWinningThreshold = 0.95
 # "auto" (default) uses the GEMM only in FP16, where it is slightly faster.
 # cudaUse1x1Matmul = auto
 
+# ------------------------------
+# ONNX Runtime backend settings
+# ------------------------------
+# These only apply when using the ONNX version of KataGo (USE_BACKEND=ONNX).
+# The official prebuilt ONNX Runtime packages do NOT include the OpenVINO
+# execution provider; for Intel GPU (Arc) acceleration, build ORT from source
+# with --use_openvino GPU.
+
+# Execution provider. One of:
+#   cpu (default), openvino, cuda, tensorrt, migraphx, coreml (macOS only).
+# Use "openvino" for Intel Arc/iGPU/NPU.
+# onnxProvider = cpu
+
+# Provider-specific device selection (mostly for cuda / tensorrt / migraphx).
+# onnxDeviceToUse = 0
+# onnxDeviceToUseThread0 = 0
+# onnxDeviceToUseThread1 = 1
+
+# OpenVINO EP options (only used when onnxProvider = openvino):
+# Device type: GPU, CPU, NPU, AUTO:GPU,CPU, MULTI:GPU.0,GPU.1, etc.
+# onnxOpenVINODeviceType = GPU
+# Optional explicit device id (usually unnecessary for a single GPU).
+# onnxOpenVINODeviceId = 0
+
+# Per-thread device type assignment (optional).
+# Overrides onnxOpenVINODeviceType for the specified thread.
+# This allows mixing CPU, GPU, and NPU inference within the same process.
+# onnxOpenVINODeviceTypeThread0 = NPU
+# onnxOpenVINODeviceTypeThread1 = GPU.0
+# onnxOpenVINODeviceTypeThread2 = GPU.1
+# onnxOpenVINODeviceTypeThread3 = CPU
+
+# OpenVINO EP: cache compiled graphs under cwd to skip recompile on restart; unset => full recompile every startup
+# onnxOpenVINOCacheDir = katago_ov_cache
+# Optional precision override: FP16, FP32, ACCURACY
+# onnxOpenVINOPrecision = FP16
+# Optional OpenVINO execution streams / inference threads / priority:
+# onnxOpenVINONumStreams = 1
+# onnxOpenVINONumOfThreads = 1
+# onnxOpenVINOModelPriority = DEFAULT
+
+# Per-device-type EP option overrides (optional).
+# Fine-tune streams, precision etc. per device type (NPU, GPU, CPU).
+# "GPU" matches GPU, GPU.0, GPU.1 and other GPU variants.
+# onnxOpenVINODeviceConfig_NPU_NumStreams = 4
+# onnxOpenVINODeviceConfig_NPU_Precision = FP16
+# onnxOpenVINODeviceConfig_GPU_NumStreams = 2
+# onnxOpenVINODeviceConfig_CPU_NumOfThreads = 2
+
+# Run the trunk block stack channel-last (NHWC) for transformer models.
+# Default false (NCHW). Only takes effect for models with transformer blocks.
+# onnxTransformerNHWC = false
+
 # ------------------------------
 # Metal GPU settings
 # ------------------------------

cpp/main.cpp

@@ -253,6 +253,8 @@ string Version::getKataGoVersionFullInfo() {
   out << "Using OpenCL backend" << endl;
 #elif defined(USE_EIGEN_BACKEND)
   out << "Using Eigen(CPU) backend" << endl;
+#elif defined(USE_ONNX_BACKEND)
+  out << "Using ONNX Runtime backend" << endl;
 #else
   out << "Using dummy backend" << endl;
 #endif
@@ -289,6 +291,8 @@ string Version::getGitRevisionWithBackend() {
   s += "-opencl";
 #elif defined(USE_EIGEN_BACKEND)
   s += "-eigen";
+#elif defined(USE_ONNX_BACKEND)
+  s += "-onnx";
 #else
   s += "-dummy";
 #endif

cpp/neuralnet/nneval.cpp

@@ -66,7 +66,8 @@ NNEvaluator::NNEvaluator(
   bool doRandomize,
   int defaultSymmetry,
   bool disableWarmup_,
-  ConfigParser& cfg
+  ConfigParser& cfg,
+  const vector<int>& maxBatchSizeByServerThr
 )
   :modelName(mName),
    modelFileName(mFileName),
@@ -78,6 +79,9 @@ NNEvaluator::NNEvaluator(
    usingFP16Mode(useFP16Mode),
    numThreads(numThr),
    gpuIdxByServerThread(gpuIdxByServerThr),
+   maxBatchSizeByServerThread(
+     maxBatchSizeByServerThr.empty() ? vector<int>(numThr, maxBatchSz) : maxBatchSizeByServerThr
+   ),
    randSeed(rSeed),
    debugSkipNeuralNet(skipNeuralNet),
    disableWarmup(disableWarmup_),
@@ -117,6 +121,12 @@ NNEvaluator::NNEvaluator(
     throw StringError("maxBatchSize is negative: " + Global::intToString(maxBatchSize));
   if(gpuIdxByServerThread.size() != numThreads)
     throw StringError("gpuIdxByServerThread.size() != numThreads");
+  if(maxBatchSizeByServerThread.size() != numThreads)
+    throw StringError("maxBatchSizeByServerThread.size() != numThreads");
+  for(int threadMaxBatchSize : maxBatchSizeByServerThread) {
+    if(threadMaxBatchSize <= 0 || threadMaxBatchSize > maxBatchSize)
+      throw StringError("Invalid per-server-thread max batch size: " + Global::intToString(threadMaxBatchSize));
+  }
 
   if(logger != NULL) {
     logger->write(
@@ -382,6 +392,7 @@ void NNEvaluator::setNumThreads(const vector<int>& gpuIdxByServerThr) {
     throw StringError("NNEvaluator::setNumThreads called when threads were already running!");
   numThreads = (int)gpuIdxByServerThr.size();
   gpuIdxByServerThread = gpuIdxByServerThr;
+  maxBatchSizeByServerThread.assign(numThreads, maxBatchSize);
 }
 
 void NNEvaluator::spawnServerThreads() {
@@ -566,14 +577,16 @@ void NNEvaluator::serve(
 ) {
   int64_t numBatchesHandledThisThread = 0;
   int64_t numRowsHandledThisThread = 0;
+  testAssert(serverThreadIdx >= 0 && serverThreadIdx < (int)maxBatchSizeByServerThread.size());
+  const int maxBatchSizeForThisThread = maxBatchSizeByServerThread[serverThreadIdx];
 
   ComputeHandle* gpuHandle = NULL;
   if(loadedModel != NULL) {
     gpuHandle = NeuralNet::createComputeHandle(
       computeContext,
       loadedModel,
       logger,
-      maxBatchSize,
+      maxBatchSizeForThisThread,
       requireExactNNLen,
       inputsUseNHWC,
       gpuIdxForThisThread,
@@ -594,14 +607,14 @@ void NNEvaluator::serve(
   }
 
   vector<NNResultBuf*> resultBufs;
-  resultBufs.reserve(maxBatchSize);
+  resultBufs.reserve(maxBatchSizeForThisThread);
 
   vector<NNOutput*> outputBuf;
 
   unique_lock<std::mutex> lock(bufferMutex,std::defer_lock);
   while(true) {
     resultBufs.clear();
-    int desiredBatchSize = std::min(maxBatchSize, currentBatchSize.load(std::memory_order_acquire));
+    int desiredBatchSize = std::min(maxBatchSizeForThisThread, currentBatchSize.load(std::memory_order_acquire));
     bool gotAnything = queryQueue.waitPopUpToN(resultBufs,desiredBatchSize);
     // Queue being closed is a signal that we're done.
     if(!gotAnything)

cpp/neuralnet/nneval.h

@@ -100,7 +100,10 @@ class NNEvaluator {
     int defaultSymmetry,
     bool disableWarmup,
     // Consulted by the compute backend for its own custom options; not stored.
-    ConfigParser& cfg
+    ConfigParser& cfg,
+    // Per-server-thread max batch sizes (index = serverThreadIdx).
+    // Empty = every thread uses the global maxBatchSize.
+    const std::vector<int>& maxBatchSizeByServerThread = std::vector<int>()
   );
   ~NNEvaluator();
 
@@ -224,6 +227,7 @@ class NNEvaluator {
   const enabled_t usingFP16Mode;
   int numThreads;
   std::vector<int> gpuIdxByServerThread;
+  std::vector<int> maxBatchSizeByServerThread;
   const std::string randSeed;
   const bool debugSkipNeuralNet;
   const bool disableWarmup;