path: root/PKGBUILD

# Mainteiner: piernov <piernov@piernov.org>
# Contributor: Sven-Hendrik Haase <svenstaro@archlinux.org>
# Contributor: Torsten Keßler <tpkessler@archlinux.org>
# Contributor: Jakub Klinkovský <lahwaacz at archlinux dot org>
# Contributor: Stephen Zhang <zsrkmyn at gmail dot com>

_pkgname=pytorch
pkgbase="python-${_pkgname}-cuda12.9"
pkgname=("${pkgbase}" "python-${_pkgname}-opt-cuda12.9")
# When updating pytorch, also check the compatibility table for torchvision
# https://github.com/pytorch/vision?tab=readme-ov-file#installation
pkgver=2.12.1
pkgrel=2
pkgdesc='Tensors and Dynamic neural networks in Python with strong GPU acceleration (Maxwell/Pascal/Volta support)'
arch=('x86_64')
url="https://pytorch.org"
license=('BSD-3-Clause-Modification')
depends=(
  abseil-cpp
  eigen
  gcc14-libs
  gflags
  glibc
  google-glog
  intel-oneapi-mkl
  libgcc
  libgomp
  libstdc++
  libuv
  numactl
  openmp
  openmpi
  protobuf
  pybind11
  python
  python-filelock
  python-fsspec
  python-jinja
  python-networkx
  python-numpy
  python-sympy
  python-typing_extensions
  python-yaml
  qt6-base
  vulkan-icd-loader
)
# https://github.com/ROCm/aotriton/blob/main/requirements-dev.txt
makedepends=(
  cmake
  cuda-12.9
  cudnn9.10-cuda12.9
  doxygen
  gcc14
  git
  nccl-cuda12.9
  ninja
  onednn
  pkgconfig
  python-build
  python-installer
  python-setuptools
  python-triton
  python-yaml
  shaderc
  vulkan-headers
)
optdepends=(
  'python-onnxscript: for the ONNX exporter'
)
source=("${_pkgname}::git+https://github.com/pytorch/pytorch.git#tag=v$pkgver"
        # generated using parse-submodules
        "${pkgname}-FP16::git+https://github.com/Maratyszcza/FP16.git"
        "${pkgname}-FXdiv::git+https://github.com/Maratyszcza/FXdiv.git"
        "${pkgname}-NNPACK::git+https://github.com/Maratyszcza/NNPACK.git"
        "${pkgname}-NVTX::git+https://github.com/NVIDIA/NVTX.git"
        "${pkgname}-PeachPy::git+https://github.com/malfet/PeachPy.git"
        "${pkgname}-VulkanMemoryAllocator::git+https://github.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator.git"
        "${pkgname}-aiter::git+https://github.com/ROCm/aiter.git"
        "${pkgname}-XNNPACK::git+https://github.com/google/XNNPACK.git"
        "${pkgname}-benchmark::git+https://github.com/google/benchmark.git"
        "${pkgname}-composable_kernel::git+https://github.com/ROCm/composable_kernel.git"
        "${pkgname}-cpp-httplib::git+https://github.com/yhirose/cpp-httplib.git"
        "${pkgname}-cpuinfo::git+https://github.com/pytorch/cpuinfo.git"
        "${pkgname}-cudnn-frontend::git+https://github.com/NVIDIA/cudnn-frontend.git"
        "${pkgname}-cutlass::git+https://github.com/NVIDIA/cutlass.git"
        "${pkgname}-fbgemm::git+https://github.com/pytorch/fbgemm"
        "${pkgname}-fbjni::git+https://github.com/facebookincubator/fbjni.git"
        "${pkgname}-flash-attention::git+https://github.com/Dao-AILab/flash-attention.git"
        "${pkgname}-flatbuffers::git+https://github.com/google/flatbuffers.git"
        "${pkgname}-fmt::git+https://github.com/fmtlib/fmt.git"
        "${pkgname}-gemmlowp::git+https://github.com/google/gemmlowp.git"
        "${pkgname}-gloo::git+https://github.com/pytorch/gloo"
        "${pkgname}-googletest::git+https://github.com/google/googletest.git"
        "${pkgname}-ideep::git+https://github.com/intel/ideep"
        "${pkgname}-ittapi::git+https://github.com/intel/ittapi.git"
        "${pkgname}-json::git+https://github.com/nlohmann/json.git"
        "${pkgname}-kineto::git+https://github.com/pytorch/kineto"
        "${pkgname}-kleidiai::git+https://github.com/ARM-software/kleidiai.git"
        "${pkgname}-mimalloc::git+https://github.com/microsoft/mimalloc.git"
        "${pkgname}-onnx::git+https://github.com/onnx/onnx.git"
        "${pkgname}-opentelemetry-cpp::git+https://github.com/open-telemetry/opentelemetry-cpp.git"
        "${pkgname}-pocketfft::git+https://github.com/mreineck/pocketfft"
        "${pkgname}-protobuf::git+https://github.com/protocolbuffers/protobuf.git"
        "${pkgname}-psimd::git+https://github.com/Maratyszcza/psimd.git"
        "${pkgname}-pthreadpool::git+https://github.com/Maratyszcza/pthreadpool.git"
        "${pkgname}-pybind11::git+https://github.com/pybind/pybind11.git"
        "${pkgname}-sleef::git+https://github.com/shibatch/sleef"
        "${pkgname}-tensorpipe::git+https://github.com/pytorch/tensorpipe.git"
        fix_include_system.patch
        use-system-libuv.patch
        87773.patch
        glog-0.7.patch
        fix_cmake_prefix_path.patch
        add_gpu_targets_rocm.patch
        aotriton_disable_install.patch
        pyproject.patch
        )
