Make Kernel Source Code Discoverable

src/include/proteus/impl/CoreLLVM.h

@@ -366,6 +366,35 @@ linkModules(LLVMContext &Ctx,
   return LinkedModule;
 }
 
+inline void pruneDanglingNVVMAnnotations(Module &M) {
+  NamedMDNode *Annotations = M.getNamedMetadata("nvvm.annotations");
+  if (!Annotations)
+    return;
+
+  SmallVector<MDNode *> LiveEntries;
+  for (MDNode *Entry : Annotations->operands()) {
+    // Global DCE nulls out the entry of a function it removes and stripping
+    // debug info rewrites such a node to an empty tuple, !{}.
+    // LLVM's nvvm.annotations which LLVM's later UpgradeNVVMAnnotations
+    // pass reads unconditionally, so we need to prune out such entries.
+    // Otherwise, UpgradeNVVMAnnotations runs out of bounds!
+    if (!Entry || Entry->getNumOperands() == 0)
+      continue;
+
+    if (!mdconst::dyn_extract_or_null<GlobalValue>(Entry->getOperand(0)))
+      continue;
+
+    LiveEntries.push_back(Entry);
+  }
+
+  if (LiveEntries.size() == Annotations->getNumOperands())
+    return;
+
+  Annotations->clearOperands();
+  for (MDNode *Entry : LiveEntries)
+    Annotations->addOperand(Entry);
+}
+
 inline void runCleanupPassPipeline(Module &M) {
   TIMESCOPE("proteus::runCleanupPassPipeline");
   PassBuilder PB;
@@ -387,7 +416,11 @@ inline void runCleanupPassPipeline(Module &M) {
 
   Passes.run(M, MAM);
 
-  StripDebugInfo(M);
+  pruneDanglingNVVMAnnotations(M);
+
+  // Keep line-table debug info so that recorded device modules retain the
+  // kernel's source file and line information.
+  stripNonLineTableDebugInfo(M);
 }
 
 inline void findFunctionsWithU64Metadata(

src/pass/Helpers.h

@@ -81,6 +81,16 @@ inline std::string getUniqueFileID(Module &M) {
   return std::string(Out);
 }
 
+// Resolve the module source path so consumers of the embedded bitcode can
+// locate the translation unit independently of the compile working directory.
+inline std::string getCanonicalSourceFileName(Module &M) {
+  SmallString<256> RealPath;
+  if (llvm::sys::fs::real_path(M.getSourceFileName(), RealPath))
+    return M.getSourceFileName();
+
+  return std::string(RealPath);
+}
+
 inline bool isDeviceKernel(const Function *F) {
   if (!F)
     reportFatalError("Expected non-null function");

src/pass/ProteusPass.cpp

@@ -276,7 +276,11 @@ class ProteusPassImpl {
 
     Passes.run(M, MAM);
 
-    StripDebugInfo(M);
+    proteus::pruneDanglingNVVMAnnotations(M);
+
+    // Keep line-table debug info so that recorded device modules retain the
+    // kernel's source file and line information.
+    stripNonLineTableDebugInfo(M);
   }
 
   void runOptimizationPassPipeline(Module &M) {
@@ -989,7 +993,11 @@ class ProteusPassImpl {
         if (G.hasInternalLinkage())
           G.setLinkage(GlobalValue::ExternalLinkage);
 
-      StripDebugInfo(*PrunedLTOModule);
+      proteus::pruneDanglingNVVMAnnotations(*PrunedLTOModule);
+
+      // Keep line-table debug info so that recorded device modules retain the
+      // kernel's source file and line information.
+      stripNonLineTableDebugInfo(*PrunedLTOModule);
 
       if (verifyModule(*PrunedLTOModule, &errs()))
         reportFatalError(
@@ -1072,6 +1080,7 @@ class ProteusPassImpl {
       };
       auto EmitM = CloneModule(M, VMap, ShouldClone);
       runCleanupPassPipeline(*EmitM);
+      EmitM->setSourceFileName(getCanonicalSourceFileName(M));
 
       emitModuleDevice(M, *EmitM, "tu", true);
     }