Support -export-shaders-only for linker. (#3571)

* Support -export-shaders-only for linker.

include/dxc/HLSL/DxilExportMap.h

@@ -38,10 +38,13 @@ namespace dxilutil {
     typedef ExportMapByString::iterator iterator;
     typedef ExportMapByString::const_iterator const_iterator;
 
-    ExportMap() {}
+    ExportMap():m_ExportShadersOnly(false) {}
     void clear();
     bool empty() const;
 
+    void setExportShadersOnly(bool v) { m_ExportShadersOnly = v; }
+    bool isExportShadersOnly() const { return m_ExportShadersOnly; }
+
     // Iterate export map by string name
     iterator begin() { return m_ExportMap.begin(); }
     const_iterator begin() const { return m_ExportMap.begin(); }
@@ -100,6 +103,7 @@ namespace dxilutil {
     NameSet m_ExportNames;
     NameSet m_NameCollisions;
     NameSet m_UnusedExports;
+    bool    m_ExportShadersOnly;
   };
 }
 

lib/HLSL/DxilLinker.cpp

@@ -119,6 +119,7 @@ public:
     return m_functionNameMap;
   }
   bool IsInitFunc(llvm::Function *F);
+  bool IsEntry(llvm::Function *F);
   bool IsResourceGlobal(const llvm::Constant *GV);
   DxilResourceBase *GetResource(const llvm::Constant *GV);
 
@@ -135,6 +136,7 @@ private:
   DxilModule &m_DM;
   // Map from name to Link info for extern functions.
   llvm::StringMap<std::unique_ptr<DxilFunctionLinkInfo>> m_functionNameMap;
+  llvm::SmallPtrSet<llvm::Function*,4>  m_entrySet;
   // Map from resource link global to resource. MapVector for deterministic iteration.
   llvm::MapVector<const llvm::Constant *, DxilResourceBase *> m_resourceMap;
   // Set of initialize functions for global variable. SetVector for deterministic iteration.
@@ -202,6 +204,8 @@ DxilLib::DxilLib(std::unique_ptr<llvm::Module> pModule)
     }
     m_functionNameMap[F.getName()] =
         llvm::make_unique<DxilFunctionLinkInfo>(&F);
+    if (m_DM.IsEntry(&F))
+      m_entrySet.insert(&F);
   }
 
   // Update internal global name.
@@ -211,6 +215,7 @@ DxilLib::DxilLib(std::unique_ptr<llvm::Module> pModule)
       GV.setName(MID + GV.getName());
     }
   }
