Merged PR 29: Fix linking to library with unresolved function declarations.

Fix linking to library with unresolved function declarations.

lib/HLSL/DxilLinker.cpp

@@ -145,7 +145,8 @@ private:
   bool DetachLib(DxilLib *lib);
   bool AddFunctions(SmallVector<StringRef, 4> &workList,
                     DenseSet<DxilLib *> &libSet, StringSet<> &addedFunctionSet,
-                    DxilLinkJob &linkJob, bool bLazyLoadDone);
+                    DxilLinkJob &linkJob, bool bLazyLoadDone,
+                    bool bAllowFuncionDecls);
   // Attached libs to link.
   std::unordered_set<DxilLib *> m_attachedLibs;
   // Owner of all DxilLib.
@@ -330,7 +331,7 @@ struct DxilLinkJob {
 
 private:
   void LinkNamedMDNodes(Module *pM, ValueToValueMapTy &vmap);
-  void AddDxilOperations(Module *pM);
+  void AddFunctionDecls(Module *pM);
   bool AddGlobals(DxilModule &DM, ValueToValueMapTy &vmap);
   void CloneFunctions(ValueToValueMapTy &vmap);
   void AddFunctions(DxilModule &DM, ValueToValueMapTy &vmap,
@@ -338,7 +339,7 @@ private:
   bool AddResource(DxilResourceBase *res, llvm::GlobalVariable *GV);
   void AddResourceToDM(DxilModule &DM);
   std::unordered_map<DxilFunctionLinkInfo *, DxilLib *> m_functionDefs;
-  llvm::StringMap<llvm::Function *> m_dxilFunctions;
+  llvm::StringMap<llvm::Function *> m_functionDecls;
   // New created functions.
   llvm::StringMap<llvm::Function *> m_newFunctions;
   // New created globals.
@@ -549,8 +550,8 @@ void DxilLinkJob::LinkNamedMDNodes(Module *pM, ValueToValueMapTy &vmap) {
   }
 }
 
-void DxilLinkJob::AddDxilOperations(Module *pM) {
-  for (auto &it : m_dxilFunctions) {
+void DxilLinkJob::AddFunctionDecls(Module *pM) {
+  for (auto &it : m_functionDecls) {
     Function *F = it.second;
     Function *NewF = Function::Create(F->getFunctionType(), F->getLinkage(),
                                       F->getName(), pM);
@@ -696,7 +697,7 @@ DxilLinkJob::Link(std::pair<DxilFunctionLinkInfo *, DxilLib *> &entryLinkPair,
   // Set target.
   pM->setTargetTriple(entryDM.GetModule()->getTargetTriple());
   // Add dxil operation functions before create DxilModule.
-  AddDxilOperations(pM.get());
+  AddFunctionDecls(pM.get());
 
   // Create DxilModule.
   const bool bSkipInit = true;
@@ -784,8 +785,8 @@ DxilLinkJob::LinkToLib(const ShaderModel *pSM) {
       llvm::make_unique<Module>("merged_lib", tmpDM.GetCtx());
   // Set target.
   pM->setTargetTriple(tmpDM.GetModule()->getTargetTriple());
-  // Add dxil operation functions before create DxilModule.
-  AddDxilOperations(pM.get());
+  // Add dxil operation functions and external decls before create DxilModule.
+  AddFunctionDecls(pM.get());
 
   // Create DxilModule.
   const bool bSkipInit = true;
@@ -874,7 +875,7 @@ void DxilLinkJob::AddFunction(
 }
 
 void DxilLinkJob::AddFunction(llvm::Function *F) {
-  m_dxilFunctions[F->getName()] = F;
+  m_functionDecls[F->getName()] = F;
 }
 
 void DxilLinkJob::RunPreparePass(Module &M) {
@@ -1026,7 +1027,8 @@ bool DxilLinkerImpl::DetachLib(DxilLib *lib) {
 bool DxilLinkerImpl::AddFunctions(SmallVector<StringRef, 4> &workList,
                                   DenseSet<DxilLib *> &libSet,
                                   StringSet<> &addedFunctionSet,
-                                  DxilLinkJob &linkJob, bool bLazyLoadDone) {
+                                  DxilLinkJob &linkJob, bool bLazyLoadDone,
+                                  bool bAllowFuncionDecls) {
   while (!workList.empty()) {
     StringRef name = workList.pop_back_val();
     // Ignore added function.
@@ -1052,9 +1054,14 @@ bool DxilLinkerImpl::AddFunctions(SmallVector<StringRef, 4> &workList,
       if (hlsl::OP::IsDxilOpFunc(F) || F->isIntrinsic()) {
         // Add dxil operations directly.
         linkJob.AddFunction(F);
-      } else {
-        // Push function name to work list.
-        workList.emplace_back(F->getName());
+      } else if (addedFunctionSet.count(F->getName()) == 0) {
+        if (bAllowFuncionDecls && F->isDeclaration() && !m_functionNameMap.count(F->getName())) {
+          // When linking to lib, use of undefined function is allowed; add directly.
+          linkJob.AddFunction(F);
+        } else {
+          // Push function name to work list.
+          workList.emplace_back(F->getName());
+        }
       }
     }
 
@@ -1104,7 +1111,8 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile,
     workList.emplace_back(entry);
 
     if (!AddFunctions(workList, libSet, addedFunctionSet, linkJob,
-                      /*bLazyLoadDone*/ false))
+                      /*bLazyLoadDone*/ false,
+                      /*bAllowFuncionDecls*/ false))
       return nullptr;
 
   } else {
@@ -1125,6 +1133,19 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile,
 
         addedFunctionSet.insert(name);
       }
+      // Add every dxil function and llvm intrinsic.
+      for (auto *pLib : libSet) {
+        auto &DM = pLib->GetDxilModule();
+        DM.GetOP();
+        auto *pM = DM.GetModule();
+        for (Function &F : pM->functions()) {
+          if (hlsl::OP::IsDxilOpFunc(&F) || F.isIntrinsic() ||
+            (F.isDeclaration() && m_functionNameMap.count(F.getName()) == 0)) {
+            // Add intrinsics and function decls still not defined in any lib
+            linkJob.AddFunction(&F);
+          }
+        }
+      }
     } else {
       SmallVector<StringRef, 4> workList;
 
@@ -1138,20 +1159,10 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile,
       }
 
       if (!AddFunctions(workList, libSet, addedFunctionSet, linkJob,
-                        /*bLazyLoadDone*/ false))
+                        /*bLazyLoadDone*/ false,
+                        /*bAllowFuncionDecls*/ true))
         return nullptr;
     }
-    // Add every dxil functions and llvm intrinsic.
-    for (auto *pLib : libSet) {
-      auto &DM = pLib->GetDxilModule();
-      DM.GetOP();
-      auto *pM = DM.GetModule();
-      for (Function &F : pM->functions()) {
-        if (hlsl::OP::IsDxilOpFunc(&F) || F.isIntrinsic()) {
-          linkJob.AddFunction(&F);
-        }
-      }
-    }
   }
 
   // Save global users.
@@ -1165,10 +1176,13 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile,
     pLib->CollectUsedInitFunctions(addedFunctionSet, workList);
   }
   // Add init functions if used.
-  // All init function already loaded in BuildGlobalUsage, so set bLazyLoad
-  // false here.
+  // All init function already loaded in BuildGlobalUsage,
+  // so set bLazyLoadDone to true here.
+  // Decls should have been added to addedFunctionSet if lib,
+  // so set bAllowFuncionDecls is false here.
   if (!AddFunctions(workList, libSet, addedFunctionSet, linkJob,
-                    /*bLazyLoadDone*/ true))
+                    /*bLazyLoadDone*/ true,
+                    /*bAllowFuncionDecls*/ false))
     return nullptr;
 
   if (!bIsLib) {

tools/clang/test/CodeGenHLSL/shader-compat-suite/lib_unresolved_func1.hlsl

@@ -0,0 +1,23 @@
+// RUN: %dxc -T lib_6_1 %s | FileCheck %s
+
+// CHECK-DAG: define float @"\01?lib1_fn@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?external_fn@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?external_fn1@@YAMXZ"()
+// CHECK-DAG: define float @"\01?call_lib2@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?lib2_fn@@YAMXZ"()
+// CHECK-NOT: @"\01?unused_fn1
+
+float external_fn();
+float external_fn1();
+float lib2_fn();
+float unused_fn1();
+
+float lib1_fn() {
+  if (false)
+    return unused_fn1();
+  return 11.0 * external_fn() * external_fn1();
+}
+
+float call_lib2() {
+  return lib2_fn();
+}

tools/clang/test/CodeGenHLSL/shader-compat-suite/lib_unresolved_func2.hlsl

@@ -0,0 +1,23 @@
+// RUN: %dxc -T lib_6_1 %s | FileCheck %s
+
+// CHECK-DAG: define float @"\01?lib2_fn@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?external_fn@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?external_fn2@@YAMXZ"()
+// CHECK-DAG: define float @"\01?call_lib1@@YAMXZ"()
+// CHECK-DAG: declare float @"\01?lib1_fn@@YAMXZ"()
+// CHECK-NOT: @"\01?unused_fn2
+
+float external_fn();
+float external_fn2();
+float lib1_fn();
+float unused_fn2();
+
+float lib2_fn() {
+  if (false)
+    return unused_fn2();
+  return 22.0 * external_fn() * external_fn2();
+}
+
+float call_lib1() {
+  return lib1_fn();
+}

tools/clang/unittests/HLSL/LinkerTest.cpp

@@ -51,6 +51,8 @@ public:
   TEST_METHOD(RunLinkFailProfileMismatch);
   TEST_METHOD(RunLinkFailEntryNoProps);
   TEST_METHOD(RunLinkFailSelectRes);
+  TEST_METHOD(RunLinkToLibWithUnresolvedFunctions);
+  TEST_METHOD(RunLinkToLibWithUnresolvedFunctionsExports);
 
 
   dxc::DxcDllSupport m_dllSupport;
@@ -393,3 +395,63 @@ TEST_F(LinkerTest, RunLinkFailSelectRes) {
   LinkCheckMsg(L"main", L"ps_6_0", pLinker, {libName, libName2},
                {"Local resource must map to global resource"});
 }
+
+TEST_F(LinkerTest, RunLinkToLibWithUnresolvedFunctions) {
+  LPCWSTR option[] = { L"-Zi" };
+
+  CComPtr<IDxcBlob> pLib1;
+  CompileLib(L"..\\CodeGenHLSL\\shader-compat-suite\\lib_unresolved_func1.hlsl",
+             &pLib1, option, 1);
+  CComPtr<IDxcBlob> pLib2;
+  CompileLib(L"..\\CodeGenHLSL\\shader-compat-suite\\lib_unresolved_func2.hlsl",
+             &pLib2, option, 1);
+
+  CComPtr<IDxcLinker> pLinker;
+  CreateLinker(&pLinker);
+
+  LPCWSTR libName1 = L"lib1";
+  RegisterDxcModule(libName1, pLib1, pLinker);
+
+  LPCWSTR libName2 = L"lib2";
+  RegisterDxcModule(libName2, pLib2, pLinker);
+
+  Link(L"", L"lib_6_2", pLinker, { libName1, libName2 }, {
+    "declare float @\"\\01?external_fn1@@YAMXZ\"()",
+    "declare float @\"\\01?external_fn2@@YAMXZ\"()",
+    "declare float @\"\\01?external_fn@@YAMXZ\"()",
+    "define float @\"\\01?lib1_fn@@YAMXZ\"()",
+    "define float @\"\\01?lib2_fn@@YAMXZ\"()",
+    "define float @\"\\01?call_lib1@@YAMXZ\"()",
+    "define float @\"\\01?call_lib2@@YAMXZ\"()"
+    }, {"declare float @\"\\01?unused_fn1", "declare float @\"\\01?unused_fn2"});
+}
+
+TEST_F(LinkerTest, RunLinkToLibWithUnresolvedFunctionsExports) {
+  LPCWSTR option[] = { L"-Zi" };
+
+  CComPtr<IDxcBlob> pLib1;
+  CompileLib(L"..\\CodeGenHLSL\\shader-compat-suite\\lib_unresolved_func1.hlsl",
+    &pLib1, option, 1);
+  CComPtr<IDxcBlob> pLib2;
+  CompileLib(L"..\\CodeGenHLSL\\shader-compat-suite\\lib_unresolved_func2.hlsl",
+    &pLib2, option, 1);
+
+  CComPtr<IDxcLinker> pLinker;
+  CreateLinker(&pLinker);
+
+  LPCWSTR libName1 = L"lib1";
+  RegisterDxcModule(libName1, pLib1, pLinker);
+
+  LPCWSTR libName2 = L"lib2";
+  RegisterDxcModule(libName2, pLib2, pLinker);
+
+  DxcDefine exports[] = { { L"call_lib1", L"" }, { L"call_lib2", L"" } };
+  LinkWithExports(pLinker, { libName1, libName2 }, exports, {
+    "declare float @\"\\01?external_fn1@@YAMXZ\"()",
+    "declare float @\"\\01?external_fn2@@YAMXZ\"()",
+    "declare float @\"\\01?external_fn@@YAMXZ\"()",
+    "define float @\"\\01?call_lib1@@YAMXZ\"()",
+    "define float @\"\\01?call_lib2@@YAMXZ\"()"
+  }, { "declare float @\"\\01?unused_fn1", "declare float @\"\\01?unused_fn2",
+       "declare float @\"\\01?lib1_fn", "declare float @\"\\01?lib2_fn" });
+}