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

EAPI=8

PYTHON_COMPAT=( python3_{12..14} )

inherit cuda cmake edo flag-o-matic python-r1

EIGEN_COMMIT="1d8b82b0740839c0de7f1242a3585e3390ff5f33"
ABSEIL_VERSION="20250814.1"
CUTLASS_VERSION="4.4.2"
CUDNN_FRONTEND_VERSION="1.24.0"
GTEST_VERSION="1.17.0"

DESCRIPTION="Cross-platform, high performance ML inferencing and training accelerator"
HOMEPAGE="
	https://onnxruntime.ai
	https://github.com/microsoft/onnxruntime
"
SRC_URI="
	https://github.com/microsoft/onnxruntime/archive/refs/tags/v${PV}.tar.gz -> ${P}.tar.gz
	https://gitlab.com/libeigen/eigen/-/archive/${EIGEN_COMMIT}/eigen-${EIGEN_COMMIT}.tar.bz2 ->
		eigen-3.4.0_p20250216.tar.bz2
	cuda? (
		https://github.com/abseil/abseil-cpp/archive/refs/tags/${ABSEIL_VERSION}.tar.gz ->
			abseil-cpp-${ABSEIL_VERSION}.tar.gz
		https://github.com/NVIDIA/cutlass/archive/refs/tags/v${CUTLASS_VERSION}.tar.gz ->
			cutlass-${CUTLASS_VERSION}.tar.gz
		https://github.com/NVIDIA/cudnn-frontend/archive/refs/tags/v${CUDNN_FRONTEND_VERSION}.tar.gz ->
			cudnn-frontend-${CUDNN_FRONTEND_VERSION}.tar.gz
	)
	test? (
		https://github.com/google/googletest/archive/refs/tags/v${GTEST_VERSION}.tar.gz ->
			googletest-${GTEST_VERSION}.tar.gz
	)
"

LICENSE="Apache-2.0 BSD MIT"
SLOT="0"
KEYWORDS="~amd64 ~arm64"
IUSE="cuda python test"
REQUIRED_USE="${PYTHON_REQUIRED_USE}"
RESTRICT="!test? ( test )"

# cmake/deps.txt for v1.28.0 pins onnx v1.22.0, and the
# use-system-libraries patch rewrites that FetchContent entry to
# FIND_PACKAGE_ARGS ... REQUIRED, so the system onnx is what gets used with
# no version guard of its own. That floor is not expressible: sci-ml/onnx
# exists only in ::gentoo and tops out at 1.20.1. Floored at 1.20.1 to keep
# 1.18.0-r1 out, which is further still from the pin. Whether 1.20.1 is
# actually sufficient has NOT been established -- raise this to
# >=sci-ml/onnx-1.22.0 once onnx is packaged that far.
# verified 2026-07-27
RDEPEND="
	!cuda? ( dev-cpp/abseil-cpp:= )
	dev-libs/cpuinfo
	dev-libs/protobuf:=
	dev-libs/re2:=
	>=sci-ml/onnx-1.20.1[disableStaticReg]
	cuda? (
		~dev-cpp/abseil-cpp-20250814.1:=
		dev-libs/cudnn:=
		dev-util/nvidia-cuda-toolkit:=
	)

	python? (
		${PYTHON_DEPS}
		dev-python/coloredlogs[${PYTHON_USEDEP}]
		dev-python/flatbuffers[${PYTHON_USEDEP}]
		>=dev-python/numpy-2[${PYTHON_USEDEP}]
		dev-python/packaging[${PYTHON_USEDEP}]
		dev-python/protobuf[${PYTHON_USEDEP}]
		dev-python/sympy[${PYTHON_USEDEP}]
	)
"
DEPEND="
	${RDEPEND}
	dev-cpp/ms-gsl
	dev-cpp/nlohmann_json
	dev-cpp/safeint
	dev-libs/boost
	dev-libs/date
	dev-libs/flatbuffers

	python? (
		dev-python/pybind11[${PYTHON_USEDEP}]
		sci-libs/dlpack
	)
"
BDEPEND="
	${PYTHON_DEPS}
	cuda? ( sys-devel/gcc:15 )

	test? (
		python? ( dev-python/pytest[${PYTHON_USEDEP}] )
	)
"

PATCHES=(
	"${FILESDIR}/${PN}-1.22.2-relax-the-dependency-on-flatbuffers.patch"
	"${FILESDIR}/${PN}-1.24.4-no-werror.patch"
	"${FILESDIR}/${PN}-1.28.0-use-system-libraries.patch"
)

CMAKE_USE_DIR="${S}/cmake"

src_prepare() {
	cmake_src_prepare

	if use cuda; then
		pushd "${WORKDIR}/abseil-cpp-${ABSEIL_VERSION}" >/dev/null || die
		eapply "${FILESDIR}/${PN}-1.28.0-abseil-nvcc.patch"
		popd >/dev/null || die
	fi
}