b2sums=('98f5a0196764acd7b8595f43473a544f3b515f25546a526a58b0fa7f57854025990860efa4e3d765dadb94b308b40c5b0b5344a13dd23f62207d15f0956d6214'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        'SKIP'
        '8c91e931f34fcc1cc5d916572e9808435bbc4a853d80af92a5484b4bb1111ef3b2a1c58283b97c2c3071de3177a124669004855f7ebad0734b0941714e3baf4f'
        'af8c724ed80898ae3875a295ad6bd4d18d90f8a9124f6cff6d1b2f525bf7806fe61306e739c1f7362fbd8d0e4f8ba57d0e3bf925ea3f7a78a0a98f26722db147'
        '0a8fc110a306e81beeb9ddfb3a1ddfd26aeda5e3f7adfb0f7c9bc3fd999c2dde62e0b407d3eca573097a53fd97329214e30e8767fb38d770197c7ec2b53daf18'
        '20d044c5c80354af5ed63847fa4332e96cbfc32a351788f6458fb92b322de7f64b10c188ff26e4f34e422cfe30e082c3ca23ee3e9094616c142aa53588dd451e'
        'eb1a4305c9e753774ce27256f8e7f35ae52986c8dfefddb71062f7abc71eec04eaae80cd03b9cb362150465000728390b7bfd0e539f772761c0a8d5dd8dbe980'
        '007fc33064c55b1a080f8c3dcb0c03acc21629d7034426d0622b56ace3936ae07e0f4bca578327542fa3333cc127ef2e2379ebc8e1f97b561ee54de58ce84d3c'
        'ec9aea1481c6ae85288d7ab7c709af80ab919face22c17710cfadd80f07111fe53c3241f278fc76c43f28813581a4be0280a5590f8a8fd6dd6b46bc8d2ea25e0'
        '21234592e20b5ff1bf43f926bdda72b089cc2b32b92d4287e5aa6b20dc8ebbb2e30135ba8b881d64c35d98be457366bcfd9982cf2e38fd3fc13901955fa571da')
options=('!lto' '!debug')

get_pyver () {
  python -c 'import sys; print(str(sys.version_info[0]) + "." + str(sys.version_info[1]))'
}

