_pkgname=pytorch
pkgbase="python-${_pkgname}"
pkgname=("${pkgbase}-rocm")
pkgver=2.1.0
_pkgver=${pkgver}
pkgrel=1
_pkgdesc='Tensors and Dynamic neural networks in Python with strong GPU acceleration'
pkgdesc="${_pkgdesc}"
arch=('x86_64')
url="https://pytorch.org"
license=('BSD')
source=("${_pkgname}::git+https://github.com/pytorch/pytorch.git#tag=v$_pkgver")
#source=("${_pkgname}::git+https://github.com/pytorch/pytorch.git")
b2sums=('SKIP')
options=('!lto' '!debug')
conflicts=(python-pytorch)
provides=(python-pytorch)

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

prepare() {
  cd "${srcdir}/${_pkgname}"
  git submodule sync
  git submodule update --init --recursive
}
build() {
  cd "${srcdir}/${_pkgname}"
  export VERBOSE=1
  #export PYTORCH_BUILD_VERSION="${pkgver}"
  #export PYTORCH_BUILD_NUMBER=1

  #export _GLIBCXX_USE_CXX11_ABI=1
  export HIP_ROOT_DIR=/opt/rocm
  export ROCM_HOME=/opt/rocm/
  #export USE_CUDA=OFF
  #export USE_ROCM=1
  export PYTORCH_ROCM_ARCH="gfx1030"
  export AMDGPU_TARGETS=${PYTORCH_ROCM_ARCH}
  export GPU_TARGETS=${PYTORCH_ROCM_ARCH}
  export MAGMA_HOME=/opt/rocm
  
  python tools/amd_build/build_amd.py
  #python setup.py build

  python setup.py build --cmake-only
  ccmake build  # or cmake-gui build
}

package() {
  cd "${srcdir}/${_pkgname}"
  python setup.py install --root="${pkgdir}"/
  install -Dm644 LICENSE "${pkgdir}/usr/share/licenses/${pkgname}/LICENSE"
}