Native CUDA mesh SDF backend + prebuilt CUDA wheels (0.0.2)

.github/workflows/publish-witwin.yml

@@ -1,19 +1,22 @@
-name: Publish witwin
+name: Build and publish witwin
 
 on:
+  push:
+    branches:
+      - main
+      - "codex/**"
+  pull_request:
   release:
     types: [published]
+  workflow_dispatch:
 
 jobs:
-  publish:
-    if: startsWith(github.event.release.tag_name, 'witwin-v') && !github.event.release.prerelease
+  validate_release:
+    if: github.event_name != 'release' || (startsWith(github.event.release.tag_name, 'witwin-v') && !github.event.release.prerelease)
+    name: Validate release
     runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      id-token: write
-    environment:
-      name: pypi
-      url: https://pypi.org/project/witwin/
+    outputs:
+      package_dir: ${{ steps.package_dir.outputs.path }}
 
     steps:
       - name: Check out repository
@@ -38,6 +41,7 @@ jobs:
           fi
 
       - name: Validate release tag and package version
+        if: github.event_name == 'release'
         shell: bash
         env:
           PACKAGE_DIR: ${{ steps.package_dir.outputs.path }}
@@ -78,15 +82,206 @@ jobs:
           print(f"Validated {package_name} {package_version} from {pyproject_path}.")
           PY
 
+  build_cuda_wheels:
+    name: Build CUDA wheel / ${{ matrix.os }} / py${{ matrix.python_version }}
+    needs: validate_release
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-22.04, windows-2022]
+        python_version: ["3.10", "3.11", "3.12"]
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python_version }}
+
+      - name: Free disk space (Linux)
+        if: runner.os == 'Linux'
+        shell: bash
+        run: |
+          echo "Disk before cleanup:"; df -h /
+          sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/lib/android \
+            /usr/local/share/boost /opt/hostedtoolcache/CodeQL \
+            /usr/local/share/powershell /usr/share/swift || true
+          sudo docker image prune --all --force || true
+          echo "Disk after cleanup:"; df -h /
+
+      - name: Install CUDA Toolkit
+        uses: Jimver/cuda-toolkit@v0.2.29
+        with:
+          cuda: "12.5.0"
+          method: local
+          log-file-suffix: "${{ runner.os }}-py${{ matrix.python_version }}"
+
+      - name: Set up MSVC
+        if: runner.os == 'Windows'
+        uses: ilammy/msvc-dev-cmd@v1
+        with:
+          arch: x64
+
+      - name: Export Windows CUDA_HOME
+        if: runner.os == 'Windows'
+        shell: pwsh
+        run: |
+          $cudaRoot = $env:CUDA_PATH
+          if (-not $cudaRoot) {
+            $nvcc = (Get-Command nvcc -ErrorAction Stop).Source
+            $cudaRoot = Split-Path (Split-Path $nvcc -Parent) -Parent
+          }
+          "CUDA_HOME=$cudaRoot" >> $env:GITHUB_ENV
+          "CUDA_PATH=$cudaRoot" >> $env:GITHUB_ENV
+          "CUDA_ROOT=$cudaRoot" >> $env:GITHUB_ENV
+          "CUDA_BIN_PATH=$cudaRoot" >> $env:GITHUB_ENV
+
+      - name: Show toolchain
+        shell: bash
+        run: |
+          python --version
+          nvcc --version
+          command -v cl || true
+
+      - name: Install CUDA extension build dependencies
+        run: |
+          python -m pip install --upgrade pip setuptools wheel ninja numpy
+          python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124
+          python -m pip install build
+
+      - name: Build packaged mesh SDF CUDA extension
+        shell: bash
+        env:
+          MAX_JOBS: "2"
+          TORCH_CUDA_ARCH_LIST: "7.0;8.0;8.6;8.9;9.0+PTX"
+          WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/witwin_core_mesh_sdf_cuda
+          PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }}
+        run: |
+          python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose
+
+      - name: Build platform wheel
+        shell: bash
+        env:
+          PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }}
+        run: |
+          python -m build --wheel "$PACKAGE_DIR"
+
+      - name: Normalize Linux wheel platform tag
+        if: runner.os == 'Linux'
+        shell: bash
+        env:
+          PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }}
+        run: |
+          # PyPI rejects the raw linux_x86_64 tag; the extension is built on the
+          # ubuntu-22.04 runner (glibc 2.35), so retag to the matching manylinux
+          # tag. torch/CUDA runtime libraries are runtime dependencies and stay
+          # unbundled, so no auditwheel repair is required.
+          shopt -s nullglob
+          for whl in "$PACKAGE_DIR"/dist/*-linux_x86_64.whl; do
+            echo "Retagging $whl"
+            python -m wheel tags --platform-tag manylinux_2_35_x86_64 --remove "$whl"
+          done
+          ls -l "$PACKAGE_DIR"/dist
+
+      - name: Smoke install prebuilt wheel
+        shell: bash
+        env:
+          PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }}
+          WITWIN_CORE_MESH_SDF_CUDA_PREBUILT: "1"
+          WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/missing_prebuilt_build_dir
+        run: |
+          python -m pip uninstall -y witwin || true
+          python -m pip install --no-deps "$PACKAGE_DIR"/dist/*.whl
+          python - <<'PY'
+          import importlib.util
+          import site
+          from pathlib import Path
+
+          candidates = []
+          for root in site.getsitepackages() + [site.getusersitepackages()]:
+              candidates.append(Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py")
+          build_path = next((path for path in candidates if path.exists()), None)
+          assert build_path is not None, f"installed build.py not found in {candidates}"
+
+          spec = importlib.util.spec_from_file_location("witwin_core_mesh_sdf_cuda_build", build_path)
+          module = importlib.util.module_from_spec(spec)
+          spec.loader.exec_module(module)
+
+          path = module.prebuilt_extension_path()
+          assert path.exists(), f"missing packaged prebuilt extension: {path}"
+          extension = module.build_extension()
+          print(f"Loaded packaged mesh SDF CUDA extension from {extension.__file__}")
+          PY
+
+      - name: Upload wheel artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheel-${{ matrix.os }}-py${{ matrix.python_version }}
+          path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.whl
+          if-no-files-found: error
+
+  build_sdist:
+    name: Build source distribution
+    needs: validate_release
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
       - name: Build distributions
         shell: bash
         env:
-          PACKAGE_DIR: ${{ steps.package_dir.outputs.path }}
+          PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }}
         run: |
           python -m pip install --upgrade pip build
-          python -m build "$PACKAGE_DIR"
+          python -m build --sdist "$PACKAGE_DIR"
+
+      - name: Upload source distribution artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.tar.gz
+          if-no-files-found: error
+
+  publish:
+    name: Publish to PyPI
+    if: github.event_name == 'release'
+    needs:
+      - build_cuda_wheels
+      - build_sdist
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      id-token: write
+    environment:
+      name: pypi
+      url: https://pypi.org/project/witwin/
+
+    steps:
+      - name: Download wheel artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: wheel-*
+          path: dist
+          merge-multiple: true
+
+      - name: Download source distribution artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: sdist
+          path: dist
 
       - name: Publish distributions to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          packages-dir: ${{ steps.package_dir.outputs.path }}/dist
+          packages-dir: dist
+          skip-existing: true

.gitignore

@@ -228,3 +228,7 @@ __marimo__/
 
 # Streamlit
 .streamlit/secrets.toml
+
+# Native CUDA mesh SDF prebuilt extensions (built per-platform in CI)
+witwin/core/geometry/cuda/prebuilt/*.pyd
+witwin/core/geometry/cuda/prebuilt/*.so

hatch_build.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from hatchling.builders.hooks.plugin.interface import BuildHookInterface
+
+
+class CustomBuildHook(BuildHookInterface):
+    def initialize(self, version: str, build_data: dict) -> None:
+        del version
+        if self.target_name != "wheel":
+            return
+        prebuilt_dir = Path(self.root) / "witwin" / "core" / "geometry" / "cuda" / "prebuilt"
+        has_prebuilt_extension = any(prebuilt_dir.glob("witwin_core_mesh_sdf_cuda.*"))
+        if has_prebuilt_extension:
+            build_data["infer_tag"] = True
+            build_data["pure_python"] = False

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "witwin"
-version = "0.0.1"
+version = "0.0.2"
 description = "Witwin - Multi-physics simulation platform"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -29,3 +29,14 @@ all     = [
 
 [tool.hatch.build.targets.wheel]
 packages = ["witwin"]
+artifacts = [
+    "witwin/core/geometry/cuda/**/*.cpp",
+    "witwin/core/geometry/cuda/**/*.cu",
+    "witwin/core/geometry/cuda/**/*.cuh",
+    "witwin/core/geometry/cuda/**/*.h",
+    "witwin/core/geometry/cuda/prebuilt/*.pyd",
+    "witwin/core/geometry/cuda/prebuilt/*.so",
+]
+
+[tool.hatch.build.hooks.custom]
+path = "hatch_build.py"

