# Copyright 2022-2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2
# ROCm implementation of caffe2/PyTorch backend

EAPI=8

PYTHON_COMPAT=( python3_{11..14} )
ROCM_VERSION=6.1
LLVM_COMPAT=( {18..22} )
inherit python-single-r1 cmake flag-o-matic llvm-r1 llvm-utils prefix rocm

# Source package is pytorch, not caffe2-rocm
MY_PN=caffe2
MYPN=pytorch
MYP=${MYPN}-${PV}

CK_COMMIT=f1746955fdaf80a3414de814bf32437686dac347
CK_P=composable_kernel-${CK_COMMIT:0:8}

DESCRIPTION="A deep learning framework (ROCm backend)"
HOMEPAGE="https://pytorch.org/"
SRC_URI="
	https://github.com/pytorch/${MYPN}/archive/refs/tags/v${PV}.tar.gz -> ${MYP}.tar.gz
	https://github.com/ROCm/composable_kernel/archive/${CK_COMMIT}.tar.gz -> ${CK_P}.tar.gz
"

S="${WORKDIR}"/${MYP}

LICENSE="BSD"
SLOT="0"
KEYWORDS="~amd64"
IUSE="cusparselt distributed fbgemm flash gloo kineto memefficient
	mimalloc mkl mpi nccl nnpack +numpy onednn openblas opencl openmp qnnpack
	+xnnpack"
RESTRICT="test"
REQUIRED_USE="
	${PYTHON_REQUIRED_USE}
	mpi? ( distributed )
	gloo? ( distributed )
	${ROCM_REQUIRED_USE}
	memefficient? ( flash )
"

RDEPEND="
	${PYTHON_DEPS}
	app-eselect/eselect-caffe2
	dev-cpp/abseil-cpp:=
	dev-cpp/gflags:=
	>=dev-cpp/glog-0.5.0:=
	>=dev-libs/cpuinfo-2025.11.14
	dev-libs/libfmt:=
	dev-libs/protobuf:=
	dev-libs/sleef
	sci-ml/onnx
	virtual/lapack
	nccl? ( >=dev-libs/rccl-6.3:= <dev-libs/rccl-7.3:= )
	>=dev-util/hip-6.3:=[video_cards_amdgpu] <dev-util/hip-7.3:=[video_cards_amdgpu]
	>=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.12b:= )
	distributed? (
		>=dev-util/rocm-smi-6.3:= <dev-util/rocm-smi-7.3:=
		>=dev-util/amdsmi-6.3:= <dev-util/amdsmi-7.3:=
		dev-cpp/cpp-httplib:=
	)
	cusparselt? ( >=sci-libs/hipsparselt-6.3:= <sci-libs/hipsparselt-7.3:= )
	fbgemm? ( >=sci-ml/FBGEMM-1.4 )
	gloo? ( >=sci-ml/gloo-2025.06.04[rocm] )
	kineto? ( ~sci-ml/kineto-0.4.0_p20260323 )
	mimalloc? ( dev-libs/mimalloc )
	mpi? ( virtual/mpi )
	nnpack? ( sci-ml/NNPACK dev-libs/pthreadpool )
	numpy? ( $(python_gen_cond_dep 'dev-python/numpy[${PYTHON_USEDEP}]') )
	onednn? ( sci-ml/oneDNN )
	opencl? ( virtual/opencl )
	qnnpack? ( !sci-libs/QNNPACK sci-ml/gemmlowp dev-libs/pthreadpool )
	xnnpack? ( >=sci-ml/XNNPACK-2024.11 dev-libs/pthreadpool )
	mkl? ( sci-libs/mkl )
	openblas? ( sci-libs/openblas )
"

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-3.0.5[${PYTHON_USEDEP}]
		dev-python/pyyaml[${PYTHON_USEDEP}]
		dev-python/typing-extensions[${PYTHON_USEDEP}]
	')
	onednn? ( sci-ml/ideep )
	>=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 )
"