prepare() {
  cd "${srcdir}/${_pkgname}"

  # generated using parse-submodules
  git submodule init

  git config submodule."android/libs/fbjni".url "${srcdir}/${pkgname}"-fbjni
  git config submodule."third_party/NNPACK".url "${srcdir}/${pkgname}"-NNPACK
  git config submodule."third_party/NNPACK_deps/FP16".url "${srcdir}/${pkgname}"-FP16
  git config submodule."third_party/NNPACK_deps/FXdiv".url "${srcdir}/${pkgname}"-FXdiv
  git config submodule."third_party/NNPACK_deps/psimd".url "${srcdir}/${pkgname}"-psimd
  git config submodule."third_party/NNPACK_deps/pthreadpool".url "${srcdir}/${pkgname}"-pthreadpool
  git config submodule."third_party/NVTX".url "${srcdir}/${pkgname}"-NVTX
  git config submodule."third_party/VulkanMemoryAllocator".url "${srcdir}/${pkgname}"-VulkanMemoryAllocator
  git config submodule."third_party/XNNPACK".url "${srcdir}/${pkgname}"-XNNPACK
  git config submodule."third_party/aiter".url "${srcdir}/${pkgname}"-aiter
  git config submodule."third_party/benchmark".url "${srcdir}/${pkgname}"-benchmark
  git config submodule."third_party/composable_kernel".url "${srcdir}/${pkgname}"-composable_kernel
  git config submodule."third_party/cpp-httplib".url "${srcdir}/${pkgname}"-cpp-httplib
  git config submodule."third_party/cpuinfo".url "${srcdir}/${pkgname}"-cpuinfo
  git config submodule."third_party/cudnn_frontend".url "${srcdir}/${pkgname}"-cudnn-frontend
  git config submodule."third_party/cutlass".url "${srcdir}/${pkgname}"-cutlass
  git config submodule."third_party/fbgemm".url "${srcdir}/${pkgname}"-fbgemm
  git config submodule."third_party/flash-attention".url "${srcdir}/${pkgname}"-flash-attention
  git config submodule."third_party/flatbuffers".url "${srcdir}/${pkgname}"-flatbuffers
  git config submodule."third_party/fmt".url "${srcdir}/${pkgname}"-fmt
  git config submodule."third_party/gemmlowp/gemmlowp".url "${srcdir}/${pkgname}"-gemmlowp
  git config submodule."third_party/gloo".url "${srcdir}/${pkgname}"-gloo
  git config submodule."third_party/googletest".url "${srcdir}/${pkgname}"-googletest
  git config submodule."third_party/ideep".url "${srcdir}/${pkgname}"-ideep
  git config submodule."third_party/ittapi".url "${srcdir}/${pkgname}"-ittapi
  git config submodule."third_party/kineto".url "${srcdir}/${pkgname}"-kineto
  git config submodule."third_party/kleidiai".url "${srcdir}/${pkgname}"-kleidiai
  git config submodule."third_party/mimalloc".url "${srcdir}/${pkgname}"-mimalloc
  git config submodule."third_party/nlohmann".url "${srcdir}/${pkgname}"-json
  git config submodule."third_party/onnx".url "${srcdir}/${pkgname}"-onnx
  git config submodule."third_party/opentelemetry-cpp".url "${srcdir}/${pkgname}"-opentelemetry-cpp
  git config submodule."third_party/pocketfft".url "${srcdir}/${pkgname}"-pocketfft
  git config submodule."third_party/protobuf".url "${srcdir}/${pkgname}"-protobuf
  git config submodule."third_party/pybind11".url "${srcdir}/${pkgname}"-pybind11
  git config submodule."third_party/python-peachpy".url "${srcdir}/${pkgname}"-PeachPy
  git config submodule."third_party/sleef".url "${srcdir}/${pkgname}"-sleef
  git config submodule."third_party/tensorpipe".url "${srcdir}/${pkgname}"-tensorpipe

  git -c protocol.file.allow=always submodule update --init --recursive

  # Fix cmake prefix path (FS#78665)
  patch -Np1 -i "${srcdir}"/fix_cmake_prefix_path.patch

  # Use system libuv
  patch -Np1 -i "${srcdir}"/use-system-libuv.patch

  # Fix building against glog 0.6
  patch -Np1 -i "${srcdir}/87773.patch"

  # Fix building against glog 0.7
  patch -p1 -i "${srcdir}/glog-0.7.patch"

  # ROCm 6.4.2+ requires architectures to appear as cmake arguments too
  patch -p1 -i "${srcdir}/add_gpu_targets_rocm.patch"

  # If using prebuilt aotriton, pytorch attempts to copy /usr/lib and /user/include
  # into the torch folder. Disable this behavior.
  patch -p1 -i "${srcdir}/aotriton_disable_install.patch"

  # Avoid using /usr/include along with -isystem
  # https://bugs.archlinux.org/task/64981
  # https://gitlab.archlinux.org/archlinux/packaging/packages/python-pytorch/-/issues/37
  patch -N torch/utils/cpp_extension.py "${srcdir}"/fix_include_system.patch

  # patch python dependencies in pyproject.toml
  patch -p1 -i "${srcdir}/pyproject.patch"

  cd "${srcdir}"

  cp -r "${_pkgname}" "${_pkgname}-cuda"
  mv "${_pkgname}" "${_pkgname}-opt-cuda"
}

