torchref-kernels 0.1.0__tar.gz

torchref_kernels-0.1.0/Cargo.lock

@@ -0,0 +1,223 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
+
+[[package]]
+name = "either"
+version = "1.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "libc"
+version = "0.2.190"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ce5d3ddc6d3fa000eb1536d85e147bfe31aacaba692ed6a876f95cb7c855be78"
+
+[[package]]
+name = "multiversion"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b4ca4bea16ffc3f443cf7d866912118196bfef4c6a1556ca00f9f9b00bb43f7c"
+dependencies = [
+ "multiversion-macros",
+]
+
+[[package]]
+name = "multiversion-macros"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0d416831a7317ef4b08bee00b69cbbb9c8763da7959a7026244d6266869f9c83"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "rustversion",
+ "syn 3.0.7",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.29.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "700d18fa267b73b9b521fd7e13580e2f446916f176cacee1ab63fcc8191f1655"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.29.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7b3fc0c4d08f6bb10e71fe39dfb9e2f59c6eb6854e22ec8092f50c69a4499adb"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.29.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dfc0b8e19df29aad7086cf977bb0c2a2f143e30567eb113e9cf72b62ca698330"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.29.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6100e8a4b5eba53afaa5ed078364851a0b2499c553a44c31026b929049b49dc6"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.29.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6143877a16e82b5a727b7127ff4cd86858a24a28d745d72f43e6f227c7b1bdb3"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.119",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "rayon"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "rustversion"
+version = "1.0.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f"
+
+[[package]]
+name = "syn"
+version = "2.0.119"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "syn"
+version = "3.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d62a2e0561533f2ca2561d0cf27fd9fedb640a1bf2616ff5d5c80d99017faadc"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
+[[package]]
+name = "torchref_kernels"
+version = "0.1.0"
+dependencies = [
+ "multiversion",
+ "pyo3",
+ "rayon",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d245f478577f809a851594d02313b640fb437e0bb33866753cff937863096954"

torchref_kernels-0.1.0/Cargo.toml

@@ -0,0 +1,35 @@
+[package]
+name = "torchref_kernels"
+version = "0.1.0"
+edition = "2024"
+# `slice::as_chunks`.
+rust-version = "1.88"
+license = "MIT"
+description = "Prebuilt CPU kernels for TorchRef"
+repository = "https://github.com/HatPdotS/TorchRef"
+build = "build.rs"
+readme = "README.md"
+
+[lib]
+name = "torchref_kernels"
+# cdylib is the Python extension; rlib lets `cargo test` link the kernels.
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+multiversion = "0.9.0"
+pyo3 = { version = "0.29.3", features = ["abi3-py310"] }
+rayon = "1.12.0"
+
+[features]
+# Enabled by maturin (pyproject.toml). Off for `cargo test`, which must link libpython.
+extension-module = ["pyo3/extension-module"]
+
+[profile.release]
+lto = "fat"
+codegen-units = 1
+opt-level = 3
+
+# A forgotten `--release` on `maturin develop` must not produce a kernel that is
+# silently 10-100x slower; debug assertions stay on so dev builds still check bounds.
+[profile.dev]
+opt-level = 3

torchref_kernels-0.1.0/PKG-INFO

@@ -0,0 +1,33 @@
+Metadata-Version: 2.4
+Name: torchref-kernels
+Version: 0.1.0
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Summary: Prebuilt CPU kernels for TorchRef
+Author: HansPeterSeidel
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+
+# torchref-kernels
+
+Prebuilt CPU kernels for [TorchRef](https://github.com/HatPdotS/TorchRef), written in Rust.
+
+The wheels link neither libtorch nor a specific CPython version (abi3), so one wheel per
+platform works with every torch and Python version TorchRef supports, and nothing is
+compiled when TorchRef is imported. On x86-64 each kernel is built for several
+microarchitecture levels and the best one for the running CPU is chosen on first use.
+
+This package is an implementation detail of TorchRef and is installed with it; its entry
+points take raw tensor pointers and are not a public API.
+
+## Development
+
+```bash
+pip install -e ./kernels          # or: maturin develop --release -m kernels/Cargo.toml
+cargo test --manifest-path kernels/Cargo.toml
+```
+
+Rust changes are picked up only after rebuilding. `torchref_kernels.build_info()` reports
+the build profile, the selected CPU variant and a hash of the sources it was built from.
+

torchref_kernels-0.1.0/README.md

@@ -0,0 +1,21 @@
+# torchref-kernels
+
+Prebuilt CPU kernels for [TorchRef](https://github.com/HatPdotS/TorchRef), written in Rust.
+
+The wheels link neither libtorch nor a specific CPython version (abi3), so one wheel per
+platform works with every torch and Python version TorchRef supports, and nothing is
+compiled when TorchRef is imported. On x86-64 each kernel is built for several
+microarchitecture levels and the best one for the running CPU is chosen on first use.
+
+This package is an implementation detail of TorchRef and is installed with it; its entry
+points take raw tensor pointers and are not a public API.
+
+## Development
+
+```bash
+pip install -e ./kernels          # or: maturin develop --release -m kernels/Cargo.toml
+cargo test --manifest-path kernels/Cargo.toml
+```
+
+Rust changes are picked up only after rebuilding. `torchref_kernels.build_info()` reports
+the build profile, the selected CPU variant and a hash of the sources it was built from.

torchref_kernels-0.1.0/build.rs

@@ -0,0 +1,61 @@
+//! Embed the build profile, optimisation level and a hash of the kernel sources, so the
+//! Python side can refuse a debug build and detect a stale one.
+
+use std::fs;
+use std::path::{Path, PathBuf};
+
+/// FNV-1a 64; `torchref.utils.native.source_hash` computes the same digest in Python.
+fn fnv1a(hash: &mut u64, bytes: &[u8]) {
+    for b in bytes {
+        *hash ^= u64::from(*b);
+        *hash = hash.wrapping_mul(0x100_0000_01b3);
+    }
+}
+
+fn collect(dir: &Path, out: &mut Vec<PathBuf>) {
+    for entry in fs::read_dir(dir).expect("readable source dir") {
+        let path = entry.expect("dir entry").path();
+        if path.is_dir() {
+            collect(&path, out);
+        } else if path.extension().is_some_and(|e| e == "rs") {
+            out.push(path);
+        }
+    }
+}
+
+fn main() {
+    let root = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
+    let mut files = Vec::new();
+    collect(&root.join("src"), &mut files);
+    // Relative, '/'-separated and sorted, so the digest is the same on every OS.
+    let mut rel: Vec<(String, PathBuf)> = files
+        .into_iter()
+        .map(|p| {
+            let r = p
+                .strip_prefix(&root)
+                .unwrap()
+                .to_string_lossy()
+                .replace('\\', "/");
+            (r, p)
+        })
+        .collect();
+    rel.sort();
+    let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
+    for (name, path) in &rel {
+        fnv1a(&mut hash, name.as_bytes());
+        fnv1a(&mut hash, &[0]);
+        fnv1a(&mut hash, &fs::read(path).expect("readable source"));
+        fnv1a(&mut hash, &[0]);
+        println!("cargo:rerun-if-changed={}", path.display());
+    }
+    println!("cargo:rerun-if-changed=src");
+    println!("cargo:rustc-env=TORCHREF_KERNELS_SOURCE_HASH={hash:016x}");
+    println!(
+        "cargo:rustc-env=TORCHREF_KERNELS_PROFILE={}",
+        std::env::var("PROFILE").unwrap()
+    );
+    println!(
+        "cargo:rustc-env=TORCHREF_KERNELS_OPT_LEVEL={}",
+        std::env::var("OPT_LEVEL").unwrap()
+    );
+}

torchref_kernels-0.1.0/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["maturin>=1.7,<2"]
+build-backend = "maturin"
+
+[project]
+name = "torchref-kernels"
+dynamic = ["version"]
+description = "Prebuilt CPU kernels for TorchRef"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [{name = "HansPeterSeidel"}]
+classifiers = [
+    "Programming Language :: Rust",
+    "Programming Language :: Python :: Implementation :: CPython",
+]
+
+[tool.maturin]
+python-source = "python"
+module-name = "torchref_kernels._native"
+features = ["extension-module"]
+# `pip install ./kernels` and the wheels build in release; `maturin develop` uses the dev
+# profile, which Cargo.toml also optimises, so neither path yields a slow kernel.
+profile = "release"
+
+[tool.uv]
+# uv only rebuilds a local package when these change; without them a Rust edit is ignored.
+cache-keys = [
+    {file = "pyproject.toml"},
+    {file = "Cargo.toml"},
+    {file = "Cargo.lock"},
+    {file = "build.rs"},
+    {file = "src/**/*.rs"},
+]

torchref_kernels-0.1.0/python/torchref_kernels/__init__.py

@@ -0,0 +1,19 @@
+"""Prebuilt CPU kernels for TorchRef.
+
+Low-level entry points taking ``(data_ptr, numel)`` pairs of CPU tensors. They are not
+meant to be called directly: TorchRef's wrappers (``torchref.utils.native``) validate
+device, dtype, contiguity and shapes, own the autograd, and choose the thread count.
+Nothing here imports torch.
+"""
+
+from torchref_kernels._native import ABI_VERSION, build_info
+from torchref_kernels import _native
+
+__version__ = build_info()["version"]
+
+__all__ = ["ABI_VERSION", "build_info", "native", "__version__"]
+
+
+def native():
+    """The compiled extension module holding the kernel entry points."""
+    return _native

torchref_kernels-0.1.0/src/ffi.rs

@@ -0,0 +1,83 @@
+//! The one unsafe boundary: raw `(address, length)` pairs from Python become slices.
+//!
+//! Python passes `tensor.data_ptr()` and `tensor.numel()` for CPU, contiguous tensors of the
+//! kernel's dtype (`torchref.utils.native.buf` checks all three). Everything past this
+//! module works on safe slices, and every kernel checks slice lengths against each other
+//! before touching memory, so a wrong shape is a `ValueError`, not a segfault.
+
+use pyo3::PyResult;
+use pyo3::exceptions::PyValueError;
+use std::ptr::NonNull;
+
+/// `(data_ptr, numel)` of a CPU tensor, as passed from Python.
+pub type Buf = (usize, usize);
+
+fn checked<T>(buf: Buf, what: &str) -> PyResult<*mut T> {
+    let (addr, len) = buf;
+    if len == 0 {
+        // An empty tensor may report address 0, which is not a valid slice pointer.
+        return Ok(NonNull::<T>::dangling().as_ptr());
+    }
+    if addr == 0 {
+        return Err(PyValueError::new_err(format!("{what}: null data pointer")));
+    }
+    if addr % std::mem::align_of::<T>() != 0 {
+        return Err(PyValueError::new_err(format!(
+            "{what}: data pointer is not aligned for its dtype"
+        )));
+    }
+    Ok(addr as *mut T)
+}
+
+/// Borrow an input tensor's storage.
+///
+/// # Safety
+/// `buf` must describe `buf.1` initialised elements of `T` that stay alive and are not
+/// written by anything else while the returned slice is in use.
+pub unsafe fn input<'a, T>(buf: Buf, what: &str) -> PyResult<&'a [T]> {
+    let ptr = checked::<T>(buf, what)?;
+    Ok(unsafe { std::slice::from_raw_parts(ptr, buf.1) })
+}
+
+/// Borrow an output tensor's storage for writing.
+///
+/// # Safety
+/// As [`input`], and additionally no other slice (input or output) may overlap it.
+pub unsafe fn output<'a, T>(buf: Buf, what: &str) -> PyResult<&'a mut [T]> {
+    let ptr = checked::<T>(buf, what)?;
+    Ok(unsafe { std::slice::from_raw_parts_mut(ptr, buf.1) })
+}
+
+/// Fail with a `ValueError` unless `len == expected`.
+pub fn expect_len(len: usize, expected: usize, what: &str) -> PyResult<()> {
+    if len != expected {
+        return Err(PyValueError::new_err(format!(
+            "{what}: expected {expected} elements, got {len}"
+        )));
+    }
+    Ok(())
+}
+
+/// Address ranges of two buffers overlap (empty buffers never do).
+fn overlaps(a: Buf, b: Buf, elem: usize) -> bool {
+    a.1 != 0 && b.1 != 0 && a.0 < b.0 + b.1 * elem && b.0 < a.0 + a.1 * elem
+}
+
+/// Refuse outputs that alias each other or any input; the kernels assume exclusive access.
+pub fn no_aliasing(outputs: &[(Buf, &str)], inputs: &[Buf], elem: usize) -> PyResult<()> {
+    for (i, (o, name)) in outputs.iter().enumerate() {
+        for (p, _) in &outputs[i + 1..] {
+            if overlaps(*o, *p, elem) {
+                return Err(PyValueError::new_err(format!(
+                    "{name} aliases another output"
+                )));
+            }
+        }
+        for p in inputs {
+            if overlaps(*o, *p, elem) {
+                return Err(PyValueError::new_err(format!("{name} aliases an input")));
+            }
+        }
+    }
+    Ok(())
+}