From 17d5e93e6d5d9ff1f9f46b39ec2ce89ace2d60ed Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Wed, 29 Jan 2025 16:55:52 -0500
Subject: [PATCH] cmLinkLineDeviceComputer: Support fatbin device linking

Fixes #26651
---
 Source/cmGeneratorTarget_Sources.cxx |  4 ++-
 Source/cmLinkLineDeviceComputer.cxx  | 47 ++++++++++++++++++++++++++--
 Source/cmLinkLineDeviceComputer.h    |  3 ++
 3 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/Source/cmGeneratorTarget_Sources.cxx b/Source/cmGeneratorTarget_Sources.cxx
index 742d0d9f25d..b7e54aafda0 100644
--- a/Source/cmGeneratorTarget_Sources.cxx
+++ b/Source/cmGeneratorTarget_Sources.cxx
@@ -28,6 +28,7 @@
 #include "cmGeneratorExpressionDAGChecker.h"
 #include "cmGlobalGenerator.h"
 #include "cmLinkItem.h"
+#include "cmLinkLineDeviceComputer.h"
 #include "cmList.h"
 #include "cmListFileCache.h"
 #include "cmLocalGenerator.h"
@@ -55,7 +56,8 @@ void AddObjectEntries(cmGeneratorTarget const* headTarget,
     entries.HadContextSensitiveCondition = impl->HadContextSensitiveCondition;
     for (cmLinkImplItem const& lib : impl->Libraries) {
       if (lib.Target &&
-          lib.Target->GetType() == cmStateEnums::OBJECT_LIBRARY) {
+          lib.Target->GetType() == cmStateEnums::OBJECT_LIBRARY &&
+          !isObjectLibraryJustForDeviceLinking(*lib.Target, config)) {
         std::string uniqueName =
           headTarget->GetGlobalGenerator()->IndexGeneratorTargetUniquely(
             lib.Target);
diff --git a/Source/cmLinkLineDeviceComputer.cxx b/Source/cmLinkLineDeviceComputer.cxx
index dac37acb592..34b90c461df 100644
--- a/Source/cmLinkLineDeviceComputer.cxx
+++ b/Source/cmLinkLineDeviceComputer.cxx
@@ -113,10 +113,12 @@ void cmLinkLineDeviceComputer::ComputeLinkLibraries(
       switch (item.Target->GetType()) {
         case cmStateEnums::SHARED_LIBRARY:
         case cmStateEnums::MODULE_LIBRARY:
-        case cmStateEnums::OBJECT_LIBRARY:
         case cmStateEnums::INTERFACE_LIBRARY:
           skip = true;
           break;
+        case cmStateEnums::OBJECT_LIBRARY:
+          skip = !isObjectLibraryJustForDeviceLinking(*item.Target, config);
+          break;
         case cmStateEnums::STATIC_LIBRARY:
           skip = item.Target->GetPropertyAsBool("CUDA_RESOLVE_DEVICE_SYMBOLS");
           break;
@@ -148,7 +150,19 @@ void cmLinkLineDeviceComputer::ComputeLinkLibraries(
       // name is specified as a following item.  Ignore both.
       skipItemAfterFramework = true;
       continue;
-    } else if (cmLinkItemValidForDevice(item.Value.Value)) {
+    } else if (item.Target && item.Target->GetType() == cmStateEnums::OBJECT_LIBRARY) {
+      std::vector<std::string> objects;
+      item.Target->GetTargetObjectNames(config, objects);
+      std::string const obj_dir = item.Target->GetObjectDirectory(config);
+      for (std::string& o : objects) {
+        o = cmStrCat(obj_dir, o, " ");
+        if (emitted.insert(o).second) {
+          std::cout << "\t" << o << std::endl;
+          linkLibraries.emplace_back(o);
+        }
+      }
+    }
+     else if (cmLinkItemValidForDevice(item.Value.Value)) {
       linkLib.Value = item.Value.Value;
     }
 
@@ -237,3 +251,32 @@ bool requireDeviceLinking(cmGeneratorTarget& target, cmLocalGenerator& lg,
   }
   return false;
 }
+
+bool isObjectLibraryJustForDeviceLinking(cmGeneratorTarget const& target,
+                                         std::string const& config)
+
+{
+  if (target.GetType() != cmStateEnums::OBJECT_LIBRARY) {
+    return false;
+  }
+
+  bool has_rdc = target.GetPropertyAsBool("CUDA_SEPARABLE_COMPILATION");
+  bool has_ipo = target.GetPropertyAsBool("INTERPROCEDURAL_OPTIMIZATION");
+  if (!(has_rdc && has_ipo)) {
+    return false;
+  }
+
+  // Using an array so we can easily expand as other types
+  // become supported by nvlink
+  static std::array<cm::string_view, 4> const compileModes {
+    {"FATBIN"_s};
+    for (cm::string_view mode : compileModes) {
+      auto propName = cmStrCat("CUDA_", mode, "_COMPILATION");
+      bool enabled = target.GetPropertyAsBool(propName);
+      if (enabled) {
+        return true;
+      }
+    }
+
+    return false;
+  }
diff --git a/Source/cmLinkLineDeviceComputer.h b/Source/cmLinkLineDeviceComputer.h
index c7c5c1755cd..361d0694766 100644
--- a/Source/cmLinkLineDeviceComputer.h
+++ b/Source/cmLinkLineDeviceComputer.h
@@ -42,3 +42,6 @@ public:
 
 bool requireDeviceLinking(cmGeneratorTarget& target, cmLocalGenerator& lg,
                           std::string const& config);
+
+bool isObjectLibraryJustForDeviceLinking(cmGeneratorTarget const& target,
+                                         std::string const& config);
-- 
GitLab