From fe46687835786191e68107647acabe1e0d068ca0 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 12 Aug 2026 07:28:52 -0700 Subject: [PATCH] Add MIGraphX backend for AMD GPUs (ROCm) Adds a third GPU backend targeting AMD via MIGraphX, ROCm's graph compiler, alongside the existing CUDA/TensorRT and OpenCL paths. The backend reuses the ONNX ModelProto that OnnxModelBuilder already emits for the TensorRT path and hands the identical bytes to MIGraphX's parse_onnx_buffer, so network construction is shared and onnxmodelbuilder.cpp is untouched. Measured on MI300X (gfx942), ROCm 7.2.0, b18c384nbt, 19x19, FP16, with both backends built and benchmarked in a single job on one node: visits=3200 OpenCL (tuned) 1564.38 MIGraphX 4599.00 2.94x visits=800 OpenCL (tuned) 1548.94 MIGraphX 4239.44 2.74x KataGo's OpenCL tuner reports canUseFP16TensorCores=0 on gfx942, so the OpenCL path never issues MFMA; MIGraphX routes convolutions through rocMLIR/MIOpen, which do. Validated with runnnonmanyposestest over 254 positions across all 5 nets in cpp/tests/models: FP32 agrees with OpenCL to 2.6e-11..5.6e-10 policyProbSquerr, and MIGraphX's FP16 is 2.2x-51x closer to the FP32 reference than OpenCL's FP16. Notes: - MIGraphX compiles one static shape, so the program is compiled at maxBatchSize and short batches are zero-padded. Padding rows get an all-ones mask, since the graph divides by maskSum for masked means and a zero mask row is a division by zero that propagates NaN into real rows. - Graph outputs are exposed as positional main:#output_N parameters while inputs keep their ONNX names; the mapping is asserted against declared shapes so an emitter reordering fails loudly rather than silently swapping tensors. - migraphxTransformerNHWC defaults to false, unlike TensorRT's trtTransformerNHWC. The channel-last trunk produces wrong policy output on transformer nets under MIGraphX (policySqErr 136 vs 6e-10) while value heads stay correct; root cause is still open, so the safe NCHW default ships. - Protobuf must be linked statically with -Wl,--exclude-libs,ALL, because libmigraphx_onnx exports its bundled protobuf as weak symbols that a shared libprotobuf would preempt. Documented in Compiling.md. --- Compiling.md | 28 +- cpp/CMakeLists.txt | 65 ++- cpp/main.cpp | 4 + cpp/neuralnet/migraphxbackend.cpp | 900 ++++++++++++++++++++++++++++++ cpp/program/setup.cpp | 4 +- cpp/tests/testcommon.cpp | 9 + 6 files changed, 1007 insertions(+), 3 deletions(-) create mode 100644 cpp/neuralnet/migraphxbackend.cpp diff --git a/Compiling.md b/Compiling.md index abe7de36fc..aa05985469 100644 --- a/Compiling.md +++ b/Compiling.md @@ -33,6 +33,7 @@ As also mentioned in the instructions below but repeated here for visibility, if * If using the OpenCL backend, a modern GPU that supports OpenCL 1.2 or greater, or else something like [this](https://software.intel.com/en-us/opencl-sdk) for CPU. But if using CPU, Eigen should be better. * If using the CUDA backend, CUDA 11 or later and a compatible version of CUDNN based on your CUDA version (https://developer.nvidia.com/cuda-toolkit) (https://developer.nvidia.com/cudnn) and a GPU capable of supporting them. * If using the TensorRT backend, in addition to a compatible CUDA Toolkit (https://developer.nvidia.com/cuda-toolkit), you also need TensorRT (https://developer.nvidia.com/tensorrt) that is at least version 8.5. + * If using the MIGraphX backend (AMD GPUs), ROCm with MIGraphX and its headers - with Debian packages this is `migraphx` and `migraphx-dev`. Set `-DROCM_PATH=...` if ROCm is not at `/opt/rocm`. You also need the **static** protobuf library `libprotobuf.a` (Debian: `libprotobuf-dev`); see the note below for why a shared libprotobuf does not work. * If using the Eigen backend, Eigen3. With Debian packages, (i.e. apt or apt-get), this should be `libeigen3-dev`. * zlib, libzip. With Debian packages (i.e. apt or apt-get), these should be `zlib1g-dev`, `libzip-dev`. * If you want to do self-play training and research, probably Google perftools `libgoogle-perftools-dev` for TCMalloc or some other better malloc implementation. For unknown reasons, the allocation pattern in self-play with large numbers of threads and parallel games causes a lot of memory fragmentation under glibc malloc that will eventually run your machine out of memory, but better mallocs handle it fine. @@ -41,7 +42,7 @@ As also mentioned in the instructions below but repeated here for visibility, if * `git clone https://github.com/lightvector/KataGo.git` * Compile using CMake and make in the cpp directory: * `cd KataGo/cpp` - * `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=CUDA` or `cmake . -DUSE_BACKEND=TENSORRT` or `cmake . -DUSE_BACKEND=EIGEN` depending on which backend you want. + * `cmake . -DUSE_BACKEND=OPENCL` or `cmake . -DUSE_BACKEND=CUDA` or `cmake . -DUSE_BACKEND=TENSORRT` or `cmake . -DUSE_BACKEND=MIGRAPHX` or `cmake . -DUSE_BACKEND=EIGEN` depending on which backend you want. * Specify also `-DUSE_TCMALLOC=1` if using TCMalloc. * Compiling will also call git commands to embed the git hash into the compiled executable, specify also `-DNO_GIT_REVISION=1` to disable it if this is causing issues for you. * Specify `-DUSE_AVX2=1` to also compile Eigen with AVX2 and FMA support, which will make it incompatible with old CPUs but much faster. (If you want to go further, you can also add `-DCMAKE_CXX_FLAGS='-march=native'` which will specialize to precisely your machine's CPU, but the exe might not run on other machines at all). @@ -54,6 +55,31 @@ As also mentioned in the instructions below but repeated here for visibility, if * 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`). +### Note on the MIGraphX backend and protobuf + +The MIGraphX backend links protobuf **statically** and builds with `-Wl,--exclude-libs,ALL`. This is +required, not a preference, and CMake will stop with an error if `libprotobuf.a` is not found. + +`libmigraphx_onnx` bundles its own copy of protobuf and exports roughly 160 protobuf symbols as +*weak* template instantiations. If KataGo links a shared `libprotobuf`, the dynamic linker resolves +those weak symbols to whichever definition is global — KataGo's — so MIGraphX's ONNX parser ends up +running against a protobuf whose object layout it was not compiled against. The failure appears at +model load as an abort inside protobuf rather than as a link error: + +``` +CHECK failed: (total_size_) > (0) ... google/protobuf/repeated_field.h +``` + +Linking the static archive and marking its symbols local keeps the two copies apart. You can confirm +a correct build exports none: + +``` +nm -D --defined-only ./katago | grep -c protobuf # must print 0 +``` + +The TensorRT backend does not need this because `nvonnxparser` statically links its own protobuf and +the only thing crossing the boundary is a serialized byte buffer. + ## Windows * TLDR: * Building from source on Windows is actually a bit tricky, depending on what version you're building, there's not necessarily a super-fast way. diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index fb8bb130fa..5e749295ef 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 MIGRAPHX OPENCL EIGEN METAL) set(USE_TCMALLOC 0 CACHE BOOL "Use TCMalloc") set(NO_GIT_REVISION 0 CACHE BOOL "Disable embedding the git revision into the compiled exe") @@ -114,6 +114,11 @@ elseif(USE_BACKEND STREQUAL "TENSORRT") elseif(USE_CACHE_TENSORRT_PLAN AND BUILD_DISTRIBUTED) message(FATAL_ERROR "Combining USE_CACHE_TENSORRT_PLAN with BUILD_DISTRIBUTED is not supported - it would consume excessive disk space and might worsen performance every time models are updated. Use only one at a time in a given build of KataGo.") endif() +elseif(USE_BACKEND STREQUAL "MIGRAPHX") + message(STATUS "-DUSE_BACKEND=MIGRAPHX, using AMD ROCm MIGraphX backend.") + set(NEURALNET_BACKEND_SOURCES + neuralnet/migraphxbackend.cpp + ) elseif(USE_BACKEND STREQUAL "METAL") message(STATUS "-DUSE_BACKEND=METAL, using Metal backend with hybrid MPSGraph + CoreML execution.") if(NOT "${CMAKE_GENERATOR}" STREQUAL "Ninja") @@ -476,6 +481,64 @@ elseif(USE_BACKEND STREQUAL "TENSORRT") # CMake package config (e.g. vcpkg), the variable can resolve to the DLL itself rather # than the import lib, and it also omits protobuf's own dependencies such as abseil. target_link_libraries(katago ${TENSORRT_ONNXPARSER_LIBRARY} protobuf::libprotobuf) +elseif(USE_BACKEND STREQUAL "MIGRAPHX") + target_compile_definitions(katago PRIVATE USE_MIGRAPHX_BACKEND) + + # ROCm ships MIGraphX and HIP under the same prefix; ROCM_PATH lets a user point at a + # non-default or side-by-side install (e.g. /opt/rocm-6.4.1). + if(NOT DEFINED ROCM_PATH) + if(DEFINED ENV{ROCM_PATH}) + set(ROCM_PATH $ENV{ROCM_PATH}) + else() + set(ROCM_PATH "/opt/rocm") + endif() + endif() + list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH} ${ROCM_PATH}/hip) + + find_package(hip REQUIRED) + + find_path(MIGRAPHX_INCLUDE_DIR migraphx/migraphx.hpp HINTS ${ROCM_PATH} PATH_SUFFIXES include) + if(NOT MIGRAPHX_INCLUDE_DIR) + message(FATAL_ERROR "${ColorBoldRed} migraphx/migraphx.hpp was NOT found. Install migraphx-dev, or set ROCM_PATH to your ROCm install. ${ColorReset}") + endif() + # The C++ header migraphx.hpp is a header-only wrapper over the C API in libmigraphx_c, so that + # is the only MIGraphX library we need to link. + find_library(MIGRAPHX_C_LIBRARY NAMES migraphx_c HINTS ${ROCM_PATH} PATH_SUFFIXES lib lib64) + if(NOT MIGRAPHX_C_LIBRARY) + message(FATAL_ERROR "${ColorBoldRed} libmigraphx_c was NOT found. Install migraphx, or set ROCM_PATH to your ROCm install. ${ColorReset}") + endif() + + # Like the TensorRT backend, this backend builds its network by emitting an ONNX ModelProto and + # handing the serialized bytes to the inference engine's ONNX parser, so it needs the same + # vendored ONNX schema compiled with protoc. + 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") + # protoc-generated code is not ours to lint; silence its warnings to keep build output readable. + set_source_files_properties(${ONNX_PROTO_SRCS} PROPERTIES COMPILE_OPTIONS "-w") + target_sources(katago PRIVATE ${ONNX_PROTO_SRCS} neuralnet/onnxmodelbuilder.cpp) + # Generated onnx.pb.h lands in the build dir; let backend code include it. + target_include_directories(katago SYSTEM PRIVATE ${CMAKE_CURRENT_BINARY_DIR} ${Protobuf_INCLUDE_DIRS} ${MIGRAPHX_INCLUDE_DIR}) + + # Protobuf must be linked STATICALLY and its symbols kept out of the dynamic symbol table. + # + # libmigraphx_onnx bundles its own protobuf and exports ~160 protobuf symbols as *weak* template + # instantiations. If KataGo also pulls in a shared libprotobuf, the dynamic linker resolves those + # weak symbols to whichever copy is global — ours — and MIGraphX's parser then runs against a + # protobuf whose object layout it was not compiled for. That fails at parse time with + # "CHECK failed: (total_size_) > (0)" inside repeated_field.h. + # + # Linking the static archive with -Wl,--exclude-libs makes every protobuf symbol we pull in local + # to the katago binary, so MIGraphX resolves its bundled copy and the two never interact. This is + # the same isolation the TensorRT backend gets for free (nvonnxparser statically links its own + # protobuf and the handoff is serialized bytes). + find_library(PROTOBUF_STATIC_LIBRARY NAMES libprotobuf.a HINTS ${Protobuf_LIBRARY_DIRS} /usr/lib/x86_64-linux-gnu) + if(NOT PROTOBUF_STATIC_LIBRARY) + message(FATAL_ERROR "${ColorBoldRed} libprotobuf.a (static) was NOT found, but the MIGraphX backend requires it to avoid a protobuf symbol collision with libmigraphx_onnx. Install libprotobuf-dev. ${ColorReset}") + endif() + target_link_libraries(katago ${MIGRAPHX_C_LIBRARY} hip::host ${PROTOBUF_STATIC_LIBRARY}) + target_link_options(katago PRIVATE -Wl,--exclude-libs,ALL) elseif(USE_BACKEND STREQUAL "METAL") target_compile_definitions(katago PRIVATE USE_METAL_BACKEND) target_link_libraries(katago KataGoSwift katagocoreml diff --git a/cpp/main.cpp b/cpp/main.cpp index f9c95e09a6..f7d1b43718 100644 --- a/cpp/main.cpp +++ b/cpp/main.cpp @@ -247,6 +247,8 @@ string Version::getKataGoVersionFullInfo() { #endif #elif defined(USE_TENSORRT_BACKEND) out << "Using TensorRT backend" << endl; +#elif defined(USE_MIGRAPHX_BACKEND) + out << "Using MIGraphX(ROCm) backend" << endl; #elif defined(USE_METAL_BACKEND) out << "Using Metal backend" << endl; #elif defined(USE_OPENCL_BACKEND) @@ -283,6 +285,8 @@ string Version::getGitRevisionWithBackend() { s += "-cuda"; #elif defined(USE_TENSORRT_BACKEND) s += "-trt"; +#elif defined(USE_MIGRAPHX_BACKEND) + s += "-migraphx"; #elif defined(USE_METAL_BACKEND) s += "-metal"; #elif defined(USE_OPENCL_BACKEND) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp new file mode 100644 index 0000000000..5e3643403b --- /dev/null +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -0,0 +1,900 @@ +#ifdef USE_MIGRAPHX_BACKEND + +#include +#include + +#include +#include +#include +#include +#include + +#include "../core/fileutils.h" +#include "../core/makedir.h" +#include "../core/sha2.h" +#include "../core/test.h" +#include "../dataio/homedata.h" +#include "../neuralnet/desc.h" +#include "../neuralnet/modelversion.h" +#include "../neuralnet/nneval.h" +#include "../neuralnet/nninputs.h" +#include "../neuralnet/nninterface.h" +#include "../neuralnet/onnxmodelbuilder.h" + +using namespace std; + +// AMD ROCm backend for KataGo, built on MIGraphX. +// +// This is the AMD analogue of the TensorRT backend and it deliberately reuses that backend's +// network-construction path: OnnxModelBuilder emits a self-contained ONNX ModelProto (weights +// baked in as initializers) with RAW-head outputs, and MIGraphX parses/compiles it into a GPU +// program. Because the emitted graph is identical to the one TensorRT consumes, the getOutput +// decode below is the same decode the TensorRT backend does, and the two backends agree +// numerically up to precision. +// +// Two MIGraphX specifics drive the design here: +// +// 1. MIGraphX compiles for one static shape. There is no TensorRT-style optimization profile with +// a dynamic batch dimension, so the program is compiled at exactly maxBatchSize and smaller +// batches are run by zero-padding up to maxBatchSize. MCTS batches are near-full in practice, +// and a fixed shape lets MIGraphX pick the best kernels and fuse aggressively. +// +// 2. Manual device buffers (set_offload_copy(false)). With offload copy MIGraphX would allocate +// and copy every input and output on each eval; instead we hipMalloc each parameter once and +// hand MIGraphX raw device pointers, so the steady-state eval does only the H2D copies of the +// inputs that actually changed and the D2H copies of the outputs. + +static void checkHipError(const hipError_t status, const char* opName, const char* file, const char* func, int line) { + if(status != hipSuccess) + throw StringError( + string("HIP Error, for ") + opName + " file " + file + ", func " + func + ", line " + Global::intToString(line) + + ", error " + hipGetErrorString(status)); +} +#define HIP_ERR(opName, x) \ + { checkHipError((x), opName, __FILE__, #x, __LINE__); } + +void NeuralNet::globalInitialize() { + // Nothing to do, MIGraphX and HIP initialize lazily. +} + +void NeuralNet::globalCleanup() { + (void)hipDeviceReset(); +} + +struct ComputeContext { + int nnXLen; + int nnYLen; + enabled_t useFP16Mode; + string homeDataDirOverride; + bool transformerNHWC; // ONNX emitter: run transformer blocks channel-last + string dumpDebugModelToDir; + bool useExhaustiveTune; // MIGraphX exhaustive_tune: slower compile, faster kernels +}; + +ComputeContext* NeuralNet::createComputeContext( + const vector& gpuIdxs, + Logger* logger, + int nnXLen, + int nnYLen, + const string& homeDataDirOverride, + enabled_t useFP16Mode, + const LoadedModel* loadedModel, + ConfigParser& cfg +) { + (void)gpuIdxs; + (void)logger; + + ComputeContext* context = new ComputeContext(); + context->nnXLen = nnXLen; + context->nnYLen = nnYLen; + context->useFP16Mode = useFP16Mode; + context->homeDataDirOverride = homeDataDirOverride; + // Mirrors the TensorRT backend's trtTransformerNHWC, but defaults to FALSE here, unlike + // TensorRT which defaults it to true. + // + // The channel-last trunk produces wrong POLICY output under MIGraphX on transformer models + // while the value heads stay correct. Measured against the OpenCL backend over KataGo's own + // runnnonmanyposestest (254 positions), FP32: + // + // model NHWC=true NHWC=false + // b7c96h3tfrs-test5-cnorm policySqErr 136.1 policySqErr 6.0e-10 + // b7c96h6kv3qk32v16tflrs-fson-bnh policySqErr 125.4 policySqErr 1.4e-10 + // + // Every board position on every test position is affected, with the logits collapsing toward + // a near-flat distribution, so this is a wrong computation rather than a layout permutation. + // Root cause is still open (it is either MIGraphX's lowering of an op the channel-last path + // emits, or an emitter assumption that only holds for TensorRT); until that is resolved the + // safe default is the NCHW trunk, which is correct on every model tested. + // + // Convnets never take this path at all: the emitter only goes channel-last when the model + // actually has transformer blocks. + context->transformerNHWC = + (cfg.contains("migraphxTransformerNHWC") ? cfg.getBool("migraphxTransformerNHWC") : false) && + NeuralNet::getModelDesc(loadedModel).hasAnyTransformerBlocks(); + context->dumpDebugModelToDir = + cfg.contains("migraphxDumpDebugModelToDir") ? cfg.getString("migraphxDumpDebugModelToDir") : ""; + // Exhaustive tuning searches more kernel candidates (notably for the trunk convolutions) at + // compile time. It costs minutes per compile, so it is off unless asked for. + context->useExhaustiveTune = + cfg.contains("migraphxExhaustiveTune") ? cfg.getBool("migraphxExhaustiveTune") : false; + return context; +} + +void NeuralNet::freeComputeContext(ComputeContext* computeContext) { + delete computeContext; +} + +struct LoadedModel { + ModelDesc modelDesc; + + LoadedModel(const string& fileName, const string& expectedSha256) { + 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; +} + +// MIGraphX compilation is not thread-safe against itself in all ROCm versions, and KataGo creates +// one ComputeHandle per server thread, all of which compile the same model at startup. Serialize +// compiles so that N threads do not race inside the compiler. +static mutex compileMutex; + +struct ComputeHandle { + ComputeContext* ctx; + + bool usingFP16; + int maxBatchSize; + int modelVersion; + bool hasInputMeta; + + // All work for this handle is ordered on one non-default stream: the input H2D copies, the + // program itself (via run_async), and the output D2H copies. Using the default stream instead + // would not order correctly against MIGraphX, which runs on its own internal stream. + hipStream_t stream; + + // hipGraph capture was prototyped and set aside. This path is GPU-bound, not launch-bound: + // instrumenting the eval measured 10.294 ms blocked in hipStreamSynchronize against 0.028 ms + // of host-side output decode per batch (avgRows 60.8), i.e. the host is 0.3% of the time. + // Collapsing the ~12 driver calls per eval into one graph launch cannot beat that 0.3%. + migraphx::program prog; + migraphx::program_parameters params; + // Device allocations for every program parameter and output, keyed by name. Owned here. + map buffers; + map bufferBytes; + map bufferRowElts; + // Output parameter names, in the order MIGraphX returns them from eval(). + vector outputNames; + + ComputeHandle( + Logger* logger, + ComputeContext* context, + const LoadedModel* loadedModel, + int maxBatchSz, + bool requireExactNNLen + ) { + ctx = context; + maxBatchSize = maxBatchSz; + modelVersion = loadedModel->modelDesc.modelVersion; + hasInputMeta = loadedModel->modelDesc.numInputMetaChannels > 0; + + HIP_ERR("ComputeHandle", hipStreamCreate(&stream)); + + const ModelDesc& desc = loadedModel->modelDesc; + + // Emit the same ONNX graph the TensorRT backend builds. Weights are baked in as initializers, + // so the returned bytes are fully self-contained. + OnnxModelBuilder::Result onnxResult = + OnnxModelBuilder::build(desc, ctx->nnXLen, ctx->nnYLen, requireExactNNLen, ctx->transformerNHWC, logger); + const string& onnxBytes = onnxResult.serializedModel; + + if(!ctx->dumpDebugModelToDir.empty()) { + MakeDir::make(ctx->dumpDebugModelToDir); + string onnxPath = ctx->dumpDebugModelToDir + "/model_" + Global::intToString(ctx->nnXLen) + "x" + + Global::intToString(ctx->nnYLen) + "_bs" + Global::intToString(maxBatchSize) + ".onnx"; + ofstream dumpOut; + FileUtils::open(dumpOut, onnxPath, ios::binary); + dumpOut.write(onnxBytes.data(), (std::streamsize)onnxBytes.size()); + dumpOut.close(); + if(logger != NULL) + logger->write("MIGraphX backend: dumped emitted ONNX to " + onnxPath); + } + + // MIGraphX compiles a static shape, so pin the batch dimension of every input to maxBatchSize. + // The ONNX emitter declares batch as a dynamic dim; set_input_parameter_shape fixes it. + migraphx::onnx_options onnxOptions; + const size_t bs = (size_t)maxBatchSize; + onnxOptions.set_input_parameter_shape("InputMask", {bs, 1, (size_t)ctx->nnYLen, (size_t)ctx->nnXLen}); + onnxOptions.set_input_parameter_shape( + "InputSpatial", {bs, (size_t)desc.numInputChannels, (size_t)ctx->nnYLen, (size_t)ctx->nnXLen}); + onnxOptions.set_input_parameter_shape( + "InputGlobal", {bs, (size_t)desc.numInputGlobalChannels, 1, 1}); + if(hasInputMeta) + onnxOptions.set_input_parameter_shape("InputMeta", {bs, (size_t)desc.numInputMetaChannels, 1, 1}); + + { + lock_guard lock(compileMutex); + + prog = migraphx::parse_onnx_buffer(onnxBytes, onnxOptions); + + usingFP16 = false; + if(ctx->useFP16Mode == enabled_t::True || ctx->useFP16Mode == enabled_t::Auto) { + // quantize_fp16 converts convolutions and dots to FP16 while leaving the reductions that + // feed RMSNorm and the policy/value heads in FP32, which is the same split the TensorRT + // backend enforces via per-layer setPrecision. All CDNA parts have fast FP16 so Auto + // enables it, matching the TensorRT backend's platformHasFastFp16 behavior. + migraphx::quantize_fp16(prog); + usingFP16 = true; + } + + migraphx::compile_options options; + // Manage device memory ourselves; see the note at the top of this file. + options.set_offload_copy(false); + options.set_fast_math(true); + options.set_exhaustive_tune_flag(ctx->useExhaustiveTune); + prog.compile(migraphx::target("gpu"), options); + } + + // Allocate a device buffer for every program parameter. This covers the graph inputs, the + // graph outputs (MIGraphX exposes each output as an "outputName" parameter when offload copy + // is off), and the internal scratch parameter. + migraphx::program_parameter_shapes paramShapes = prog.get_parameter_shapes(); + // names() hands back pointers into MIGraphX-owned storage; copy them into strings we own. + vector paramNames; + for(const char* n: paramShapes.names()) + paramNames.emplace_back(n); + for(const string& name: paramNames) { + migraphx::shape s = paramShapes[name.c_str()]; + size_t bytes = s.bytes(); + void* devPtr = nullptr; + HIP_ERR("ComputeHandle", hipMalloc(&devPtr, bytes)); + HIP_ERR("ComputeHandle", hipMemset(devPtr, 0, bytes)); + buffers[name] = devPtr; + bufferBytes[name] = bytes; + + // Row elements: elements per batch element. The scratch parameter has no batch dim, so guard. + vector lens = s.lengths(); + size_t rowElts = 1; + if(lens.size() >= 1 && lens[0] == (size_t)maxBatchSize) { + for(size_t i = 1; i < lens.size(); i++) + rowElts *= lens[i]; + } else { + rowElts = s.elements(); + } + bufferRowElts[name] = rowElts; + + params.add(name.c_str(), migraphx::argument(s, devPtr)); + } + + // Inputs are addressable by their ONNX names directly. + for(const char* n: {"InputMask", "InputSpatial", "InputGlobal"}) + aliasName[n] = n; + if(hasInputMeta) + aliasName["InputMeta"] = "InputMeta"; + + // Outputs are positional. OnnxModelBuilder declares them in this fixed order (see the + // markOutput calls in onnxmodelbuilder.cpp), so index i corresponds to outputOrder[i]. + static const char* outputOrder[] = { + "OutputPolicyPass", "OutputPolicy", "OutputValue", "OutputScoreValue", "OutputOwnership"}; + const size_t numOutputs = sizeof(outputOrder) / sizeof(outputOrder[0]); + for(size_t i = 0; i < numOutputs; i++) { + string param = "main:#output_" + Global::uint64ToString((uint64_t)i); + if(buffers.find(param) == buffers.end()) + throw StringError( + "MIGraphX backend: expected output parameter " + param + " for " + outputOrder[i] + + " but the compiled program does not have it. MIGraphX's output parameter naming may have " + "changed; the program has these parameters: " + [&] { + string all; + for(const auto& kv: buffers) all += kv.first + " "; + return all; + }()); + aliasName[outputOrder[i]] = param; + } + + // Sanity-check the positional mapping against the shapes the model actually declares, so a + // reordering in the emitter surfaces here rather than as silently swapped policy/value data. + auto expectRowElts = [&](const char* name, size_t expected) { + size_t actual = bufferRowElts.at(aliasName.at(name)); + if(actual != expected) + throw StringError(Global::strprintf( + "MIGraphX backend: output %s mapped to %s has %llu elts per row, expected %llu — the " + "ONNX graph output order does not match this backend's assumed order", + name, aliasName.at(name).c_str(), (unsigned long long)actual, (unsigned long long)expected)); + }; + const size_t area = (size_t)ctx->nnXLen * ctx->nnYLen; + expectRowElts("OutputPolicyPass", (size_t)desc.numPolicyChannels); + expectRowElts("OutputPolicy", (size_t)desc.numPolicyChannels * area); + expectRowElts("OutputValue", (size_t)desc.numValueChannels); + expectRowElts("OutputScoreValue", (size_t)desc.numScoreValueChannels); + expectRowElts("OutputOwnership", (size_t)desc.numOwnershipChannels * area); + + if(logger != NULL) { + logger->write( + "MIGraphX backend: compiled model at batch size " + Global::intToString(maxBatchSize) + + " board " + Global::intToString(ctx->nnXLen) + "x" + Global::intToString(ctx->nnYLen) + + " FP16 = " + Global::boolToString(usingFP16)); + } + } + + ~ComputeHandle() { + // Destructors must not throw, so free errors are swallowed rather than routed through HIP_ERR. + (void)hipStreamSynchronize(stream); + for(auto& kv: buffers) { + (void)hipFree(kv.second); + } + (void)hipStreamDestroy(stream); + } + + ComputeHandle() = delete; + ComputeHandle(const ComputeHandle&) = delete; + ComputeHandle& operator=(const ComputeHandle&) = delete; + + // Inputs keep their ONNX names as parameter names, but MIGraphX does NOT: graph outputs become + // positional parameters "main:#output_0", "main:#output_1", ... in graph-declaration order. + // aliasName maps the ONNX tensor name the rest of this file uses onto the actual parameter name. + map aliasName; + + const string& resolveName(const char* name) const { + auto it = aliasName.find(name); + if(it != aliasName.end()) + return it->second; + throw StringError(Global::strprintf("MIGraphX ComputeHandle: unknown tensor name %s", name)); + } + + void* getBuffer(const char* name) const { + return buffers.at(resolveName(name)); + } + + size_t getBufferBytes(const char* name) const { + return bufferBytes.at(resolveName(name)); + } + + size_t getBufferRowElts(const char* name) const { + return bufferRowElts.at(resolveName(name)); + } +}; + +ComputeHandle* NeuralNet::createComputeHandle( + ComputeContext* context, + const LoadedModel* loadedModel, + Logger* logger, + int maxBatchSize, + bool requireExactNNLen, + bool inputsUseNHWC, + int gpuIdxForThisThread, + int serverThreadIdx +) { + if(inputsUseNHWC) { + throw StringError("MIGraphX backend: inputsUseNHWC = false required, other configurations not supported"); + } + + if(gpuIdxForThisThread == -1) + gpuIdxForThisThread = 0; + HIP_ERR("createComputeHandle", hipSetDevice(gpuIdxForThisThread)); + + hipDeviceProp_t prop; + HIP_ERR("createComputeHandle", hipGetDeviceProperties(&prop, gpuIdxForThisThread)); + + if(logger != NULL) { + logger->write( + "MIGraphX backend thread " + Global::intToString(serverThreadIdx) + ": Found GPU " + string(prop.name) + + " (" + string(prop.gcnArchName) + ") memory " + Global::uint64ToString(prop.totalGlobalMem)); + logger->write( + "MIGraphX backend thread " + Global::intToString(serverThreadIdx) + ": Initializing (may take a long time)"); + } + + auto handle = new ComputeHandle(logger, context, loadedModel, maxBatchSize, requireExactNNLen); + + if(logger != NULL) { + logger->write( + "MIGraphX backend thread " + Global::intToString(serverThreadIdx) + ": Model version " + + Global::intToString(loadedModel->modelDesc.modelVersion) + + " useFP16 = " + Global::boolToString(handle->usingFP16)); + logger->write( + "MIGraphX backend thread " + Global::intToString(serverThreadIdx) + + ": Model name: " + loadedModel->modelDesc.name + + " (" + loadedModel->modelDesc.getShortInfoString() + ")"); + } + + return handle; +} + +void NeuralNet::freeComputeHandle(ComputeHandle* gpuHandle) { + delete gpuHandle; +} + +bool NeuralNet::isUsingFP16(const ComputeHandle* gpuHandle) { + return gpuHandle->usingFP16; +} + +bool NeuralNet::setIsWarmup(const ComputeHandle* gpuHandle, bool isWarmup) { + (void)gpuHandle; + (void)isWarmup; + return false; +} + +void NeuralNet::printDevices() { + int numDevices = 0; + HIP_ERR("printDevices", hipGetDeviceCount(&numDevices)); + for(int i = 0; i < numDevices; i++) { + hipDeviceProp_t prop; + HIP_ERR("printDevices", hipGetDeviceProperties(&prop, i)); + cout << "Found GPU device " << i << ": " << prop.name << " (" << prop.gcnArchName << ")" << endl; + } +} + +struct InputBuffers { + int maxBatchSize; + + size_t singleMaskElts; + size_t singleMaskBytes; + size_t singleInputElts; + size_t singleInputBytes; + size_t singleInputGlobalElts; + size_t singleInputGlobalBytes; + size_t singleInputMetaElts; + size_t singleInputMetaBytes; + size_t singlePolicyPassResultElts; + size_t singlePolicyPassResultBytes; + size_t singlePolicyResultElts; + size_t singlePolicyResultBytes; + size_t singleValueResultElts; + size_t singleValueResultBytes; + size_t singleScoreValueResultElts; + size_t singleScoreValueResultBytes; + size_t singleOwnershipResultElts; + size_t singleOwnershipResultBytes; + + size_t inputMaskBufferBytes; + size_t inputSpatialBufferBytes; + size_t inputGlobalBufferBytes; + size_t inputMetaBufferBytes; + size_t policyPassResultBufferBytes; + size_t policyResultBufferBytes; + size_t valueResultBufferBytes; + size_t scoreValueResultBufferBytes; + size_t ownershipResultBufferBytes; + + // Host staging buffers. Allocated as pinned memory so the H2D/D2H copies run on the DMA engines + // rather than through a pageable-memory bounce buffer; at MCTS batch sizes these copies are + // frequent enough that the difference is measurable. + float* maskInputs; + float* spatialInputs; + float* globalInputs; + float* metaInputs; + float* policyPassResults; + float* policyResults; + float* valueResults; + float* scoreValueResults; + float* ownershipResults; + + // All-ones mask rows used to pad a short batch up to maxBatchSize. See the note in getOutput: + // an all-zero mask row divides by zero in the graph's masked-mean ops. Sized lazily. + std::vector paddingMaskOnes; + + 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; + singleMaskBytes = singleMaskElts * sizeof(float); + singleInputElts = (size_t)m.numInputChannels * nnXLen * nnYLen; + singleInputBytes = singleInputElts * sizeof(float); + singleInputGlobalElts = m.numInputGlobalChannels; + singleInputGlobalBytes = singleInputGlobalElts * sizeof(float); + singleInputMetaElts = m.numInputMetaChannels; + singleInputMetaBytes = singleInputMetaElts * sizeof(float); + singlePolicyPassResultElts = (size_t)m.numPolicyChannels; + singlePolicyPassResultBytes = singlePolicyPassResultElts * sizeof(float); + singlePolicyResultElts = (size_t)m.numPolicyChannels * nnXLen * nnYLen; + singlePolicyResultBytes = singlePolicyResultElts * sizeof(float); + singleValueResultElts = m.numValueChannels; + singleValueResultBytes = singleValueResultElts * sizeof(float); + singleScoreValueResultElts = m.numScoreValueChannels; + singleScoreValueResultBytes = singleScoreValueResultElts * sizeof(float); + singleOwnershipResultElts = (size_t)m.numOwnershipChannels * nnXLen * nnYLen; + singleOwnershipResultBytes = singleOwnershipResultElts * sizeof(float); + + 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); + } + + inputMaskBufferBytes = maxBatchSize * singleMaskBytes; + inputSpatialBufferBytes = maxBatchSize * singleInputBytes; + inputGlobalBufferBytes = maxBatchSize * singleInputGlobalBytes; + inputMetaBufferBytes = maxBatchSize * singleInputMetaBytes; + policyPassResultBufferBytes = maxBatchSize * singlePolicyPassResultBytes; + policyResultBufferBytes = maxBatchSize * singlePolicyResultBytes; + valueResultBufferBytes = maxBatchSize * singleValueResultBytes; + scoreValueResultBufferBytes = maxBatchSize * singleScoreValueResultBytes; + ownershipResultBufferBytes = maxBatchSize * singleOwnershipResultBytes; + + auto allocHost = [](float** ptr, size_t bytes) { + if(bytes == 0) { + *ptr = nullptr; + return; + } + HIP_ERR("InputBuffers", hipHostMalloc((void**)ptr, bytes, hipHostMallocDefault)); + memset(*ptr, 0, bytes); + }; + allocHost(&maskInputs, inputMaskBufferBytes); + allocHost(&spatialInputs, inputSpatialBufferBytes); + allocHost(&globalInputs, inputGlobalBufferBytes); + allocHost(&metaInputs, inputMetaBufferBytes); + allocHost(&policyPassResults, policyPassResultBufferBytes); + allocHost(&policyResults, policyResultBufferBytes); + allocHost(&valueResults, valueResultBufferBytes); + allocHost(&scoreValueResults, scoreValueResultBufferBytes); + allocHost(&ownershipResults, ownershipResultBufferBytes); + } + + ~InputBuffers() { + for(float* p: {maskInputs, spatialInputs, globalInputs, metaInputs, policyPassResults, + policyResults, valueResults, scoreValueResults, ownershipResults}) { + if(p != nullptr) + (void)hipHostFree(p); + } + } + + 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::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 = inputBuffers->singleInputMetaElts; + assert((size_t)numSpatialFeatures * nnXLen * nnYLen == inputBuffers->singleInputElts); + assert(numGlobalFeatures == inputBuffers->singleInputGlobalElts); + + for(int nIdx = 0; nIdx < batchSize; nIdx++) { + float* rowMaskInput = &inputBuffers->maskInputs[inputBuffers->singleMaskElts * nIdx]; + float* rowSpatialInput = &inputBuffers->spatialInputs[inputBuffers->singleInputElts * nIdx]; + float* rowGlobalInput = &inputBuffers->globalInputs[inputBuffers->singleInputGlobalElts * nIdx]; + float* rowMetaInput = &inputBuffers->metaInputs[inputBuffers->singleInputMetaElts * nIdx]; + + const float* rowGlobal = inputBufs[nIdx]->rowGlobalBuf.data(); + const float* rowSpatial = inputBufs[nIdx]->rowSpatialBuf.data(); + const float* rowMeta = inputBufs[nIdx]->rowMetaBuf.data(); + const bool hasRowMeta = inputBufs[nIdx]->hasRowMeta; + std::copy(rowGlobal, rowGlobal + numGlobalFeatures, rowGlobalInput); + if(numMetaFeatures > 0) { + testAssert(rowMeta != NULL); + testAssert(hasRowMeta); + std::copy(rowMeta, rowMeta + numMetaFeatures, rowMetaInput); + } else { + testAssert(!hasRowMeta); + } + SymmetryHelpers::copyInputsWithSymmetry( + rowSpatial, rowSpatialInput, 1, nnYLen, nnXLen, numSpatialFeatures, false, inputBufs[nIdx]->symmetry); + std::copy(rowSpatialInput, rowSpatialInput + inputBuffers->singleMaskElts, rowMaskInput); + } + + assert(inputBuffers->singleMaskElts == gpuHandle->getBufferRowElts("InputMask")); + assert(inputBuffers->singleInputElts == gpuHandle->getBufferRowElts("InputSpatial")); + assert(inputBuffers->singleInputGlobalElts == gpuHandle->getBufferRowElts("InputGlobal")); + if(numMetaFeatures > 0) + assert(inputBuffers->singleInputMetaElts == gpuHandle->getBufferRowElts("InputMeta")); + assert(inputBuffers->singlePolicyPassResultElts == gpuHandle->getBufferRowElts("OutputPolicyPass")); + assert(inputBuffers->singlePolicyResultElts == gpuHandle->getBufferRowElts("OutputPolicy")); + assert(inputBuffers->singleValueResultElts == gpuHandle->getBufferRowElts("OutputValue")); + assert(inputBuffers->singleScoreValueResultElts == gpuHandle->getBufferRowElts("OutputScoreValue")); + assert(inputBuffers->singleOwnershipResultElts == gpuHandle->getBufferRowElts("OutputOwnership")); + + const int numPolicyChannels = inputBuffers->singlePolicyPassResultElts; + assert(inputBuffers->singlePolicyResultElts == (size_t)numPolicyChannels * nnXLen * nnYLen); + + // The program is compiled for exactly maxBatchSize, so only the first batchSize rows are copied + // in and read back; the padding rows' outputs are ignored. + // + // Padding rows must NOT be left as all-zero. When requireExactNNLen is false the emitted graph + // takes masked means as Div(ReduceSum(x), maskSum), where maskSum is the per-row count of + // on-board cells. An all-zero mask row makes that a 0/0 division, so the padding rows produce + // NaN/Inf rather than harmless garbage. + // + // Give every padding row a fully on-board mask (all ones) so maskSum == H*W. The rows then + // compute finite values from zero spatial input and are discarded. + // + // Note: this does NOT fix the transformer policy discrepancy — that was the original hypothesis + // and it was disproved (the error was bit-identical afterwards, because the single-position test + // path never pads at all). This guard matters for the MCTS path, where short batches are real. + // Re-padded on every call rather than cached: a larger batch overwrites this region with real + // data, so a later smaller batch would otherwise inherit stale rows. The copy is one contiguous + // memcpy of (maxBatchSize-batchSize) mask rows and is negligible next to the forward pass. + hipStream_t stream = gpuHandle->stream; + + if(batchSize < inputBuffers->maxBatchSize) { + const int padRows = inputBuffers->maxBatchSize - batchSize; + if(inputBuffers->paddingMaskOnes.size() != inputBuffers->singleMaskElts * (size_t)padRows) + inputBuffers->paddingMaskOnes.assign(inputBuffers->singleMaskElts * (size_t)padRows, 1.0f); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + (char*)gpuHandle->getBuffer("InputMask") + inputBuffers->singleMaskBytes * batchSize, + inputBuffers->paddingMaskOnes.data(), inputBuffers->singleMaskBytes * (size_t)padRows, + hipMemcpyHostToDevice, stream)); + } + HIP_ERR( + "getOutput", + hipMemcpyAsync( + gpuHandle->getBuffer("InputMask"), inputBuffers->maskInputs, + inputBuffers->singleMaskBytes * batchSize, hipMemcpyHostToDevice, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + gpuHandle->getBuffer("InputSpatial"), inputBuffers->spatialInputs, + inputBuffers->singleInputBytes * batchSize, hipMemcpyHostToDevice, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + gpuHandle->getBuffer("InputGlobal"), inputBuffers->globalInputs, + inputBuffers->singleInputGlobalBytes * batchSize, hipMemcpyHostToDevice, stream)); + if(numMetaFeatures > 0) { + HIP_ERR( + "getOutput", + hipMemcpyAsync( + gpuHandle->getBuffer("InputMeta"), inputBuffers->metaInputs, + inputBuffers->singleInputMetaBytes * batchSize, hipMemcpyHostToDevice, stream)); + } + + // run_async rather than eval: eval() runs on MIGraphX's own internal stream, which is not + // ordered against the copies above, so the program could read inputs before they land. + gpuHandle->prog.run_async(gpuHandle->params, stream); + + HIP_ERR( + "getOutput", + hipMemcpyAsync( + inputBuffers->policyPassResults, gpuHandle->getBuffer("OutputPolicyPass"), + inputBuffers->singlePolicyPassResultBytes * batchSize, hipMemcpyDeviceToHost, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + inputBuffers->policyResults, gpuHandle->getBuffer("OutputPolicy"), + inputBuffers->singlePolicyResultBytes * batchSize, hipMemcpyDeviceToHost, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + inputBuffers->valueResults, gpuHandle->getBuffer("OutputValue"), + inputBuffers->singleValueResultBytes * batchSize, hipMemcpyDeviceToHost, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + inputBuffers->scoreValueResults, gpuHandle->getBuffer("OutputScoreValue"), + inputBuffers->singleScoreValueResultBytes * batchSize, hipMemcpyDeviceToHost, stream)); + HIP_ERR( + "getOutput", + hipMemcpyAsync( + inputBuffers->ownershipResults, gpuHandle->getBuffer("OutputOwnership"), + inputBuffers->singleOwnershipResultBytes * batchSize, hipMemcpyDeviceToHost, stream)); + + // One sync per eval, after all the D2H copies are queued, rather than an implicit sync per copy. + HIP_ERR("getOutput", hipStreamSynchronize(stream)); + + assert(outputs.size() == batchSize); + + 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; + + const float* policyPassSrcBuf = &inputBuffers->policyPassResults[row * inputBuffers->singlePolicyPassResultElts]; + const float* policySrcBuf = &inputBuffers->policyResults[row * inputBuffers->singlePolicyResultElts]; + float* policyProbs = output->policyProbs; + + // These are in logits, the client does the postprocessing to turn them into + // policy probabilities and white game outcome probabilities + // Also we don't fill in the nnHash here either + // Handle version >= 12 policy optimism + if(numPolicyChannels == 2 || (numPolicyChannels == 4 && modelVersion >= 16)) { + // MIGraphX outputs are NCHW, same as TensorRT + 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]; + } + + int numValueChannels = inputBuffers->singleValueResultElts; + assert(numValueChannels == 3); + output->whiteWinProb = inputBuffers->valueResults[row * numValueChannels]; + output->whiteLossProb = inputBuffers->valueResults[row * numValueChannels + 1]; + output->whiteNoResultProb = inputBuffers->valueResults[row * numValueChannels + 2]; + + // As above, these are NOT actually from white's perspective, but rather the player to move. + // As usual the client does the postprocessing. + if(output->whiteOwnerMap != NULL) { + const float* ownershipSrcBuf = &inputBuffers->ownershipResults[row * nnXLen * nnYLen]; + assert(inputBuffers->singleOwnershipResultElts == (size_t)nnXLen * nnYLen); + SymmetryHelpers::copyOutputsWithSymmetry( + ownershipSrcBuf, output->whiteOwnerMap, 1, nnYLen, nnXLen, inputBufs[row]->symmetry); + } + + int numScoreValueChannels = inputBuffers->singleScoreValueResultElts; + if(modelVersion >= 9) { + assert(numScoreValueChannels == 6); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[row * numScoreValueChannels + 1]; + output->whiteLead = inputBuffers->scoreValueResults[row * numScoreValueChannels + 2]; + output->varTimeLeft = inputBuffers->scoreValueResults[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = inputBuffers->scoreValueResults[row * numScoreValueChannels + 4]; + output->shorttermScoreError = inputBuffers->scoreValueResults[row * numScoreValueChannels + 5]; + } else if(modelVersion >= 8) { + assert(numScoreValueChannels == 4); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[row * numScoreValueChannels + 1]; + output->whiteLead = inputBuffers->scoreValueResults[row * numScoreValueChannels + 2]; + output->varTimeLeft = inputBuffers->scoreValueResults[row * numScoreValueChannels + 3]; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } else if(modelVersion >= 4) { + assert(numScoreValueChannels == 2); + output->whiteScoreMean = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + output->whiteScoreMeanSq = inputBuffers->scoreValueResults[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 = inputBuffers->scoreValueResults[row * numScoreValueChannels]; + // Version 3 neural nets don't have any second moment output, implicitly already folding it in, so we just use the + // mean squared + output->whiteScoreMeanSq = output->whiteScoreMean * output->whiteScoreMean; + output->whiteLead = output->whiteScoreMean; + output->varTimeLeft = 0; + output->shorttermWinlossError = 0; + output->shorttermScoreError = 0; + } else { + ASSERT_UNREACHABLE; + } + } +} + +// These per-layer test entry points exist for the CUDA/Eigen backends which build the net layer by +// layer. This backend hands a whole ONNX graph to MIGraphX and has no per-layer handles, so like +// the TensorRT backend it declines all of them. +bool NeuralNet::testEvaluateConv( + const ConvLayerDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)outputBuffer; + return false; +} + +// Mask should be in 'NHW' format (no "C" channel). +bool NeuralNet::testEvaluateBatchNorm( + const BatchNormLayerDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +bool NeuralNet::testEvaluateResidualBlock( + const ResidualBlockDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +bool NeuralNet::testEvaluateGlobalPoolingResidualBlock( + const GlobalPoolingResidualBlockDesc* desc, + int batchSize, + int nnXLen, + int nnYLen, + bool useFP16, + bool useNHWC, + const vector& inputBuffer, + const vector& maskBuffer, + vector& outputBuffer) { + (void)desc; + (void)batchSize; + (void)nnXLen; + (void)nnYLen; + (void)useFP16; + (void)useNHWC; + (void)inputBuffer; + (void)maskBuffer; + (void)outputBuffer; + return false; +} + +#endif // USE_MIGRAPHX_BACKEND diff --git a/cpp/program/setup.cpp b/cpp/program/setup.cpp index bf7950315e..a96d763a54 100644 --- a/cpp/program/setup.cpp +++ b/cpp/program/setup.cpp @@ -83,6 +83,8 @@ vector Setup::initializeNNEvaluators( string backendPrefix = "cuda"; #elif defined(USE_TENSORRT_BACKEND) string backendPrefix = "trt"; + #elif defined(USE_MIGRAPHX_BACKEND) + string backendPrefix = "migraphx"; #elif defined(USE_METAL_BACKEND) string backendPrefix = "metal"; #elif defined(USE_OPENCL_BACKEND) @@ -142,7 +144,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 == "migraphx" || backendPrefix == "metal" ? false : true; if(cfg.contains(backendPrefix+"InputsUseNHWC"+idxStr)) inputsUseNHWC = cfg.getBool(backendPrefix+"InputsUseNHWC"+idxStr); else if(cfg.contains("inputsUseNHWC"+idxStr)) diff --git a/cpp/tests/testcommon.cpp b/cpp/tests/testcommon.cpp index fa1ce3532e..4c08996958 100644 --- a/cpp/tests/testcommon.cpp +++ b/cpp/tests/testcommon.cpp @@ -83,6 +83,15 @@ void TestCommon::overrideForBackends(bool& inputsNHWC, bool& useNHWC) { cout << "Backend is TensorRT, ignoring args and forcing useNHWC=false" << endl; useNHWC = false; } +#elif defined(USE_MIGRAPHX_BACKEND) + if(inputsNHWC != false) { + cout << "Backend is MIGraphX, ignoring args and forcing inputsNHWC=false" << endl; + inputsNHWC = false; + } + if(useNHWC != false) { + cout << "Backend is MIGraphX, ignoring args and forcing useNHWC=false" << endl; + useNHWC = false; + } #else (void)inputsNHWC; (void)useNHWC;