# Patches use upstream names (caffe2, not caffe2-rocm)
PATCHES=(
	# caffe2-2.5.1-unbundle_fmt.patch dropped for 2.13.0 — handled via sed in src_prepare
	"${FILESDIR}"/${MY_PN}-2.5.1-unbundle_kineto.patch
	"${FILESDIR}"/${MY_PN}-2.8.0-unbundle_pocketfft.patch
	"${FILESDIR}"/${MY_PN}-2.5.1-cudnn_include_fix.patch
	"${FILESDIR}"/${MY_PN}-2.4.0-cpp-httplib.patch
	"${FILESDIR}"/${MY_PN}-2.5.1-glog-0.6.0.patch
	"${FILESDIR}"/${MY_PN}-2.7.0-glog-0.7.1.patch
	"${FILESDIR}"/${MY_PN}-2.12.0-aotriton-fixes.patch
	"${FILESDIR}"/${MY_PN}-2.8.0-rocm-minus-flash.patch
	"${FILESDIR}"/${MY_PN}-2.12.0-rocm-distributed-link.patch
	"${FILESDIR}"/${MY_PN}-2.9.1-torch_cpu.patch
	"${FILESDIR}"/${MY_PN}-2.10.0-gentoo.patch
	"${FILESDIR}"/${MY_PN}-2.11.0-mimalloc.patch
	"${FILESDIR}"/${MY_PN}-2.12.0-removekineto-pr178960.patch
)

# Install prefix for this backend
CAFFE2_PREFIX="/usr/lib/caffe2/rocm"
PYTORCH_PREFIX="/usr/lib/pytorch/rocm"

