File PKGBUILD of Package python-torchvision-rocm

# Maintainer: wuxxin <wuxxin@gmail.com>
# Contributor: Lubosz Sarnecki @lubosz
# Based on python-torchvision; original contributors:
# Contributor: Christian Heusel
# Contributor: Konstantin Gizdov
# Contributor: Torsten Keßler <tpkessler@archlinux.org>
# Contributor: Konstantin Gizdov <arch at kge dot pw>
# Contributor: Butui Hu <hot123tea123@gmail.com>
# Contributor: Chih-Hsuan Yen <yan12125@archlinux.org>
# Contributor: Jean Lucas <jean@4ray.co>
# Based on python-torchvision-git; original contributors:
# Contributor: Stephen Zhang <zsrkmyn at gmail dot com>

_pkgname='vision'
pkgbase='python-torchvision-rocm'
pkgname=('torchvision-rocm' 'python-torchvision-rocm')
pkgver=0.24.1
pkgrel=1
pkgdesc='Datasets, transforms, and models specific to computer vision (with ROCM support)'
arch=('x86_64')
url='https://github.com/pytorch/vision'
license=('BSD-3-Clause')
depends=(
    numactl
    python-numpy
    python-pillow
    python-requests
    python-scipy
    python-sympy
    python-pytorch-rocm
    ffmpeg
    libjpeg-turbo
    libpng
)
makedepends=(
    cmake
    ninja
    python-setuptools
    rocm-hip-sdk
    rocm-toolchain
    qt6-base
    miopen-hip
)
source=(
    "${_pkgname}-${pkgver}.tar.gz::https://github.com/pytorch/vision/archive/v${pkgver}.tar.gz"
    "pytorch-vision-9231.patch"
    "torchvision-0_17_1-fix-build.patch"
    "ffmpeg-8.patch"
)
sha256sums=('071da2078600bfec4886efab77358c9329abfedcf1488b05879b556cb9b84ba7'
            'ac2b78100923115e7aa090c0e6f936d1580c0cdb7a19ecbfcce5c5c6fa3e5c14'
            'ed715ca202d2b010c50414e370ebc0492f0f42b298a8e6e03f9fe80b7ce60331'
            '1c466d4ae0874a8ab6518cb3b0c7747f43060fb6803352ac9e13f2b5ecae1931')
prepare() {
    cd "${srcdir}/${_pkgname}-${pkgver}"

    # https://github.com/pytorch/vision/issues/8307
    patch -N -i "${srcdir}"/torchvision-0_17_1-fix-build.patch
    # FFmpeg 7.1+ support
    patch -p1 -N -i "${srcdir}"/pytorch-vision-9231.patch
    # FFmpeg 8
    patch -p1 -N -i "${srcdir}"/ffmpeg-8.patch
}

build() {
    cd "${srcdir}/${_pkgname}-${pkgver}"

    if test -e build; then rm -rf build; fi
    mkdir build

    # populate build architecture list similar to pkg arch:python-pytorch
    # 2.8.0-5: gfx950 lacks support for 128 bit atomics
    _PYTORCH_ROCM_ARCH="$(rocm-supported-gfx -e gfx950)"
    if test -n "$GPU_TARGETS"; then _PYTORCH_ROCM_ARCH="$GPU_TARGETS"; fi
    if test -n "$AMDGPU_TARGETS"; then _PYTORCH_ROCM_ARCH="$AMDGPU_TARGETS"; fi
    if test -n "$PYTORCH_ROCM_ARCH"; then _PYTORCH_ROCM_ARCH="$PYTORCH_ROCM_ARCH"; fi
    export PYTORCH_ROCM_ARCH="${_PYTORCH_ROCM_ARCH}"
    echo "building for PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"

    # if ROCM_HOME is not set, hardcode ROCM_HOME, ROCM_PATH, HIP_ROOT_DIR to /opt/rocm
    export ROCM_HOME="${ROCM_HOME:-/opt/rocm}"
    export ROCM_PATH="$ROCM_HOME"
    export HIP_ROOT_DIR="$ROCM_HOME"

    # XXX we need to exclude this in favor of the hipified source, setup.py, make_C_extension
    # add after: shutil.copy(str(header), str(CSRS_DIR / "ops/hip"))
    # cuda_sources.remove(CSRS_DIR / "vision.cpp")

    cmake \
        -G Ninja \
        -Wno-dev \
        -B build \
        -DCMAKE_INSTALL_PREFIX=/usr \
        -DCMAKE_BUILD_TYPE=None \
        -DCMAKE_CXX_FLAGS="${CXXFLAGS} -O3" \
        -DCMAKE_SHARED_LINKER_FLAGS="${LDFLAGS}" \
        -DWITH_PNG=ON \
        -DWITH_JPEG=ON \
        -DCMAKE_POLICY_VERSION_MINIMUM=3.5

    cmake --build build

    TORCHVISION_INCLUDE=${srcdir} \
        TORCHVISION_LIBRARY=/usr/lib \
        TORCHVISION_USE_NVJPEG=0 \
        TORCHVISION_USE_VIDEO_CODEC=0 \
        TORCHVISION_USE_FFMPEG=0 \
        python setup.py build
}

package_torchvision-rocm() {
    pkgdesc='Datasets, transforms, and models specific to computer vision (C++ library only, with ROCM support)'
    provides+=(torchvision torchvision=${pkgver})
    conflicts+=(torchvision)

    cd "${srcdir}/${_pkgname}-${pkgver}"
    DESTDIR="${pkgdir}" cmake --install build
    install -m644 -Dt "$pkgdir/usr/share/licenses/$pkgname" LICENSE
}

package_python-torchvision-rocm() {
    pkgdesc='Datasets, transforms, and models specific to computer vision (with ROCM support)'
    provides+=(python-torchvision python-torchvision=${pkgver})
    conflicts+=(python-torchvision)

    cd "${srcdir}/${_pkgname}-${pkgver}"
    TORCHVISION_INCLUDE=${srcdir} \
        TORCHVISION_LIBRARY=/usr/lib \
        FORCE_CUDA=1 \
        python setup.py install --root="${pkgdir}" --optimize=1 --skip-build
    install -Dm644 LICENSE -t "${pkgdir}/usr/share/licenses/${pkgname}"
}