[ROCm] Add AMD GPU (HIP) support to the RPUCuda backend

CMakeLists.txt

@@ -9,6 +9,7 @@ project(aihwkit C CXX)
 option(BUILD_TEST "Build C++ test binaries" OFF)
 option(BUILD_EXTENSION "Build additional C++ tools" OFF)
 option(USE_CUDA "Build with CUDA support" $ENV{USE_CUDA})
+option(USE_HIP "Build with ROCm/HIP support" $ENV{USE_HIP})
 
 # experimental precision flags
 option(RPU_USE_FP16 "EXPERIMENTAL: Build FP16 support (only available with CUDA)" OFF)
@@ -22,14 +23,23 @@ option(RPU_USE_FASTRAND "Use fastrand" OFF)
 option(RPU_USE_TORCH_BUFFERS "Use torch buffers for RPUCuda" ON)
 
 
+# The C++ standard tracks the linked PyTorch headers; newer Torch releases
+# require C++20. Default stays 17 to match older Torch.
+set(RPU_CXX_STANDARD "17" CACHE STRING "C++ standard used for all RPU targets")
+
 set(RPU_BLAS "OpenBLAS" CACHE STRING "BLAS backend of choice (OpenBLAS, MKL)")
 set(RPU_CUDA_ARCHITECTURES "75;80;89" CACHE STRING "Target CUDA architectures")
+# CMAKE_HIP_ARCHITECTURES is left to enable_language(HIP) to auto-detect the host
+# GPU when not set explicitly via -DCMAKE_HIP_ARCHITECTURES.
 
 # Internal variables.
 set(CUDA_TARGET_PROPERTIES POSITION_INDEPENDENT_CODE ON
                            CUDA_RESOLVE_DEVICE_SYMBOLS ON
                            CUDA_SEPARABLE_COMPILATION ON
-                           CXX_STANDARD 17)
+                           CXX_STANDARD ${RPU_CXX_STANDARD})
+set(HIP_TARGET_PROPERTIES POSITION_INDEPENDENT_CODE ON
+                          INTERPROCEDURAL_OPTIMIZATION OFF
+                          CXX_STANDARD ${RPU_CXX_STANDARD})
 
 # Append the virtualenv library path to cmake.
 if(DEFINED ENV{VIRTUAL_ENV})
@@ -41,13 +51,16 @@ endif()
 # Check for dependencies.
 include(cmake/dependencies.cmake)
 include(cmake/dependencies_cuda.cmake)
+include(cmake/dependencies_hip.cmake)
 include(cmake/dependencies_test.cmake)
 
 # Set compilation flags.
 if(WIN32)
-  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /O2")
-  set(CMAKE_CXX_STANDARD 17)
+  set(CMAKE_CXX_STANDARD ${RPU_CXX_STANDARD})
   set(CMAKE_CXX_STANDARD_REQUIRED ON)
+endif()
+if(MSVC)
+  set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} /O2")
 else()
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-narrowing -Wno-strict-overflow")
   set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -O3 -ftree-vectorize")
@@ -68,7 +81,7 @@ if(WIN32)
   target_link_libraries(RPU_CPU c10.lib torch_cpu.lib)
 endif()
 
