From 46433088c4d90ff159d13524becabcb577b919a1 Mon Sep 17 00:00:00 2001 From: samthakur587 Date: Thu, 27 Aug 2026 15:30:18 +0530 Subject: [PATCH] Restructure docs into PyTorch-style Get Started and API sections. Split installation from getting started, add hand-written API pages with mkdocstrings for Python modules, and teach docs CI to build with a stub native extension. Co-authored-by: Cursor --- .github/workflows/ci.yml | 7 ++- .github/workflows/docs.yml | 7 ++- CONTRIBUTING.md | 6 ++- README.md | 3 +- docs/api/cpp.md | 44 ++++++++++++++++++ docs/api/index.md | 12 +++++ docs/api/kernels.md | 42 +++++++++++++++++ docs/api/ops.md | 46 ++++++++++++++++++ docs/api/torch.md | 65 ++++++++++++++++++++++++++ docs/getting-started.md | 51 ++++---------------- docs/index.md | 5 +- docs/install.md | 58 +++++++++++++++++++++++ docs/operations.md | 75 ++---------------------------- mkdocs.yml | 36 ++++++++++---- pyproject.toml | 15 +++++- torch_tvarant/__init__.py | 26 +++++++++-- torch_tvarant/_native_docs_stub.py | 48 +++++++++++++++++++ torch_tvarant/compiler.py | 1 + torch_tvarant/tvarant/__init__.py | 34 +++++++++++--- 19 files changed, 444 insertions(+), 137 deletions(-) create mode 100644 docs/api/cpp.md create mode 100644 docs/api/index.md create mode 100644 docs/api/kernels.md create mode 100644 docs/api/ops.md create mode 100644 docs/api/torch.md create mode 100644 docs/install.md create mode 100644 torch_tvarant/_native_docs_stub.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 11ac26a..0eb12c2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,7 +96,12 @@ jobs: cache: pip - name: Install docs dependencies - run: pip install mkdocs-material mkdocs + run: | + pip install mkdocs-material mkdocs "mkdocstrings[python]>=0.25" + pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Build MkDocs site + env: + MKDOCS_BUILD: "1" + TORCH_DEVICE_BACKEND_AUTOLOAD: "0" run: mkdocs build --strict diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 8e1f81e..923e7b9 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -26,9 +26,14 @@ jobs: cache: pip - name: Install MkDocs - run: pip install mkdocs-material mkdocs + run: | + pip install mkdocs-material mkdocs "mkdocstrings[python]>=0.25" + pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Build site + env: + MKDOCS_BUILD: "1" + TORCH_DEVICE_BACKEND_AUTOLOAD: "0" run: mkdocs build --strict - name: Upload artifact diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 771417e..accc44f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -45,9 +45,13 @@ TVARANT_BACKEND=opencl pytest tests/test_opencl.py -v ```bash pip install -e ".[docs]" -mkdocs serve +# Optional: CPU torch so live imports work +pip install torch --index-url https://download.pytorch.org/whl/cpu +MKDOCS_BUILD=1 TORCH_DEVICE_BACKEND_AUTOLOAD=0 mkdocs serve ``` +`MKDOCS_BUILD=1` loads a stub `_C` when the native extension is not built (for mkdocstrings). + Open http://127.0.0.1:8000 to preview the documentation site. ## Project layout diff --git a/README.md b/README.md index 34bc900..2b58d75 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,8 @@ pip install -e . pytest tests/ -v ``` -See [Getting Started](docs/getting-started.md) for full install instructions. +See [Installation](docs/install.md) and [Getting Started](docs/getting-started.md). +API: [docs/api](docs/api/index.md). ## Features diff --git a/docs/api/cpp.md b/docs/api/cpp.md new file mode 100644 index 0000000..c6142a6 --- /dev/null +++ b/docs/api/cpp.md @@ -0,0 +1,44 @@ +# C++ extension + +`torch_tvarant` is a PyTorch C++ extension (`PrivateUse1`) built from `csrc/`. + +## Layout + +| Path | Role | +|---|---| +| `csrc/Module.cpp` | pybind entry: device helpers, force-link registrations | +| `csrc/aten/` | ATen `TORCH_LIBRARY_IMPL(aten, PrivateUse1, …)` | +| `csrc/runtime/` | CPU sim + OpenCL runtimes | +| `csrc/kernels/host/` | Host fp32 kernels | +| `csrc/kernels/opencl/` | OpenCL `.cl` sources | +| `csrc/jit/` | Pointwise / GEMM fusion helpers | + +## Registering an ATen op + +```cpp +TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) { + m.impl("relu", TORCH_FN(tvarant::ops::relu)); +} +``` + +Implement the kernel (host + optional OpenCL), wire `runtime().launch`, then +register the impl in `Ops.cpp`, `LlmOps.cpp`, or another `*Ops.cpp` TU that is +force-linked from `Module.cpp`. + +Qualify calls into `tvarant::ops::…` when names collide with ATen ADL +(see existing `silu` / unary helpers). + +## Custom library + +`TORCH_LIBRARY(tvarant, m)` in `CustomOps.cpp` defines +`torch.ops.tvarant.linear_act` and `torch.ops.tvarant.pointwise`. + +## Build + +```bash +pip install -e . --no-build-isolation +# OpenCL: +USE_OPENCL=1 pip install -e . --no-build-isolation +``` + +See [Installation](../install.md) and [Contributing](../contributing.md). diff --git a/docs/api/index.md b/docs/api/index.md new file mode 100644 index 0000000..2471b24 --- /dev/null +++ b/docs/api/index.md @@ -0,0 +1,12 @@ +# API reference + +This section mirrors a PyTorch-style API layout for **torch_tvarant**. + +| Page | Contents | +|---|---| +| [torch.tvarant](torch.md) | Device module helpers + Python package / compiler (auto-generated) | +| [Supported ops](ops.md) | ATen and custom ops registered for `PrivateUse1` | +| [Kernels](kernels.md) | Host and OpenCL kernel catalog | +| [C++](cpp.md) | Extension layout, `TORCH_LIBRARY_IMPL`, adding ops | + +Unimplemented ATen ops fall back to CPU; see [Supported ops](ops.md). diff --git a/docs/api/kernels.md b/docs/api/kernels.md new file mode 100644 index 0000000..4ccddfe --- /dev/null +++ b/docs/api/kernels.md @@ -0,0 +1,42 @@ +# Kernels + +Kernels are the device implementations behind ATen ops. The CPU simulator calls +host functions in `csrc/kernels/host/`; the OpenCL path uses sources under +`csrc/kernels/opencl/` (also embedded in `OpenCLRuntime.cpp`). + +## Host kernels (`tvarant::host`) + +| Function | Role | +|---|---| +| `fill_f32` / `copy_f32` | Init and memcpy | +| `add_f32` / `mul_f32` | Binary elementwise | +| `add_scalar_f32` / `scale_f32` | Scalar elementwise | +| `relu_f32` / `silu_f32` | Activations | +| `gemm_f32` / `gemm_bias_act_f32` | GEMM (+ bias / act) | +| `bmm_f32` | Batched matmul | +| `softmax_f32` | Softmax along a dim | +| `layer_norm_f32` | Layer norm forward | +| `embedding_f32` | Embedding lookup | + +Declared in `csrc/kernels/host/HostKernels.h`. + +## OpenCL kernel catalog + +| Kernel | File | +|---|---| +| `fill_kernel` | `fill.cl` | +| `copy_kernel` | `copy.cl` | +| `add_kernel`, `mul_kernel` | `binary.cl` | +| `relu_kernel` | `relu.cl` | +| `silu_kernel`, `scale_kernel`, `add_scalar_kernel` | `elementwise.cl` | +| `gemm_kernel`, `gemm_bias_act_kernel`, `bmm_kernel` | `gemm.cl` | +| `softmax_kernel`, `layernorm_kernel`, `embedding_kernel` | `reduce.cl` | + +Pointwise JIT kernels are generated at runtime (not checked in as `.cl` files). +See [JIT Compiler](../jit-compiler.md). + +## Dispatch + +`CpuSimRuntime::launch` / `OpenCLRuntime::launch` select a kernel by name from +`LaunchParams.kernel`. ATen ops in `csrc/aten/` fill those params and call +`tvarant::runtime().launch(...)`. diff --git a/docs/api/ops.md b/docs/api/ops.md new file mode 100644 index 0000000..368fd48 --- /dev/null +++ b/docs/api/ops.md @@ -0,0 +1,46 @@ +# Supported ops + +All listed ops run on `torch.device("tvarant")` unless they fall back to CPU. + +## Core ATen ops + +| Op | Notes | +|---|---| +| `empty`, `empty_strided` | Strided layout only | +| `as_strided`, `view`, `resize_` | View/resize support | +| `_copy_from` | H2D, D2H, D2D via runtime | +| `fill_` | Scalar fill | +| `add`, `mul` | Broadcast elementwise | +| `relu` | Elementwise ReLU | +| `mm`, `addmm` | Matrix multiply; uses fused GEMM path | +| `linear` | `nn.Linear` backend; transposed weight | + +## LLM ops + +| Op | Notes | +|---|---| +| `silu` | SiLU / Swish activation | +| `softmax` | Last-dim or arbitrary dim | +| `matmul` | 2D GEMM or batched BMM | +| `bmm` | Batch matrix multiply | +| `native_layer_norm` | Forward layer norm + mean/rstd | +| `embedding` | Token lookup (int32 indices) | +| `mul.Scalar`, `add.Scalar` | Scalar elementwise | + +## Custom ops (`torch.ops.tvarant.*`) + +| Op | Signature | Description | +|---|---|---| +| `linear_act` | `(x, weight, bias?, act, trans_b=False)` | Fused GEMM + bias + relu/silu | +| `pointwise` | `(inputs, ops, a, b, input_ids, alphas, consts)` | JIT fused elementwise program | + +Activation strings for `linear_act`: `"none"`, `"relu"`, `"silu"`. + +## CPU fallback + +Any ATen op not registered for `PrivateUse1` executes on CPU via the fallback +handler. Tensors are copied to/from device memory as needed. + +## Adding a new op + +See [C++ API](cpp.md) and [Contributing](../contributing.md). diff --git a/docs/api/torch.md b/docs/api/torch.md new file mode 100644 index 0000000..4cc7d73 --- /dev/null +++ b/docs/api/torch.md @@ -0,0 +1,65 @@ +# torch.tvarant + +Importing `torch_tvarant` registers the **tvarant** PrivateUse1 backend and exposes +helpers under `torch.tvarant` (same surface as other PyTorch device modules). + +```python +import torch +import torch_tvarant + +torch.tvarant.is_available() +torch.tvarant.device_count() +x = torch.empty(2, 3, device="tvarant") +``` + +## Device helpers + +| Function | Description | +|---|---| +| `is_available()` | Backend extension loaded and usable | +| `is_initialized()` | Lazy init has run | +| `device_count()` | Number of logical devices (currently 1) | +| `current_device()` | Current device index | +| `set_device(index)` | Set current device index | +| `synchronize(device=None)` | Wait for outstanding kernels | +| `backend()` | Runtime name: `'sim'` or `'opencl'` | +| `manual_seed(seed)` | Seed the Tvarant RNG | +| `manual_seed_all(seed)` | Alias of `manual_seed` | +| `get_rng_state(device=...)` | RNG state tensor | +| `set_rng_state(state, device=...)` | Restore RNG state | + +These map to bindings in `csrc/Module.cpp`. + +## Package modules + +Auto-generated from Python docstrings (requires `MKDOCS_BUILD=1` when the +native extension is not built). + +::: torch_tvarant + options: + members: false + show_root_heading: true + show_source: false + +::: torch_tvarant.tvarant + options: + show_root_heading: true + show_source: false + members_order: source + docstring_style: google + filters: + - "!^_" + +::: torch_tvarant.compiler + options: + show_root_heading: true + show_source: false + members: + - compile + - compile_fx + - register + - trace_module + - fuse_gemm_epilogue + - fuse_pointwise + - TvarantTracer + members_order: source diff --git a/docs/getting-started.md b/docs/getting-started.md index 32e863b..1cf4eb7 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -1,28 +1,6 @@ # Getting started -## Requirements - -| Component | Version | -|---|---| -| Python | 3.9+ | -| PyTorch | 2.1+ | -| C++ compiler | C++17 (GCC, Clang, or MSVC) | -| Optional | OpenCL SDK for FPGA builds | - -## Install (CPU simulator) - -```bash -git clone https://github.com/samthakur587/Tvarant.git -cd Tvarant -python -m venv .venv -source .venv/bin/activate -pip install -U pip wheel setuptools ninja -pip install torch pytest -pip install -e . -pytest tests/ -v -``` - -On **Windows**, use an x64 Native Tools prompt (or run `vcvars64.bat` first). +After [installing](install.md) `torch_tvarant`, verify the device and run a small model. ## Verify installation @@ -31,7 +9,7 @@ import torch import torch_tvarant print(torch.tvarant.is_available()) # True -print(torch.tvarant.backend()) # 'sim' +print(torch.tvarant.backend()) # 'sim' or 'opencl' x = torch.ones(4, device="tvarant") print(x.device) # tvarant:0 ``` @@ -65,21 +43,12 @@ y = compiled(x) See [JIT Compiler](jit-compiler.md) for fusion details. -## OpenCL / FPGA +## Learn more -```bash -USE_OPENCL=1 pip install -e . -TVARANT_BACKEND=opencl python your_script.py -``` - -See [FPGA / OpenCL](fpga-opencl.md) for environment variables and kernel layout. - -## Development install - -```bash -pip install -e ".[dev,docs]" -pytest tests/ -v -mkdocs serve -``` - -See [Contributing](contributing.md) for the full contributor workflow. +| Topic | Page | +|---|---| +| Device helpers | [torch.tvarant](api/torch.md) | +| Supported ops | [Supported ops](api/ops.md) | +| Kernels | [Kernels](api/kernels.md) | +| C++ extension | [C++ API](api/cpp.md) | +| FPGA path | [FPGA / OpenCL](fpga-opencl.md) | diff --git a/docs/index.md b/docs/index.md index 6ecaf9d..3539ddc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -27,12 +27,13 @@ y = torch.nn.functional.relu(x @ w) | Resource | Link | |---|---| -| GitHub repository | [samthakur587/Tvarant](https://github.com/samthakur587/Tvarant) | +| Installation | [install.md](install.md) | | Getting started | [getting-started.md](getting-started.md) | +| API reference | [api/index.md](api/index.md) | | Architecture | [architecture.md](architecture.md) | | JIT compiler | [jit-compiler.md](jit-compiler.md) | | Contributing | [contributing.md](contributing.md) | -| Issue tracker | [GitHub Issues](https://github.com/samthakur587/Tvarant/issues) | +| GitHub | [samthakur587/Tvarant](https://github.com/samthakur587/Tvarant) | | Changelog | [CHANGELOG.md](https://github.com/samthakur587/Tvarant/blob/main/CHANGELOG.md) | ## License diff --git a/docs/install.md b/docs/install.md new file mode 100644 index 0000000..b645023 --- /dev/null +++ b/docs/install.md @@ -0,0 +1,58 @@ +# Installation + +## Requirements + +| Component | Version | +|---|---| +| Python | 3.9+ | +| PyTorch | 2.1+ | +| C++ compiler | C++17 (GCC, Clang, or MSVC) | +| Optional | OpenCL SDK for FPGA builds | + +## Install from source (CPU simulator) + +```bash +git clone https://github.com/samthakur587/Tvarant.git +cd Tvarant +python -m venv .venv +source .venv/bin/activate +pip install -U pip wheel setuptools ninja +pip install torch pytest +pip install -e . +pytest tests/ -v +``` + +On **Windows**, use an x64 Native Tools prompt (or run `vcvars64.bat` first). + +## Install from PyPI + +When a release is published: + +```bash +pip install torch-tvarant +``` + +Use a matching CPU or CUDA PyTorch build from [pytorch.org](https://pytorch.org/get-started/locally/) first. + +## OpenCL / FPGA build + +```bash +USE_OPENCL=1 pip install -e . +TVARANT_BACKEND=opencl python your_script.py +``` + +See [FPGA / OpenCL](fpga-opencl.md) for environment variables and kernel layout. + +## Development install + +```bash +pip install -e ".[dev,docs]" +pytest tests/ -v +mkdocs serve +``` + +See [Contributing](contributing.md) for the contributor workflow. + +## Next steps + +Continue with [Getting Started](getting-started.md). diff --git a/docs/operations.md b/docs/operations.md index d2b89a1..b6b0b96 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -1,74 +1,7 @@ # Operations -All ops run on `torch.device("tvarant")` unless they fall back to CPU. +This page moved into the API reference: -## Core ATen ops - -| Op | Notes | -|---|---| -| `empty`, `empty_strided` | Strided layout only | -| `as_strided`, `view`, `resize_` | View/resize support | -| `_copy_from` | H2D, D2H, D2D via runtime | -| `fill_` | Scalar fill | -| `add`, `mul` | Broadcast elementwise | -| `relu` | Elementwise ReLU | -| `mm`, `addmm` | Matrix multiply; uses fused GEMM path | -| `linear` | `nn.Linear` backend; transposed weight | - -## LLM ops - -| Op | Notes | -|---|---| -| `silu` | SiLU / Swish activation | -| `softmax` | Last-dim or arbitrary dim | -| `matmul` | 2D GEMM or batched BMM | -| `bmm` | Batch matrix multiply | -| `native_layer_norm` | Forward layer norm + mean/rstd | -| `embedding` | Token lookup (int32 indices) | -| `mul.Scalar`, `add.Scalar` | Scalar elementwise | - -## Custom ops (`torch.ops.tvarant.*`) - -| Op | Signature | Description | -|---|---|---| -| `linear_act` | `(x, weight, bias?, act, trans_b=False)` | Fused GEMM + bias + relu/silu | -| `pointwise` | `(inputs, ops, a, b, input_ids, alphas, consts)` | JIT fused elementwise program | - -Activation strings for `linear_act`: `"none"`, `"relu"`, `"silu"`. - -## CPU fallback - -Any ATen op not registered for `PrivateUse1` executes on CPU via the fallback -handler. Tensors are copied to/from device memory as needed. - -Check for unexpected CPU fallback during profiling: - -```python -import torch -torch.autograd.set_detect_anomaly(False) -# Watch for slow ops or CPU device in stack traces -``` - -## Adding a new op - -1. Implement host kernel in `csrc/kernels/host/HostKernels.{h,cpp}` -2. Add OpenCL kernel in `csrc/kernels/opencl/` and embed in `OpenCLRuntime.cpp` -3. Wire dispatch in `CpuSimRuntime::launch` and `OpenCLRuntime::launch` -4. Register ATen impl in `Ops.cpp` or `LlmOps.cpp` -5. Add pytest in `tests/` - -See [Contributing](contributing.md) for the full workflow. - -## OpenCL kernel catalog - -| Kernel | File | -|---|---| -| `fill_kernel` | `fill.cl` | -| `copy_kernel` | `copy.cl` | -| `add_kernel`, `mul_kernel` | `binary.cl` | -| `relu_kernel` | `relu.cl` | -| `silu_kernel`, `scale_kernel` | `elementwise.cl` | -| `gemm_kernel`, `gemm_bias_act_kernel`, `bmm_kernel` | `gemm.cl` | -| `softmax_kernel`, `layernorm_kernel`, `embedding_kernel` | `reduce.cl` | - -Pointwise JIT kernels are generated at runtime (not in `.cl` files). +- [Supported ops](api/ops.md) +- [Kernels](api/kernels.md) +- [C++ extension](api/cpp.md) diff --git a/mkdocs.yml b/mkdocs.yml index d573f1a..cdd6c6d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -27,14 +27,26 @@ theme: - search.suggest nav: - - Home: index.md - - Getting Started: getting-started.md - - Architecture: architecture.md - - JIT Compiler: jit-compiler.md - - Operations: operations.md - - FPGA / OpenCL: fpga-opencl.md - - Roadmap: roadmap.md - - Contributing: contributing.md + - Get Started: + - Home: index.md + - Installation: install.md + - Getting Started: getting-started.md + - Tutorials: + - Getting Started: getting-started.md + - How-Tos: + - JIT Compiler: jit-compiler.md + - FPGA / OpenCL: fpga-opencl.md + - Notes: + - Architecture: architecture.md + - Roadmap: roadmap.md + - Contributing: contributing.md + - API: + - Overview: api/index.md + - torch.tvarant: api/torch.md + - Supported ops: api/ops.md + - Kernels: api/kernels.md + - C++: api/cpp.md + - Operations (redirect): operations.md markdown_extensions: - admonition @@ -47,6 +59,14 @@ markdown_extensions: plugins: - search + - mkdocstrings: + handlers: + python: + options: + docstring_style: google + show_source: false + separate_signature: true + merge_init_into_class: true extra: social: diff --git a/pyproject.toml b/pyproject.toml index 1797fad..2ffeb9a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,8 +50,19 @@ torch_tvarant = "torch_tvarant:_autoload" [project.optional-dependencies] dev = ["pytest>=7", "ruff>=0.4", "build>=1.0"] -docs = ["mkdocs>=1.5", "mkdocs-material>=9"] -all = ["pytest>=7", "ruff>=0.4", "build>=1.0", "mkdocs>=1.5", "mkdocs-material>=9"] +docs = [ + "mkdocs>=1.5", + "mkdocs-material>=9", + "mkdocstrings[python]>=0.25", +] +all = [ + "pytest>=7", + "ruff>=0.4", + "build>=1.0", + "mkdocs>=1.5", + "mkdocs-material>=9", + "mkdocstrings[python]>=0.25", +] [tool.setuptools.packages.find] include = ["torch_tvarant", "torch_tvarant.*"] diff --git a/torch_tvarant/__init__.py b/torch_tvarant/__init__.py index 24e4e9a..2cf100a 100644 --- a/torch_tvarant/__init__.py +++ b/torch_tvarant/__init__.py @@ -1,10 +1,30 @@ -"""Out-of-tree PyTorch backend that registers torch.device('tvarant').""" +"""Out-of-tree PyTorch backend that registers ``torch.device("tvarant")``. + +Importing this package: + +* loads the native ``_C`` extension (or a docs stub when ``MKDOCS_BUILD=1``) +* renames PrivateUse1 to ``tvarant`` +* exposes device helpers under ``torch.tvarant`` +""" from __future__ import annotations +import os +import sys + import torch -from . import _C, compiler, tvarant +try: + from . import _C +except ImportError: + if os.environ.get("MKDOCS_BUILD") == "1": + from . import _native_docs_stub as _C + + sys.modules[__name__ + "._C"] = _C + else: + raise + +from . import compiler, tvarant torch.utils.rename_privateuse1_backend("tvarant") torch._register_device_module("tvarant", tvarant) @@ -19,4 +39,4 @@ def _autoload() -> None: - """Entry point for TORCH_DEVICE_BACKEND_AUTOLOAD.""" + """Entry point for ``TORCH_DEVICE_BACKEND_AUTOLOAD``.""" diff --git a/torch_tvarant/_native_docs_stub.py b/torch_tvarant/_native_docs_stub.py new file mode 100644 index 0000000..75c147a --- /dev/null +++ b/torch_tvarant/_native_docs_stub.py @@ -0,0 +1,48 @@ +"""Docs-only stand-in for the native ``torch_tvarant._C`` extension. + +Loaded when ``MKDOCS_BUILD=1`` and the real extension is not importable. +""" + +from __future__ import annotations + +import torch + + +def is_available() -> bool: + return True + + +def device_count() -> int: + return 1 + + +def current_device() -> int: + return 0 + + +def set_device(index: int) -> None: + _ = index + + +def synchronize() -> None: + return None + + +def backend() -> str: + return "sim" + + +def device() -> torch.device: + return torch.device("cpu") + + +def manual_seed(seed: int) -> None: + _ = seed + + +def get_rng_state() -> torch.Tensor: + return torch.zeros(1, dtype=torch.uint8) + + +def set_rng_state(state: torch.Tensor) -> None: + _ = state diff --git a/torch_tvarant/compiler.py b/torch_tvarant/compiler.py index 227ec3c..018df06 100644 --- a/torch_tvarant/compiler.py +++ b/torch_tvarant/compiler.py @@ -49,6 +49,7 @@ def is_leaf_module(self, m: nn.Module, module_qualified_name: str) -> bool: def trace_module(module: nn.Module) -> torch.fx.GraphModule: + """FX-trace ``module``, inlining Linear / common activations.""" tracer = TvarantTracer() graph = tracer.trace(module) return torch.fx.GraphModule(module, graph) diff --git a/torch_tvarant/tvarant/__init__.py b/torch_tvarant/tvarant/__init__.py index 7190201..1182bd7 100644 --- a/torch_tvarant/tvarant/__init__.py +++ b/torch_tvarant/tvarant/__init__.py @@ -1,15 +1,28 @@ -"""Python device module registered as torch.tvarant.""" +"""Python device module registered as ``torch.tvarant``. + +Functions here mirror the usual PyTorch device-module surface +(``is_available``, ``synchronize``, RNG helpers, …) and forward to the +native extension. +""" from __future__ import annotations import importlib +import os +import sys from typing import Union import torch -# Import the extension module directly so this package can load while -# torch_tvarant/__init__.py is still initializing (avoids circular import). -_C = importlib.import_module("torch_tvarant._C") +try: + _C = importlib.import_module("torch_tvarant._C") +except ImportError: + if os.environ.get("MKDOCS_BUILD") == "1": + from torch_tvarant import _native_docs_stub as _C + + sys.modules["torch_tvarant._C"] = _C + else: + raise _initialized = False @@ -23,31 +36,38 @@ def _lazy_init() -> None: def is_available() -> bool: + """Return whether the Tvarant backend extension is usable.""" return bool(_C.is_available()) def is_initialized() -> bool: + """Return whether lazy device initialization has completed.""" return _initialized def device_count() -> int: + """Return the number of logical Tvarant devices.""" return int(_C.device_count()) def current_device() -> int: + """Return the current device index.""" return int(_C.current_device()) def set_device(device: int) -> None: + """Set the current device index.""" _C.set_device(int(device)) def synchronize(device: int | None = None) -> None: + """Wait for outstanding work on the runtime (``device`` reserved).""" _ = device _C.synchronize() def backend() -> str: + """Return the active runtime name (``sim`` or ``opencl``).""" return str(_C.backend()) @@ -56,16 +76,17 @@ def _is_in_bad_fork() -> bool: def manual_seed(seed: int) -> None: - # Seed only the Tvarant generator. Do not call torch.manual_seed — that - # re-enters torch.tvarant.manual_seed_all and recurses. + """Seed the Tvarant generator only (does not call ``torch.manual_seed``).""" _C.manual_seed(int(seed)) def manual_seed_all(seed: int) -> None: + """Seed all Tvarant devices (currently one device).""" manual_seed(seed) def get_rng_state(device: Union[int, str, torch.device] = "tvarant") -> torch.Tensor: + """Return the RNG state tensor for ``device``.""" _ = device return _C.get_rng_state() @@ -73,5 +94,6 @@ def get_rng_state(device: Union[int, str, torch.device] = "tvarant") -> torch.Te def set_rng_state( new_state: torch.Tensor, device: Union[int, str, torch.device] = "tvarant" ) -> None: + """Restore RNG state for ``device``.""" _ = device _C.set_rng_state(new_state)