Allow to use a differently suffixed jemalloc (#147)

Author: xhochy

CHANGELOG.rst

@@ -7,6 +7,13 @@
 Changelog
 =========
 
+3.0.2 - 2021-10-14
+------------------
+
+**Bug fix**
+
+- Allow to link to alternatively suffixed jemalloc installation to workaround `#113 <https://github.com/Quantco/tabmat/issues/113>`_ .
+
 3.0.1 - 2021-10-07
 ------------------
 

conda.recipe/meta.yaml

@@ -27,7 +27,7 @@ requirements:
   host:
     - python
     - cython
-    - jemalloc     # [not win]
+    - jemalloc-local  # [not win]
     - llvm-openmp  # [osx]
     - mako
     - numpy

setup.py

@@ -1,8 +1,10 @@
 import io
 import os
 import platform
+import shutil
 import sys
 from os import path
+from pathlib import Path
 
 import mako.template
 import numpy as np
@@ -55,12 +57,25 @@
     # make sure we can find xsimd headers
     include_dirs.append(os.path.join(sys.prefix, "Library", "include"))
 else:
-    allocator_libs = ["jemalloc"]
+    jemalloc_config = shutil.which("jemalloc-config")
+    if jemalloc_config is None:
+        je_install_suffix = ""
+    else:
+        pkg_info = (
+            Path(jemalloc_config).parent.parent / "lib" / "pkgconfig" / "jemalloc.pc"
+        ).read_text()
+        je_install_suffix = [
+            i.split("=")[1]
+            for i in pkg_info.split("\n")
+            if i.startswith("install_suffix=")
+        ].pop()
+    allocator_libs = [f"jemalloc{je_install_suffix}"]
     extra_compile_args = [
         "-fopenmp",
         "-O3",
         "-ffast-math",
         "--std=c++17",
+        f"-DJEMALLOC_INSTALL_SUFFIX={je_install_suffix}",
     ]
     extra_link_args = ["-fopenmp"]
 

src/tabmat/ext/alloc.h

@@ -3,9 +3,17 @@
 #ifndef _WIN32
   #define JEMALLOC_NO_DEMANGLE
   #if __APPLE__
-    #define JEMALLOC_NO_RENAME
+    #if !(!JEMALLOC_INSTALL_SUFFIX)
+      #define JEMALLOC_NO_RENAME
+    #endif
   #endif
-  #include <jemalloc/jemalloc.h>
+  // Compute jemalloc include path
+  #define STRINGIFY(X) STRINGIFY2(X)
+  #define STRINGIFY2(X) #X
+  #define CAT(X,Y) CAT2(X,Y)
+  #define CAT2(X,Y) X##Y
+  #define JE_INCLUDE STRINGIFY(CAT(jemalloc/jemalloc,JEMALLOC_INSTALL_SUFFIX).h)
+  #include JE_INCLUDE
 #endif
 
 #include <functional>