# Copyright 2022-2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2

EAPI=8

DISTUTILS_USE_PEP517=scikit-build-core
PYTHON_COMPAT=( python3_{12..15} )
DISTUTILS_SINGLE_IMPL=1
DISTUTILS_EXT=1
ROCM_VERSION=6.1
MYPN=pytorch
MYP=${MYPN}-${PV}
# pytorch-2.13.0 depends on specific commit of composable kernel
# TODO: replace it with DEPEND in the future
CK_COMMIT=f1746955fdaf80a3414de814bf32437686dac347
CK_P=composable_kernel-${CK_COMMIT:0:8}

inherit distutils-r1 cmake cuda cuda-extra flag-o-matic prefix rocm toolchain-funcs

DESCRIPTION="Tensors and Dynamic neural networks in Python with strong GPU acceleration"
HOMEPAGE="https://pytorch.org/"
SRC_URI="https://github.com/pytorch/pytorch/releases/download/v${PV}/pytorch-v${PV}.tar.gz
	-> ${MYP}-full.tar.gz
	rocm? (
		https://github.com/ROCm/composable_kernel/archive/${CK_COMMIT}.tar.gz
		-> ${CK_P}.tar.gz
	)
"
S="${WORKDIR}"/pytorch-v${PV}
LICENSE="BSD"
SLOT="0"
KEYWORDS="~amd64 ~arm64"
IUSE="cuda cusparselt cudss distributed fbgemm flash gloo memefficient mimalloc mkl
	mpi nccl nnpack +numpy onednn openblas openmp qnnpack rocm vulkan xnnpack
	cpu_flags_x86_avx cpu_flags_x86_avx2 cpu_flags_x86_avx512f"
RESTRICT="test"
REQUIRED_USE="
	${PYTHON_REQUIRED_USE}
	fbgemm? ( cpu_flags_x86_avx cpu_flags_x86_avx2 cpu_flags_x86_avx512f )
	mpi? ( distributed )
	gloo? ( distributed )
	?? ( cuda rocm )
	rocm? (
		|| ( ${ROCM_REQUIRED_USE} )
		memefficient? ( flash )
	)
	cusparselt? ( || ( cuda rocm ) )
	flash? ( || ( cuda rocm ) )
	memefficient? ( || ( cuda rocm ) )
	nccl? ( || ( cuda rocm ) )
"