src_prepare() {
	filter-lto

	sed -i \
		-e 's|::fmt-header-only||' \
		c10/CMakeLists.txt \
		cmake/Dependencies.cmake \
		torch/CMakeLists.txt \
		|| die

	sed -e '/target_compile_options_if_supported(tensorpipe/d' -i cmake/Dependencies.cmake || die

	# pytorch 2.13.0: remove fmt target_compile_definitions (system fmt has no cmake target)
	sed -i \
		-e '/target_compile_definitions(fmt /d' \
		-e '/target_compile_definitions(fmt-header-only /d' \
		-e '/set_target_properties(fmt-header-only/d' \
		-e '/_fmt_no_unique_address/d' \
		cmake/Dependencies.cmake || die

	# pytorch 2.13.0: CuTeDSL is prototype, don't fail if cutlass submodule missing
	sed -i 's/FATAL_ERROR/WARNING/' cmake/FileMirroring.cmake || die

	# pytorch 2.13.0: fix glog 0.6.0 compat — internal namespace doesn't exist in 0.6.0
	sed -i \
		-e '/namespace glog_internal_namespace_/d' \
		-e 's/glog_internal_namespace_:://g' \
		c10/util/Exception.cpp || die

	# pytorch 2.13.0: wrap_headers.py ignores DESTDIR, writes to live system → sandbox violation
	# The wrapping adds TORCH_STABLE_ONLY guards we don't need. Remove the install(CODE) call.
	sed -i '/wrap_headers\.py/,/COMMAND_ERROR_IS_FATAL/d' cmake/PostBuildSteps.cmake || die

	# Fix ambiguous pow() call under clang++/hipcc — cast args to resolve overload
	sed -i 's/compat_pow(1 - standard_uniform.sample(), 1.0f \/ alpha)/compat_pow(static_cast<accscalar_t>(1 - standard_uniform.sample()), static_cast<accscalar_t>(1.0f \/ alpha))/' \
		aten/src/ATen/native/Distributions.h || die

	sed -i \
		-e '/add_subdirectory.*third_party/d' \
		CMakeLists.txt \
		cmake/Dependencies.cmake \
		cmake/ProtoBuf.cmake \
		aten/src/ATen/CMakeLists.txt \
		|| die

	sed -i \
		-e "/EXPORT/s|DESTINATION lib)|DESTINATION $(get_libdir))|" \
		c10/CMakeLists.txt \
		c10/hip/CMakeLists.txt \
		|| die

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

	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

	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 \
		torch/CMakeLists.txt \
		CMakeLists.txt

	# ROCm-specific patches + fix slotted LLVM path for HIP compiler
	sed -e "s:/opt/rocm:/usr:" \
		-e "s:lib/cmake:$(get_libdir)/cmake:g" \
		-i cmake/public/LoadHIP.cmake || die
	sed -i 's|"${ROCM_PATH}/lib/llvm/bin"|"/usr/lib/llvm/'"${LLVM_SLOT}"'/bin"|' \
		cmake/public/LoadHIP.cmake || die

	sed -e "s:third_party/composable_kernel:../composable_kernel-${CK_COMMIT}:g" \
		-i aten/src/ATen/CMakeLists.txt || die

	# CK gfx101 fixes are upstream in f174695+ — no patch needed

	sed -e 's/std::memcpy/memcpy/g' -i torch/headeronly/util/Half.h || die

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

	# Fix gemm<float> undefined symbol: clang++ mismatches explicit specialization
	# with primary template due to static_assert(false&&sizeof) being non-dependent.
	# Replace with sizeof(Dtype)==0 which is always false but type-dependent.
	sed -i 's/static_assert(false&&sizeof(Dtype)/static_assert(sizeof(Dtype) == 0/' \
		aten/src/ATen/hip/HIPBlas.h || die

	# hip/hip_runtime.h defines __noinline__ as empty macro, breaking GCC 16's
	# <format> header where [[__gnu__::__noinline__]] becomes [[__gnu__::]]
	# hip/amd_detail/host_defines.h:260 defines __noinline__ as EMPTY for non-HCC
	# compilers, which breaks GCC 16's <format> header where
	# [[__gnu__::__noinline__]] becomes [[__gnu__::]]. This affects ALL headers
	# that include hip_runtime.h (directly or via hipblas, hipsparse, etc.)
	# Fix: overlay header that wraps the real one and undefs the broken macro
	mkdir -p "${S}/hip_fix/hip/amd_detail" || die
	cat > "${S}/hip_fix/hip/amd_detail/host_defines.h" <<-'HIPFIX' || die
	#include_next <hip/amd_detail/host_defines.h>
	#ifdef __noinline__
	#undef __noinline__
	#endif
	HIPFIX
	# Store hip_fix path for src_configure to pick up
	HIP_FIX_DIR="${S}/hip_fix"

	# Force CPU-only kernel DEFAULT variants to compile with host gcc instead of hipcc.
	# hipcc fails on ALSO_REGISTER_AVX512_DISPATCH in generated build-dir context.
	# These are pure CPU kernels with no GPU code.
	sed -i '/set_source_files_properties(${NEW_IMPL} PROPERTIES COMPILE_FLAGS/s|COMPILE_FLAGS|COMPILE_OPTIONS "SHELL:-x c++" COMPILE_FLAGS|' \
		cmake/Codegen.cmake 2>/dev/null || true
}

