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

Compiling.md

@@ -34,6 +34,7 @@ As also mentioned in the instructions below but repeated here for visibility, if
       * 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 ROCm backend, ROCm 6.4 or later (https://rocm.docs.amd.com/projects/install-on-linux/en/latest/) and a GPU capable of supporting it. Install the ROCm developer packages, not just the ROCm runtime packages.
+      * 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.
@@ -43,6 +44,7 @@ As also mentioned in the instructions below but repeated here for visibility, if
    * 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` or `cmake . -DUSE_BACKEND=ROCM` 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).
@@ -90,6 +92,30 @@ As also mentioned in the instructions below but repeated here for visibility, if
         RDNA1/RDNA2 always use the built-in kernel, as does any GPU if the fused path isn't
         available or is explicitly disabled with `rocmDisableFusedAttention = true` in the config.
         See `cpp/external/composable_kernel_fmha/README.md` for details.
+### 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:

cpp/CMakeLists.txt

@@ -553,7 +553,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 ONNX ROCM)
+set_property(CACHE USE_BACKEND PROPERTY STRINGS "" CUDA TENSORRT OPENCL EIGEN METAL ONNX ROCM MIGRAPHX)
 
 set(USE_TCMALLOC 0 CACHE BOOL "Use TCMalloc")
 set(NO_GIT_REVISION 0 CACHE BOOL "Disable embedding the git revision into the compiled exe")
@@ -737,6 +737,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")
@@ -1151,6 +1156,64 @@ elseif(USE_BACKEND STREQUAL "TENSORRT")
     message(FATAL_ERROR "${ColorBoldRed} libnvonnxparser was NOT found in the TensorRT lib dir. ${ColorReset}")
   endif()
   target_link_libraries(katago ${TENSORRT_ONNXPARSER_LIBRARY})
+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

@@ -260,6 +260,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)
@@ -305,6 +307,8 @@ string Version::getGitRevisionWithBackend() {
   s += "-trt";
 #elif defined(USE_ROCM_BACKEND)
   s += "-rocm";
+#elif defined(USE_MIGRAPHX_BACKEND)
+  s += "-migraphx";
 #elif defined(USE_METAL_BACKEND)
   s += "-metal";
 #elif defined(USE_OPENCL_BACKEND)