scripts/build_mesh_sdf_cuda_prebuilt.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import torch
+
+
+def _ensure_current_device_arch() -> None:
+    if not torch.cuda.is_available():
+        return
+    major, minor = torch.cuda.get_device_capability()
+    current_arch = f"{major}.{minor}"
+    arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", "")
+    entries = [entry.strip() for entry in arch_list.split(";") if entry.strip()]
+    normalized = {entry.removesuffix("+PTX") for entry in entries}
+    if current_arch not in normalized:
+        entries.append(current_arch)
+        os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(entries)
+
+
+def _load_cuda_build_module():
+    repo_root = Path(__file__).resolve().parents[1]
+    build_path = repo_root / "witwin" / "core" / "geometry" / "cuda" / "build.py"
+    spec = importlib.util.spec_from_file_location("witwin_core_mesh_sdf_cuda_build", build_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load mesh SDF CUDA build module from {build_path}.")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Build the packaged mesh SDF CUDA extension.")
+    parser.add_argument("--verbose", action="store_true")
+    args = parser.parse_args()
+
+    build_dir = Path(
+        os.environ.get(
+            "WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR",
+            Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel",
+        )
+    )
+    os.environ["WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR"] = str(build_dir)
+    os.environ["WITWIN_CORE_MESH_SDF_CUDA_SKIP_PREBUILT"] = "1"
+    _ensure_current_device_arch()
+
+    cuda_build = _load_cuda_build_module()
+    module = cuda_build.build_extension(verbose=args.verbose)
+    module_file = Path(module.__file__).resolve()
+
+    target_dir = cuda_build.prebuilt_root()
+    target_dir.mkdir(parents=True, exist_ok=True)
+    for suffix in (".pyd", ".so"):
+        existing = target_dir / f"witwin_core_mesh_sdf_cuda{suffix}"
+        if existing.exists():
+            existing.unlink()
+    target = cuda_build.prebuilt_extension_path()
+    shutil.copy2(module_file, target)
+    print(f"Built prebuilt mesh SDF CUDA extension: {target}")
+
+
+if __name__ == "__main__":
+    main()