src_configure() {
	# Python is used at build time unconditionally
	python_setup

	local mycmakeargs=(
		-Donnxruntime_BUILD_SHARED_LIB=on

		-Donnxruntime_BUILD_UNIT_TESTS=$(usex test)
		-Donnxruntime_ENABLE_PYTHON=$(usex python)
		-Donnxruntime_USE_CUDA=$(usex cuda)

		# Use vendored Eigen at a specific 3.4-branch commit (2025-02-15).
		# ::gentoo's dev-cpp/eigen-3.4.0-r3 (Aug 2021) lacks 3+ years of
		# fixes onnxruntime depends on; eigen-3.4.9999 (live) would work
		# but a live ebuild as a build-dep is fragile. Eigen 5.0.1
		# (released 2026) is a major API break; onnxruntime's CMakeLists
		# doesn't yet support it. Drop the vendor when ::gentoo carries
		# a tagged 3.4.x release post-2025-02 or when upstream supports
		# Eigen 5.x. Verified 2026-05-16.
		-DFETCHCONTENT_SOURCE_DIR_EIGEN3="${WORKDIR}/eigen-${EIGEN_COMMIT}"

		# This makes it possible for `find_path` to find the `onnx-ml.proto` file
		-DCMAKE_INCLUDE_PATH="$(python_get_sitedir)"

		-Wno-dev
	)

	if use cuda; then
		# nvcc rejects gcc newer than the active CUDA toolkit supports
		# (CUDA 13 tops out at gcc 15). cuda_gccdir picks the newest
		# supported slot; pin ordinary C++, CUDA host compilation, and
		# final linking to it so all three share one libstdc++ ABI. The
		# cuda? sys-devel/gcc:15 BDEPEND guarantees a compatible slot is
		# installed for cuda_gccdir to find.
		local cuda_gcc_bindir
		cuda_gcc_bindir="$(cuda_gccdir)" || die
		local -x CC="${cuda_gcc_bindir}/gcc"
		local -x CXX="${cuda_gcc_bindir}/g++"
		local -x CUDAHOSTCXX="${CXX}"
		cuda_add_sandbox -w
		mycmakeargs+=(
			-DCMAKE_CUDA_ARCHITECTURES="${CUDAARCHS:-all-major}"
			-DCMAKE_CUDA_COMPILER="/opt/cuda/bin/nvcc"
			-DCMAKE_CUDA_FLAGS="-I${WORKDIR}/abseil-cpp-${ABSEIL_VERSION}"
			-DCMAKE_CUDA_HOST_COMPILER="${CUDAHOSTCXX}"
			-DFETCHCONTENT_SOURCE_DIR_CUDNN_FRONTEND="${WORKDIR}/cudnn-frontend-${CUDNN_FRONTEND_VERSION}"
			-DFETCHCONTENT_SOURCE_DIR_CUTLASS="${WORKDIR}/cutlass-${CUTLASS_VERSION}"
			-Donnxruntime_CUDA_HOME="/opt/cuda"
			-Donnxruntime_CUDNN_HOME="/opt/cuda"
		)
	fi
	use test && mycmakeargs+=(
		-DFETCHCONTENT_SOURCE_DIR_GOOGLETEST="${WORKDIR}/googletest-${GTEST_VERSION}"
	)

	append-ldflags -Wl,-z,noexecstack
	cmake_src_configure
}

# Adapted from `run_onnxruntime_tests` in `tools/ci_build/build.py`
python_test() {
	cd "${S}/cmake_build" || die
	epytest --pyargs \
		onnxruntime_test_python.py \
		onnxruntime_test_python_backend.py \
		onnxruntime_test_python_mlops.py \
		onnxruntime_test_python_sparse_matmul.py
}

src_test() {
	local -x GTEST_FILTER="*:-ActivationOpNoInfTest.Softsign:LayoutTransformationPotentiallyAddedOpsTests.OpsHaveLatestVersions"
	cmake_src_test

	if use python ; then
		python_test
	fi
}

# There is some custom logic in `setup.py`
python_install() {
	cd "${S}/cmake_build" || die
	edo "${EPYTHON}" ../setup.py install \
		--prefix="${EPREFIX}/usr" \
		--root="${D}"

	local libs=(
		"libonnxruntime.so.${PV}"
		"libonnxruntime_providers_shared.so"
	)
	use cuda && libs+=( "libonnxruntime_providers_cuda.so" )
	for lib in "${libs[@]}"; do
		ln -fsr "${ED}/usr/$(get_libdir)/${lib}" "${D}/$(python_get_sitedir)/onnxruntime/capi/${lib}" || die
	done

	rm -rf "${D}/$(python_get_sitedir)"/*.egg-info || die
	python_optimize
}

src_install() {
	cmake_src_install

	if use python ; then
		python_foreach_impl python_install
	fi

	dodoc "${S}/"{README.md,LICENSE}
}