[lld][WebAssebmly] Preserve custom import attributes with LTO

[sbc100]

Summary

Undefined symbols in WebAssembly can come with custom import-module
and import-field attributes. However when reading symbols from
bitcode object files during LTO those curtom attributes are not
available.

Once we compile the LTO object and read in the symbol table from the
object file we have access to these custom attributes. In this case,
when undefined symbols are added and a symbol already exists in the
SymbolTable we can't simple return it, we may need to update the
symbol's attributes.

Fixes: PR43211

[sbc100]

• .

[sbc100]

• revert debugging
• cleanup

[ruiu]

It seems that the patch description is too terse. Please explain the bug that this patch is fixing in the commit message.

lld/wasm/Driver.cpp

580–589

Can you add a comment to explain what this wrapper provides?

lld/wasm/SymbolTable.cpp

396

Can you add a comment to explain what this function is supposed to do?

[sbc100]

• feedback

[sbc100]

• feedback

[sbc100]

Sorry, I've added a lot more information to the description and in comments.

[sbc100]

Also see https://bugs.llvm.org/show_bug.cgi?id=43211

[ruiu]

lld/wasm/Driver.cpp

581

Is this code tested? (i.e. is there any test that will fail if we undo this change?)

lld/wasm/SymbolTable.cpp

453

nit: auto* -> auto *

[sbc100]

• clang-format

[sbc100]

• revert part

[sbc100]

lld/wasm/Driver.cpp

581

Hmmm, you are right, this change was needed as part of an earlier version of this change. Reverting it works fine.

[sbc100]

• revert

[ruiu]

LGTM

lld/test/wasm/Inputs/import-attributes.ll

+target triple = "wasm32-unknown-unknown"
+
+define void @call_foo() {
+  call void @foo();
+  ret void
+}
+
+declare void @foo() #0
+
+attributes #0 = { "wasm-import-module"="baz" }

lld/test/wasm/import-attribute-mismatch.ll

+; RUN: llc -filetype=obj %s -o %t1.o
+; RUN: llc -filetype=obj %S/Inputs/import-attributes.ll -o %t2.o
+; RUN: not wasm-ld --export call_foo --allow-undefined -o %t.wasm %t1.o %t2.o 2>&1 | FileCheck %s
+
+target triple = "wasm32-unknown-unknown-wasm"
+
+define void @_start() {
+  call void @foo();
+  ret void
+}
+
+declare void @foo() #0
+
+attributes #0 = { "wasm-import-module"="bar" }
+
+; CHECK: wasm-ld: error: import module mismatch for symbol: foo
+; CHECK: >>> defined as bar in {{.*}}1.o
+; CHECK: >>> defined as baz in {{.*}}2.o

lld/test/wasm/lto/import-attributes.ll

+; RUN: llvm-as %s -o %t.o
+; RUN: wasm-ld --allow-undefined -o %t.wasm %t.o
+; RUN: obj2yaml %t.wasm | FileCheck %s
+
+target triple = "wasm32-unknown-unknown-wasm"
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+
+define void @_start() {
+  call void @foo();
+  ret void
+}
+
+declare void @foo() #0
+
+attributes #0 = { "wasm-import-module"="bar" "wasm-import-name"="customfoo" }
+
+; CHECK:        - Type:            IMPORT
+; CHECK-NEXT:     Imports:
+; CHECK-NEXT:       - Module:          bar
+; CHECK-NEXT:         Field:           customfoo
+; CHECK-NEXT:         Kind:            FUNCTION
+; CHECK-NEXT:         SigIndex:        0

lld/wasm/Driver.cpp

 //
 // This data structure is instantiated for each -wrap option.
 struct WrappedSymbol {
   Symbol *sym;
   Symbol *real;
   Symbol *wrap;
 };
 
