clingo-bootstrap: pgo, lto, allocator optimizations (#34926)

Add support for PGO and LTO for gcc, clang and apple-clang, and add a
patch to allow mimalloc as an allocator in operator new/delete, give
reduces clingo runtime by about 30%.

6 files changed, 238 insertions, 26 deletions

diff --git a/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc-pre-5.5.0.patch b/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc-pre-5.5.0.patch
new file mode 100644
index 0000000000..54491c8c1a
--- /dev/null
+++ b/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc-pre-5.5.0.patch
@@ -0,0 +1,39 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index f11e6e2..209970b 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -164,6 +164,7 @@ if (CLINGO_BUILD_WITH_LUA)
+         set_property(TARGET Lua::Lua PROPERTY INTERFACE_INCLUDE_DIRECTORIES "${LUA_INCLUDE_DIR}")
+     endif()
+ endif()
++find_package(mimalloc REQUIRED)
+ find_package(BISON "2.5")
+ find_package(RE2C "0.13")
+ if (PYCLINGO_USE_CFFI AND Python_Development_FOUND)
+diff --git a/libclingo/CMakeLists.txt b/libclingo/CMakeLists.txt
+index 83acc22..51d5762 100644
+--- a/libclingo/CMakeLists.txt
++++ b/libclingo/CMakeLists.txt
+@@ -50,7 +50,7 @@ else()
+ endif()
+ 
+ add_library(libclingo ${clingo_lib_type} ${header} ${source})
+-target_link_libraries(libclingo PRIVATE libgringo libclasp)
++target_link_libraries(libclingo PRIVATE mimalloc-static libgringo libclasp)
+ target_include_directories(libclingo
+     PUBLIC
+     "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
+diff --git a/libclingo/src/clingo_app.cc b/libclingo/src/clingo_app.cc
+index 3e4d14c..fcfc9ea 100644
+--- a/libclingo/src/clingo_app.cc
++++ b/libclingo/src/clingo_app.cc
+@@ -27,6 +27,9 @@
+ #include <clasp/parser.h>
+ #include <climits>
+ 
++#include <mimalloc-new-delete.h>
++
++
+ namespace Gringo {
+ 
+ // {{{ declaration of ClingoApp
diff --git a/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc.patch b/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc.patch
new file mode 100644
index 0000000000..b1de684591
--- /dev/null
+++ b/var/spack/repos/builtin/packages/clingo-bootstrap/mimalloc.patch
@@ -0,0 +1,39 @@
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 7fbe16bc..78539519 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -224,6 +224,7 @@ if (CLINGO_BUILD_WITH_LUA)
+         set_property(TARGET Lua::Lua PROPERTY INTERFACE_INCLUDE_DIRECTORIES "${LUA_INCLUDE_DIR}")
+     endif()
+ endif()
++find_package(mimalloc REQUIRED)
+ find_package(BISON "2.5")
+ find_package(RE2C "0.101")
+ if (Python_Development_FOUND)
+diff --git a/libclingo/CMakeLists.txt b/libclingo/CMakeLists.txt
+index 1d70ba56..de2f2766 100644
+--- a/libclingo/CMakeLists.txt
++++ b/libclingo/CMakeLists.txt
+@@ -51,7 +51,7 @@ endif()
+ 
+ add_library(libclingo ${clingo_lib_type})
+ target_sources(libclingo ${clingo_private_scope_} ${header} ${source})
+-target_link_libraries(libclingo ${clingo_private_scope_} libgringo libclasp)
++target_link_libraries(libclingo ${clingo_private_scope_} mimalloc-static libgringo libclasp)
+ target_include_directories(libclingo
+     ${clingo_public_scope_}
+     "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
+diff --git a/libclingo/src/clingo_app.cc b/libclingo/src/clingo_app.cc
+index 13980efa..3c3b404b 100644
+--- a/libclingo/src/clingo_app.cc
++++ b/libclingo/src/clingo_app.cc
+@@ -26,6 +26,9 @@
+ #include <clasp/parser.h>
+ #include <climits>
+ 
++#include <mimalloc-new-delete.h>
++
++
+ namespace Gringo {
+ 
+ // {{{ declaration of ClingoApp
diff --git a/var/spack/repos/builtin/packages/clingo-bootstrap/package.py b/var/spack/repos/builtin/packages/clingo-bootstrap/package.py
index 0e77dde84b..e3aed932a9 100644
--- a/var/spack/repos/builtin/packages/clingo-bootstrap/package.py
+++ b/var/spack/repos/builtin/packages/clingo-bootstrap/package.py
@@ -2,8 +2,15 @@
 # Spack Project Developers. See the top-level COPYRIGHT file for details.
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
+import glob
+import os
+
+import spack.compilers
+import spack.paths
+import spack.user_environment
 from spack.package import *
 from spack.pkg.builtin.clingo import Clingo
+from spack.util.environment import EnvironmentModifications
 
 
 class ClingoBootstrap(Clingo):
@@ -13,24 +20,51 @@ class ClingoBootstrap(Clingo):
 
     variant("build_type", default="Release", values=("Release",), description="CMake build type")
 
-    variant("static_libstdcpp", default=False, description="Require a static version of libstdc++")
+    variant(
+        "static_libstdcpp",
+        default=False,
+        when="platform=linux",
+        description="Require a static version of libstdc++",
+    )
+
+    variant(
+        "optimized",
+        default=False,
+        description="Enable a series of Spack-specific optimizations (PGO, LTO, mimalloc)",
+    )
+
+    # Enable LTO
+    conflicts("~ipo", when="+optimized")
+
+    with when("+optimized platform=linux"):
+        # Statically linked. Don't use ~override so we don't duplicate malloc/free, they
+        # get resolved to Python's libc's malloc in our case anyway.
+        depends_on("mimalloc +ipo libs=static ~override", type="build")
+        conflicts("~static_libstdcpp", msg="Custom allocator requires static libstdc++")
+        # Override new/delete with mimalloc.
+        patch("mimalloc.patch", when="@5.5.0:")
+        patch("mimalloc-pre-5.5.0.patch", when="@:5.4")
+        # ensure we hide libstdc++ with custom operator new/delete symbols
+        patch("version-script.patch")
 
     # CMake at version 3.16.0 or higher has the possibility to force the
     # Python interpreter, which is crucial to build against external Python
     # in environment where more than one interpreter is in the same prefix
     depends_on("cmake@3.16.0:", type="build")
 
-    # On Linux we bootstrap with GCC
-    for compiler_spec in [c for c in spack.compilers.supported_compilers() if c != "gcc"]:
+    # On Linux we bootstrap with GCC or clang
+    for compiler_spec in [
+        c for c in spack.compilers.supported_compilers() if c not in ("gcc", "clang")
+    ]:
         conflicts(
             "%{0}".format(compiler_spec),
             when="platform=linux",
-            msg="GCC is required to bootstrap clingo on Linux",
+            msg="GCC or clang are required to bootstrap clingo on Linux",
         )
         conflicts(
             "%{0}".format(compiler_spec),
             when="platform=cray",
-            msg="GCC is required to bootstrap clingo on Cray",
+            msg="GCC or clang are required to bootstrap clingo on Cray",
         )
     conflicts("%gcc@:5", msg="C++14 support is required to bootstrap clingo")
 
@@ -51,28 +85,70 @@ class ClingoBootstrap(Clingo):
 
     def cmake_args(self):
         args = super().cmake_args()
-        args.extend(
-            [
-                # Avoid building the clingo executable
-                self.define("CLINGO_BUILD_APPS", "OFF")
-            ]
-        )
+        args.append(self.define("CLINGO_BUILD_APPS", False))
         return args
 
-    def setup_build_environment(self, env):
-        opts = None
-        if "%apple-clang platform=darwin" in self.spec:
-            opts = "-mmacosx-version-min=10.13"
-        elif "%gcc" in self.spec:
-            if "+static_libstdcpp" in self.spec:
-                # This is either linux or cray
-                opts = "-static-libstdc++ -static-libgcc -Wl,--exclude-libs,ALL"
-        elif "platform=windows" in self.spec:
-            pass
+    @run_before("cmake", when="+optimized")
+    def pgo_train(self):
+        if self.spec.compiler.name == "clang":
+            llvm_profdata = which("llvm-profdata", required=True)
+        elif self.spec.compiler.name == "apple-clang":
+            llvm_profdata = Executable(
+                Executable("xcrun")("-find", "llvm-profdata", output=str).strip()
+            )
+
+        # First configure with PGO flags, and do build apps.
+        reports = os.path.abspath("reports")
+        sources = os.path.abspath(self.root_cmakelists_dir)
+        cmake_options = self.std_cmake_args + self.cmake_args() + [sources]
+
+        # Set PGO training flags.
+        generate_mods = EnvironmentModifications()
+        generate_mods.append_flags("CFLAGS", "-fprofile-generate={}".format(reports))
+        generate_mods.append_flags("CXXFLAGS", "-fprofile-generate={}".format(reports))
+        generate_mods.append_flags("LDFLAGS", "-fprofile-generate={} --verbose".format(reports))
+
+        with working_dir(self.build_directory, create=True):
+            cmake(*cmake_options, sources, extra_env=generate_mods)
+            make()
+            make("install")
+
+        # Clean the reports dir.
+        rmtree(reports, ignore_errors=True)
+
+        # Run spack solve --fresh hdf5 with instrumented clingo.
+        python_runtime_env = EnvironmentModifications()
+        for s in self.spec.traverse(deptype=("run", "link"), order="post"):
+            python_runtime_env.extend(spack.user_environment.environment_modifications_for_spec(s))
+        python_runtime_env.unset("SPACK_ENV")
+        python_runtime_env.unset("SPACK_PYTHON")
+        self.spec["python"].command(
+            spack.paths.spack_script, "solve", "--fresh", "hdf5", extra_env=python_runtime_env
+        )
+
+        # Clean the build dir.
+        rmtree(self.build_directory, ignore_errors=True)
+
+        if self.spec.compiler.name in ("clang", "apple-clang"):
+            # merge reports
+            use_report = join_path(reports, "merged.prof")
+            raw_files = glob.glob(join_path(reports, "*.profraw"))
+            llvm_profdata("merge", "--output={}".format(use_report), *raw_files)
+            use_flag = "-fprofile-instr-use={}".format(use_report)
         else:
-            msg = 'unexpected compiler for spec "{0}"'.format(self.spec)
-            raise RuntimeError(msg)
+            use_flag = "-fprofile-use={}".format(reports)
 
-        if opts:
-            env.set("CXXFLAGS", opts)
-            env.set("LDFLAGS", opts)
+        # Set PGO use flags for next cmake phase.
+        use_mods = EnvironmentModifications()
+        use_mods.append_flags("CFLAGS", use_flag)
+        use_mods.append_flags("CXXFLAGS", use_flag)
+        use_mods.append_flags("LDFLAGS", use_flag)
+        cmake.add_default_envmod(use_mods)
+
+    def setup_build_environment(self, env):
+        if "%apple-clang" in self.spec:
+            env.append_flags("CFLAGS", "-mmacosx-version-min=10.13")
+            env.append_flags("CXXFLAGS", "-mmacosx-version-min=10.13")
+            env.append_flags("LDFLAGS", "-mmacosx-version-min=10.13")
+        elif self.spec.compiler.name in ("gcc", "clang") and "+static_libstdcpp" in self.spec:
+            env.append_flags("LDFLAGS", "-static-libstdc++ -static-libgcc -Wl,--exclude-libs,ALL")
diff --git a/var/spack/repos/builtin/packages/clingo-bootstrap/version-script.patch b/var/spack/repos/builtin/packages/clingo-bootstrap/version-script.patch
new file mode 100644
index 0000000000..df9d90cd65
--- /dev/null
+++ b/var/spack/repos/builtin/packages/clingo-bootstrap/version-script.patch
@@ -0,0 +1,48 @@
+From 59859b8896e527bbd4a727beb798776d2716a8b3 Mon Sep 17 00:00:00 2001
+From: Harmen Stoppels <me@harmenstoppels.nl>
+Date: Thu, 10 Aug 2023 18:53:17 +0200
+Subject: [PATCH] version script
+
+---
+ libclingo/CMakeLists.txt | 12 ++++++++++++
+ libclingo/clingo.map     |  4 ++++
+ 2 files changed, 16 insertions(+)
+ create mode 100644 libclingo/clingo.map
+
+diff --git a/libclingo/CMakeLists.txt b/libclingo/CMakeLists.txt
+index 1d70ba56..0fd3bf49 100644
+--- a/libclingo/CMakeLists.txt
++++ b/libclingo/CMakeLists.txt
+@@ -58,6 +58,18 @@ target_include_directories(libclingo
+     "$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>")
+ target_compile_definitions(libclingo ${clingo_private_scope_} CLINGO_BUILD_LIBRARY)
+ 
++# Hide private symbols on Linux.
++include(CheckCSourceCompiles)
++file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/version.map" "{ global: f; local: *;};")
++set(CMAKE_REQUIRED_FLAGS_SAVE ${CMAKE_REQUIRED_FLAGS})
++set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS} "-Wl,--version-script='${CMAKE_CURRENT_BINARY_DIR}/version.map'")
++check_c_source_compiles("void f(void) {} int main(void) {return 0;}" HAVE_LD_VERSION_SCRIPT)
++set(CMAKE_REQUIRED_FLAGS ${CMAKE_REQUIRED_FLAGS_SAVE})
++file(REMOVE "${CMAKE_CURRENT_BINARY_DIR}/version.map")
++if(HAVE_LD_VERSION_SCRIPT)
++set_target_properties(libclingo PROPERTIES LINK_FLAGS "-Wl,--version-script='${CMAKE_CURRENT_SOURCE_DIR}/clingo.map'")
++endif()
++
+ if (NOT CLINGO_BUILD_SHARED)
+     target_compile_definitions(libclingo PUBLIC CLINGO_NO_VISIBILITY)
+ endif()
+diff --git a/libclingo/clingo.map b/libclingo/clingo.map
+new file mode 100644
+index 00000000..a665456c
+--- /dev/null
++++ b/libclingo/clingo.map
+@@ -0,0 +1,4 @@
++{
++    global: clingo_*; gringo_*; g_clingo_*;
++    local: *;
++};
+\ No newline at end of file
+-- 
+2.39.2
+
diff --git a/var/spack/repos/builtin/packages/clingo/package.py b/var/spack/repos/builtin/packages/clingo/package.py
index babf556cc6..73797762b5 100644
--- a/var/spack/repos/builtin/packages/clingo/package.py
+++ b/var/spack/repos/builtin/packages/clingo/package.py
@@ -120,6 +120,11 @@ class Clingo(CMakePackage):
         else:
             args += ["-DCLINGO_BUILD_WITH_PYTHON=OFF"]
 
+        # Use LTO also for non-Intel compilers please. This can be removed when they
+        # bump cmake_minimum_required to VERSION 3.9.
+        if "+ipo" in self.spec:
+            args.append("-DCMAKE_POLICY_DEFAULT_CMP0069=NEW")
+
         return args
 
     def win_add_library_dependent(self):
diff --git a/var/spack/repos/builtin/packages/mimalloc/package.py b/var/spack/repos/builtin/packages/mimalloc/package.py
index 08cd3fd847..fed6e5bf78 100644
--- a/var/spack/repos/builtin/packages/mimalloc/package.py
+++ b/var/spack/repos/builtin/packages/mimalloc/package.py
@@ -115,4 +115,9 @@ class Mimalloc(CMakePackage):
             for lib in self.libs_values
         ]
         args += [self.define_from_variant("MI_%s" % k.upper(), k) for k in self.mimalloc_options]
+
+        # Use LTO also for non-Intel compilers please. This can be removed when they
+        # bump cmake_minimum_required to VERSION 3.9.
+        if "+ipo" in self.spec:
+            args.append("-DCMAKE_POLICY_DEFAULT_CMP0069=NEW")
         return args