User/joporto/dxcmem refactor (#2196)

* Moves the implementation of DxcThreadMalloc to dxcmem.cpp; removes DxcSwapThreadMalloc

* Remove DxcSetThreadMalloc(OrDefault) -- they were always been used for installing the default allocator

* Deletes copy and move ctors and assignment

* stores the allocator in the TLS slot.

* DxcSwapThreadMalloc should be able to install a null allocator

* Marks the DxcThreadMalloc members as hidden (linux only)

include/dxc/Support/Global.h

@@ -48,26 +48,34 @@ void DxcCleanupThreadMalloc() throw();
 // Used by APIs entry points to set up per-thread/invocation allocator.
 // Setting the IMalloc on the thread increases the reference count,
 // clearing it decreases it.
+void DxcSetThreadMallocToDefault() throw();
 void DxcClearThreadMalloc() throw();
-void DxcSetThreadMalloc(IMalloc *pMalloc) throw();
-void DxcSetThreadMallocOrDefault(IMalloc *pMalloc) throw();
-
-// Swapping does not AddRef or Release new or prior. The pattern is to keep both alive,
-// either in TLS, or on the stack to restore later. The returned value is the effective
-// IMalloc also available in TLS.
-IMalloc *DxcSwapThreadMalloc(IMalloc *pMalloc, IMalloc **ppPrior) throw();
-IMalloc *DxcSwapThreadMallocOrDefault(IMalloc *pMalloc, IMalloc **ppPrior) throw();
 
 // Used to retrieve the current invocation's allocator or perform an alloc/free/realloc.
 IMalloc *DxcGetThreadMallocNoRef() throw();
 
-struct DxcThreadMalloc {
+#if defined(LLVM_ON_UNIX)
-  DxcThreadMalloc(IMalloc *pMallocOrNull) throw() {
+#define DXC_HIDDEN_LINKAGE __attribute__(( visibility("hidden") ))
-    p = DxcSwapThreadMallocOrDefault(pMallocOrNull, &pPrior);
+#else  // LLVM_ON_UNIX
-  }
+#define DXC_HIDDEN_LINKAGE
-  ~DxcThreadMalloc() {
+#endif  // LLVM_ON_UNIX
-    DxcSwapThreadMalloc(pPrior, nullptr);
+
-  }
+class DxcThreadMalloc {
+public:
+  explicit DXC_HIDDEN_LINKAGE DxcThreadMalloc(IMalloc *pMallocOrNull) throw();
+  DXC_HIDDEN_LINKAGE ~DxcThreadMalloc();
+
+  IMalloc *GetInstalledAllocator() const { return p; }
+
+private:
+  // Copy constructor and assignment are dangerous and should always be
+  // deleted...
+  DxcThreadMalloc(const DxcThreadMalloc &) = delete;
+  DxcThreadMalloc &operator =(const DxcThreadMalloc &) = delete;
+  // Move constructor and assignment should be OK to be added if needed.
+  DxcThreadMalloc(DxcThreadMalloc &&) = delete;
+  DxcThreadMalloc &operator =(DxcThreadMalloc &&) = delete;
+
   IMalloc *p;
   IMalloc *pPrior;
 };

lib/DxcSupport/dxcmem.cpp

@@ -66,16 +66,15 @@ void DxcClearThreadMalloc() throw() {
   g_ThreadMallocTls->erase();
   pMalloc->Release();
 }
-void DxcSetThreadMalloc(IMalloc *pMalloc) throw() {
+
+void DxcSetThreadMallocToDefault() throw() {
   DXASSERT(g_ThreadMallocTls != nullptr, "else prior to DxcInitThreadMalloc or after DxcCleanupThreadMalloc");
   DXASSERT(DxcGetThreadMallocNoRef() == nullptr, "else nested allocation invoked");
-  g_ThreadMallocTls->set(pMalloc);
+  g_ThreadMallocTls->set(g_pDefaultMalloc);
-  pMalloc->AddRef();
+  g_pDefaultMalloc->AddRef();
 }
-void DxcSetThreadMallocOrDefault(IMalloc *pMalloc) throw() {
+
-  DxcSetThreadMalloc(pMalloc ? pMalloc : g_pDefaultMalloc);
+static IMalloc *DxcSwapThreadMalloc(IMalloc *pMalloc, IMalloc **ppPrior) throw() {
-}
-IMalloc *DxcSwapThreadMalloc(IMalloc *pMalloc, IMalloc **ppPrior) throw() {
   DXASSERT(g_ThreadMallocTls != nullptr, "else prior to DxcInitThreadMalloc or after DxcCleanupThreadMalloc");
   IMalloc *pPrior = DxcGetThreadMallocNoRef();
   if (ppPrior) {
@@ -84,8 +83,13 @@ IMalloc *DxcSwapThreadMalloc(IMalloc *pMalloc, IMalloc **ppPrior) throw() {
   g_ThreadMallocTls->set(pMalloc);
   return pMalloc;
 }
-IMalloc *DxcSwapThreadMallocOrDefault(IMalloc *pMallocOrNull, IMalloc **ppPrior) throw() {
+
-  return DxcSwapThreadMalloc(pMallocOrNull ? pMallocOrNull : g_pDefaultMalloc, ppPrior);
+DXC_HIDDEN_LINKAGE DxcThreadMalloc::DxcThreadMalloc(IMalloc *pMallocOrNull) throw() {
+    p = DxcSwapThreadMalloc(pMallocOrNull ? pMallocOrNull : g_pDefaultMalloc, &pPrior);
+}
+
+DXC_HIDDEN_LINKAGE DxcThreadMalloc::~DxcThreadMalloc() {
+    DxcSwapThreadMalloc(pPrior, nullptr);
 }
 
 #ifndef _WIN32

tools/clang/tools/dxc/dxclib/dxc.cpp

@@ -1095,7 +1095,7 @@ int dxc::main(int argc, const char **argv_) {
   const char *pStage = "Operation";
   int retVal = 0;
   if (FAILED(DxcInitThreadMalloc())) return 1;
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
   try {
     pStage = "Argument processing";
     if (initHlslOptTable()) throw std::bad_alloc();

tools/clang/tools/dxcompiler/DXCompiler.cpp

@@ -51,7 +51,7 @@ static HRESULT InitMaybeFail() throw() {
   HRESULT hr;
   bool fsSetup = false, memSetup = false;
   IFC(DxcInitThreadMalloc());
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
   memSetup = true;
   if (::llvm::sys::fs::SetupPerThreadFileSystem()) {
     hr = E_FAIL;
@@ -86,7 +86,7 @@ HRESULT __attribute__ ((constructor)) DllMain() {
 }
 
 void __attribute__ ((destructor)) DllShutdown() {
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
   ::hlsl::options::cleanupHlslOptTable();
   ::llvm::sys::fs::CleanupPerThreadFileSystem();
   ::llvm::llvm_shutdown();
@@ -105,7 +105,7 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD Reason, LPVOID reserved) {
     result = SUCCEEDED(hr) ? TRUE : FALSE;
   } else if (Reason == DLL_PROCESS_DETACH) {
     DxcEtw_DXCompilerShutdown_Start();
-    DxcSetThreadMallocOrDefault(nullptr);
+    DxcSetThreadMallocToDefault();
     ::hlsl::options::cleanupHlslOptTable();
     ::llvm::sys::fs::CleanupPerThreadFileSystem();
     ::llvm::llvm_shutdown();

tools/clang/tools/dxcompiler/dxcassembler.cpp

@@ -101,7 +101,7 @@ HRESULT STDMETHODCALLTYPE DxcAssembler::AssembleToContainer(
         parseIR(memBuf->getMemBufferRef(), Err, Context);
 
     CComPtr<AbstractMemoryStream> pOutputStream;
-    IFT(CreateMemoryStream(TM.p, &pOutputStream));
+    IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pOutputStream));
     raw_stream_ostream outStream(pOutputStream.p);
 
     // Check for success.
@@ -148,7 +148,7 @@ HRESULT STDMETHODCALLTYPE DxcAssembler::AssembleToContainer(
       flags |= SerializeDxilFlags::DebugNameDependOnSource;
     }
     dxcutil::AssembleToContainer(std::move(M), pResultBlob,
-                                         TM.p, flags,
+                                         TM.GetInstalledAllocator(), flags,
                                          pOutputStream);
 
     IFT(DxcOperationResult::CreateFromResultErrorStatus(pResultBlob, nullptr, S_OK, ppResult));

tools/clang/tools/dxl/dxl.cpp

@@ -124,7 +124,7 @@ int __cdecl main(int argc, _In_reads_z_(argc) char **argv) {
     return 1;
   llvm::sys::fs::AutoCleanupPerThreadFileSystem auto_cleanup_fs;
   if (FAILED(DxcInitThreadMalloc())) return 1;
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
   try {
     llvm::sys::fs::MSFileSystem *msfPtr;
     IFT(CreateMSFileSystemForDisk(&msfPtr));

tools/clang/tools/dxlib-sample/dxlib_sample.cpp

@@ -45,7 +45,7 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD Reason, LPVOID) {
     DisableThreadLibraryCalls(hinstDLL);
 
     DxcInitThreadMalloc();
-    DxcSetThreadMallocOrDefault(nullptr);
+    DxcSetThreadMallocToDefault();
 
     if (hlsl::options::initHlslOptTable()) {
     DxcClearThreadMalloc();
@@ -55,7 +55,7 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD Reason, LPVOID) {
       return TRUE;
     }
   } else if (Reason == DLL_PROCESS_DETACH) {
-    DxcSetThreadMallocOrDefault(nullptr);
+    DxcSetThreadMallocToDefault();
     libshare::LibCacheManager::ReleaseLibCacheManager();
     ::hlsl::options::cleanupHlslOptTable();
     DxcClearThreadMalloc();

tools/clang/tools/dxrfallbackcompiler/DXCompiler.cpp

@@ -44,7 +44,7 @@ static HRESULT InitMaybeFail() throw() {
   HRESULT hr;
   bool fsSetup = false, memSetup = false;
   IFC(DxcInitThreadMalloc());
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
   memSetup = true;
   if (::llvm::sys::fs::SetupPerThreadFileSystem()) {
     hr = E_FAIL;
@@ -84,7 +84,7 @@ BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD Reason, LPVOID reserved) {
     result = SUCCEEDED(hr) ? TRUE : FALSE;
   } else if (Reason == DLL_PROCESS_DETACH) {
     DxcEtw_DXCompilerShutdown_Start();
-    DxcSetThreadMallocOrDefault(nullptr);
+    DxcSetThreadMallocToDefault();
     ::hlsl::options::cleanupHlslOptTable();
     ::llvm::sys::fs::CleanupPerThreadFileSystem();
     ::llvm::llvm_shutdown();

tools/clang/tools/dxrfallbackcompiler/dxcdxrfallbackcompiler.cpp

@@ -251,7 +251,7 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::RenameAndLink(
 
         // Create a diagnostic printer
         CComPtr<AbstractMemoryStream> pDiagStream;
-        IFT(CreateMemoryStream(TM.p, &pDiagStream));
+        IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pDiagStream));
         raw_stream_ostream DiagStream(pDiagStream);
         DiagnosticPrinterRawOStream DiagPrinter(DiagStream);
         PrintDiagnosticContext DiagContext(DiagPrinter);
@@ -318,14 +318,14 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::RenameAndLink(
         if (M)
         {
             CComPtr<AbstractMemoryStream> pOutputStream;
-            IFT(CreateMemoryStream(TM.p, &pOutputStream));
+            IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pOutputStream));
             raw_stream_ostream outStream(pOutputStream.p);
             WriteBitcodeToFile(M.get(), outStream);
             outStream.flush();
 
             // Validation.
             dxcutil::AssembleToContainer(
-                std::move(M), pResultBlob, TM.p, SerializeDxilFlags::None,
+                std::move(M), pResultBlob, TM.GetInstalledAllocator(), SerializeDxilFlags::None,
                 pOutputStream
 #if !DISABLE_GET_CUSTOM_DIAG_ID
                 , Diag
@@ -371,7 +371,7 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::PatchShaderBindingTables(
 
         // Create a diagnostic printer
         CComPtr<AbstractMemoryStream> pDiagStream;
-        IFT(CreateMemoryStream(TM.p, &pDiagStream));
+        IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pDiagStream));
         raw_stream_ostream DiagStream(pDiagStream);
         DiagnosticPrinterRawOStream DiagPrinter(DiagStream);
         PrintDiagnosticContext DiagContext(DiagPrinter);
@@ -404,14 +404,14 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::PatchShaderBindingTables(
         if (M)
         {
             CComPtr<AbstractMemoryStream> pOutputStream;
-            IFT(CreateMemoryStream(TM.p, &pOutputStream));
+            IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pOutputStream));
             raw_stream_ostream outStream(pOutputStream.p);
             WriteBitcodeToFile(M.get(), outStream);
             outStream.flush();
             dxcutil::AssembleToContainer(
                 std::move(M),
                 pResultBlob,
-                TM.p,
+                TM.GetInstalledAllocator(),
                 SerializeDxilFlags::None,
                 pOutputStream);
         }
@@ -459,7 +459,7 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::Link(
 
         // Create a diagnostic printer
         CComPtr<AbstractMemoryStream> pDiagStream;
-        IFT(CreateMemoryStream(TM.p, &pDiagStream));
+        IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pDiagStream));
         raw_stream_ostream DiagStream(pDiagStream);
         DiagnosticPrinterRawOStream DiagPrinter(DiagStream);
         PrintDiagnosticContext DiagContext(DiagPrinter);
@@ -558,14 +558,14 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::Link(
         if (M)
         {
             CComPtr<AbstractMemoryStream> pOutputStream;
-            IFT(CreateMemoryStream(TM.p, &pOutputStream));
+            IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pOutputStream));
             raw_stream_ostream outStream(pOutputStream.p);
             WriteBitcodeToFile(M.get(), outStream);
             outStream.flush();
 
             // Validation.
             HRESULT valHR = dxcutil::ValidateAndAssembleToContainer(
-                std::move(M), pResultBlob, TM.p, SerializeDxilFlags::None,
+                std::move(M), pResultBlob, TM.GetInstalledAllocator(), SerializeDxilFlags::None,
                 pOutputStream,
                 /*bDebugInfo*/ false
 #if !DISABLE_GET_CUSTOM_DIAG_ID
@@ -625,7 +625,7 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::Compile(
 
     // Create a diagnostic printer
     CComPtr<AbstractMemoryStream> pDiagStream;
-    IFT(CreateMemoryStream(TM.p, &pDiagStream));
+    IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pDiagStream));
     raw_stream_ostream DiagStream(pDiagStream);
     DiagnosticPrinterRawOStream DiagPrinter(DiagStream);
     PrintDiagnosticContext DiagContext(DiagPrinter);
@@ -719,14 +719,14 @@ HRESULT STDMETHODCALLTYPE DxcDxrFallbackCompiler::Compile(
     if (M)
     {
       CComPtr<AbstractMemoryStream> pOutputStream;
-      IFT(CreateMemoryStream(TM.p, &pOutputStream));
+      IFT(CreateMemoryStream(TM.GetInstalledAllocator(), &pOutputStream));
       raw_stream_ostream outStream(pOutputStream.p);
       WriteBitcodeToFile(M.get(), outStream);
       outStream.flush();
       dxcutil::AssembleToContainer(
           std::move(M), 
           pResultBlob, 
-          TM.p, 
+          TM.GetInstalledAllocator(),
           SerializeDxilFlags::None,
           pOutputStream);
     }

tools/clang/unittests/HLSL/DxcTestUtils.cpp

@@ -31,7 +31,7 @@ bool TestModuleSetup() {
     return false;
   if (FAILED(DxcInitThreadMalloc()))
     return false;
-  DxcSetThreadMallocOrDefault(nullptr);
+  DxcSetThreadMallocToDefault();
 
   if (hlsl::options::initHlslOptTable()) {
     return false;