-static Symbol *addUndefined(StringRef name) {
-  return symtab->addUndefinedFunction(name, "", "", WASM_SYMBOL_UNDEFINED,
-                                      nullptr, nullptr, false);
+static Symbol *addUndefinedWrapper(Symbol* target, StringRef name) {
+  StringRef importModule = defaultModule;
+  StringRef importName = name;
+  if (auto uf = dyn_cast<UndefinedFunction>(target)) {
+    importModule = uf->importModule;
+    importName = uf->importName;
+  }
+  return symtab->addUndefinedFunction(name, importName, importModule,
+                                      WASM_SYMBOL_UNDEFINED, nullptr, nullptr,
+                                      false);
 }
 
 // Handles -wrap option.
 //
 // This function instantiates wrapper symbols. At this point, they seem
 // like they are not being used at all, so we explicitly set some flags so
 // that LTO won't eliminate them.
 static std::vector<WrappedSymbol> addWrappedSymbols(opt::InputArgList &args) {
   std::vector<WrappedSymbol> v;
   DenseSet<StringRef> seen;
 
   for (auto *arg : args.filtered(OPT_wrap)) {
     StringRef name = arg->getValue();
     if (!seen.insert(name).second)
       continue;
 
     Symbol *sym = symtab->find(name);
     if (!sym)
       continue;
 
-    Symbol *real = addUndefined(saver.save("__real_" + name));
-    Symbol *wrap = addUndefined(saver.save("__wrap_" + name));
+    Symbol *real = addUndefinedWrapper(sym, saver.save("__real_" + name));
+    Symbol *wrap = addUndefinedWrapper(sym, saver.save("__wrap_" + name));
     v.push_back({sym, real, wrap});
 
     // We want to tell LTO not to inline symbols to be overwritten
     // because LTO doesn't know the final symbol contents after renaming.
     real->canInline = false;
     sym->canInline = false;
 
     // Tell LTO not to eliminate these symbols.

lld/wasm/InputFiles.cpp

static Symbol *createBitcodeSymbol(const std::vector<bool> &keptComdats,
   flags |= mapVisibility(objSym.getVisibility());
 
   int c = objSym.getComdatIndex();
   bool excludedByComdat = c != -1 && !keptComdats[c];
 
   if (objSym.isUndefined() || excludedByComdat) {
     flags |= WASM_SYMBOL_UNDEFINED;
     if (objSym.isExecutable())
-      return symtab->addUndefinedFunction(name, name, defaultModule, flags, &f,
-                                          nullptr, true);
+      return symtab->addUndefinedFunction(name, "", "", flags, &f, nullptr,
+                                          true);
     return symtab->addUndefinedData(name, flags, &f);
   }
 
   if (objSym.isExecutable())
     return symtab->addDefinedFunction(name, flags, &f, nullptr);
   return symtab->addDefinedData(name, flags, &f, nullptr, 0, 0);
 }
 

lld/wasm/LTO.cpp

 }
 
 BitcodeCompiler::BitcodeCompiler() : ltoObj(createLTO()) {}
 
 BitcodeCompiler::~BitcodeCompiler() = default;
 
 static void undefine(Symbol *s) {
   if (auto f = dyn_cast<DefinedFunction>(s))
-    replaceSymbol<UndefinedFunction>(f, f->getName(), f->getName(),
-                                     defaultModule, 0,
+    replaceSymbol<UndefinedFunction>(f, f->getName(), f->getName(), "", 0,
                                      f->getFile(), f->signature);
   else if (isa<DefinedData>(s))
     replaceSymbol<UndefinedData>(s, s->getName(), 0, s->getFile());
   else
     llvm_unreachable("unexpected symbol kind");
 }
 
 void BitcodeCompiler::add(BitcodeFile &f) {

lld/wasm/SymbolTable.cpp

Symbol *SymbolTable::addDefinedEvent(StringRef name, uint32_t flags,
 
   checkEventType(s, file, &event->getType(), &event->signature);
 
   if (shouldReplace(s, file, flags))
     replaceSym();
   return s;
 }
 
+template <typename T>
+static void setImportAttributes(T *sym, StringRef importName,
+                                StringRef importModule, InputFile *file) {
+  // Silently allow for empty atttibutes to be replaced with non-empty ones.
+  // This occurs during LTO when the initial undefined symbol in the bitcode
+  // is replaced with the one in the LTO object which contains the extra
+  // symbol attributes.
+  if (sym->importName.empty() && !importName.empty()) {
+    sym->importName = importName;
+    dbgs() << "setting import name: " << importName << "\n";
+  }
+  if (sym->importModule.empty() && !importModule.empty()) {
+    sym->importModule = importModule;
+    dbgs() << "setting import module: " << importName << "\n";
+  }
+  if (sym->importModule != importModule)
+    error("import module mismatch for symbol: " + toString(*sym) +
+          "\n>>> defined as " + sym->importModule + " in " +
+          toString(sym->getFile()) + "\n>>> defined as " + importModule +
+          " in " + toString(file));
+  if (sym->importName != importName)
+    error("import name mismatch for symbol: " + toString(*sym) + "\n>>> defined as " +
+          sym->importName + " in " + toString(sym->getFile()) +
+          "\n>>> defined as " + importName + " in " + toString(file));
+}
+
 Symbol *SymbolTable::addUndefinedFunction(StringRef name, StringRef importName,
                                           StringRef importModule,
                                           uint32_t flags, InputFile *file,
                                           const WasmSignature *sig,
                                           bool isCalledDirectly) {
   LLVM_DEBUG(dbgs() << "addUndefinedFunction: " << name << " ["
                     << (sig ? toString(*sig) : "none")
                     << "] IsCalledDirectly:" << isCalledDirectly << "\n");
Symbol *SymbolTable::addUndefinedFunction(StringRef name, StringRef importName,
   else if (auto *lazy = dyn_cast<LazySymbol>(s))
     lazy->fetch();
   else {
     auto existingFunction = dyn_cast<FunctionSymbol>(s);
     if (!existingFunction) {
       reportTypeError(s, file, WASM_SYMBOL_TYPE_FUNCTION);
       return s;
     }
+    auto* existingUndefined = dyn_cast<UndefinedFunction>(existingFunction);
     if (!existingFunction->signature && sig)
       existingFunction->signature = sig;
     if (isCalledDirectly && !signatureMatches(existingFunction, sig)) {
-      auto* existingUndefined = dyn_cast<UndefinedFunction>(existingFunction);
       // If the existing undefined functions is not called direcltly then let
       // this one take precedence.  Otherwise the existing function is either
       // direclty called or defined, in which case we need a function variant.
       if (existingUndefined && !existingUndefined->isCalledDirectly)
         replaceSym();
       else if (getFunctionVariant(s, sig, file, &s))
         replaceSym();
     }
+    if (existingUndefined)
+      setImportAttributes(existingUndefined, importName, importModule, file);
   }
 
   return s;
 }
 
 Symbol *SymbolTable::addUndefinedData(StringRef name, uint32_t flags,
                                       InputFile *file) {
   LLVM_DEBUG(dbgs() << "addUndefinedData: " << name << "\n");
Symbol *SymbolTable::addUndefinedGlobal(StringRef name, StringRef importName,
   if (s->traced)
     printTraceSymbolUndefined(name, file);
 
   if (wasInserted)
     replaceSymbol<UndefinedGlobal>(s, name, importName, importModule, flags,
                                    file, type);
   else if (auto *lazy = dyn_cast<LazySymbol>(s))
     lazy->fetch();
-  else if (s->isDefined())
+  else if (s->isDefined()) {
     checkGlobalType(s, file, type);
+    if (auto existingUndefined = dyn_cast<UndefinedGlobal>(s))
+      setImportAttributes(existingUndefined, importName, importModule, file);
+  }
   return s;
 }
 
 void SymbolTable::addLazy(ArchiveFile *file, const Archive::Symbol *sym) {
   LLVM_DEBUG(dbgs() << "addLazy: " << sym->getName() << "\n");
   StringRef name = sym->getName();
 
   Symbol *s;