Plug-and-play wheels: vendored FKL headers + arch/CUDA auto-detection (pip install and go)

.github/workflows/wheels.yml

@@ -24,6 +24,9 @@ jobs:
       - name: Install build tooling
         run: python -m pip install --upgrade build twine
 
+      - name: Vendor FKL headers (self-contained wheel)
+        run: python scripts/vendor_fkl.py --ref LTS-C++17
+
       - name: Build sdist and wheel
         run: python -m build
 
@@ -34,8 +37,13 @@ jobs:
         run: |
           python -m pip install dist/*.whl
           python - <<'EOF'
-          # Import must work WITHOUT CUDA/FKL present (lazy backend init).
+          # Import must work WITHOUT CUDA present (lazy backend init).
           import fkl
+          # vendored headers must be inside the installed wheel
+          from pathlib import Path
+          vend = Path(fkl.__file__).parent / "_vendor" / "FusedKernelLibrary" / "include" / "fused_kernel" / "fused_kernel.h"
+          assert vend.exists(), f"vendored headers missing: {vend}"
+          print("vendored FKL headers present in wheel")
           expected = {"compose", "compose_divergent", "CircularTensor",
                       "TensorRead", "TensorWrite", "Crop", "Resize",
                       "Warping", "Mul", "Add", "DeviceBuffer"}
@@ -170,3 +178,33 @@ jobs:
         with:
           files: dist/*
           generate_release_notes: true
+
+  publish-pypi:
+    name: Publish to PyPI
+    needs: [build, codegen-compile-check]
+    if: startsWith(github.ref, 'refs/tags/v')
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/fkl-python
+    permissions:
+      id-token: write   # OIDC token for PyPI Trusted Publishing (no API keys)
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist
+
+      # Trusted Publishing: PyPI trusts this repo+workflow+environment via
+      # OIDC. One-time setup on pypi.org (no secrets stored in GitHub):
+      #   PyPI -> Your projects -> Publishing -> Add a new pending publisher
+      #     project:     fkl-python
+      #     owner:       <repo owner, e.g. Libraries-Openly-Fused>
+      #     repository:  fkl_python
+      #     workflow:    wheels.yml
+      #     environment: pypi
+      - name: Publish to PyPI (trusted publishing)
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist
+          skip-existing: true

.gitignore

@@ -11,3 +11,4 @@ build/
 err.log
 examples/assets/
 .venv-yolo/
+fkl/_vendor/

README.md

@@ -58,14 +58,33 @@ No torch? Use the built-in dependency-free `DeviceBuffer` (CUDA driver via ctype
 
 ## Setup
 
+### Plug and play (wheel with vendored headers)
+
 ```bash
-export FKL_INCLUDE=/path/to/FusedKernelLibrary/include
-export FKL_ROOT=/path/to/FusedKernelLibrary
-export FKL_ARCH=sm_120          # your GPU arch
+pip install fkl-python      # once published; or the wheel from GitHub Releases
+python -c "import fkl; ..."  # just works: no env vars, no FKL checkout
+```
+
+The wheel ships the FKL headers inside (`fkl/_vendor/`, header-only,
+Apache-2.0, upstream commit recorded in `VENDOR_INFO.txt`). At first use it
+auto-detects your GPU arch (CUDA driver API) and CUDA toolkit, JIT-compiles
+the chain, and caches the .so under `~/.cache/fkl`. Requirements on the
+machine: an NVIDIA driver + a CUDA toolkit (nvcc) or clang++.
+
+### Development (use your own FKL checkout)
+
+```bash
+export FKL_INCLUDE=/path/to/FusedKernelLibrary/include   # overrides vendored
+export FKL_ARCH=sm_120                                   # optional override
 pip install -e .
 python tests/test_e2e.py
 ```
 
+Resolution order: `FKL_INCLUDE` env var > vendored headers in the wheel >
+sibling dev checkout. Arch: `FKL_ARCH` > driver query of GPU 0 > sm_75.
+To (re)vendor headers before building a wheel:
+`python scripts/vendor_fkl.py --ref LTS-C++17`.
+
 ## Status
 
 Verified end-to-end on RTX PRO 6000 Blackwell (sm_120), CUDA 13.3, on BOTH

fkl/backend.py

@@ -19,17 +19,84 @@
 
 # ---- configuration -------------------------------------------------------
 
-_FKL_INCLUDE = os.environ.get(
-    "FKL_INCLUDE",
-    "/home/johnny/Projects/oscar/FusedKernelLibrary/include",
-)
-_FKL_ROOT = os.environ.get(
-    "FKL_ROOT",
-    "/home/johnny/Projects/oscar/FusedKernelLibrary",
-)
-_CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")
+def _resolve_fkl_include() -> str:
+    """FKL headers resolution, in priority order:
+    1. FKL_INCLUDE env var (explicit override / developer checkout)
+    2. vendored headers shipped inside the wheel (fkl/_vendor/...)
+    3. well-known sibling checkout (development convenience)
+    """
+    env = os.environ.get("FKL_INCLUDE")
+    if env:
+        return env
+    vendored = Path(__file__).parent / "_vendor" / "FusedKernelLibrary" / "include"
+    if (vendored / "fused_kernel" / "fused_kernel.h").exists():
+        return str(vendored)
+    dev = Path.home() / "Projects" / "oscar" / "FusedKernelLibrary" / "include"
+    if (dev / "fused_kernel" / "fused_kernel.h").exists():
+        return str(dev)
+    raise RuntimeError(
+        "FusedKernelLibrary headers not found. Either:\n"
+        "  - pip install a wheel with vendored headers (run scripts/vendor_fkl.py"
+        " before building), or\n"
+        "  - set FKL_INCLUDE to <FusedKernelLibrary>/include")
+
+
+_FKL_INCLUDE = None  # resolved lazily on first compile (import must not fail)
+
+
+def fkl_include() -> str:
+    global _FKL_INCLUDE
+    if _FKL_INCLUDE is None:
+        _FKL_INCLUDE = _resolve_fkl_include()
+    return _FKL_INCLUDE
+
+
+def fkl_root() -> str:
+    # FKL_ROOT is only needed for in-repo test utilities (<tests/main.h>);
+    # generated chains only need include/. Default: include's parent.
+    return os.environ.get("FKL_ROOT", str(Path(fkl_include()).parent))
+
+
+def _detect_cuda_home() -> str:
+    env = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH")
+    if env:
+        return env
+    if Path("/usr/local/cuda").exists():
+        return "/usr/local/cuda"
+    nvcc = shutil.which("nvcc")
+    if nvcc:
+        return str(Path(nvcc).parent.parent)
+    return "/usr/local/cuda"  # last resort; clear error surfaces at compile
+
+
+_CUDA_HOME = _detect_cuda_home()
 _CACHE_DIR = Path(os.environ.get("FKL_CACHE", str(Path.home() / ".cache" / "fkl")))
-_ARCH = os.environ.get("FKL_ARCH", "sm_120")
+def _detect_arch() -> str:
+    """GPU arch resolution for plug-and-play installs:
+    1. FKL_ARCH env var (explicit override)
+    2. compute capability of GPU 0 via the CUDA driver API (ctypes; no
+       dependencies; works without nvcc/nvidia-smi on PATH)
+    3. sm_75 floor as a last resort (compiles everywhere Turing+)
+    """
+    env = os.environ.get("FKL_ARCH")
+    if env:
+        return env
+    try:
+        import ctypes
+        cuda = ctypes.CDLL("libcuda.so.1")
+        if cuda.cuInit(0) == 0:
+            major, minor, dev = ctypes.c_int(), ctypes.c_int(), ctypes.c_int()
+            if cuda.cuDeviceGet(ctypes.byref(dev), 0) == 0:
+                cuda.cuDeviceGetAttribute(ctypes.byref(major), 75, dev)   # COMPUTE_CAPABILITY_MAJOR
+                cuda.cuDeviceGetAttribute(ctypes.byref(minor), 76, dev)   # COMPUTE_CAPABILITY_MINOR
+                if major.value > 0:
+                    return f"sm_{major.value}{minor.value}"
+    except Exception:
+        pass
+    return "sm_75"
+
+
+_ARCH = _detect_arch()
 _STD = os.environ.get("FKL_STD", "c++20")
 
 
@@ -50,7 +117,7 @@ def _cuda_version() -> str:
 def _fkl_version() -> str:
     # cheap content fingerprint of the public header to invalidate cache on bumps
     h = hashlib.sha1()
-    fk_h = Path(_FKL_INCLUDE) / "fused_kernel" / "fused_kernel.h"
+    fk_h = Path(fkl_include()) / "fused_kernel" / "fused_kernel.h"
     try:
         h.update(fk_h.read_bytes())
     except Exception:
@@ -156,15 +223,15 @@ def _cpu_cmd(self, src: Path, so: Path):
         if not cpp.exists():
             cpp.write_text(src.read_text())
         return [cxx, f"-std={_STD}", "-O2", "-shared", "-fPIC",
-                "-I", _FKL_INCLUDE, "-I", _FKL_ROOT,
+                "-I", fkl_include(), "-I", fkl_root(),
                 str(cpp), "-o", str(so)]
 
     def _nvcc_cmd(self, src: Path, so: Path):
         nvcc = str(Path(_CUDA_HOME) / "bin" / "nvcc")
         return [
             nvcc, f"-std={_STD}", f"-arch={_ARCH}",
             "-shared", "-Xcompiler", "-fPIC",
-            "-I", _FKL_INCLUDE, "-I", _FKL_ROOT,
+            "-I", fkl_include(), "-I", fkl_root(),
             str(src), "-o", str(so),
         ]
 
@@ -178,7 +245,7 @@ def _clang_cmd(self, src: Path, so: Path):
             # clang's wrapper chain leaves undefined -> define it explicitly.
             "-D_NV_RSQRT_SPECIFIER=noexcept(true)",
             "-shared", "-fPIC",
-            "-I", _FKL_INCLUDE, "-I", _FKL_ROOT,
+            "-I", fkl_include(), "-I", fkl_root(),
             str(src), "-o", str(so),
             f"-L{_CUDA_HOME}/lib64", "-lcudart",
         ]

pyproject.toml

@@ -40,9 +40,17 @@ Repository = "https://github.com/johnnynunez/fkl_python"
 [tool.setuptools.packages.find]
 include = ["fkl*"]
 
+[tool.setuptools.package-data]
+# vendored FKL headers (header-only, Apache-2.0): makes wheels self-contained.
+# populate with scripts/vendor_fkl.py before building.
+"fkl._vendor" = ["FusedKernelLibrary/LICENSE",
+                 "FusedKernelLibrary/VENDOR_INFO.txt",
+                 "FusedKernelLibrary/include/**/*.h",
+                 "FusedKernelLibrary/include/**/*.cuh"]
+
 # Environment variables understood at runtime:
-#   FKL_INCLUDE  path to FusedKernelLibrary/include   (required)
-#   FKL_ROOT     path to FusedKernelLibrary repo root (required)
+#   FKL_INCLUDE  override FKL headers (default: vendored copy in the wheel)
+#   FKL_ROOT     override FKL repo root (only needed for in-repo test utils)
 #   CUDA_HOME    path to CUDA toolkit  (default /usr/local/cuda)
 #   FKL_ARCH     GPU arch              (default sm_120)
 #   FKL_STD      C++ standard          (default c++20)

scripts/vendor_fkl.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""Vendor the FusedKernelLibrary headers into fkl/_vendor/ so the wheel is
+self-contained: `pip install fkl-python` works with no FKL checkout and no
+FKL_INCLUDE/FKL_ROOT environment variables.
+
+FKL is header-only and Apache-2.0, same license as this package, so
+redistributing the headers inside the wheel is clean. The LICENSE and the
+exact upstream commit are recorded alongside.
+
+Usage:
+    python scripts/vendor_fkl.py [--source /path/to/FusedKernelLibrary]
+                                 [--ref LTS-C++17]
+
+With --source, copies from a local checkout (records its HEAD commit).
+Without it, shallow-clones the given ref from GitHub into a temp dir.
+"""
+import argparse
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+REPO_URL = "https://github.com/Libraries-Openly-Fused/FusedKernelLibrary.git"
+PKG_ROOT = Path(__file__).resolve().parent.parent
+VENDOR_DIR = PKG_ROOT / "fkl" / "_vendor" / "FusedKernelLibrary"
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--source", type=Path, default=None,
+                    help="local FusedKernelLibrary checkout to copy from")
+    ap.add_argument("--ref", default="LTS-C++17",
+                    help="git ref to clone when --source is not given")
+    args = ap.parse_args()
+
+    tmp = None
+    if args.source is not None:
+        src = args.source.resolve()
+    else:
+        tmp = tempfile.mkdtemp(prefix="fkl_vendor_")
+        subprocess.run(["git", "clone", "--depth", "1", "--branch", args.ref,
+                        REPO_URL, tmp], check=True)
+        src = Path(tmp)
+
+    inc = src / "include"
+    lic = src / "LICENSE"
+    if not (inc / "fused_kernel" / "fused_kernel.h").exists():
+        print(f"error: {inc} does not look like FKL's include dir", file=sys.stderr)
+        return 1
+
+    commit = subprocess.run(["git", "-C", str(src), "rev-parse", "HEAD"],
+                            capture_output=True, text=True, check=True).stdout.strip()
+    branch = subprocess.run(["git", "-C", str(src), "rev-parse", "--abbrev-ref", "HEAD"],
+                            capture_output=True, text=True, check=True).stdout.strip()
+
+    if VENDOR_DIR.exists():
+        shutil.rmtree(VENDOR_DIR)
+    VENDOR_DIR.mkdir(parents=True)
+    shutil.copytree(inc, VENDOR_DIR / "include")
+    if lic.exists():
+        shutil.copy2(lic, VENDOR_DIR / "LICENSE")
+    (VENDOR_DIR / "VENDOR_INFO.txt").write_text(
+        f"FusedKernelLibrary headers vendored for self-contained wheels.\n"
+        f"upstream: {REPO_URL}\nbranch: {branch}\ncommit: {commit}\n"
+        f"license: Apache-2.0 (see LICENSE in this directory)\n")
+    # package marker so setuptools treats it as data within the fkl package
+    (PKG_ROOT / "fkl" / "_vendor" / "__init__.py").touch()
+
+    n = sum(1 for _ in (VENDOR_DIR / "include").rglob("*.h"))
+    print(f"vendored {n} headers from {branch}@{commit[:9]} -> {VENDOR_DIR}")
+    if tmp:
+        shutil.rmtree(tmp, ignore_errors=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())