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:
@@ -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)" ""
|
||||
|
||||
Reference in New Issue
Block a user