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

EAPI=8

PYTHON_COMPAT=( python3_{12..14} )

inherit cmake git-r3 python-any-r1 toolchain-funcs

DESCRIPTION="Ahead of Time (AOT) Triton Math Library"
HOMEPAGE="https://github.com/ROCm/aotriton"
EGIT_REPO_URI="https://github.com/ROCm/aotriton.git"
EGIT_COMMIT="${PV}"
EGIT_SUBMODULES=( '*' )
# FlyDSL kernel sources (third_party/flydsl-kernel.txt); upstream clones them
# at configure time.
FLYDSL_KERNEL_PV="0.3.1"
SRC_URI="
	https://github.com/ROCm/FlyDSL/archive/refs/tags/v${FLYDSL_KERNEL_PV}.tar.gz
		-> flydsl-${FLYDSL_KERNEL_PV}.gh.tar.gz
"

LICENSE="MIT"
SLOT="0/${PV%b}"
KEYWORDS="-* ~amd64"

IUSE_TARGETS=(
	gfx90a
	gfx942
	gfx950
	gfx1100
	gfx1101
	gfx1102
	gfx1103
	gfx1150
	gfx1151
	gfx1200
	gfx1201
	gfx1250
)
IUSE_TARGETS=( "${IUSE_TARGETS[@]/#/amdgpu_targets_}" )
IUSE="${IUSE_TARGETS[*]/#/+}"

REQUIRED_USE="|| ( ${IUSE_TARGETS[*]} )"

# Triton's wheel build downloads its backend and LLVM; AOTriton fetches aiter.
PROPERTIES="live"
RESTRICT="network-sandbox test"

# Source builds have no binary shim constraint; upstream only requires HIP.
# Match aotriton-bin's 6.4 floor and next-major cap. verified 2026-09-26
RDEPEND="
	!!sci-libs/aotriton-bin
	sys-libs/glibc
	sys-devel/gcc
	app-arch/xz-utils
	>=dev-util/hip-6.4:=
	<dev-util/hip-11:=
"
DEPEND="${RDEPEND}"
BDEPEND="
	>=dev-build/cmake-3.27
	>=dev-build/ninja-1.11
	${PYTHON_DEPS}
	virtual/zlib
	virtual/pkgconfig
	$(python_gen_any_dep '
		dev-python/filelock[${PYTHON_USEDEP}]
		dev-python/iniconfig[${PYTHON_USEDEP}]
		dev-python/numpy[${PYTHON_USEDEP}]
		dev-python/packaging[${PYTHON_USEDEP}]
		dev-python/pandas[${PYTHON_USEDEP}]
		dev-python/pip[${PYTHON_USEDEP}]
		dev-python/pluggy[${PYTHON_USEDEP}]
		dev-python/pybind11[${PYTHON_USEDEP}]
		dev-python/pyyaml[${PYTHON_USEDEP}]
		>=dev-python/setuptools-64[${PYTHON_USEDEP}]
		dev-python/wheel[${PYTHON_USEDEP}]
	')
	$(python_gen_any_dep '
		dev-python/flydsl[${PYTHON_SINGLE_USEDEP}]
	')
"

PATCHES=(
	"${FILESDIR}/${PN}-0.13b-include-iostream.patch"
	"${FILESDIR}/${P}-system-flydsl.patch"
)

python_check_deps() {
	python_has_version "dev-python/filelock[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/iniconfig[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/numpy[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/packaging[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/pandas[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/pip[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/pluggy[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/pybind11[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/pyyaml[${PYTHON_USEDEP}]" &&
		python_has_version ">=dev-python/setuptools-64[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/wheel[${PYTHON_USEDEP}]" &&
		python_has_version "dev-python/flydsl[${PYTHON_SINGLE_USEDEP}]"
}

src_unpack() {
	git-r3_src_unpack
	default
}

src_prepare() {
	cmake_src_prepare

	# Expose declared build dependencies and the system flydsl in the
	# disposable venv; install the local Triton wheel into it below.
	[[ $(grep -c 'execute_process(COMMAND.*-m venv "${VENV_DIR}")' CMakeLists.txt) == 1 &&
		$(grep -c -- '-m pip install "${CMAKE_CURRENT_LIST_DIR}"' CMakeLists.txt) == 1 ]] ||
		die "venv setup lines moved; re-audit the sed below"
	sed -i \
		-e '/execute_process(COMMAND.*-m venv "${VENV_DIR}")/s/-m venv /-m venv --system-site-packages /' \
		-e '/-m pip install "${CMAKE_CURRENT_LIST_DIR}"/s/pip install/pip install --no-build-isolation --no-deps/' \
		CMakeLists.txt || die
}

src_configure() {
	python_setup
	tc-export CC CXX
	local CMAKE_BUILD_TYPE=Release

	local wheel_dir="${T}/triton-wheel"
	mkdir -p "${wheel_dir}" "${T}/home" || die

	# Cap LLVM/MLIR wheel compilation to avoid exhausting RAM.
	local -x MAX_JOBS=4
	local -x HOME="${T}/home"
	local -x PIP_CACHE_DIR="${T}/pip-cache"
	local -x TRITON_BUILD_PROTON=OFF
	local -x TRITON_HOME="${T}/triton-cache"
	"${EPYTHON}" -m pip wheel --no-deps \
		--wheel-dir "${wheel_dir}" "${S}/third_party/triton" ||
		die "failed to build the vendored Triton wheel"

	local triton_wheels=( "${wheel_dir}"/triton-*.whl )
	[[ ${#triton_wheels[@]} -eq 1 && -f ${triton_wheels[0]} ]] ||
		die "expected exactly one Triton wheel"

	local targets=() target
	for target in "${IUSE_TARGETS[@]}"; do
		use "${target}" && targets+=( "${target#amdgpu_targets_}" )
	done
	local target_arch
	printf -v target_arch '%s;' "${targets[@]}"
	target_arch=${target_arch%;}

	# Leave ROCM_PATH unset here: hip-config-amd.cmake derives its clang root
	# and device libraries from it, and the flydsl toolkit tree set below
	# for the kernel build would break hip::host. verified 2026-09-26
	local -x PIP_NO_INDEX=1
	local mycmakeargs=(
		-DCMAKE_INSTALL_PREFIX="${BUILD_DIR}/install_dir"
		-DAOTRITON_GPU_BUILD_TIMEOUT=0
		-DAOTRITON_NO_PYTHON=ON
		-DAOTRITON_TARGET_ARCH="${target_arch}"
		-DAOTRITON_USE_LOCAL_TRITON_WHEEL="${triton_wheels[0]}"
		-DAOTRITON_USE_TORCH=OFF
		-DAOTRITON_USE_SYSTEM_FLYDSL=ON
		-DAOTRITON_FLYDSL_KERNEL_ROOT="${WORKDIR}/FlyDSL-${FLYDSL_KERNEL_PV}"
		-DPython3_EXECUTABLE="${PYTHON}"
	)
	cmake_src_configure
}

src_compile() {
	# The install target owns the generated-kernel build graph.
	:
}

src_install() {
	local -x PIP_NO_INDEX=1
	# The kernels are compiled here. flyc_compile needs a ROCm root holding
	# llvm/bin/ld.lld and llvm/bin/llvm-readelf: dev-python/flydsl's toolkit.
	local -x ROCM_PATH="${EPREFIX}/usr/lib/flydsl/rocm"
	cmake_build install

	local install_dir="${BUILD_DIR}/install_dir"
	doheader -r "${install_dir}"/include/*
	insinto /usr/$(get_libdir)
	doins -r "${install_dir}"/lib/*
}