diff --git a/Makefile b/Makefile index 9f5239d88..b19055d87 100644 --- a/Makefile +++ b/Makefile @@ -3,7 +3,7 @@ export check_dirs := kernels/src kernels/tests kernels-data/bindings/python -all: src/kernels/python_depends.json +all: kernels/src/kernels/python_depends.json kernels/src/kernels/python_depends.json: kernels-data/src/python_dependencies.json cp $< $@ diff --git a/docs/source/builder-cli.md b/docs/source/builder-cli.md index dbdba2bf4..63c7a6ef7 100644 --- a/docs/source/builder-cli.md +++ b/docs/source/builder-cli.md @@ -83,7 +83,7 @@ Initialize a new kernel project from template Default value: `Apache-2.0` * `--name ` — Name of the kernel repo (e.g. `drbh/my-kernel`) -* `--backends ` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`) +* `--backends ` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`) * `--overwrite` — Overwrite existing scaffold files (preserves other files) diff --git a/docs/source/kernel-requirements.md b/docs/source/kernel-requirements.md index 9d01fd5a4..eb1667ec5 100644 --- a/docs/source/kernel-requirements.md +++ b/docs/source/kernel-requirements.md @@ -149,7 +149,7 @@ The `backend` specifies a dictionary of the following form: ``` The backend `type` must be one of `cann`, `cpu`, `cuda`, `metal`, `neuron`, -`rocm`, or `xpu`. For CUDA and ROCm, the supported architectures must +`rocm`, `tpu`, or `xpu`. For CUDA and ROCm, the supported architectures must be specified in the `archs` field. ### Python dependencies diff --git a/examples/kernels/relu-tpu/CARD.md b/examples/kernels/relu-tpu/CARD.md new file mode 100644 index 000000000..92ecc33da --- /dev/null +++ b/examples/kernels/relu-tpu/CARD.md @@ -0,0 +1,61 @@ +--- +library_name: kernels +{% if license %}license: {{ license }} +{% endif %}--- + +This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated. + +## How to use +{% if functions %} + +```python +# make sure `kernels` is installed: `pip install -U kernels` +from kernels import get_kernel + +kernel_module = get_kernel("{{ repo_id }}", version={{ version }}) +{{ functions[0] }} = kernel_module.{{ functions[0] }} + +{{ functions[0] }}(...) +``` +{% else %} + +Usage example not available. +{% endif %} + +## Available functions +{% if functions %} +{% for func in functions %} +- `{{ func }}` +{% endfor %} +{% else %} + +Function list not available. +{% endif %} +{% if layers %} + +## Available layers +{% for layer in layers %} +- `{{ layer }}` +{% endfor %} +{% endif %} + +## Benchmarks +{% if has_benchmark %} + +Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`. +{% else %} + +No benchmark available yet. +{% endif %} +{% if upstream %} + +## Upstream + +The original source code for this kernel comes from {{ upstream }}. +{% endif %} +{% if source %} + +## Source + +The kernel-builder formatted source for this kernel is available at {{ source }}. +{% endif %} diff --git a/examples/kernels/relu-tpu/build.toml b/examples/kernels/relu-tpu/build.toml new file mode 100644 index 000000000..a51fee218 --- /dev/null +++ b/examples/kernels/relu-tpu/build.toml @@ -0,0 +1,13 @@ +[general] +name = "relu-tpu" +version = 1 +edition = 5 +license = "Apache-2.0" +backends = ["tpu"] + +[general.hub] +repo-id = "kernels-test/relu-tpu" + +[torch-noarch] + +[kernel] diff --git a/examples/kernels/relu-tpu/flake.nix b/examples/kernels/relu-tpu/flake.nix new file mode 100644 index 000000000..a30872561 --- /dev/null +++ b/examples/kernels/relu-tpu/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for TPU ReLU kernel"; + + inputs = { + kernel-builder.url = "path:../../.."; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py new file mode 100644 index 000000000..677bca5ff --- /dev/null +++ b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py @@ -0,0 +1,14 @@ +import jax +import jax.numpy as jnp +from torch_tpu._internal.pallas import jax_op + + +def _jax_relu(x: jax.Array) -> jax.Array: + return jnp.maximum(x, 0) + + +relu = jax_op("relu_tpu::relu", _jax_relu) + +from . import layers # noqa: E402 + +__all__ = ["layers", "relu"] diff --git a/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py b/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py new file mode 100644 index 000000000..7800a2ef7 --- /dev/null +++ b/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py @@ -0,0 +1,9 @@ +import torch +import torch.nn as nn + +from .. import relu + + +class ReLU(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return relu(x) diff --git a/kernel-builder/src/init.rs b/kernel-builder/src/init.rs index 8ce37f48a..393de2dd9 100644 --- a/kernel-builder/src/init.rs +++ b/kernel-builder/src/init.rs @@ -44,7 +44,7 @@ pub struct InitArgs { #[arg(long, value_name = "OWNER/REPO")] pub name: Option, - /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`). + /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`). #[arg(long, num_args = 1.., default_values_t = default_init_backends())] pub backends: Vec, diff --git a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py index 75fd8e086..e8d30b5ff 100644 --- a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py +++ b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py @@ -4,7 +4,13 @@ def get_backend() -> str: """Detect the backend by inspecting torch.""" import torch - if hasattr(torch, "neuron"): + if hasattr(torch.backends, "tpu"): + # torch_tpu sets torch.backends.tpu when it is imported (via + # torch's device-backend autoload), regardless of whether TPU + # hardware is present — analogous to torch.version.cuda being + # set on CUDA builds without a GPU. + return "tpu" + elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron # extension can be loaded into e.g. CUDA Torch builds. return "neuron" diff --git a/kernels-data/bindings/python/kernels_data.pyi b/kernels-data/bindings/python/kernels_data.pyi index 20e29619f..43b7af6a6 100644 --- a/kernels-data/bindings/python/kernels_data.pyi +++ b/kernels-data/bindings/python/kernels_data.pyi @@ -32,6 +32,7 @@ class Backend(Enum): Metal = "Metal" Neuron = "Neuron" ROCm = "ROCm" + TPU = "TPU" XPU = "XPU" @staticmethod @@ -40,7 +41,7 @@ class Backend(Enum): Args: s: One of `"cann"`, `"cpu"`, `"cuda"`, `"metal"`, - `"neuron"`, `"rocm"`, `"xpu"`. + `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`. Raises: ValueError: If the backend name is unknown. diff --git a/kernels-data/bindings/python/src/lib.rs b/kernels-data/bindings/python/src/lib.rs index 9aa29f06d..3cfcae33e 100644 --- a/kernels-data/bindings/python/src/lib.rs +++ b/kernels-data/bindings/python/src/lib.rs @@ -88,6 +88,8 @@ enum PyBackend { Neuron, #[pyo3(name = "ROCm")] Rocm, + #[pyo3(name = "TPU")] + Tpu, #[pyo3(name = "XPU")] Xpu, } @@ -101,6 +103,7 @@ impl From for PyBackend { Backend::Metal => PyBackend::Metal, Backend::Neuron => PyBackend::Neuron, Backend::Rocm => PyBackend::Rocm, + Backend::Tpu => PyBackend::Tpu, Backend::Xpu => PyBackend::Xpu, } } @@ -115,6 +118,7 @@ impl From for Backend { PyBackend::Metal => Backend::Metal, PyBackend::Neuron => Backend::Neuron, PyBackend::Rocm => Backend::Rocm, + PyBackend::Tpu => Backend::Tpu, PyBackend::Xpu => Backend::Xpu, } } @@ -123,7 +127,7 @@ impl From for Backend { #[pymethods] impl PyBackend { /// Parse a backend name (`"cann"`, `"cpu"`, `"cuda"`, `"metal"`, - /// `"neuron"`, `"rocm"`, `"xpu"`). + /// `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`). #[staticmethod] #[pyo3(name = "from_str")] fn py_from_str(s: &str) -> PyResult { @@ -144,6 +148,7 @@ impl PyBackend { PyBackend::Metal => "Metal", PyBackend::Neuron => "Neuron", PyBackend::Rocm => "ROCm", + PyBackend::Tpu => "TPU", PyBackend::Xpu => "XPU", }; format!("Backend.{variant}") diff --git a/kernels-data/bindings/python/tests/test_kernels_data.py b/kernels-data/bindings/python/tests/test_kernels_data.py index 29bb52fe5..87f9d1119 100644 --- a/kernels-data/bindings/python/tests/test_kernels_data.py +++ b/kernels-data/bindings/python/tests/test_kernels_data.py @@ -68,7 +68,7 @@ def test_backend_hash(): def test_backend_unknown(): with pytest.raises(ValueError): - Backend.from_str("tpu") + Backend.from_str("dsp") def test_backend_all_variants_and_casing(): @@ -80,8 +80,11 @@ def test_backend_all_variants_and_casing(): assert repr(Backend.ROCm) == "Backend.ROCm" assert repr(Backend.XPU) == "Backend.XPU" assert repr(Backend.CANN) == "Backend.CANN" + assert str(Backend.TPU) == "tpu" + assert repr(Backend.TPU) == "Backend.TPU" assert Backend.from_str("cann") == Backend.CANN assert Backend.from_str("ROCM") == Backend.ROCm + assert Backend.from_str("TPU") == Backend.TPU assert Backend.from_str("metal") == Backend.Metal diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs index bc7855f71..93ba480b0 100644 --- a/kernels-data/src/config/mod.rs +++ b/kernels-data/src/config/mod.rs @@ -113,6 +113,7 @@ pub struct General { pub cuda: Option, pub neuron: Option, + pub tpu: Option, pub xpu: Option, } @@ -144,6 +145,10 @@ impl General { .cuda .as_ref() .and_then(|cuda| cuda.python_depends.as_ref()), + Backend::Tpu => self + .tpu + .as_ref() + .and_then(|tpu| tpu.python_depends.as_ref()), Backend::Xpu => self .xpu .as_ref() @@ -191,6 +196,10 @@ pub struct NeuronGeneral { pub python_depends: Option>, } +pub struct TpuGeneral { + pub python_depends: Option>, +} + pub struct Hub { pub repo_id: Option, pub branch: Option, @@ -379,11 +388,12 @@ pub enum Backend { Metal, Neuron, Rocm, + Tpu, Xpu, } impl Backend { - pub const fn all() -> [Backend; 7] { + pub const fn all() -> [Backend; 8] { [ Backend::Cann, Backend::Cpu, @@ -391,6 +401,7 @@ impl Backend { Backend::Metal, Backend::Neuron, Backend::Rocm, + Backend::Tpu, Backend::Xpu, ] } @@ -403,6 +414,7 @@ impl Backend { Backend::Metal => "metal", Backend::Neuron => "neuron", Backend::Rocm => "rocm", + Backend::Tpu => "tpu", Backend::Xpu => "xpu", } } @@ -417,6 +429,7 @@ impl Display for Backend { Backend::Metal => write!(f, "metal"), Backend::Neuron => write!(f, "neuron"), Backend::Rocm => write!(f, "rocm"), + Backend::Tpu => write!(f, "tpu"), Backend::Xpu => write!(f, "xpu"), } } @@ -433,6 +446,7 @@ impl FromStr for Backend { "metal" => Ok(Backend::Metal), "neuron" => Ok(Backend::Neuron), "rocm" => Ok(Backend::Rocm), + "tpu" => Ok(Backend::Tpu), "xpu" => Ok(Backend::Xpu), _ => Err(format!("Unknown backend: {s}")), } @@ -444,3 +458,53 @@ pub enum ConfigError { #[error("Cannot migrate configuration: {reason:?}")] Migration { reason: String }, } + +#[cfg(test)] +mod tests { + use super::*; + + fn general_with_cuda_deps() -> General { + General { + name: KernelName::new("test-kernel").unwrap(), + version: 1, + license: "apache-2.0".to_string(), + upstream: None, + source: None, + backends: vec![Backend::Tpu], + hub: None, + python_depends: None, + cuda: Some(CudaGeneral { + minver: None, + maxver: None, + python_depends: Some(vec!["nvidia-cutlass-dsl".to_string()]), + }), + neuron: None, + tpu: None, + xpu: None, + } + } + + #[test] + fn backend_python_depends_resolves_cuda() { + let general = general_with_cuda_deps(); + let deps = general + .backend_python_depends(Backend::Cuda) + .map(|dep| dep.map(|(name, _)| name.to_string())) + .collect::>>() + .unwrap(); + + assert_eq!(deps, vec!["nvidia-cutlass-dsl".to_string()]); + } + + #[test] + fn backend_python_depends_empty_for_backend_without_deps() { + let general = general_with_cuda_deps(); + + assert!( + general + .backend_python_depends(Backend::Cpu) + .next() + .is_none() + ); + } +} diff --git a/kernels-data/src/config/v2.rs b/kernels-data/src/config/v2.rs index eed543ba6..f79204a8b 100644 --- a/kernels-data/src/config/v2.rs +++ b/kernels-data/src/config/v2.rs @@ -179,6 +179,7 @@ impl General { hub: general.hub.map(Into::into), neuron: None, python_depends: None, + tpu: None, xpu: None, } } diff --git a/kernels-data/src/config/v3.rs b/kernels-data/src/config/v3.rs index 590d16e6a..b9aa400a6 100644 --- a/kernels-data/src/config/v3.rs +++ b/kernels-data/src/config/v3.rs @@ -155,6 +155,7 @@ pub enum Backend { Metal, Neuron, Rocm, + Tpu, Xpu, } @@ -205,6 +206,7 @@ impl TryFrom for super::General { hub: general.hub.map(Into::into), neuron: general.neuron.map(Into::into), python_depends: general.python_depends, + tpu: None, xpu: general.xpu.map(Into::into), }) } @@ -279,6 +281,7 @@ impl From for super::Backend { Backend::Metal => super::Backend::Metal, Backend::Neuron => super::Backend::Neuron, Backend::Rocm => super::Backend::Rocm, + Backend::Tpu => super::Backend::Tpu, Backend::Xpu => super::Backend::Xpu, } } diff --git a/kernels-data/src/config/v4.rs b/kernels-data/src/config/v4.rs index a88db69c9..1624cf983 100644 --- a/kernels-data/src/config/v4.rs +++ b/kernels-data/src/config/v4.rs @@ -171,6 +171,7 @@ pub enum Backend { Metal, Neuron, Rocm, + Tpu, Xpu, } @@ -203,6 +204,7 @@ impl From for super::General { hub: general.hub.map(Into::into), neuron: general.neuron.map(Into::into), python_depends: general.python_depends, + tpu: None, xpu: general.xpu.map(Into::into), } } @@ -300,6 +302,7 @@ impl From for super::Backend { Backend::Metal => super::Backend::Metal, Backend::Neuron => super::Backend::Neuron, Backend::Rocm => super::Backend::Rocm, + Backend::Tpu => super::Backend::Tpu, Backend::Xpu => super::Backend::Xpu, } } diff --git a/kernels-data/src/config/v5.rs b/kernels-data/src/config/v5.rs index 0154eeef6..3e8a99960 100644 --- a/kernels-data/src/config/v5.rs +++ b/kernels-data/src/config/v5.rs @@ -62,6 +62,8 @@ pub struct General { pub python_depends: Option>, + pub tpu: Option, + pub xpu: Option, } @@ -79,6 +81,12 @@ pub struct NeuronGeneral { pub python_depends: Option>, } +#[derive(Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields, rename_all = "kebab-case")] +pub struct TpuGeneral { + pub python_depends: Option>, +} + #[derive(Debug, Deserialize, Serialize)] #[serde(deny_unknown_fields, rename_all = "kebab-case")] pub struct XpuGeneral { @@ -185,6 +193,7 @@ pub enum Backend { Metal, Neuron, Rocm, + Tpu, Xpu, } @@ -217,6 +226,7 @@ impl From for super::General { hub: general.hub.map(Into::into), neuron: general.neuron.map(Into::into), python_depends: general.python_depends, + tpu: general.tpu.map(Into::into), xpu: general.xpu.map(Into::into), } } @@ -252,6 +262,14 @@ impl From for super::NeuronGeneral { } } +impl From for super::TpuGeneral { + fn from(tpu: TpuGeneral) -> Self { + Self { + python_depends: tpu.python_depends, + } + } +} + impl From for super::XpuGeneral { fn from(xpu: XpuGeneral) -> Self { Self { @@ -313,6 +331,7 @@ impl From for super::Backend { Backend::Metal => super::Backend::Metal, Backend::Neuron => super::Backend::Neuron, Backend::Rocm => super::Backend::Rocm, + Backend::Tpu => super::Backend::Tpu, Backend::Xpu => super::Backend::Xpu, } } @@ -420,6 +439,7 @@ impl From for General { hub: general.hub.map(Into::into), neuron: general.neuron.map(Into::into), python_depends: general.python_depends, + tpu: general.tpu.map(Into::into), xpu: general.xpu.map(Into::into), } } @@ -455,6 +475,14 @@ impl From for NeuronGeneral { } } +impl From for TpuGeneral { + fn from(tpu: super::TpuGeneral) -> Self { + Self { + python_depends: tpu.python_depends, + } + } +} + impl From for XpuGeneral { fn from(xpu: super::XpuGeneral) -> Self { Self { @@ -516,6 +544,7 @@ impl From for Backend { super::Backend::Metal => Backend::Metal, super::Backend::Neuron => Backend::Neuron, super::Backend::Rocm => Backend::Rocm, + super::Backend::Tpu => Backend::Tpu, super::Backend::Xpu => Backend::Xpu, } } diff --git a/kernels-data/src/metadata.rs b/kernels-data/src/metadata.rs index d895f7fbb..3da04dc46 100644 --- a/kernels-data/src/metadata.rs +++ b/kernels-data/src/metadata.rs @@ -159,6 +159,7 @@ mod tests { python_depends: None, cuda: None, neuron: None, + tpu: None, xpu: None, }, kernels: HashMap::new(), @@ -214,6 +215,7 @@ mod tests { python_depends: None, cuda: None, neuron: None, + tpu: None, xpu: None, }, kernels: HashMap::new(), diff --git a/kernels/pytest.ini b/kernels/pytest.ini index a94211637..22164adcc 100644 --- a/kernels/pytest.ini +++ b/kernels/pytest.ini @@ -5,4 +5,5 @@ markers = darwin_only: marks tests that should only run on macOS xpu_only: marks tests that should only run on hosts with Intel XPUs npu_only: marks tests that should only run on Ascend NPUs + tpu_only: marks tests that should only run on hosts with TPU devices token: enable tests that require a write token diff --git a/kernels/src/kernels/backends.py b/kernels/src/kernels/backends.py index 5dff532a1..c5b829209 100644 --- a/kernels/src/kernels/backends.py +++ b/kernels/src/kernels/backends.py @@ -149,6 +149,24 @@ def parse(s: str) -> "ROCm": return ROCm(version=Version(f"{m.group(1)}.{m.group(2)}")) +@strict +@dataclass(unsafe_hash=True) +class TPU: + @property + def name(self) -> str: + return "tpu" + + @property + def variant_str(self) -> str: + return "tpu" + + @staticmethod + def parse(s: str) -> "TPU": + if s != "tpu": + raise ValueError(f"Invalid TPU variant string: {s!r}") + return TPU() + + @dataclass(unsafe_hash=True) class XPU: _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"xpu(\d+)(\d+)") @@ -179,6 +197,8 @@ def parse_backend(s: str) -> Backend: return Metal.parse(s) elif s == "neuron": return Neuron.parse(s) + elif s == "tpu": + return TPU.parse(s) elif s.startswith("cu"): return CUDA.parse(s) elif s.startswith("rocm"): @@ -195,7 +215,14 @@ def _backend() -> Backend: if has_torch: import torch - if hasattr(torch, "neuron"): + if hasattr(torch.backends, "tpu"): + # torch_tpu sets torch.backends.tpu when it is imported (via + # torch's device-backend autoload), regardless of whether TPU + # hardware is present — analogous to torch.version.cuda being + # set on CUDA builds without a GPU. The hardware-gated + # torch.tpu namespace only appears on hosts with TPU devices. + return TPU() + elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron # extension can be loaded into e.g. CUDA Torch builds. return Neuron() diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py index 5294a5f5d..60d221d6a 100644 --- a/kernels/src/kernels/layer/kernelize.py +++ b/kernels/src/kernels/layer/kernelize.py @@ -192,7 +192,7 @@ def kernelize( `Mode.TRAINING | Mode.TORCH_COMPILE` kernelizes the model for training with `torch.compile`. device (`Union[str, torch.device]`, *optional*): - The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "xpu". + The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "tpu", "xpu". The device type will be inferred from the model parameters when not provided. use_fallback (`bool`, *optional*, defaults to `True`): Whether to use the original forward method of modules when no compatible kernel could be found. @@ -276,7 +276,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: def _validate_device_type(device_type: str) -> None: """Validate that the device type is supported.""" - supported_devices = {"cpu", "cuda", "mps", "neuron", "npu", "rocm", "xpu"} + supported_devices = {"cpu", "cuda", "mps", "neuron", "npu", "rocm", "tpu", "xpu"} if device_type not in supported_devices: raise ValueError( f"Unsupported device type '{device_type}'. Supported device types are: {', '.join(sorted(supported_devices))}" diff --git a/kernels/src/kernels/layer/repos.py b/kernels/src/kernels/layer/repos.py index fd02707c2..c53c833bb 100644 --- a/kernels/src/kernels/layer/repos.py +++ b/kernels/src/kernels/layer/repos.py @@ -37,6 +37,8 @@ def create_repo(device: Device) -> "DeviceRepos": return _NPURepos() elif device.type == "neuron": return _NeuronRepos() + elif device.type == "tpu": + return _TPURepos() else: raise ValueError(f"Unknown device type: {device.type}") @@ -114,6 +116,26 @@ def insert(self, device: Device, repos: dict[Mode, RepositoryProtocol]): self._repos = repos +class _TPURepos(DeviceRepos): + _repos: dict[Mode, RepositoryProtocol] + + def __init__(self): + super().__init__() + self._repos = {} + + @property + def repos( + self, + ) -> dict[Mode, RepositoryProtocol] | None: + return self._repos + + def insert(self, device: Device, repos: dict[Mode, RepositoryProtocol]): + if device.type != "tpu": + raise ValueError(f"Device type must be 'tpu', got {device.type}") + + self._repos = repos + + class _NPURepos(DeviceRepos): _repos: dict[Mode, RepositoryProtocol] diff --git a/kernels/tests/conftest.py b/kernels/tests/conftest.py index a9cb3c01c..513261ffd 100644 --- a/kernels/tests/conftest.py +++ b/kernels/tests/conftest.py @@ -34,6 +34,8 @@ has_npu = torch is not None and _get_torch_privateuse_backend_name() == "npu" +has_tpu = torch is not None and hasattr(torch, "tpu") and torch.tpu.device_count() > 0 + has_jax = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("jax_tvm_ffi") is not None @@ -53,6 +55,8 @@ def device(): return "xpu" elif has_npu: return "npu" + elif has_tpu: + return "tpu" return "cpu" @@ -74,5 +78,7 @@ def pytest_runtest_setup(item): pytest.skip("skipping XPU-only test on host without XPU") if "npu_only" in item.keywords and not has_npu: pytest.skip("skipping NPU-only test on host without NPU") + if "tpu_only" in item.keywords and not has_tpu: + pytest.skip("skipping TPU-only test on host without TPU") if "token" in item.keywords and not item.config.getoption("--token"): pytest.skip("need --token option to run this test") diff --git a/kernels/tests/test_basic.py b/kernels/tests/test_basic.py index 1460f9454..5c9ce8dbd 100644 --- a/kernels/tests/test_basic.py +++ b/kernels/tests/test_basic.py @@ -229,6 +229,13 @@ def test_neuron(): torch.testing.assert_close(relu.relu(x), x.relu()) +@pytest.mark.tpu_only +def test_tpu(): + relu = get_kernel("kernels-test/relu-tpu", version=1) + x = torch.randn((16, 16), device="tpu") + torch.testing.assert_close(relu.relu(x), x.relu()) + + def test_trust_remote_code_blocks_untrusted_org(): """Kernels from untrusted orgs should be rejected by default.""" with pytest.raises(ValueError, match=r"not from a trusted publisher"): diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py index e67abfc28..790e64458 100644 --- a/kernels/tests/test_layer.py +++ b/kernels/tests/test_layer.py @@ -255,6 +255,47 @@ def forward(self, x): assert smol.relu.n_calls == 1 +@pytest.mark.tpu_only +def test_hub_forward_tpu(): + torch.manual_seed(0) + + mapping = {"ReLU": {"tpu": LayerRepository(repo_id="kernels-test/relu-tpu", version=1, layer_name="ReLU")}} + + relu = ReLU() + X = torch.randn((16, 16), device="tpu") + Y = relu(X) + + with use_kernel_mapping(mapping): + relu_with_kernel = kernelize(ReLUWithKernel(), device="tpu", mode=Mode.INFERENCE) + Y_kernel = relu_with_kernel(X) + + torch.testing.assert_close(Y_kernel, Y) + + assert relu.n_calls == 1 + assert relu_with_kernel.n_calls == 0 + + # Check that the device type can be determined automatically. + class SMOL(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(16, 16) + self.relu = ReLUWithKernel() + + def forward(self, x): + return self.relu(self.linear(x)) + + smol = SMOL().to("tpu") + + Y = smol(X) + + with use_kernel_mapping(mapping): + smol = kernelize(smol, mode=Mode.INFERENCE) + Y_kernel = smol(X) + + torch.testing.assert_close(Y, Y_kernel) + assert smol.relu.n_calls == 1 + + @pytest.mark.rocm_only def test_hub_forward_rocm(): torch.manual_seed(0) diff --git a/kernels/tests/test_variants.py b/kernels/tests/test_variants.py index e9ed9d6ec..427735549 100644 --- a/kernels/tests/test_variants.py +++ b/kernels/tests/test_variants.py @@ -75,6 +75,7 @@ "torch-metal", "torch-neuron", "torch-rocm", + "torch-tpu", "torch-xpu", "torch-npu", "torch-universal", diff --git a/nix-builder/README.md b/nix-builder/README.md index 7b31003ba..e59b7119e 100644 --- a/nix-builder/README.md +++ b/nix-builder/README.md @@ -54,8 +54,9 @@ The compiled kernel will then be available in the local `build/` directory. | Metal | ✓ | ✓ | ✗ | 2 | | Huawei NPU | ✓ | ✗ | ✗ | 3 | | Neuron | ✓ | x | x | 3 | +| TPU | ✓ | ✗ | ✗ | 3 | -**Warning:** Neuron support is experimental and currently requires pre-release packages. +**Warning:** Neuron and TPU support are experimental and packages are not public yet. # 📚 Documentation diff --git a/nix-builder/lib/build-variants.nix b/nix-builder/lib/build-variants.nix index f7dfb7776..8948911f1 100644 --- a/nix-builder/lib/build-variants.nix +++ b/nix-builder/lib/build-variants.nix @@ -26,7 +26,10 @@ rec { throw "No compute framework set in Torch version"; in version: - if version.system == "aarch64-darwin" then + # TPU builds are noarch: one variant, independent of Torch version and system. + if backend version == "tpu" then + "torch-tpu" + else if version.system == "aarch64-darwin" then "torch${flattenVersion version.torchVersion}-${computeString version}-${version.system}" else "torch${flattenVersion version.torchVersion}-cxx11-${computeString version}-${version.system}"; diff --git a/nix-builder/lib/gen-flake-outputs.nix b/nix-builder/lib/gen-flake-outputs.nix index 99fba216e..9e553cb3d 100644 --- a/nix-builder/lib/gen-flake-outputs.nix +++ b/nix-builder/lib/gen-flake-outputs.nix @@ -58,6 +58,8 @@ let "xpu" else if buildConfig.metal or false then "metal" + else if buildConfig.tpu or false then + "tpu" else throw "Cannot determine framework for build set"; @@ -249,6 +251,7 @@ in ++ (headOrEmpty (setsWithFramework "cuda")) ++ (headOrEmpty (setsWithFramework "metal")) ++ (headOrEmpty (setsWithFramework "rocm")) + ++ (headOrEmpty (setsWithFramework "tpu")) ++ (headOrEmpty (setsWithFramework "xpu")); in build.mkExtensionBundle { diff --git a/nix-builder/lib/kernel-config.nix b/nix-builder/lib/kernel-config.nix index 0b274135f..4770404a6 100644 --- a/nix-builder/lib/kernel-config.nix +++ b/nix-builder/lib/kernel-config.nix @@ -48,6 +48,7 @@ in cuda = false; metal = false; rocm = false; + tpu = false; xpu = false; }; in @@ -63,6 +64,7 @@ in cuda = false; metal = false; rocm = false; + tpu = false; xpu = false; }; in diff --git a/nix-builder/lib/mk-build-set.nix b/nix-builder/lib/mk-build-set.nix index 920d51cc1..74a4ff573 100644 --- a/nix-builder/lib/mk-build-set.nix +++ b/nix-builder/lib/mk-build-set.nix @@ -82,6 +82,15 @@ let rocmSupport = true; }; + tpu = { + # torch_tpu is Apache-2.0, but libtpu's wheel METADATA declares + # its license as "Google Cloud Platform Terms of Service" + # (unfree), so the tpu buildSet needs allowUnfree just like the + # cuda/rocm/xpu sets. See pkgs/python-modules/libtpu/default.nix. + allowUnfree = true; + tpuSupport = true; + }; + xpu = { allowUnfree = true; xpuSupport = true; @@ -102,6 +111,7 @@ buildConfig@{ ptxasVersion ? cudaVersion, metal ? false, rocmVersion ? null, + tpu ? false, xpuVersion ? null, torchVersion, system, @@ -117,6 +127,8 @@ let [ (overlayForCudaVersion cudaVersion ptxasVersion) ] else if buildConfig.backend == "rocm" then [ (overlayForRocmVersion rocmVersion) ] + else if buildConfig.backend == "tpu" then + [ ] else if buildConfig.backend == "metal" then [ ] else if buildConfig.backend == "xpu" then diff --git a/nix-builder/lib/torch-version-utils.nix b/nix-builder/lib/torch-version-utils.nix index 5b43482cb..b7009de5f 100644 --- a/nix-builder/lib/torch-version-utils.nix +++ b/nix-builder/lib/torch-version-utils.nix @@ -4,6 +4,7 @@ let isCuda = version: version ? cudaVersion; isMetal = version: version.metal or false; isRocm = version: version ? rocmVersion; + isTpu = version: version.tpu or false; isXpu = version: version ? xpuVersion; in @@ -25,6 +26,8 @@ rec { "metal" else if isRocm version then "rocm" + else if isTpu version then + "tpu" else if isXpu version then "xpu" else diff --git a/nix-builder/lib/variants/torch.nix b/nix-builder/lib/variants/torch.nix index 406a75fd4..7b9469277 100644 --- a/nix-builder/lib/variants/torch.nix +++ b/nix-builder/lib/variants/torch.nix @@ -13,6 +13,8 @@ let "metal" else if buildConfig.backend == "rocm" then "rocm${flattenVersion (lib.versions.majorMinor buildConfig.rocmVersion)}" + else if buildConfig.backend == "tpu" then + "tpu" else if buildConfig.backend == "xpu" then "xpu${flattenVersion (lib.versions.majorMinor buildConfig.xpuVersion)}" else diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix index 547074c76..bb49d2fa0 100644 --- a/nix-builder/overlay.nix +++ b/nix-builder/overlay.nix @@ -161,6 +161,8 @@ final: prev: jax-tvm-ffi = python-self.callPackage ./pkgs/python-modules/jax-tvm-ffi { }; + libtpu = python-self.callPackage ./pkgs/python-modules/libtpu { }; + jupyter-server = python-super.jupyter-server.overrideAttrs ( _: prevAttrs: { # Gets stuck sometimes, already tested in nixpkgs. diff --git a/nix-builder/pkgs/get-kernel-check/default.nix b/nix-builder/pkgs/get-kernel-check/default.nix index 51e492832..dc1d8eb98 100644 --- a/nix-builder/pkgs/get-kernel-check/default.nix +++ b/nix-builder/pkgs/get-kernel-check/default.nix @@ -1,9 +1,22 @@ -{ makeSetupHook, python3 }: +{ + config, + lib, + makeSetupHook, + stdenv, + proot, + python3, +}: + +let + useFakeSys = config.tpuSupport or false; +in makeSetupHook { name = "get-kernel-check-hook"; substitutions = { python3 = "${python3}/bin/python"; kernels = "${with python3.pkgs; makePythonPath [ kernels ]}"; + proot = lib.optionalString useFakeSys "${proot}/bin/proot"; + useFakeSys = lib.optionalString useFakeSys "1"; }; } ./get-kernel-check-hook.sh diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh index 23d4501a7..5d229f8ef 100755 --- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh +++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh @@ -21,11 +21,27 @@ _getKernelCheckHook() { export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}" fi - # Some kernels want to write stuff (especially when they use Triton). HOME=$(mktemp -d -t test.XXXXXX) || exit 1 trap "rm -rf '$HOME'" EXIT + # Prepare fake /sys for tcmalloc. Without this path, tcmalloc will crash: + # + # https://github.com/google/tcmalloc/issues/245 + # + # tcmalloc is used by the TPU libraries. + local prootCmd="" + if [[ -n "@useFakeSys@" ]]; then + echo "Faking /sys for tcmalloc" + local fakeSys + fakeSys="$(mktemp -d)" + trap 'rm -rf -- "${fakeSys}"' EXIT + mkdir -p "${fakeSys}/devices/system/cpu" + echo "0-1" > "${fakeSys}/devices/system/cpu/possible" + prootCmd="@proot@ -b ${fakeSys}:/sys" + fi + PYTHONPATH="@kernels@" \ + ${prootCmd} \ @python3@ -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${out}'))" } diff --git a/nix-builder/pkgs/python-modules/hooks/default.nix b/nix-builder/pkgs/python-modules/hooks/default.nix index 246fc3682..f4eae33ef 100644 --- a/nix-builder/pkgs/python-modules/hooks/default.nix +++ b/nix-builder/pkgs/python-modules/hooks/default.nix @@ -4,8 +4,24 @@ let inherit (python) pythonOnBuildForHost; pythonInterpreter = pythonOnBuildForHost.interpreter; pythonSitePackages = python.sitePackages; + pythonCheckInterpreter = python.interpreter; + useFakeSys = pkgs.config.tpuSupport or false; in { + # Ideally we'd just call this pythonImportsCheckHook, but it would cause + # a rebuild of most nixpkgs dependencies. + pythonKernelImportsCheckHook = callPackage ( + { makePythonHook }: + makePythonHook { + name = "python-kernel-imports-check-hook.sh"; + substitutions = { + inherit pythonCheckInterpreter pythonSitePackages; + proot = lib.optionalString useFakeSys "${pkgs.proot}/bin/proot"; + useFakeSys = lib.optionalString useFakeSys "1"; + }; + } ./python-kernel-imports-check-hook.sh + ) { }; + pythonRelaxWheelDepsHook = callPackage ( { makePythonHook, wheel }: makePythonHook { diff --git a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh new file mode 100644 index 000000000..49e41ae86 --- /dev/null +++ b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh @@ -0,0 +1,49 @@ +# shellcheck shell=bash + +# This is the imports check hook from nixpkgs, modified to fake a /sys path +# required for tcmalloc. Without this path, tcmalloc will crash: +# +# https://github.com/google/tcmalloc/issues/245 + +# Setup hook for checking whether Python imports succeed +echo "Sourcing python-kernel-imports-check-hook.sh" + +pythonKernelImportsCheckPhase() { + echo "Executing pythonKernelImportsCheckPhase" + + if [[ -n "${pythonKernelImportsCheck[*]-}" ]]; then + echo "Check whether the following modules can be imported: ${pythonKernelImportsCheck[*]}" + # shellcheck disable=SC2154 + pythonKernelImportsCheckOutput="$out" + if [[ -n "${python-}" ]]; then + echo "Using python specific output \$python for imports check" + pythonKernelImportsCheckOutput=$python + fi + export PYTHONPATH="$pythonKernelImportsCheckOutput/@pythonSitePackages@:$PYTHONPATH" + + local prootCmd="" + if [[ -n "@useFakeSys@" ]]; then + echo "Faking /sys for tcmalloc" + local fakeSys + fakeSys="$(mktemp -d)" + trap 'rm -rf -- "${fakeSys}"' EXIT + mkdir -p "${fakeSys}/devices/system/cpu" + echo "0-1" > "${fakeSys}/devices/system/cpu/possible" + prootCmd="@proot@ -b ${fakeSys}:/sys" + fi + + # Python modules and namespaces names are Python identifiers, which must not contain spaces. + # See https://docs.python.org/3/reference/lexical_analysis.html + # shellcheck disable=SC2048,SC2086 + ( + cd "$pythonKernelImportsCheckOutput" && + ${prootCmd} \ + @pythonCheckInterpreter@ -c 'import sys; import importlib; list(map(lambda mod: importlib.import_module(mod), sys.argv[1:]))' ${pythonKernelImportsCheck[*]} + ) + fi +} + +if [[ -z "${dontUsePythonImportsCheck-}" ]]; then + echo "Using pythonKernelImportsCheckPhase" + appendToVar preDistPhases pythonKernelImportsCheckPhase +fi diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix new file mode 100644 index 000000000..e3cf3c198 --- /dev/null +++ b/nix-builder/pkgs/python-modules/libtpu/default.nix @@ -0,0 +1,64 @@ +{ + lib, + buildPythonPackage, + requireFile, + python, +}: + +# libtpu is served from Google's gated Artifact Registry, which requires +# an OAuth2 bearer token. That token can't be handed to a build sandbox +# without making the derivation impure, so instead the wheel is +# prefetched out-of-band and handed to `requireFile`, which only checks +# that a file matching `name` + `hash` already exists in the store: +# +# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ +# nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43 +# +# The script prints the `hash` line to paste in below. To bump the +# version, update `version`, then re-run the script with the new +# version. +# +# libtpu ships the `libtpu/libtpu.so` runtime that both torch_tpu and +# jaxlib dlopen. Pinned to 0.0.x to match the jax 0.10.x / torch_tpu +# 0.1.x stack (torch_tpu requires libtpu>=0.0.40; jax's "tpu" extra +# pins libtpu==0.0.43.*), so bump it together with jaxlib/torch_tpu. +# +# NOTE on license: libtpu's wheel METADATA declares its license as +# "Google Cloud Platform Terms of Service" — i.e. unfree, NOT +# Apache-2.0 (unlike torch_tpu). The tpu buildSet therefore sets +# allowUnfree = true (see lib/mk-build-set.nix). + +let + # The wheel ships per-CPython-ABI builds (cp311..cp314); pick the tag + # matching the python this package set is built for. The hash below is + # for cp313 (the nixpkgs default python); if either moves, recompute + # via the prefetch script above. + abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}"; +in +buildPythonPackage rec { + pname = "libtpu"; + version = "0.0.44.1"; + format = "wheel"; + + src = requireFile { + name = "libtpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl"; + hash = "sha256-qd34JzD0TpphtqLj67/HqYgVg/bXT3gF6E+pYQyaCsA="; # cp313 + message = '' + libtpu is served from a gated Google Artifact Registry and cannot + be fetched by a pure Nix build. Fetch and register it with: + + GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ + nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu ${version} ${abi} + ''; + }; + + pythonImportsCheck = [ "libtpu" ]; + doInstallCheck = false; # requires actual /dev/accel + + meta = with lib; { + description = "TPU runtime shared library, dlopened by torch_tpu and jaxlib"; + homepage = "https://cloud.google.com/tpu"; + license = licenses.unfree; # "Google Cloud Platform Terms of Service" + platforms = [ "x86_64-linux" ]; + }; +} diff --git a/nix-builder/pkgs/python-modules/torch/binary/generic.nix b/nix-builder/pkgs/python-modules/torch/binary/generic.nix index c66c6cb97..1b8183511 100644 --- a/nix-builder/pkgs/python-modules/torch/binary/generic.nix +++ b/nix-builder/pkgs/python-modules/torch/binary/generic.nix @@ -7,10 +7,12 @@ buildPythonPackage, fetchFromGitHub, fetchurl, + requireFile, cudaSupport ? config.cudaSupport, metalSupport ? config.metalSupport, rocmSupport ? config.rocmSupport, + tpuSupport ? (config.tpuSupport or false), tritonSupport ? (!stdenv.hostPlatform.isDarwin), xpuSupport ? (config.xpuSupport or false), @@ -18,6 +20,7 @@ autoAddDriverRunpath, autoPatchelfHook, python, + pythonKernelImportsCheckHook, pythonRelaxWheelDepsHook, pythonWheelDepsCheckHook, @@ -31,15 +34,21 @@ # Python dependencies cuda-bindings, filelock, + frozendict, fsspec, + immutabledict, + jax, jinja2, + libtpu, networkx, numpy, + portpicker, pyyaml, pyzes, requests, setuptools, sympy, + tensorboard, triton, triton-cuda, triton-rocm, @@ -122,6 +131,24 @@ let ''; }; + # Once we require multiple versions, factor this out. + torch-tpu = + let + abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}"; + version = "0.1.1.dev20260707090224"; + in + requireFile { + name = "torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl"; + hash = "sha256-eCjwoX0UKd/L/IccxXn88p0GdyQjFdPl2Er1luZz4H0="; # cp313 + message = '' + torch_tpu is served from a gated Google Artifact Registry and + cannot be fetched by a pure Nix build. Fetch and register it with: + + GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ + nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi} + ''; + }; + in buildPythonPackage.override { stdenv = effectiveStdenv; } { pname = "torch"; @@ -139,10 +166,20 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } { inherit url hash; }; + # If we are building for TPU, merge torch-tpu into the output path. + # Note that we cannot simply depend on torch-tpu, because torch-tpu + # itself depends on torch. The wheel unpacking hook does not seemt + # to support `srcs`, so we'll just copy over the wheel on TPU. + preUnpack = lib.optionals tpuSupport '' + mkdir -p dist + cp "${torch-tpu}" "dist/$(stripHash "${torch-tpu}")" + ''; + # We use our own pythonWheelDepsHook. dontCheckRuntimeDeps = true; nativeBuildInputs = [ + pythonKernelImportsCheckHook pythonRelaxWheelDepsHook pythonWheelDepsCheckHook ] @@ -225,6 +262,14 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } { sympy typing-extensions ] + ++ lib.optionals tpuSupport [ + libtpu + jax + frozendict + immutabledict + portpicker + tensorboard + ] ++ lib.optionals tritonSupport [ effectiveTriton ] @@ -367,7 +412,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } { dontStrip = true; - pythonImportsCheck = [ "torch" ]; + pythonKernelImportsCheck = if tpuSupport then [ "torch_tpu" ] else [ "torch" ]; passthru = { inherit diff --git a/nix-builder/scripts/gen_variants_markdown.py b/nix-builder/scripts/gen_variants_markdown.py index b097cebcc..f528144ce 100755 --- a/nix-builder/scripts/gen_variants_markdown.py +++ b/nix-builder/scripts/gen_variants_markdown.py @@ -8,6 +8,7 @@ "cuda": "CUDA", "metal": "Metal", "rocm": "ROCm", + "tpu": "TPU", "xpu": "XPU", } diff --git a/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh b/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh new file mode 100755 index 000000000..31c9f9edb --- /dev/null +++ b/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Fetch a wheel from Google's gated torch-tpu-virtual-registry and +# register it in the Nix store under the exact (name, sha256) that +# `requireFile` expects in pkgs/python-modules/{libtpu,torch_tpu}/default.nix. +# +# The fetch (and the OAuth2 token it needs) happens here, out-of-band, +# so the actual Nix build stays pure and never touches the network. +# +# Usage: +# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ +# scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43 [cp313] +# GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \ +# scripts/helpers/prefetch-tpu-wheel.sh torch_tpu 0.1.1.dev20260707090224 [cp313] +# +# Prints a `hash = "sha256-...";` line to paste into the matching +# default.nix when bumping a version. The third argument is the +# CPython ABI tag (default cp313, matching the nixpkgs default python +# used by the builder). +# +# Requires `gcloud auth login` once, plus nix (with the `nix-command` +# experimental feature) and jq. +set -euo pipefail + +pkg="${1:?usage: $0 [abi]}" +version="${2:?usage: $0 [abi]}" +abi="${3:-cp313}" +: "${GCLOUD_ACCESS_TOKEN:?set GCLOUD_ACCESS_TOKEN to \$(gcloud auth print-access-token)}" + +registry_pkg="${pkg//_/-}" +name="${pkg}-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl" +url="https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/${registry_pkg}/${name}" +netrc="$(mktemp)" +trap 'rm -f "${netrc}"' EXIT +printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \ + "${GCLOUD_ACCESS_TOKEN}" >"${netrc}" + +hash="$(nix store prefetch-file --json --hash-type sha256 --name "${name}" \ + --option netrc-file "${netrc}" "${url}" | jq -r .hash)" + +echo "hash = \"${hash}\"; # ${abi}" diff --git a/nix-builder/versions.nix b/nix-builder/versions.nix index 4e15c3b5d..6587d31e9 100644 --- a/nix-builder/versions.nix +++ b/nix-builder/versions.nix @@ -7,6 +7,15 @@ ]; bundleBuild = false; } + { + torchVersion = "2.11"; + tpu = true; + systems = [ "x86_64-linux" ]; + # Excluded from the shared build cache: torch_tpu's autoPatchelfHook + # step produces an output that is itself a modified copy of Google's + # gated libtpu/torch_tpu wheels, so it can't be published publicly. + bundleBuild = false; + } { torchVersion = "2.12";