# Common build configuration, called in the build() function.
_prepare() {
  export VERBOSE=1
  export PYTORCH_BUILD_VERSION="${pkgver}"
  export PYTORCH_BUILD_NUMBER=1

  # Check tools/setup_helpers/cmake.py, setup.py and CMakeLists.txt for a list of flags that can be set via env vars.
  export BUILD_TEST=OFF  # do not build tests
  export ATEN_NO_TEST=ON  # do not build ATen tests
  export USE_MKLDNN=ON
  export BUILD_CUSTOM_PROTOBUF=OFF
  export USE_GFLAGS=ON
  export USE_GLOG=ON
  export USE_VULKAN=ON
  export USE_OBSERVERS=ON
  # export USE_SYSTEM_LIBS=ON  # experimental, not all libs present in repos
  # USE_SYSTEM_ONNX=ON does not work and onnx itself should be removed from pytorch: https://github.com/pytorch/pytorch/issues/166546#issuecomment-3463370459
  export USE_NCCL=ON
  export USE_SYSTEM_NCCL=ON
  export USE_SYSTEM_PYBIND11=ON
  export USE_SYSTEM_EIGEN_INSTALL=ON
  export USE_GOLD_LINKER=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 BUILD_SPLIT_CUDA=ON  # modern preferred build, but splits libs and symbols, ABI break
  export USE_CUPTI_SO=ON  # make sure cupti.so is used as shared lib
  export CC=/usr/bin/gcc-14
  export CXX=/usr/bin/g++-14
  export CXXFLAGS+=" -Wno-error=maybe-uninitialized"
  export CUDAHOSTCXX="${NVCC_CCBIN}"
  export CUDA_HOST_COMPILER="${CUDAHOSTCXX}"
  export CUDA_HOME=/opt/cuda
  # hide build-time CUDA devices
  export CUDA_VISIBLE_DEVICES=""
  export CUDNN_LIB_DIR=/usr/lib
  export CUDNN_INCLUDE_DIR=/usr/include
  export TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
  # the list of supported archs is in torch/utils/cpp_extension.py
  # https://github.com/pytorch/pytorch/blob/f9074c7332c3dfd43fe39e8733ec98f7f29b3e61/torch/utils/cpp_extension.py#L2417-L2420
  # (note that 8.8 is not supported)
  export TORCH_CUDA_ARCH_LIST="5.2 5.3 6.0 6.1 6.2 7.0 7.2"

  # Composable kernels is not supported for all architectures.
  # https://github.com/pytorch/pytorch/issues/150187
  export USE_ROCM_CK_GEMM=OFF
  export USE_ROCM_CK_SDPA=OFF
  export USE_FBGEMM_GENAI=OFF

  # Compile source code for supported GPU archs in parallel (but using too many jobs is not helpful)
  export HIPCC_COMPILE_FLAGS_APPEND="-parallel-jobs=4 --gcc-install-dir=$(dirname $(gcc-14 -print-libgcc-file-name))"
  export HIPCC_LINK_FLAGS_APPEND="-parallel-jobs=4"

  export AOTRITON_INSTALLED_PREFIX=/usr

  # Fix build issues for onnx with cmake 4.0
  export CMAKE_POLICY_VERSION_MINIMUM=3.5

  # Fix CUDA and ROCm builds with glog (these macros are defined in /usr/lib/cmake/glog/glog-targets.cmake
  # but for some reason this target is not applied when building some *.cu and *.hip files)
  export NVCC_APPEND_FLAGS="-DGLOG_USE_GLOG_EXPORT -DGLOG_USE_GFLAGS"
  HIPCC_COMPILE_FLAGS_APPEND+=" -DGLOG_USE_GLOG_EXPORT -DGLOG_USE_GFLAGS"

  # Limit number of threads to avoid running out of memory
  export MAX_JOBS=20
}