RDEPEND="
	${PYTHON_DEPS}
	!sci-ml/caffe2
	dev-cpp/abseil-cpp:=
	dev-cpp/gflags:=
	>=dev-cpp/glog-0.6.0:=
	dev-cpp/opentelemetry-cpp
	>=dev-libs/cpuinfo-2025.11.14
	dev-libs/libfmt:=
	dev-libs/protobuf:=
	dev-libs/sleef
	sci-ml/foxi
	sci-ml/onnx:=
	$(python_gen_cond_dep '
		dev-python/networkx[${PYTHON_USEDEP}]
		dev-python/sympy[${PYTHON_USEDEP}]
		dev-python/typing-extensions[${PYTHON_USEDEP}]
	')
	virtual/lapack
	!mkl? ( !openblas? ( virtual/blas ) )
	cuda? (
		dev-libs/cudnn:=
		>=sci-ml/cudnn-frontend-1.12.0:0=
		cudss? ( <dev-libs/cudss-0.8 )
		>=dev-util/nvidia-cuda-toolkit-12.9:=[profiler]
		cusparselt? ( dev-libs/cusparselt )
	)
	distributed? (
		!rocm? ( sci-ml/tensorpipe[cuda?] )
		dev-cpp/cpp-httplib:=
	)
	fbgemm? ( >=sci-ml/FBGEMM-1.4 )
	gloo? ( >=sci-ml/gloo-2025.06.04[cuda?,rocm?] )
	mimalloc? ( dev-libs/mimalloc )
	mkl? ( sci-libs/mkl )
	mpi? ( virtual/mpi )
	nnpack? (
		sci-ml/NNPACK
		dev-libs/pthreadpool
	)
	numpy? ( $(python_gen_cond_dep '
		dev-python/numpy[${PYTHON_USEDEP}]
	') )
	onednn? ( sci-ml/oneDNN )
	openblas? ( sci-libs/openblas )
	qnnpack? (
		sci-ml/gemmlowp
		dev-libs/pthreadpool
	)
	rocm? (
		nccl? ( >=dev-libs/rccl-6.3:= <dev-libs/rccl-7.3:= )
		>=dev-util/hip-6.3:=       <dev-util/hip-7.3:=
		>=dev-util/roctracer-6.3:= <dev-util/roctracer-7.3:=
		>=sci-libs/hipBLAS-6.3:=   <sci-libs/hipBLAS-7.3:=[rocsolver(+)]
		>=sci-libs/hipBLASLt-6.3:= <sci-libs/hipBLASLt-7.3:=
		>=sci-libs/hipFFT-6.3:=    <sci-libs/hipFFT-7.3:=
		>=sci-libs/hipRAND-6.3:=   <sci-libs/hipRAND-7.3:=
		>=sci-libs/hipSOLVER-6.3:= <sci-libs/hipSOLVER-7.3:=
		>=sci-libs/hipSPARSE-6.3:= <sci-libs/hipSPARSE-7.3:=
		>=sci-libs/miopen-6.3:=    <sci-libs/miopen-7.3:=
		>=sci-libs/rocBLAS-6.3:=   <sci-libs/rocBLAS-7.3:=
		>=sci-libs/rocRAND-6.3:=   <sci-libs/rocRAND-7.3:=
		>=sci-libs/rocSOLVER-6.3:= <sci-libs/rocSOLVER-7.3:=
		memefficient? ( =sci-libs/aotriton-bin-0.13*:= )
		distributed? (
			>=dev-util/rocm-smi-6.3:= <dev-util/rocm-smi-7.3:=
			>=dev-util/amdsmi-6.3:= <dev-util/amdsmi-7.3:=
		)
		cusparselt? ( >=sci-libs/hipsparselt-6.3:= <sci-libs/hipsparselt-7.3:= )
	)
	xnnpack? (
		>=sci-ml/XNNPACK-2024.11
		dev-libs/pthreadpool
	)
	vulkan? ( media-libs/vulkan-loader )
"

DEPEND="
	${RDEPEND}
	dev-cpp/nlohmann_json
	dev-libs/flatbuffers
	dev-libs/FXdiv
	dev-libs/pocketfft
	dev-libs/psimd
	sci-ml/FP16
	$(python_gen_cond_dep '
		dev-python/pybind11[${PYTHON_USEDEP}]
		dev-python/pyyaml[${PYTHON_USEDEP}]
		dev-python/typing-extensions[${PYTHON_USEDEP}]
	')
	cuda? ( >=dev-libs/cutlass-4.6.1[tools(+)] )
	onednn? ( sci-ml/ideep )
	rocm? (
		>=sci-libs/hipCUB-6.3:=    <sci-libs/hipCUB-7.3:=
		>=sci-libs/rocPRIM-6.3:=   <sci-libs/rocPRIM-7.3:=
		>=sci-libs/rocThrust-6.3:= <sci-libs/rocThrust-7.3:=
	)
	qnnpack? ( dev-libs/clog )
"

BDEPEND="dev-build/cmake"

PATCHES=(
	"${FILESDIR}"/caffe2-2.13.0-gentoo.patch
	"${FILESDIR}"/caffe2-2.6.0-install-dirs.patch
	"${FILESDIR}"/caffe2-2.13.0-glog-0.6.0.patch
	"${FILESDIR}"/caffe2-2.9.0-tensorpipe.patch
	"${FILESDIR}"/caffe2-2.3.0-cudnn_include_fix.patch
	"${FILESDIR}"/caffe2-2.1.2-fix-rpath.patch
	"${FILESDIR}"/caffe2-2.4.0-fix-openmp-link.patch
	"${FILESDIR}"/caffe2-2.11.0-rocm-fix-std-cpp17.patch
	"${FILESDIR}"/caffe2-2.4.0-cpp-httplib.patch
	"${FILESDIR}"/caffe2-2.7.0-glog-0.7.1.patch
	"${FILESDIR}"/caffe2-2.13.0-aotriton-fixes.patch
	"${FILESDIR}"/caffe2-2.13.0-kineto.patch
	"${FILESDIR}"/caffe2-2.8.0-rocm-minus-flash.patch
	"${FILESDIR}"/caffe2-2.9.0-CUDA-13.patch
	"${FILESDIR}"/caffe2-2.10.0-nvrtc.patch
	"${FILESDIR}"/caffe2-2.11.0-mimalloc.patch
	"${FILESDIR}"/caffe2-2.12.0-cudss.patch
	"${FILESDIR}"/caffe2-2.12.0-rocm-distributed-link.patch
	"${FILESDIR}"/caffe2-2.12.0-rocm-assert-fix.patch
	"${FILESDIR}"/caffe2-2.13.0-cmake-install-fix.patch

	"${FILESDIR}"/pytorch-2.10.0-global-dlopen.patch
	"${FILESDIR}"/pytorch-2.5.0-torch_shm_manager.patch
	"${FILESDIR}"/pytorch-2.10.0-cpp-extension-multilib.patch
	"${FILESDIR}"/pytorch-2.14.0-cuda.patch
	"${FILESDIR}"/pytorch-2.13.0-nvcc.patch
)

src_prepare() {
	filter-lto #bug 862672

	sed -i \
		-e "/third_party\/gloo/d" \
		cmake/Dependencies.cmake \
		|| die

	# Change libaotriton path
	sed -i \
		-e "s|}/lib|}/\${CMAKE_INSTALL_LIBDIR}|g" \
		-e "/set(__AOTRITON_LIB/s|lib/|\${CMAKE_INSTALL_LIBDIR}/|g" \
		cmake/External/aotriton.cmake \
		|| die

	# Noisy warnings from Logging.h
	sed -i 's/-Wextra-semi//' cmake/public/utils.cmake || die

	cmake_src_prepare
	pushd torch/csrc/jit/serialization > /dev/null || die
	flatc --cpp --gen-mutable --scoped-enums mobile_bytecode.fbs || die
	popd > /dev/null || die

	# prefixify the hardcoded paths, after all patches are applied
	hprefixify \
		aten/CMakeLists.txt \
		caffe2/CMakeLists.txt \
		cmake/Metal.cmake \
		cmake/Modules/*.cmake \
		cmake/Modules_CUDA_fix/FindCUDNN.cmake \
		cmake/Modules_CUDA_fix/upstream/FindCUDA/make2cmake.cmake \
		cmake/Modules_CUDA_fix/upstream/FindPackageHandleStandardArgs.cmake \
		cmake/public/LoadHIP.cmake \
		cmake/public/cuda.cmake \
		cmake/Dependencies.cmake \
		tools/setup_helpers/env.py \
		torch/CMakeLists.txt \
		CMakeLists.txt

	if use rocm; then
		sed -e "s:/opt/rocm:/usr:" \
			-e "s:lib/cmake:$(get_libdir)/cmake:g" \
			-i cmake/public/LoadHIP.cmake || die

		# TODO: delete, when pytorch depends on systemwide composable_kernel
		sed -e "s:third_party/composable_kernel:../composable_kernel-${CK_COMMIT}:g" \
			-i aten/src/ATen/CMakeLists.txt || die

		# Workaround for libc++ issue https://github.com/llvm/llvm-project/issues/100802
		sed -e 's/std::memcpy/memcpy/g' \
			-i torch/headeronly/util/Half.h \
			-i aten/src/ATen/native/cuda/int4mm.cu || die

		ebegin "HIPifying cuda sources"
		FBCODE_BUILD_TOOL="buck" ${EPYTHON} tools/amd_build/build_amd.py || die
		eend $?
	fi

	if use cuda; then
		cuda_src_prepare_extra
	fi

	rm -rf third_party/flatbuffers

	PATCHES=() # Already applied.
	distutils-r1_src_prepare

	# Replace placeholders added by cpp-extension.patch
	sed -i \
		-e "s/%LIB_DIR%/$(get_libdir)/g" \
		torch/utils/cpp_extension.py \
		torch/__init__.py \
		|| die

	# Drop legacy from pyproject.toml
	sed -i \
		-e "/build-backend/s|:__legacy__||" \
		pyproject.toml \
		|| die
}

src_configure() {
	if use cuda && [[ -z ${TORCH_CUDA_ARCH_LIST} ]]; then
		ewarn "WARNING: pytorch is being built with its default CUDA compute capabilities: 5.0 and 7.0."
		ewarn "These may not be optimal for your GPU."
		ewarn ""
		ewarn "To configure pytorch with the CUDA compute capability that is optimal for your GPU,"
		ewarn "set TORCH_CUDA_ARCH_LIST in your make.conf, and re-emerge pytorch."
		ewarn "For example, to use CUDA capability 7.5, add: TORCH_CUDA_ARCH_LIST=7.5"
		ewarn "For a Maxwell model GPU, an example value would be: TORCH_CUDA_ARCH_LIST=Maxwell"
		ewarn ""
		ewarn "You can look up your GPU's CUDA compute capability at https://developer.nvidia.com/cuda-gpus"
		ewarn "or by running /opt/cuda/extras/demo_suite/deviceQuery | grep 'CUDA Capability'"
	fi

	local mycmakeargs=(
		-DBUILD_CUSTOM_PROTOBUF=OFF
		-DBUILD_SHARED_LIBS=ON
		-DLIBSHM_INSTALL_LIB_SUBDIR="${EPREFIX}"/usr/$(get_libdir)
		-DPython_EXECUTABLE="${PYTHON}"
		-DTORCH_INSTALL_LIB_DIR="${EPREFIX}"/usr/$(get_libdir)
		-DATEN_NO_TEST=ON
		-DBUILD_TEST=OFF
		-DUSE_CCACHE=OFF
		-DUSE_CUDA=$(usex cuda)
		-DUSE_DISTRIBUTED=$(usex distributed)
		-DUSE_FBGEMM=$(usex fbgemm)
		-DUSE_FLASH_ATTENTION=$(usex flash)
		-DUSE_GFLAGS=ON
		-DUSE_GLOG=ON
		-DUSE_GLOO=$(usex gloo)
		-DUSE_ITT=OFF
		-DUSE_KINETO=ON
		-DUSE_KLEIDIAI=OFF # TODO
		-DUSE_MAGMA=OFF # deprecated
		-DUSE_MEM_EFF_ATTENTION=$(usex memefficient)
		-DUSE_MIMALLOC=$(usex mimalloc)
		-DUSE_MKLDNN=$(usex onednn)
		-DUSE_MPI=$(usex mpi)
		-DUSE_NCCL=OFF
		-DUSE_NNPACK=$(usex nnpack)
		-DUSE_NUMA=OFF
		-DUSE_NUMPY=$(usex numpy)
		-DUSE_NVRTC=$(usex cuda)
		-DUSE_OPENMP=$(usex openmp)
		-DUSE_PYTORCH_QNNPACK=$(usex qnnpack)
		-DUSE_PYTORCH_METAL=OFF
		-DUSE_ROCM=$(usex rocm)
		-DUSE_SYSTEM_CPUINFO=ON
		-DUSE_SYSTEM_EIGEN_INSTALL=ON
		-DUSE_SYSTEM_FP16=ON
		-DUSE_SYSTEM_FXDIV=ON
		-DUSE_SYSTEM_GLOO=ON
		-DUSE_SYSTEM_NVTX=ON
		-DUSE_SYSTEM_ONNX=ON
		-DUSE_SYSTEM_PSIMD=ON
		-DUSE_SYSTEM_PTHREADPOOL=ON
		-DUSE_SYSTEM_PYBIND11=ON
		-DUSE_SYSTEM_SLEEF=ON
		-DUSE_TENSORPIPE=$(usex distributed $(usex !rocm))
		-DUSE_UCC=OFF
		-DUSE_VALGRIND=OFF
		-DUSE_VULKAN=$(usex vulkan)
		-DUSE_XNNPACK=$(usex xnnpack)
		-DUSE_SYSTEM_XNNPACK=$(usex xnnpack)
		-DUSE_XPU=OFF
		-DC_AVX_FOUND=$(usex cpu_flags_x86_avx)
		-DC_AVX2_FOUND=$(usex cpu_flags_x86_avx2)
		-DC_AVX512_FOUND=$(usex cpu_flags_x86_avx512f)
		-DCXX_AVX_FOUND=$(usex cpu_flags_x86_avx)
		-DCXX_AVX2_FOUND=$(usex cpu_flags_x86_avx2)
		-DCXX_AVX512_FOUND=$(usex cpu_flags_x86_avx512f)

		-Wno-dev
	)

	if use mkl; then
		mycmakeargs+=(-DBLAS=MKL)
	elif use openblas; then
		mycmakeargs+=(-DBLAS=OpenBLAS)
	else
		mycmakeargs+=(-DBLAS=Generic -DBLAS_LIBRARIES=)
	fi

	if use cuda; then
		# bug 867706 926116
		cuda_add_sandbox
		addpredict "/dev/char/"

		mycmakeargs+=(
			-DUSE_CUDNN=ON
			-DTORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-5.0 7.0}"
			-DUSE_NCCL=OFF # TODO: NVIDIA Collective Communication Library
			-DCUDA_TOOLKIT_ROOT_DIR="${EPREFIX}/opt/cuda"
			-DCMAKE_CUDA_COMPILER="nvcc"
			-DCMAKE_CUDA_HOST_COMPILER="$(cuda_gcc)"
			-DCMAKE_CUDA_FLAGS="-forward-unknown-opts -fno-lto ${NVCCFLAGS//\"/\'}"
			-DUSE_CUDSS=$(usex cudss)
			-DUSE_CUSPARSELT=$(usex cusparselt)
		)

		[[ -v CUDACXX ]] && export PYTORCH_NVCC="${CUDACXX}"
		export CMAKE_CUDA_HOST_COMPILER="$(cuda_gcc)" # for python_compile()

		if use flash; then
			export FLASH_ATTENTION_FORCE_BUILD="TRUE"
			export FLASH_ATTN_CUDA_ARCHS="${CUDAARCHS:-${TORCH_CUDA_ARCH_LIST:-5.0 7.0}}"
		fi
	elif use rocm; then
		export PYTORCH_ROCM_ARCH="$(get_amdgpu_flags)"

		if use memefficient; then
			export AOTRITON_INSTALLED_PREFIX="${ESYSROOT}/usr"
		fi

		mycmakeargs+=(
			-DUSE_NCCL=$(usex nccl)
			-DUSE_SYSTEM_NCCL=ON
			-DCMAKE_REQUIRE_FIND_PACKAGE_HIP=ON
			-DCMAKE_DISABLE_FIND_PACKAGE_hipsparselt=$(usex !cusparselt) # disable automagic
			-DUSE_ROCM_CK_SDPA=OFF # requires flash + aiter, works only on gfx90a/gfx942/gfx950
		)

		# ROCm libraries produce too much warnings
		append-cxxflags -Wno-deprecated-declarations -Wno-unused-result -Wno-unused-value
	fi

	if use onednn; then
		mycmakeargs+=(
			-DMKLDNN_FOUND=ON
			-DMKLDNN_LIBRARIES=dnnl
			-DMKLDNN_INCLUDE_DIR="${ESYSROOT}/usr/include/oneapi/dnnl"
		)
	fi

	cmake_src_configure
}

python_compile() {
	export BUILD_DIR="${WORKDIR}/pytorch-v${PV}_build"
	# Don't run CMake again.
	export SKBUILD_WHEEL_CMAKE="false"

	distutils-r1_python_compile
}

src_compile() {
	PYTORCH_BUILD_VERSION=${PV} \
	PYTORCH_BUILD_NUMBER=0 \
	cmake_src_compile

	distutils-r1_src_compile
}

python_install() {
	export BUILD_DIR="${WORKDIR}/pytorch-v${PV}_build"
	distutils-r1_python_install

	mkdir "${D}"$(python_get_sitedir)/torch/bin || die
	mkdir "${D}"$(python_get_sitedir)/torch/include || die
	ln -s ../../../../../include/torch \
		"${D}$(python_get_sitedir)"/torch/include/torch || die # bug 923269
	ln -s ../../../../../bin/torch_shm_manager \
		"${D}"$(python_get_sitedir)/torch/bin/torch_shm_manager || die
	ln -s ../../../../../$(get_libdir)/libtorch_global_deps.so \
		"${D}"$(python_get_sitedir)/torch/lib/libtorch_global_deps.so || die
}

src_install() {
	cmake_src_install
	distutils-r1_src_install

	# Clean up third-party software installs.
	rm -rf "${ED}"/usr/share/cmake/{kineto,fbgemm} \
		"${ED}"/usr/include/{kineto,fbgemm,asmjit} \
		"${ED}"/usr/$(get_libdir)/cmake/asmjit \
		"${ED}"/usr/$(get_libdir)/pkgconfig/mimalloc.pc

	cp -a "${ED}"/usr/torch "${ED}"$(python_get_sitedir)/ || die
	rm -rf "${ED}"/usr/torch

	python_optimize
}