Allow new/delete operators to be called before DllMain (#5454)

Allow new/delete operators to be called before DllMain and allocator
initialization in DxcInitThreadMalloc().

The operators can be called before DllMain from CRT libraries when
static linking is enabled. If that happens, the new/delete operators
will fallback to the standard allocator and use CoTaskMemAlloc/Free
directly instead of CoGetMalloc, Alloc/Free & Release for better perf.

Enables fixing of
[#5163](#5163)
and
[#5178](#5178).

include/dxc/Support/Global.h

@@ -54,6 +54,10 @@ void DxcClearThreadMalloc() throw();
 // Used to retrieve the current invocation's allocator or perform an alloc/free/realloc.
 IMalloc *DxcGetThreadMallocNoRef() throw();
 
+// Common implementation of operators new and delete
+void *DxcNew(std::size_t size) throw();
+void DxcDelete(void* ptr) throw();
+
 class DxcThreadMalloc {
 public:
   explicit DxcThreadMalloc(IMalloc *pMallocOrNull) throw();

lib/DxcSupport/dxcmem.cpp

@@ -98,3 +98,31 @@ DxcThreadMalloc::DxcThreadMalloc(IMalloc *pMallocOrNull) throw() {
 DxcThreadMalloc::~DxcThreadMalloc() {
     DxcSwapThreadMalloc(pPrior, nullptr);
 }
+
+void* DxcNew(std::size_t size) throw() {
+  void *ptr;
+  IMalloc* iMalloc = DxcGetThreadMallocNoRef();
+  if (iMalloc != nullptr) {
+    ptr = iMalloc->Alloc(size);
+  } else {
+    // DxcGetThreadMallocNoRef() returning null means the operator is called before DllMain
+    // where the g_pDefaultMalloc is initialized, for example from CRT libraries when
+    // static linking is enabled. In that case fallback to the standard allocator
+    // and use CoTaskMemAlloc directly instead of CoGetMalloc, Alloc & Release for better perf.
+    ptr = CoTaskMemAlloc(size);
+  }
+  return ptr;
+}
+
+void DxcDelete(void *ptr) throw() {
+  IMalloc* iMalloc = DxcGetThreadMallocNoRef();
+  if (iMalloc != nullptr) {
+    iMalloc->Free(ptr);
+  } else {
+    // DxcGetThreadMallocNoRef() returning null means the operator is called before DllMain
+    // where the g_pDefaultMalloc is initialized, for example from CRT libraries when
+    // static linking is enabled. In that case fallback to the standard allocator
+    // and use CoTaskMemFree directly instead of CoGetMalloc, Free & Release for better perf.
+    CoTaskMemFree(ptr);
+  }
+}

tools/clang/tools/dxcompiler/DXCompiler.cpp

@@ -31,21 +31,24 @@ HRESULT SetupRegistryPassForPIX();
 
 #if defined(LLVM_ON_WIN32) && !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)
 // operator new and friends.
-void *  __CRTDECL operator new(std::size_t size) noexcept(false) {
-  void * ptr = DxcGetThreadMallocNoRef()->Alloc(size);
+void*  __CRTDECL operator new(std::size_t size) noexcept(false) {
+  void *ptr = DxcNew(size);
   if (ptr == nullptr)
     throw std::bad_alloc();
   return ptr;
 }
+
 void * __CRTDECL operator new(std::size_t size,
   const std::nothrow_t &nothrow_value) throw() {
-  return DxcGetThreadMallocNoRef()->Alloc(size);
+  return DxcNew(size);
 }
+
 void  __CRTDECL operator delete (void* ptr) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
+
 void  __CRTDECL operator delete (void* ptr, const std::nothrow_t& nothrow_constant) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
 #endif
 

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

@@ -22,20 +22,20 @@ using namespace hlsl;
 
 // operator new and friends.
 void *  __CRTDECL operator new(std::size_t size) noexcept(false) {
-  void * ptr = DxcGetThreadMallocNoRef()->Alloc(size);
+  void *ptr = DxcNew(size);
   if (ptr == nullptr)
     throw std::bad_alloc();
   return ptr;
 }
 void *  __CRTDECL operator new(std::size_t size,
   const std::nothrow_t &nothrow_value) throw() {
-  return DxcGetThreadMallocNoRef()->Alloc(size);
+  return DxcNew(size);
 }
 void  __CRTDECL operator delete (void* ptr) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
 void  __CRTDECL operator delete (void* ptr, const std::nothrow_t& nothrow_constant) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
 // Finish of new delete.
 

tools/clang/tools/dxrfallbackcompiler/DXCompiler.cpp

@@ -26,21 +26,21 @@ namespace hlsl { HRESULT SetupRegistryPassForHLSL(); }
 
 #if !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)
 // operator new and friends.
-void *  __CRTDECL operator new(std::size_t size) noexcept(false) {
-  void * ptr = DxcGetThreadMallocNoRef()->Alloc(size);
+void * __CRTDECL operator new(std::size_t size) noexcept(false) {
+  void *ptr = DxcNew(size);
   if (ptr == nullptr)
     throw std::bad_alloc();
   return ptr;
 }
 void * __CRTDECL operator new(std::size_t size,
   const std::nothrow_t &nothrow_value) throw() {
-  return DxcGetThreadMallocNoRef()->Alloc(size);
+  return DxcNew(size);
 }
 void  __CRTDECL operator delete (void* ptr) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
 void  __CRTDECL operator delete (void* ptr, const std::nothrow_t& nothrow_constant) throw() {
-  DxcGetThreadMallocNoRef()->Free(ptr);
+  DxcDelete(ptr);
 }
 #endif