Add MIGraphX backend for AMD GPUs (ROCm) — 2.94x over OpenCL on MI300X

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.

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

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)