src_configure() {
	# pytorch 2.13.0 uses enable_language(HIP) — needs clang++ in PATH
	llvm_prepend_path ${LLVM_SLOT}
	# Use slotted clang++ as CXX compiler — avoids both GCC 16 <format> crash
	# AND hipcc int8mm dispatch stub issue. Combined with __noinline__ wrapper
	# header and gemm static_assert fix for correct template instantiation.
	export CC="/usr/lib/llvm/${LLVM_SLOT}/bin/clang"
	export CXX="/usr/lib/llvm/${LLVM_SLOT}/bin/clang++"

	export PYTORCH_ROCM_ARCH="$(get_amdgpu_flags)"

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

	local mycmakeargs=(
		-DCMAKE_INSTALL_PREFIX="${EPREFIX}${CAFFE2_PREFIX}"
		-DCMAKE_INSTALL_RPATH="${EPREFIX}${CAFFE2_PREFIX}/$(get_libdir)"
		-DCMAKE_BUILD_WITH_INSTALL_RPATH=ON
		-DUSE_SYSTEM_LIBS=ON
		-DBUILD_CUSTOM_PROTOBUF=OFF
		-DBUILD_TEST=OFF
		-DLIBSHM_INSTALL_LIB_SUBDIR="${EPREFIX}${CAFFE2_PREFIX}/$(get_libdir)"
		-DPython_EXECUTABLE="${PYTHON}"
		-DTORCH_INSTALL_LIB_DIR="${EPREFIX}${CAFFE2_PREFIX}/$(get_libdir)"
		-DUSE_CCACHE=OFF
		-DUSE_CUDA=OFF
		-DUSE_ROCM=ON
		-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=$(usex kineto)
		-DUSE_KLEIDIAI=OFF
		-DUSE_MAGMA=OFF
		-DUSE_MEM_EFF_ATTENTION=$(usex memefficient)
		-DUSE_MIMALLOC=$(usex mimalloc)
		-DUSE_MKLDNN=$(usex onednn)
		-DUSE_MPI=$(usex mpi)
		-DUSE_NCCL=$(usex nccl)
		-DUSE_SYSTEM_NCCL=ON
		-DUSE_NNPACK=$(usex nnpack)
		-DUSE_NUMA=OFF
		-DUSE_NUMPY=$(usex numpy)
		-DUSE_OPENCL=$(usex opencl)
		-DUSE_OPENMP=$(usex openmp)
		-DUSE_PYTORCH_QNNPACK=$(usex qnnpack)
		-DUSE_PYTORCH_METAL=OFF
		-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_SYSTEM_XNNPACK=$(usex xnnpack)
		$(usex xnnpack '' '-DCMAKE_DISABLE_FIND_PACKAGE_XNNPACK=ON')
		-DUSE_TENSORPIPE=OFF
		-DUSE_UCC=OFF
		-DUSE_VALGRIND=OFF
		-DUSE_XNNPACK=$(usex xnnpack)
		-DUSE_XPU=OFF
		-DCMAKE_REQUIRE_FIND_PACKAGE_HIP=ON
		-DCMAKE_DISABLE_FIND_PACKAGE_hipsparselt=$(usex !cusparselt)
		-DUSE_ROCM_CK_SDPA=OFF
		-Wno-dev
		# hip_fix wrapper header to undef __noinline__ (breaks GCC 16 <format>)
		-DCMAKE_CXX_FLAGS="-isystem ${HIP_FIX_DIR}"
	)

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

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

	append-cxxflags -Wno-deprecated-declarations -Wno-unused-result -Wno-unused-value -Wno-template-body

	cmake_src_configure
}

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

src_install() {
	cmake_src_install

	# CMake cache for pytorch-rocm
	insinto "/var/lib/caffe2-rocm"
	doins "${BUILD_DIR}"/CMakeCache.txt

	# Python torch module to pytorch prefix
	rm -rf python
	mkdir -p python/torch || die
	cp torch/version.py python/torch/ || die

	local pyver="${EPYTHON#python}"
	local torch_dest="${PYTORCH_PREFIX}"
	insinto "${torch_dest}"
	doins -r python/torch

	# Create required subdirs in torch module
	dodir "${torch_dest}/torch/bin"
	dodir "${torch_dest}/torch/lib"
	dodir "${torch_dest}/torch/include"
	dosym "${CAFFE2_PREFIX}/include/torch" "${torch_dest}/torch/include/torch"
	dosym "${CAFFE2_PREFIX}/bin/torch_shm_manager" "${torch_dest}/torch/bin/torch_shm_manager"
	dosym "${CAFFE2_PREFIX}/$(get_libdir)/libtorch_global_deps.so" "${torch_dest}/torch/lib/libtorch_global_deps.so"
}

pkg_postinst() {
	local active
	active=$(eselect caffe2 show 2>/dev/null)
	if [[ "${active}" == "(unset)" || -z "${active}" ]]; then
		eselect caffe2 set rocm
		elog "caffe2 backend set to: rocm"
	fi
}