v0.4.2 minor bug fixes

Author: HatPdotS

docs/changelog.rst

@@ -5,6 +5,7 @@ Version 0.4.2
 -------------
 
 - Fixed bug where reflection data object in the refinement was not created on cuda when specified. 
+- Fixed macos crash not catching compilation error in c++ extension for scatter add
 
 Version 0.4.1
 -------------

torchref/base/electron_density/main.py

@@ -83,12 +83,18 @@ def _get_separable_triton():
 
 
 def _get_cpp_scatter():
-    """Return the C++ parallel scatter_add, or None if unavailable."""
+    """Return the C++ parallel scatter_add, or None if unavailable.
+
+    Eagerly triggers the C++ compilation so that failures (missing ninja,
+    unsupported compiler flags, etc.) are caught here rather than mid-calculation.
+    """
     global _cpp_scatter_fn, _cpp_scatter_checked
     if not _cpp_scatter_checked:
         try:
-            from torchref.base.kernels.cpu_scatter import structured_scatter_add
-            _cpp_scatter_fn = structured_scatter_add
+            from torchref.base.kernels.cpu_scatter import structured_scatter_add, _get_module
+            # Trigger compilation now — _get_module returns None on failure
+            if _get_module() is not None:
+                _cpp_scatter_fn = structured_scatter_add
         except Exception:
             pass
         _cpp_scatter_checked = True

torchref/base/kernels/cpu_scatter.py

@@ -165,15 +165,27 @@
 # First process compiles; all others (same node or different) reuse cache.
 # ---------------------------------------------------------------------------
 _module = None
+_module_failed = False
 
 
 def _get_module():
-    global _module
-    if _module is None:
+    global _module, _module_failed
+    if _module is not None:
+        return _module
+    if _module_failed:
+        return None
+
+    import os
+    import sys
+
+    try:
         import fcntl
-        import os
-        import sys
+    except ImportError:
+        # fcntl is not available on non-POSIX platforms (e.g. Windows)
+        _module_failed = True
+        return None
 
+    try:
         # Ensure ninja (installed via pip) is on PATH for compute nodes
         bin_dir = os.path.dirname(sys.executable)
         if bin_dir not in os.environ.get("PATH", ""):
@@ -233,6 +245,10 @@ def _get_module():
         finally:
             fcntl.lockf(lock_fd, fcntl.LOCK_UN)
             os.close(lock_fd)
+    except Exception:
+        _module_failed = True
+        return None
+
     return _module
 
 
@@ -257,6 +273,8 @@ def forward(ctx, density_cube, wa, wbwc, map_size):
         ctx.cube_shape = density_cube.shape
 
         mod = _get_module()
+        if mod is None:
+            raise RuntimeError("C++ cpu_scatter module not available")
         result = torch.zeros(map_size, dtype=density_cube.dtype,
                              device=density_cube.device)
         mod.structured_scatter_add(
@@ -273,6 +291,8 @@ def backward(ctx, grad_output):
         wa, wbwc = ctx.saved_tensors
         C, nx, ny, nz = ctx.cube_shape
         mod = _get_module()
+        if mod is None:
+            raise RuntimeError("C++ cpu_scatter module not available")
         grad_cube = mod.structured_gather(
             grad_output.contiguous(),
             wa.contiguous(),