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

EAPI=8

DISTUTILS_USE_PEP517=setuptools
DISTUTILS_EXT=1
PYTHON_COMPAT=( python3_{12..15} )

inherit distutils-r1 flag-o-matic

LLVM_REV="b010a18d"
# The AMD backend links its own LLVM, pinned in cmake/amd-llvm-info.json.
AMD_CODEGEN_REV="6bc4aaf6"

DESCRIPTION="Language and compiler for custom deep-learning primitives"
HOMEPAGE="https://github.com/triton-lang/triton"
SRC_URI="https://github.com/triton-lang/triton/archive/refs/tags/v${PV}.tar.gz
	-> ${P}.gh.tar.gz
	https://oaitriton.blob.core.windows.net/public/llvm-builds/llvm-${LLVM_REV}-ubuntu-x64-1.tar.gz
	-> ${P}.llvm.tar.gz
	https://oaitriton.blob.core.windows.net/public/llvm-builds/amd-codegen-${AMD_CODEGEN_REV}-ubuntu-x64-1.tar.gz
	-> ${P}.amd-codegen.tar.gz"

# Prebuilt LLVM is installed statically in libtriton.so and as FileCheck; the
# AMD backend also installs UoI-NCSA HSA headers. All license terms are
# cumulative. # verified 2026-09-09 against the staged image
LICENSE="MIT Apache-2.0-with-LLVM-exceptions UoI-NCSA"
SLOT="0"
KEYWORDS="~amd64"
IUSE="cuda"

# The upstream suite requires supported NVIDIA or AMD accelerator hardware.
RESTRICT="test"

# Offline mode omits NVIDIA tools and headers, and Triton has no PATH fallback.
# Without the installed toolkit symlinks, the first CUDA JIT fails to find
# ptxas, and before that fails to compile driver.c for want of cuda.h.
# verified 2026-09-09 against the staged image and knobs.py:193-217
# verified 2026-10-05 by running JIT kernels from the staged image on CUDA 13.4
# USE=cuda adds the toolkit and links; only the NVIDIA backend uses them.
RDEPEND="
	!!dev-python/triton-bin
	virtual/zlib:=
	cuda? ( dev-util/nvidia-cuda-toolkit )
"
# Since 3.8.0 the build always resolves nlohmann/json and, offline, insists on
# JSON_SYSPATH; the system header-only copy replaces the v3.11.3 download.
# verified 2026-11-05 against 3.9.0's setup.py and build_helpers.py
DEPEND="dev-cpp/nlohmann_json"
# 3.9.0 moved its bindings from pybind11 to nanobind (find_package(nanobind
# CONFIG REQUIRED)); upstream pins nanobind==2.10.2, which ::gentoo lacks, and
# the haarmek build links nanobind 3.1.0. verified 2026-11-05
BDEPEND="
	dev-build/cmake
	dev-build/ninja
	dev-python/nanobind[${PYTHON_USEDEP}]
"

# Use upstream's pinned LLVM and forbid setup.py downloads.
export TRITON_OFFLINE_BUILD=1
export TRITON_BUILD_PROTON=OFF

src_prepare() {
	# The combined AMD/NVIDIA library violates GCC's ODR checks under LTO.
	filter-lto
	distutils-r1_src_prepare

	# Ship upstream's prebuilt AMD code-generation library where the backend
	# loads it; TRITON_AMD_CODEGEN_PATH below stops the build helper fetching it.
	mkdir -p third_party/amd/backend/lib || die
	cp "${WORKDIR}/llvm/amd-codegen-${AMD_CODEGEN_REV}-ubuntu-x64-1/lib/libtriton_amd_codegen.so" \
		third_party/amd/backend/lib/ || die
}

src_unpack() {
	unpack "${P}.gh.tar.gz"
	mkdir "${WORKDIR}/llvm" || die
	cd "${WORKDIR}/llvm" || die
	unpack "${P}.llvm.tar.gz" "${P}.amd-codegen.tar.gz"
}

python_compile() {
	local -x LLVM_SYSPATH="${WORKDIR}/llvm/llvm-${LLVM_REV}-ubuntu-x64-1"
	local -x JSON_SYSPATH="${ESYSROOT}/usr"
	local -x TRITON_AMD_CODEGEN_PATH="${S}/third_party/amd/backend/lib/libtriton_amd_codegen.so"
	# LLVM translation units use multiple GiB each; cap default concurrency.
	local -x MAX_JOBS="${MAX_JOBS:-4}"
	distutils-r1_python_compile
}

python_install() {
	distutils-r1_python_install
	rm "${D}$(python_get_sitedir)/triton/instrumentation/libGPUInstrumentationTestLib.so" || die

	if use cuda; then
		# Link toolkit tools into Triton's lookup path. ptxas-blackwell has
		# no toolkit counterpart, so sm_100+ needs TRITON_PTXAS_BLACKWELL_PATH.
		# Strip EPREFIX because dodir/dosym add it themselves.
		local tool nvdir="$(python_get_sitedir)"
		nvdir="${nvdir#"${EPREFIX}"}/triton/backends/nvidia"
		dodir "${nvdir}/bin"
		for tool in ptxas cuobjdump nvdisasm; do
			dosym -r "/opt/cuda/bin/${tool}" "${nvdir}/bin/${tool}"
		done

		# driver.py compiles driver.c on first use with this directory as its
		# only CUDA include path; the wheel fills it with a copy of the toolkit
		# headers.
		dosym -r /opt/cuda/include "${nvdir}/include"
	fi
}