make changes based on community python package

6 files changed, 106 insertions, 11 deletions

diff --git a/.SRCINFO b/.SRCINFO
index 3b685e340cb2..60a4d86973c8 100644
--- a/.SRCINFO
+++ b/.SRCINFO
@@ -1,7 +1,7 @@
 pkgbase = python-pytorch-rocm
 	pkgdesc = Tensors and Dynamic neural networks in Python with strong GPU acceleration
 	pkgver = 1.7.0
-	pkgrel = 1
+	pkgrel = 2
 	url = https://pytorch.org
 	arch = x86_64
 	license = BSD
@@ -36,9 +36,17 @@ pkgbase = python-pytorch-rocm
 	depends = intel-mkl
 	source = pytorch-1.7.0::git+https://github.com/pytorch/pytorch.git#tag=v1.7.0
 	source = fix_include_system.patch
+	source = use-system-libuv.patch
+	source = use-system-libuv2.patch
+	source = nccl_version.patch
+	source = disable_non_x86_64.patch
 	source = find-hsa-runtime.patch::https://patch-diff.githubusercontent.com/raw/pytorch/pytorch/pull/45550.patch
 	sha256sums = SKIP
 	sha256sums = 147bdaeac8ec46ea46382e6146878bd8f8d51e05d5bd6f930dfd8e2b520859b9
+	sha256sums = 26b1dd596f1e21a011ee18cab939924483d6c6d4d98e543bf76f5a9312d54d67
+	sha256sums = 7b65c3b209fc39f92ba58a58be6d3da40799f1922910b1171ccd9209eda1f9eb
+	sha256sums = e4a96887b41cbdfd4204ce5f16fcb16a23558d23126331794ab6aa30a66f2e0d
+	sha256sums = d3ef8491718ed7e814fe63e81df2f49862fffbea891d2babbcb464796a1bd680
 	sha256sums = SKIP
 
 pkgname = python-pytorch-rocm
@@ -60,6 +68,8 @@ pkgname = python-pytorch-rocm
 	depends = onednn
 	depends = intel-mkl
 	depends = rocm
+	depends = rocm-libs
+	depends = miopen
 	depends = magma
 	provides = python-pytorch
 	conflicts = python-pytorch
@@ -83,6 +93,8 @@ pkgname = python-pytorch-opt-rocm
 	depends = onednn
 	depends = intel-mkl
 	depends = rocm
+	depends = rocm-libs
+	depends = miopen
 	depends = magma
 	provides = python-pytorch
 	provides = python-pytorch-rocm
diff --git a/PKGBUILD b/PKGBUILD
index c43dbc4a35eb..8feb46fa375b 100644
--- a/PKGBUILD
+++ b/PKGBUILD
@@ -7,20 +7,28 @@ pkgname=("python-pytorch-rocm" "python-pytorch-opt-rocm")
 _pkgname="pytorch"
 pkgver=1.7.0
 _pkgver=1.7.0
-pkgrel=1
+pkgrel=2
 pkgdesc="Tensors and Dynamic neural networks in Python with strong GPU acceleration"
 arch=('x86_64')
 url="https://pytorch.org"
 license=('BSD')
 depends=('google-glog' 'gflags' 'opencv' 'openmp' 'rccl' 'pybind11' 'python' 'python-yaml' 'libuv'
          'python-numpy' 'protobuf' 'ffmpeg' 'python-future' 'qt5-base' 'onednn' 'intel-mkl')
