python3Packages.torchWithRocm: use vendored composable_kernel when no supported targets are set

Allows building torchWithRocm with clr's localGpuTargets set to gfx906
and other arches which composable_kernel does not support.
This commit is contained in:
Luna Nova
2025-09-14 17:45:25 -07:00
parent 43911cb842
commit 05ddfcff65
@@ -192,40 +192,47 @@ let
throw "No GPU targets specified"
);
# Use vendored CK as header only dep if rocmPackages' CK doesn't properly support targets
vendorComposableKernel = rocmSupport && !rocmPackages.composable_kernel.anyMfmaTarget;
rocmtoolkit_joined = symlinkJoin {
name = "rocm-merged";
paths = with rocmPackages; [
rocm-core
clr
rccl
miopen
aotriton
composable_kernel
rocrand
rocblas
rocsparse
hipsparse
rocthrust
rocprim
hipcub
roctracer
rocfft
rocsolver
hipfft
hiprand
hipsolver
hipblas-common
hipblas
hipblaslt
rocminfo
rocm-comgr
rocm-device-libs
rocm-runtime
rocm-smi
clr.icd
hipify
];
paths =
with rocmPackages;
[
rocm-core
clr
rccl
miopen
aotriton
rocrand
rocblas
rocsparse
hipsparse
rocthrust
rocprim
hipcub
roctracer
rocfft
rocsolver
hipfft
hiprand
hipsolver
hipblas-common
hipblas
hipblaslt
rocminfo
rocm-comgr
rocm-device-libs
rocm-runtime
rocm-smi
clr.icd
hipify
]
++ lib.optionals (!vendorComposableKernel) [
composable_kernel
];
# Fix `setuptools` not being found
postBuild = ''
@@ -369,8 +376,9 @@ buildPythonPackage rec {
# Doesn't pick up the environment variable?
substituteInPlace third_party/kineto/libkineto/CMakeLists.txt \
--replace-fail "\''$ENV{ROCM_SOURCE_DIR}" "${rocmtoolkit_joined}"
# Use composable kernel as dependency, rather than built-in third-party
''
# When possible, composable kernel as dependency, rather than built-in third-party
+ lib.optionalString (rocmSupport && !vendorComposableKernel) ''
substituteInPlace aten/src/ATen/CMakeLists.txt \
--replace-fail "list(APPEND ATen_HIP_INCLUDE \''${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/include)" "" \
--replace-fail "list(APPEND ATen_HIP_INCLUDE \''${CMAKE_CURRENT_SOURCE_DIR}/../../../third_party/composable_kernel/library/include)" ""