Merge pull request #29 from Arcascope/4.3.0-macos

updates to get working on macOS (again)

Author: ericcanton

CMakeLists.txt

@@ -37,14 +37,20 @@ endif()
 # Flags must be semicolon-separated CMake lists, not space-separated strings,
 # so that each flag is passed as a separate argument to the compiler.
 if(APPLE)
-    # Only use Homebrew-specific hints if Homebrew is available.
-    find_program(HOMEBREW_EXECUTABLE brew)
-    if(HOMEBREW_EXECUTABLE)
+    if(DEFINED ENV{LIBOMP_PREFIX})
+        set(LIBOMP_PREFIX "$ENV{LIBOMP_PREFIX}")
+    else()
+        # Only use Homebrew-specific hints if Homebrew is available.
+        find_program(HOMEBREW_EXECUTABLE brew)
+    endif()
+    if(NOT LIBOMP_PREFIX AND HOMEBREW_EXECUTABLE)
         execute_process(
             COMMAND "${HOMEBREW_EXECUTABLE}" --prefix libomp
             OUTPUT_VARIABLE LIBOMP_PREFIX
             OUTPUT_STRIP_TRAILING_WHITESPACE
         )
+    endif()
+    if(LIBOMP_PREFIX)
         set(OpenMP_C_FLAGS "-Xpreprocessor;-fopenmp;-I${LIBOMP_PREFIX}/include")
         set(OpenMP_CXX_FLAGS "-Xpreprocessor;-fopenmp;-I${LIBOMP_PREFIX}/include")
         set(OpenMP_C_LIB_NAMES "omp")

README.md

@@ -24,6 +24,8 @@ python -m pip install https://github.com/Arcascope/sensor_preprocessing_cpp/rele
 
 Installing from source, `pip install git+https://github.com/Arcascope/sensor_preprocessing_cpp.git@4.0.1`,
 compiles the native extension locally and requires a C++17 toolchain, CMake, and pybind11.
+On macOS it also requires Homebrew's OpenMP runtime (`brew install libomp`). senpy discovers
+that prefix itself; custom installations can set `LIBOMP_PREFIX`.
 
 Release wheels are built by `.github/workflows/release-wheel.yml` (attached to a published
 release, or run manually to backfill a tag) and published to PyPI via OpenID Connect trusted
@@ -186,6 +188,12 @@ one per recording) for the full grid and its sample counts. Batch sizes remain
 a hardware-specific throughput setting: measure with `block_until_ready()` and
 a CUDA profiler before claiming GPU saturation.
 
+On macOS, the standard JAX wheel uses the CPU backend. Pass `enable_x64=True`
+to `compute_nustft_many` when double-precision parity with the native CPU
+transform is more important than float32 throughput. This enables JAX's
+process-wide x64 setting before senpy allocates its buffers. Leave it off for
+Metal, whose kernels do not support float64.
+
 ## Streaming NUSTFT
 
 `StreamingNUSTFT` computes the **same coefficients** from a live stream: 
@@ -257,4 +265,3 @@ band at 30 s windows costs about 150 000 multiply-accumulates per second of stre
 * **Timestamp precision.** Absolute unix seconds in float64 resolve to about half a microsecond,
   which is a ~1e-5 relative phase error at the top of a 5 Hz band. Pass times relative to a recent
   origin when sub-microsecond timing matters.
-

senpy/CMakeLists.txt

@@ -45,7 +45,11 @@ target_include_directories(_core PRIVATE ${CPP_DIR})
 target_include_directories(_core PRIVATE ${FINUFFT_INCLUDE_DIR})
 
 # Link against required libraries
