Add mimalloc to improve multithreaded performance (KhronosGroup#6188)

Co-authored-by: Zackery Mason-Blaug <[email protected]>

Author: dneto0

.gitignore

@@ -11,6 +11,7 @@ compile_commands.json
 /external/effcee
 /external/re2
 /external/protobuf
+/external/mimalloc
 /out
 /TAGS
 /third_party/llvm-build/

BUILD.bazel

@@ -131,6 +131,8 @@ cc_library(
         "source/*.cpp",
         "source/util/*.cpp",
         "source/val/*.cpp",
+    ], exclude = [
+        "source/mimalloc.cpp"
     ]) + [
         ":build_version_inc",
         ":gen_compressed_tables",

DEPS

@@ -15,6 +15,8 @@ vars = {
   're2_revision': 'c84a140c93352cdabbfb547c531be34515b12228',
 
   'spirv_headers_revision': '2a611a970fdbc41ac2e3e328802aed9985352dca',
+
+  'mimalloc_revision': '51c09e7b6a0ac5feeba998710f00c7dd7aa67bbf',
 }
 
 deps = {
@@ -36,5 +38,8 @@ deps = {
   'external/spirv-headers':
       Var('github') +  '/KhronosGroup/SPIRV-Headers.git@' +
           Var('spirv_headers_revision'),
+
+  'external/mimalloc':
+      Var('github') + '/microsoft/mimalloc.git@' + Var('mimalloc_revision'),
 }
 

README.md

@@ -300,6 +300,7 @@ For some kinds of development, you may need the latest sources from the third-pa
     git clone https://github.com/google/effcee.git              spirv-tools/external/effcee
     git clone https://github.com/google/re2.git                 spirv-tools/external/re2
     git clone https://github.com/abseil/abseil-cpp.git          spirv-tools/external/abseil_cpp
+    git clone https://github.com/microsoft/mimalloc.git         spirv-tools/external/mimalloc
 
 #### Dependency on Effcee
 
@@ -312,6 +313,23 @@ Effcee itself depends on [RE2][re2], and RE2 depends on [Abseil][abseil-cpp].
   RE2 sources to appear in `external/re2`, and Abseil sources to appear in 
   `external/abseil_cpp`.
 
+#### Dependency on mimalloc
+
+SPIRV-Tools may be configured to use the [mimalloc][mimalloc] library to improve memory
+allocation performance. In order to avoid unexpectedly changing allocation behavior of
+applications that link SPIRV-Tools libraries statically, this option has no effect on
+the static libraries.
+
+In the CMake build, usage of mimalloc is controlled by the `SPIRV_TOOLS_USE_MIMALLOC`
+option. This variable defaults on `ON` when building for Windows and `OFF` when building
+for other platforms. Enabling this option on non-Windows platforms is supported and is
+expected to work normally, but this has not been tested as thoroughly and extensively as
+the Windows version. In the future, the `SPIRV_TOOLS_USE_MIMALLOC` option may default to
+`ON` for non-Windows platforms as well.
+
+*Note*: mimalloc is currently only supported when building with CMake. When using Bazel,
+mimalloc is not used.
+
 ### Source code organization
 
 * `example`: demo code of using SPIRV-Tools APIs
@@ -325,6 +343,7 @@ Effcee itself depends on [RE2][re2], and RE2 depends on [Abseil][abseil-cpp].
 * `external/abseil_cpp`: Location of [Abseil][abseil-cpp] sources, if Abseil is
    not already configured by an enclosing project.
   (The RE2 project already requires Abseil.)
+* `external/mimalloc`: Intended location for [mimalloc][mimalloc] sources, not provided
 * `include/`: API clients should add this directory to the include search path
 * `external/spirv-headers`: Intended location for
   [SPIR-V headers][spirv-headers], not provided
@@ -801,6 +820,7 @@ limitations under the License.
 [effcee]: https://github.com/google/effcee
 [re2]: https://github.com/google/re2
 [abseil-cpp]: https://github.com/abseil/abseil-cpp
+[mimalloc]: https://github.com/microsoft/mimalloc
 [CMake]: https://cmake.org/
 [cpp-style-guide]: https://google.github.io/styleguide/cppguide.html
 [clang-sanitizers]: http://clang.llvm.org/docs/UsersManual.html#controlling-code-generation

external/CMakeLists.txt

@@ -26,6 +26,46 @@ function(pop_variable var)
     set(${var} ${val} PARENT_SCOPE)
 endfunction()
 
+if (DEFINED mimalloc_SOURCE_DIR)
+  # This allows flexible position of the mimalloc repo.
+  set(MIMALLOC_DIR ${mimalloc_SOURCE_DIR})
+else()
+  if (IS_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mimalloc)
+    set(MIMALLOC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/mimalloc)
+  endif()
+endif()
+
+# Used on Windows by default, but allow opt-in on other platforms
+if (WIN32)
+  set(SPIRV_TOOLS_USE_MIMALLOC_DEFAULT_VALUE ON)
+else()
+  set(SPIRV_TOOLS_USE_MIMALLOC_DEFAULT_VALUE OFF)
+endif()
+
+# To avoid unexpected side effects on users of the static library, mimalloc
+# may only be used when building executables and shared libraries.
+include(CMakeDependentOption)
+cmake_dependent_option(SPIRV_TOOLS_USE_MIMALLOC
+  "Executables and shared libraries use mimalloc instead of the default allocator"
+  ${SPIRV_TOOLS_USE_MIMALLOC_DEFAULT_VALUE} "MIMALLOC_DIR" OFF)
+
+if (SPIRV_TOOLS_USE_MIMALLOC)
+  if (NOT WIN32)
+    push_variable(MI_OVERRIDE 0)
+  endif()
+  push_variable(MI_BUILD_TESTS 0)
+
+  add_subdirectory(${MIMALLOC_DIR} ${CMAKE_BINARY_DIR}/external/mimalloc EXCLUDE_FROM_ALL)
+  if (${CMAKE_CXX_COMPILER_ID} MATCHES Clang)
+    target_compile_options(mimalloc-static PRIVATE -Wno-int-conversion)
+  endif()
+
+  if (NOT WIN32)
+    pop_variable(MI_OVERRIDE)
+  endif()
+  pop_variable(MI_BUILD_TESTS)
+endif()
+
 if (DEFINED SPIRV-Headers_SOURCE_DIR)
   # This allows flexible position of the SPIRV-Headers repo.
   set(SPIRV_HEADER_DIR ${SPIRV-Headers_SOURCE_DIR})

source/CMakeLists.txt

@@ -297,6 +297,12 @@ if (${SPIRV_TIMER_ENABLED})
     ${CMAKE_CURRENT_SOURCE_DIR}/util/timer.cpp)
 endif()
 
+if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+  set(SPIRV_SOURCES
+      ${SPIRV_SOURCES}
+      ${CMAKE_CURRENT_SOURCE_DIR}/mimalloc.cpp)
+endif()
+
 # The software_version.cpp file includes build-version.inc.
 # Rebuild the software_version.cpp object file if it is older than
 # build-version.inc or whenever build-version.inc itself is out of
@@ -333,6 +339,9 @@ endfunction()
 # Always build ${SPIRV_TOOLS}-shared. This is expected distro packages, and
 # unlike the other SPIRV_TOOLS target, defaults to hidden symbol visibility.
 add_library(${SPIRV_TOOLS}-shared SHARED ${SPIRV_SOURCES})
+if (SPIRV_TOOLS_USE_MIMALLOC)
+  target_link_libraries(${SPIRV_TOOLS}-shared PRIVATE mimalloc-static)
+endif()
 spirv_tools_default_target_options(${SPIRV_TOOLS}-shared)
 set_target_properties(${SPIRV_TOOLS}-shared PROPERTIES CXX_VISIBILITY_PRESET hidden)
 target_compile_definitions(${SPIRV_TOOLS}-shared
@@ -357,6 +366,9 @@ if(SPIRV_TOOLS_BUILD_STATIC)
   set(SPIRV_TOOLS_TARGETS ${SPIRV_TOOLS}-static ${SPIRV_TOOLS}-shared)
 else()
   add_library(${SPIRV_TOOLS} ${SPIRV_TOOLS_LIBRARY_TYPE} ${SPIRV_SOURCES})
+  if (SPIRV_TOOLS_USE_MIMALLOC)
+    target_link_libraries(${SPIRV_TOOLS} PRIVATE mimalloc-static)
+  endif()
   spirv_tools_default_target_options(${SPIRV_TOOLS})
   set(SPIRV_TOOLS_TARGETS ${SPIRV_TOOLS} ${SPIRV_TOOLS}-shared)
 endif()
@@ -365,12 +377,15 @@ if("${CMAKE_SYSTEM_NAME}" STREQUAL "Linux")
   find_library(LIBRT rt)
   if(LIBRT)
     foreach(target ${SPIRV_TOOLS_TARGETS})
-      target_link_libraries(${target} rt)
+      target_link_libraries(${target} PUBLIC rt)
     endforeach()
   endif()
 endif()
 
 if(ENABLE_SPIRV_TOOLS_INSTALL)
+  if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+    list(APPEND SPIRV_TOOLS_TARGETS mimalloc-static)
+  endif()
   install(TARGETS ${SPIRV_TOOLS_TARGETS} EXPORT ${SPIRV_TOOLS}Targets)
   export(EXPORT ${SPIRV_TOOLS}Targets FILE ${SPIRV_TOOLS}Target.cmake)
 

source/diff/CMakeLists.txt

@@ -18,8 +18,16 @@ set(SPIRV_TOOLS_DIFF_SOURCES
   diff.cpp
 )
 
+if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+  list(APPEND SPIRV_TOOLS_DIFF_SOURCES ${spirv-tools_SOURCE_DIR}/source/mimalloc.cpp)
+endif()
+
 add_library(SPIRV-Tools-diff ${SPIRV_TOOLS_LIBRARY_TYPE} ${SPIRV_TOOLS_DIFF_SOURCES})
 
+if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+  target_link_libraries(SPIRV-Tools-diff PRIVATE mimalloc-static)
+endif()
+
 spvtools_default_compile_options(SPIRV-Tools-diff)
 target_include_directories(SPIRV-Tools-diff
   PUBLIC
@@ -39,7 +47,13 @@ set_property(TARGET SPIRV-Tools-diff PROPERTY FOLDER "SPIRV-Tools libraries")
 spvtools_check_symbol_exports(SPIRV-Tools-diff)
 
 if(ENABLE_SPIRV_TOOLS_INSTALL)
-  install(TARGETS SPIRV-Tools-diff EXPORT SPIRV-Tools-diffTargets)
+  set(SPIRV-Tools-diff-InstallTargets SPIRV-Tools-diff)
+
+  if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+    list(APPEND SPIRV-Tools-diff-InstallTargets mimalloc-static)
+  endif()
+
+  install(TARGETS ${SPIRV-Tools-diff-InstallTargets} EXPORT SPIRV-Tools-diffTargets)
   export(EXPORT SPIRV-Tools-diffTargets FILE SPIRV-Tools-diffTargets.cmake)
 
   spvtools_config_package_dir(SPIRV-Tools-diff PACKAGE_DIR)

source/fuzz/CMakeLists.txt

@@ -438,8 +438,16 @@ if(SPIRV_BUILD_FUZZER)
 
   spvtools_pch(SPIRV_TOOLS_FUZZ_SOURCES pch_source_fuzz)
 
+  if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+    list(APPEND SPIRV_TOOLS_DIFF_SOURCES ${spirv-tools_SOURCE_DIR}/source/mimalloc.cpp)
+  endif()
+
   add_library(SPIRV-Tools-fuzz ${SPIRV_TOOLS_FUZZ_SOURCES})
 
+  if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+    target_link_libraries(SPIRV-Tools-fuzz PRIVATE mimalloc-static)
+  endif()
+
   spvtools_default_compile_options(SPIRV-Tools-fuzz)
 
   # Compilation of the auto-generated protobuf source file will yield warnings,
@@ -470,7 +478,13 @@ if(SPIRV_BUILD_FUZZER)
   spvtools_check_symbol_exports(SPIRV-Tools-fuzz)
 
   if(ENABLE_SPIRV_TOOLS_INSTALL)
-      install(TARGETS SPIRV-Tools-fuzz EXPORT SPIRV-Tools-fuzzTargets)
+      set(SPIRV-Tools-fuzz-InstallTargets SPIRV-Tools-fuzz)
+
+      if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+        list(APPEND SPIRV-Tools-fuzz-InstallTargets mimalloc-static)
+      endif()
+
+      install(TARGETS ${SPIRV-Tools-fuzz-InstallTargets} EXPORT SPIRV-Tools-fuzzTargets)
       export(EXPORT SPIRV-Tools-fuzzTargets FILE SPIRV-Tools-fuzzTarget.cmake)
 
       spvtools_config_package_dir(SPIRV-Tools-fuzz PACKAGE_DIR)

source/mimalloc.cpp

@@ -0,0 +1,15 @@
+// Copyright (c) 2025 The Khronos Group Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "mimalloc-new-delete.h"

source/opt/CMakeLists.txt

@@ -260,8 +260,16 @@ endif()
 
 spvtools_pch(SPIRV_TOOLS_OPT_SOURCES pch_source_opt)
 
+if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+  list(APPEND SPIRV_TOOLS_OPT_SOURCES ${spirv-tools_SOURCE_DIR}/source/mimalloc.cpp)
+endif()
+
 add_library(SPIRV-Tools-opt ${SPIRV_TOOLS_LIBRARY_TYPE} ${SPIRV_TOOLS_OPT_SOURCES})
 
+if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+  target_link_libraries(SPIRV-Tools-opt PRIVATE mimalloc-static)
+endif()
+
 spvtools_default_compile_options(SPIRV-Tools-opt)
 target_include_directories(SPIRV-Tools-opt
   PUBLIC
@@ -278,7 +286,13 @@ set_property(TARGET SPIRV-Tools-opt PROPERTY FOLDER "SPIRV-Tools libraries")
 spvtools_check_symbol_exports(SPIRV-Tools-opt)
 
 if(ENABLE_SPIRV_TOOLS_INSTALL)
-  install(TARGETS SPIRV-Tools-opt EXPORT SPIRV-Tools-optTargets)
+  set(SPIRV-Tools-opt-InstallTargets SPIRV-Tools-opt)
+
+  if (SPIRV_TOOLS_USE_MIMALLOC AND NOT SPIRV_TOOLS_BUILD_STATIC)
+    list(APPEND SPIRV-Tools-opt-InstallTargets mimalloc-static)
+  endif()
+
+  install(TARGETS ${SPIRV-Tools-opt-InstallTargets} EXPORT SPIRV-Tools-optTargets)
   export(EXPORT SPIRV-Tools-optTargets FILE SPIRV-Tools-optTargets.cmake)
 
   spvtools_config_package_dir(SPIRV-Tools-opt PACKAGE_DIR)