[LTO] Prevent asm undefined references to be dropped from the output

[davide]

Summary

As discussed on the list (incidentally I already had a patch for that).

[mehdi_amini]

Test?
ThinLTO Path?

[tejohnson]

Thanks! With a few changes noted below this can be fixed for ThinLTO as well.

lib/LTO/LTO.cpp

354

Do this in addSymbolToGlobalRes instead, so it is added for ThinLTO as well

lib/LTO/LTOBackend.cpp

291

Ditto in thinBackend below

[mehdi_amini]

Also it is not clear to me why we track this globally instead of per-file? What is an ASM reference a local symbol for instance?

[davide]

lib/LTO/LTOBackend.cpp

291

thinBackend() is also used by clang and it's not entirely clear to me what should I pass to it as AsmUndefinedReferences. If you know that off of your mind, I can change it. Otherwise I can commit the patch as is for full LTO and you can modify for thin accordingly (sorry, my knowledge of ThinLTO is still rather limited, and I don't want to delay this further).

[tejohnson]

In D24617#543921, @mehdi_amini wrote:

Also it is not clear to me why we track this globally instead of per-file? What is an ASM reference a local symbol for instance?

Yeah at least for ThinLTO perhaps we should call IRObjectFile::CollectAsmUndefinedRefs in the thinBackend to collect this for each module? That would address the issue Davide mentioned about that being called from clang (distributed backends case). For regular LTO it is essentially global since we merge the module.

lib/LTO/LTO.cpp

310

Shouldn't this be: an undefined reference "of" a symbol "used" in asm?

[mehdi_amini]

The problem is that after merge, if two local symbols have a name conflict in LTO, one of them will be renamed.
I'd like to see a test with such cases.

[tejohnson]

In D24617#543995, @mehdi_amini wrote:

The problem is that after merge, if two local symbols have a name conflict in LTO, one of them will be renamed.
I'd like to see a test with such cases.

Hmm, so how would we catch/handle this? The uses from module level assembly won't be renamed, right? (I know that is an issue for inline assembly). That seems orthogonal to the issue of making sure their defs are not dropped by updating llvm.compiler.used appropriately (which could be fixed by doing what I suggested for ThinLTO - call CollectAsmUndefinedRefs from the backend).

[mehdi_amini]

That seems orthogonal to the issue of making sure their defs are not dropped by updating llvm.compiler.used appropriately which could be fixed by doing what I suggested for ThinLTO - call CollectAsmUndefinedRefs from the backend).

Right, this would solves it (but for the lack of renaming).

Another way is to do it on the individual modules when they're loaded, pre-merge.

[davide]

Teresa, is this what you had in mind? Also added a ThinLTO test, hopefully it's correct.

[mehdi_amini]

lib/LTO/LTO.cpp

368

Why don't you update compiler.used here before merging? This way the StringSet could be local to this function and not global to the LTO class/lifetime I believe.

[tejohnson]

In D24617#544082, @davide wrote:

Teresa, is this what you had in mind? Also added a ThinLTO test, hopefully it's correct.

Yes, thanks. One suggestion to merge the new tests below.

test/tools/gold/X86/asm_undefined2_thin.ll

6 ↗

(On Diff #71552)

I suppose since the test is basically the same you could copy these RUN lines into asm_undefined2.ll so it would test both regular and thin LTO.

[davide]

Unified tests as per Teresa's request.

[davide]

Update after Teresa's comments.

[mehdi_amini]

LGTM.

With a llvm-lto2 version of the test, it'll be perfect ;)

[davide]

In D24617#544936, @mehdi_amini wrote:

LGTM.

With a llvm-lto2 version of the test, it'll be perfect ;)

I tried to write that but there are issues providing the resolution because IRObjectFile will report the same symbol twice. I added a comment to https://llvm.org/bugs/show_bug.cgi?id=30396 so that we can add a test for this once we fix IRObjectFile =)

[mehdi_amini]

Thanks for adding the tracking of this!

[tejohnson]

Thanks!

include/llvm/LTO/LTO.h

struct RegularLTOState {
     };
     std::map<std::string, CommonResolution> Commons;
 
     unsigned ParallelCodeGenParallelismLevel;
     LTOLLVMContext Ctx;
     bool HasModule = false;
     std::unique_ptr<Module> CombinedModule;
     std::unique_ptr<IRMover> Mover;
+    StringSet<> AsmUndefinedRefs;
   } RegularLTO;
 
   struct ThinLTOState {
     ThinLTOState(ThinBackend Backend);
 
     ThinBackend Backend;
     ModuleSummaryIndex CombinedIndex;
     MapVector<StringRef, MemoryBufferRef> ModuleMap;

include/llvm/LTO/LTOBackend.h

 // to implement a standalone ThinLTO backend.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LTO_LTOBACKEND_H
 #define LLVM_LTO_LTOBACKEND_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/ModuleSummaryIndex.h"
 #include "llvm/LTO/Config.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/IPO/FunctionImport.h"
 
 namespace llvm {
 
 class Error;
 class Module;
 class Target;
 
 namespace lto {
 
 /// Runs a regular LTO backend.
 Error backend(Config &C, AddOutputFn AddStream,
               unsigned ParallelCodeGenParallelismLevel,
-              std::unique_ptr<Module> M);
+              std::unique_ptr<Module> M, StringSet<> AsmUndefinedRefs);
 
 /// Runs a ThinLTO backend.
 Error thinBackend(Config &C, unsigned Task, AddOutputFn AddStream, Module &M,
                   ModuleSummaryIndex &CombinedIndex,
                   const FunctionImporter::ImportMapTy &ImportList,
                   const GVSummaryMapTy &DefinedGlobals,
                   MapVector<StringRef, MemoryBufferRef> &ModuleMap);
 }
 }
 
 #endif

