#!/bin/bash

# Slackware build script for pytorch

# Copyright 2024  Johannes Schoepfer, Germany
# All rights reserved.
#
# Redistribution and use of this script, with or without modification, is
# permitted provided that the following conditions are met:
#
# 1. Redistributions of this script must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR "AS IS" AND ANY EXPRESS OR IMPLIED
#  WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#  MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO
#  EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
#  OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
#  WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
#  OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
#  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

cd $(dirname $0) ; CWD=$(pwd)

PRGNAM=pytorch
VERSION=${VERSION:-2.3.1}
VERSION_SIX=1.11.0
BUILD=${BUILD:-1}
TAG=${TAG:-_SBo}
PKGTYPE=${PKGTYPE:-tgz}

if [ -z "$ARCH" ]; then
  case "$( uname -m )" in
    i?86) ARCH=i586 ;;
    arm*) ARCH=arm ;;
       *) ARCH=$( uname -m ) ;;
  esac
fi

if [ ! -z "${PRINT_PACKAGE_NAME}" ]; then
  echo "$PRGNAM-$VERSION-$ARCH-$BUILD$TAG.$PKGTYPE"
  exit 0
fi

TMP=${TMP:-/tmp/SBo}
PKG=$TMP/package-$PRGNAM
OUTPUT=${OUTPUT:-/tmp}

if [ "$ARCH" = "i586" ]; then
  SLKCFLAGS="-O2 -march=i586 -mtune=i686"
elif [ "$ARCH" = "i686" ]; then
  SLKCFLAGS="-O2 -march=i686 -mtune=i686"
elif [ "$ARCH" = "aarch64" ]; then
  SLKCFLAGS="-O2 -fPIC"
  LIBDIRSUFFIX="64"
elif [ "$ARCH" = "x86_64" ]; then
  SLKCFLAGS="-O2 -fPIC"
  LIBDIRSUFFIX="64"
elif [ "$ARCH" = "arm" ]; then
  SLKCFLAGS="-O2 -march=armv7-a -mfpu=vfpv3-d16 -mfloat-abi=hard"
else
  SLKCFLAGS="-O2"
fi

set -e

rm -rf $PKG
mkdir -p $TMP $PKG $OUTPUT
cd $TMP
rm -rf $PRGNAM-v$VERSION
tar xvf $CWD/$PRGNAM-v$VERSION.tar.gz
cd $PRGNAM-v$VERSION
chown -R root:root .
chmod -R u+w,go-w,a+rX-st $PKG

# Needed for NNPACK
mkdir -p build/confu-deps/six-download/six-prefix/src/
cp $CWD/six-$VERSION_SIX.tar.gz build/confu-deps/six-download/six-prefix/src/

# Fix cmake libdir location
sed -i c10/CMakeLists.txt \
  -e "s;DESTINATION lib;DESTINATION lib$LIBDIRSUFFIX;g"

# Make Vulkan Warnings not to make the build fail
sed -i "/Werror=return-type/d" CMakeLists.txt

# Fix compilation with nvidia-cuda-toolkit-12.4.0
# https://github.com/pytorch/pytorch/issues/122169
patch -p0 < $CWD/boxing.patch

cd build
  unshare -n cmake \
    -G Ninja \
    -DCMAKE_C_FLAGS:STRING="$SLKCFLAGS" \
    -DCMAKE_CXX_FLAGS:STRING="$SLKCFLAGS" \
    -DCMAKE_CXX_STANDARD=17 \
    -DCMAKE_INSTALL_PREFIX="/usr" \
    -DLIBSHM_INSTALL_LIB_SUBDIR="lib$LIBDIRSUFFIX" \
    -DTORCH_INSTALL_LIB_DIR="lib$LIBDIRSUFFIX" \
    -DPYTHON_EXECUTABLE=$(which python3) \
    -DBUILD_CUSTOM_PROTOBUF=OFF \
    -DBUILD_TEST=OFF \
    -DUSE_FFMPEG=ON \
    -DUSE_GOLD_LINKER=ON \
    -DUSE_OPENCL=ON \
    -DUSE_OPENCV=ON \
    -DUSE_VULKAN=ON \
    -DCMAKE_BUILD_TYPE=Release ..
  "${NINJA:=ninja}"
  DESTDIR=tmpxxx $NINJA install/strip

  mkdir -p $PKG/usr/{share,lib$LIBDIRSUFFIX}
  mv tmpxxx/usr/bin $PKG/usr
  mv tmpxxx/usr/include $PKG/usr
  mv tmpxxx/usr/share/cmake $PKG/usr/share
  mv tmpxxx/usr/lib$LIBDIRSUFFIX/*.so $PKG/usr/lib$LIBDIRSUFFIX
cd ..
python3 setup.py install --root=$PKG

SITE="$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"
mkdir -p $PKG/$SITE/torch/{bin,lib}
BIN=torch_shm_manager; ln -s /usr/bin/$BIN $PKG/$SITE/torch/bin/$BIN
LIB=libtorch_global_deps.so; ln -s /usr/lib$LIBDIRSUFFIX/$LIB $PKG/$SITE/torch/lib/$LIB

find $PKG -print0 | xargs -0 file | grep -e "executable" -e "shared object" | grep ELF \
  | cut -f 1 -d : | xargs strip --strip-unneeded 2> /dev/null || true

mkdir -p $PKG/usr/doc/$PRGNAM-$VERSION
cp -a LICENSE NOTICE README.md RELEASE.md $PKG/usr/doc/$PRGNAM-$VERSION
cat $CWD/$PRGNAM.SlackBuild > $PKG/usr/doc/$PRGNAM-$VERSION/$PRGNAM.SlackBuild

mkdir -p $PKG/install
cat $CWD/slack-desc > $PKG/install/slack-desc

cd $PKG
/sbin/makepkg -l y -c n $OUTPUT/$PRGNAM-$VERSION-$ARCH-$BUILD$TAG.$PKGTYPE