feat: use thrust::inclusive_scan for recursion

.github/workflows/python-package.yml

@@ -46,7 +46,7 @@ jobs:
 
   build-macos:
     if: github.event_name == 'pull_request'
-    runs-on: macos-latest
+    runs-on: macos-15
     strategy:
       fail-fast: false
       matrix:
@@ -74,7 +74,7 @@ jobs:
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Build CPP extension with clang++
       run: |
-        export CXX=$(brew --prefix llvm@15)/bin/clang++
+        export CXX=$(brew --prefix llvm@18)/bin/clang++
         export LDFLAGS="-L/usr/local/opt/libomp/lib"
         export CPPFLAGS="-I/usr/local/opt/libomp/include"
         python -m pip install -e .

setup.py

@@ -1,6 +1,7 @@
 import setuptools
 import os
 import glob
+from platform import system
 import torch
 from torch.utils.cpp_extension import (
     CppExtension,
@@ -36,6 +37,13 @@ def get_extensions():
         extra_compile_args["cxx"] = ["-fopenmp"]
         extra_link_args.append("-fopenmp")
 
+    if system() == "Darwin":
+        extra_compile_args["cxx"] = (
+            ["-fexperimental-library"]
+            if "cxx" not in extra_compile_args
+            else extra_compile_args["cxx"] + ["-fexperimental-library"]
+        )
+
     this_dir = os.path.abspath(os.path.dirname(__file__))
     extensions_dir = os.path.join(this_dir, library_name, "csrc")
     sources = list(glob.glob(os.path.join(extensions_dir, "*.cpp")))
@@ -45,6 +53,7 @@ def get_extensions():
 
     if use_cuda:
         sources += cuda_sources
+        extra_compile_args["nvcc"] = ["--extended-lambda"]
 
     ext_modules = [
         extension(