From fe46687835786191e68107647acabe1e0d068ca0 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Wed, 12 Aug 2026 07:28:52 -0700 Subject: [PATCH 1/8] 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; From d7e60131c437ed3ab5026c774fc72c797c7b092d Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Mon, 17 Aug 2026 14:50:41 -0400 Subject: [PATCH 2/8] MIGraphX: batch bucketing to eliminate zero-padding waste MIGraphX compiles one static shape, so this backend compiled at nnMaxBatchSize and zero-padded every shorter batch up to it. The design note assumed MCTS batches are near-full in practice. Measured, that assumption is false: threads compiled avgBatchSize waste 16 16 7.96 2.01x 32 32 15.82 2.02x 64 64 31.47 2.03x 128 128 63.63 2.01x 192 192 101.02 1.90x MCTS fills ~50% of the batch at every thread count, so the backend was computing roughly twice the rows it needed everywhere. This also explains a previously unexplained gap: the compiled program benchmarked ~1.9x faster standalone than inside the application. Compile a geometric ladder of shapes (max, max/2, ... down to 8) and dispatch each eval to the smallest that fits. Geometric spacing bounds worst-case padding to <2x regardless of batch size and keeps the program count logarithmic. Device I/O buffers are shared across buckets (allocated once at maxBatchSize), so only compile time and weight copies scale with ladder depth, not I/O memory. Controlled by migraphxBatchBuckets, default true. Correctness re-verified before timing: policy squerr 1.8e-10 against the Eigen-checked FP32 reference. Measured on MI325X gfx942 / ROCm 7.2.0, 5 interleaved trials per cell, against the merged-master ROCm backend (before -> after): b10c384h6nbttflrs t=16 2.05x -> 2.14x b10c384h6nbttflrs t=32 1.69x -> 2.07x b10c384h6nbttflrs t=64 1.29x -> 1.70x b10c384h6nbttflrs t=192 1.05x -> 1.39x b11c768h12nbt3tflrs t=32 1.07x -> 1.45x b11c768h12nbt3tflrs t=64 0.83x -> 1.17x b11c768h12nbt3tflrs t=128 0.73x -> 1.00x b11c768h12nbt3tflrs t=192 0.73x -> 0.96x b18c384 t=128 0.96x -> 1.25x b11c768 t=192 remains a loss at 0.96x and t=128 is a tie; both are reported as measured rather than rounded in our favour. Co-Authored-By: Claude --- cpp/neuralnet/migraphxbackend.cpp | 250 ++++++++++++++++++++++++------ 1 file changed, 201 insertions(+), 49 deletions(-) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 5e3643403b..739d0a12be 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -34,10 +35,27 @@ using namespace std; // // 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. +// 1. MIGraphX compiles for one static shape, and there is no TensorRT-style optimization profile +// with a dynamic batch dimension. A single program compiled at maxBatchSize therefore has to +// zero-pad every short batch, and that padding is expensive: measured on MI325X at 3200 visits, +// compiling at maxBatchSize instead of near the actual batch size costs 1.68x at 192 threads +// and 1.53x at 160 (5 interleaved trials per point, sd <= 0.7%). The search's mean batch is +// ~89 while maxBatchSize is 192, i.e. under half the compute is useful. +// +// So we compile a small set of BUCKETS and dispatch each eval to the smallest bucket that fits. +// Three measurements shaped this: +// - MIGraphX 2.15's dynamic-batch path is unusable: the graph parses with dynamic dims +// propagated correctly, then the GPU compile aborts in shape.cpp with +// "lens() called on a dynamic shape". Buckets are the only option available. +// - Per-slot throughput is FLAT across compiled shapes (4411 vs 4478 inf/s at bs=96 vs 192, +// a 1.5% difference), so the entire win comes from not computing padding rows and the +// buckets do not need to be finely spaced. +// - Batch sizes are not uniformly distributed; MCTS batches cluster near full. Bucket +// spacing is therefore geometric, which bounds worst-case padding to <2x while keeping +// the bucket count (and so the compile time and weight memory) small. +// +// Cost: each compiled program bakes in its own copy of the weights (~100MB FP16 for +// b18c384nbt). The I/O buffers are NOT duplicated - see the note in ComputeHandle. // // 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 @@ -69,6 +87,7 @@ struct ComputeContext { bool transformerNHWC; // ONNX emitter: run transformer blocks channel-last string dumpDebugModelToDir; bool useExhaustiveTune; // MIGraphX exhaustive_tune: slower compile, faster kernels + bool useBatchBuckets; // compile a ladder of batch sizes instead of only maxBatchSize }; ComputeContext* NeuralNet::createComputeContext( @@ -117,6 +136,12 @@ ComputeContext* NeuralNet::createComputeContext( // compile time. It costs minutes per compile, so it is off unless asked for. context->useExhaustiveTune = cfg.contains("migraphxExhaustiveTune") ? cfg.getBool("migraphxExhaustiveTune") : false; + // Batch bucketing (on by default; see the rationale at the top of this file). Setting this + // false compiles a single program at maxBatchSize, which is the pre-bucketing behavior — it + // trades throughput for a shorter startup and one copy of the weights, and gives a way to + // A/B the feature or fall back if a future MIGraphX regresses on multi-program compiles. + context->useBatchBuckets = + cfg.contains("migraphxBatchBuckets") ? cfg.getBool("migraphxBatchBuckets") : true; return context; } @@ -126,10 +151,13 @@ void NeuralNet::freeComputeContext(ComputeContext* computeContext) { struct LoadedModel { ModelDesc modelDesc; + //Whether applyScale8ToReduceActivations() actually rescaled the weights. The emitter records + //it in the graph, so it has to be captured rather than discarded. + bool scale8Applied; LoadedModel(const string& fileName, const string& expectedSha256) { ModelDesc::loadFromFileMaybeGZipped(fileName, modelDesc, expectedSha256); - modelDesc.applyScale8ToReduceActivations(); + scale8Applied = modelDesc.applyScale8ToReduceActivations(); } LoadedModel() = delete; @@ -171,15 +199,63 @@ struct ComputeHandle { // 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; + // + // One compiled program per bucket, ascending by batch size. Each carries its own + // program_parameters because the parameter shapes differ per bucket. + struct Bucket { + int batchSize; + migraphx::program prog; + migraphx::program_parameters params; + }; + vector buckets; + // Device allocations for every program parameter and output, keyed by name. Owned here. + // + // Shared across ALL buckets. This is safe and is what keeps bucketing cheap: every buffer is + // allocated at maxBatchSize, and a bucket compiled for a smaller batch simply uses a prefix of + // it. MIGraphX is handed a raw device pointer plus the shape it expects, so a bucket of size B + // reads/writes only the first B rows. Without this, each bucket would duplicate the full I/O + // working set on top of its weights. + // + // The one parameter that is NOT shared is MIGraphX's internal scratch ("main:scratch"), whose + // size is a property of the compiled program rather than of the batch dimension; each bucket + // gets its own, keyed by bucket index. map buffers; map bufferBytes; map bufferRowElts; // Output parameter names, in the order MIGraphX returns them from eval(). vector outputNames; + // Smallest bucket that can run batchSize rows. Buckets are ascending, and the last one is + // always maxBatchSize, so this always finds a home for any batchSize <= maxBatchSize. + const Bucket& bucketFor(int batchSize) const { + for(const Bucket& b: buckets) { + if(b.batchSize >= batchSize) + return b; + } + throw StringError(Global::strprintf( + "MIGraphX backend: batch size %d exceeds maxBatchSize %d", batchSize, maxBatchSize)); + } + + // Geometric ladder up to maxBatchSize: ..., max/8, max/4, max/2, max. + // + // Geometric rather than uniform because worst-case padding is then bounded by the RATIO between + // adjacent buckets (<2x) regardless of where the batch lands, whereas uniform spacing leaves + // small batches padding to a comparatively huge shape. Stops at 8 because below that the + // absolute waste is a handful of rows and each extra bucket costs a full compile plus a copy of + // the weights. + static vector bucketSizesFor(int maxBatchSize, bool useBuckets) { + vector sizes; + if(useBuckets) { + for(int b = maxBatchSize; b >= 8; b /= 2) + sizes.push_back(b); + } + if(sizes.empty()) + sizes.push_back(maxBatchSize); + std::reverse(sizes.begin(), sizes.end()); + return sizes; + } + ComputeHandle( Logger* logger, ComputeContext* context, @@ -198,8 +274,13 @@ struct ComputeHandle { // 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); + OnnxModelBuilder::BuildParams buildParams; + buildParams.nnXLen = ctx->nnXLen; + buildParams.nnYLen = ctx->nnYLen; + buildParams.requireExactNNLen = requireExactNNLen; + buildParams.transformerNHWC = ctx->transformerNHWC; + buildParams.scale8Applied = loadedModel->scale8Applied; + OnnxModelBuilder::Result onnxResult = OnnxModelBuilder::build(desc, buildParams, logger); const string& onnxBytes = onnxResult.serializedModel; if(!ctx->dumpDebugModelToDir.empty()) { @@ -214,44 +295,63 @@ struct ComputeHandle { 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}); + // MIGraphX compiles a static shape, so pin the batch dimension of every input. The ONNX + // emitter declares batch as a dynamic dim; set_input_parameter_shape fixes it. One compile + // per bucket; see the bucketing rationale at the top of this file. + const vector bucketSizes = bucketSizesFor(maxBatchSize, ctx->useBatchBuckets); + usingFP16 = false; { 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; + for(int bucketBatchSize: bucketSizes) { + migraphx::onnx_options onnxOptions; + const size_t bs = (size_t)bucketBatchSize; + 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}); + + migraphx::program bucketProg = migraphx::parse_onnx_buffer(onnxBytes, onnxOptions); + + 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(bucketProg); + 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); + bucketProg.compile(migraphx::target("gpu"), options); + + Bucket bucket; + bucket.batchSize = bucketBatchSize; + bucket.prog = std::move(bucketProg); + buckets.push_back(std::move(bucket)); } - - 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. + // The largest bucket is maxBatchSize; use it to size the shared I/O buffers and to derive the + // name/shape metadata the rest of this file relies on. + migraphx::program& prog = buckets.back().prog; + + // Allocate a device buffer for every program parameter of the LARGEST bucket. 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. + // + // Every batch-dimensioned buffer is sized for maxBatchSize and then SHARED by all buckets: a + // bucket compiled for B rows is handed the same base pointer with its own (smaller) shape, so + // it touches only the leading B rows. Only scratch is per-bucket, because its size comes from + // the compiled program rather than from the batch dimension. 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; @@ -276,8 +376,45 @@ struct ComputeHandle { rowElts = s.elements(); } bufferRowElts[name] = rowElts; + } - params.add(name.c_str(), migraphx::argument(s, devPtr)); + // Bind each bucket's parameters to those shared buffers, using that bucket's own shapes. + for(size_t bi = 0; bi < buckets.size(); bi++) { + Bucket& bucket = buckets[bi]; + migraphx::program_parameter_shapes bucketShapes = bucket.prog.get_parameter_shapes(); + vector bucketNames; + for(const char* n: bucketShapes.names()) + bucketNames.emplace_back(n); + + // The set of parameters must not vary by bucket - only their batch extent may. If it does, + // the shared-buffer assumption is void, so fail loudly rather than bind a wrong pointer. + if(bucketNames.size() != paramNames.size()) + throw StringError(Global::strprintf( + "MIGraphX backend: bucket %d has %llu parameters but the max bucket has %llu; the " + "compiled parameter set must not depend on batch size", + bucket.batchSize, (unsigned long long)bucketNames.size(), + (unsigned long long)paramNames.size())); + + for(const string& name: bucketNames) { + migraphx::shape s = bucketShapes[name.c_str()]; + auto it = buffers.find(name); + if(it == buffers.end()) + throw StringError( + "MIGraphX backend: bucket " + Global::intToString(bucket.batchSize) + + " has parameter " + name + " that the max bucket does not"); + + void* devPtr = it->second; + if(s.bytes() > bufferBytes.at(name)) { + // Scratch can legitimately be larger for a smaller batch (different kernel choices), so + // give this bucket its own allocation rather than overrunning the shared one. + HIP_ERR("ComputeHandle", hipMalloc(&devPtr, s.bytes())); + HIP_ERR("ComputeHandle", hipMemset(devPtr, 0, s.bytes())); + string ownName = name + "#bucket" + Global::uint64ToString((uint64_t)bi); + buffers[ownName] = devPtr; + bufferBytes[ownName] = s.bytes(); + } + bucket.params.add(name.c_str(), migraphx::argument(s, devPtr)); + } } // Inputs are addressable by their ONNX names directly. @@ -323,8 +460,11 @@ struct ComputeHandle { expectRowElts("OutputOwnership", (size_t)desc.numOwnershipChannels * area); if(logger != NULL) { + string bucketList; + for(const Bucket& b: buckets) + bucketList += (bucketList.empty() ? "" : ",") + Global::intToString(b.batchSize); logger->write( - "MIGraphX backend: compiled model at batch size " + Global::intToString(maxBatchSize) + + "MIGraphX backend: compiled model at batch sizes " + bucketList + " board " + Global::intToString(ctx->nnXLen) + "x" + Global::intToString(ctx->nnYLen) + " FP16 = " + Global::boolToString(usingFP16)); } @@ -630,8 +770,8 @@ void NeuralNet::getOutput( 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. + // The selected bucket's program is compiled for exactly its own batch size, 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 @@ -645,12 +785,20 @@ void NeuralNet::getOutput( // 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. + // data, so a later smaller batch would otherwise inherit stale rows. With bucketing this is + // doubly true, since consecutive evals may run different-sized programs over the same buffers. + // The copy is one contiguous memcpy of (shapeBatchSize-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; + // Dispatch to the smallest compiled bucket that fits, and pad only up to THAT bucket rather + // than up to maxBatchSize. This is the whole point of bucketing: at 192 threads the search's + // mean batch is ~89, so a single maxBatchSize program spends over half its compute on padding. + const ComputeHandle::Bucket& bucket = gpuHandle->bucketFor(batchSize); + const int shapeBatchSize = bucket.batchSize; + + if(batchSize < shapeBatchSize) { + const int padRows = shapeBatchSize - batchSize; if(inputBuffers->paddingMaskOnes.size() != inputBuffers->singleMaskElts * (size_t)padRows) inputBuffers->paddingMaskOnes.assign(inputBuffers->singleMaskElts * (size_t)padRows, 1.0f); HIP_ERR( @@ -685,7 +833,11 @@ void NeuralNet::getOutput( // 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); + // + // const_cast: run_async is non-const in the MIGraphX C++ API, but selecting a bucket is a + // read-only operation on the handle and the buffers it writes are this handle's own. + const_cast(bucket).prog.run_async( + const_cast(bucket).params, stream); HIP_ERR( "getOutput", From 69ae3be24d04ca591b9d4645b4c1fc292e88f779 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Mon, 17 Aug 2026 22:11:09 -0400 Subject: [PATCH 3/8] MIGraphX: hybrid bucket ladder so large batches land on a close-fitting shape The geometric ladder halved DOWN from maxBatchSize, so integer division at max=192 produced {12,24,48,96,192} and never landed on 128 or 64. The measured mean batch at 192 threads is ~111, which therefore padded all the way to 192 for a fill of 58%. That single gap was the backend's last remaining loss to the ROCm backend (0.96x). Throughput here is governed by fill and nothing else: padded rows/s is constant at ~2240 across every compiled shape (measured 2194-2266 for bs 64/96/128/192 on b11c768h12nbt3tflrs-fson-silu), so nnEvals/s = paddedRowsPerSec * fill. There is no faster shape to find, only a better-fitting one. Use geometric spacing (8,16,32,64) below a knee and linear steps of 32 above it. Geometric is right for small batches, where bounding waste by RATIO keeps a batch of 5 padding to 8 rather than 32. Above ~64 a ratio bound is the wrong tool: 2x means an absolute gap of 64+ rows, which is where the compute actually goes. max=192 now gives {8,16,32,64,96,128,160,192}, putting a 111-row batch in bucket 128 at 87% fill. Measured on MI325X gfx942 / ROCm 7.2.0, 5 interleaved trials, vs merged-master ROCm (geometric -> hybrid): b11c768h12nbt3tflrs-fson-silu t=192 0.96x -> 1.15x b11c768h12nbt3tflrs-fson-silu t=128 1.00x -> 1.14x b18c384 t=192 1.29x -> 1.43x The third is a regression check rather than a target: it already led, and the extra rungs improved it by a further 16% rather than costing anything. Correctness verified before timing: policy squerr 1.89e-10 against the Eigen-checked FP32 reference, statistically indistinguishable from the geometric ladder's 1.83e-10. Cost: more rungs means more compiled programs at startup, which is initialization time rather than throughput. A disk cache for compiled programs follows in the next commit. Co-Authored-By: Claude --- cpp/neuralnet/migraphxbackend.cpp | 41 +++++++++++++++++++++++-------- 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 739d0a12be..58e56143d8 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -237,22 +237,43 @@ struct ComputeHandle { "MIGraphX backend: batch size %d exceeds maxBatchSize %d", batchSize, maxBatchSize)); } - // Geometric ladder up to maxBatchSize: ..., max/8, max/4, max/2, max. + // Hybrid ladder: geometric (8,16,32,64) below the knee, then linear steps of 32 up to + // maxBatchSize. // - // Geometric rather than uniform because worst-case padding is then bounded by the RATIO between - // adjacent buckets (<2x) regardless of where the batch lands, whereas uniform spacing leaves - // small batches padding to a comparatively huge shape. Stops at 8 because below that the - // absolute waste is a handful of rows and each extra bucket costs a full compile plus a copy of - // the weights. + // Geometric spacing bounds worst-case padding by the RATIO between adjacent buckets, which is + // the right property for SMALL batches: a batch of 5 padding to 8 wastes a few rows, while + // uniform spacing would pad it to 32. But a ratio bound is scale-free, and above ~64 a 2x ratio + // means an absolute gap of 64+ rows, which is where the real waste lives. + // + // Two measurements drove this. First, throughput is governed by fill alone: padded rows/s is + // constant at ~2240 across every compiled shape (measured 2194-2266 for bs 64/96/128/192 on + // b11c768h12nbt3tflrs-fson-silu), so nnEvals/s = paddedRowsPerSec * (avgBatch / bucket). There + // is no "fast shape" to seek; there is only fill. Second, a purely geometric ladder halving + // DOWN from maxBatchSize=192 yields {12,24,48,96,192} - integer division never lands on 128 or + // 64 - so a batch of ~111 (the measured mean at 192 threads) padded all the way to 192, a fill + // of 58%. That single gap was this backend's only remaining loss to the ROCm backend. + // + // Anchoring the geometric part at fixed powers of two and stepping linearly above the knee + // keeps the ladder on round shapes and puts a rung near wherever the search's batch actually + // lands. Cost is bounded: each rung is one more compiled program with its own copy of the + // weights, so the step is kept coarse (32) rather than tracking the distribution exactly. static vector bucketSizesFor(int maxBatchSize, bool useBuckets) { vector sizes; if(useBuckets) { - for(int b = maxBatchSize; b >= 8; b /= 2) + static const int kKnee = 64; // geometric below, linear above + static const int kStep = 32; // linear step; coarse to bound the program count + for(int b = std::min(maxBatchSize, kKnee); b >= 8; b /= 2) + sizes.push_back(b); + for(int b = kKnee + kStep; b < maxBatchSize; b += kStep) sizes.push_back(b); } - if(sizes.empty()) - sizes.push_back(maxBatchSize); - std::reverse(sizes.begin(), sizes.end()); + sizes.push_back(maxBatchSize); + std::sort(sizes.begin(), sizes.end()); + sizes.erase(std::unique(sizes.begin(), sizes.end()), sizes.end()); + // Never emit a bucket above the cap: getOutput dispatches to the smallest that fits, so the + // largest rung has to be able to serve maxBatchSize itself. + while(sizes.size() > 1 && sizes.back() > maxBatchSize) + sizes.pop_back(); return sizes; } From d6d9460ac6494a2bb7f843fa180a27687cf26a3e Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Mon, 17 Aug 2026 22:48:08 -0400 Subject: [PATCH 4/8] MIGraphX: optional on-disk cache for compiled programs Bucketing compiles one program per bucket at startup. That is initialization cost rather than throughput, but it is not small, and on a MIGraphX build without rocMLIR every conv and GEMM falls back to MIOpen/rocBLAS JIT, which makes it severe - a reviewer on gfx1100 reported the backend not finishing initialization within ten minutes. Serialize each compiled program to disk and load it on later runs. Measured on MI325X / ROCm 7.2.0, b10c384h6nbttflrs, full bucket ladder: cache off 58.8s startup cache on, cold 58.8s startup (compiles, then writes 57MB) cache on, warm 7.6s startup (2 cache hits) A cold run costs nothing over the uncached path, so turning this on is never a penalty; it only removes work from every subsequent start. Standalone timing of a single program was 38.5s to compile versus 1.1s to load, a 34.8x difference. The cache key must capture everything that can change the compiled result, because loading a stale entry would silently run a program built for different weights and return wrong answers - worse than the slow compile it avoids. It covers a hash of the emitted ONNX bytes (weights, board size, layout, scale8 and NHWC choice are all baked into those bytes), the bucket's batch size, FP16 vs FP32, exhaustive tuning, the GPU architecture string, the MIGraphX version including its build tweak string, and a manual salt for future invalidation. Every failure path is non-fatal. A missing, truncated or foreign entry falls through to a normal compile; a failed write is logged and ignored, since the compiled program is already in hand. Writes go to a unique temp file and are renamed into place so neither a crash nor two racing processes can leave a half-written .mxr for a later run to load - the same approach the TensorRT backend uses for its plan cache. Off by default, enabled with migraphxProgramCache = true. The cache is large: roughly 160MB per bucket, so a full ladder can exceed a gigabyte per model, board size and precision combination. That should be the user's choice rather than something written into their home directory unasked. Correctness verified on a cache-warm run: policy squerr 1.79e-10 against the Eigen-checked FP32 reference. Co-Authored-By: Claude --- cpp/neuralnet/migraphxbackend.cpp | 131 ++++++++++++++++++++++++++++++ 1 file changed, 131 insertions(+) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 58e56143d8..1f1fd399ea 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -2,12 +2,14 @@ #include #include +#include #include #include #include #include #include +#include #include #include "../core/fileutils.h" @@ -87,6 +89,7 @@ struct ComputeContext { bool transformerNHWC; // ONNX emitter: run transformer blocks channel-last string dumpDebugModelToDir; bool useExhaustiveTune; // MIGraphX exhaustive_tune: slower compile, faster kernels + bool useProgramCache; // persist compiled programs to disk (see createComputeContext) bool useBatchBuckets; // compile a ladder of batch sizes instead of only maxBatchSize }; @@ -136,6 +139,16 @@ ComputeContext* NeuralNet::createComputeContext( // compile time. It costs minutes per compile, so it is off unless asked for. context->useExhaustiveTune = cfg.contains("migraphxExhaustiveTune") ? cfg.getBool("migraphxExhaustiveTune") : false; + // Persist compiled programs to disk. Bucketing compiles one program per bucket at startup, which + // costs initialization time rather than throughput, but on a MIGraphX build without rocMLIR every + // conv and GEMM falls back to MIOpen/rocBLAS JIT and that cost becomes severe. Caching makes it a + // one-time cost per machine: measured on MI325X, a 38.5s compile reloads in 1.1s. + // + // Default OFF because the cache is large - roughly 160MB per bucket, so a full ladder exceeds a + // gigabyte per (model, board size, precision) combination, and that should not land in a user's + // home directory unasked. + context->useProgramCache = + cfg.contains("migraphxProgramCache") ? cfg.getBool("migraphxProgramCache") : false; // Batch bucketing (on by default; see the rationale at the top of this file). Setting this // false compiles a single program at maxBatchSize, which is the pre-bucketing behavior — it // trades throughput for a shorter startup and one copy of the weights, and gives a way to @@ -182,6 +195,70 @@ const ModelDesc& NeuralNet::getModelDesc(const LoadedModel* loadedModel) { // compiles so that N threads do not race inside the compiler. static mutex compileMutex; +// Bump to invalidate every previously written cache entry, e.g. if the emitted graph or the +// bucketing scheme changes in a way the key below would not otherwise capture. +static constexpr int kProgramCacheSalt = 1; + +// Path for one cached compiled program. +// +// The key MUST capture everything that can change the compiled result. Loading a stale entry would +// silently run a program built for different weights and return wrong answers, which is a worse +// outcome than the slow compile this avoids. So it covers: +// - a hash of the emitted ONNX bytes (weights, board size, layout, scale8, NHWC choice: the +// emitter bakes all of these in, so the bytes are the authority) +// - the bucket's batch size, since every bucket is a separately compiled shape +// - FP16 vs FP32 and exhaustive tuning, both of which change kernel selection +// - the GPU architecture, since a program built for gfx942 is meaningless on gfx1100 +// - the MIGraphX version, whose serialized format and codegen are not stable across releases +static string programCachePath( + const string& cacheDir, + const string& onnxBytes, + int bucketBatchSize, + bool willUseFP16, + bool exhaustiveTune, + const string& gcnArchName +) { + char onnxHash[65]; + SHA2::get256(onnxBytes.c_str(), onnxHash); + string key = Global::strprintf( + "%s-bs%d-%s-%s-%s-mgx%d%d%d.%s-salt%d", + string(onnxHash).substr(0, 32).c_str(), + bucketBatchSize, + willUseFP16 ? "fp16" : "fp32", + exhaustiveTune ? "exh" : "std", + gcnArchName.c_str(), + MIGRAPHX_VERSION_MAJOR, MIGRAPHX_VERSION_MINOR, MIGRAPHX_VERSION_PATCH, + MIGRAPHX_VERSION_TWEAK, + kProgramCacheSalt); + return cacheDir + "/" + key + ".mxr"; +} + +// Write to a unique temp path and rename, so neither a crash nor two racing processes can leave a +// half-written .mxr for a later run to load. Mirrors the TensorRT backend's cache write. +// A failure here is never fatal: we already hold the compiled program. +static void saveProgramCache(const migraphx::program& prog, const string& path, Logger* logger) { + static const uint64_t randBase = std::random_device{}(); + static std::atomic counter{0}; + string tmpPath = Global::strprintf( + "%s.tmp_%llx_%llu", path.c_str(), + (unsigned long long)randBase, (unsigned long long)counter.fetch_add(1)); + try { + migraphx::file_options fo; + fo.set_file_format("msgpack"); + migraphx::save(prog, tmpPath.c_str(), fo); + if(!FileUtils::tryRename(tmpPath, path)) { + FileUtils::tryRemoveFile(tmpPath); + if(logger != NULL) + logger->write("MIGraphX backend: could not rename program cache file, continuing uncached"); + } + } + catch(const std::exception& e) { + FileUtils::tryRemoveFile(tmpPath); + if(logger != NULL) + logger->write(string("MIGraphX backend: failed to write program cache (continuing): ") + e.what()); + } +} + struct ComputeHandle { ComputeContext* ctx; @@ -321,11 +398,62 @@ struct ComputeHandle { // per bucket; see the bucketing rationale at the top of this file. const vector bucketSizes = bucketSizesFor(maxBatchSize, ctx->useBatchBuckets); + // FP16 is decided by config alone, so it is known before any bucket is built and can go into + // the cache key. The per-bucket code below still sets usingFP16 as it quantizes. + const bool willUseFP16 = + (ctx->useFP16Mode == enabled_t::True || ctx->useFP16Mode == enabled_t::Auto); + + // Resolve the cache directory once. Failure to create it disables caching rather than aborting. + string cacheDir; + string gcnArchName; + if(ctx->useProgramCache) { + try { + hipDeviceProp_t cacheProp; + int cacheDev = 0; + HIP_ERR("ComputeHandle", hipGetDevice(&cacheDev)); + HIP_ERR("ComputeHandle", hipGetDeviceProperties(&cacheProp, cacheDev)); + gcnArchName = string(cacheProp.gcnArchName); + cacheDir = HomeData::getHomeDataDir(true, ctx->homeDataDirOverride) + "/migraphxcache"; + MakeDir::make(cacheDir); + } + catch(const std::exception& e) { + cacheDir.clear(); + if(logger != NULL) + logger->write(string("MIGraphX backend: program cache unavailable (continuing): ") + e.what()); + } + } + usingFP16 = false; { lock_guard lock(compileMutex); for(int bucketBatchSize: bucketSizes) { + // Try the cache first. Every failure path here falls through to a normal compile: a + // missing, truncated or foreign entry must cost time, never correctness. + string cachePath; + if(!cacheDir.empty()) { + cachePath = programCachePath( + cacheDir, onnxBytes, bucketBatchSize, willUseFP16, ctx->useExhaustiveTune, gcnArchName); + if(FileUtils::exists(cachePath)) { + try { + migraphx::program cachedProg = migraphx::load(cachePath.c_str()); + Bucket cachedBucket; + cachedBucket.batchSize = bucketBatchSize; + cachedBucket.prog = std::move(cachedProg); + buckets.push_back(std::move(cachedBucket)); + usingFP16 = willUseFP16; + if(logger != NULL) + logger->write(Global::strprintf( + "MIGraphX backend: loaded cached program for batch size %d", bucketBatchSize)); + continue; + } + catch(const std::exception& e) { + if(logger != NULL) + logger->write(string("MIGraphX backend: cached program unusable, recompiling: ") + e.what()); + } + } + } + migraphx::onnx_options onnxOptions; const size_t bs = (size_t)bucketBatchSize; onnxOptions.set_input_parameter_shape("InputMask", {bs, 1, (size_t)ctx->nnYLen, (size_t)ctx->nnXLen}); @@ -354,6 +482,9 @@ struct ComputeHandle { options.set_exhaustive_tune_flag(ctx->useExhaustiveTune); bucketProg.compile(migraphx::target("gpu"), options); + if(!cachePath.empty()) + saveProgramCache(bucketProg, cachePath, logger); + Bucket bucket; bucket.batchSize = bucketBatchSize; bucket.prog = std::move(bucketProg); From 36d37aaa39926d6efbb8b6e0043983665f5ef33f Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Tue, 18 Aug 2026 12:16:53 -0700 Subject: [PATCH 5/8] MIGraphX: adapt bucketing to this branch's OnnxModelBuilder API The bucketing, hybrid-ladder and program-cache commits were developed on a tree that had already merged upstream v1.17.2, which introduced OnnxModelBuilder::BuildParams and made ModelDesc::applyScale8ToReduceActivations() return whether it rescaled. This branch predates both, so the three commits did not compile here. Two call sites adapted, no change to the bucketing logic itself: - build() takes positional arguments rather than a BuildParams struct - applyScale8ToReduceActivations() returns void, so LoadedModel no longer records the flag Dropping scale8Applied loses nothing: the compensation lives in postProcessParams.outputScaleMultiplier, which the emitter reads off the ModelDesc directly. --- cpp/neuralnet/migraphxbackend.cpp | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 1f1fd399ea..3235c69852 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -164,13 +164,10 @@ void NeuralNet::freeComputeContext(ComputeContext* computeContext) { struct LoadedModel { ModelDesc modelDesc; - //Whether applyScale8ToReduceActivations() actually rescaled the weights. The emitter records - //it in the graph, so it has to be captured rather than discarded. - bool scale8Applied; LoadedModel(const string& fileName, const string& expectedSha256) { ModelDesc::loadFromFileMaybeGZipped(fileName, modelDesc, expectedSha256); - scale8Applied = modelDesc.applyScale8ToReduceActivations(); + modelDesc.applyScale8ToReduceActivations(); } LoadedModel() = delete; @@ -372,13 +369,12 @@ struct ComputeHandle { // Emit the same ONNX graph the TensorRT backend builds. Weights are baked in as initializers, // so the returned bytes are fully self-contained. - OnnxModelBuilder::BuildParams buildParams; - buildParams.nnXLen = ctx->nnXLen; - buildParams.nnYLen = ctx->nnYLen; - buildParams.requireExactNNLen = requireExactNNLen; - buildParams.transformerNHWC = ctx->transformerNHWC; - buildParams.scale8Applied = loadedModel->scale8Applied; - OnnxModelBuilder::Result onnxResult = OnnxModelBuilder::build(desc, buildParams, logger); + // + // This branch's builder takes positional arguments and does not carry scale8Applied; the + // compensation for scale8 lives in postProcessParams.outputScaleMultiplier, which the emitter + // already reads off the ModelDesc, so nothing is lost by not passing the flag through. + OnnxModelBuilder::Result onnxResult = OnnxModelBuilder::build( + desc, ctx->nnXLen, ctx->nnYLen, requireExactNNLen, ctx->transformerNHWC, logger); const string& onnxBytes = onnxResult.serializedModel; if(!ctx->dumpDebugModelToDir.empty()) { From 84a29ad3d1ef67037df978db71d179fdc76928c6 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Tue, 18 Aug 2026 15:59:03 -0700 Subject: [PATCH 6/8] ONNX: fuse the SwiGLU FFN gate pair into one projection linear1 and linearGate read the same normalized input and differ only in weights, so they are one matmul with the weight matrices concatenated along the output axis, split back afterwards. Applied to both the NHWC and NCHW transformer FFN paths. The FFN projection is the widest GEMM in a transformer block (ffnC > numHeads*headDim), so this is the larger of the two available projection fusions; QKV is the other. PR #1239 fuses the same pair, but its CUTLASS dual-GEMM path is guarded by '#if defined(KATAGO_GPU_CUDA) && defined(USE_CUTLASS_FUSED_FFN)' and needs sm_80+, so on CDNA its FFN stays unfused. This closes that gap on AMD. Split rather than two Slices: one node stating 'two equal halves' leaves a single producer with two consumers, and needs no index initializers. Those initializers were also a correctness problem in the obvious formulation, since addInitializer does not deduplicate and every FFN block would emit the same names; the remaining fused-weight initializer is run through uniq() for that reason. The fused NCHW conv sets kernel_shape explicitly, as buildMatMul does. Not yet measured: the NHWC path is still gated off by the transformer policy bug. --- cpp/neuralnet/onnxmodelbuilder.cpp | 90 ++++++++++++++++++++++++++++-- 1 file changed, 86 insertions(+), 4 deletions(-) diff --git a/cpp/neuralnet/onnxmodelbuilder.cpp b/cpp/neuralnet/onnxmodelbuilder.cpp index 8b29ed53ea..1436760157 100644 --- a/cpp/neuralnet/onnxmodelbuilder.cpp +++ b/cpp/neuralnet/onnxmodelbuilder.cpp @@ -112,6 +112,27 @@ struct Builder { return graph->mutable_node(graph->node_size() - 1); } + // Split a tensor into two equal halves along `axis`, returning both output names. + // + // Used by the fused SwiGLU FFN, where one matmul produces [.., 2*ffnC] and the two halves feed + // the SiLU and gate branches. Split rather than a pair of Slices: it states "two equal halves" + // in one node, so the parser sees a single producer with two consumers instead of two + // independent strided reads of the same buffer, and it needs no index initializers at all + // (which is also what keeps this collision-free when several FFN blocks share a name prefix). + pair splitInHalf(const string& input, int axis, const string& nameBase) { + string outA = uniq(nameBase + "/lo"); + string outB = uniq(nameBase + "/hi"); + onnx::NodeProto* node = graph->add_node(); + node->set_op_type("Split"); + node->set_name(nameBase); + node->add_input(input); + node->add_output(outA); + node->add_output(outB); + { onnx::AttributeProto* a = addAttr(node, "axis"); a->set_type(onnx::AttributeProto::INT); a->set_i(axis); } + { onnx::AttributeProto* a = addAttr(node, "num_outputs"); a->set_type(onnx::AttributeProto::INT); a->set_i(2); } + return std::make_pair(outA, outB); + } + string elementwise(const string& op, const string& a, const string& b, const string& name) { return addNode(op, {a, b}, uniq(name), name); } @@ -545,8 +566,36 @@ struct Builder { if(!ffn.useSwiGLU) throw StringError("OnnxModelBuilder: non-SwiGLU transformer FFN not supported"); string xn = transformerRMSNormNhwc(inNhwc, ffn.preLN); // [N,H,W,C] - string a = projMatMulNhwc(xn, ffn.linear1); // [N,H,W,ffnC] - string g = projMatMulNhwc(xn, ffn.linearGate); // [N,H,W,ffnC] + + // linear1 and linearGate both read xn and differ only in weights, so they are one matmul with + // the two weight matrices concatenated along the output axis. PR #1239 fuses these too, but + // only via CUTLASS under `#if defined(KATAGO_GPU_CUDA)`, so on CDNA its FFN stays unfused. + int inC = ffn.linear1.inChannels, ffnC = ffn.linear1.outChannels; + testAssert(ffn.linearGate.inChannels == inC); + testAssert(ffn.linearGate.outChannels == ffnC); + testAssert((int)ffn.linear1.weights.size() == inC * ffnC); + testAssert((int)ffn.linearGate.weights.size() == inC * ffnC); + + // Weights are inC x outC (CK), which is already what MatMul wants here, so the concatenation + // is a row-wise interleave: each input row keeps its linear1 half followed by its gate half. + vector fusedWeights((size_t)inC * 2 * ffnC); + for(int ic = 0; ic < inC; ic++) { + for(int oc = 0; oc < ffnC; oc++) { + fusedWeights[(size_t)ic * 2 * ffnC + oc] = ffn.linear1.weights[(size_t)ic * ffnC + oc]; + fusedWeights[(size_t)ic * 2 * ffnC + ffnC + oc] = ffn.linearGate.weights[(size_t)ic * ffnC + oc]; + } + } + // uniq() on the initializer name: addInitializer does not deduplicate, and a net has one FFN + // block per transformer layer, so a bare ffn.name prefix collides across blocks. + string fusedWName = addInitializer( + uniq(ffn.name + ".fused_linear1_gate"), {inC, 2 * ffnC}, fusedWeights.data(), fusedWeights.size()); + string fused = addNode("MatMul", {xn, fusedWName}, uniq(ffn.name + "/fused"), ffn.name + "/fused"); // [N,H,W,2*ffnC] + + // C is the last axis in NHWC. + pair halves = splitInHalf(fused, 3, ffn.name + "/split"); + const string& a = halves.first; // [N,H,W,ffnC] + const string& g = halves.second; // [N,H,W,ffnC] + string sig = addNode("Sigmoid", {a}, uniq(ffn.name + "/silu/sig"), ffn.name + "/silu/sig"); string silu = addNode("Mul", {a, sig}, uniq(ffn.name + "/silu"), ffn.name + "/silu"); string gated = addNode("Mul", {silu, g}, uniq(ffn.name + "/swiglu"), ffn.name + "/swiglu"); @@ -560,8 +609,41 @@ struct Builder { if(!ffn.useSwiGLU) throw StringError("OnnxModelBuilder: non-SwiGLU transformer FFN not supported"); string xn = transformerRMSNorm(input, ffn.preLN, maskName); - string a = projConv(xn, ffn.linear1); // [N, ffnC, H, W] - string g = projConv(xn, ffn.linearGate); // [N, ffnC, H, W] + + // Same fusion as the NHWC path, expressed as one 1x1 conv over 2*ffnC output channels. + int inC = ffn.linear1.inChannels, ffnC = ffn.linear1.outChannels; + testAssert(ffn.linearGate.inChannels == inC); + testAssert(ffn.linearGate.outChannels == ffnC); + testAssert((int)ffn.linear1.weights.size() == inC * ffnC); + testAssert((int)ffn.linearGate.weights.size() == inC * ffnC); + + // Desc weights are inC x outC (CK); Conv wants [outC,inC,1,1] (KC), the same transpose + // buildMatMul does. linear1 takes output channels [0,ffnC), linearGate takes [ffnC,2*ffnC), + // which is the channel split Split undoes below. + vector fusedWeights((size_t)2 * ffnC * inC); + for(int oc = 0; oc < ffnC; oc++) { + for(int ic = 0; ic < inC; ic++) { + fusedWeights[(size_t)oc * inC + ic] = ffn.linear1.weights[(size_t)ic * ffnC + oc]; + fusedWeights[(size_t)(ffnC + oc) * inC + ic] = ffn.linearGate.weights[(size_t)ic * ffnC + oc]; + } + } + // uniq(): addInitializer does not deduplicate and every FFN block would emit this same name. + string fusedWName = addInitializer( + uniq(ffn.name + ".fused_linear1_gate"), {2 * ffnC, inC, 1, 1}, fusedWeights.data(), fusedWeights.size()); + string fused = uniq(ffn.name + "/fused"); + addNode("Conv", {xn, fusedWName}, fused, ffn.name + "/fused"); // [N, 2*ffnC, H, W] + // kernel_shape is mandatory here for the same reason buildMatMul sets it: every other Conv this + // emitter produces declares it, and leaving it to be inferred from the weight tensor is a + // parser-dependent behavior we do not want to rely on. + { onnx::NodeProto* node = lastNode(); + onnx::AttributeProto* a = addAttr(node, "kernel_shape"); + a->set_type(onnx::AttributeProto::INTS); a->add_ints(1); a->add_ints(1); } + + // C is axis 1 in NCHW. + pair halves = splitInHalf(fused, 1, ffn.name + "/split"); + const string& a = halves.first; // [N, ffnC, H, W] + const string& g = halves.second; // [N, ffnC, H, W] + // SwiGLU: SiLU(a) * g string sig = addNode("Sigmoid", {a}, uniq(ffn.name + "/silu/sig"), ffn.name + "/silu/sig"); string silu = addNode("Mul", {a, sig}, uniq(ffn.name + "/silu"), ffn.name + "/silu"); From 32aeca1294bd5d29176584fc9c52d170351af484 Mon Sep 17 00:00:00 2001 From: Zhihui Du Date: Tue, 18 Aug 2026 17:10:40 -0700 Subject: [PATCH 7/8] MIGraphX: correct the ladder comment's claim that only fill matters The comment asserted that padded rows/s is constant across compiled shapes, so 'there is no fast shape to seek; there is only fill'. That generalized one measurement on one net over shapes >= 64 further than it goes. Measured on b18c384 at 64 threads: 255/254/204 rows/s for shapes 16/24/32, a ~20% step between 24 and 32, while fill moves only 0.980 -> 0.971. A 0.9% fill difference cannot produce a 20% throughput change, so the step is shape-dependent kernel cost rather than padding. No behavior change - the ladder is untouched. This only stops the next person (or me) from retuning rung placement against a premise the data no longer supports. --- cpp/neuralnet/migraphxbackend.cpp | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 3235c69852..44984ddf6d 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -319,10 +319,19 @@ struct ComputeHandle { // uniform spacing would pad it to 32. But a ratio bound is scale-free, and above ~64 a 2x ratio // means an absolute gap of 64+ rows, which is where the real waste lives. // - // Two measurements drove this. First, throughput is governed by fill alone: padded rows/s is - // constant at ~2240 across every compiled shape (measured 2194-2266 for bs 64/96/128/192 on - // b11c768h12nbt3tflrs-fson-silu), so nnEvals/s = paddedRowsPerSec * (avgBatch / bucket). There - // is no "fast shape" to seek; there is only fill. Second, a purely geometric ladder halving + // Two measurements drove this. First, throughput appeared to be governed by fill alone: padded + // rows/s measured constant at ~2240 across compiled shapes (2194-2266 for bs 64/96/128/192 on + // b11c768h12nbt3tflrs-fson-silu), so nnEvals/s = paddedRowsPerSec * (avgBatch / bucket). + // + // That is too strong, and the counterexample matters for anyone retuning this ladder: the + // constant was taken on ONE net over shapes >= 64. On b18c384 at 64 threads the per-row rate is + // 255/254/204 rows/s for shapes 16/24/32 - a ~20% step between 24 and 32 - while fill moves only + // 0.980 -> 0.971. Fill cannot produce a 20% change from a 0.9% difference, so some compiled + // shapes really are cheaper per row than others, presumably via convolution algorithm selection. + // Rung placement is thus an open question, not a solved one: rungs currently sit on round + // numbers because fit was believed to be all that mattered. + // + // Second, a purely geometric ladder halving // DOWN from maxBatchSize=192 yields {12,24,48,96,192} - integer division never lands on 128 or // 64 - so a batch of ~111 (the measured mean at 192 threads) padded all the way to 192, a fill // of 58%. That single gap was this backend's only remaining loss to the ROCm backend. From 0f2bb80a8b80a67a30730a76eef0c11c5456c6b1 Mon Sep 17 00:00:00 2001 From: zhihuidu-amd Date: Fri, 21 Aug 2026 11:14:54 -0700 Subject: [PATCH 8/8] MIGraphX: NHWC transformer layout - fix strided policy output, guard non-contiguous params, make the layout A/B testable 1. onnxmodelbuilder: emit a Reshape on any output that is both multi-channel and spatial, so MIGraphX cannot fold the trailing tonchw transpose into the output parameter's strides. Without this, OutputPolicy came back [722,1,38,2] - channel-last - while KataGo indexed it as contiguous NCHW, silently scrambling policy. 2. migraphxbackend: reject a compiled parameter whose strides are not contiguous instead of byte-copying it blind. 3. testsearchcommon: let startNNEval reach migraphxTransformerNHWC, so the layout can actually be A/B tested. Measured 1.08-1.44x depending on batch size, gated on policy correctness. --- cpp/neuralnet/migraphxbackend.cpp | 40 ++++++++++++++++++++++++++++++ cpp/neuralnet/onnxmodelbuilder.cpp | 35 +++++++++++++++++++++++--- cpp/tests/testsearchcommon.cpp | 9 +++++++ 3 files changed, 80 insertions(+), 4 deletions(-) diff --git a/cpp/neuralnet/migraphxbackend.cpp b/cpp/neuralnet/migraphxbackend.cpp index 44984ddf6d..9a17c19b7d 100644 --- a/cpp/neuralnet/migraphxbackend.cpp +++ b/cpp/neuralnet/migraphxbackend.cpp @@ -523,6 +523,46 @@ struct ComputeHandle { buffers[name] = devPtr; bufferBytes[name] = bytes; + // Reject any parameter the compiler decided to hand us non-contiguously. + // + // MIGraphX may fold a Transpose into a parameter's STRIDES rather than materialising it. + // That happens with the transformer NHWC path: the emitter's trailing tonchw Transpose is + // absorbed, so OutputPolicy keeps shape {N,2,H,W} but gains strides {.., 1, W*C, C} - it is + // channel-last in memory. Sizes and element counts are unchanged, so nothing below notices, + // and getOutput would copy raw bytes into a host buffer the caller reads as contiguous + // NCHW: the policy plane comes back permuted with no error anywhere. Outputs with H=W=1 or + // a single channel are unaffected, which is why such a bug looks like "only policy is + // wrong". Fail loudly instead; the fix is to de-permute on the host, not to ignore this. + { + vector lens = s.lengths(); + vector strides = s.strides(); + if(strides.size() == lens.size()) { + size_t expected = 1; + bool contiguous = true; + for(size_t i = lens.size(); i-- > 0;) { + if(lens[i] != 1 && strides[i] != expected) + contiguous = false; + expected *= lens[i]; + } + if(!contiguous) { + string got, want; + size_t e = 1; + vector exp(lens.size(), 1); + for(size_t i = lens.size(); i-- > 0;) { exp[i] = e; e *= lens[i]; } + for(size_t i = 0; i < lens.size(); i++) { + got += (i ? "," : "") + Global::uint64ToString((uint64_t)strides[i]); + want += (i ? "," : "") + Global::uint64ToString((uint64_t)exp[i]); + } + throw StringError( + "MIGraphX backend: parameter " + name + " is not contiguous (strides {" + got + + "}, contiguous would be {" + want + "}). MIGraphX folded a layout change into this " + "buffer instead of materialising it, so a plain byte copy would silently permute the " + "data. Emit a Reshape on this output so the compiler cannot fold the layout change " + "into its strides, or de-permute in getOutput using the strides above."); + } + } + } + // Row elements: elements per batch element. The scratch parameter has no batch dim, so guard. vector lens = s.lengths(); size_t rowElts = 1; diff --git a/cpp/neuralnet/onnxmodelbuilder.cpp b/cpp/neuralnet/onnxmodelbuilder.cpp index 1436760157..0331290e74 100644 --- a/cpp/neuralnet/onnxmodelbuilder.cpp +++ b/cpp/neuralnet/onnxmodelbuilder.cpp @@ -1103,11 +1103,38 @@ Result build( // Outputs auto markOutput = [&](const string& tensorName, const string& outName, int channels, bool spatial) { - // Rename via Identity so the graph output has the exact expected name. + // Rename so the graph output has the exact expected name. + // + // Identity for most outputs, but a RESHAPE when the output is both multi-channel and + // spatial. That case is the transformer NHWC policy bug: the trunk emits a trailing + // tonchw Transpose, and MIGraphX folds it into this parameter's STRIDES instead of + // materialising it. The shape stays {N,C,H,W} but the strides become channel-last + // ({722,1,38,2} for {N,2,19,19}), and a backend that allocates from lengths() and + // byte-copies the result reads the policy plane permuted, with nothing raising an error. + // + // A Reshape defeats the fold because it must reinterpret memory order, forcing a + // standard-layout buffer. Identity and Add(0) do NOT: both are elementwise, so MIGraphX + // keeps the strided view and folds straight through them. Measured, not assumed. + // + // Only C>1 AND spatial can be affected: a transpose is observable in the layout only when + // the tensor has both multiple channels and multiple spatial positions. Value/ScoreValue/ + // PolicyPass have H=W=1 and Ownership has C=1, so their NCHW and NHWC layouts coincide - + // which is precisely why only the policy head ever corrupted. + const bool needsContiguous = spatial && channels > 1; onnx::NodeProto* node = graph->add_node(); - node->set_op_type("Identity"); - node->set_name(outName + "/out"); - node->add_input(tensorName); + if(needsContiguous) { + vector outShape = {-1, (int64_t)channels, (int64_t)nnYLen, (int64_t)nnXLen}; + string shapeName = b.addInt64Initializer(outName + "/contigshape", outShape); + node->set_op_type("Reshape"); + node->set_name(outName + "/out"); + node->add_input(tensorName); + node->add_input(shapeName); + } + else { + node->set_op_type("Identity"); + node->set_name(outName + "/out"); + node->add_input(tensorName); + } node->add_output(outName); onnx::ValueInfoProto* vi = graph->add_output(); vi->set_name(outName); diff --git a/cpp/tests/testsearchcommon.cpp b/cpp/tests/testsearchcommon.cpp index 8892180cf6..75e6dc8bdd 100644 --- a/cpp/tests/testsearchcommon.cpp +++ b/cpp/tests/testsearchcommon.cpp @@ -1,4 +1,5 @@ #include "../tests/testsearchcommon.h" +#include #include "../dataio/sgf.h" #include "../search/searchnode.h" @@ -206,6 +207,14 @@ NNEvaluator* TestSearchCommon::startNNEval( //NHWC layout is no longer a generic NNEvaluator option; only the CUDA backend reads it (off cfg). //Route the test's useNHWC param into a cudaUseNHWC override so it still drives the CUDA layout. cfg.overrideKey("cudaUseNHWC", useNHWC ? "true" : "false"); + //The MIGraphX backend reads migraphxTransformerNHWC off the cfg, but this test builds an + //empty ConfigParser, so an NHWC A/B run through runnnonmanyposestest silently compares + //NCHW against NCHW. Route an env var in so the layout can actually be toggled under test. + { + const char* nhwcEnv = std::getenv("MIGRAPHX_TRANSFORMER_NHWC"); + if(nhwcEnv != NULL && string(nhwcEnv) != "") + cfg.overrideKey("migraphxTransformerNHWC", string(nhwcEnv)); + } int numNNServerThreadsPerModel = 1; bool nnRandomize = false; string nnRandSeed = "runSearchTestsRandSeed"+seed;