+
 }
 
 void DxilLib::FixIntrinsicOverloads() {
@@ -327,6 +332,7 @@ bool DxilLib::HasFunction(std::string &name) {
   return m_functionNameMap.count(name);
 }
 
+bool DxilLib::IsEntry(llvm::Function *F) { return m_entrySet.count(F); }
 bool DxilLib::IsInitFunc(llvm::Function *F) { return m_initFuncSet.count(F); }
 bool DxilLib::IsResourceGlobal(const llvm::Constant *GV) {
   return m_resourceMap.count(GV);
@@ -1466,7 +1472,7 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile, dxilutil::ExportMap &ex
       return nullptr;
 
   } else {
-    if (exportMap.empty()) {
+    if (exportMap.empty() && !exportMap.isExportShadersOnly()) {
       // Add every function for lib profile.
       for (auto &it : m_functionNameMap) {
         StringRef name = it.getKey();
@@ -1496,6 +1502,28 @@ DxilLinkerImpl::Link(StringRef entry, StringRef profile, dxilutil::ExportMap &ex
           }
         }
       }
+    } else if (exportMap.isExportShadersOnly()) {
+      SmallVector<StringRef, 4> workList;
+      for (auto *pLib : m_attachedLibs) {
+        auto &DM = pLib->GetDxilModule();
+        auto *pM = DM.GetModule();
+        for (Function &F : pM->functions()) {
+          if (!pLib->IsEntry(&F)) {
+            if (!F.isDeclaration()) {
+              // Set none entry to be internal so they could be removed.
+              F.setLinkage(GlobalValue::LinkageTypes::InternalLinkage);
+            }
+            continue;
+          }
+          workList.emplace_back(F.getName());
+        }
+        libSet.insert(pLib);
+      }
+
+      if (!AddFunctions(workList, libSet, addedFunctionSet, linkJob,
+                        /*bLazyLoadDone*/ false,
+                        /*bAllowFuncionDecls*/ false))
+        return nullptr;
     } else {
       SmallVector<StringRef, 4> workList;
 

tools/clang/tools/dxcompiler/dxclinker.cpp

@@ -236,6 +236,9 @@ HRESULT STDMETHODCALLTYPE DxcLinker::Link(
     dxilutil::ExportMap exportMap;
     bSuccess = exportMap.ParseExports(opts.Exports, DiagStream);
 
+    if (opts.ExportShadersOnly)
+      exportMap.setExportShadersOnly(true);
+
     bool hasErrorOccurred = !bSuccess;
     if (bSuccess) {
       std::unique_ptr<Module> pM = m_pLinker->Link(

tools/clang/tools/dxl/dxl.cpp

@@ -47,6 +47,7 @@ static cl::opt<std::string> OutputFilename("Fo",
                                            cl::desc("Override output filename"),
                                            cl::value_desc("filename"));
 
+static cl::opt<bool> ExportShadersOnly("export-shaders-only", cl::desc("Export entry function only"));
 
 class DxlContext {
 
@@ -87,9 +88,12 @@ int DxlContext::Link() {
   }
 
   CComPtr<IDxcOperationResult> pLinkResult;
-
+  std::vector<LPCWSTR> args;
+  if (ExportShadersOnly) {
+    args.emplace_back(L"-export-shaders-only");
+  }
   IFT(pLinker->Link(StringRefUtf16(entry), StringRefUtf16(profile),
-                wpInputFiles.data(), wpInputFiles.size(), nullptr, 0,
+                wpInputFiles.data(), wpInputFiles.size(), args.data(), args.size(),
                 &pLinkResult));
 
   HRESULT status;

tools/clang/unittests/HLSL/ExecutionTest.cpp

@@ -9324,7 +9324,7 @@ bool VerifyHelperLaneWaveResults(ExecutionTest::D3D_SHADER_MODEL sm, HelperLaneW
     passed &= HelperLaneResultLogAndVerify(L"QuadReadAcross* - lane 0 / pixel (0,0) - IsHelperLane()", quad_tr_exp.is_helper_across_Diag, quad_tr.is_helper_across_Diag);
   }
 
-  if (sm >= D3D_SHADER_MODEL_6_5) {
+  if (sm >= ExecutionTest::D3D_SHADER_MODEL_6_5) {
     HelperLaneWaveTestResult65& tr65 = testResults.sm65;
     HelperLaneWaveTestResult65& tr65exp = expectedResults.sm65;
     

tools/clang/unittests/HLSL/LinkerTest.cpp

@@ -51,6 +51,7 @@ public:
   TEST_METHOD(RunLinkResRet);
   TEST_METHOD(RunLinkToLib);
   TEST_METHOD(RunLinkToLibExport);
+  TEST_METHOD(RunLinkToLibExportShadersOnly);
   TEST_METHOD(RunLinkFailReDefineGlobal);
   TEST_METHOD(RunLinkFailProfileMismatch);
   TEST_METHOD(RunLinkFailEntryNoProps);
@@ -496,6 +497,28 @@ TEST_F(LinkerTest, RunLinkToLibExport) {
     {L"-exports", L"renamed_test,cloned_test=\\01?mat_test@@YA?AV?$vector@M$02@@V?$vector@M$03@@0AIAV?$matrix@M$03$02@@@Z;main"});
 }
 
+TEST_F(LinkerTest, RunLinkToLibExportShadersOnly) {
+  CComPtr<IDxcBlob> pEntryLib;
+  CompileLib(L"..\\CodeGenHLSL\\linker\\lib_mat_entry2.hlsl",
+             &pEntryLib);
+  CComPtr<IDxcBlob> pLib;
+  CompileLib(L"..\\CodeGenHLSL\\linker\\lib_mat_cast2.hlsl",
+             &pLib);
+
+  CComPtr<IDxcLinker> pLinker;
+  CreateLinker(&pLinker);
+
+  LPCWSTR libName = L"ps_main";
+  RegisterDxcModule(libName, pEntryLib, pLinker);
+
+  LPCWSTR libName2 = L"test";
+  RegisterDxcModule(libName2, pLib, pLinker);
+  Link(L"", L"lib_6_3", pLinker, {libName, libName2},
+    { "@main" },
+    { "@\"\\01?mat_test" },
+    {L"-export-shaders-only"});
+}
+
 TEST_F(LinkerTest, RunLinkFailSelectRes) {
   if (m_ver.SkipDxilVersion(1, 3)) return;
   CComPtr<IDxcBlob> pEntryLib;