gentoo-overlay/sci-ml/caffe2/caffe2-2.11.0-r3.ebuild
2026-04-20 19:30:35 +03:00

550 lines
22 KiB
Bash

# Copyright 2022-2026 Gentoo Authors
# Distributed under the terms of the GNU General Public License v2
EAPI=8
PYTHON_COMPAT=( python3_{11..14} )
ROCM_VERSION=6.1
inherit python-single-r1 cmake cuda flag-o-matic prefix rocm toolchain-funcs
MYPN=pytorch
MYP=${MYPN}-${PV}
# caffe2-2.9.0 depends on future version of composable kernel
# TODO: replace it with DEPEND in the future
CK_COMMIT=7fe50dc3da2069d6645d9deb8c017a876472a977
CK_P=composable_kernel-${CK_COMMIT:0:8}
FLASH_PV=2.7.4
FLASH_PN=flash-attention
FLASH_P=${FLASH_PN}-${FLASH_PV}
FLASH_ATT_URI="https://github.com/Dao-AILab/${FLASH_PN}/archive/refs/tags/v${FLASH_PV}.tar.gz -> ${FLASH_P}.gh.tar.gz"
AOTRITON_PV=0.9.2b
AOTRITON_PN=aotriton
AOTRITON_P=${AOTRITON_PN}-${AOTRITON_PV}
AOTRITON_tar=${AOTRITON_P}-manylinux_2_28_x86_64-rocm6.3-shared.tar.gz
DESCRIPTION="A deep learning framework"
HOMEPAGE="https://pytorch.org/"
SRC_URI="
https://github.com/pytorch/${MYPN}/archive/refs/tags/v${PV}.tar.gz -> ${MYP}.tar.gz
rocm? (
https://github.com/ROCm/composable_kernel/archive/${CK_COMMIT}.tar.gz
-> ${CK_P}.tar.gz
)
cuda? (
flash? ( ${FLASH_ATT_URI} )
memefficient? ( ${FLASH_ATT_URI} )
)
"
S="${WORKDIR}"/${MYP}
LICENSE="BSD"
SLOT="0"
KEYWORDS="~amd64 ~arm64"
IUSE="cuda cusparselt distributed fbgemm flash gloo intel_targets_gfx-bmg kineto
memefficient mimalloc mkl mpi nccl nnpack +numpy onednn openblas opencl
openmp qnnpack rocm xnnpack xpu"
RESTRICT="test"
REQUIRED_USE="
${PYTHON_REQUIRED_USE}
mpi? ( distributed )
gloo? ( distributed )
?? ( cuda rocm xpu )
rocm? (
|| ( ${ROCM_REQUIRED_USE} )
memefficient? ( flash )
)
xpu? ( amd64 )
cusparselt? ( || ( cuda rocm ) )
flash? ( || ( cuda rocm ) )
memefficient? ( || ( cuda rocm ) )
nccl? ( rocm )
"
RDEPEND="
${PYTHON_DEPS}
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
cuda? (
dev-libs/cudnn
>=sci-ml/cudnn-frontend-1.12.0:=
>=dev-util/nvidia-cuda-toolkit-12.9:=[profiler]
cusparselt? ( dev-libs/cusparselt )
)
fbgemm? ( >=sci-ml/FBGEMM-1.4 )
gloo? ( >=sci-ml/gloo-2025.06.04[cuda?,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
)
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.11*:= )
distributed? ( >=dev-util/rocm-smi-6.3:= <dev-util/rocm-smi-7.3:= )
cusparselt? ( >=sci-libs/hipsparselt-6.3:= <sci-libs/hipsparselt-7.3:= )
)
distributed? (
!rocm? ( sci-ml/tensorpipe[cuda?] )
dev-cpp/cpp-httplib:=
)
xnnpack? (
>=sci-ml/XNNPACK-2024.11
dev-libs/pthreadpool
)
mkl? ( sci-libs/mkl )
openblas? ( sci-libs/openblas )
xpu? (
>=dev-libs/intel-compute-runtime-26.14.37833.4[l0]
>=dev-util/intel-llvm-6.3.0
sci-ml/oneDNN[sycl]
~sci-ml/torch-xpu-ops-2.12.0_p20260414
virtual/opencl
)
"
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.2[${PYTHON_USEDEP}]
dev-python/pyyaml[${PYTHON_USEDEP}]
dev-python/typing-extensions[${PYTHON_USEDEP}]
')
cuda? ( >=dev-libs/cutlass-3.9.2[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 )
"
PATCHES=(
"${FILESDIR}"/${PN}-2.5.1-unbundle_fmt.patch
"${FILESDIR}"/${PN}-2.5.1-unbundle_kineto.patch
"${FILESDIR}"/${PN}-2.8.0-unbundle_pocketfft.patch
"${FILESDIR}"/${PN}-2.5.1-cudnn_include_fix.patch
"${FILESDIR}"/${PN}-2.4.0-cpp-httplib.patch
"${FILESDIR}"/${PN}-2.5.1-glog-0.6.0.patch
"${FILESDIR}"/${P}-rocm-fix-std-cpp17.patch
"${FILESDIR}"/${PN}-2.7.0-glog-0.7.1.patch
"${FILESDIR}"/${PN}-2.7.1-aotriton-fixes.patch
"${FILESDIR}"/${PN}-2.8.0-rocm-minus-flash.patch
"${FILESDIR}"/${PN}-2.9.0-cmake.patch
"${FILESDIR}"/${PN}-2.9.0-rocm-distributed-link.patch
"${FILESDIR}"/${PN}-2.9.1-torch_cpu.patch
"${FILESDIR}"/${PN}-2.10.0-gentoo.patch
"${FILESDIR}"/${P}-mimalloc.patch
"${FILESDIR}"/${P}-removekineto-pr178960.patch
)
src_prepare() {
if use cuda && ( use flash || use memefficient ); then
mv "${WORKDIR}"/${FLASH_P}/* third_party/${FLASH_PN}/ || die
fi
if use xpu; then
rm -rf third_party/torch-xpu-ops || die
mkdir -p third_party/torch-xpu-ops || die
cp -a "${ESYSROOT}/usr/share/torch-xpu-ops/current/." third_party/torch-xpu-ops/ || die
sed -i \
-e 's/SYCL_CMPLR_TEST_EXTRACT(${test_output}/SYCL_CMPLR_TEST_EXTRACT("${test_output}"/g' \
-e 's/CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "MSVC"/CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang" OR CMAKE_CXX_COMPILER_ID STREQUAL "MSVC"/' \
-e 's/CMAKE_CXX_COMPILER_ID STREQUAL "GNU"/CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang"/g' \
-e '/gline-tables-only/d' \
third_party/torch-xpu-ops/cmake/Modules/FindSYCLToolkit.cmake \
third_party/torch-xpu-ops/cmake/BuildFlags.cmake \
third_party/torch-xpu-ops/CMakeLists.txt || die
sed -i \
-e '44i\ target_include_directories(${sycl_lib} PRIVATE ${TORCH_XPU_OPS_INCLUDE_DIRS})\n target_include_directories(${sycl_lib} PRIVATE ${TORCH_XPU_OPS_ROOT}/src)\n target_include_directories(${sycl_lib} PRIVATE ${ATen_XPU_INCLUDE_DIRS})\n target_include_directories(${sycl_lib} PRIVATE ${SYCL_INCLUDE_DIR})' \
-e '/target_include_directories(${lib} PUBLIC ${TORCH_XPU_OPS_INCLUDE_DIRS})/a\ target_include_directories(${lib} PUBLIC ${TORCH_XPU_OPS_ROOT}/src)' \
-e '/target_link_libraries(${lib} PUBLIC ${SYCL_LIBRARY})/a\ target_link_libraries(${lib} PUBLIC c10_xpu)\n target_link_libraries(${lib} PUBLIC torch_cpu)\n target_link_libraries(${lib} PUBLIC c10)' \
third_party/torch-xpu-ops/src/BuildOnLinux.cmake || die
sed -i \
-e '/SYCL_OFFLINE_COMPILER_CG_OPTIONS .*cl-poison-unsupported-fp64-kernels/d' \
-e '/SYCL_OFFLINE_COMPILER_CG_OPTIONS .*cl-intel-enable-auto-large-GRF-mode/d' \
-e '/SYCL_OFFLINE_COMPILER_CG_OPTIONS .*cl-fp32-correctly-rounded-divide-sqrt/d' \
-e '/SYCL_OFFLINE_COMPILER_CG_OPTIONS .*cl-intel-greater-than-4GB-buffer-required/d' \
third_party/torch-xpu-ops/cmake/BuildFlags.cmake || die
perl -0pi -e 's@return WelfordDataLN\{\n sycl::select_from_group\(sg, wd.mean, 0\),\n sycl::select_from_group\(sg, wd.sigma2, 0\) \/ float\(N\),\n 0\.f\};@if (sg.get_local_linear_id() == 0) {\n buf[0] = wd.mean;\n buf[1] = wd.sigma2 / float(N);\n }\n sycl::group_barrier(item_id.get_group());\n return WelfordDataLN{\n static_cast<float>(buf[0]), static_cast<float>(buf[1]), 0.f};@' \
third_party/torch-xpu-ops/src/ATen/native/xpu/sycl/LayerNormKernels.cpp || die
perl -0pi -e 's@if \(sub_group_num == 1\) \{\n sum1 = sycl::group_broadcast\(sg, sum1, 0\);\n sum2 = sycl::group_broadcast\(sg, sum2, 0\);\n return;\n \}@if (sub_group_num == 1) {\n if (sg.get_local_linear_id() == 0) {\n local_data1[0] = sum1;\n local_data2[0] = sum2;\n }\n sycl::group_barrier(item.get_group());\n sum1 = local_data1[0];\n sum2 = local_data2[0];\n return;\n }@' \
third_party/torch-xpu-ops/src/ATen/native/xpu/sycl/Norm.h || die
perl -0pi -e 's@if \(sub_group_num == 1\) \{\n val = sycl::group_broadcast\(sg, val, 0\);\n return;\n \}@if (sub_group_num == 1) {\n if (sg.get_local_linear_id() == 0) {\n local_data[lid_row][0] = val;\n }\n sycl::group_barrier(item.get_group());\n val = local_data[lid_row][0];\n return;\n }@' \
third_party/torch-xpu-ops/src/ATen/native/xpu/sycl/SoftMaxKernels.cpp || die
perl -0pi -e 's@auto\* device_prop = at::xpu::getCurrentDeviceProperties\(\);\n auto device_architecture = device_prop->architecture;@auto& device = c10::xpu::get_raw_device(at::xpu::current_device());\n auto device_architecture = device.get_info<sycl::ext::oneapi::experimental::info::device::architecture>();@' \
aten/src/ATen/native/transformers/xpu/sdp_utils.cpp || die
perl -0pi -e 's@list\(APPEND ATen_XPU_DEPENDENCY_LIBS xpu_mkldnn\)@list(APPEND ATen_XPU_DEPENDENCY_LIBS ${XPU_MKLDNN_LIBRARIES})@' \
aten/src/ATen/CMakeLists.txt || die
# XPU increases generated TraceType size enough that some shards become
# pathological compile units on Clang; split them further.
perl -0pi -e 's@num_shards=5,@num_shards=10,@' \
tools/autograd/gen_trace_type.py || die
sed -i \
-e 's@set(SYCL_EXECUTABLE \${SYCL_COMPILER} CACHE FILEPATH "SYCL compiler")@set(SYCL_EXECUTABLE "icpx" CACHE FILEPATH "SYCL compiler")@' \
-e '67,68c\set(SYCL_HOST_COMPILER "clang++"\n CACHE FILEPATH "Host side compiler used by SYCL")' \
-e '/set(_sycl_c_or_cxx_flags ${CMAKE_${SYCL_C_OR_CXX}_FLAGS})/c\
set(_sycl_c_or_cxx_flags ${CMAKE_${SYCL_C_OR_CXX}_FLAGS})\
string(APPEND _sycl_c_or_cxx_flags " -stdlib=libstdc++")\
string(REPLACE "-fma" "-mfma" _sycl_c_or_cxx_flags "${_sycl_c_or_cxx_flags}")\
string(REPLACE "-no-ftz" "" _sycl_c_or_cxx_flags "${_sycl_c_or_cxx_flags}")' \
-e '405,409c\ COMMAND ${CMAKE_SYCL_COMPILER_LAUNCHER} ${SYCL_EXECUTABLE}\n -fsycl ${SYCL_device_link_flags}\n -fsycl-link ${object_files}\n -Xs ${SYCL_OFFLINE_COMPILER_FLAGS}\n -o ${output_file}' \
third_party/torch-xpu-ops/cmake/Modules/FindSYCL.cmake || die
perl -0pi -e 's@list\(APPEND SYCL_include_dirs "\$<TARGET_PROPERTY:\$\{sycl_target\},INCLUDE_DIRECTORIES>"\)@list(APPEND SYCL_include_dirs "${TORCH_XPU_OPS_ROOT}/src")\n list(APPEND SYCL_include_dirs "\$<TARGET_PROPERTY:\${sycl_target},INCLUDE_DIRECTORIES>")@' \
third_party/torch-xpu-ops/cmake/Modules/FindSYCL.cmake || die
sed -i \
-e '37c\set(SYCL_executable "icpx") # path' \
-e '38i\set(SYCL_launcher "@CMAKE_SYCL_COMPILER_LAUNCHER@") # path' \
-e '39i\list(REMOVE_ITEM SYCL_compile_flags -fma -no-ftz -gline-tables-only)\nlist(INSERT SYCL_compile_flags 0 -fsycl -mfma -stdlib=libstdc++)' \
-e '43i\string(REGEX REPLACE "/ATen/.*" "" TORCH_XPU_OPS_SRC_ROOT "${source_file}")\nlist(APPEND SYCL_include_dirs "${TORCH_XPU_OPS_SRC_ROOT}")' \
-e '66i\list(REMOVE_ITEM CMAKE_HOST_FLAGS -fma -no-ftz)\nlist(APPEND CMAKE_HOST_FLAGS -mfma -stdlib=libstdc++)' \
third_party/torch-xpu-ops/cmake/Modules/FindSYCL/run_sycl.cmake || die
perl -0pi -e 's@SYCL_execute_process\(\n "Generating \$\{generated_file\}"\n COMMAND "\$\{SYCL_executable\}"\n \$\{SYCL_dependency_file_args\}\n -c\n "\$\{source_file\}"\n -o "\$\{generated_file\}"\n \$\{SYCL_include_args\}\n \$\{SYCL_host_compiler\}\n \$\{SYCL_host_compiler_flags\}\n \$\{SYCL_compile_flags\}\n \)@SYCL_execute_process(\n "Generating \$\{generated_file\}"\n COMMAND \$\{SYCL_launcher\} "\$\{SYCL_executable\}"\n \$\{SYCL_dependency_file_args\}\n -c\n "\$\{source_file\}"\n -o "\$\{generated_file\}"\n \$\{SYCL_include_args\}\n \$\{SYCL_host_compiler\}\n \$\{SYCL_host_compiler_flags\}\n \$\{SYCL_compile_flags\}\n )@s' \
third_party/torch-xpu-ops/cmake/Modules/FindSYCL/run_sycl.cmake || die
sed -i \
-e '10i\if(SYCL_FOUND)\n set(SYCLTOOLKIT_FOUND TRUE)\nendif()' \
third_party/torch-xpu-ops/cmake/SYCL.cmake || die
sed -i \
-e '93,96c\if(TORCH_XPU_OPS_PYTORCH_DEPS AND COMMAND sycl_include_external_dependencies)\n sycl_include_external_dependencies(${TORCH_XPU_OPS_PYTORCH_DEPS})\nendif()' \
third_party/torch-xpu-ops/CMakeLists.txt || die
fi
filter-lto #bug 862672
# Unbundle fmt
sed -i \
-e 's|::fmt-header-only||' \
c10/CMakeLists.txt \
cmake/Dependencies.cmake \
torch/CMakeLists.txt \
|| die
# tensorpipe is in system, not a build target of caffe2
sed -e '/target_compile_options_if_supported(tensorpipe/d' -i cmake/Dependencies.cmake || die
# Drop third_party from CMake tree
sed -i \
-e '/add_subdirectory.*third_party/d' \
CMakeLists.txt \
cmake/Dependencies.cmake \
cmake/ProtoBuf.cmake \
aten/src/ATen/CMakeLists.txt \
|| die
# Change libc10* path
sed -i \
-e "/EXPORT/s|DESTINATION lib)|DESTINATION $(get_libdir))|" \
c10/cuda/CMakeLists.txt \
c10/CMakeLists.txt \
c10/hip/CMakeLists.txt \
|| 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
# Gentoo provides torch-xpu-ops from installed source payload, so do not clone/fetch.
sed -i -e '1156,1179c\ # Gentoo: torch-xpu-ops vendored in distfiles.' \
caffe2/CMakeLists.txt || die
sed -i -e '1157d' caffe2/CMakeLists.txt || die
# Upstream XPU finder hardcodes OpenCL lookup to SYCL_ROOT/lib.
# Gentoo ships libOpenCL in system libdirs, so use default search there.
perl -0pi -e 's@find_library\(\n OCL_LIBRARY\n NAMES OpenCL\n HINTS \$\{SYCL_LIBRARY_DIR\}\n NO_DEFAULT_PATH\n\)@find_library(\n OCL_LIBRARY\n NAMES OpenCL\n)@' \
cmake/Modules/FindSYCLToolkit.cmake || die
# intel-llvm 6.3.0 reports `Intel SYCL compiler 6.3.0 ...`, while upstream
# only matches the older `Intel(R) ... Compiler ...` form.
sed -i \
-e '47c\ string(REGEX MATCH "([0-9]+\\.[0-9]+\\.[0-9]+)" SYCL_VERSION_STRING_MATCH "${SYCL_VERSION_STRING}")' \
-e '48c\ set(SYCL_VERSION_STRING_MATCH "${CMAKE_MATCH_1}")' \
cmake/Modules/FindSYCLToolkit.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 \
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 caffe2 depends on systemwide composable_kernel
sed -e "s:third_party/composable_kernel:../composable_kernel-${CK_COMMIT}:g" \
-i aten/src/ATen/CMakeLists.txt || die
# Bug 959808: fix for gfx101x targets
pushd "${WORKDIR}/composable_kernel-${CK_COMMIT}" > /dev/null || die
eapply "${FILESDIR}"/composable-kernel-7fe50dc-expand-isa.patch
popd > /dev/null || die
if tc-is-clang; then
# Systemwide gcc (for absl and at::TensorBase) + hipcc (llvm>=18) need abi-compat=17.
# But systemwide clang>=18 + hipcc (>=llvm-18) need opposite!
# See also: https://github.com/llvm/llvm-project/issues/102443#issuecomment-2329726287
sed -e '/-fclang-abi-compat=17/d' -i cmake/Dependencies.cmake || die
fi
# 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 || die
ebegin "HIPifying cuda sources"
FBCODE_BUILD_TOOL="buck" ${EPYTHON} tools/amd_build/build_amd.py || die
eend $?
fi
}
src_configure() {
if use xpu; then
# Intel's libsycl in intel-llvm 6.3.0 exposes libstdc++ mangled symbols.
# Keep the whole C++ build on the same stdlib ABI to avoid runtime loader failures.
append-cxxflags -stdlib=libstdc++
fi
if use cuda && [[ -z ${TORCH_CUDA_ARCH_LIST} ]]; then
ewarn "WARNING: caffe2 is being built with its default CUDA compute capabilities: 3.5 and 7.0."
ewarn "These may not be optimal for your GPU."
ewarn ""
ewarn "To configure caffe2 with the CUDA compute capability that is optimal for your GPU,"
ewarn "set TORCH_CUDA_ARCH_LIST in your make.conf, and re-emerge caffe2."
ewarn "For example, to use CUDA capability 7.5 & 3.5, add: TORCH_CUDA_ARCH_LIST=7.5 3.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_TEST=OFF
-DLIBSHM_INSTALL_LIB_SUBDIR="${EPREFIX}"/usr/$(get_libdir)
-DPython_EXECUTABLE="${PYTHON}"
-DTORCH_INSTALL_LIB_DIR="${EPREFIX}"/usr/$(get_libdir)
-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=$(usex kineto)
-DUSE_KLEIDIAI=OFF # TODO
-DUSE_MAGMA=OFF # TODO: In GURU as sci-libs/magma
-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_OPENCL=$(usex opencl)
-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_SYSTEM_XNNPACK=$(usex xnnpack)
-DUSE_TENSORPIPE=$(usex distributed $(usex !rocm))
-DUSE_UCC=OFF
-DUSE_VALGRIND=OFF
-DUSE_XNNPACK=$(usex xnnpack)
-DUSE_XCCL=OFF
-DUSE_XPU=$(usex xpu)
-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:-3.5 7.0}"
-DUSE_NCCL=OFF # TODO: NVIDIA Collective Communication Library
-DCMAKE_CUDA_FLAGS="$(cuda_gccdir -f | tr -d \")"
-DUSE_CUSPARSELT=$(usex cusparselt)
)
[[ -v CUDACXX ]] && export PYTORCH_NVCC="${CUDACXX}"
if use flash; then
export FLASH_ATTENTION_FORCE_BUILD="TRUE"
export FLASH_ATTN_CUDA_ARCHS="${CUDAARCHS:-${TORCH_CUDA_ARCH_LIST:-3.5 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
elif use xpu; then
local _host_clangxx=$(type -ap clang++ | grep -v '^/usr/lib/ccache/bin/clang++$' | head -n1)
local _host_clang_dir=${_host_clangxx%/*}
export SYCL_ROOT="${ESYSROOT}/opt/intel-llvm"
export CCACHE_PATH="${_host_clang_dir}:/usr/bin:${SYCL_ROOT}/bin"
export PATH="/usr/lib/ccache/bin:${_host_clang_dir}:/usr/bin:${SYCL_ROOT}/bin${PATH:+:${PATH}}"
export LD_LIBRARY_PATH="${SYCL_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
export LIBRARY_PATH="${SYCL_ROOT}/lib${LIBRARY_PATH:+:${LIBRARY_PATH}}"
# TODO: try enabling again after XPU ebuild compiles reliably.
export USE_ONEMKL_XPU=OFF
if use intel_targets_gfx-bmg; then
export TORCH_XPU_ARCH_LIST=bmg
else
export TORCH_XPU_ARCH_LIST=none
fi
[[ -n ${TORCH_XPU_ARCH_LIST} ]] && mycmakeargs+=(
-DTORCH_XPU_ARCH_LIST="${TORCH_XPU_ARCH_LIST}"
)
if has ccache ${FEATURES}; then
mycmakeargs+=(
-DCMAKE_SYCL_COMPILER_LAUNCHER=ccache
)
fi
fi
if use onednn; then
mycmakeargs+=(
-DMKLDNN_FOUND=ON
-DMKLDNN_LIBRARIES=dnnl
-DMKLDNN_INCLUDE_DIR="${ESYSROOT}/usr/include/oneapi/dnnl"
)
fi
cmake_src_configure
}
src_compile() {
PYTORCH_BUILD_VERSION=${PV} \
PYTORCH_BUILD_NUMBER=0 \
cmake_src_compile
}
python_install() {
python_domodule python/torch
mkdir "${D}"$(python_get_sitedir)/torch/bin || die
mkdir "${D}"$(python_get_sitedir)/torch/lib || 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
if [[ -e ${ED}/usr/lib/libc10_xpu.so ]]; then
mkdir -p "${ED}/usr/$(get_libdir)" || die
mv "${ED}/usr/lib/libc10_xpu.so" "${ED}/usr/$(get_libdir)/libc10_xpu.so" || die
rmdir --ignore-fail-on-non-empty "${ED}/usr/lib" || die
fi
# Used by pytorch ebuild
insinto "/var/lib/${PN}"
doins "${BUILD_DIR}"/CMakeCache.txt
rm -rf python
mkdir -p python/torch || die
cp torch/version.py python/torch/ || die
python_install
}