Build release runners with platform accelerators

Author: fffonion

.github/scripts/patch-koharu-diffusion.py

@@ -29,6 +29,60 @@ def koharu_enums_path() -> Path:
     return Path(matches[0]["manifest_path"]).parent / "src" / "enums.rs"
 
 
+def koharu_libtorch_path() -> Path:
+    metadata = subprocess.run(
+        ["cargo", "metadata", "--locked", "--format-version", "1"],
+        check=True,
+        stdout=subprocess.PIPE,
+        text=True,
+        encoding="utf-8",
+    )
+    packages = json.loads(metadata.stdout)["packages"]
+    matches = [
+        package
+        for package in packages
+        if package["name"] == "koharu-runtime"
+        and package["source"].startswith("git+")
+    ]
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"expected one Git koharu-runtime package, found {len(matches)}"
+        )
+    return Path(matches[0]["manifest_path"]).parent / "src" / "package" / "libtorch.rs"
+
+
+def patch_libtorch_selection() -> None:
+    libtorch_path = koharu_libtorch_path()
+    libtorch_source = libtorch_path.read_text(encoding="utf-8")
+    libtorch_marker = "// Patched by Flint: allow release CI to select a LibTorch device."
+    if libtorch_marker in libtorch_source:
+        print(f"already patched {libtorch_path}")
+        return
+
+    libtorch_old = "    pub fn for_current_target() -> Result<Self> {\n"
+    libtorch_new = libtorch_old + f"""        {libtorch_marker}
+        if let Some(device) = std::env::var_os("FLINT_LIBTORCH_DEVICE") {{
+            match device.to_str() {{
+                Some("cpu") => return Ok(Self::Cpu),
+                Some("cu126") => return Ok(Self::Cuda126),
+                Some("cu129") => return Ok(Self::Cuda129),
+                Some("cu130") => return Ok(Self::Cuda130),
+                Some(device) => bail!("unsupported FLINT_LIBTORCH_DEVICE '{{device}}'"),
+                None => bail!("FLINT_LIBTORCH_DEVICE is not valid UTF-8"),
+            }}
+        }}
+"""
+    if libtorch_source.count(libtorch_old) != 1:
+        raise RuntimeError("could not locate Koharu LibTorch target selection")
+    libtorch_path.write_text(
+        libtorch_source.replace(libtorch_old, libtorch_new), encoding="utf-8"
+    )
+    print(f"patched {libtorch_path}")
+
+
+patch_libtorch_selection()
+
+
 path = koharu_enums_path()
 source = path.read_text(encoding="utf-8")
 

.github/workflows/release.yml

@@ -21,6 +21,7 @@ concurrency:
 env:
   CARGO_TERM_COLOR: always
   CARGO_TARGET_DIR: target
+  LIBTORCH_VERSION: "2.12.1"
 
 jobs:
   build-runner:
@@ -33,12 +34,15 @@ jobs:
           - os: ubuntu-latest
             platform: linux-x86_64
             exe_suffix: ""
+            libtorch_device: cpu
           - os: windows-latest
             platform: windows-x86_64
             exe_suffix: ".exe"
+            libtorch_device: cu130
           - os: macos-14
             platform: macos-arm64
             exe_suffix: ""
+            libtorch_device: cpu
 
     steps:
       - name: Checkout Flint
@@ -55,12 +59,11 @@ jobs:
       - name: Install Rust toolchain
         uses: dtolnay/rust-toolchain@stable
 
-      - name: Patch Koharu C enum ABI compatibility
+      - name: Patch Koharu release dependencies
         working-directory: flint
         run: python .github/scripts/patch-koharu-diffusion.py
 
-      - name: Preload Unix LibTorch
-        if: runner.os != 'Windows'
+      - name: Prefetch platform LibTorch
         working-directory: flint
         shell: python
         env:
@@ -73,19 +76,55 @@ jobs:
           import urllib.request
           import zipfile
 
-          version = "2.12.1"
+          version = os.environ["LIBTORCH_VERSION"]
           runner_os = os.environ["RUNNER_OS"]
           if runner_os == "Linux":
+              device = "cpu"
               archive_name = f"libtorch-shared-with-deps-{version}%2Bcpu.zip"
-              required = ("libc10.so", "libtorch_cpu.so", "libtorch.so")
+              required = (
+                  "libgomp.so.1",
+                  "libc10.so",
+                  "libshm.so",
+                  "libtorch_global_deps.so",
+                  "libtorch_cpu.so",
+                  "libtorch.so",
+              )
           elif runner_os == "macOS":
+              # Apple Silicon LibTorch includes the MPS backend implemented on Metal.
+              device = "cpu"
               archive_name = f"libtorch-macos-arm64-{version}.zip"