-makedepends=('python' 'python-setuptools' 'python-yaml' 'python-numpy' 'cmake' 'rocm' 'rocm-libs' 'miopen'
-             'git' 'magma' 'ninja' 'pkgconfig' 'doxygen')
+makedepends=('python' 'python-setuptools' 'python-yaml' 'python-numpy' 'cmake' 'rocm'
+             'rocm-libs' 'miopen' 'git' 'magma' 'ninja' 'pkgconfig' 'doxygen')
 source=("${_pkgname}-${pkgver}::git+https://github.com/pytorch/pytorch.git#tag=v$_pkgver"
         fix_include_system.patch
+        use-system-libuv.patch
+        use-system-libuv2.patch
+        nccl_version.patch
+        disable_non_x86_64.patch
         "find-hsa-runtime.patch::https://patch-diff.githubusercontent.com/raw/pytorch/pytorch/pull/45550.patch")
 sha256sums=('SKIP'
             '147bdaeac8ec46ea46382e6146878bd8f8d51e05d5bd6f930dfd8e2b520859b9'
+            '26b1dd596f1e21a011ee18cab939924483d6c6d4d98e543bf76f5a9312d54d67'
+            '7b65c3b209fc39f92ba58a58be6d3da40799f1922910b1171ccd9209eda1f9eb'
+            'e4a96887b41cbdfd4204ce5f16fcb16a23558d23126331794ab6aa30a66f2e0d'
+            'd3ef8491718ed7e814fe63e81df2f49862fffbea891d2babbcb464796a1bd680'
             'SKIP')
 
 prepare() {
@@ -36,6 +44,12 @@ prepare() {
   # https://bugs.archlinux.org/task/64981
   patch -N torch/utils/cpp_extension.py "${srcdir}"/fix_include_system.patch
 
+  # Use system libuv
+  patch -Np1 -i "${srcdir}"/use-system-libuv.patch
+
+  # FindNCCL patch to export correct nccl version
+  # patch -Np1 -i "${srcdir}"/nccl_version.patch
+
   # https://github.com/pytorch/pytorch/pull/45550
   patch -Np1 -i "${srcdir}"/find-hsa-runtime.patch
 
@@ -53,22 +67,23 @@ prepare() {
 
   # Check tools/setup_helpers/cmake.py, setup.py and CMakeLists.txt for a list of flags that can be set via env vars.
   export USE_MKLDNN=ON
-  # export BUILD_CUSTOM_PROTOBUF=OFF
+  export BUILD_CUSTOM_PROTOBUF=OFF
   # export BUILD_SHARED_LIBS=OFF
   export USE_FFMPEG=ON
   export USE_GFLAGS=ON
   export USE_GLOG=ON
   export BUILD_BINARY=ON
-  #export USE_OPENCV=ON
+  # export USE_OPENCV=ON
   export USE_SYSTEM_NCCL=ON
+  # export USE_SYSTEM_LIBS=ON
   export NCCL_VERSION=$(pkg-config nccl --modversion)
   export NCCL_VER_CODE=$(sed -n 's/^#define NCCL_VERSION_CODE\s*\(.*\).*/\1/p' /usr/include/nccl.h)
-  export CUDAHOSTCXX=g++-9
+  export CUDAHOSTCXX=g++
   export CUDA_HOME=/opt/cuda
   export CUDNN_LIB_DIR=/usr/lib
   export CUDNN_INCLUDE_DIR=/usr/include
-  export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
-  export TORCH_CUDA_ARCH_LIST="5.2;5.3;6.0;6.0+PTX;6.1;6.1+PTX;6.2;6.2+PTX;7.0;7.0+PTX;7.2;7.2+PTX;7.5;7.5+PTX;8.0;8.0+PTX;"
+  # export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+  export TORCH_CUDA_ARCH_LIST="5.2;5.3;6.0;6.1;6.2;7.0;7.0+PTX;7.2;7.2+PTX;7.5;7.5+PTX;8.0;8.0+PTX;8.6;8.6+PTX"
 }
 
 build() {
@@ -76,6 +91,8 @@ build() {
   export USE_CUDA=OFF
   export USE_ROCM=ON
   cd "${srcdir}/${_pkgname}-${pkgver}-rocm"
+  patch -Np1 -i "${srcdir}/disable_non_x86_64.patch"
+  echo "add_definitions(-march=x86-64)" >> cmake/MiscCheck.cmake
 
   # Apply changes needed for ROCm
   python tools/amd_build/build_amd.py
@@ -128,7 +145,7 @@ _package() {
 
 package_python-pytorch-rocm() {
   pkgdesc="Tensors and Dynamic neural networks in Python with strong GPU acceleration (with ROCM)"
-  depends+=(rocm magma)
+  depends+=(rocm rocm-libs miopen magma)
   conflicts=(python-pytorch)
   provides=(python-pytorch)
 
@@ -138,7 +155,7 @@ package_python-pytorch-rocm() {
 
 package_python-pytorch-opt-rocm() {
   pkgdesc="Tensors and Dynamic neural networks in Python with strong GPU acceleration (with ROCM and CPU optimizations)"
-  depends+=(rocm magma)
+  depends+=(rocm rocm-libs miopen magma)
   conflicts=(python-pytorch)
   provides=(python-pytorch python-pytorch-rocm)
 
diff --git a/disable_non_x86_64.patch b/disable_non_x86_64.patch
new file mode 100644
index 000000000000..3b1d380a8502
--- /dev/null
+++ b/disable_non_x86_64.patch
@@ -0,0 +1,15 @@
+diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
+index d5db749d15..fd54cca6b8 100644
+--- a/tools/setup_helpers/cmake.py
++++ b/tools/setup_helpers/cmake.py
+@@ -295,6 +295,10 @@ class CMake:
+         build_options.update(cmake__options)
+ 
+         CMake.defines(args,
++                      DISABLE_AVX2=1,
++                      DISABLE_AVX512F=1,
++                      DISABLE_FMA4=1,
++                      DISABLE_SSE4=1,
+                       PYTHON_EXECUTABLE=sys.executable,
+                       PYTHON_LIBRARY=cmake_python_library,
+                       PYTHON_INCLUDE_DIR=distutils.sysconfig.get_python_inc(),
diff --git a/nccl_version.patch b/nccl_version.patch
new file mode 100644
index 000000000000..c0cd0005c215
--- /dev/null
+++ b/nccl_version.patch
@@ -0,0 +1,25 @@
+diff --git a/cmake/Modules/FindNCCL.cmake b/cmake/Modules/FindNCCL.cmake
+index a16c9aca67..51eb54f14c 100644
+--- a/cmake/Modules/FindNCCL.cmake
++++ b/cmake/Modules/FindNCCL.cmake
+@@ -55,9 +55,10 @@ if(NCCL_FOUND)  # obtaining NCCL version and some sanity checks
+   set (OLD_CMAKE_REQUIRED_INCLUDES ${CMAKE_REQUIRED_INCLUDES})
+   list (APPEND CMAKE_REQUIRED_INCLUDES ${NCCL_INCLUDE_DIRS})
+   include(CheckCXXSymbolExists)
+-  check_cxx_symbol_exists(NCCL_VERSION_CODE nccl.h NCCL_VERSION_DEFINED)
++  set(NCCL_VERSION_CODE $ENV{NCCL_VER_CODE})
++  set(NCCL_VERSION_DEFINED $ENV{NCCL_VER_CODE})
+ 
+-  if (NCCL_VERSION_DEFINED)
++  if (DEFINED NCCL_VERSION_DEFINED)
+     set(file "${PROJECT_BINARY_DIR}/detect_nccl_version.cc")
+     file(WRITE ${file} "
+       #include <iostream>
+@@ -72,6 +73,7 @@ if(NCCL_FOUND)  # obtaining NCCL version and some sanity checks
+       }
+ ")
+     try_run(NCCL_VERSION_MATCHED compile_result ${PROJECT_BINARY_DIR} ${file}
++          CMAKE_FLAGS -DINCLUDE_DIRECTORIES=/opt/cuda/include
+           RUN_OUTPUT_VARIABLE NCCL_VERSION_FROM_HEADER
+           CMAKE_FLAGS  "-DINCLUDE_DIRECTORIES=${NCCL_INCLUDE_DIRS}"
+           LINK_LIBRARIES ${NCCL_LIBRARIES})
diff --git a/use-system-libuv.patch b/use-system-libuv.patch
new file mode 100644
index 000000000000..919d278310bd
--- /dev/null
+++ b/use-system-libuv.patch
@@ -0,0 +1,13 @@
+diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
+index 023bbe9e8d..e25aa0aa7f 100644
+--- a/cmake/Dependencies.cmake
++++ b/cmake/Dependencies.cmake
+@@ -1308,7 +1308,7 @@ if(USE_DISTRIBUTED AND USE_TENSORPIPE)
+   if(MSVC)
+     message(WARNING "Tensorpipe cannot be used on Windows.")
+   else()
+-    set(TP_BUILD_LIBUV ON CACHE BOOL "" FORCE)
++    set(TP_BUILD_LIBUV OFF CACHE BOOL "" FORCE)
+     set(TP_ENABLE_SHM OFF CACHE BOOL "" FORCE)
+     set(TP_ENABLE_CMA OFF CACHE BOOL "" FORCE)
+     set(TP_STATIC_OR_SHARED STATIC CACHE STRING "" FORCE)
diff --git a/use-system-libuv2.patch b/use-system-libuv2.patch
new file mode 100644
index 000000000000..b9a3cc190c05
--- /dev/null
+++ b/use-system-libuv2.patch
@@ -0,0 +1,13 @@
+diff --git a/cmake/pytorch.cmake b/cmake/pytorch.cmake
+index e4d6b9c..48e0669 100644
+--- a/cmake/pytorch.cmake
++++ b/cmake/pytorch.cmake
+@@ -138,7 +138,7 @@ endif()
+ 
+ add_library(tensorpipe ${TENSORPIPE_SRC})
+ 
+-set(TP_BUILD_LIBUV ON)
++set(TP_BUILD_LIBUV OFF)
+ find_package(uv REQUIRED)
+ target_link_libraries(tensorpipe PRIVATE uv::uv)
+