lib/LTO/LTO.cpp

Error LTO::add(std::unique_ptr<InputFile> Input,
   bool HasThinLTOSummary = hasGlobalValueSummary(MBRef, Conf.DiagHandler);
 
   if (HasThinLTOSummary)
     return addThinLTO(std::move(Input), Res);
   else
     return addRegularLTO(std::move(Input), Res);
 }
 
+static void handleUndefinedAsmRefs(const InputFile::Symbol &Sym,
+                                   GlobalValue *GV,
+                                   StringSet<> &AsmUndefinedRefs) {
+  // GV associated => not an assembly symbol, bail out.
+  if (GV)
+    return;
+
+  // This is an undefined reference of a symbol used in asm. We put that in
+  // compiler.used, so that we can preserve it from being dropped from
+  // the output, without necessarily preventing its internalization.
+  AsmUndefinedRefs.insert(Sym.getName().str());
+}
+
 // Add a regular LTO object to the link.
 Error LTO::addRegularLTO(std::unique_ptr<InputFile> Input,
                          ArrayRef<SymbolResolution> Res) {
   if (!RegularLTO.CombinedModule) {
     RegularLTO.CombinedModule =
         llvm::make_unique<Module>("ld-temp.o", RegularLTO.Ctx);
     RegularLTO.Mover = llvm::make_unique<IRMover>(*RegularLTO.CombinedModule);
   }
Error LTO::addRegularLTO(std::unique_ptr<InputFile> Input,
   for (const InputFile::Symbol &Sym :
        make_range(InputFile::symbol_iterator(Obj->symbol_begin()),
                   InputFile::symbol_iterator(Obj->symbol_end()))) {
     assert(ResI != Res.end());
     SymbolResolution Res = *ResI++;
     addSymbolToGlobalRes(Obj.get(), Used, Sym, Res, 0);
 
     GlobalValue *GV = Obj->getSymbolGV(Sym.I->getRawDataRefImpl());
-    if (Sym.getFlags() & object::BasicSymbolRef::SF_Undefined)
+    if (Sym.getFlags() & object::BasicSymbolRef::SF_Undefined) {
+      handleUndefinedAsmRefs(Sym, GV, RegularLTO.AsmUndefinedRefs);
       continue;
+    }
     if (Res.Prevailing && GV) {
       Keep.push_back(GV);
       switch (GV->getLinkage()) {
       default:
         break;
       case GlobalValue::LinkOnceAnyLinkage:
         GV->setLinkage(GlobalValue::WeakAnyLinkage);
         break;
if (Sym.getFlags() & object::BasicSymbolRef::SF_Common) {
       CommonRes.Size = std::max(CommonRes.Size, Sym.getCommonSize());
       CommonRes.Align = std::max(CommonRes.Align, Sym.getCommonAlignment());
       CommonRes.Prevailing |= Res.Prevailing;
     }
 
     // FIXME: use proposed local attribute for FinalDefinitionInLinkageUnit.
   }
   assert(ResI == Res.end());
 
   return RegularLTO.Mover->move(Obj->takeModule(), Keep,
                                 [](GlobalValue &, IRMover::ValueAdder) {});
 }
 
 // Add a ThinLTO object to the link.
 Error LTO::addThinLTO(std::unique_ptr<InputFile> Input,
                       ArrayRef<SymbolResolution> Res) {
   Module &M = Input->Obj->getModule();
for (const auto &R : GlobalResolutions) {
         GV->setLinkage(GlobalValue::InternalLinkage);
     }
 
     if (Conf.PostInternalizeModuleHook &&
         !Conf.PostInternalizeModuleHook(0, *RegularLTO.CombinedModule))
       return Error();
   }
   return backend(Conf, AddOutput, RegularLTO.ParallelCodeGenParallelismLevel,
-                 std::move(RegularLTO.CombinedModule));
+                 std::move(RegularLTO.CombinedModule),
+                 RegularLTO.AsmUndefinedRefs);
 }
 
 /// This class defines the interface to the ThinLTO backend.
 class lto::ThinBackendProc {
 protected:
   Config &Conf;
   ModuleSummaryIndex &CombinedIndex;
   const StringMap<GVSummaryMapTy> &ModuleToDefinedGVSummaries;

lib/LTO/LTOBackend.cpp

 #include "llvm/Analysis/LoopPassManager.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Bitcode/ReaderWriter.h"
 #include "llvm/IR/LegacyPassManager.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/LTO/LTO.h"
+#include "llvm/LTO/legacy/UpdateCompilerUsed.h"
 #include "llvm/MC/SubtargetFeature.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/ThreadPool.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/IPO.h"
if (!T)
     return make_error<StringError>(Msg, inconvertibleErrorCode());
   return T;
 }
 
 }
 
 Error lto::backend(Config &C, AddOutputFn AddOutput,
                    unsigned ParallelCodeGenParallelismLevel,
-                   std::unique_ptr<Module> Mod) {
+                   std::unique_ptr<Module> Mod, StringSet<> AsmUndefinedRefs) {
   Expected<const Target *> TOrErr = initAndLookupTarget(C, *Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
   std::unique_ptr<TargetMachine> TM =
       createTargetMachine(C, Mod->getTargetTriple(), *TOrErr);
 
+  // Update llvm.compiler.used to that optimizations won't strip
+  // off asm undefined references.
+  updateCompilerUsed(*Mod, *TM, AsmUndefinedRefs);
+
   if (!C.CodeGenOnly)
     if (!opt(C, TM.get(), 0, *Mod, /*IsThinLto=*/false))
       return Error();
 
   if (ParallelCodeGenParallelismLevel == 1) {
     codegen(C, TM.get(), getUncachedOutputWrapper(AddOutput, 0), 0, *Mod);
   } else {
     splitCodeGen(C, TM.get(), AddOutput, ParallelCodeGenParallelismLevel,
                  std::move(Mod));
   }
   return Error();
 }
 
 Error lto::thinBackend(Config &Conf, unsigned Task, AddOutputFn AddOutput,
                        Module &Mod, ModuleSummaryIndex &CombinedIndex,
                        const FunctionImporter::ImportMapTy &ImportList,
                        const GVSummaryMapTy &DefinedGlobals,
                        MapVector<StringRef, MemoryBufferRef> &ModuleMap) {
   Expected<const Target *> TOrErr = initAndLookupTarget(Conf, Mod);
   if (!TOrErr)
     return TOrErr.takeError();
 
+  StringRef TheTriple = Mod.getTargetTriple();
+
   std::unique_ptr<TargetMachine> TM =
-      createTargetMachine(Conf, Mod.getTargetTriple(), *TOrErr);
+      createTargetMachine(Conf, TheTriple, *TOrErr);
+
+  // Collect the list of undefined symbols used in asm and update
+  // llvm.compiler.used to prevent optimization to drop these from the output.
+  StringSet<> AsmUndefinedRefs;
+  object::IRObjectFile::CollectAsmUndefinedRefs(
+      Triple(TheTriple), Mod.getModuleInlineAsm(),
+      [&AsmUndefinedRefs](StringRef Name, object::BasicSymbolRef::Flags Flags) {
+        if (Flags & object::BasicSymbolRef::SF_Undefined)
+          AsmUndefinedRefs.insert(Name);
+      });
+  updateCompilerUsed(Mod, *TM, AsmUndefinedRefs);
 
   if (Conf.CodeGenOnly) {
     codegen(Conf, TM.get(), AddOutput, Task, Mod);
     return Error();
   }
 
   if (Conf.PreOptModuleHook && !Conf.PreOptModuleHook(Task, Mod))
     return Error();

test/tools/gold/X86/asm_undefined2.ll

+; RegularLTO testcase
+; RUN: llvm-as %s -o %t.o
+; RUN: %gold -shared -m elf_x86_64 -o %t2 -plugin %llvmshlibdir/LLVMgold.so \
+; RUN: %t.o --plugin-opt=save-temps -upatatino
+; RUN: llvm-dis < %t2.0.5.precodegen.bc | FileCheck %s
+
+; ThinLTO testcase
+; RUN: opt -module-summary %s -o %t.o
+; RUN: %gold -m elf_x86_64 -plugin %llvmshlibdir/LLVMgold.so \
+; RUN:     --plugin-opt=save-temps \
+; RUN:     --plugin-opt=thinlto -o %t2 %t.o
+; RUN: llvm-dis < %t.o.5.precodegen.bc | FileCheck %s
+
+; Check that foo is properly appended to llvm.compiler.used
+; CHECK: @llvm.compiler.used = appending global [1 x i8*] [i8* bitcast (void ()* @foo to i8*)], section "llvm.metadata"
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+module asm ".global patatino"
+module asm ".equ patatino, foo"
+
+declare void @patatino()
+
+define void @foo() {
+  call void @patatino()
+  ret void
+}