-              required = ("libomp.dylib", "libc10.dylib", "libtorch_cpu.dylib", "libtorch.dylib")
+              required = (
+                  "libomp.dylib",
+                  "libc10.dylib",
+                  "libshm.dylib",
+                  "libtorch_global_deps.dylib",
+                  "libtorch_cpu.dylib",
+                  "libtorch.dylib",
+              )
+          elif runner_os == "Windows":
+              # GitHub's Windows runner has no CUDA driver. Prefetch CUDA LibTorch
+              # explicitly so the released archive can use CUDA on a user's machine.
+              device = "cu130"
+              archive_name = f"libtorch-win-shared-with-deps-{version}%2B{device}.zip"
+              required = (
+                  "libiomp5md.dll",
+                  "libiompstubs5md.dll",
+                  "zlibwapi.dll",
+                  "uv.dll",
+                  "c10.dll",
+                  "c10_cuda.dll",
+                  "caffe2_nvrtc.dll",
+                  "torch_global_deps.dll",
+                  "torch_cpu.dll",
+                  "torch_cuda.dll",
+                  "shm.dll",
+                  "torch.dll",
+              )
           else:
               raise RuntimeError(f"unsupported runner OS: {runner_os}")
-          url = f"https://download.pytorch.org/libtorch/cpu/{archive_name}"
+          url = f"https://download.pytorch.org/libtorch/{device}/{archive_name}"
           target_dir = pathlib.Path(os.environ.get("CARGO_TARGET_DIR", "target"))
-          package_dir = target_dir / "store" / "libtorch" / version / "cpu"
+          package_dir = target_dir / "store" / "libtorch" / version / device
           lib_dir = package_dir / "libtorch" / "lib"
 
           if all((lib_dir / name).is_file() for name in required):
@@ -128,6 +167,8 @@ jobs:
 
       - name: Build release runner
         working-directory: flint
+        env:
+          FLINT_LIBTORCH_DEVICE: ${{ matrix.libtorch_device }}
         run: cargo build --release --locked --bin flint
 
       - name: Package runner
@@ -136,6 +177,8 @@ jobs:
         env:
           PLATFORM: ${{ matrix.platform }}
           EXE_SUFFIX: ${{ matrix.exe_suffix }}
+          RUNNER_OS: ${{ runner.os }}
+          LIBTORCH_VERSION: ${{ env.LIBTORCH_VERSION }}
         run: |
           import os
           import pathlib
@@ -161,13 +204,40 @@ jobs:
               raise FileNotFoundError(f"missing built runner binary: {binary}")
           shutil.copy2(binary, package_dir / binary.name)
 
+          native_library_patterns = {
+              "Windows": "koharu*.dll",
+              "macOS": "libkoharu*.dylib",
+              "Linux": "libkoharu*.so",
+          }
+          pattern = native_library_patterns[os.environ["RUNNER_OS"]]
+          for source in release_dir.glob(pattern):
+              shutil.copy2(source, package_dir / source.name)
+
           for extra in ("README.md", "LICENSE"):
               source = root / extra
               if source.exists():
                   shutil.copy2(source, package_dir / extra)
 
+          platform = os.environ["PLATFORM"]
+          if platform == "windows-x86_64":
+              required = (
+                  "flint.exe",
+                  "koharu-torch.dll",
+              )
+          elif platform == "macos-arm64":
+              required = (
+                  "flint",
+                  "libkoharu-torch.dylib",
+              )
+          else:
+              required = ("flint", "libkoharu-torch.so")
+          missing = [name for name in required if not (package_dir / name).is_file()]
+          if missing:
+              raise RuntimeError(f"release package is missing: {', '.join(missing)}")
+
+          all_files = sorted(path for path in package_dir.rglob("*") if path.is_file())
           with zipfile.ZipFile(archive, "w", compression=zipfile.ZIP_DEFLATED) as zf:
-              for path in package_dir.rglob("*"):
+              for path in all_files:
                   zf.write(path, path.relative_to(package_dir.parent))
 
           with open(os.environ["GITHUB_OUTPUT"], "a", encoding="utf-8") as output:

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "rs-flint"
-version = "0.1.0"
+version = "0.1.8"
 edition = "2024"
 license = "Apache-2.0"
 description = "RustScript-native AI inference with Torch, llama.cpp, and stable-diffusion.cpp"

Cargo.toml

@@ -51,6 +51,16 @@ imported as `flint_ai`.
 repository checkout so the referenced files under `scripts/` are available;
 `--script` also accepts an absolute path to an RSS program.
 
+## Release Archives
+
+Release assets are zip archives containing the runner and its native Koharu
+shims. LibTorch, llama.cpp, and stable-diffusion.cpp are fetched by
+`koharu-runtime` on first use and stored next to the executable under `store`.
+The `macos-arm64` runner uses Apple Silicon LibTorch with its MPS backend,
+which uses Metal; run Torch scripts with `--device mps`. The `windows-x86_64`
+runner is built against CUDA 13 LibTorch; with a compatible NVIDIA driver,
+Torch scripts can use `--device cuda` or `--device cuda:N`.
+
 ## CLI
 
 The `flint` binary has explicit modes: