diff --git a/docs/api/compiler/TvarantTracer.md b/docs/api/compiler/TvarantTracer.md new file mode 100644 index 0000000..5864d80 --- /dev/null +++ b/docs/api/compiler/TvarantTracer.md @@ -0,0 +1,18 @@ +# TvarantTracer + +```python +class torch_tvarant.compiler.TvarantTracer +``` + +::: torch_tvarant.compiler.TvarantTracer + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/compile.md b/docs/api/compiler/compile.md new file mode 100644 index 0000000..9d81a5d --- /dev/null +++ b/docs/api/compiler/compile.md @@ -0,0 +1,18 @@ +# compile + +```python +torch_tvarant.compiler.compile(module, *, dynamic=False) → Callable +``` + +::: torch_tvarant.compiler.compile + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/compile_fx.md b/docs/api/compiler/compile_fx.md new file mode 100644 index 0000000..0fc131d --- /dev/null +++ b/docs/api/compiler/compile_fx.md @@ -0,0 +1,18 @@ +# compile_fx + +```python +torch_tvarant.compiler.compile_fx(gm, example_inputs=None) → GraphModule +``` + +::: torch_tvarant.compiler.compile_fx + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/fuse_gemm_epilogue.md b/docs/api/compiler/fuse_gemm_epilogue.md new file mode 100644 index 0000000..5422aab --- /dev/null +++ b/docs/api/compiler/fuse_gemm_epilogue.md @@ -0,0 +1,18 @@ +# fuse_gemm_epilogue + +```python +torch_tvarant.compiler.fuse_gemm_epilogue(gm) → GraphModule +``` + +::: torch_tvarant.compiler.fuse_gemm_epilogue + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/fuse_pointwise.md b/docs/api/compiler/fuse_pointwise.md new file mode 100644 index 0000000..57bbc2c --- /dev/null +++ b/docs/api/compiler/fuse_pointwise.md @@ -0,0 +1,18 @@ +# fuse_pointwise + +```python +torch_tvarant.compiler.fuse_pointwise(gm) → GraphModule +``` + +::: torch_tvarant.compiler.fuse_pointwise + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/index.md b/docs/api/compiler/index.md new file mode 100644 index 0000000..48eba59 --- /dev/null +++ b/docs/api/compiler/index.md @@ -0,0 +1,20 @@ +# torch_tvarant.compiler + +Graph compiler that fuses GEMM epilogues and pointwise chains for inference on +`device="tvarant"`. + +```python +import torch_tvarant +compiled = torch_tvarant.compiler.compile(model) +y = compiled(x) +``` + +| API | Summary | +|---|---| +| [`compile`](compile.md) | Compile a module for inference | +| [`compile_fx`](compile_fx.md) | Fuse an FX graph | +| [`trace_module`](trace_module.md) | FX-trace with inlined Linear/activations | +| [`fuse_gemm_epilogue`](fuse_gemm_epilogue.md) | Fuse mm/linear + act into `linear_act` | +| [`fuse_pointwise`](fuse_pointwise.md) | Collapse pointwise chains | +| [`register`](register.md) | Register Dynamo backend `"tvarant"` | +| [`TvarantTracer`](TvarantTracer.md) | Custom FX tracer | diff --git a/docs/api/compiler/register.md b/docs/api/compiler/register.md new file mode 100644 index 0000000..0fa91cf --- /dev/null +++ b/docs/api/compiler/register.md @@ -0,0 +1,18 @@ +# register + +```python +torch_tvarant.compiler.register() → None +``` + +::: torch_tvarant.compiler.register + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/compiler/trace_module.md b/docs/api/compiler/trace_module.md new file mode 100644 index 0000000..67aeba9 --- /dev/null +++ b/docs/api/compiler/trace_module.md @@ -0,0 +1,18 @@ +# trace_module + +```python +torch_tvarant.compiler.trace_module(module) → GraphModule +``` + +::: torch_tvarant.compiler.trace_module + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + merge_init_into_class: true + members_order: source + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/index.md b/docs/api/index.md index 2471b24..6164953 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -1,12 +1,24 @@ # API reference -This section mirrors a PyTorch-style API layout for **torch_tvarant**. +PyTorch-style reference for **torch_tvarant**. Each entry is a dedicated page +with signature, parameters, and examples. -| 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 | +## torch.tvarant -Unimplemented ATen ops fall back to CPU; see [Supported ops](ops.md). +Device module — [`torch.tvarant` overview](torch/index.md) + +- [`is_available`](torch/is_available.md) · [`device_count`](torch/device_count.md) · [`backend`](torch/backend.md) +- [`set_device`](torch/set_device.md) · [`synchronize`](torch/synchronize.md) +- [`manual_seed`](torch/manual_seed.md) · [`get_rng_state`](torch/get_rng_state.md) · … + +## torch_tvarant.compiler + +- [`compile`](compiler/compile.md) · [`compile_fx`](compiler/compile_fx.md) +- [`trace_module`](compiler/trace_module.md) · [`register`](compiler/register.md) +- [Compiler overview](compiler/index.md) + +## Supported ops + +- [Ops index](ops/index.md) — `relu`, `add`, `mm`, `linear`, `linear_act`, … +- [Kernels](kernels.md) — host / OpenCL catalog +- [C++](cpp.md) — extension layout & `TORCH_LIBRARY_IMPL` diff --git a/docs/api/ops.md b/docs/api/ops.md deleted file mode 100644 index 368fd48..0000000 --- a/docs/api/ops.md +++ /dev/null @@ -1,46 +0,0 @@ -# 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/ops/add.md b/docs/api/ops/add.md new file mode 100644 index 0000000..8ac50ad --- /dev/null +++ b/docs/api/ops/add.md @@ -0,0 +1,33 @@ +# add + +```python +torch.add(input, other, *, alpha=1, out=None) → Tensor +input + other +``` + +Adds `other`, scaled by `alpha`, to `input`: + +$$ +\text{out}_i = \text{input}_i + \alpha \times \text{other}_i +$$ + +## Parameters + +- **input** (`Tensor`) – First operand (Tvarant). +- **other** (`Tensor or Number`) – Second operand (broadcasting). +- **alpha** (`Number`) – Multiplier for `other`. Default: `1`. + +## Shape + +- Input: $(*)$ +- Output: broadcast shape of `input` and `other` + +## Example + +```python +y = torch.add(a, b, alpha=1.0) # or a + b +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/clone.md b/docs/api/ops/clone.md new file mode 100644 index 0000000..d9345ff --- /dev/null +++ b/docs/api/ops/clone.md @@ -0,0 +1,27 @@ +# clone / contiguous + +```python +Tensor.clone(memory_format=torch.preserve_format) → Tensor +Tensor.contiguous(memory_format=torch.contiguous_format) → Tensor +``` + +Materialize a copy (`clone`) or a contiguous storage (`contiguous`) on Tvarant via `_copy_from`. + +## Parameters + +- **memory_format** (`torch.memory_format`) – Desired memory format. + +## Shape + +- Output: same logical shape as input + +## Example + +```python +y = x.clone() +c = x.t().contiguous() +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/embedding.md b/docs/api/ops/embedding.md new file mode 100644 index 0000000..0b9fb57 --- /dev/null +++ b/docs/api/ops/embedding.md @@ -0,0 +1,27 @@ +# embedding + +```python +torch.nn.functional.embedding(input, weight, ...) → Tensor +``` + +Lookup embeddings. Indices should be convertible to int32 on the Tvarant path. + +## Parameters + +- **input** (`LongTensor`) – Token indices. +- **weight** (`Tensor`) – Embedding table `(V, D)` on Tvarant. + +## Shape + +- Input: $(*)$ +- Output: $(*, D)$ + +## Example + +```python +e = torch.nn.functional.embedding(tokens, weight) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/empty.md b/docs/api/ops/empty.md new file mode 100644 index 0000000..edcedbd --- /dev/null +++ b/docs/api/ops/empty.md @@ -0,0 +1,28 @@ +# empty + +```python +torch.empty(*size, *, dtype=None, layout=torch.strided, device=None, pin_memory=False, memory_format=torch.contiguous_format) → Tensor +``` + +Returns a tensor filled with uninitialized data on `device="tvarant"`. + +## Parameters + +- **size** (`int...`) – Size of the returned tensor. +- **dtype** (`torch.dtype, optional`) – Desired data type. Defaults to the global default (prefer `float32`). +- **device** (`torch.device, optional`) – Use `"tvarant"` / `"tvarant:0"`. +- **memory_format** (`torch.memory_format, optional`) – Memory format for the result. + +## Shape + +- Output: $(\textit{size})$ + +## Example + +```python +x = torch.empty(2, 3, device='tvarant', dtype=torch.float32) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/index.md b/docs/api/ops/index.md new file mode 100644 index 0000000..842f008 --- /dev/null +++ b/docs/api/ops/index.md @@ -0,0 +1,42 @@ +# Supported ops + +ATen and custom ops registered for `PrivateUse1` (`device="tvarant"`). +Unlisted ops **fall back to CPU** (tensors are copied as needed). + +Click an op for signature, parameters, shape, and examples. + +## Factories & memory + +| Op | Page | +|---|---| +| `empty` / `empty_strided` | [empty](empty.md) | +| `zeros` / `ones` / `full` | [zeros / ones / full](zeros.md) | +| `clone` / `contiguous` | [clone / contiguous](clone.md) | + +## Pointwise + +| Op | Page | +|---|---| +| `add` | [add](add.md) | +| `mul` | [mul](mul.md) | +| `sub` | [sub](sub.md) | +| `relu` | [relu](relu.md) | +| `silu` | [silu](silu.md) | +| `neg` / `abs` / `exp` / `sqrt` / … | [unary](neg.md) | + +## Linear algebra & NN + +| Op | Page | +|---|---| +| `mm` / `addmm` / `matmul` / `bmm` | [mm family](mm.md) | +| `linear` | [linear](linear.md) | +| `softmax` | [softmax](softmax.md) | +| `layer_norm` | [layer_norm](layer_norm.md) | +| `embedding` | [embedding](embedding.md) | + +## Custom (`torch.ops.tvarant`) + +| Op | Page | +|---|---| +| `linear_act` | [linear_act](linear_act.md) | +| `pointwise` | [pointwise](pointwise.md) | diff --git a/docs/api/ops/layer_norm.md b/docs/api/ops/layer_norm.md new file mode 100644 index 0000000..c08fc0e --- /dev/null +++ b/docs/api/ops/layer_norm.md @@ -0,0 +1,28 @@ +# layer_norm + +```python +torch.nn.functional.layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05) → Tensor +``` + +Applies Layer Normalization over the last dimensions described by `normalized_shape` (via `native_layer_norm`). + +## Parameters + +- **input** (`Tensor`) – Input tensor. +- **normalized_shape** (`list of int`) – Shape over which to normalize. +- **weight, bias** (`Tensor, optional`) – Affine parameters. +- **eps** (`float`) – Numerical stability epsilon. + +## Shape + +- Output: same as input + +## Example + +```python +y = torch.nn.functional.layer_norm(x, x.shape[-1:]) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/linear.md b/docs/api/ops/linear.md new file mode 100644 index 0000000..231e37a --- /dev/null +++ b/docs/api/ops/linear.md @@ -0,0 +1,32 @@ +# linear + +```python +torch.nn.functional.linear(input, weight, bias=None) → Tensor +``` + +Applies an affine linear transformation with a **transposed** weight layout matching `nn.Linear`: + +$$ +y = x A^T + b +$$ + +## Parameters + +- **input** (`Tensor`) – Input of shape $(*, in\_features)$. +- **weight** (`Tensor`) – Weight of shape $(out\_features, in\_features)$. +- **bias** (`Tensor, optional`) – Bias of shape $(out\_features)$. + +## Shape + +- Input: $(*, in)$ +- Output: $(*, out)$ + +## Example + +```python +y = torch.nn.functional.linear(x, w, b) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/linear_act.md b/docs/api/ops/linear_act.md new file mode 100644 index 0000000..feec5dd --- /dev/null +++ b/docs/api/ops/linear_act.md @@ -0,0 +1,29 @@ +# torch.ops.tvarant.linear_act + +```python +torch.ops.tvarant.linear_act(x, weight, bias, act, trans_b=False) → Tensor +``` + +Fused GEMM + optional bias + activation. Used by the graph compiler for Linear+ReLU/SiLU epilogues. + +## Parameters + +- **x** (`Tensor`) – Activations `(..., K)`. +- **weight** (`Tensor`) – Weight matrix. +- **bias** (`Tensor or None`) – Optional bias. +- **act** (`str`) – `"none"`, `"relu"`, or `"silu"`. +- **trans_b** (`bool`) – If `True`, treat weight like `nn.Linear` (`A^T`). Default: `False`. + +## Shape + +- Output: `(..., N)` where `N` is the output feature size + +## Example + +```python +y = torch.ops.tvarant.linear_act(x, w, b, 'relu', True) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/mm.md b/docs/api/ops/mm.md new file mode 100644 index 0000000..391d9d2 --- /dev/null +++ b/docs/api/ops/mm.md @@ -0,0 +1,31 @@ +# mm / addmm / matmul / bmm + +```python +torch.mm(input, mat2) → Tensor +torch.addmm(input, mat1, mat2, *, beta=1, alpha=1) → Tensor +torch.matmul(input, other) → Tensor +torch.bmm(input, mat2) → Tensor +``` + +Matrix products. `mm`/`addmm` use the fused GEMM path; `matmul`/`bmm` cover 2D and batched cases. + +## Parameters + +- **input / mat1** (`Tensor`) – Left matrix / batch. +- **mat2 / other** (`Tensor`) – Right matrix / batch. +- **beta, alpha** (`Number`) – `addmm`: `beta*input + alpha*(mat1@mat2)`. + +## Shape + +- `mm`: $(n,m) \times (m,p) \rightarrow (n,p)$ +- `bmm`: $(B,n,m) \times (B,m,p) \rightarrow (B,n,p)$ + +## Example + +```python +y = torch.mm(a, b) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/mul.md b/docs/api/ops/mul.md new file mode 100644 index 0000000..deefc07 --- /dev/null +++ b/docs/api/ops/mul.md @@ -0,0 +1,27 @@ +# mul + +```python +torch.mul(input, other, *, out=None) → Tensor +input * other +``` + +Multiplies `input` by `other` element-wise (tensor or scalar). + +## Parameters + +- **input** (`Tensor`) – First operand. +- **other** (`Tensor or Number`) – Second operand. + +## Shape + +- Output: broadcast shape + +## Example + +```python +y = a * b +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/neg.md b/docs/api/ops/neg.md new file mode 100644 index 0000000..ad10ef6 --- /dev/null +++ b/docs/api/ops/neg.md @@ -0,0 +1,32 @@ +# neg / abs / exp / sqrt + +```python +torch.neg(input) → Tensor +torch.abs(input) → Tensor +torch.exp(input) → Tensor +torch.exp2(input) → Tensor +torch.sqrt(input) → Tensor +torch.rsqrt(input) → Tensor +``` + +Unary element-wise math ops on fp32 Tvarant tensors (host + OpenCL kernels). + +## Parameters + +- **input** (`Tensor`) – Input tensor. + +## Shape + +- Input: $(*)$ +- Output: $(*)$ + +## Example + +```python +y = torch.exp(x) +z = torch.rsqrt(torch.abs(x) + 1e-6) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/pointwise.md b/docs/api/ops/pointwise.md new file mode 100644 index 0000000..54ae10b --- /dev/null +++ b/docs/api/ops/pointwise.md @@ -0,0 +1,31 @@ +# torch.ops.tvarant.pointwise + +```python +torch.ops.tvarant.pointwise(inputs, ops, a, b, input_ids, alphas, consts) → Tensor +``` + +Runs a JIT pointwise SSA program over equal-numel Tvarant tensors. Prefer `torch_tvarant.compiler.compile` rather than building programs by hand. + +!!! note + + Opcode IDs must match `csrc/jit/Jit.h` (`LOAD`, `ADD`, `MUL`, `RELU`, …). + +## Parameters + +- **inputs** (`List[Tensor]`) – Input tensors (same numel). +- **ops, a, b, input_ids** (`List[int]`) – SSA node opcodes and operands. +- **alphas, consts** (`List[float]`) – Per-node scalars. + +## Shape + +- Output: same shape as `inputs[0]` + +## Example + +```python +# Emitted by torch_tvarant.compiler.fuse_pointwise — not usually called directly. +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/relu.md b/docs/api/ops/relu.md new file mode 100644 index 0000000..6009162 --- /dev/null +++ b/docs/api/ops/relu.md @@ -0,0 +1,31 @@ +# relu + +```python +torch.relu(input) → Tensor +torch.nn.functional.relu(input, inplace=False) → Tensor +``` + +Applies the rectified linear unit function element-wise: + +$$ +\text{ReLU}(x) = \max(0, x) +$$ + +## Parameters + +- **input** (`Tensor`) – Input tensor on `device="tvarant"`. + +## Shape + +- Input: $(*)$ +- Output: $(*)$, same as input + +## Example + +```python +y = torch.relu(x) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/silu.md b/docs/api/ops/silu.md new file mode 100644 index 0000000..f6b1109 --- /dev/null +++ b/docs/api/ops/silu.md @@ -0,0 +1,30 @@ +# silu + +```python +torch.nn.functional.silu(input, inplace=False) → Tensor +``` + +Applies the Sigmoid Linear Unit (SiLU / Swish): + +$$ +\text{SiLU}(x) = x \cdot \sigma(x) +$$ + +## Parameters + +- **input** (`Tensor`) – Input tensor. + +## Shape + +- Input: $(*)$ +- Output: $(*)$ + +## Example + +```python +y = torch.nn.functional.silu(x) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/softmax.md b/docs/api/ops/softmax.md new file mode 100644 index 0000000..fe257f1 --- /dev/null +++ b/docs/api/ops/softmax.md @@ -0,0 +1,27 @@ +# softmax + +```python +torch.nn.functional.softmax(input, dim=None, dtype=None) → Tensor +``` + +Applies softmax along `dim` on Tvarant tensors. + +## Parameters + +- **input** (`Tensor`) – Input tensor. +- **dim** (`int`) – Dimension along which softmax is computed. + +## Shape + +- Input: $(*)$ +- Output: $(*)$ + +## Example + +```python +p = torch.nn.functional.softmax(logits, dim=-1) +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/sub.md b/docs/api/ops/sub.md new file mode 100644 index 0000000..b6b2b55 --- /dev/null +++ b/docs/api/ops/sub.md @@ -0,0 +1,28 @@ +# sub + +```python +torch.sub(input, other, *, alpha=1, out=None) → Tensor +input - other +``` + +Subtracts `alpha * other` from `input` (implemented via negated `add` when registered). + +## Parameters + +- **input** (`Tensor`) – Minuend. +- **other** (`Tensor or Number`) – Subtrahend. +- **alpha** (`Number`) – Multiplier for `other`. Default: `1`. + +## Shape + +- Output: broadcast shape + +## Example + +```python +y = a - b +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/ops/zeros.md b/docs/api/ops/zeros.md new file mode 100644 index 0000000..25d76bd --- /dev/null +++ b/docs/api/ops/zeros.md @@ -0,0 +1,35 @@ +# zeros / ones / full + +```python +torch.zeros(*size, *, dtype=None, device=None, ...) → Tensor +torch.ones(*size, *, dtype=None, device=None, ...) → Tensor +torch.full(size, fill_value, *, dtype=None, device=None, ...) → Tensor +``` + +Factory ops that allocate on Tvarant then fill via `fill_`. `*_like` variants match another tensor's shape/device. + +!!! note + + On older builds without PrivateUse1 factory impls, factories may allocate via empty+fill after the extended-ops merge. + +## Parameters + +- **size** (`int... or size`) – Output shape. +- **fill_value** (`Number`) – Constant for `full` / `full_like`. +- **dtype** (`torch.dtype, optional`) – Desired dtype (fp32 path is fully accelerated). + +## Shape + +- Output: same as `size` / like-tensor + +## Example + +```python +z = torch.zeros(2, 3, device='tvarant') +o = torch.ones_like(z) +f = torch.full((2, 3), 2.5, device='tvarant') +``` + +--- + +[← All supported ops](index.md) diff --git a/docs/api/torch.md b/docs/api/torch.md deleted file mode 100644 index 4cc7d73..0000000 --- a/docs/api/torch.md +++ /dev/null @@ -1,65 +0,0 @@ -# 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/api/torch/backend.md b/docs/api/torch/backend.md new file mode 100644 index 0000000..ca187bf --- /dev/null +++ b/docs/api/torch/backend.md @@ -0,0 +1,18 @@ +# backend + +```python +torch.tvarant.backend() → str +``` + +Returns the active runtime name (`"sim"` or `"opencl"`). + +::: torch_tvarant.tvarant.backend + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/current_device.md b/docs/api/torch/current_device.md new file mode 100644 index 0000000..5ebd6a1 --- /dev/null +++ b/docs/api/torch/current_device.md @@ -0,0 +1,18 @@ +# current_device + +```python +torch.tvarant.current_device() → int +``` + +Returns the current device index. + +::: torch_tvarant.tvarant.current_device + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/device_count.md b/docs/api/torch/device_count.md new file mode 100644 index 0000000..e5d1e55 --- /dev/null +++ b/docs/api/torch/device_count.md @@ -0,0 +1,18 @@ +# device_count + +```python +torch.tvarant.device_count() → int +``` + +Returns the number of logical Tvarant devices. + +::: torch_tvarant.tvarant.device_count + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/get_rng_state.md b/docs/api/torch/get_rng_state.md new file mode 100644 index 0000000..4cdaafb --- /dev/null +++ b/docs/api/torch/get_rng_state.md @@ -0,0 +1,18 @@ +# get_rng_state + +```python +torch.tvarant.get_rng_state(device='tvarant') → Tensor +``` + +Returns the RNG state tensor for the Tvarant generator. + +::: torch_tvarant.tvarant.get_rng_state + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/index.md b/docs/api/torch/index.md new file mode 100644 index 0000000..65318d2 --- /dev/null +++ b/docs/api/torch/index.md @@ -0,0 +1,37 @@ +# torch.tvarant + +Importing `torch_tvarant` registers the **tvarant** PrivateUse1 backend and +exposes helpers under `torch.tvarant` (same idea as `torch.cuda` / `torch.mps`). + +```python +import torch +import torch_tvarant + +assert torch.tvarant.is_available() +x = torch.empty(2, 3, device="tvarant") +``` + +Click an API below for the full signature, parameters, and examples. + +## Device + +| API | Summary | +|---|---| +| [`is_available`](is_available.md) | Whether the backend can run | +| [`is_initialized`](is_initialized.md) | Lazy init completed | +| [`device_count`](device_count.md) | Number of logical devices | +| [`current_device`](current_device.md) | Current device index | +| [`set_device`](set_device.md) | Select device index | +| [`synchronize`](synchronize.md) | Wait for kernels | +| [`backend`](backend.md) | `"sim"` or `"opencl"` | + +## RNG + +| API | Summary | +|---|---| +| [`manual_seed`](manual_seed.md) | Seed Tvarant RNG | +| [`manual_seed_all`](manual_seed_all.md) | Seed all devices | +| [`get_rng_state`](get_rng_state.md) | Read RNG state | +| [`set_rng_state`](set_rng_state.md) | Restore RNG state | + +See also: [`torch_tvarant.compiler`](../compiler/index.md). diff --git a/docs/api/torch/is_available.md b/docs/api/torch/is_available.md new file mode 100644 index 0000000..f0c5e8a --- /dev/null +++ b/docs/api/torch/is_available.md @@ -0,0 +1,18 @@ +# is_available + +```python +torch.tvarant.is_available() → bool +``` + +Returns whether the Tvarant backend extension is loaded and usable. + +::: torch_tvarant.tvarant.is_available + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/is_initialized.md b/docs/api/torch/is_initialized.md new file mode 100644 index 0000000..307cbb2 --- /dev/null +++ b/docs/api/torch/is_initialized.md @@ -0,0 +1,18 @@ +# is_initialized + +```python +torch.tvarant.is_initialized() → bool +``` + +Returns whether lazy device initialization has completed. + +::: torch_tvarant.tvarant.is_initialized + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/manual_seed.md b/docs/api/torch/manual_seed.md new file mode 100644 index 0000000..eafdb2b --- /dev/null +++ b/docs/api/torch/manual_seed.md @@ -0,0 +1,18 @@ +# manual_seed + +```python +torch.tvarant.manual_seed(seed) → None +``` + +Seeds the Tvarant generator only (does not call `torch.manual_seed`). + +::: torch_tvarant.tvarant.manual_seed + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/manual_seed_all.md b/docs/api/torch/manual_seed_all.md new file mode 100644 index 0000000..b6dbf5f --- /dev/null +++ b/docs/api/torch/manual_seed_all.md @@ -0,0 +1,18 @@ +# manual_seed_all + +```python +torch.tvarant.manual_seed_all(seed) → None +``` + +Seeds all Tvarant devices (currently one). + +::: torch_tvarant.tvarant.manual_seed_all + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/set_device.md b/docs/api/torch/set_device.md new file mode 100644 index 0000000..1b5d87b --- /dev/null +++ b/docs/api/torch/set_device.md @@ -0,0 +1,18 @@ +# set_device + +```python +torch.tvarant.set_device(device) → None +``` + +Sets the current device index. + +::: torch_tvarant.tvarant.set_device + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/set_rng_state.md b/docs/api/torch/set_rng_state.md new file mode 100644 index 0000000..d6fee13 --- /dev/null +++ b/docs/api/torch/set_rng_state.md @@ -0,0 +1,18 @@ +# set_rng_state + +```python +torch.tvarant.set_rng_state(new_state, device='tvarant') → None +``` + +Restores RNG state previously returned by `get_rng_state`. + +::: torch_tvarant.tvarant.set_rng_state + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/api/torch/synchronize.md b/docs/api/torch/synchronize.md new file mode 100644 index 0000000..78cadac --- /dev/null +++ b/docs/api/torch/synchronize.md @@ -0,0 +1,18 @@ +# synchronize + +```python +torch.tvarant.synchronize(device=None) → None +``` + +Waits for outstanding work on the Tvarant runtime. + +::: torch_tvarant.tvarant.synchronize + options: + show_root_heading: false + show_root_toc_entry: false + show_source: true + show_signature: true + show_signature_annotations: true + separate_signature: true + docstring_section_style: table + heading_level: 2 diff --git a/docs/getting-started.md b/docs/getting-started.md index 1cf4eb7..7cfb037 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -47,8 +47,8 @@ See [JIT Compiler](jit-compiler.md) for fusion details. | Topic | Page | |---|---| -| Device helpers | [torch.tvarant](api/torch.md) | -| Supported ops | [Supported ops](api/ops.md) | +| Device helpers | [torch.tvarant](api/torch/index.md) | +| Supported ops | [Supported ops](api/ops/index.md) | | Kernels | [Kernels](api/kernels.md) | | C++ extension | [C++ API](api/cpp.md) | | FPGA path | [FPGA / OpenCL](fpga-opencl.md) | diff --git a/docs/operations.md b/docs/operations.md index b6b0b96..ef91f5b 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -2,6 +2,6 @@ This page moved into the API reference: -- [Supported ops](api/ops.md) +- [Supported ops](api/ops/index.md) - [Kernels](api/kernels.md) - [C++ extension](api/cpp.md) diff --git a/mkdocs.yml b/mkdocs.yml index cdd6c6d..d6ee0a1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -23,6 +23,9 @@ theme: features: - navigation.tabs - navigation.sections + - navigation.expand + - navigation.indexes + - toc.follow - content.code.copy - search.suggest @@ -42,8 +45,46 @@ nav: - Contributing: contributing.md - API: - Overview: api/index.md - - torch.tvarant: api/torch.md - - Supported ops: api/ops.md + - torch.tvarant: + - api/torch/index.md + - is_available: api/torch/is_available.md + - is_initialized: api/torch/is_initialized.md + - device_count: api/torch/device_count.md + - current_device: api/torch/current_device.md + - set_device: api/torch/set_device.md + - synchronize: api/torch/synchronize.md + - backend: api/torch/backend.md + - manual_seed: api/torch/manual_seed.md + - manual_seed_all: api/torch/manual_seed_all.md + - get_rng_state: api/torch/get_rng_state.md + - set_rng_state: api/torch/set_rng_state.md + - compiler: + - api/compiler/index.md + - compile: api/compiler/compile.md + - compile_fx: api/compiler/compile_fx.md + - trace_module: api/compiler/trace_module.md + - fuse_gemm_epilogue: api/compiler/fuse_gemm_epilogue.md + - fuse_pointwise: api/compiler/fuse_pointwise.md + - register: api/compiler/register.md + - TvarantTracer: api/compiler/TvarantTracer.md + - Supported ops: + - api/ops/index.md + - empty: api/ops/empty.md + - zeros / ones / full: api/ops/zeros.md + - clone / contiguous: api/ops/clone.md + - add: api/ops/add.md + - mul: api/ops/mul.md + - sub: api/ops/sub.md + - relu: api/ops/relu.md + - silu: api/ops/silu.md + - unary: api/ops/neg.md + - mm family: api/ops/mm.md + - linear: api/ops/linear.md + - softmax: api/ops/softmax.md + - layer_norm: api/ops/layer_norm.md + - embedding: api/ops/embedding.md + - linear_act: api/ops/linear_act.md + - pointwise: api/ops/pointwise.md - Kernels: api/kernels.md - C++: api/cpp.md - Operations (redirect): operations.md @@ -54,6 +95,8 @@ markdown_extensions: - pymdownx.superfences - pymdownx.highlight - pymdownx.inlinehilite + - pymdownx.arithmatex: + generic: true - toc: permalink: true @@ -64,9 +107,17 @@ plugins: python: options: docstring_style: google - show_source: false + show_source: true + show_signature: true + show_signature_annotations: true separate_signature: true merge_init_into_class: true + docstring_section_style: table + show_symbol_type_heading: true + show_symbol_type_toc: true + +extra_javascript: + - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js extra: social: diff --git a/torch_tvarant/compiler.py b/torch_tvarant/compiler.py index 018df06..7e276f4 100644 --- a/torch_tvarant/compiler.py +++ b/torch_tvarant/compiler.py @@ -49,7 +49,15 @@ 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.""" + """FX-trace ``module``, inlining Linear / common activations. + + Args: + module: ``torch.nn.Module`` to trace. + + Returns: + torch.fx.GraphModule: Graph with inlined leaf modules so the fuser + sees ATen ops (``linear``, ``relu``, …). + """ tracer = TvarantTracer() graph = tracer.trace(module) return torch.fx.GraphModule(module, graph) @@ -315,7 +323,17 @@ def rec_input(src: object) -> int: def compile_fx(gm: torch.fx.GraphModule, example_inputs: Any = None) -> torch.fx.GraphModule: - """Optimize a captured FX graph for Tvarant.""" + """Optimize a captured FX graph for Tvarant. + + Runs GEMM-epilogue fusion then pointwise fusion. + + Args: + gm: FX graph module to optimize in place / recompile. + example_inputs: Unused; accepted for Dynamo backend compatibility. + + Returns: + torch.fx.GraphModule: Fused graph ready to run on ``device="tvarant"``. + """ _ = example_inputs gm = fuse_gemm_epilogue(gm) gm = fuse_pointwise(gm) @@ -330,7 +348,11 @@ def _dynamo_backend(gm: torch.fx.GraphModule, example_inputs: Any): def register() -> None: - """Register ``backend="tvarant"`` with torch.compile / Dynamo.""" + """Register ``backend="tvarant"`` with ``torch.compile`` / Dynamo. + + Safe to call multiple times; subsequent calls are no-ops. Invoked + automatically on ``import torch_tvarant.compiler``. + """ global _registered if _registered: return @@ -357,6 +379,18 @@ def compile(module: torch.nn.Module, *, dynamic: bool = False) -> Callable: Prefers FX tracing that inlines Linear / activations (stable on this custom device). Falls back to ``torch.compile(..., backend="tvarant")``. + + Args: + module: Module to optimize. Switched to ``eval()`` mode. + dynamic: Forwarded to ``torch.compile`` on the fallback path only. + + Returns: + Callable: Callable module / optimized callable producing the same + outputs as ``module`` on ``device="tvarant"``. + + Examples: + >>> compiled = torch_tvarant.compiler.compile(model) + >>> y = compiled(x) """ module.eval() try: diff --git a/torch_tvarant/tvarant/__init__.py b/torch_tvarant/tvarant/__init__.py index 1182bd7..f399bab 100644 --- a/torch_tvarant/tvarant/__init__.py +++ b/torch_tvarant/tvarant/__init__.py @@ -36,38 +36,86 @@ def _lazy_init() -> None: def is_available() -> bool: - """Return whether the Tvarant backend extension is usable.""" + """Return whether the Tvarant backend extension is usable. + + After a successful ``import torch_tvarant``, this is normally ``True``. + + Returns: + bool: ``True`` if kernels can run on ``device="tvarant"``. + + Examples: + >>> import torch_tvarant + >>> torch.tvarant.is_available() + True + """ return bool(_C.is_available()) def is_initialized() -> bool: - """Return whether lazy device initialization has completed.""" + """Return whether lazy device initialization has completed. + + Returns: + bool: ``True`` after the first device API call that triggers init. + """ return _initialized def device_count() -> int: - """Return the number of logical Tvarant devices.""" + """Return the number of logical Tvarant devices. + + Returns: + int: Device count (currently ``1`` for both sim and OpenCL backends). + + Examples: + >>> torch.tvarant.device_count() + 1 + """ return int(_C.device_count()) def current_device() -> int: - """Return the current device index.""" + """Return the current device index. + + Returns: + int: Active device ordinal (``0`` today). + """ return int(_C.current_device()) def set_device(device: int) -> None: - """Set the current device index.""" + """Set the current device index. + + Args: + device: Device ordinal. Must be in ``[0, device_count())``. + + Examples: + >>> torch.tvarant.set_device(0) + """ _C.set_device(int(device)) def synchronize(device: int | None = None) -> None: - """Wait for outstanding work on the runtime (``device`` reserved).""" + """Wait for outstanding work on the Tvarant runtime. + + Args: + device: Reserved for API compatibility with other PyTorch device + modules. Ignored; all work shares one runtime queue today. + """ _ = device _C.synchronize() def backend() -> str: - """Return the active runtime name (``sim`` or ``opencl``).""" + """Return the active runtime name. + + Returns: + str: ``"sim"`` for the CPU simulator, or ``"opencl"`` when built and + selected via ``TVARANT_BACKEND=opencl``. + + Examples: + >>> torch.tvarant.backend() + 'sim' + """ return str(_C.backend()) @@ -76,17 +124,36 @@ def _is_in_bad_fork() -> bool: def manual_seed(seed: int) -> None: - """Seed the Tvarant generator only (does not call ``torch.manual_seed``).""" + """Seed the Tvarant generator only. + + Does **not** call ``torch.manual_seed``, which would re-enter + ``torch.tvarant.manual_seed_all`` and recurse. + + Args: + seed: Integer seed for the PrivateUse1 default generator. + """ _C.manual_seed(int(seed)) def manual_seed_all(seed: int) -> None: - """Seed all Tvarant devices (currently one device).""" + """Seed all Tvarant devices. + + Args: + seed: Integer seed applied to every logical device (currently one). + """ manual_seed(seed) def get_rng_state(device: Union[int, str, torch.device] = "tvarant") -> torch.Tensor: - """Return the RNG state tensor for ``device``.""" + """Return the RNG state tensor for ``device``. + + Args: + device: Device specifier. Accepted for API compatibility; only the + single Tvarant generator is used today. + + Returns: + torch.Tensor: Opaque byte state suitable for ``set_rng_state``. + """ _ = device return _C.get_rng_state() @@ -94,6 +161,11 @@ 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``.""" + """Restore RNG state previously returned by ``get_rng_state``. + + Args: + new_state: State tensor from ``get_rng_state``. + device: Device specifier (compatibility; ignored beyond validation). + """ _ = device _C.set_rng_state(new_state)