Restructure docs into PyTorch-style Get Started and API sections

.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

.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

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

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
 

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).

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).

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(...)`.

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).

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

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) |

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