Unverified Commit 05ddfcff authored by Luna Nova's avatar Luna Nova
Browse files

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.
parent 43911cb8

pkgs/development/python-modules/torch/source/default.nix

+41 −33
Original line number Diff line number Diff line
@@ -192,16 +192,20 @@ 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; [

    paths =

      with rocmPackages;

      [

        rocm-core

        clr

        rccl

        miopen

        aotriton

      composable_kernel

        rocrand

        rocblas

        rocsparse
@@ -225,6 +229,9 @@ let 
        rocm-smi

        clr.icd

        hipify

      ]

      ++ lib.optionals (!vendorComposableKernel) [

        composable_kernel

      ];



    # Fix `setuptools` not being found
@@ -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)" ""