[WIP] Build PyTorch with ROCm

pytorch-cms-aotriton.patch

@@ -0,0 +1,15 @@
+diff --git a/cmake/External/aotriton.cmake b/cmake/External/aotriton.cmake
+index de64370b37..986ad6c178 100644
+--- a/cmake/External/aotriton.cmake
++++ b/cmake/External/aotriton.cmake
+@@ -5,8 +5,8 @@ if(NOT __AOTRITON_INCLUDED)
+   set(__AOTRITON_BUILD_DIR "${CMAKE_CURRENT_BINARY_DIR}/aotriton/build")
+   set(__AOTRITON_INSTALL_DIR "${PROJECT_SOURCE_DIR}/torch")
+   ExternalProject_Add(aotriton_external
+-    GIT_REPOSITORY https://github.com/ROCm/aotriton.git
+-    GIT_TAG 24a3fe9cb57e5cda3c923df29743f9767194cc27
++    GIT_REPOSITORY https://github.com/cms-externals/aotriton.git
++    GIT_TAG 763f12b4cfb4fcb1460a635dd76094589efb10e9
+     SOURCE_DIR ${__AOTRITON_SOURCE_DIR}
+     BINARY_DIR ${__AOTRITON_BUILD_DIR}
+     PREFIX ${__AOTRITON_INSTALL_DIR}

pytorch-hipcc-clang-flags.patch

@@ -0,0 +1,14 @@
+diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
+index a96075245a..aa4e22391e 100644
+--- a/cmake/Dependencies.cmake
++++ b/cmake/Dependencies.cmake
+@@ -1305,6 +1305,9 @@ if(USE_ROCM)
+     endif(CMAKE_BUILD_TYPE MATCHES Debug)
+
+     set(HIP_CLANG_FLAGS ${HIP_CXX_FLAGS})
++    # TODO: dynamic target
++    list(APPEND HIP_CLANG_FLAGS "--gcc-toolchain=$ENV{GCC_ROOT}")
++    list(APPEND HIP_CLANG_FLAGS "--target=x86_64-redhat-linux-gnu")
+     # Ask hcc to generate device code during compilation so we can use
+     # host linker to link.
+     list(APPEND HIP_CLANG_FLAGS -fno-gpu-rdc)

pytorch-rocm-version.patch

@@ -0,0 +1,14 @@
+diff --git a/aten/src/ATen/cuda/tunable/TunableGemm.h b/aten/src/ATen/cuda/tunable/TunableGemm.h
+index 3ba0d76127..775a04b8e8 100644
+--- a/aten/src/ATen/cuda/tunable/TunableGemm.h
++++ b/aten/src/ATen/cuda/tunable/TunableGemm.h
+@@ -22,7 +22,7 @@
+ #include <c10/util/StringUtil.h>
+
+ #ifdef USE_ROCM
+-#include <rocm-core/rocm_version.h>
++#include <rocm_version.h>
+ #endif
+
+ #define STRINGIFY(s) #s
+

pytorch-system-fmt.patch

@@ -1,15 +1,15 @@
 diff --git a/c10/CMakeLists.txt b/c10/CMakeLists.txt
 index feebad7cbb7..8a2ad840286 100644
---- a/c10/CMakeLists.txt
-+++ b/c10/CMakeLists.txt
-@@ -87,6 +87,7 @@ endif()
- if(${USE_GLOG})
-     target_link_libraries(c10 PUBLIC glog::glog)
+--- a/c10/CMakeLists.txt 2024-07-10 13:31:50.729017827 +0200
++++ b/c10/CMakeLists.txt  2024-07-10 13:32:36.649765887 +0200
+@@ -87,6 +87,7 @@
+ if(C10_USE_GLOG)
+   target_link_libraries(c10 PUBLIC glog::glog)
  endif()
 +find_package(fmt REQUIRED)
  target_link_libraries(c10 PRIVATE fmt::fmt-header-only)
 
- find_package(Backtrace)
+ if(C10_USE_NUMA)
 diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
 index 74d0d557190..ff3e94b1c4b 100644
 --- a/caffe2/CMakeLists.txt

pytorch.spec

@@ -1,40 +1,50 @@
-### RPM external pytorch 2.1.1
-## INCLUDE cuda-flags
+### RPM external pytorch 2.3.1
 ## INCLUDE microarch_flags
 
-%define cuda_arch_float $(echo %{cuda_arch} | tr ' ' '\\n' | sed -E 's|([0-9])$|.\\1|' | tr '\\n' ' ')
-%define tag bb938bbe9f53414dda1e1159795b7536dbffd041
+%define tag 63d5e9221bedd1546b7d364b5ce4171547db12a9
 %define branch cms/v%{realversion}
 
 Source: git+https://github.com/cms-externals/pytorch.git?obj=%{branch}/%{tag}&export=%{n}-%{realversion}&submodules=1&output=/%{n}-%{realversion}.tgz
 Source1: FindEigen3.cmake
 Source2: FindFMT.cmake
 Source99: scram-tools.file/tools/eigen/env
-Patch0: pytorch-ignore-different-cuda-include-dir
 Patch1: pytorch-missing-braces
 Patch2: pytorch-system-fmt
+Patch3: pytorch-cms-aotriton
+Patch4: pytorch-hipcc-clang-flags
+Patch5: pytorch-rocm-version
 
-BuildRequires: cmake ninja
+BuildRequires: cmake ninja py3-pip
 Requires: eigen fxdiv numactl openmpi protobuf psimd python3 py3-PyYAML
-Requires: OpenBLAS zlib protobuf fmt py3-pybind11
-%{!?without_cuda:Requires: cuda cudnn}
+Requires: OpenBLAS zlib protobuf fmt py3-pybind11 py3-typing-extensions
+Requires: py3-filelock py3-iniconfig py3-packaging py3-packaging py3-pluggy py3-numpy py3-setuptools
+%{!?without_rocm:Requires: rocm rocm-rocrand}
 
 %prep
 %setup -n %{n}-%{realversion}
-%patch0 -p1
 %patch1 -p1
 %patch2 -p1
+%patch3 -p1
+%patch4 -p1
+%patch5 -p1
 
 %build
+# Pregenerate some files
+python3 tools/amd_build/build_amd.py
+
 cp %{_sourcedir}/FindEigen3.cmake %{_sourcedir}/FindFMT.cmake cmake/Modules/
 rm -rf ../build && mkdir ../build && cd ../build
 source %{_sourcedir}/env
 
-USE_CUDA=OFF
-%if "%{cmsos}" != "slc7_aarch64"
-if [ "%{cuda_gcc_support}" = "true" ] ; then
-USE_CUDA=%{!?without_cuda:ON}
-fi
+%if 0%{!?without_rocm:1}
+# Notice: must be environment variables
+export ROCM_PATH=${ROCM_ROOT}
+export ROCM_SOURCE_DIR=${ROCM_ROOT}
+export PYTORCH_ROCM_ARCH=gfx900,gfx906,gfx908,gfx90a,gfx1030
+export PYBIND11_SYSPATH=${PY3_PYBIND11_ROOT}
+export TRITON_CACHE_DIR=$WORKSPACE/.triton/cache
+export GCC_ROOT=${GCC_ROOT}
+hipcc -v
 %endif
 
 cmake ../%{n}-%{realversion} \
@@ -44,11 +54,9 @@ cmake ../%{n}-%{realversion} \
     -DBUILD_TEST=OFF \
     -DBUILD_BINARY=OFF \
     -DBUILD_PYTHON=OFF \
-%if 0%{!?without_cuda:1}
-    -DUSE_CUDA=${USE_CUDA} \
-    -DTORCH_CUDA_ARCH_LIST="%{cuda_arch_float}" \
-    -DCUDNN_INCLUDE_DIR=${CUDNN_ROOT}/include \
-    -DCUDNN_LIBRARY=${CUDNN_ROOT}/lib64/libcudnn.so \
+    -DUSE_CUDA=OFF \
+%if 0%{!?without_rocm:1}
+    -DUSE_ROCM=ON \
 %endif
     -DUSE_NCCL=OFF \
     -DUSE_FBGEMM=OFF \
@@ -94,6 +102,4 @@ ninja -v  %{makeprocesses} install
 %{relocateConfig}include/caffe2/core/macros.h
 %{relocateConfig}share/cmake/ATen/ATenConfig.cmake
 
-# For ROCm, pre-build
-# NOTICE: can't build with both cuda and rocm
-# python @{_builddir}/tools/amd_build/build_amd.py
+# NOTICE: can't build with both cuda and rocm - see aten/CMakeLists.txt

rocm.spec

@@ -29,6 +29,31 @@ Source12: https://%{repository}/%{repoversion}/main/openmp-extras-runtime-16.56.
 Source13: https://%{repository}/%{repoversion}/main/rocm-openmp-sdk-5.6.1.50601-93.el%{rhel}.%{_arch}.rpm
 Source14: https://%{repository}/%{repoversion}/main/rocprim-devel-2.13.0.50601-93.el%{rhel}.%{_arch}.rpm
 Source15: https://%{repository}/%{repoversion}/main/rocthrust-devel-2.18.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source16: https://%{repository}/%{repoversion}/main/hsa-rocr-devel-1.9.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source17: https://%{repository}/%{repoversion}/main/rocblas-devel-3.0.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source18: https://%{repository}/%{repoversion}/main/rocblas-3.0.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source19: https://%{repository}/%{repoversion}/main/hipblas-1.0.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source20: https://%{repository}/%{repoversion}/main/hipblas-devel-1.0.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source21: https://%{repository}/%{repoversion}/main/hipblaslt-0.2.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source22: https://%{repository}/%{repoversion}/main/hipblaslt-devel-0.2.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source23: https://%{repository}/%{repoversion}/main/miopen-hip-2.20.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source24: https://%{repository}/%{repoversion}/main/miopen-hip-devel-2.20.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source25: https://%{repository}/%{repoversion}/main/rocfft-1.0.23.50601-93.el%{rhel}.%{_arch}.rpm
+Source26: https://%{repository}/%{repoversion}/main/rocfft-devel-1.0.23.50601-93.el%{rhel}.%{_arch}.rpm
+Source27: https://%{repository}/%{repoversion}/main/hipfft-1.0.12.50601-93.el%{rhel}.%{_arch}.rpm
+Source28: https://%{repository}/%{repoversion}/main/hipfft-devel-1.0.12.50601-93.el%{rhel}.%{_arch}.rpm
+Source29: https://%{repository}/%{repoversion}/main/hipsparse-2.3.7.50601-93.el%{rhel}.%{_arch}.rpm
+Source30: https://%{repository}/%{repoversion}/main/hipsparse-devel-2.3.7.50601-93.el%{rhel}.%{_arch}.rpm
+Source31: https://%{repository}/%{repoversion}/main/rccl-2.16.5.50601-93.el%{rhel}.%{_arch}.rpm
+Source32: https://%{repository}/%{repoversion}/main/rccl-devel-2.16.5.50601-93.el%{rhel}.%{_arch}.rpm
+Source33: https://%{repository}/%{repoversion}/main/rocprim-devel-2.13.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source34: https://%{repository}/%{repoversion}/main/hipcub-devel-2.13.1.50601-93.el%{rhel}.%{_arch}.rpm
+Source35: https://%{repository}/%{repoversion}/main/rocthrust-devel-2.18.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source36: https://%{repository}/%{repoversion}/main/hipsolver-1.8.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source37: https://%{repository}/%{repoversion}/main/hipsolver-devel-1.8.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source38: https://%{repository}/%{repoversion}/main/roctracer-4.1.0.50601-93.el%{rhel}.%{_arch}.rpm
+Source39: https://%{repository}/%{repoversion}/main/roctracer-devel-4.1.0.50601-93.el%{rhel}.%{_arch}.rpm
+
 Requires: numactl zstd
 Requires: python3
 AutoReq: no
@@ -52,6 +77,30 @@ rpm2cpio %{SOURCE12} | cpio -idmv
 rpm2cpio %{SOURCE13} | cpio -idmv
 rpm2cpio %{SOURCE14} | cpio -idmv
 rpm2cpio %{SOURCE15} | cpio -idmv
+rpm2cpio %{SOURCE16} | cpio -idmv
+rpm2cpio %{SOURCE17} | cpio -idmv
+rpm2cpio %{SOURCE18} | cpio -idmv
+rpm2cpio %{SOURCE19} | cpio -idmv
+rpm2cpio %{SOURCE20} | cpio -idmv
+rpm2cpio %{SOURCE21} | cpio -idmv
+rpm2cpio %{SOURCE22} | cpio -idmv
+rpm2cpio %{SOURCE23} | cpio -idmv
+rpm2cpio %{SOURCE24} | cpio -idmv
+rpm2cpio %{SOURCE25} | cpio -idmv
+rpm2cpio %{SOURCE26} | cpio -idmv
+rpm2cpio %{SOURCE27} | cpio -idmv
+rpm2cpio %{SOURCE28} | cpio -idmv
+rpm2cpio %{SOURCE29} | cpio -idmv
+rpm2cpio %{SOURCE30} | cpio -idmv
+rpm2cpio %{SOURCE31} | cpio -idmv
+rpm2cpio %{SOURCE32} | cpio -idmv
+rpm2cpio %{SOURCE33} | cpio -idmv
+rpm2cpio %{SOURCE34} | cpio -idmv
+rpm2cpio %{SOURCE35} | cpio -idmv
+rpm2cpio %{SOURCE36} | cpio -idmv
+rpm2cpio %{SOURCE37} | cpio -idmv
+rpm2cpio %{SOURCE38} | cpio -idmv
+rpm2cpio %{SOURCE39} | cpio -idmv
 
 %install
 rmdir %{i}