Add macro to disable usage of allocator overrides (#5296)

Add macro DXC_DISABLE_ALLOCATOR_OVERRIDES which, if defined, will use
the default CRT allocator for new/delete and malloc/realloc/free.

We need this on Chrome for our GN build of DXC, because our component
builds use libc++as a shared library, and we end up with mismatched
new/delete calls between libc++ and these custom overrides.

---------

Co-authored-by: Chris Bieneman <[email protected]>

CMakeLists.txt

@@ -101,6 +101,9 @@ if ( HLSL_SUPPORT_QUERY_GIT_COMMIT_INFO )
   add_definitions(-DSUPPORT_QUERY_GIT_COMMIT_INFO)
 endif()
 
+option(DXC_DISABLE_ALLOCATOR_OVERRIDES "Disable usage of allocator overrides" OFF)
+mark_as_advanced(DXC_DISABLE_ALLOCATOR_OVERRIDES)
+
 # adjust link option to enable debugging from kernel mode; not compatible with incremental linking
 if(NOT CMAKE_VERSION VERSION_LESS "3.13" AND WIN32 AND NOT CMAKE_C_COMPILER_ARCHITECTURE_ID STREQUAL "ARM64EC")
   add_link_options(/DEBUGTYPE:CV,FIXUP,PDATA /INCREMENTAL:NO)

include/dxc/CMakeLists.txt

@@ -10,6 +10,11 @@ configure_file(
   ${LLVM_INCLUDE_DIR}/dxc/Test/TestConfig.h
   )
 
+configure_file(
+  ${CMAKE_CURRENT_SOURCE_DIR}/config.h.cmake
+  ${LLVM_INCLUDE_DIR}/dxc/config.h
+  )
+
 add_subdirectory(DXIL)
 add_subdirectory(DxilContainer)
 add_subdirectory(HLSL)
@@ -17,6 +22,7 @@ add_subdirectory(Support)
 add_subdirectory(Tracing)
 
 set(DXC_PUBLIC_HEADERS
+    ${CMAKE_CURRENT_BINARY_DIR}/config.h
     ${CMAKE_CURRENT_SOURCE_DIR}/dxcapi.h
     ${CMAKE_CURRENT_SOURCE_DIR}/dxcerrors.h
     ${CMAKE_CURRENT_SOURCE_DIR}/dxcisense.h)

include/dxc/Support/WinFunctions.h

@@ -31,7 +31,6 @@ int _wcsicmp(const wchar_t *str1, const wchar_t *str2);
 int _wcsnicmp(const wchar_t *str1, const wchar_t *str2, size_t n);
 int wsprintf(wchar_t *wcs, const wchar_t *fmt, ...);
 unsigned char _BitScanForward(unsigned long * Index, unsigned long Mask);
-HRESULT CoGetMalloc(DWORD dwMemContext, IMalloc **ppMalloc);
 
 HANDLE CreateFile2(_In_ LPCWSTR lpFileName, _In_ DWORD dwDesiredAccess,
                    _In_ DWORD dwShareMode, _In_ DWORD dwCreationDisposition,

include/dxc/Support/WinIncludes.h

@@ -47,6 +47,8 @@
 #include <intsafe.h>
 #include <ObjIdl.h>
 
+#include "dxc/config.h"
+
 // Support older atlbase.h if needed
 #ifndef _ATL_DECLSPEC_ALLOCATOR
 #define _ATL_DECLSPEC_ALLOCATOR
@@ -112,3 +114,26 @@ inline ENUMTYPE &operator ^= (ENUMTYPE &a, ENUMTYPE b) { return (ENUMTYPE &)(((_
 #endif
 
 #endif // _MSC_VER
+
+/// DxcCoGetMalloc
+#if defined(_WIN32) && !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)
+
+#define DxcCoGetMalloc CoGetMalloc
+
+#else // defined(_WIN32) && !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)
+
+#ifndef _WIN32
+CROSS_PLATFORM_UUIDOF(IMalloc, "00000002-0000-0000-C000-000000000046")
+struct IMalloc : public IUnknown {
+  virtual void *Alloc(SIZE_T size) = 0;
+  virtual void *Realloc(void *ptr, SIZE_T size) = 0;
+  virtual void Free(void *ptr) = 0;
+  virtual SIZE_T GetSize(void *pv) = 0;
+  virtual int DidAlloc(void *pv) = 0;
+  virtual void HeapMinimize(void) = 0;
+};
+#endif
+
+HRESULT DxcCoGetMalloc(DWORD dwMemContext, IMalloc **ppMalloc);
+
+#endif // defined(_WIN32) && !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)

include/dxc/Support/microcom.h

@@ -12,8 +12,9 @@
 #ifndef __DXC_MICROCOM__
 #define __DXC_MICROCOM__
 
-#include <atomic>
+#include "dxc/Support/WinIncludes.h"
 #include "llvm/Support/Atomic.h"
+#include <atomic>
 
 template <typename TIface>
 class CComInterfaceArray {

include/dxc/WinAdapter.h

@@ -642,16 +642,6 @@ struct IUnknown {
 CROSS_PLATFORM_UUIDOF(INoMarshal, "ECC8691B-C1DB-4DC0-855E-65F6C551AF49")
 struct INoMarshal : public IUnknown {};
 
-CROSS_PLATFORM_UUIDOF(IMalloc, "00000002-0000-0000-C000-000000000046")
-struct IMalloc : public IUnknown {
-  virtual void *Alloc(size_t size) = 0;
-  virtual void *Realloc(void *ptr, size_t size) = 0;
-  virtual void Free(void *ptr) = 0;
-  virtual size_t GetSize(void *pv) = 0;
-  virtual int DidAlloc(void *pv) = 0;
-  virtual void HeapMinimize(void) = 0;
-};
-
 CROSS_PLATFORM_UUIDOF(ISequentialStream, "0C733A30-2A1C-11CE-ADE5-00AA0044773D")
 struct ISequentialStream : public IUnknown {
   virtual HRESULT Read(void *pv, ULONG cb, ULONG *pcbRead) = 0;

include/dxc/config.h.cmake

@@ -0,0 +1,2 @@
+/* Disable overriding memory allocators. */
+#cmakedefine DXC_DISABLE_ALLOCATOR_OVERRIDES

lib/DxcSupport/CMakeLists.txt

@@ -8,6 +8,7 @@ add_llvm_library(LLVMDxcSupport
   HLSLOptions.cpp
   Unicode.cpp
   WinAdapter.cpp
+  WinIncludes.cpp
   WinFunctions.cpp
   )
 
@@ -17,4 +18,6 @@ ${LLVM_MAIN_SRC_DIR}/lib/DxcSupport/SharedLibAffix.inc
 ${LLVM_INCLUDE_DIR}/dxc/Support/SharedLibAffix.h
 )
 
-add_dependencies(LLVMDxcSupport TablegenHLSLOptions)
+target_link_libraries(LLVMDxcSupport PUBLIC LLVMSupport)
+
+add_dependencies(LLVMDxcSupport LLVMSupport TablegenHLSLOptions)

lib/DxcSupport/WinFunctions.cpp

@@ -139,32 +139,6 @@ unsigned char _BitScanForward(unsigned long * Index, unsigned long Mask) {
   return 1;
 }
 
-struct CoMalloc : public IMalloc {
-  CoMalloc() : m_dwRef(0) {};
-
-  DXC_MICROCOM_ADDREF_RELEASE_IMPL(m_dwRef)
-  STDMETHODIMP QueryInterface(REFIID riid, void **ppvObject) override {
-    assert(false && "QueryInterface not implemented for CoMalloc.");
-    return E_NOINTERFACE;
-  }
-
-  void *STDMETHODCALLTYPE Alloc(size_t size) override { return malloc(size); }
-  void *STDMETHODCALLTYPE Realloc(void *ptr, size_t size) override { return realloc(ptr, size); }
-  void STDMETHODCALLTYPE Free(void *ptr) override { free(ptr); }
-  size_t STDMETHODCALLTYPE GetSize(void *pv) override { return -1; }
-  int STDMETHODCALLTYPE DidAlloc(void *pv) override { return -1; }
-  void STDMETHODCALLTYPE HeapMinimize(void) override {}
-
-private:
-  DXC_MICROCOM_REF_FIELD(m_dwRef)
-};
-
-HRESULT CoGetMalloc(DWORD dwMemContext, IMalloc **ppMalloc) {
-  *ppMalloc = new CoMalloc;
-  (*ppMalloc)->AddRef();
-  return S_OK;
-}
-
 HANDLE CreateFile2(_In_ LPCWSTR lpFileName, _In_ DWORD dwDesiredAccess,
                    _In_ DWORD dwShareMode, _In_ DWORD dwCreationDisposition,
                    _In_opt_ void *pCreateExParams) {

lib/DxcSupport/WinIncludes.cpp

@@ -0,0 +1,46 @@
+//===- WinIncludes.cpp -----------------------------------------*- C++ -*-===//
+///////////////////////////////////////////////////////////////////////////////
+//                                                                           //
+// WinIncludes.cpp                                                           //
+// Copyright (C) Microsoft Corporation. All rights reserved.                 //
+// This file is distributed under the University of Illinois Open Source     //
+// License. See LICENSE.TXT for details.                                     //
+//                                                                           //
+///////////////////////////////////////////////////////////////////////////////
+
+#include "dxc/Support/WinIncludes.h"
+
+#include "assert.h"
+#include "dxc/Support/microcom.h"
+
+#if defined(_WIN32) && !defined(DXC_DISABLE_ALLOCATOR_OVERRIDES)
+// CoGetMalloc from combaseapi.h is used
+#else
+struct DxcCoMalloc : public IMalloc {
+  DxcCoMalloc() : m_dwRef(0){};
+
+  DXC_MICROCOM_ADDREF_RELEASE_IMPL(m_dwRef)
+  STDMETHODIMP QueryInterface(REFIID riid, void **ppvObject) override {
+    assert(false && "QueryInterface not implemented for DxcCoMalloc.");
+    return E_NOINTERFACE;
+  }
+
+  void *STDMETHODCALLTYPE Alloc(SIZE_T size) override { return malloc(size); }
+  void *STDMETHODCALLTYPE Realloc(void *ptr, SIZE_T size) override {
+    return realloc(ptr, size);
+  }
+  void STDMETHODCALLTYPE Free(void *ptr) override { free(ptr); }
+  SIZE_T STDMETHODCALLTYPE GetSize(void *pv) override { return -1; }
+  int STDMETHODCALLTYPE DidAlloc(void *pv) override { return -1; }
+  void STDMETHODCALLTYPE HeapMinimize(void) override {}
+
+private:
+  DXC_MICROCOM_REF_FIELD(m_dwRef)
+};
+
+HRESULT DxcCoGetMalloc(DWORD dwMemContext, IMalloc **ppMalloc) {
+  *ppMalloc = new DxcCoMalloc;
+  (*ppMalloc)->AddRef();
+  return S_OK;
+}
+#endif

lib/DxcSupport/dxcmem.cpp

@@ -31,7 +31,7 @@ HRESULT DxcInitThreadMalloc() throw() {
   }
   else {
     // We capture the default malloc early to avoid potential failures later on.
-    HRESULT hrMalloc = CoGetMalloc(1, &g_pDefaultMalloc);
+    HRESULT hrMalloc = DxcCoGetMalloc(1, &g_pDefaultMalloc);
     if (FAILED(hrMalloc)) return hrMalloc;
   }
   DXASSERT(g_ThreadMallocTls == nullptr, "else InitThreadMalloc already called");

projects/dxilconv/tools/dxilconv/dxilconv.cpp

@@ -18,10 +18,11 @@
 
 #define DXC_API_IMPORT
 
+#include "dxc/config.h"
 #include "dxc/dxcisense.h"
 #include "dxc/dxctools.h"
 #include "dxcetw.h"
-#include "Tracing/DxcRuntimeEtw.h"
+
 #include "DxbcConverter.h"
 
 // Defined in DxbcConverter.lib (projects/dxilconv/lib/DxbcConverter/DxbcConverter.cpp)
@@ -67,6 +68,12 @@ DxcCreateInstance2(_In_ IMalloc *pMalloc,
     if (ppv == nullptr) {
         return E_POINTER;
     }
+#ifdef DXC_DISABLE_ALLOCATOR_OVERRIDES
+    if (pMalloc != DxcGetThreadMallocNoRef()) {
+      return E_INVALIDARG;
+    }
+#endif // DXC_DISABLE_ALLOCATOR_OVERRIDES
+
     HRESULT hr = S_OK;
     DxcEtw_DXCompilerCreateInstance_Start();
     DxcThreadMalloc TM(pMalloc);

tools/clang/tools/dxclib/dxc.cpp

@@ -555,10 +555,8 @@ void DxcContext::ExtractRootSignature(IDxcBlob *pBlob, IDxcBlob **ppResult) {
   hlsl::DxilContainerHeader newHeader;
   uint32_t containerSize = hlsl::GetDxilContainerSizeFromParts(1, pPartHeader->PartSize);
   hlsl::InitDxilContainer(&newHeader, 1, containerSize); 
-  CComPtr<IMalloc> pMalloc;
   CComPtr<hlsl::AbstractMemoryStream> pMemoryStream;
-  IFT(CoGetMalloc(1, &pMalloc));
-  IFT(hlsl::CreateMemoryStream(pMalloc, &pMemoryStream));
+  IFT(hlsl::CreateMemoryStream(DxcGetThreadMallocNoRef(), &pMemoryStream));
   ULONG cbWritten;
 
   // Write Container Header

tools/clang/tools/dxcompiler/DXCompiler.cpp

@@ -9,11 +9,13 @@
 //                                                                           //
 ///////////////////////////////////////////////////////////////////////////////
 
-#include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/FileSystem.h"
-#include "dxc/Support/Global.h"
 #include "dxc/Support/WinIncludes.h"
+
+#include "dxc/Support/Global.h"
 #include "dxc/Support/HLSLOptions.h"
+#include "dxc/config.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/ManagedStatic.h"
 #ifdef LLVM_ON_WIN32
 #include "dxcetw.h"
 #endif
@@ -27,7 +29,7 @@ HRESULT SetupRegistryPassForPIX();
 // C++ exception specification ignored except to indicate a function is not __declspec(nothrow)
 #pragma warning( disable : 4290 )
 
-#ifdef LLVM_ON_WIN32
+#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);

tools/clang/tools/dxcompiler/dxcapi.cpp

@@ -17,9 +17,10 @@
 #define DXC_API_IMPORT __attribute__ ((visibility ("default")))
 #endif
 
+#include "dxc/Support/Global.h"
+#include "dxc/config.h"
 #include "dxc/dxcisense.h"
 #include "dxc/dxctools.h"
-#include "dxc/Support/Global.h"
 #ifdef _WIN32
 #include "dxcetw.h"
 #endif
@@ -158,6 +159,11 @@ DxcCreateInstance2(
   if (ppv == nullptr) {
     return E_POINTER;
   }
+#ifdef DXC_DISABLE_ALLOCATOR_OVERRIDES
+  if (pMalloc != DxcGetThreadMallocNoRef()) {
+    return E_INVALIDARG;
+  }
+#endif // DXC_DISABLE_ALLOCATOR_OVERRIDES
 
   HRESULT hr = S_OK;
   DxcEtw_DXCompilerCreateInstance_Start();

tools/clang/tools/dxcompiler/dxclinker.cpp

@@ -128,11 +128,8 @@ DxcLinker::RegisterLibrary(_In_opt_ LPCWSTR pLibName, // Name of the library.
   try {
     std::unique_ptr<llvm::Module> pModule, pDebugModule;
 
-    CComPtr<IMalloc> pMalloc;
     CComPtr<AbstractMemoryStream> pDiagStream;
-
-    IFT(CoGetMalloc(1, &pMalloc));
-    IFT(CreateMemoryStream(pMalloc, &pDiagStream));
+    IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pDiagStream));
 
     raw_stream_ostream DiagStream(pDiagStream);
 
@@ -180,12 +177,10 @@ HRESULT STDMETHODCALLTYPE DxcLinker::Link(
 
   HRESULT hr = S_OK;
   try {
-    CComPtr<IMalloc> pMalloc;
     CComPtr<IDxcBlob> pOutputBlob;
     CComPtr<AbstractMemoryStream> pDiagStream;
 
-    IFT(CoGetMalloc(1, &pMalloc));
-    IFT(CreateMemoryStream(pMalloc, &pOutputStream));
+    IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pOutputStream));
 
     // Read and validate options.
     int argCountInt;
@@ -207,7 +202,7 @@ HRESULT STDMETHODCALLTYPE DxcLinker::Link(
 
     std::string warnings;
     //llvm::raw_string_ostream w(warnings);
-    IFT(CreateMemoryStream(pMalloc, &pDiagStream));
+    IFT(CreateMemoryStream(DxcGetThreadMallocNoRef(), &pDiagStream));
     raw_stream_ostream DiagStream(pDiagStream);
     llvm::DiagnosticPrinterRawOStream DiagPrinter(DiagStream);
     PrintDiagnosticContext DiagContext(DiagPrinter);
@@ -273,9 +268,8 @@ HRESULT STDMETHODCALLTYPE DxcLinker::Link(
         // Validation.
         HRESULT valHR = S_OK;
         dxcutil::AssembleInputs inputs(
-          std::move(pM), pOutputBlob, pMalloc, SerializeFlags,
-          pOutputStream,
-          opts.DebugFile, &Diag);
+            std::move(pM), pOutputBlob, DxcGetThreadMallocNoRef(),
+            SerializeFlags, pOutputStream, opts.DebugFile, &Diag);
         if (needsValidation) {
           valHR = dxcutil::ValidateAndAssembleToContainer(inputs);
         } else {

tools/clang/tools/dxrfallbackcompiler/DXCompiler.cpp

@@ -9,19 +9,22 @@
 //                                                                           //
 ///////////////////////////////////////////////////////////////////////////////
 
-#include "llvm/Support/ManagedStatic.h"
-#include "llvm/Support/FileSystem.h"
-#include "dxc/Support/Global.h"
 #include "dxc/Support/WinIncludes.h"
+
+#include "dxc/Support/Global.h"
 #include "dxc/Support/HLSLOptions.h"
+#include "dxc/config.h"
 #include "dxcetw.h"
 #include "dxillib.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/ManagedStatic.h"
 
 namespace hlsl { HRESULT SetupRegistryPassForHLSL(); }
 
 // C++ exception specification ignored except to indicate a function is not __declspec(nothrow)
 #pragma warning( disable : 4290 )
 
+#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);
@@ -39,6 +42,7 @@ void  __CRTDECL operator delete (void* ptr) throw() {
 void  __CRTDECL operator delete (void* ptr, const std::nothrow_t& nothrow_constant) throw() {
   DxcGetThreadMallocNoRef()->Free(ptr);
 }
+#endif
 
 static HRESULT InitMaybeFail() throw() {
   HRESULT hr;

tools/clang/unittests/HLSL/ValidationTest.cpp

@@ -15,11 +15,12 @@
 #include <string>
 #include <algorithm>
 
-#include "llvm/ADT/StringRef.h"
-#include "llvm/Support/Regex.h"
-#include "llvm/ADT/ArrayRef.h"
 #include "dxc/DxilContainer/DxilContainer.h"
 #include "dxc/DxilContainer/DxilContainerAssembler.h"
+#include "dxc/Support/WinIncludes.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Regex.h"
 
 #ifdef _WIN32
 #include <atlbase.h>
@@ -600,7 +601,7 @@ public:
 
     // Write the container
     CComPtr<IMalloc> pMalloc;
-    VERIFY_SUCCEEDED(CoGetMalloc(1, &pMalloc));
+    VERIFY_SUCCEEDED(DxcCoGetMalloc(1, &pMalloc));
     CComPtr<AbstractMemoryStream> pOutputStream;
     VERIFY_SUCCEEDED(CreateMemoryStream(pMalloc, &pOutputStream));
     pOutputStream->Reserve(pContainerWriter->size());

tools/clang/unittests/HLSLExec/ExecutionTest.cpp

@@ -165,7 +165,6 @@ static void SavePixelsToFile(LPCVOID pPixels, DXGI_FORMAT format, UINT32 m_width
   CComPtr<IWICBitmapEncoder> pEncoder;
   CComPtr<IWICBitmapFrameEncode> pFrameEncode;
   CComPtr<IStream> pStream;
-  CComPtr<IMalloc> pMalloc;
 
   struct PF {
     DXGI_FORMAT Format;
@@ -182,7 +181,6 @@ static void SavePixelsToFile(LPCVOID pPixels, DXGI_FORMAT format, UINT32 m_width
 
   VERIFY_SUCCEEDED(ctx.Init());
   VERIFY_SUCCEEDED(CoCreateInstance(CLSID_WICImagingFactory, NULL, CLSCTX_INPROC_SERVER, IID_IWICImagingFactory, (LPVOID*)&pFactory));
-  VERIFY_SUCCEEDED(CoGetMalloc(1, &pMalloc));
   VERIFY_ARE_NOT_EQUAL(pFormat, Vals + _countof(Vals));
   VERIFY_SUCCEEDED(pFactory->CreateBitmapFromMemory(m_width, m_height, pFormat->PixelFormat, m_width * pFormat->PixelSize, m_width * m_height * pFormat->PixelSize, (BYTE *)pPixels, &pBitmap));
   VERIFY_SUCCEEDED(pFactory->CreateEncoder(GUID_ContainerFormatBmp, nullptr, &pEncoder));

tools/clang/unittests/dxc_batch/dxc_batch.cpp

@@ -443,7 +443,7 @@ void DxcContext::ExtractRootSignature(IDxcBlob *pBlob, IDxcBlob **ppResult) {
   hlsl::InitDxilContainer(&newHeader, 1, containerSize);
   CComPtr<IMalloc> pMalloc;
   CComPtr<hlsl::AbstractMemoryStream> pMemoryStream;
-  IFT(CoGetMalloc(1, &pMalloc));
+  IFT(DxcCoGetMalloc(1, &pMalloc));
   IFT(hlsl::CreateMemoryStream(pMalloc, &pMemoryStream));
   ULONG cbWritten;