build() {
  cd "${srcdir}/${_pkgname}-cuda"
  echo "Building with cuda and without non-x86-64 optimizations"
  _prepare
  export USE_CUDA=1
  export USE_CUDNN=1
  export USE_ROCM=0
  export USE_XPU=0
  cd "${srcdir}/${_pkgname}-cuda"
  echo "add_definitions(-march=x86-64)" >> cmake/MiscCheck.cmake
  python -m build --wheel --no-isolation

  cd "${srcdir}/${_pkgname}-opt-cuda"
  echo "Building with cuda and with non-x86-64 optimizations"
  _prepare
  export USE_CUDA=1
  export USE_CUDNN=1
  export USE_ROCM=0
  export USE_XPU=0
  echo "add_definitions(-march=x86-64-v3)" >> cmake/MiscCheck.cmake
  python -m build --wheel --no-isolation
}

_package() {
  python -m installer --destdir="$pkgdir" dist/*.whl

  install -Dm644 LICENSE "${pkgdir}/usr/share/licenses/${pkgname}/LICENSE"

  pytorchpath="usr/lib/python$(get_pyver)/site-packages/torch"
  install -d "${pkgdir}/usr/lib"

  # put CMake files in correct place
  mv "${pkgdir}/${pytorchpath}/share/cmake" "${pkgdir}/usr/lib/cmake"

  # put C++ API in correct place
  mv "${pkgdir}/${pytorchpath}/include" "${pkgdir}/usr/include"
  find "${pkgdir}/${pytorchpath}"/lib/ -type f,l \( -iname '*.so' -or -iname '*.so*' \) -print0 | while read -rd $'\0' _lib; do
    mv "${_lib}" "${pkgdir}"/usr/lib/
  done

  # Clean up duplicates with Arch packages
  rm -vrf $(ls -d "${pkgdir}"/usr/include/* | grep -v "/ATen$\|/caffe2$\|/c10$\|/torch$\|/tensorpipe$")

  # Python module is hardcoded so look there at runtime
  ln -s /usr/include "${pkgdir}/${pytorchpath}/include"
  find "${pkgdir}"/usr/lib -maxdepth 1 -type f,l \( -iname '*.so' -or -iname '*.so*' \) -print0 | while read -rd $'\0' _lib; do
    ln -s ${_lib#"$pkgdir"} "${pkgdir}/${pytorchpath}/lib/"
  done
}

package_python-pytorch-cuda12.9() {
  pkgdesc+=" (with CUDA 12.9)"
  depends+=(cuda-12.9 nccl-cuda12.9 cudnn9.10-cuda12.9 onednn)
  conflicts=(python-pytorch)
  provides=(python-pytorch=${pkgver} python-pytorch-cuda=${pkgver})

  cd "${srcdir}/${_pkgname}-cuda"
  _package
}

package_python-pytorch-opt-cuda12.9() {
  pkgdesc+=" (with CUDA 12.9 and AVX2 CPU optimizations)"
  depends+=(cuda-12.9 nccl-cuda12.9 cudnn9.10-cuda12.9 onednn)
  conflicts=(python-pytorch)
  provides=(python-pytorch=${pkgver} python-pytorch-cuda=${pkgver} python-pytorch-cuda12.9=${pkgver})

  cd "${srcdir}/${_pkgname}-opt-cuda"
  _package
}

# vim:set ts=2 sw=2 et: