From ff3c105fa717bfe0eb076bec8558537d94032689 Mon Sep 17 00:00:00 2001 From: seniorfish Date: Fri, 31 Jul 2026 12:17:35 +0800 Subject: [PATCH 01/23] Add ONNX Runtime neural net backend Add a new USE_BACKEND=ONNX option that runs inference through ONNX Runtime, reusing the existing OnnxModelBuilder graph emitter (shared with the TensorRT backend) so the IO protocol and output decode are identical to TensorRT. Only onnxbackend.cpp is new; the rest is wiring (CMake, setup, version info, config keys, example config). The backend supports any ONNX Runtime execution provider. It is primarily useful for running KataGo on non-NVIDIA accelerators that have an EP, e.g. Intel GPUs/NPUs via the OpenVINO EP. Because the official prebuilt ONNX Runtime packages do not ship those EPs, the Compiling.md section documents building ONNX Runtime from source with the desired provider enabled. onnxmodelbuilder.cpp: declare graph inputs in consumption order (InputSpatial, InputGlobal, InputMask) rather than the previous InputMask-first order. This is purely cosmetic for backends that bind inputs by name (TensorRT), but is required for the OpenVINO EP, which builds its name->index map from declaration order while the ORT runtime feeds the EP kernel inputs in consumption order -- with InputMask first the EP misroutes the mask tensor into the InputSpatial port. --- Compiling.md | 29 ++ cpp/CMakeLists.txt | 61 ++- cpp/configs/gtp_example.cfg | 38 ++ cpp/main.cpp | 4 + cpp/neuralnet/onnxbackend.cpp | 784 +++++++++++++++++++++++++++++ cpp/neuralnet/onnxmodelbuilder.cpp | 28 +- cpp/neuralnet/onnxmodelbuilder.h | 9 +- cpp/program/gtpconfig.cpp | 3 + cpp/program/setup.cpp | 5 +- 9 files changed, 953 insertions(+), 8 deletions(-) create mode 100644 cpp/neuralnet/onnxbackend.cpp diff --git a/Compiling.md b/Compiling.md index abe7de36fc..a0fd1cab59 100644 --- a/Compiling.md +++ b/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= ^ + -DProtobuf_PROTOC_EXECUTABLE= ^ + -DProtobuf_INCLUDE_DIR= ^ + -DProtobuf_LIBRARY= + 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. diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 52ba702f02..5be2381865 100644 --- a/cpp/CMakeLists.txt +++ b/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=.") + 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}" + $) + endforeach() + endif() endif() if(USE_BIGGER_BOARDS_EXPENSIVE) diff --git a/cpp/configs/gtp_example.cfg b/cpp/configs/gtp_example.cfg index 7c3a8f8341..67992c12a6 100644 --- a/cpp/configs/gtp_example.cfg +++ b/cpp/configs/gtp_example.cfg @@ -462,6 +462,44 @@ 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 +# Optional compiled-graph cache dir (speeds up repeated session creation). +# onnxOpenVINOCacheDir = C:\temp\katago_ov_cache +# Optional precision override: FP16, FP32, ACCURACY +# onnxOpenVINOPrecision = FP16 +# Optional OpenVINO execution streams / inference threads / priority: +# onnxOpenVINONumStreams = 1 +# onnxOpenVINONumOfThreads = 1 +# onnxOpenVINOModelPriority = DEFAULT +# Optional NPU fast-compile (may be ignored if unsupported by your ORT build): +# onnxOpenVINOEnableNPUFastCompile = true + +# 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 # ------------------------------ diff --git a/cpp/main.cpp b/cpp/main.cpp index 9bc545ea09..517414b1ab 100644 --- a/cpp/main.cpp +++ b/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 diff --git a/cpp/neuralnet/onnxbackend.cpp b/cpp/neuralnet/onnxbackend.cpp new file mode 100644 index 0000000000..deecab8752 --- /dev/null +++ b/cpp/neuralnet/onnxbackend.cpp @@ -0,0 +1,784 @@ +// ONNX Runtime backend for KataGo. +// +// Loads standard .bin.gz KataGo model files, converts the ModelDesc to a serialized +// ONNX ModelProto via the same OnnxModelBuilder that the TensorRT backend uses, and +// hands the bytes to an Ort::Session. Inference is run through ONNX Runtime with a +// configurable execution provider (CPU, OpenVINO, CUDA, TensorRT, MIGraphX, CoreML) +// selected at runtime via the onnxProvider config key. +// +// The IO tensor protocol is identical to the TensorRT ONNX-emitter path (see +// onnxmodelbuilder.h): four NCHW float32 inputs InputMask / InputSpatial / +// InputGlobal / InputMeta and five NCHW float32 outputs OutputPolicyPass / +// OutputPolicy / OutputValue / OutputScoreValue / OutputOwnership, all raw logits. +// The C++ getOutput below reproduces the TensorRT backend's post-processing exactly +// (per-row optimism blend, inverse-symmetry, version-branched score-value decode) so +// that the same downstream decode path is shared. + +#ifdef USE_ONNX_BACKEND + +#include "../neuralnet/nninterface.h" +#include "../neuralnet/nneval.h" +#include "../neuralnet/nninputs.h" +#include "../neuralnet/modelversion.h" +#include "../neuralnet/onnxmodelbuilder.h" + +#include +#ifdef __APPLE__ +#include +#endif + +#include + +using namespace std; + +//-------------------------------------------------------------- + +struct LoadedModel { + ModelDesc modelDesc; + + LoadedModel(const string& fileName, const string& expectedSha256) { + if(Global::isSuffix(fileName, ".onnx")) + throw StringError( + "ONNX backend: loading a raw .onnx file is not supported by this backend. " + "Feed a standard KataGo .bin.gz model instead (this backend builds the ONNX " + "graph from the model weights internally)."); + ModelDesc::loadFromFileMaybeGZipped(fileName, modelDesc, expectedSha256); + modelDesc.applyScale8ToReduceActivations(); + } + + LoadedModel() = delete; + LoadedModel(const LoadedModel&) = delete; + LoadedModel& operator=(const LoadedModel&) = delete; +}; + +LoadedModel* NeuralNet::loadModelFile(const string& file, const string& expectedSha256) { + return new LoadedModel(file, expectedSha256); +} + +void NeuralNet::freeLoadedModel(LoadedModel* loadedModel) { + delete loadedModel; +} + +const ModelDesc& NeuralNet::getModelDesc(const LoadedModel* loadedModel) { + return loadedModel->modelDesc; +} + +//-------------------------------------------------------------- + +struct ComputeContext { + Ort::Env env; + int nnXLen; + int nnYLen; + bool requireExactNNLenStored; // not used (per-handle), kept for clarity + string providerName; + string openvinoDeviceType; + string openvinoDeviceId; + bool openvinoEnableNPUFastCompile; + string openvinoCacheDir; + // Optional OpenVINO provider options (empty = not passed to ORT) + string openvinoPrecision; // FP16 / FP32 / ACCURACY + string openvinoNumStreams; // 1-8 + string openvinoNumOfThreads; // positive int (infer requests per session) + string openvinoModelPriority; // LOW / MEDIUM / HIGH / DEFAULT + bool transformerNHWC; // run the trunk block stack channel-last (NHWC) + + ComputeContext(int xLen, int yLen) + : env(ORT_LOGGING_LEVEL_WARNING, "KataGoOnnx"), + nnXLen(xLen), + nnYLen(yLen), + requireExactNNLenStored(false), + providerName("cpu"), + openvinoDeviceType("GPU"), + openvinoDeviceId(""), + openvinoEnableNPUFastCompile(false), + openvinoCacheDir(""), + openvinoPrecision(""), + openvinoNumStreams(""), + openvinoNumOfThreads(""), + openvinoModelPriority(""), + transformerNHWC(false) + {} +}; + +ComputeContext* NeuralNet::createComputeContext( + const std::vector& gpuIdxs, + Logger* logger, + int nnXLen, + int nnYLen, + const string& homeDataDirOverride, + enabled_t useFP16Mode, + const LoadedModel* loadedModel, + ConfigParser& cfg +) { + (void)gpuIdxs; + (void)homeDataDirOverride; + // ONNX Runtime and the selected execution provider decide precision internally (e.g. + // OpenVINO/CUDA pick FP16 themselves), so the global useFP16 flag is intentionally ignored. + (void)useFP16Mode; + (void)loadedModel; + + ComputeContext* ctx = new ComputeContext(nnXLen, nnYLen); + + // Provider selection. Default CPU; OpenVINO is the EP used for Intel Arc GPUs. + string providerName = cfg.contains("onnxProvider") ? cfg.getString("onnxProvider") : "cpu"; + ctx->providerName = Global::toLower(providerName); + + // OpenVINO EP options. + ctx->openvinoDeviceType = cfg.contains("onnxOpenVINODeviceType") ? cfg.getString("onnxOpenVINODeviceType") : "GPU"; + ctx->openvinoDeviceId = cfg.contains("onnxOpenVINODeviceId") ? cfg.getString("onnxOpenVINODeviceId") : ""; + if(cfg.contains("onnxOpenVINOEnableNPUFastCompile")) { + string v = Global::toLower(cfg.getString("onnxOpenVINOEnableNPUFastCompile")); + ctx->openvinoEnableNPUFastCompile = (v == "1" || v == "true" || v == "yes" || v == "on"); + } + ctx->openvinoCacheDir = cfg.contains("onnxOpenVINOCacheDir") ? cfg.getString("onnxOpenVINOCacheDir") : ""; + ctx->openvinoPrecision = cfg.contains("onnxOpenVINOPrecision") ? cfg.getString("onnxOpenVINOPrecision") : ""; + ctx->openvinoNumStreams = cfg.contains("onnxOpenVINONumStreams") ? cfg.getString("onnxOpenVINONumStreams") : ""; + ctx->openvinoNumOfThreads = cfg.contains("onnxOpenVINONumOfThreads") ? cfg.getString("onnxOpenVINONumOfThreads") : ""; + ctx->openvinoModelPriority = cfg.contains("onnxOpenVINOModelPriority") ? cfg.getString("onnxOpenVINOModelPriority") : ""; + + // Trunk layout for transformer models (NCHW by default; NHWC only when opted in). + ctx->transformerNHWC = cfg.contains("onnxTransformerNHWC") ? cfg.getBool("onnxTransformerNHWC") : false; + + if(ctx->providerName != "cpu" && ctx->providerName != "openvino" && ctx->providerName != "cuda" && + ctx->providerName != "tensorrt" && ctx->providerName != "migraphx" && ctx->providerName != "coreml") + throw StringError( + "ONNX backend: unknown onnxProvider '" + ctx->providerName + + "', expected one of 'cpu','openvino','cuda','tensorrt','migraphx','coreml'"); + + if(logger != NULL) + logger->write("ONNX backend: creating compute context for " + + Global::intToString(nnXLen) + "x" + Global::intToString(nnYLen) + + " with provider '" + ctx->providerName + "'"); + + return ctx; +} + +void NeuralNet::freeComputeContext(ComputeContext* computeContext) { + delete computeContext; +} + +//-------------------------------------------------------------- + +struct ComputeHandle { + ComputeContext* ctx; + std::unique_ptr session; + int modelVersion; + int numInputChannels; + int numInputGlobalChannels; + int numInputMetaChannels; + int numPolicyChannels; + int numValueChannels; + int numScoreValueChannels; + int numOwnershipChannels; + + // Queried graph input/output names (and raw-char pointer views for Run). + vector inputNames; + vector outputNames; + vector inputNamePtrs; + vector outputNamePtrs; + + ComputeHandle(ComputeContext* context, const LoadedModel& loadedModel, Logger* logger, int deviceIdxForThread, bool requireExactNNLen) + : ctx(context), + modelVersion(loadedModel.modelDesc.modelVersion), + numInputChannels(loadedModel.modelDesc.numInputChannels), + numInputGlobalChannels(loadedModel.modelDesc.numInputGlobalChannels), + numInputMetaChannels(loadedModel.modelDesc.numInputMetaChannels), + numPolicyChannels(loadedModel.modelDesc.numPolicyChannels), + numValueChannels(loadedModel.modelDesc.numValueChannels), + numScoreValueChannels(loadedModel.modelDesc.numScoreValueChannels), + numOwnershipChannels(loadedModel.modelDesc.numOwnershipChannels) + { + if(logger != NULL) + logger->write("ONNX backend: building ONNX graph from model weights..."); + + // Reuse the same ONNX emitter as the TensorRT backend. The serialized ModelProto is + // a standard ONNX graph that Ort::Session can parse directly; the TRT-only FP32 + // node-name lists in the Result are ignored (ORT has no per-node precision API). + OnnxModelBuilder::Result onnxResult = OnnxModelBuilder::build( + loadedModel.modelDesc, ctx->nnXLen, ctx->nnYLen, requireExactNNLen, ctx->transformerNHWC, logger, + // OpenVINO EP misroutes the mask tensor unless graph inputs are declared in consumption + // order (see onnxmodelbuilder.cpp). Other ORT providers bind by name and are unaffected. + ctx->providerName == "openvino"); + const string& onnxBytes = onnxResult.serializedModel; + (void)onnxResult.trunkTipAndHeadNodeNames; + (void)onnxResult.rmsNormNodeNames; + + if(logger != NULL) + logger->write("ONNX backend: ONNX graph built (" + Global::uint64ToString(onnxBytes.size()) + " bytes)"); + + Ort::SessionOptions sessionOpts; + sessionOpts.SetIntraOpNumThreads(1); + + // Select execution provider based on providerName. + const string& provider = ctx->providerName; + if(provider == "coreml") { +#ifdef __APPLE__ + uint32_t coremlFlags = COREML_FLAG_CREATE_ML_PROGRAM; + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CoreML(sessionOpts, coremlFlags)); + if(logger != NULL) + logger->write("ONNX backend: CoreML execution provider enabled (MLProgram mode)"); +#else + throw StringError("ONNX backend: CoreML is only available on Apple platforms"); +#endif + } + else if(provider == "cuda") { + OrtCUDAProviderOptions cudaOpts{}; + cudaOpts.device_id = (unsigned int)(deviceIdxForThread >= 0 ? deviceIdxForThread : 0); + sessionOpts.AppendExecutionProvider_CUDA(cudaOpts); + if(logger != NULL) + logger->write("ONNX backend: CUDA execution provider enabled, device_id=" + Global::intToString((int)cudaOpts.device_id)); + } + else if(provider == "tensorrt") { + OrtTensorRTProviderOptions trtOpts{}; + trtOpts.device_id = (unsigned int)(deviceIdxForThread >= 0 ? deviceIdxForThread : 0); + sessionOpts.AppendExecutionProvider_TensorRT(trtOpts); + if(logger != NULL) + logger->write("ONNX backend: TensorRT execution provider enabled, device_id=" + Global::intToString((int)trtOpts.device_id)); + } + else if(provider == "migraphx") { + OrtMIGraphXProviderOptions migraphxOpts{}; + migraphxOpts.device_id = (unsigned int)(deviceIdxForThread >= 0 ? deviceIdxForThread : 0); + sessionOpts.AppendExecutionProvider_MIGraphX(migraphxOpts); + if(logger != NULL) + logger->write("ONNX backend: MIGraphX execution provider enabled, device_id=" + Global::intToString((int)migraphxOpts.device_id)); + } + else if(provider == "openvino") { + std::unordered_map openvinoOpts; + openvinoOpts["device_type"] = ctx->openvinoDeviceType; + // Fall back to the per-thread device index for OpenVINO device_id when not explicitly set + // (only meaningful when > 0, so a single-default-GPU setup passes no device_id at all). + string deviceId = ctx->openvinoDeviceId; + if(deviceId.empty() && deviceIdxForThread > 0) + deviceId = Global::intToString(deviceIdxForThread); + if(!deviceId.empty()) + openvinoOpts["device_id"] = deviceId; + if(!ctx->openvinoCacheDir.empty()) + openvinoOpts["cache_dir"] = ctx->openvinoCacheDir; + if(!ctx->openvinoPrecision.empty()) + openvinoOpts["precision"] = ctx->openvinoPrecision; + if(!ctx->openvinoNumStreams.empty()) + openvinoOpts["num_streams"] = ctx->openvinoNumStreams; + if(!ctx->openvinoNumOfThreads.empty()) + openvinoOpts["num_of_threads"] = ctx->openvinoNumOfThreads; + if(!ctx->openvinoModelPriority.empty()) + openvinoOpts["model_priority"] = ctx->openvinoModelPriority; + + if(ctx->openvinoEnableNPUFastCompile && logger != NULL) { + logger->write( + "ONNX backend: onnxOpenVINOEnableNPUFastCompile requested, but this ORT build may not " + "accept 'enable_npu_fast_compile'; currently ignoring this option for compatibility." + ); + } + + // Some ORT OpenVINO builds reject optional keys (cache_dir, precision, num_streams, + // num_of_threads, model_priority). Retry with only the core device keys if optional keys + // are rejected, so that setting onnxOpenVINOCacheDir on an EP that doesn't support it + // degrades gracefully instead of crashing. + static const char* optionalKeys[] = { + "cache_dir", "precision", "num_streams", "num_of_threads", "model_priority" + }; + try { + sessionOpts.AppendExecutionProvider_OpenVINO_V2(openvinoOpts); + } + catch(const Ort::Exception& e) { + bool hadOptionalKeys = false; + for(const char* k : optionalKeys) { + if(openvinoOpts.count(k) > 0) { + hadOptionalKeys = true; + break; + } + } + if(!hadOptionalKeys) + throw; + + if(logger != NULL) { + logger->write( + string("ONNX backend: OpenVINO optional provider options rejected, retrying without optional keys. Error: ") + + e.what() + ); + } + for(const char* k : optionalKeys) + openvinoOpts.erase(k); + sessionOpts.AppendExecutionProvider_OpenVINO_V2(openvinoOpts); + } + + if(logger != NULL) { + string devId = openvinoOpts.count("device_id") > 0 ? openvinoOpts["device_id"] : ""; + string extras; + for(const char* k : optionalKeys) { + if(openvinoOpts.count(k) > 0) + extras += string(", ") + k + "=" + openvinoOpts[k]; + } + logger->write( + "ONNX backend: OpenVINO execution provider enabled, device_type=" + ctx->openvinoDeviceType + + (devId.empty() ? "" : (", device_id=" + devId)) + extras + ); + } + } + else if(provider == "cpu" || provider.empty()) { + if(logger != NULL) + logger->write("ONNX backend: using CPU execution provider"); + } + else { + throw StringError("ONNX backend: unknown onnxProvider '" + provider + "'"); + } + + // Create session from in-memory bytes. + session = std::make_unique(ctx->env, onnxBytes.data(), onnxBytes.size(), sessionOpts); + + // Query and store graph input names. + Ort::AllocatorWithDefaultOptions allocator; + size_t numInputs = session->GetInputCount(); + for(size_t i = 0; i < numInputs; i++) { + Ort::AllocatedStringPtr name = session->GetInputNameAllocated(i, allocator); + inputNames.push_back(name.get()); + } + for(auto& n : inputNames) + inputNamePtrs.push_back(n.c_str()); + + // Query and store graph output names. + size_t numOutputs = session->GetOutputCount(); + for(size_t i = 0; i < numOutputs; i++) { + Ort::AllocatedStringPtr name = session->GetOutputNameAllocated(i, allocator); + outputNames.push_back(name.get()); + } + for(auto& n : outputNames) + outputNamePtrs.push_back(n.c_str()); + + if(logger != NULL) { + string inList = "ONNX backend: graph input order:"; + for(size_t i = 0; i < inputNames.size(); i++) + inList += " [" + Global::uint64ToString(i) + "]" + inputNames[i]; + logger->write(inList); + string outList = "ONNX backend: graph output order:"; + for(size_t i = 0; i < outputNames.size(); i++) + outList += " [" + Global::uint64ToString(i) + "]" + outputNames[i]; + logger->write(outList); + logger->write("ONNX backend: session created, inputs=" + Global::uint64ToString(numInputs) + + " outputs=" + Global::uint64ToString(numOutputs)); + } + } + + ComputeHandle() = delete; + ComputeHandle(const ComputeHandle&) = delete; + ComputeHandle& operator=(const ComputeHandle&) = delete; +}; + +ComputeHandle* NeuralNet::createComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + Logger* logger, + int maxBatchSize, + bool requireExactNNLen, + bool inputsUseNHWC, + int gpuIdxForThisThread, + int serverThreadIdx +) { + // ONNX Runtime sessions support dynamic batch sizes; the InputBuffers maxBatchSize + // field still enforces the upper bound at inference time. + (void)maxBatchSize; + if(inputsUseNHWC) + throw StringError("ONNX backend: inputsUseNHWC = true not supported, must use NCHW"); + + if(logger != NULL) { + logger->write("ONNX backend thread " + Global::intToString(serverThreadIdx) + + ": Model version " + Global::intToString(loadedModel->modelDesc.modelVersion)); + logger->write("ONNX backend thread " + Global::intToString(serverThreadIdx) + + ": Model name: " + loadedModel->modelDesc.name + + " (" + loadedModel->modelDesc.getShortInfoString() + ")"); + string deviceInfo = + context->providerName == "openvino" + ? "n/a (use onnxOpenVINODeviceType/onnxOpenVINODeviceId)" + : Global::intToString(gpuIdxForThisThread); + logger->write("ONNX backend thread " + Global::intToString(serverThreadIdx) + + ": provider=" + context->providerName + " deviceIdx=" + deviceInfo); + } + + return new ComputeHandle(context, *loadedModel, logger, gpuIdxForThisThread, requireExactNNLen); +} + +void NeuralNet::freeComputeHandle(ComputeHandle* computeHandle) { + delete computeHandle; +} + +bool NeuralNet::isUsingFP16(const ComputeHandle* handle) { + (void)handle; + // The emitted ONNX graph is fp32; precision is delegated to the execution provider + // (e.g. OpenVINO may downcast internally), so from KataGo's perspective this is fp32. + return false; +} + +bool NeuralNet::setIsWarmup(const ComputeHandle* handle, bool isWarmup) { + (void)handle; + (void)isWarmup; + return false; +} + +//-------------------------------------------------------------- + +struct InputBuffers { + int maxBatchSize; + + size_t singleMaskElts; + size_t singleInputElts; + size_t singleInputGlobalElts; + size_t singleInputMetaElts; + + size_t singlePolicyPassResultElts; + size_t singlePolicyResultElts; + size_t singleValueResultElts; + size_t singleScoreValueResultElts; + size_t singleOwnershipResultElts; + + vector maskInput; + vector spatialInput; + vector globalInput; + vector metaInput; + + InputBuffers(const LoadedModel* loadedModel, int maxBatchSz, int nnXLen, int nnYLen) { + const ModelDesc& m = loadedModel->modelDesc; + + if(nnXLen > NNPos::MAX_BOARD_LEN) + throw StringError( + Global::strprintf("nnXLen (%d) is greater than NNPos::MAX_BOARD_LEN (%d)", nnXLen, NNPos::MAX_BOARD_LEN)); + if(nnYLen > NNPos::MAX_BOARD_LEN) + throw StringError( + Global::strprintf("nnYLen (%d) is greater than NNPos::MAX_BOARD_LEN (%d)", nnYLen, NNPos::MAX_BOARD_LEN)); + + maxBatchSize = maxBatchSz; + singleMaskElts = (size_t)nnXLen * nnYLen; + singleInputElts = (size_t)m.numInputChannels * nnXLen * nnYLen; + singleInputGlobalElts = (size_t)m.numInputGlobalChannels; + singleInputMetaElts = (size_t)m.numInputMetaChannels; + singlePolicyPassResultElts = (size_t)m.numPolicyChannels; + singlePolicyResultElts = (size_t)m.numPolicyChannels * nnXLen * nnYLen; + singleValueResultElts = (size_t)m.numValueChannels; + singleScoreValueResultElts = (size_t)m.numScoreValueChannels; + singleOwnershipResultElts = (size_t)m.numOwnershipChannels * nnXLen * nnYLen; + + testAssert(NNModelVersion::getNumSpatialFeatures(m.modelVersion) == m.numInputChannels); + testAssert(NNModelVersion::getNumGlobalFeatures(m.modelVersion) == m.numInputGlobalChannels); + if(m.numInputMetaChannels > 0) + testAssert(SGFMetadata::METADATA_INPUT_NUM_CHANNELS == m.numInputMetaChannels); + + maskInput.assign(singleMaskElts * maxBatchSize, 0.0f); + spatialInput.assign(singleInputElts * maxBatchSize, 0.0f); + globalInput.assign(singleInputGlobalElts * maxBatchSize, 0.0f); + if(singleInputMetaElts > 0) + metaInput.assign(singleInputMetaElts * maxBatchSize, 0.0f); + } + + InputBuffers() = delete; + InputBuffers(const InputBuffers&) = delete; + InputBuffers& operator=(const InputBuffers&) = delete; +}; + +InputBuffers* NeuralNet::createInputBuffers(const LoadedModel* loadedModel, int maxBatchSize, int nnXLen, int nnYLen) { + return new InputBuffers(loadedModel, maxBatchSize, nnXLen, nnYLen); +} + +void NeuralNet::freeInputBuffers(InputBuffers* inputBuffers) { + delete inputBuffers; +} + +//-------------------------------------------------------------- + +void NeuralNet::globalInitialize() { +} + +void NeuralNet::globalCleanup() { +} + +//-------------------------------------------------------------- + +// Find the index of a name in the graph's name list, matching any of the target alternatives. +static int findNameIndex(const vector& names, std::initializer_list targets) { + for(size_t i = 0; i < names.size(); i++) { + for(const char* t : targets) { + if(names[i] == t) + return (int)i; + } + } + return -1; +} + +void NeuralNet::getOutput( + ComputeHandle* gpuHandle, + InputBuffers* inputBuffers, + int numBatchEltsFilled, + NNResultBuf** inputBufs, + vector& outputs +) { + assert(numBatchEltsFilled <= inputBuffers->maxBatchSize); + assert(numBatchEltsFilled > 0); + + const int batchSize = numBatchEltsFilled; + const int nnXLen = gpuHandle->ctx->nnXLen; + const int nnYLen = gpuHandle->ctx->nnYLen; + const int modelVersion = gpuHandle->modelVersion; + + const int numSpatialFeatures = NNModelVersion::getNumSpatialFeatures(modelVersion); + const int numGlobalFeatures = NNModelVersion::getNumGlobalFeatures(modelVersion); + const int numMetaFeatures = (int)inputBuffers->singleInputMetaElts; + assert(numSpatialFeatures * nnXLen * nnYLen == inputBuffers->singleInputElts); + assert(numGlobalFeatures == inputBuffers->singleInputGlobalElts); + + // Fill host input buffers, mirroring the TensorRT backend exactly: + // - global / meta are straight copies (no symmetry) + // - spatial is symmetry-transformed (NCHW, useNHWC=false) + // - mask = channel 0 of the symmetry-transformed spatial input + for(int nIdx = 0; nIdx < batchSize; nIdx++) { + float* rowMaskInput = inputBuffers->maskInput.data() + inputBuffers->singleMaskElts * nIdx; + float* rowSpatialInput = inputBuffers->spatialInput.data() + inputBuffers->singleInputElts * nIdx; + float* rowGlobalInput = inputBuffers->globalInput.data() + inputBuffers->singleInputGlobalElts * nIdx; + + const float* rowGlobal = inputBufs[nIdx]->rowGlobalBuf.data(); + const float* rowSpatial = inputBufs[nIdx]->rowSpatialBuf.data(); + std::copy(rowGlobal, rowGlobal + numGlobalFeatures, rowGlobalInput); + SymmetryHelpers::copyInputsWithSymmetry( + rowSpatial, rowSpatialInput, 1, nnYLen, nnXLen, numSpatialFeatures, false, inputBufs[nIdx]->symmetry); + std::copy(rowSpatialInput, rowSpatialInput + inputBuffers->singleMaskElts, rowMaskInput); + + if(numMetaFeatures > 0) { + float* rowMetaInput = inputBuffers->metaInput.data() + inputBuffers->singleInputMetaElts * nIdx; + const float* rowMeta = inputBufs[nIdx]->rowMetaBuf.data(); + testAssert(inputBufs[nIdx]->hasRowMeta); + std::copy(rowMeta, rowMeta + numMetaFeatures, rowMetaInput); + } + else { + testAssert(!inputBufs[nIdx]->hasRowMeta); + } + } + + // Build Ort::Value views over the host buffers (CPU memory; the execution provider + // copies to device internally and returns outputs in CPU memory). + Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + + std::array maskShape = {batchSize, 1, nnYLen, nnXLen}; + Ort::Value maskTensor = Ort::Value::CreateTensor( + memInfo, inputBuffers->maskInput.data(), inputBuffers->singleMaskElts * batchSize, + maskShape.data(), maskShape.size()); + + std::array spatialShape = {batchSize, numSpatialFeatures, nnYLen, nnXLen}; + Ort::Value spatialTensor = Ort::Value::CreateTensor( + memInfo, inputBuffers->spatialInput.data(), inputBuffers->singleInputElts * batchSize, + spatialShape.data(), spatialShape.size()); + + std::array globalShape = {batchSize, numGlobalFeatures, 1, 1}; + Ort::Value globalTensor = Ort::Value::CreateTensor( + memInfo, inputBuffers->globalInput.data(), inputBuffers->singleInputGlobalElts * batchSize, + globalShape.data(), globalShape.size()); + + Ort::Value metaTensor(nullptr); + std::array metaShape; + if(numMetaFeatures > 0) { + metaShape = {batchSize, numMetaFeatures, 1, 1}; + metaTensor = Ort::Value::CreateTensor( + memInfo, inputBuffers->metaInput.data(), inputBuffers->singleInputMetaElts * batchSize, + metaShape.data(), metaShape.size()); + } + + // Bind tensors in the graph's declared input order (ORT matches by pointer array + name array). + int maskIdx = findNameIndex(gpuHandle->inputNames, {"InputMask"}); + int spatialIdx = findNameIndex(gpuHandle->inputNames, {"InputSpatial"}); + int globalIdx = findNameIndex(gpuHandle->inputNames, {"InputGlobal"}); + if(maskIdx < 0 || spatialIdx < 0 || globalIdx < 0) + throw StringError("ONNX backend: graph is missing expected inputs InputMask/InputSpatial/InputGlobal"); + int metaIdx = -1; + if(numMetaFeatures > 0) { + metaIdx = findNameIndex(gpuHandle->inputNames, {"InputMeta"}); + if(metaIdx < 0) + throw StringError("ONNX backend: model has metadata channels but the graph has no InputMeta input"); + } + + vector inputTensors; + inputTensors.reserve(gpuHandle->inputNames.size()); + for(size_t i = 0; i < gpuHandle->inputNames.size(); i++) { + if((int)i == maskIdx) + inputTensors.push_back(std::move(maskTensor)); + else if((int)i == spatialIdx) + inputTensors.push_back(std::move(spatialTensor)); + else if((int)i == globalIdx) + inputTensors.push_back(std::move(globalTensor)); + else if((int)i == metaIdx) + inputTensors.push_back(std::move(metaTensor)); + else + throw StringError("ONNX backend: unexpected graph input '" + gpuHandle->inputNames[i] + + "' -- only InputMask/InputSpatial/InputGlobal/InputMeta are supported"); + } + + // Run inference. + auto outputTensors = gpuHandle->session->Run( + Ort::RunOptions{nullptr}, + gpuHandle->inputNamePtrs.data(), + inputTensors.data(), + inputTensors.size(), + gpuHandle->outputNamePtrs.data(), + gpuHandle->outputNamePtrs.size()); + + // Locate outputs by name. + int policyPassIdx = findNameIndex(gpuHandle->outputNames, {"OutputPolicyPass"}); + int policyIdx = findNameIndex(gpuHandle->outputNames, {"OutputPolicy"}); + int valueIdx = findNameIndex(gpuHandle->outputNames, {"OutputValue"}); + int scoreValueIdx = findNameIndex(gpuHandle->outputNames, {"OutputScoreValue"}); + int ownershipIdx = findNameIndex(gpuHandle->outputNames, {"OutputOwnership"}); + if(policyPassIdx < 0 || policyIdx < 0 || valueIdx < 0 || scoreValueIdx < 0 || ownershipIdx < 0) + throw StringError( + "ONNX backend: graph is missing expected outputs " + "(OutputPolicyPass/OutputPolicy/OutputValue/OutputScoreValue/OutputOwnership)"); + + const float* policyPassData = outputTensors[policyPassIdx].GetTensorData(); + const float* policyData = outputTensors[policyIdx].GetTensorData(); + const float* valueData = outputTensors[valueIdx].GetTensorData(); + const float* scoreValueData = outputTensors[scoreValueIdx].GetTensorData(); + const float* ownershipData = outputTensors[ownershipIdx].GetTensorData(); + + assert(policyPassData != nullptr); + assert(policyData != nullptr); + assert(valueData != nullptr); + assert(scoreValueData != nullptr); + assert(ownershipData != nullptr); + assert((int)outputs.size() == batchSize); + + const int numPolicyChannels = (int)inputBuffers->singlePolicyPassResultElts; + assert(inputBuffers->singlePolicyResultElts == (size_t)numPolicyChannels * nnXLen * nnYLen); + const int numValueChannels = (int)inputBuffers->singleValueResultElts; + const int numScoreValueChannels = (int)inputBuffers->singleScoreValueResultElts; + + // Per-row decode, reproducing the TensorRT backend's post-processing exactly. + // Outputs are raw logits; the client applies softmax / tanh / etc. + float policyProbsTmp[NNPos::MAX_NN_POLICY_SIZE]; + + for(int row = 0; row < batchSize; row++) { + NNOutput* output = outputs[row]; + assert(output->nnXLen == nnXLen); + assert(output->nnYLen == nnYLen); + float policyOptimism = (float)inputBufs[row]->policyOptimism; + + // Policy: OutputPolicyPass is [N, numPolicyChannels, 1, 1]; OutputPolicy is [N, numPolicyChannels, H, W]. + { + const float* policyPassSrcBuf = policyPassData + row * numPolicyChannels; + const float* policySrcBuf = policyData + row * numPolicyChannels * nnXLen * nnYLen; + float* policyProbs = output->policyProbs; + + if(numPolicyChannels == 2 || (numPolicyChannels == 4 && modelVersion >= 16)) { + // NCHW: channel 0 = base logits, channel 1 = optimism logits. + for(int i = 0; i < nnXLen * nnYLen; i++) { + float p = policySrcBuf[i]; + float pOpt = policySrcBuf[i + nnXLen * nnYLen]; + policyProbsTmp[i] = p + (pOpt - p) * policyOptimism; + } + SymmetryHelpers::copyOutputsWithSymmetry( + policyProbsTmp, policyProbs, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + policyProbs[nnXLen * nnYLen] = + policyPassSrcBuf[0] + (policyPassSrcBuf[1] - policyPassSrcBuf[0]) * policyOptimism; + } + else { + assert(numPolicyChannels == 1); + SymmetryHelpers::copyOutputsWithSymmetry( + policySrcBuf, policyProbs, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + policyProbs[nnXLen * nnYLen] = policyPassSrcBuf[0]; + } + } + + // Value: [N, 3] raw categorical logits (win/loss/noresult). + { + assert(numValueChannels == 3); + output->whiteWinProb = valueData[row * numValueChannels]; + output->whiteLossProb = valueData[row * numValueChannels + 1]; + output->whiteNoResultProb = valueData[row * numValueChannels + 2]; + } + + // Ownership: [N, 1, H, W] raw; inverse-symmetry back to canonical orientation. + if(output->whiteOwnerMap != NULL) { + assert(inputBuffers->singleOwnershipResultElts == (size_t)nnXLen * nnYLen); + const float* ownershipSrcBuf = ownershipData + row * nnXLen * nnYLen; + SymmetryHelpers::copyOutputsWithSymmetry( + ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + } + + // ScoreValue: [N, numScoreValueChannels] raw, version-dependent channel interpretation. + { + if(modelVersion >= 9) { + assert(numScoreValueChannels == 6); + output->whiteScoreMean = scoreValueData[row * numScoreValueChannels]; + output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1]; + output->whiteLead = scoreValueData[row * numScoreValueChannels + 2]; + output->varTimeLeft = scoreValueData[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = scoreValueData[row * numScoreValueChannels + 4]; + output->shorttermScoreError = scoreValueData[row * numScoreValueChannels + 5]; + } + else if(modelVersion >= 8) { + assert(numScoreValueChannels == 4); + output->whiteScoreMean = scoreValueData[row * numScoreValueChannels]; + output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1]; + output->whiteLead = scoreValueData[row * numScoreValueChannels + 2]; + output->varTimeLeft = scoreValueData[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else if(modelVersion >= 4) { + assert(numScoreValueChannels == 2); + output->whiteScoreMean = scoreValueData[row * numScoreValueChannels]; + output->whiteScoreMeanSq = scoreValueData[row * numScoreValueChannels + 1]; + output->whiteLead = output->whiteScoreMean; + output->varTimeLeft = 0; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else if(modelVersion >= 3) { + assert(numScoreValueChannels == 1); + output->whiteScoreMean = scoreValueData[row * numScoreValueChannels]; + output->whiteScoreMeanSq = output->whiteScoreMean * output->whiteScoreMean; + output->whiteLead = output->whiteScoreMean; + output->varTimeLeft = 0; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } + else { + ASSERT_UNREACHABLE; + } + } + } +} + +void NeuralNet::printDevices() { + cout << "ONNX backend: device enumeration is execution-provider-specific." << endl; + cout << "Set onnxProvider (e.g. 'openvino') plus provider-specific options" << endl; + cout << "(onnxOpenVINODeviceType, onnxOpenVINODeviceId, ...) in the config." << endl; +} + +//-------------------------------------------------------------- +// The layer-level test entry points are not implemented for this backend. Returning +// false tells the test harness this configuration is unsupported (not a failure). +// (The TensorRT backend likewise returns false for all of these.) + +bool NeuralNet::testEvaluateConv( + const ConvLayerDesc*, int, int, int, bool, bool, + const std::vector&, std::vector& +) { + return false; +} + +bool NeuralNet::testEvaluateBatchNorm( + const BatchNormLayerDesc*, int, int, int, bool, bool, + const std::vector&, const std::vector&, std::vector& +) { + return false; +} + +bool NeuralNet::testEvaluateResidualBlock( + const ResidualBlockDesc*, int, int, int, bool, bool, + const std::vector&, const std::vector&, std::vector& +) { + return false; +} + +bool NeuralNet::testEvaluateGlobalPoolingResidualBlock( + const GlobalPoolingResidualBlockDesc*, int, int, int, bool, bool, + const std::vector&, const std::vector&, std::vector& +) { + return false; +} + +#endif // USE_ONNX_BACKEND diff --git a/cpp/neuralnet/onnxmodelbuilder.cpp b/cpp/neuralnet/onnxmodelbuilder.cpp index 6469237bdb..10799b2468 100644 --- a/cpp/neuralnet/onnxmodelbuilder.cpp +++ b/cpp/neuralnet/onnxmodelbuilder.cpp @@ -812,7 +812,8 @@ Result build( int nnYLen, bool requireExactNNLen, bool transformerNHWC, - Logger* logger + Logger* logger, + bool alignInputsToConsumptionOrder ) { if(desc.metaEncoderVersion > 0) throw StringError("OnnxModelBuilder: SGF metadata encoder not yet supported"); @@ -872,9 +873,28 @@ Result build( shape->add_dim()->set_dim_value(1); shape->add_dim()->set_dim_value(1); }; - addInput("InputMask", 1); - addInput("InputSpatial", numInputChannels); - addInputNC11("InputGlobal", numInputGlobalChannels); + // Declaration order matters for the OpenVINO execution provider under ONNX Runtime: ORT + // feeds the EP kernel's input ports in consumption order (the order the graph first + // references each input), but the EP's own name->index map is built from this declaration + // order. With master's default order (InputMask first), the two disagree and the EP + // misroutes the [N,1,H,W] mask tensor into the InputSpatial port, failing at runtime: + // "can't handle input tensor ...:InputSpatial, because model input (shape=[?,22,19,19]) + // and tensor (shape=[1,1,19,19]) are incompatible" + // (measured on ORT 1.29 + OpenVINO 2026.2, Intel Arc B580). Declaring inputs in + // consumption order -- InputSpatial (trunk conv), InputMask (first applyMask), InputGlobal + // (gpool merge) -- aligns the two and is verified to fix it (~60-74 visits/s recovered). + // This is purely cosmetic for backends that bind inputs by name (TensorRT, CUDA, CoreML), + // so it is opt-in: only the ONNX backend requests it, and only for the OpenVINO provider. + if(alignInputsToConsumptionOrder) { + addInput("InputSpatial", numInputChannels); + addInputNC11("InputGlobal", numInputGlobalChannels); + addInput("InputMask", 1); + } + else { + addInput("InputMask", 1); + addInput("InputSpatial", numInputChannels); + addInputNC11("InputGlobal", numInputGlobalChannels); + } // ---- Mask-derived features ---- if(!requireExactNNLen) { diff --git a/cpp/neuralnet/onnxmodelbuilder.h b/cpp/neuralnet/onnxmodelbuilder.h index 10d3819150..12dccbc922 100644 --- a/cpp/neuralnet/onnxmodelbuilder.h +++ b/cpp/neuralnet/onnxmodelbuilder.h @@ -32,13 +32,20 @@ namespace OnnxModelBuilder { }; // Build a serialized ONNX ModelProto for the given model. + // alignInputsToConsumptionOrder reorders the graph's declared inputs to match the order + // in which the graph first consumes them (InputSpatial, InputGlobal, InputMask) instead of + // the default declaration order (InputMask, InputSpatial, InputGlobal). This is a no-op for + // backends that bind inputs by name (TensorRT) but is required for the OpenVINO execution + // provider under ONNX Runtime, which builds its name->index map from declaration order while + // the runtime feeds the EP kernel input ports in consumption order. See onnxmodelbuilder.cpp. Result build( const ModelDesc& desc, int nnXLen, int nnYLen, bool requireExactNNLen, bool transformerNHWC, - Logger* logger + Logger* logger, + bool alignInputsToConsumptionOrder = false ); } diff --git a/cpp/program/gtpconfig.cpp b/cpp/program/gtpconfig.cpp index 3fa8651e5b..741c8cb2e3 100644 --- a/cpp/program/gtpconfig.cpp +++ b/cpp/program/gtpconfig.cpp @@ -537,6 +537,9 @@ string GTPConfig::makeConfig( #endif #ifdef USE_OPENCL_BACKEND replacement += "openclDeviceToUseThread" + Global::intToString(i) + " = " + Global::intToString(deviceIdxs[i]) + "\n"; +#endif +#ifdef USE_ONNX_BACKEND + replacement += "onnxDeviceToUseThread" + Global::intToString(i) + " = " + Global::intToString(deviceIdxs[i]) + "\n"; #endif } replace("$$MULTIPLE_GPUS", replacement); diff --git a/cpp/program/setup.cpp b/cpp/program/setup.cpp index bf7950315e..04c04dcff2 100644 --- a/cpp/program/setup.cpp +++ b/cpp/program/setup.cpp @@ -21,6 +21,7 @@ std::vector Setup::getBackendPrefixes() { prefixes.push_back("metal"); prefixes.push_back("opencl"); prefixes.push_back("eigen"); + prefixes.push_back("onnx"); prefixes.push_back("dummybackend"); return prefixes; } @@ -89,6 +90,8 @@ vector Setup::initializeNNEvaluators( string backendPrefix = "opencl"; #elif defined(USE_EIGEN_BACKEND) string backendPrefix = "eigen"; + #elif defined(USE_ONNX_BACKEND) + string backendPrefix = "onnx"; #else string backendPrefix = "dummybackend"; #endif @@ -142,7 +145,7 @@ vector Setup::initializeNNEvaluators( requireExactNNLen = cfg.getBool("requireMaxBoardSize"); } - bool inputsUseNHWC = backendPrefix == "opencl" || backendPrefix == "trt" || backendPrefix == "metal" ? false : true; + bool inputsUseNHWC = backendPrefix == "opencl" || backendPrefix == "trt" || backendPrefix == "metal" || backendPrefix == "onnx" ? false : true; if(cfg.contains(backendPrefix+"InputsUseNHWC"+idxStr)) inputsUseNHWC = cfg.getBool(backendPrefix+"InputsUseNHWC"+idxStr); else if(cfg.contains("inputsUseNHWC"+idxStr)) From 46d18427fcac9407ba628c1778c0560ffb80dd51 Mon Sep 17 00:00:00 2001 From: seniorfish Date: Sat, 1 Aug 2026 11:51:08 +0800 Subject: [PATCH 02/23] @ onnxmodelbuilder: use Pow(x,2.0) instead of Mul(x,x) for RMSNorm square OpenVINO RMSFusion (rms_fusion.cpp:38) matches Power(x, const(2)) but not Mul(x,x). Without this, all 66 RMSNorm nodes in the b11c768 transformer run as unfused ReduceMean->Sqrt->Div chains. Benchmark (Arc B580, NHWC, numStreams=2, 10 threads, 800 visits): Before: 347.49 visits/s 296.04 nnEvals/s After: 416.52 visits/s 353.14 nnEvals/s (+19.9%) Also add KATAGO_DUMP_ONNX env-var debug aid to dump the serialized ONNX model before session creation. Co-Authored-By: Claude @ --- cpp/neuralnet/onnxbackend.cpp | 19 +++++++++++++++++++ cpp/neuralnet/onnxmodelbuilder.cpp | 11 +++++++++-- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/cpp/neuralnet/onnxbackend.cpp b/cpp/neuralnet/onnxbackend.cpp index deecab8752..66aae61bf8 100644 --- a/cpp/neuralnet/onnxbackend.cpp +++ b/cpp/neuralnet/onnxbackend.cpp @@ -28,6 +28,8 @@ #endif #include +#include +#include using namespace std; @@ -206,6 +208,23 @@ struct ComputeHandle { if(logger != NULL) logger->write("ONNX backend: ONNX graph built (" + Global::uint64ToString(onnxBytes.size()) + " bytes)"); + // Dump the ONNX model to a file when KATAGO_DUMP_ONNX is set (debug aid). + { + const char* dumpPath = getenv("KATAGO_DUMP_ONNX"); + if(dumpPath != nullptr && dumpPath[0] != '\0') { + ofstream dumpFile(dumpPath, ios::binary); + if(dumpFile.is_open()) { + dumpFile.write(onnxBytes.data(), (streamsize)onnxBytes.size()); + dumpFile.close(); + if(logger != NULL) + logger->write(string("ONNX backend: dumped ONNX model to ") + dumpPath + + " (" + Global::uint64ToString(onnxBytes.size()) + " bytes)"); + } else if(logger != NULL) { + logger->write(string("ONNX backend: WARNING - could not open dump path ") + dumpPath); + } + } + } + Ort::SessionOptions sessionOpts; sessionOpts.SetIntraOpNumThreads(1); diff --git a/cpp/neuralnet/onnxmodelbuilder.cpp b/cpp/neuralnet/onnxmodelbuilder.cpp index 10799b2468..1bcd90fc88 100644 --- a/cpp/neuralnet/onnxmodelbuilder.cpp +++ b/cpp/neuralnet/onnxmodelbuilder.cpp @@ -405,7 +405,12 @@ struct Builder { testAssert((int)desc.weight.size() == C); int rmsStart = graph->node_size(); // meanSq over channels (axis 1), keepdims -> [N,1,H,W] - string sq = addNode("Mul", {input, input}, uniq(desc.name + "/sq"), desc.name + "/sq"); + // Pow(x, 2.0) instead of Mul(x, x): OpenVINO RMSFusion matcher requires + // Power(x, const(2)) (rms_fusion.cpp:38) and silently skips Mul(x,x). + // Without this, all 66 RMSNorm nodes run as unfused ReduceMean→Sqrt→Div + // chains, costing ~0.5–1.0 ms/frame on GPU. + string twoName = addScalarInitializer(uniq(desc.name + "/pow2"), 2.0f); + string sq = addNode("Pow", {input, twoName}, uniq(desc.name + "/sq"), desc.name + "/sq"); string meanSq = addNode("ReduceMean", {sq, addInt64Initializer(uniq(desc.name + "/axC"), {1})}, uniq(desc.name + "/meansq"), desc.name + "/meansq"); { onnx::NodeProto* n = lastNode(); onnx::AttributeProto* a = addAttr(n, "keepdims"); a->set_type(onnx::AttributeProto::INT); a->set_i(1); } @@ -440,7 +445,9 @@ struct Builder { int C = desc.numChannels; testAssert((int)desc.weight.size() == C); int rmsStart = graph->node_size(); - string sq = addNode("Mul", {input, input}, uniq(desc.name + "/sq"), desc.name + "/sq"); + // Pow(x, 2.0) instead of Mul(x, x): see NCHW variant above. + string twoName = addScalarInitializer(uniq(desc.name + "/pow2"), 2.0f); + string sq = addNode("Pow", {input, twoName}, uniq(desc.name + "/sq"), desc.name + "/sq"); string meanSq = addNode("ReduceMean", {sq, addInt64Initializer(uniq(desc.name + "/axC"), {3})}, // C is axis 3 of [N,H,W,C] uniq(desc.name + "/meansq"), desc.name + "/meansq"); { onnx::NodeProto* n = lastNode(); onnx::AttributeProto* a = addAttr(n, "keepdims"); a->set_type(onnx::AttributeProto::INT); a->set_i(1); } From 4f86efa3261f28d1cafa740e6f196c1a19eb576b Mon Sep 17 00:00:00 2001 From: seniorfish Date: Sat, 1 Aug 2026 22:25:57 +0800 Subject: [PATCH 03/23] Skip applyScale8ToReduceActivations in ONNX backend to avoid MISH_SCALE8 subgraphs hurting NPU performance --- cpp/neuralnet/onnxbackend.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/cpp/neuralnet/onnxbackend.cpp b/cpp/neuralnet/onnxbackend.cpp index 66aae61bf8..18c2a4cb5b 100644 --- a/cpp/neuralnet/onnxbackend.cpp +++ b/cpp/neuralnet/onnxbackend.cpp @@ -45,7 +45,11 @@ struct LoadedModel { "Feed a standard KataGo .bin.gz model instead (this backend builds the ONNX " "graph from the model weights internally)."); ModelDesc::loadFromFileMaybeGZipped(fileName, modelDesc, expectedSha256); - modelDesc.applyScale8ToReduceActivations(); + // Skip applyScale8ToReduceActivations() for ONNX backend: + // NPU/ONNX Runtime execution providers don't benefit from the fp16 + // dynamic-range workaround; removing this avoids MISH_SCALE8 subgraphs + // that block operator fusion and cost ~25% extra ops per activation. + // modelDesc.applyScale8ToReduceActivations(); } LoadedModel() = delete; From 815378d849488a26b336e9cbfd79dbd7640d9183 Mon Sep 17 00:00:00 2001 From: seniorfish Date: Sun, 2 Aug 2026 23:26:34 +0800 Subject: [PATCH 04/23] @ Add per-thread OpenVINO device type and batch size config for multi-device inference - onnxOpenVINODeviceTypeThread: assign different device types (CPU/GPU/NPU) to individual server threads, enabling simultaneous heterogeneous inference - onnxOpenVINODeviceConfig__