From 05ddfcff656c4bab645f7804db5fe2487873513b Mon Sep 17 00:00:00 2001
From: Luna Nova <git@lunnova.dev>
Date: Fri, 12 Sep 2025 21:30:05 -0700
Subject: [PATCH] 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.
---
 .../python-modules/torch/source/default.nix   | 74 ++++++++++---------
 1 file changed, 41 insertions(+), 33 deletions(-)

diff --git a/pkgs/development/python-modules/torch/source/default.nix b/pkgs/development/python-modules/torch/source/default.nix
index e8d2cf720165..d9b25ed09bb9 100644
--- a/pkgs/development/python-modules/torch/source/default.nix
+++ b/pkgs/development/python-modules/torch/source/default.nix
@@ -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)" ""