-set_target_properties(RPU_CPU PROPERTIES CXX_STANDARD 17
+set_target_properties(RPU_CPU PROPERTIES CXX_STANDARD ${RPU_CXX_STANDARD}
   POSITION_INDEPENDENT_CODE ON)
 
 add_compile_definitions(RPU_USE_WITH_TORCH)
@@ -78,9 +91,12 @@ if (RPU_USE_DOUBLE)
 endif(RPU_USE_DOUBLE)
 
 
-if(USE_CUDA)
+if(USE_CUDA OR USE_HIP)
   add_subdirectory(src/rpucuda/cuda)
   include_directories(SYSTEM src/rpucuda/cuda)
+  if (USE_HIP)
+    set_source_files_properties(${RPU_GPU_SRCS} PROPERTIES LANGUAGE HIP)
+  endif()
   add_library(RPU_GPU ${RPU_GPU_SRCS})
 
   if (RPU_USE_FP16)
@@ -99,13 +115,19 @@ if(USE_CUDA)
     message(STATUS "Use FP16 parameters for CUDA (for all RPU number types).")
   endif(RPU_PARAM_FP16)
 
-  target_link_libraries(RPU_GPU RPU_CPU cublas curand ${RPU_DEPENDENCY_LIBS})
-  if(WIN32)
-    target_link_libraries(RPU_GPU c10_cuda.lib torch_cuda.lib)
-  endif(WIN32)
+  if (USE_HIP)
+    target_link_libraries(RPU_GPU RPU_CPU roc::hipblas hip::hiprand hip::hipcub ${RPU_DEPENDENCY_LIBS})
+    set_target_properties(RPU_GPU PROPERTIES ${HIP_TARGET_PROPERTIES})
+    set_property(TARGET RPU_GPU PROPERTY HIP_ARCHITECTURES ${CMAKE_HIP_ARCHITECTURES})
+  else()
+    target_link_libraries(RPU_GPU RPU_CPU cublas curand ${RPU_DEPENDENCY_LIBS})
+    if(WIN32)
+      target_link_libraries(RPU_GPU c10_cuda.lib torch_cuda.lib)
+    endif(WIN32)
 
-  set_target_properties(RPU_GPU PROPERTIES ${CUDA_TARGET_PROPERTIES})
-  set_property(TARGET RPU_GPU PROPERTY CUDA_ARCHITECTURES ${RPU_CUDA_ARCHITECTURES})
+    set_target_properties(RPU_GPU PROPERTIES ${CUDA_TARGET_PROPERTIES})
+    set_property(TARGET RPU_GPU PROPERTY CUDA_ARCHITECTURES ${RPU_CUDA_ARCHITECTURES})
+  endif()
 
   if (RPU_USE_TORCH_BUFFERS)
     if (BUILD_TEST)
@@ -114,18 +136,20 @@ if(USE_CUDA)
       set(RPU_USE_TORCH_BUFFERS OFF)
     else (BUILD_TEST)
       add_compile_definitions(RPU_TORCH_CUDA_BUFFERS)
-      set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
-      set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcudafe --diag_suppress=186")
+      if (NOT USE_HIP)
+        set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
+        set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcudafe --diag_suppress=186")
+      endif()
     endif(BUILD_TEST)
   endif(RPU_USE_TORCH_BUFFERS)
 
-  if(${CUDAToolkit_VERSION_MAJOR} LESS 11)
+  if(USE_CUDA AND CUDAToolkit_VERSION_MAJOR LESS 11)
     # The "cub" target only exists if cub was downloaded during build.
     if(TARGET cub)
         add_dependencies(RPU_GPU cub)
     endif()
   endif()
-endif(USE_CUDA)
+endif()
 
 # Add aihwkit targets.
 add_subdirectory(src/aihwkit/simulator)
@@ -136,7 +160,7 @@ if (BUILD_EXTENSION)
 
   add_library(AIHWKIT_EXTENSION_OPS ${AIHWKIT_EXTENSION_OPS_CPU_SRCS})
 
-  set_target_properties(AIHWKIT_EXTENSION_OPS PROPERTIES CXX_STANDARD 17
+  set_target_properties(AIHWKIT_EXTENSION_OPS PROPERTIES CXX_STANDARD ${RPU_CXX_STANDARD}
     POSITION_INDEPENDENT_CODE ON)
 
   target_link_libraries(AIHWKIT_EXTENSION_OPS
@@ -178,7 +202,7 @@ if (BUILD_EXTENSION)
   target_link_libraries(${extension_module_name} PRIVATE torch_python)
   target_include_directories(${extension_module_name} PRIVATE src/aihwkit/extension/extension_src)
   target_include_directories(${extension_module_name} PRIVATE src/aihwkit/extension/extension_src/ops)
-  set_target_properties(${extension_module_name} PROPERTIES CXX_STANDARD 17
+  set_target_properties(${extension_module_name} PROPERTIES CXX_STANDARD ${RPU_CXX_STANDARD}
     POSITION_INDEPENDENT_CODE ON)
 
   if (USE_CUDA)
@@ -211,7 +235,7 @@ if(BUILD_TEST)
     add_executable(${test_name} ${test_src})
     target_link_libraries(${test_name} gtest gmock)
     target_link_libraries(${test_name} torch_python c10 torch_cpu)
-    set_target_properties(${test_name} PROPERTIES CXX_STANDARD 17
+    set_target_properties(${test_name} PROPERTIES CXX_STANDARD ${RPU_CXX_STANDARD}
       POSITION_INDEPENDENT_CODE ON)
 
     if(WIN32)
@@ -222,14 +246,20 @@ if(BUILD_TEST)
     target_link_libraries(${test_name} RPU_CPU ${RPU_DEPENDENCY_LIBS})
 
     if(${test_src} IN_LIST RPU_GPU_TEST_SRCS)
-      target_link_libraries(${test_name} torch_cuda c10_cuda cudart)
-      target_link_libraries(${test_name} RPU_GPU RPU_CPU cublas curand ${RPU_DEPENDENCY_LIBS})
-      set_target_properties(${test_name} PROPERTIES ${CUDA_TARGET_PROPERTIES})
-      set_property(TARGET ${test_name} PROPERTY CUDA_ARCHITECTURES ${RPU_CUDA_ARCHITECTURES})
-
-      if(WIN32)
-	target_link_libraries(${test_name} c10_cuda.lib torch_cuda.lib)
-      endif(WIN32)
+      if(USE_HIP)
+        target_link_libraries(${test_name} RPU_GPU RPU_CPU roc::hipblas hip::hiprand hip::hipcub ${RPU_DEPENDENCY_LIBS})
+        set_target_properties(${test_name} PROPERTIES ${HIP_TARGET_PROPERTIES})
+        set_property(TARGET ${test_name} PROPERTY HIP_ARCHITECTURES ${CMAKE_HIP_ARCHITECTURES})
+      else()
+        target_link_libraries(${test_name} torch_cuda c10_cuda cudart)
+        target_link_libraries(${test_name} RPU_GPU RPU_CPU cublas curand ${RPU_DEPENDENCY_LIBS})
+        set_target_properties(${test_name} PROPERTIES ${CUDA_TARGET_PROPERTIES})
+        set_property(TARGET ${test_name} PROPERTY CUDA_ARCHITECTURES ${RPU_CUDA_ARCHITECTURES})
+
+        if(WIN32)
+	  target_link_libraries(${test_name} c10_cuda.lib torch_cuda.lib)
+        endif(WIN32)
+      endif()
 
     endif()
 

README.md

@@ -37,7 +37,7 @@ A series of primitives and features that allow using the toolkit within
 
 ### Analog devices simulator
 
-A high-performant (CUDA-capable) C++ simulator that allows for
+A high-performant (CUDA- and ROCm/HIP-capable) C++ simulator that allows for
 simulating a wide range of analog devices and crossbar configurations
 by using abstract functional models of material characteristics with
 adjustable parameters. Features include:

cmake/dependencies_hip.cmake

@@ -0,0 +1,47 @@
+# (C) Copyright 2020, 2021, 2022, 2023, 2024 IBM. All Rights Reserved.
+#
+# Licensed under the MIT license. See LICENSE file in the project root for details.
+
+# ROCm/HIP support, mirroring dependencies_cuda.cmake. The CUDA-spelled device
+# sources compile under HIP through the cuda_to_hip.h shim; RPU_USE_CUDA stays
+# defined so the device tile code and CudaAnalogTile pybind exposure are built.
+if(USE_HIP)
+  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})
+
+  enable_language(HIP)
+
+  find_package(hip REQUIRED)
+  find_package(hipblas REQUIRED)
+  find_package(hiprand REQUIRED)
+  find_package(hipcub REQUIRED)
+
+  add_compile_definitions(RPU_USE_CUDA)
+  add_compile_definitions(USE_HIP)
+
+  # Detect the torch hipify generation. hipify v1 renames CUDA spellings to HIP
+  # (c10::cuda::getCurrentCUDAStream is removed, only c10::hip::getCurrentHIPStream
+  # exists); hipify v2 masquerades, keeping c10::cuda::getCurrentCUDAStream as the
+  # public API while c10::hip::getCurrentHIPStream stays guarded by USE_ROCM (which
+  # this build does not define). rpu_base_tiles_cuda.cpp selects the c10 stream
+  # symbol from this, so probe the torch the build already uses.
+  execute_process(COMMAND "${RPU_PYTHON_EXECUTABLE}" -c "from torch.utils.hipify import __version__ as v; print(v)"
+      RESULT_VARIABLE TORCH_HIPIFY_VERSION_RESULT
+      OUTPUT_VARIABLE TORCH_HIPIFY_VERSION
+      OUTPUT_STRIP_TRAILING_WHITESPACE
+      ERROR_QUIET)
+  if(TORCH_HIPIFY_VERSION_RESULT EQUAL 0 AND NOT TORCH_HIPIFY_VERSION STREQUAL "")
+    message(STATUS "torch hipify version: ${TORCH_HIPIFY_VERSION}")
+    if(TORCH_HIPIFY_VERSION VERSION_GREATER_EQUAL "2.0.0")
+      set(RPU_TORCH_HIPIFY_V2 ON)
+    endif()
+  else()
+    message(STATUS "torch hipify version: unknown (assuming v1 stream API)")
+  endif()
+endif()

docs/source/developer_install.rst

@@ -102,6 +102,10 @@ Via python command
 
       $ python setup.py build_ext -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=TRUE --inplace -DRPU_BLAS=MKL -j16 -DUSE_CUDA=ON -DRPU_CUDA_ARCHITECTURES="60;70" -DCMAKE_PREFIX_PATH=$CONDA_PREFIX
 
+    - GPU with MKL (AMD, ROCm/HIP)::
+
+      $ python setup.py build_ext -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=TRUE --inplace -DRPU_BLAS=MKL -j16 -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx90a -DCMAKE_PREFIX_PATH=$CONDA_PREFIX
+
 
 If you want to use ``OpenBLAS`` instead ``MKL``, you need to set ``-DRPU_BLAS=OpenBLAS``.
 
@@ -111,6 +115,11 @@ To identify the ``CUDA_ARCH`` for your GPU using ``nvidia-smi`` in your system::
     $ export CUDA_ARCH=$(nvidia-smi --query-gpu=compute_cap --format=csv | sed -n '2 p' | tr -d '.')
     $ echo $CUDA_ARCH
 
+For an AMD GPU, set ``-DCMAKE_HIP_ARCHITECTURES`` to its ``gfx`` target instead
+(for example ``gfx90a`` or ``gfx1100``). To identify it::
+
+    $ rocminfo | grep -m1 -o 'gfx[0-9a-f]*'
+
 This will produce a shared library under the ``src/aihwkit/simulator``
 directory, without installing the package.  For how to use the shared library see the :ref:`use-the-library` section below.
 
@@ -222,6 +231,7 @@ compilation process:
 Flag                        Description                                       Default
 ==========================  ================================================  =======
 ``USE_CUDA``                Build with CUDA support                           ``OFF``
+``USE_HIP``                 Build with ROCm/HIP support (AMD GPUs)            ``OFF``
 ``BUILD_TEST``              Build the C++ test binaries                       ``OFF``
 ``RPU_BLAS``                BLAS backend of choice (``OpenBLAS`` or ``MKL``)  ``OpenBLAS``
 ``RPU_USE_FASTMOD``         Use fast mod                                      ``ON``
@@ -238,6 +248,18 @@ or if using ``cmake`` directly::
 
     build$ cmake -DUSE_CUDA=ON -DRPU_CUDA_ARCHITECTURES="60;70" ..
 
+For AMD GPUs, build with ROCm/HIP support instead (``USE_HIP`` and ``USE_CUDA``
+are mutually exclusive)::
+
+    $ python setup.py build_ext --inplace -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx90a
+
+or if using ``cmake`` directly::
+
+    build$ cmake -DUSE_HIP=ON -DCMAKE_HIP_ARCHITECTURES=gfx90a ..
+
+Set ``-DCMAKE_HIP_ARCHITECTURES`` to your GPU's ``gfx`` target (for example
+``gfx90a`` or ``gfx1100``).
+
 Passing other ``cmake`` flags
 """""""""""""""""""""""""""""
 

src/aihwkit/simulator/CMakeLists.txt

@@ -5,16 +5,35 @@
 set(python_module_name rpu_base)
 
 file(GLOB RPU_BINDINGS_SRCS rpu_base_src/*.cpp)
+if (USE_HIP)
+  # The binding TUs pull in HIP device headers transitively from RPU_GPU and
+  # inherit its --offload-arch usage requirement, which only the HIP (clang)
+  # driver understands; compile them with the HIP toolchain.
+  set_source_files_properties(${RPU_BINDINGS_SRCS} PROPERTIES LANGUAGE HIP)
+endif()
 pybind11_add_module(${python_module_name} MODULE ${RPU_BINDINGS_SRCS})
 target_link_libraries(${python_module_name} PRIVATE torch_python)
-set_target_properties(${python_module_name} PROPERTIES CXX_STANDARD 17)
+set_target_properties(${python_module_name} PROPERTIES CXX_STANDARD ${RPU_CXX_STANDARD})
 
-if (USE_CUDA)
+if (USE_CUDA OR USE_HIP)
   target_link_libraries(${python_module_name} PRIVATE RPU_GPU)
 else()
   target_link_libraries(${python_module_name} PRIVATE RPU_CPU)
 endif()
 
+if (USE_HIP)
+  # HIP device-code linking does not finalize LTO, which yields an empty
+  # PyInit_ and an ImportError; keep IPO off for the pybind module.
+  set_target_properties(${python_module_name} PROPERTIES INTERPROCEDURAL_OPTIMIZATION OFF)
+
+  # rpu_base_tiles_cuda.cpp picks the c10 current-stream symbol by hipify
+  # generation (see cmake/dependencies_hip.cmake); v2 keeps the c10::cuda
+  # spelling, v1 (undefined) uses c10::hip.
+  if(RPU_TORCH_HIPIFY_V2)
+    target_compile_definitions(${python_module_name} PRIVATE TORCH_HIPIFY_V2)
+  endif()
+endif()
+
 install(TARGETS ${python_module_name} DESTINATION "src/aihwkit/simulator")
 
 add_custom_command(TARGET ${python_module_name}

src/aihwkit/simulator/rpu_base_src/rpu_base_tiles_cuda.cpp

@@ -6,13 +6,39 @@
 
 #ifdef RPU_USE_CUDA
 #define __RPU_CUDA_HALF_DEFINED
+#if !defined(USE_HIP)
 #include "cuda.h"
+#endif
 #include "cuda_util.h"
 #include "rpu_base.h"
 
 #include "rpucuda.h"
 #include "rpucuda_pulsed.h"
+#if defined(USE_HIP)
+// The full ATen HIP context header transitively pulls in hipSOLVER / hipSPARSE
+// headers that are not part of this build; only the current-stream accessor is
+// needed, so include the lightweight c10 stream header and surface it under the
+// at::cuda namespace the call sites use.
+#include <c10/hip/HIPStream.h>
+namespace at {
+namespace cuda {
+inline auto getCurrentCUDAStream(c10::DeviceIndex device_index = -1) {
+#if defined(TORCH_HIPIFY_V2)
+  // torch hipify v2 masquerades: c10::cuda::getCurrentCUDAStream is the public
+  // stream API, while c10::hip::getCurrentHIPStream stays guarded by USE_ROCM,
+  // which this build does not define.
+  return c10::cuda::getCurrentCUDAStream(device_index);
+#else
+  // torch hipify v1 renames CUDA spellings to HIP: getCurrentCUDAStream is gone
+  // from c10::cuda, so use the hip-prefixed equivalent.
+  return c10::hip::getCurrentHIPStream(device_index);
+#endif
+}
+}
+} // namespace at
+#else
 #include <ATen/cuda/CUDAContext.h>
+#endif
 
 #define CHECK_CUDA(x)                                                                              \
   TORCH_CHECK(                                                                                     \