-target_link_libraries(_core PRIVATE ${FINUFFT_LIB} ${FFTW3_LIB} ${FFTW3F_LIB} m gomp)
+find_package(OpenMP REQUIRED)
+target_link_libraries(_core PRIVATE ${FINUFFT_LIB} ${FFTW3_LIB} ${FFTW3F_LIB} OpenMP::OpenMP_CXX)
+if(UNIX AND NOT APPLE)
+    target_link_libraries(_core PRIVATE m)
+endif()
 
 # Set RPATH for the module
 set_target_properties(_core PROPERTIES
@@ -56,4 +60,4 @@ set_target_properties(_core PROPERTIES
 install(TARGETS _core 
     LIBRARY DESTINATION senpy
     RUNTIME DESTINATION senpy
-)
+)

senpy/_version.py

@@ -1 +1 @@
-__version__ = "4.2.0"
+__version__ = "4.3.0"

senpy/jax_backend.py

@@ -992,6 +992,7 @@ def compute_nustft_many(
     rows_per_call: int = DEFAULT_ROWS_PER_CALL,
     max_in_flight: int = 3,
     build_threads: int = 4,
+    enable_x64: bool = False,
 ) -> List[List[Any]]:
     """Transform many multi-channel recordings at device throughput.
 
@@ -1013,6 +1014,12 @@ def compute_nustft_many(
     is enabled. Local sample coordinates are computed on the host in float64
     before that cast, so long recordings keep their phase precision.
 
+    ``enable_x64=True`` enables JAX's process-wide 64-bit mode before buffers
+    are allocated. This is useful on macOS's CPU backend when close numerical
+    agreement with the double-precision native transform matters. It is opt-in
+    because 64-bit mode affects all JAX work in the process and is unavailable
+    on some accelerators, including Metal.
+
     Like the other transforms here, a recording with nothing to report gets
     zero rows rather than an error.
     """
@@ -1041,6 +1048,8 @@ def compute_nustft_many(
     _timestamp_scale(ts_unit)
 
     jax, jnp = _dependencies()
+    if enable_x64:
+        jax.config.update("jax_enable_x64", True)
     host_dtype = np.float64 if jnp.asarray(0.0).dtype == np.float64 else np.float32
     complex_dtype = np.complex128 if host_dtype == np.float64 else np.complex64
 

senpy/setup.py

@@ -57,6 +57,41 @@ def __str__(self):
 fftw3_lib = os.path.join(BUILD_DIR, "_deps/fftw3-build/libfftw3.a")
 fftw3f_lib = os.path.join(BUILD_DIR, "_deps/fftw3f-build/libfftw3f.a")
 
+
+def find_libomp_prefix():
+    """Return the macOS libomp prefix without relying on shell flags."""
+    if platform.system() != "Darwin":
+        return None
+    configured = os.environ.get("LIBOMP_PREFIX")
+    if configured:
+        prefix = configured
+    else:
+        brew = shutil.which("brew")
+        if brew is None:
+            raise RuntimeError(
+                "Building senpy on macOS requires libomp. Install it with "
+                "'brew install libomp', or set LIBOMP_PREFIX to its install prefix."
+            )
+        try:
+            prefix = subprocess.check_output(
+                [brew, "--prefix", "libomp"], text=True
+            ).strip()
+        except subprocess.CalledProcessError as exc:
+            raise RuntimeError(
+                "Homebrew could not locate libomp. Install it with "
+                "'brew install libomp', or set LIBOMP_PREFIX to its install prefix."
+            ) from exc
+    library = os.path.join(prefix, "lib", "libomp.dylib")
+    if not os.path.exists(library):
+        raise RuntimeError(
+            f"libomp was not found at {library}. Install it with "
+            "'brew install libomp', or correct LIBOMP_PREFIX."
+        )
+    return prefix
+
+
+LIBOMP_PREFIX = find_libomp_prefix()
+
 def required_artifacts():
     return [finufft_include_dir, finufft_lib_so, fftw3_lib, fftw3f_lib]
 
@@ -111,20 +146,29 @@ def _copy_runtime_libraries(self):
 
 include_dirs.append(finufft_include_dir)
 
+extra_link_args = [
+    f'-L{finufft_lib_dir}',
+    '-lfinufft',
+    fftw3_lib, fftw3f_lib,
+    '-Wl,-rpath,@loader_path' if platform.system() == 'Darwin' else '-Wl,-rpath,$ORIGIN',
+]
+if LIBOMP_PREFIX is not None:
+    extra_link_args.extend([
+        f'-L{os.path.join(LIBOMP_PREFIX, "lib")}',
+        f'-Wl,-rpath,{os.path.join(LIBOMP_PREFIX, "lib")}',
+        '-lomp',
+    ])
+else:
+    extra_link_args.append('-lgomp')
+
 ext_modules = [
     Extension(
         'senpy._core',  # Full module path - creates senpy/_core.so
         sources=sources,
         include_dirs=include_dirs,
         language='c++',
         extra_compile_args=['-std=c++17', '-O3', '-march=x86-64-v3' if os.environ.get('CI') == 'true' else '-march=native', '-DPYTHON'],
-        extra_link_args=[
-            f'-L{finufft_lib_dir}',
-            '-lfinufft',
-            fftw3_lib, fftw3f_lib,
-            '-Wl,-rpath,@loader_path' if platform.system() == 'Darwin' else '-Wl,-rpath,$ORIGIN',
-            '-lomp' if platform.system() == 'Darwin' else '-lgomp',
-        ],
+        extra_link_args=extra_link_args,
     ),
 ]
 

senpy/tests/test_spectral_api.py

@@ -276,4 +276,4 @@ def test_removed_nufft_alias_and_uniform_compatibility_alias_warns():
 
 
 def test_version_is_consolidated():
-    assert senpy.__version__ == sp.__version__ == "4.2.0"
+    assert senpy.__version__ == sp.__version__ == "4.3.0"