From 21cb672ec6f67621054865e93d7f6ecab6392bbc Mon Sep 17 00:00:00 2001 From: lcy-seso Date: Thu, 1 Oct 2026 15:55:16 +0800 Subject: [PATCH] [Docs] Add the manifest and kernel dispatch user guides, and polish the site's prose The manifest guide and the kernel dispatch guide were written and reviewed in TileOpsGov. They move here under docs/user-guide/, Chinese first with English translations, with figure sources and a render script under figures/user-guide/. The old manifest page is replaced by the new guide. The kernel selection part of the new-op page and the three ways of the backends page now link to the dispatch guide instead of repeating it. The remaining Chinese pages are polished to the guides' style and terminology (op, family, kernel interface, in-tree implementation, backend), and the English pages are then aligned with the Chinese and polished. --- CLAUDE.md | 12 +- docs/api/index.md | 2 +- docs/backends.md | 531 ++++++++---------- docs/backends.zh.md | 369 ++++++------ docs/index.md | 48 +- docs/index.zh.md | 21 +- docs/manifest.md | 291 ---------- docs/manifest.zh.md | 230 -------- docs/new-op.md | 284 +++++----- docs/new-op.zh.md | 140 ++--- docs/performance-guides/index.md | 8 +- docs/performance-guides/index.zh.md | 8 +- .../memory-bound/global-memory-access.md | 108 ++-- .../memory-bound/global-memory-access.zh.md | 56 +- docs/performance-guides/memory-bound/index.md | 5 +- .../memory-bound/index.zh.md | 4 +- .../memory-bound/shared-memory-access.md | 115 ++-- .../memory-bound/shared-memory-access.zh.md | 64 +-- docs/timing.md | 185 +++--- docs/timing.zh.md | 80 +-- docs/torch-compile.md | 350 ++++++------ docs/torch-compile.zh.md | 160 +++--- docs/user-guide/dispatch/backends.md | 111 ++++ docs/user-guide/dispatch/backends.zh.md | 111 ++++ docs/user-guide/dispatch/img/call.svg | 1 + docs/user-guide/dispatch/img/call.zh.svg | 1 + docs/user-guide/dispatch/img/extension.svg | 1 + docs/user-guide/dispatch/img/extension.zh.svg | 1 + docs/user-guide/dispatch/index.md | 74 +++ docs/user-guide/dispatch/index.zh.md | 74 +++ docs/user-guide/dispatch/writing.md | 201 +++++++ docs/user-guide/dispatch/writing.zh.md | 201 +++++++ docs/user-guide/index.md | 11 +- docs/user-guide/index.zh.md | 11 +- docs/user-guide/manifest/calls.md | 135 +++++ docs/user-guide/manifest/calls.zh.md | 135 +++++ docs/user-guide/manifest/concepts.md | 154 +++++ docs/user-guide/manifest/concepts.zh.md | 154 +++++ docs/user-guide/manifest/examples.md | 226 ++++++++ docs/user-guide/manifest/examples.zh.md | 226 ++++++++ docs/user-guide/manifest/extensions.md | 253 +++++++++ docs/user-guide/manifest/extensions.zh.md | 253 +++++++++ docs/user-guide/manifest/img/call.svg | 1 + docs/user-guide/manifest/img/call.zh.svg | 1 + docs/user-guide/manifest/img/device.svg | 1 + docs/user-guide/manifest/img/device.zh.svg | 1 + docs/user-guide/manifest/img/new-op.svg | 1 + docs/user-guide/manifest/img/new-op.zh.svg | 1 + docs/user-guide/manifest/img/overview.svg | 1 + docs/user-guide/manifest/img/overview.zh.svg | 1 + docs/user-guide/manifest/index.md | 152 +++++ docs/user-guide/manifest/index.zh.md | 152 +++++ docs/user-guide/manifest/writing.md | 220 ++++++++ docs/user-guide/manifest/writing.zh.md | 216 +++++++ figures/user-guide/dispatch/call.en.puml | 65 +++ figures/user-guide/dispatch/call.zh.puml | 65 +++ figures/user-guide/dispatch/extension.en.puml | 54 ++ figures/user-guide/dispatch/extension.zh.puml | 54 ++ figures/user-guide/manifest/call.en.puml | 92 +++ figures/user-guide/manifest/call.zh.puml | 91 +++ figures/user-guide/manifest/device.en.puml | 56 ++ figures/user-guide/manifest/device.zh.puml | 55 ++ figures/user-guide/manifest/new-op.en.puml | 68 +++ figures/user-guide/manifest/new-op.zh.puml | 62 ++ figures/user-guide/manifest/overview.py | 257 +++++++++ figures/user-guide/render.sh | 42 ++ mkdocs.yml | 32 +- 67 files changed, 5310 insertions(+), 1836 deletions(-) delete mode 100644 docs/manifest.md delete mode 100644 docs/manifest.zh.md create mode 100644 docs/user-guide/dispatch/backends.md create mode 100644 docs/user-guide/dispatch/backends.zh.md create mode 100644 docs/user-guide/dispatch/img/call.svg create mode 100644 docs/user-guide/dispatch/img/call.zh.svg create mode 100644 docs/user-guide/dispatch/img/extension.svg create mode 100644 docs/user-guide/dispatch/img/extension.zh.svg create mode 100644 docs/user-guide/dispatch/index.md create mode 100644 docs/user-guide/dispatch/index.zh.md create mode 100644 docs/user-guide/dispatch/writing.md create mode 100644 docs/user-guide/dispatch/writing.zh.md create mode 100644 docs/user-guide/manifest/calls.md create mode 100644 docs/user-guide/manifest/calls.zh.md create mode 100644 docs/user-guide/manifest/concepts.md create mode 100644 docs/user-guide/manifest/concepts.zh.md create mode 100644 docs/user-guide/manifest/examples.md create mode 100644 docs/user-guide/manifest/examples.zh.md create mode 100644 docs/user-guide/manifest/extensions.md create mode 100644 docs/user-guide/manifest/extensions.zh.md create mode 100644 docs/user-guide/manifest/img/call.svg create mode 100644 docs/user-guide/manifest/img/call.zh.svg create mode 100644 docs/user-guide/manifest/img/device.svg create mode 100644 docs/user-guide/manifest/img/device.zh.svg create mode 100644 docs/user-guide/manifest/img/new-op.svg create mode 100644 docs/user-guide/manifest/img/new-op.zh.svg create mode 100644 docs/user-guide/manifest/img/overview.svg create mode 100644 docs/user-guide/manifest/img/overview.zh.svg create mode 100644 docs/user-guide/manifest/index.md create mode 100644 docs/user-guide/manifest/index.zh.md create mode 100644 docs/user-guide/manifest/writing.md create mode 100644 docs/user-guide/manifest/writing.zh.md create mode 100644 figures/user-guide/dispatch/call.en.puml create mode 100644 figures/user-guide/dispatch/call.zh.puml create mode 100644 figures/user-guide/dispatch/extension.en.puml create mode 100644 figures/user-guide/dispatch/extension.zh.puml create mode 100644 figures/user-guide/manifest/call.en.puml create mode 100644 figures/user-guide/manifest/call.zh.puml create mode 100644 figures/user-guide/manifest/device.en.puml create mode 100644 figures/user-guide/manifest/device.zh.puml create mode 100644 figures/user-guide/manifest/new-op.en.puml create mode 100644 figures/user-guide/manifest/new-op.zh.puml create mode 100644 figures/user-guide/manifest/overview.py create mode 100755 figures/user-guide/render.sh diff --git a/CLAUDE.md b/CLAUDE.md index 32c55b82..1be1a475 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -77,21 +77,23 @@ implementation of the same op on that workload. English at the site root, Chinese under `/zh/`. A Chinese page is a `.zh.md` beside the English `.md`, full prose, never an -`include-markdown` shell. `backends.md`, `torch-compile.md` and everything under -`performance-guides/memory-bound/` were authored in Chinese: edit the `.zh.md` -first, then bring the English page in line. Everything else goes the other way. +`include-markdown` shell. `backends.md`, `torch-compile.md`, everything under +`performance-guides/memory-bound/` and the two guides under `user-guide/manifest/` +and `user-guide/dispatch/` were authored in Chinese: edit the `.zh.md` first, then +bring the English page in line. Everything else goes the other way. | Rule | Detail | |------|--------| | Coverage | Whichever pages have a `.zh.md` — `ls docs/**/*.zh.md` | | Never translate | `api/` and `benchmarks/`, both generated; `design/`, mirrored English | | Missing translation | Falls back to English at the same URL, and `hooks.py` prepends a "本页暂无中文版" notice. The fallback runs zh → en only: a page that exists only as `.zh.md` leaves its `nav` entry on a missing file and the English sidebar renders a dead link | -| Figures | A figure with text needs one SVG per language: translate the `` nodes and the `aria-label`, keep the geometry. English runs longer than Chinese — grow the `viewBox` rather than let text overflow | +| Figures | A figure with text needs one SVG per language: `img/.svg` for English and `img/.zh.svg` beside it, which the `zh` build picks up for the same reference. Translate the `` nodes and the `aria-label`, keep the geometry. English runs longer than Chinese — grow the `viewBox` rather than let text overflow. The user-guide figures are drawn from sources under `figures/user-guide/` (`.zh.puml`, `.en.puml`, `manifest/overview.py`): edit the source, then run `figures/user-guide/render.sh` | | Nav labels | `nav_translations` in the `i18n` plugin block; keep an entry for every `nav` title | | Chinese search | Requires `jieba` | | Punctuation | Full-width in Chinese prose: `，。：；（）`. Latin quotes and brackets stay half-width inside code spans | | Latin in Chinese | A space either side of a Latin token: `由 spec 驱动`, `形状和 dtype`. Not inside code spans | -| Keep in English | kernel, spec, agent, dtype, roofline, GEMM, target, and every op name. Why: translating them loses the link to the API | +| Keep in English | kernel, spec, agent, dtype, roofline, GEMM, target, op, family, backend, and every op name. Why: translating them loses the link to the API | +| Terms | 「kernel 接口」 for kernel interface, 「in-tree 实现」 for an in-tree implementation, 「build identity」 and 「构建函数」 for what `entry_for` returns | | Inline code | Real identifiers only (`GemmOp`, `eval_roofline`, paths, flags). A concept mentioned in prose is not code | | Type | `extra.css` gives `html[lang="zh"]` looser leading and headings at 700, not 800 — at 800 a CJK fallback face closes up the strokes. Scoped away from fallback pages, whose body text is English. No CJK webfont: Han glyphs come from the platform UI face (`--tf-cjk`) | diff --git a/docs/api/index.md b/docs/api/index.md index f4f021a7..ad23c654 100644 --- a/docs/api/index.md +++ b/docs/api/index.md @@ -45,7 +45,7 @@ Two things this reference does not carry: - **What each op is allowed to receive.** The authoritative dtype domains, shape rules and measured workloads are in the op's spec; see [Writing a - Spec](../manifest.md). + Spec](../user-guide/manifest/index.md). - **How fast it is.** Device time against the fastest alternative on each workload is on the [Benchmarks](../benchmarks/index.md) pages. diff --git a/docs/backends.md b/docs/backends.md index 499a10b1..141d1c13 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -1,148 +1,96 @@ # Adding a hardware backend TileLang is a multi-backend DSL: each kind of hardware has its own set of kernels, -shipped as its own Python package. TileOPs therefore defines a protocol under which a -package outside the repository takes over an op's kernel, replacing the implementation -TileOPs ships — with no change to TileOPs itself. +distributed as its own Python package. TileOPs therefore defines a protocol under which a +package outside the repository takes over an op's kernel in place of the in-tree +implementation, without any change to TileOPs. -This page is how a new class of hardware gets brought in, so that the ops on those -devices run your kernels. +This page describes how to bring in a new class of hardware, so that the ops on those +devices run that hardware's own kernels. -**A backend supplies one thing: something callable that computes this call.** -Everything else is the op layer's. +**A backend supplies one thing: a callable that computes this call.** The op layer does +everything else. -This page is about a target, the extension mechanism that covers a whole op. The first -half is the -work, in the order it is done: the four things to write, the protocol's four functions, -how one call reaches them, a backend that installs and runs as it stands, how to turn the -template into a backend for real hardware, the four rules for writing a kernel, what each -phase may do, and — after install — which state each op is in and what each error means. +This page covers the target, the mechanism that takes over a whole op. The first half +describes the backend author's work, in the order it is done: -The second half is why the protocol looks like this: the two layers of selection, the op -layer's contract, when a kernel is rebuilt, what a caller can reach for, and what the -protocol deliberately leaves out. +1. the four things to write; +1. the four functions of the protocol; +1. how one call reaches those functions; +1. a backend that installs and runs as it stands; +1. how to turn the template into a backend for real hardware; +1. the four rules for writing a kernel; +1. what each phase may do; +1. the state each op is in after install, and the cause behind each error message. -## Three ways to extend dispatch {#three-ways} - -A package outside TileOPs picks one of three mechanisms by how much of an op it takes -over. -The two smaller ones write a kernel class against a [kernel -interface](new-op.md#kernel-selection), the same contract the in-tree -implementations are written against; a target writes a `build_kernel` against the op's -manifest signature instead. - -| # | | `kernel_map=` | `register_implementation` | target | -| --- | --- | --- | --- | --- | -| 1 | Changes | the class registered under one key; which calls that key serves is unchanged | adds a key, with its own applicability and precedence | every call of the op | -| 2 | Applies to | the one op instance the caller constructed it on | every instance of that op constructed afterwards | every instance that settles on the target | -| 3 | Written against | the kernel interface | the kernel interface | the op's manifest signature | -| 4 | Calls the new class does not serve | an error when that key is selected | still served by the in-tree implementations | none: a target serves them all | - -**`kernel_map=`** is a constructor argument of every op, a mapping from key to class. It -replaces the class registered under that key in this instance; the key keeps the registered -implementation's `applies`, `general` and `preferred_over`, and the replacement, like -every implementation, inherits the key's interface and is built through its own -`entry_for`. A selected key whose replacement cannot serve the call is an error, never a -fall back to what it replaced. From -[`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py): - -```python -class _TorchLayerNorm(Kernel, LayerNormFwdInterface): - """A replacement written against ``LayerNormFwdInterface`` alone.""" - - devices = frozenset({torch.device(run_device()).type}) - - def __init__(self, n: int, eps: float) -> None: - super().__init__() - self.n, self.eps = n, eps - - @classmethod - def entry_for(cls, call: LayerNormCall): - return (call.n, call.eps), lambda: cls(call.n, call.eps) - - def forward(self, x, weight, bias): - return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( - x.dtype - ) - - -op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) -``` - -A key the op does not have, but another op does, is ignored, which is how a composite op -passes one mapping down to its sub-ops; a key no op has raises -`was given kernel_map keys no op has` at construction. +The second half explains why the protocol is designed this way: -**`register_implementation(op, key, implementation)`** adds an implementation instead of -replacing one. `op` is the op's manifest key, `key` the new implementation's dispatch key, -and the interface it joins is the one the class inherits. It declares its own region, so -it needs `preferred_over` where it overlaps an in-tree implementation that is not -`general`: - -```python -class _NarrowTorchLayerNorm(_TorchLayerNorm): - """An added implementation for short rows, which wins over the in-tree one there.""" - - preferred_over = frozenset({"layer_norm"}) - - @classmethod - def applies(cls, call: LayerNormCall) -> bool: - return call.n <= 64 +- the two layers of selection; +- the op layer's contract; +- when a kernel is rebuilt; +- the interfaces available to a caller; +- what the protocol deliberately does not support. +## Three ways to extend dispatch {#three-ways} -register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) -``` +A package outside TileOPs picks one of three mechanisms, by how much of an op it takes +over: -Here `n <= 64` goes to `_NarrowTorchLayerNorm` and `n = 1024` stays with the in-tree -`LayerNormKernel`. The added implementation applies only to op instances constructed -after the call. -Registering the same key twice under one op raises `BackendError`; a key an in-tree -implementation already uses raises `reuse keys it has` when an instance is constructed. +1. `kernel_map=`: the caller, when constructing an op, replaces the class that runs behind + one key; +1. `register_implementation`: adds an implementation to a kernel interface, which takes + part in selection alongside the in-tree implementations; +1. target: takes over every call of the op. -`register_implementation` runs when the backend module is imported, through the same -entry point a target uses; `kernel_map=` registers nothing and is passed by the caller at -construction. What follows is the target. +In the first two, the kernel class follows the kernel interface, the same contract the +in-tree implementations follow; how to write one is in +[How a backend joins TileOPs](user-guide/dispatch/backends.md). A target follows the op's +manifest signature instead, and the backend writes a `build_kernel` for it. The rest of +this page covers the target only. ## Four things to write | # | What to do | | --- | --- | -| 1 | An entry point in `pyproject.toml` pointing at the backend module | -| 2 | A target name and a `detect`, declaring which class of devices these kernels are for | -| 3 | A `build_kernel` for the first op you take over, written to its manifest signature | -| 4 | The `register_detector` and `register_kernel_builder` calls, at module top level | - -With those four written, `pip install` is all it takes. What follows: their signatures, -how one call reaches them, and a [complete backend](#runnable) written to these four steps, -installable as it stands. - -After the first op comes a `build_kernel` per op. **Every op the target model uses that -builds kernels of its own has to be covered** — a missing one is an error, with no fall -back to the implementation TileOPs ships, because those kernels cannot launch on this +| 1 | Declare an entry point in `pyproject.toml` that points at the backend module | +| 2 | Choose a target name and write a `detect`, declaring which class of devices these kernels are for | +| 3 | Pick the first op to take over, and write its `build_kernel` to the op's manifest signature | +| 4 | Call `register_detector` and `register_kernel_builder` at module top level | + +Once these four are written, `pip install` makes the backend take effect. The following +sections give: + +1. the signatures of the four functions; +1. how one call reaches them; +1. a [complete backend](#runnable) written to these four steps, which installs and runs as + it stands. + +After the first op, the backend adds a `build_kernel` per op. **Every op the target model +uses that builds kernels of its own must be covered.** A missing one is an error, with no +fall back to the in-tree implementation, because the in-tree kernels cannot launch on this target's devices. A composite op, which only runs sub-ops, needs no builder. ## The protocol: four functions -`tileops.backend` defines the outward-facing interface and contains no -implementation. The interface is expressed with Python structural typing -(`typing.Protocol`): a backend subclasses nothing and implements no abstract method, it -writes plain functions with matching signatures and registers them. The op layer checks a -backend's return value structurally too: `callable()`. +`tileops.backend` defines only the outward-facing interface, expressed with Python +structural typing (`typing.Protocol`). A backend subclasses no base class and implements +no abstract method; it writes plain functions with matching signatures and registers them. +The op layer checks a backend's return value structurally as well, with `callable()`. -A backend writes two functions (`detect`, `build_kernel`) and calls two to register -them (`register_detector`, `register_kernel_builder`), alongside the protocol's -`TensorSpec` and one entry point in `pyproject.toml`: +A backend implements `detect` and `build_kernel`, and registers them by calling +`register_detector` and `register_kernel_builder`. The protocol also defines `TensorSpec`, +and the backend declares one entry point in `pyproject.toml`: | # | Name | Written by | Called by, and when | | --- | --- | --- | --- | -| 1 | `detect` | implemented by the backend | asked of every target while the op layer picks one for a call | +| 1 | `detect` | implemented by the backend | called once for each target while the op layer settles the target for a call | | 2 | `build_kernel` | implemented by the backend | called by the op layer on a memo miss | | 3 | `register_detector` | the backend calls it | once, when the backend module is imported | -| 4 | `register_kernel_builder` | the backend calls it | likewise, once per op it takes over | +| 4 | `register_kernel_builder` | the backend calls it | when the backend module is imported, once per op it takes over | | — | `TensorSpec` | defined by the protocol | built by the op layer and passed into `build_kernel` | | — | the entry point | declared by the backend in `pyproject.toml` | enumerated by TileOPs when the first op is constructed | -Their signatures follow in that order, and the protocol's `TensorSpec` after them. +The signatures follow in that order, with the protocol's `TensorSpec` last. ### 1. `detect` @@ -150,9 +98,9 @@ Their signatures follow in that order, and the protocol's `TensorSpec` after the def detect(device: torch.device) -> bool: ... ``` -Implemented by the backend. Answers whether such a device is served by this set of -kernels: devices only, not dtypes or shapes; return `False` for someone else's device, and -do not raise. +Implemented by the backend. It answers whether this kind of device is served by this set +of kernels. It looks at the device only, not at dtypes or shapes. For a device that is not +its own, it returns `False` and does not raise. ```python # claiming a whole device type @@ -174,8 +122,8 @@ def build_kernel(*inputs: "TensorSpec | None", **params) -> Callable[..., Kernel Implemented by the backend, one per `(op, target)`. Its signature is the op's manifest signature: `inputs` correspond one-to-one to `signature.inputs` in declaration order, and -`params` are named after `signature.params`. An input declared optional that was not -passed on this call arrives as `None`. +`params` are named after `signature.params`. An optional input that was not passed on +this call arrives as `None`. ```python # GroupNormFwdOp's spec: weight and bias are optional, and arrive as None when absent @@ -191,8 +139,8 @@ def build_group_norm(x, weight, bias, *, num_groups, eps): def register_detector(target: str, detect: Callable[[torch.device], bool]) -> None: ... ``` -Called by the backend, once per target, at module import time. Registers that target's -device detection. +Called by the backend, once per target, when the backend module is imported. It registers +that target's device detection function. ### 4. `register_kernel_builder` @@ -200,7 +148,7 @@ device detection. def register_kernel_builder(op: str, target: str, build_kernel: BuildKernel) -> None: ... ``` -Called by the backend, once per op it takes over. Registers the kernel builder for +Called by the backend, once per op it takes over. It registers the kernel builder for `(op, target)`; registering the same pair twice is an error. ### `TensorSpec` @@ -212,8 +160,8 @@ class TensorSpec(NamedTuple): shape: tuple[int, ...] ``` -Defined by the protocol, built by the op layer and passed into `build_kernel`. What a -tensor is, without the tensor. +Defined by the protocol, built by the op layer and passed into `build_kernel`. It describes +a tensor's properties and does not contain the tensor. ```python # what a build_kernel argument looks like @@ -227,41 +175,42 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): ... ``` -The return value has one structural requirement: **it must be callable**, invocable -as `(*tensors)`, returning a tensor, a tuple of tensors, or `None` for a pure -in-place write. What the op layer checks is `callable()`. +The return value has one structural requirement: **it must be callable**. It is called +as `(*tensors)` and returns a tensor, a tuple of tensors, or `None` for a pure in-place +write. The op layer checks it with `callable()`. -**The protocol passes descriptions, not tensors.** That removes the need for a rule -the op layer could not enforce — "a builder must not read tensor contents or keep a -reference to a tensor". Two things are what such a rule would guard against: +**The protocol passes only descriptions of tensors.** The protocol therefore needs no +separate rule that "a builder must not read tensor contents or keep a reference to a +tensor", a rule the op layer could not enforce. Such a rule would guard against two +things: - **Reading data** would make the built kernel depend on data, while the memo table keys only on device and shape. -- **Keeping a reference** would have a tensor live as long as the cached kernel does. +- **Keeping a reference** would keep a tensor alive as long as the cached kernel. -A `TensorSpec` carries neither data nor tensor, so neither is expressible. +A `TensorSpec` carries neither data nor a tensor, so neither can happen. -When each of the four gets called during a real call is the next section. +The next section shows when each of the four functions is called during a real call. ## How one call reaches `build_kernel` {#from-op-layer} -One call, from the user's line to a backend's `build_kernel`: +The steps of one call, from user code to a backend's `build_kernel`: ```python # ── the caller ─────────────────────────────────────────────────────── op = GemmFwdOp() # no target= in the constructor, so the inputs' device decides # target="acme" skips detection and uses it directly; - # target=BUILTIN forces the kernels TileOPs ships + # target=BUILTIN forces the in-tree kernels a = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") b = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") d = op(a, b) # every input on one device: a.device == b.device # ── op layer: settle the target ────────────────────────────────────── # Every installed backend put a detect in the registry when it was imported, and the op -# layer hands a.device to each of them in turn — "is this device yours?": +# layer passes a.device to each of them in turn and asks "is this device yours?": # acme's detect(device) → True every other backend's → False # exactly one True → target = "acme", and this instance keeps it from here on -# none True → the kernels TileOPs ships run +# none True → the in-tree kernels run # two or more True → AmbiguousTargetError, asking for an explicit target= # ── op layer: run the checks generated from the manifest signature, then hand the whole op to the target ── @@ -284,7 +233,7 @@ d = op(a, b) # every input on one device: a.device == b.devi # kernel(a, b) # d = a @ b.T, computed by acme's kernel ``` -A backend writes one step of that — `build_gemm` — and registers it: +A backend writes one of these steps, `build_gemm`, and registers it: ```python def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): @@ -298,39 +247,38 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) ``` The op layer calls `build_gemm`; the backend never calls it itself. Importing the backend -module only records it in the registry, and the call comes when this target serves an op call -and it misses the external memo table — once per device and input -signature. Whatever it returns, the op layer stores and launches. +module only records it in the registry. It is called when this target serves an op call +and the external memo table misses, once per device and input signature. The op layer +stores the callable it returns in the memo table and launches it. -Four things follow from that: +This call path implies four things: - **`kernel_for` and the implementations' `entry_for` serve the in-tree path only.** They - decide which in-tree kernel is fetched, what it is looked up on and how it is built. - Once a target serves the op, it serves the whole op, and none of them runs. + decide which in-tree kernel is fetched, what it is looked up by and how it is built. + Once the op settles on a target, that target serves the whole op, and none of them runs. - **Tensors arrive positionally, params by name.** `build_kernel(*inputs, **params)`: the positional arguments are `TensorSpec`s (`None` for an optional input the call omitted), - the keywords the manifest's `params` names with the values this call settled on. -- **One builder per `(op, target)`.** Which of its kernels the in-tree path would run — - GEMM declares three in `kernel_types` — is not passed in; `build_kernel` decides from the + and the keyword arguments are the manifest's `params` names with the values settled for + this call. +- **One builder per `(op, target)`.** The backend is not told which kernel the in-tree + path would run (GEMM declares three in `kernel_types`); `build_kernel` decides from the `TensorSpec`s which kernel to return. -- **No memoisation of its own is needed.** For the same device and input signature the op - layer does not call again; for a finer split, or fewer rebuilds, add a cache inside - `build_kernel`. An op written to depend on a backend declares neither `kernel_types` nor - `interfaces`, and a call on it with no target claiming the device raises - `OpNotAvailableError`. +- **The backend needs no memoisation of its own.** For the same device and input signature + the op layer does not call `build_kernel` again. For a finer split, or fewer rebuilds, + the backend adds a cache inside `build_kernel`. An op written only for external backends + declares neither `kernel_types` nor `interfaces`; a call on it with no target claiming + the device raises `OpNotAvailableError`. ## Writing a backend that runs {#runnable} -With the four functions and one call's path in hand, the quickest start is to copy a -backend that already works. -[`tileops-backend-example`](https://github.com/lcy-seso/tileops-backend-example) is one, -written to those four steps. It implements its kernels in pure -PyTorch and claims CPU, so it installs, runs and tests anywhere; apart from the -kernels touching no dedicated hardware, every other part — entry point, -registration, the `build_kernel` signature, the memoisation rule, the error -messages — is what a backend for dedicated hardware writes. +[`tileops-backend-example`](https://github.com/lcy-seso/tileops-backend-example) is a +complete backend written to the four steps, and it can be copied as a starting point. It +implements its kernels in pure PyTorch and claims CPU, so it installs, runs and tests on +any machine. Apart from kernels that use no dedicated hardware, every part of it is what a +backend for dedicated hardware writes: the entry point, registration, the `build_kernel` +signature, the memoisation rule and the error messages. -What installing it changes: +The difference before and after installing it: ```console $ python -c "import torch; from tileops.norm import RMSNormFwdOp; \ @@ -344,7 +292,7 @@ $ python -c "...the same code..." # returns normally, bit-identical to torch.nn.functional.rms_norm ``` -Here is what it writes, step by step. +The following describes its contents, one step at a time. **Step 1, three lines of `pyproject.toml`.** The entry-point group is always `tileops.backends`, and the value is the backend's module: @@ -354,13 +302,13 @@ Here is what it writes, step by step. torch_cpu = "tileops_cpu" ``` -After `pip install` nothing initialises anything: TileOPs enumerates this group while -constructing its first op, imports the module named there, and the registration calls at -module top level fill the registry. There is no base class to inherit and no interface to -implement. +After `pip install`, no initialisation is needed. TileOPs enumerates this group while +constructing its first op and imports the module named there; the registration calls at +module top level fill the registry. The backend inherits no base class and implements no +interface. -**Step 2, a target name and a `detect`.** Both live in `target.py`, and `detect` is a -single line — it claims every CPU device: +**Step 2, a target name and a `detect`.** Both are in `target.py`. `detect` is a single +line and claims every CPU device: ```python TARGET = "torch_cpu" @@ -375,9 +323,9 @@ the backend author's to choose, while `device.type == "cpu"` is the device type defined by torch. **Step 3, a `build_kernel` written to the manifest signature.** `RMSNormFwdOp`'s spec -declares two inputs, `x` and an optional `weight`, and two params, `normalized_shape` and `eps`; the -function's parameters follow that declaration. It sits in `ops/rms_norm.py` with the -kernel class `CpuRMSNorm`: +declares two inputs, `x` and an optional `weight`, and two params, `normalized_shape` and +`eps`; the function's parameters follow that declaration. It is in `ops/rms_norm.py`, +together with the kernel class `CpuRMSNorm`: ```python def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): @@ -408,7 +356,7 @@ for _op, _build_kernel in BUILDERS.items(): ### Repository layout -Each file in the example covers one part of the work: +Each file in the example covers one part of the backend author's work: | File | Contents | | --- | --- | @@ -422,19 +370,19 @@ Each file in the example covers one part of the work: | `tests/test_errors.py` | the three error paths: an unregistered op raises rather than falling back, an unknown target raises, and a failed call binds the op to no target | | `tests/test_memoization.py` | when `build_kernel` is called again | -`CpuRMSNorm` does not receive the row count when it is constructed, which is -"compile-time parameters only" in practice. +`CpuRMSNorm` does not receive the row count when it is constructed. This is the rule +"the constructor takes compile-time parameters only" in practice. ### Running the tests -The tests need an environment with `tileops` installed: +The example's tests need an environment with `tileops` installed: ```bash pip install -e . # add --no-deps when tileops is already installed python -m pytest -q # 24 passed, both with two H200s visible and with CUDA_VISIBLE_DEVICES="" ``` -The same holds inside the TileOPs dev image, again without modifying TileOPs: +The tests also run inside the TileOPs dev image, again without modifying TileOPs: ```bash docker run --rm --gpus all -v "$PWD/..":/work -w /work \ @@ -444,11 +392,11 @@ docker run --rm --gpus all -v "$PWD/..":/work -w /work \ cd /work/tileops-backend-example && python -m pytest -q' ``` -`tileops` is deliberately absent from the example's dependencies. The package -extends an installation that already exists, and a version floor here would resolve -a release predating `tileops.backend`; the resulting `ImportError` is collected into -`load_failures()` and presents as "this backend is unusable" when the real cause is -that TileOPs is too old. +`tileops` is deliberately absent from the example's dependencies. The package extends an +installation that already exists, and a version floor here would resolve to a release +that predates `tileops.backend`. The resulting `ImportError` is collected into +`load_failures()` and appears as "this backend is unusable", while the real cause is that +TileOPs is too old. ### Turning it into a backend for real hardware @@ -464,8 +412,9 @@ that TileOPs is too old. ### The signature comes from the manifest **Writing a kernel needs the manifest, not the TileOPs source.** A builder's signature -is the op's manifest signature — `RMSNormFwdOp` in -[`src/tileops/manifest/spec/norm.yaml`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/manifest/spec/norm.yaml): +is the op's manifest signature. For example, `RMSNormFwdOp` in +[`src/tileops/manifest/spec/norm.yaml`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/manifest/spec/norm.yaml) +declares: ```yaml signature: @@ -484,36 +433,38 @@ The corresponding builder signature: def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): ``` -Two things to note about it. +The signature follows two conventions: - **Parameters arrive as the op instance holds them.** An `eps` not given at construction - arrives as the manifest default, `None`, meaning what it means in the reference API, - and the builder handles it that way; an omitted `weight` arrives as `None`. -- **The return value follows `signature.outputs`** — a tensor for a single output, - a tuple in declaration order for several, `None` for a pure in-place write. + arrives as the manifest default, `None`, with the meaning it has in the reference API, + and the builder handles it with that meaning. An omitted optional `weight` arrives as + `None`. +- **The return value follows `signature.outputs`:** a tensor for a single output, a tuple + in declaration order for several outputs, and `None` for a pure in-place write. ### The constructor takes compile-time parameters only -Values compiled into generated code — tile sizes, dimensions treated as constants, -dtypes — go in the constructor; the rest belongs to `__call__`. +Values compiled into generated code (tile sizes, dimensions treated as constants, dtypes) +go in the constructor; all other values belong to `__call__`. -Decode makes this a hard requirement: `seq_len` grows step by step and batch changes with -the running set, so putting them in the constructor means recompiling every step. +On the decode path this rule is a hard requirement: `seq_len` grows step by step and +batch changes with the running set, so putting them in the constructor recompiles the +kernel at every step. ### Shapes are the manifest's -The op layer changes no shapes: a kernel receives what the manifest declares, and -arranges whatever layout it needs inside its own call wrapper. +The op layer changes no shapes. A kernel receives the shapes the manifest declares, and +arranges any layout it needs inside its own call wrapper. -Where the code and the manifest disagree, the manifest governs: output dtype, shape -rules and parameter types are its, and a kernel does not rewrite them. +Where the code and the manifest disagree, the manifest governs. The manifest defines the +output dtype, the shape rules and the parameter types, and a kernel does not rewrite them. ### What a kernel's error has to say -A kernel that cannot serve a call raises rather than degrading, and its error says two +A kernel that cannot serve a call raises instead of degrading. Its error states two things: -- **Which item is unmet** — dtype, shape, arch, no implementation available, +- **Which item is unmet:** dtype, shape, arch, no implementation available, or compilation failed. - **The value it actually received.** @@ -521,11 +472,11 @@ things: ## What each phase may do {#phase-limits} -The decode path is captured by a CUDA graph, so each phase is bounded separately: +The decode path is captured by a CUDA graph, so each phase has its own limits: | Phase | May | May not | | --- | --- | --- | -| Memo lookup (its key and rebuild rules are [below](#memo)) | one dict lookup | anything else | +| Memo lookup (its key and rebuild rules are in [When a kernel is rebuilt](#memo)) | one dict lookup | anything else | | `detect` | one predicate | any import, any lock | | Building a kernel | select an implementation, compile, allocate, re-import, build handles | tuning that depends on real tensors | | Calling a kernel | launch a compiled kernel, allocate outputs through the torch allocator | compile, lazy init, build handles, host-side synchronisation | @@ -533,60 +484,57 @@ The decode path is captured by a CUDA graph, so each phase is bounded separately **A module-level import must not trigger compilation.** TileOPs imports the backend module while constructing the first op; compilation belongs in `build_kernel`. -A kernel call has two further stream rules: +A kernel call also follows two stream rules: -- **Launch on the current stream**, under CUDA `torch.cuda.current_stream(device)`; - never fall through to the default stream. Backends with their own launcher break - this most easily. -- **Internal allocations must outlive asynchronous execution.** Where only a raw - pointer is passed to a launch, the object has to stay alive until that stream has - finished. The protocol provides no workspace; this safety is the backend's. +- **It launches on the current stream.** Under CUDA, the current stream is + `torch.cuda.current_stream(device)`; a kernel never falls through to the default stream. + Backends with their own launcher break this rule most easily. +- **Internal allocations must outlive asynchronous execution.** When only a raw pointer is + passed to a launch, the object must stay alive until that stream has finished. The + protocol provides no workspace; the backend is responsible for this. -The caller warms up before capture — at least one non-captured call at the same -shape — because building a kernel may compile. During capture only one path is -allowed: memo hit, then call. +Because building a kernel may compile, the caller warms up before capture, with at least +one non-captured call at the same shape. During capture only one path is allowed: a memo +hit, then the call. ## After install: three states {#three-states} -Once `detect` claims a class of devices, **every** op on those devices is served by -that target, a missing one is an error, and there is no fall back to the -implementation TileOPs ships. The one exception is a composite op, which builds no -kernel of its own. +Once `detect` claims a class of devices, **every** op on those devices is served by that +target. A missing op is an error, with no fall back to the in-tree implementation. The +one exception is a composite op, which builds no kernel of its own. -The reason for not falling back: selecting a target means this device belongs to other -hardware, where the shipped kernels cannot launch at all. Falling back would trade a -clear "this target does not implement this op" for an incomprehensible launch -failure. +The op layer does not fall back because selecting a target means the device belongs to +other hardware, where the in-tree kernels cannot launch at all. Falling back would replace +a clear "this target does not implement this op" error with an obscure launch failure. -So after install, each op is in one of three states: +After install, each op is therefore in one of three states: | State | Result | | --- | --- | | The target registered a `build_kernel` for the op | it runs, the whole op on the target | -| It did not, and the op builds kernels of its own | an error naming the target and the op, with no fall back to the shipped implementation | +| It did not, and the op builds kernels of its own | an error naming the target and the op, with no fall back to the in-tree implementation | | It did not, and the op is a composite | the op runs its composition, and each sub-op settles on a target itself | -Covering every op the target model uses is therefore work on the backend's side. The op -side is settled by design: the op layer keys the external path on the call's own inputs -(see [how one call reaches `build_kernel`](#from-op-layer)). +Covering every op the target model uses is therefore the backend's work. The op side is +settled by design: the op layer computes the memo key of the external path from the call's +own inputs (see [How one call reaches `build_kernel`](#from-op-layer)). ### No hardware queried before the target is settled -Until a target is settled, the op layer queries nothing bound to specific hardware — a -CUDA SM version, say. Querying it would mean that on a machine without that driver the -call fails before it reaches `build_kernel`, for a reason that has nothing to do with the -backend. +Until a target is settled, the op layer queries nothing bound to specific hardware, such +as a CUDA SM version. With such a query, on a machine without that driver the call would +fail before it reaches `build_kernel`, for a reason unrelated to the backend. -If such a failure does show up on your hardware, the traceback stops inside TileOPs rather -than in the backend's `build_kernel`. That is a regression on the TileOPs side: file an -issue with the traceback. +If such a failure occurs on your hardware, the traceback stops inside TileOPs, not in the +backend's `build_kernel`. That is a regression on the TileOPs side: file an issue with the +traceback. -Every test in the example also passes where no GPU is visible, and none is skipped; that -run is itself the check of this premise. +Every test in the example passes where no GPU is visible, and none is skipped; that run +checks this premise. ## Error messages and what to do -All three are measured output, each with one cause and one way to handle it. +The three messages below are measured output. Each has one cause and one fix. **No builder registered for the op:** @@ -607,20 +555,19 @@ UnknownTargetError: no backend registered target 'nope'; known targets: ['torch_ The package did not install, or the target name is misspelled. Use `tileops.backend.registered_targets()` to see what actually registered. -**`target=BUILTIN` forces the implementation TileOPs ships:** +**`target=BUILTIN` forces the in-tree implementation:** ``` OpNotAvailableError: RMSNormFwdOp's in-tree kernels do not run on cpu; known targets for this op: ['torch_cpu'] ``` -`BUILTIN` bypasses backends explicitly. The shipped implementation cannot run on -CPU tensors, which is precisely the outcome the no-fall-back rule avoids. +`BUILTIN` bypasses all backends explicitly. The in-tree implementation cannot run on CPU +tensors; this error shows the outcome that the no-fall-back rule avoids. -**When a backend package fails to import**, TileOPs skips it, warns, and collects -the reason into `load_failures()`. One broken plugin does not make TileOPs -unimportable. If registration raises part way through, everything that backend -registered in that pass is **rolled back** — no half-implemented target is left in -the registry. +**When a backend package fails to import**, TileOPs skips it, issues a warning, and +collects the reason into `load_failures()`. One broken plugin does not make TileOPs fail +to import. If registration raises part way through, everything that backend registered in +that pass is **rolled back**, and no partly registered target is left in the registry. ```python from tileops.backend import load_failures @@ -635,31 +582,32 @@ print(load_failures()) | **`detect`** | A function a backend writes, one per target | How the first layer picks: it receives a `torch.device` and answers whether such a device is what its kernels are for; `False` if not | | **`build_kernel`** | A function a backend writes per op, one per `(op, target)` | The second layer: it receives a description of this call — each input's device, dtype and shape, plus the op's parameters — and picks, builds and returns a kernel from its own set | -**Selection has two layers: TileOPs picks the target, the target picks the -kernel.** The second happens inside `build_kernel`, with no protocol involvement: on this -path there is no kernel-level concept, no capability negotiation and no candidate -filtering. The candidate filtering TileOPs does run — availability, applicability, -precedence — belongs to the in-tree path and to the two smaller ways in -([three ways in](#three-ways)), which a target bypasses. - -`detect` answers only which devices belong to the backend, and that is as fine as -it gets. **Whether this call is supported — dtype, shape, parameter combination — -is answered by `build_kernel`**, the only place that sees the full input -description and the parameters; it raises there when it cannot serve the call. -Leaving those judgements to `detect` is not possible: all it receives is a -`torch.device`. - -TileOPs does not parse `torch.device` — it passes it through to `detect`. Device -types and targets are not in one-to-one correspondence: one device type can carry -several sets of kernels from different vendors; some hardware arrives through -`privateuseone`, whose string carries no vendor information at all; and some -backends have to read an environment variable or call a vendor runtime to decide. +**Selection has two layers: TileOPs picks the target, and the target picks the kernel +from its own set.** The second layer happens inside `build_kernel`, without the protocol: +this path has no kernel-level concept, no capability negotiation and no candidate +filtering. The candidate filtering TileOPs does run (availability, applicability, +precedence) belongs to the in-tree path and to the two smaller mechanisms (see +[Three ways to extend dispatch](#three-ways)); a target bypasses it. + +`detect` answers only which devices belong to the backend, and nothing finer. +**Whether this call is supported (dtype, shape, parameter combination) is answered by +`build_kernel`**, the only place that sees the full input description and the +parameters; it raises there when it cannot serve the call. `detect` cannot make these +judgements, because it receives only a `torch.device`. + +TileOPs does not parse `torch.device`; it passes it unchanged to `detect`. Device types +and targets are not in one-to-one correspondence: + +- one device type can carry several sets of kernels from different vendors; +- some hardware arrives through `privateuseone`, whose string carries no vendor + information; +- some backends must read an environment variable or call a vendor runtime to decide. ## The op layer's contract -The seven below are the op layer's contract to every target: it implements them and a -backend reuses them rather than writing its own. They are listed in the order a backend -author meets them: +The seven items below are the op layer's contract to every target. The op layer +implements them, and a backend reuses them instead of writing its own. They are listed in +the order a backend author meets them: | # | The op layer supplies | What it means for a backend | | --- | --- | --- | @@ -671,47 +619,46 @@ author meets them: | 6 | The `torch.compile` and CUDA-graph boundary | The op layer wraps a call as an opaque operator and registers a fake alongside, so the compiler can infer the output's shape and dtype without executing. **A backend's kernels do nothing for compilation**; see [Bringing an op into torch.compile](torch-compile.md) | | 7 | Roofline, profiling and numerical tests | The op layer's existing tests run once with the backend's kernel and compare against the manifest's `ref_api`; performance reports are produced as usual | -None of the seven depends on hardware and every target gets them identically; a -third-party backend neither bypasses one nor substitutes its own. +None of the seven depends on hardware, and every target gets them identically. A +third-party backend neither bypasses any of them nor substitutes its own. -The kernels TileOPs ships ([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)) are the **default -implementation**: they have no target name and are not in the registry. +The in-tree kernels ([`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)) +are the **default implementation**: they have no target name and are not in the registry. -**The in-tree implementations run by default.** With no backend installed, no `target=` named -and no process default set, calls run the shipped implementation. Nothing is preconfigured -to substitute: a backend serves an op only once it is installed and either claims the -device through its `detect`, or is named by `target=` or `set_default_target`. +**The in-tree implementations run by default.** With no backend installed, no `target=` +named and no process default set, calls run the in-tree implementation. An installed +backend serves an op only when its `detect` claims the device, or when it is named by +`target=` or `set_default_target`. ## When a kernel is rebuilt {#memo} -TileOPs remembers a builder's return value by **device plus input signature**: +TileOPs memoises a builder's return value by **device plus input signature**: > the device this call's tensors are on, plus `(dtype, shape)` taken per input in > `signature.inputs` order; an optional input not passed on this call is recorded > as `None`. -That is: **two calls agreeing on device and input signature get the same kernel -back**, with no further call to `build_kernel`; a second card under the same target -builds again, because an artefact compiled for one device need not launch on -another. Params are not part of the key — they are fixed for an op instance. +**Two calls with the same device and input signature get the same kernel**, with no +further call to `build_kernel`. A second card under the same target builds again, because +an artefact compiled for one device need not launch on another. Params are not part of +the key, because they are fixed for an op instance. -How the op layer looks that table up, and what it does on a miss, is in [how one call -reaches `build_kernel`](#from-op-layer). +How the op layer looks up that table, and what it does on a miss, is in +[How one call reaches `build_kernel`](#from-op-layer). Two consequences: - **An entry may not last.** When a call fails, the op revokes its target decision and - drops its memo table. A backend must not assume the callable it returned stays - alive; whatever resources it - depends on, it holds references to itself. -- **A finer or a coarser grain is resolved on the backend side.** Finer - distinctions happen inside the backend; to rebuild less often, add a cache inside + drops its memo table. A backend must not assume the callable it returned stays alive; + the callable holds its own references to the resources it depends on. +- **A finer or a coarser grain is handled on the backend side.** Finer distinctions are + made inside the backend; to rebuild less often, the backend adds a cache inside `build_kernel`. ## What a caller can reach for -These are for callers. A backend author does not need them, but they help while -debugging: +These interfaces are for callers. A backend author does not need to call them, but they +help while debugging: ```python from tileops.backend import ( @@ -724,14 +671,18 @@ set_default_target("torch_cpu") # process default, ahead of device detectio set_default_target(BUILTIN) # turn substitution off globally ``` -Target selection order: the `target=` constructor argument, then the process -default, then device detection. `BUILTIN` forces the implementation TileOPs ships. -A named target that is not registered, or does not implement the op, is an error; -another target is not used instead. +The target is selected in this order: + +1. the `target=` constructor argument; +1. the process default; +1. device detection. + +`BUILTIN` forces the in-tree implementation. A named target that is not registered, or +does not implement the op, is an error; another target is not used instead. ## What the protocol does not support -These are outside the protocol, each for a reason: +The following cases are outside the protocol, each for the reason given: | Not supported | Reason | | --- | --- | diff --git a/docs/backends.zh.md b/docs/backends.zh.md index 127bb29e..4e15415d 100644 --- a/docs/backends.zh.md +++ b/docs/backends.zh.md @@ -1,102 +1,73 @@ -# 接入新硬件后端 +# 接入新的硬件 backend -TileLang 是多后端 DSL，每种硬件各有一套独立的 kernel，由各自的 Python 包发行。TileOPs 因此定了一套协议：仓外的 Python 包可以接管某个算子的 kernel，取代自带的实现，且不必改 TileOPs 的任何代码。 +TileLang 是支持多种 backend 的 DSL，每种硬件各有一套独立的 kernel，由各自的 Python 包发行。TileOPs 因此定义了一套协议：仓库之外的 Python 包可以接管某个 op 的 kernel，取代 in-tree 实现，且不必修改 TileOPs 的任何代码。 -本页讲怎么把一类新硬件接进来，让这类设备上的算子由自己的 kernel 执行。 +本页说明如何接入一类新硬件，使这类设备上的 op 由该硬件自己的 kernel 执行。 -**后端只提供一件事：一个能算这次调用的可调用对象。** 其余都由算子层负责。 +**backend 只提供一件事：一个能执行这次调用的可调用对象。** 其余工作都由 op 层负责。 -本页讲的是 target，接管整个算子的那一种接入方式。前半部分按实现顺序列出要做的事：要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么，以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。 +本页讲的是 target，即接管整个 op 的那种接入方式。前半部分按实现顺序说明 backend 作者要做的事： -后半说明协议何以如此设计：两层选择、算子层的契约、kernel 的重建条件、调用方可用的接口，以及刻意不支持的情形。 +1. 要写的四样东西； +2. 协议中的四个函数； +3. 一次调用如何走到这些函数； +4. 一个可直接安装运行的 backend； +5. 如何把模板改造为面向真实硬件的 backend； +6. 编写 kernel 的四条规则； +7. 各阶段允许做什么； +8. 安装之后每个 op 所处的状态，以及各条错误信息对应的原因。 -## 三种接入方式 {#three-ways} - -一个仓库之外的包按自己的接管范围，在三种方式中选一种。范围较小的两种，写出的 kernel 类遵守 [kernel 接口](new-op.md#kernel-selection)，与自带实现遵守同一份契约；target 遵守的则是算子在 manifest 中的签名，写一个 `build_kernel`。 - -| # | | `kernel_map=` | `register_implementation` | target | -| --- | --- | --- | --- | --- | -| 1 | 改变什么 | 一个 key 登记的类；这个 key 服务哪些调用不变 | 新增一个 key，带有它自己的适用范围与优先关系 | 算子的全部调用 | -| 2 | 作用于 | 调用方构造的那一个算子实例 | 注册之后构造的所有该算子实例 | 选中该 target 的算子实例 | -| 3 | 依据的契约 | kernel 接口 | kernel 接口 | 算子在 manifest 中的签名 | -| 4 | 新类不服务的调用 | 该 key 被选中时报错 | 仍由自带实现服务 | 不存在，target 服务全部调用 | - -**`kernel_map=`** 是每个算子构造函数的参数，值是从 key 到类的映射。它只把这个实例中该 key 登记的类换掉：这个 key 仍按原来登记的实现声明的 `applies`、`general` 与 `preferred_over` 参与选择，替换者与其他实现一样继承这个 key 所属的接口，由它自己的 `entry_for` 构造。该 key 被选中而替换者不服务这次调用时报错，不退回被替换的实现。下例取自 [`tests/test_kernel_dispatch.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_kernel_dispatch.py)： - -```python -class _TorchLayerNorm(Kernel, LayerNormFwdInterface): - """A replacement written against ``LayerNormFwdInterface`` alone.""" +后半部分说明协议为何如此设计，包括： - devices = frozenset({torch.device(run_device()).type}) - - def __init__(self, n: int, eps: float) -> None: - super().__init__() - self.n, self.eps = n, eps - - @classmethod - def entry_for(cls, call: LayerNormCall): - return (call.n, call.eps), lambda: cls(call.n, call.eps) - - def forward(self, x, weight, bias): - return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( - x.dtype - ) - - -op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) -``` - -`kernel_map=` 中本算子没有、而其他算子有的 key 被忽略，复合算子由此把一份映射传给各个子算子；任何算子都没有的 key 在构造时报 `was given kernel_map keys no op has`。 - -**`register_implementation(op, key, implementation)`** 为一个接口新增一个实现。`op` 是算子在 manifest 中的 key，`key` 是新实现的名字，它加入哪个接口由这个类继承了哪个接口决定。新实现自己声明适用范围，与某个不是 `general` 的自带实现重叠时要声明 `preferred_over`： - -```python -class _NarrowTorchLayerNorm(_TorchLayerNorm): - """An added implementation for short rows, which wins over the in-tree one there.""" - - preferred_over = frozenset({"layer_norm"}) - - @classmethod - def applies(cls, call: LayerNormCall) -> bool: - return call.n <= 64 +- 两层选择； +- op 层的契约； +- kernel 的重建条件； +- 调用方可用的接口； +- 刻意不支持的情形。 +## 三种接入方式 {#three-ways} -register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) -``` +一个仓库之外的包按自己接管的范围，从三种方式中选一种： -这样 `n <= 64` 的调用由 `_NarrowTorchLayerNorm` 服务，`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增的实现只对注册之后构造的算子实例生效。同一个算子下重复注册同一个 key 报 `BackendError`；key 与自带实现的 key 相同时，构造实例时报 `reuse keys it has`。 +1. `kernel_map=`：调用方构造 op 时，替换某个 key 背后运行的类； +2. `register_implementation`：为一个 kernel 接口新增一个实现，与 in-tree 实现一起参与选择； +3. target：接管 op 的全部调用。 -`register_implementation` 在后端模块被导入时执行，通过 target 所用的同一个 entry point 触发；`kernel_map=` 则由调用方在构造算子时传入，不经过注册。以下各节讲 target。 +前两种方式写出的 kernel 类遵守 kernel 接口，与 in-tree 实现遵守同一份契约，写法见 [backend 如何接入](user-guide/dispatch/backends.md)。target 遵守的是 op 在 manifest 中的签名，backend 为它写一个 `build_kernel`。本页以下各节只讲 target。 -## 写一个后端要做的四件事 +## 写一个 backend 要做的四件事 | # | 做什么 | | --- | --- | -| 1 | 在 `pyproject.toml` 里声明一条 entry point，指向后端模块 | -| 2 | 起一个 target 名，写 `detect`，声明这套 kernel 面向哪一类设备 | -| 3 | 挑第一个要接管的算子，照它的 manifest 签名写 `build_kernel` | +| 1 | 在 `pyproject.toml` 中声明一条 entry point，指向 backend 模块 | +| 2 | 起一个 target 名，编写 `detect`，声明这套 kernel 面向哪一类设备 | +| 3 | 选定第一个要接管的 op，按它的 manifest 签名编写 `build_kernel` | | 4 | 在模块顶层调用 `register_detector` 与 `register_kernel_builder` | -四件事写完，`pip install` 即生效。下面几节依次是：这四个函数的签名、一次调用怎么走到它们，以及一个照这四步写成、可直接安装运行的[完整后端](#runnable)。 +四件事完成后，`pip install` 即生效。下面几节依次给出： + +1. 这四个函数的签名； +2. 一次调用如何走到它们； +3. 一个按这四步写成、可直接安装运行的[完整 backend](#runnable)。 -之后逐个算子增加 `build_kernel`。**目标模型用到的、自己构造 kernel 的算子必须全部覆盖** —— 缺一个就报错，不会改用自带实现，因为那些 kernel 在该 target 的设备上启动不了。只调用子算子的复合算子不需要 builder。 +之后逐个 op 增加 `build_kernel`。**目标模型用到的、自己构造 kernel 的 op 必须全部覆盖**：缺少任何一个都会报错，不会改用 in-tree 实现，因为 in-tree kernel 在该 target 的设备上无法启动。只调用子 op 的复合 op 不需要 builder。 ## 协议中的四个函数 -`tileops.backend` 只定义对外接口，用 Python 的结构化类型（`typing.Protocol`）表达：后端不继承基类，不实现抽象方法，写出签名相符的普通函数再注册进来即可。算子层检查返回值时同样只看结构，即 `callable()`。 +`tileops.backend` 只定义对外接口，用 Python 的结构化类型（`typing.Protocol`）表达：backend 不继承基类，也不实现抽象方法，只需写出签名相符的普通函数并注册。op 层检查返回值时同样只看结构，即 `callable()`。 -后端实现 `detect` 与 `build_kernel`，再调 `register_detector` 与 `register_kernel_builder` 把它们登记进来；另有协议定义的 `TensorSpec`，以及 `pyproject.toml` 里的一条 entry point： +backend 实现 `detect` 与 `build_kernel`，再调用 `register_detector` 与 `register_kernel_builder` 登记它们。此外还有协议定义的 `TensorSpec`，以及 `pyproject.toml` 中的一条 entry point： | # | 名字 | 谁写 | 谁调用，何时调用 | | --- | --- | --- | --- | -| 1 | `detect` | 后端实现 | 算子层为一次调用选定 target 时，逐个 target 调用 | -| 2 | `build_kernel` | 后端实现 | 算子层在记忆表未命中时调用 | -| 3 | `register_detector` | 后端调用 | 后端模块被 import 时执行一次 | -| 4 | `register_kernel_builder` | 后端调用 | 同上，每个要接管的算子调用一次 | -| —— | `TensorSpec` | 协议定义 | 算子层构造，作为 `build_kernel` 的实参传入 | -| —— | entry point | 后端在 `pyproject.toml` 中声明 | TileOPs 在构造第一个算子时枚举 | +| 1 | `detect` | backend 实现 | op 层为一次调用选定 target 时，对每个 target 各调用一次 | +| 2 | `build_kernel` | backend 实现 | op 层在记忆表未命中时调用 | +| 3 | `register_detector` | backend 调用 | backend 模块被 import 时执行一次 | +| 4 | `register_kernel_builder` | backend 调用 | backend 模块被 import 时执行，每个要接管的 op 调用一次 | +| —— | `TensorSpec` | 协议定义 | op 层构造，作为 `build_kernel` 的实参传入 | +| —— | entry point | backend 在 `pyproject.toml` 中声明 | TileOPs 在构造第一个 op 时枚举 | -下面按这个顺序逐个给出签名与含义，最后是协议定义的 `TensorSpec`。 +下面按这个顺序逐个给出签名与含义，最后给出协议定义的 `TensorSpec`。 ### 1. `detect` @@ -104,7 +75,7 @@ register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerN def detect(device: torch.device) -> bool: ... ``` -后端实现。回答这类设备是否由自己这套 kernel 服务：只看设备，不看 dtype 与形状；不是自己的设备返回 `False`，不要抛异常。 +由 backend 实现。它回答这类设备是否由自己这套 kernel 服务，只看设备，不看 dtype 与形状。设备不属于自己时返回 `False`，不抛异常。 ```python # 认领一整类设备 @@ -124,7 +95,7 @@ def detect(device: torch.device) -> bool: def build_kernel(*inputs: "TensorSpec | None", **params) -> Callable[..., KernelResult]: ... ``` -后端实现，一组 `(算子, target)` 一个。签名即该算子的 manifest 签名：`inputs` 与 `signature.inputs` 的条目一一对应、按声明顺序，`params` 按 `signature.params` 命名。声明为 optional 的输入本次没有传入时，对应的实参是 `None`。 +由 backend 实现，每组 `(op, target)` 一个。它的签名就是该 op 的 manifest 签名：`inputs` 按声明顺序与 `signature.inputs` 的条目一一对应，`params` 按 `signature.params` 命名。声明为 optional 的输入在本次调用中没有传入时，对应的实参是 `None`。 ```python # GroupNormFwdOp 的 spec：weight、bias 是可选输入，没传时实参是 None @@ -140,7 +111,7 @@ def build_group_norm(x, weight, bias, *, num_groups, eps): def register_detector(target: str, detect: Callable[[torch.device], bool]) -> None: ... ``` -后端调用，每个 target 一次，在后端模块被 import 时执行。登记该 target 的设备识别函数。 +由 backend 调用，每个 target 一次，在 backend 模块被 import 时执行。它登记该 target 的设备识别函数。 ### 4. `register_kernel_builder` @@ -148,7 +119,7 @@ def register_detector(target: str, detect: Callable[[torch.device], bool]) -> No def register_kernel_builder(op: str, target: str, build_kernel: BuildKernel) -> None: ... ``` -后端调用，每个要接管的算子一次。登记 `(算子, target)` 的 kernel 构造函数；同一组重复登记会报错。 +由 backend 调用，每个要接管的 op 一次。它登记 `(op, target)` 的 kernel 构造函数；同一组重复登记会报错。 ### `TensorSpec` @@ -159,7 +130,7 @@ class TensorSpec(NamedTuple): shape: tuple[int, ...] ``` -协议定义的类型，由算子层构造后传入 `build_kernel`。描述一个张量是什么，不含张量本身。 +协议定义的类型，由 op 层构造后传入 `build_kernel`。它描述一个张量的属性，不包含张量本身。 ```python # build_kernel 收到的实参长这样 @@ -173,20 +144,20 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): ... ``` -返回值只需满足一条结构约定：**它必须可调用**，能以 `(*tensors)` 的形式调用，返回一个张量、一个张量元组，或纯原地写入时的 `None`。算子层对它的检查就是 `callable()`。 +返回值只需满足一条结构约定：**它必须可调用**，能以 `(*tensors)` 的形式调用，返回一个张量、一个张量元组，或在纯原地写入时返回 `None`。op 层对它的检查就是 `callable()`。 -**协议传的是描述，不是张量。** 这样就不必再写一条「构造时不得读张量内容、不得保存对张量的引用」的规则，那条规则算子层根本无法校验。它要防的是两件事： +**协议只传入张量的描述。** 因此协议不需要另写一条「构造时不得读取张量内容、不得保存对张量的引用」的规则，op 层也无法校验这样的规则。这条规则要防止两件事： -- **读取数据**会让构造结果依赖数据，而记忆表只按设备与形状记录。 -- **保存引用**会让张量随着被缓存的 kernel 一起存活整个进程。 +- **读取数据**会使构造结果依赖数据，而记忆表只按设备与形状记录。 +- **保存引用**会使张量与被缓存的 kernel 存活得一样久。 -`TensorSpec` 上既没有数据也没有张量，这两件事于是无从写起。 +`TensorSpec` 上既没有数据也没有张量，这两件事因此无从发生。 -这四个函数在一次真实调用里各自何时被调到，见下一节。 +下一节说明这四个函数在一次真实调用中分别在何时被调用。 -## 一次调用怎么走到 `build_kernel` {#from-op-layer} +## 一次调用如何走到 `build_kernel` {#from-op-layer} -一次调用从用户代码走到后端的 `build_kernel`，中间经过的每一步： +下面列出一次调用从用户代码到 backend 的 `build_kernel` 所经过的每一步： ```python # ── 调用方 ─────────────────────────────────────────────────────────── @@ -197,40 +168,40 @@ a = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") b = torch.randn(4096, 4096, dtype=torch.float16, device="acme:0") d = op(a, b) # 所有输入必须在同一设备上：a.device == b.device -# ── 算子层：定 target ──────────────────────────────────────────────── -# 每个装好的后端在 import 时都往注册表里放了一个 detect。算子层把 a.device +# ── op 层：定 target ──────────────────────────────────────────────── +# 每个装好的 backend在 import 时都往注册表里放了一个 detect。op 层把 a.device # 这一个对象原样交给每个 detect，问「这块设备是不是你这套 kernel 的」： -# acme 的 detect(device) → True 其他后端的 → False -# 恰好一个返回 True → target = "acme"，这个算子实例此后固定用它 +# acme 的 detect(device) → True 其他 backend的 → False +# 恰好一个返回 True → target = "acme"，这个op 实例此后固定用它 # 一个都没有返回 True → 用 TileOPs 自带的 kernel # 两个以上返回 True → 抛 AmbiguousTargetError，要求显式写 target= -# ── 算子层：先跑由 manifest 签名生成的检查，再把整个算子交给 target ── -# GemmFwdOp 自己的 _eager_forward 与 kernel_for 只服务自带实现，这次不走 +# ── op 层：先跑由 manifest 签名生成的检查，再把整个 op交给 target ── +# GemmFwdOp 自己的 _eager_forward 与 kernel_for 只服务in-tree 实现，这次不走 # 传给 target 的张量顺序照 signature.inputs，不写入的输入先转成连续 -# ── 算子层：按设备与输入签名查外部记忆表 ───────────────────────────── +# ── op 层：按设备与输入签名查外部记忆表 ───────────────────────────── # ("acme:0", (float16, (4096, 4096)), (float16, (4096, 4096))) # 第一项是设备，其余每项对应一个输入的 (dtype, shape) -# 这是这个算子实例的第一次调用，表还是空的 → 未命中，往下走构造 +# 这是这个op 实例的第一次调用，表还是空的 → 未命中，往下走构造 # 同样设备、同样 dtype 与形状的下一次调用就会命中，直接跳到最后一步 -# ── 后端：算子层调 build_gemm，张量已转成 TensorSpec ───────────────── +# ── backend：op 层调 build_gemm，张量已转成 TensorSpec ───────────────── # build_gemm(TensorSpec("acme:0", float16, (4096, 4096)), # TensorSpec("acme:0", float16, (4096, 4096)), # trans_a=False, trans_b=True) # params 按 manifest 的名字传 # → 返回一个可调用对象 -# ── 算子层：存进记忆表，然后 launch ───────────────────────────────── +# ── op 层：存进记忆表，然后 launch ───────────────────────────────── # kernel(a, b) # d = a @ b.T，由 acme 的 kernel 算出 ``` -后端要写的只有其中一步 —— 那个 `build_gemm`，以及把它注册进来： +backend 只需编写其中一步，即 `build_gemm`，并把它注册进来： ```python def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): m = a.shape[1] if trans_a else a.shape[0] - if m == 1: # 情形由后端从 TensorSpec 自行判断 + if m == 1: # 情形由 backend从 TensorSpec 自行判断 return AcmeGemv(a, b, trans_a, trans_b) return AcmeGemm(a, b, trans_a, trans_b) @@ -238,18 +209,18 @@ def build_gemm(a: TensorSpec, b: TensorSpec, *, trans_a, trans_b): register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm) ``` -`build_gemm` 由算子层调用，后端自己从不调它：import 后端模块时只是把它登记进注册表，真正被调是在一次调用由这个 target 服务、且外部记忆表未命中的时候，每个「设备 + 输入签名」一次。它返回的可调用对象随后由算子层 launch，也由算子层存进记忆表。 +`build_gemm` 由 op 层调用，backend 自己从不调用它。import backend 模块时只是把它登记进注册表；它真正被调用，是在一次调用由这个 target 服务、且外部记忆表未命中的时候，每个「设备 + 输入签名」调用一次。它返回的可调用对象随后由 op 层 launch，并由 op 层存进记忆表。 -四点对应关系值得记住： +这条调用路径对应以下四点： -- **`kernel_for` 与各实现的 `entry_for` 只服务自带实现。** 它们决定取哪个自带 kernel、按什么查表、又怎么构造。算子选中某个 target 之后，整个算子由这个 target 服务，这几处都不会执行。 -- **张量按位置传，参数按名字传。** `build_kernel(*inputs, **params)`：位置实参是 `TensorSpec`（没传的可选输入是 `None`），关键字实参是 manifest 里 `params` 的名字与本次调用的确定值。 -- **一个 `(算子, target)` 只注册一个 builder。** 自带实现内部分几种 kernel（GEMM 的 `kernel_types` 里有三个）不会传进来，`build_kernel` 从 `TensorSpec` 自行判断该返回哪个 kernel。 -- **不必自己缓存构建结果。** 设备与输入签名相同时，算子层不会再调用一次 `build_kernel`；需要更细的区分或更少的重建，在 `build_kernel` 内部另加一层缓存。专为外部后端而写的算子，`kernel_types` 与 `interfaces` 都不声明，没有 target 认领设备时调用直接抛 `OpNotAvailableError`。 +- **`kernel_for` 与各实现的 `entry_for` 只服务 in-tree 实现。** 它们决定取哪个 in-tree kernel、按什么查表以及如何构造。op 选中某个 target 之后，整个 op 由这个 target 服务，这几处都不会执行。 +- **张量按位置传入，参数按名字传入。** 在 `build_kernel(*inputs, **params)` 中，位置实参是 `TensorSpec`（没有传入的可选输入是 `None`），关键字实参是 manifest 中 `params` 的名字与本次调用的确定值。 +- **一组 `(op, target)` 只注册一个 builder。** in-tree 实现内部区分的几种 kernel（GEMM 的 `kernel_types` 中有三个）不会传给 backend，`build_kernel` 从 `TensorSpec` 自行判断应返回哪个 kernel。 +- **backend 不必自己缓存构建结果。** 设备与输入签名相同时，op 层不会再次调用 `build_kernel`；需要更细的区分或更少的重建时，backend 在 `build_kernel` 内部另加一层缓存。专为外部 backend 编写的 op 既不声明 `kernel_types` 也不声明 `interfaces`，没有 target 认领设备时，调用直接抛出 `OpNotAvailableError`。 -## 实现一个可运行的后端 {#runnable} +## 实现一个可运行的 backend {#runnable} -读到这里，四个函数与一次调用的路径都齐了，可以直接照抄一个能跑的后端。[`tileops-backend-example`](https://github.com/lcy-seso/tileops-backend-example) 就是按这四步写成的完整后端。它以纯 PyTorch 实现 kernel、认领 CPU，因此在任何机器上都能安装、运行和测试；除 kernel 本身不涉及专用硬件之外，其余各部分 —— entry point、注册方式、`build_kernel` 签名、记忆规则、错误信息 —— 与一个面向专用硬件的后端完全一致。 +[`tileops-backend-example`](https://github.com/lcy-seso/tileops-backend-example) 是按这四步写成的完整 backend，可以直接复制使用。它用纯 PyTorch 实现 kernel 并认领 CPU，因此在任何机器上都能安装、运行和测试。除了 kernel 本身不涉及专用硬件，它的其余部分与面向专用硬件的 backend 完全一致，包括 entry point、注册方式、`build_kernel` 签名、记忆规则与错误信息。 安装这个包前后的差别如下： @@ -265,18 +236,18 @@ $ python -c "...同一段代码..." # 正常返回，结果与 torch.nn.functional.rms_norm 逐位相同 ``` -下面按这四步逐段看它写了什么。 +下面按这四步逐段说明它的内容。 -**第一步，`pyproject.toml` 里的三行。** entry point 组的名字固定为 `tileops.backends`，值是后端模块名： +**第一步，`pyproject.toml` 中的三行。** entry point 组的名字固定为 `tileops.backends`，值是 backend 模块名： ```toml [project.entry-points."tileops.backends"] torch_cpu = "tileops_cpu" ``` -`pip install` 之后不需要任何初始化：TileOPs 在构造第一个算子时枚举这个组、import 其中声明的模块，模块顶层的注册调用就把注册表填好。既没有需要继承的基类，也没有需要实现的接口。 +`pip install` 之后不需要任何初始化。TileOPs 在构造第一个 op 时枚举这个组并 import 其中声明的模块，模块顶层的注册调用随之填好注册表。backend 既不需要继承基类，也不需要实现接口。 -**第二步，起 target 名，写 `detect`。** 两者都在 `target.py` 里，`detect` 只有一行 —— 认领所有 CPU 设备： +**第二步，起 target 名，编写 `detect`。** 两者都在 `target.py` 中。`detect` 只有一行，认领所有 CPU 设备： ```python TARGET = "torch_cpu" @@ -286,9 +257,9 @@ def detect(device: torch.device) -> bool: return device.type == "cpu" ``` -两个名字含义不同：`TARGET = "torch_cpu"` 是这一套 kernel 的名字，由后端作者决定；`device.type == "cpu"` 是它认领的设备类型，由 torch 定义。 +两个名字含义不同：`TARGET = "torch_cpu"` 是这套 kernel 的名字，由 backend 作者决定；`device.type == "cpu"` 是它认领的设备类型，由 torch 定义。 -**第三步，照 manifest 签名写 `build_kernel`。** `RMSNormFwdOp` 的 spec 声明了两个输入 `x`、`weight`（可选）与两个参数 `normalized_shape`、`eps`，函数的形参照抄这份声明。它和 kernel 类 `CpuRMSNorm` 一起放在 `ops/rms_norm.py` 里： +**第三步，按 manifest 签名编写 `build_kernel`。** `RMSNormFwdOp` 的 spec 声明了两个输入 `x` 与 `weight`（可选），以及两个参数 `normalized_shape` 与 `eps`，函数的形参照抄这份声明。它与 kernel 类 `CpuRMSNorm` 一起放在 `ops/rms_norm.py` 中： ```python def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): @@ -297,7 +268,7 @@ def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape return CpuRMSNorm(normalized_shape, eps, x.dtype) ``` -**第四步，在模块顶层注册。** `ops/__init__.py` 里的 `BUILDERS` 列出这个 target 接管的全部算子，每个键都按 manifest 中的拼写书写；包的 `__init__.py` 登记一个 detector，再给表里每个算子登记一次 builder： +**第四步，在模块顶层注册。** `ops/__init__.py` 中的 `BUILDERS` 列出这个 target 接管的全部 op，每个键都按 manifest 中的拼写书写。包的 `__init__.py` 登记一个 detector，再为表中每个 op 登记一次 builder： ```python BUILDERS = { @@ -321,17 +292,17 @@ for _op, _build_kernel in BUILDERS.items(): | 文件 | 内容 | | --- | --- | -| `pyproject.toml` | entry point 声明，也就是全部安装机制 | +| `pyproject.toml` | entry point 声明，这就是全部的安装机制 | | `src/tileops_cpu/__init__.py` | 全部注册代码 | | `src/tileops_cpu/target.py` | target 名与 `detect` | -| `src/tileops_cpu/ops/__init__.py` | `BUILDERS`，即这个 target 接管的全部算子。键按 manifest 中的拼写书写，键写错的 builder 永远不会被调用 | -| `src/tileops_cpu/ops/rms_norm.py`、`ops/gemm.py` | 每个算子一个模块，存放 kernel 实现和构造它的 builder。真实后端在 kernel 的构造函数里编译 | +| `src/tileops_cpu/ops/__init__.py` | `BUILDERS`，即这个 target 接管的全部 op。键按 manifest 中的拼写书写，键写错的 builder 永远不会被调用 | +| `src/tileops_cpu/ops/rms_norm.py`、`ops/gemm.py` | 每个 op 一个模块，存放 kernel 实现和构造它的 builder。真实 backend 在 kernel 的构造函数中编译 | | `tests/test_takeover.py` | 数值、校验、归一与输出 | | `tests/test_discovery.py` | entry point 与注册 | -| `tests/test_errors.py` | 三条错误路径：未登记的算子报错而不回退，未知的 target 报错，调用失败后算子不固定到任何 target | +| `tests/test_errors.py` | 三条错误路径：未登记的 op 报错而不回退，未知的 target 报错，调用失败后 op 不固定到任何 target | | `tests/test_memoization.py` | `build_kernel` 在什么条件下被重新调用 | -其中 `CpuRMSNorm` 在构造时**得不到行数**，正是「构造函数只接收编译期参数」这一条的体现。 +`CpuRMSNorm` 在构造时**得不到行数**，这体现了「构造函数只接收编译期参数」这条规则。 ### 运行测试 @@ -352,22 +323,22 @@ docker run --rm --gpus all -v "$PWD/..":/work -w /work \ cd /work/tileops-backend-example && python -m pytest -q' ``` -`tileops` 刻意不写在示例的依赖列表中。这个包扩展的是一个已经存在的安装，而在依赖中写上版本下限会解析到早于 `tileops.backend` 的发行版；由此产生的 `ImportError` 会被收入 `load_failures()`，呈现出来的是「这个后端不可用」，而真正的原因是 TileOPs 版本过旧。 +`tileops` 有意不写在示例的依赖列表中。这个包扩展的是一个已经存在的安装；在依赖中写上版本下限会解析到早于 `tileops.backend` 的发行版。由此产生的 `ImportError` 会被收入 `load_failures()`，表现为「这个 backend 不可用」，而真正的原因是 TileOPs 版本过旧。 -### 改造成面向真实硬件的后端 +### 改造成面向真实硬件的 backend -1. 复制该仓库，把 `tileops_cpu` 改为 `tileops_<硬件名>`，target 名同样改写。 -2. 修改 `target.py` 里的 `detect`，认领对应的设备类型。 -3. 把 `ops/` 下各模块的 kernel 替换为真实 kernel：构造时编译，调用 `__call__` 时启动。 -4. 选定第一个要接管的算子，照它的 manifest 签名编写 `build_kernel`，并在 `BUILDERS` 里加一行。 -5. [`tests/`](https://github.com/lcy-seso/tileops-backend-example/tree/main/tests) 中的四个文件大体可以直接沿用，替换其中的算子名与 target 名即可。 -6. 之后逐个算子增加 `build_kernel`，直到覆盖目标模型用到的全部算子。 +1. 复制该仓库，把 `tileops_cpu` 改为 `tileops_<硬件名>`，并相应改写 target 名。 +2. 修改 `target.py` 中的 `detect`，认领对应的设备类型。 +3. 把 `ops/` 下各模块中的 kernel 替换为真实 kernel：构造时编译，在 `__call__` 中启动。 +4. 选定第一个要接管的 op，按它的 manifest 签名编写 `build_kernel`，并在 `BUILDERS` 中加一行。 +5. [`tests/`](https://github.com/lcy-seso/tileops-backend-example/tree/main/tests) 中的四个文件大体可以直接沿用，只需替换其中的 op 名与 target 名。 +6. 之后逐个 op 增加 `build_kernel`，直到覆盖目标模型用到的全部 op。 ## 编写 kernel ### 签名来自 manifest -**编写 kernel 只需读 manifest，不必读 TileOPs 的源码。** builder 的签名就是该算子的 manifest 签名。以 [`src/tileops/manifest/spec/norm.yaml`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/manifest/spec/norm.yaml) 里的 `RMSNormFwdOp` 为例： +**编写 kernel 只需读 manifest，不必读 TileOPs 的源码。** builder 的签名就是该 op 的 manifest 签名。以 [`src/tileops/manifest/spec/norm.yaml`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/manifest/spec/norm.yaml) 中的 `RMSNormFwdOp` 为例： ```yaml signature: @@ -386,81 +357,81 @@ signature: def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): ``` -两点要注意。 +这个签名有两点约定： -- **参数按算子实例保存的值传入。** 构造时没给 `eps`，收到的就是 manifest 的默认值 `None`，含义与参考 API 相同，由 builder 按那个语义处理；没传的可选输入 `weight` 收到 `None`。 -- **返回值按 `signature.outputs` 的声明给出** —— 单输出返回张量，多输出按声明顺序返回 tuple，纯原地写入的算子返回 `None`。 +- **参数按 op 实例保存的值传入。** 构造时没有给出 `eps` 时，builder 收到的是 manifest 的默认值 `None`，其含义与参考 API 相同，由 builder 按该语义处理；没有传入的可选输入 `weight` 收到 `None`。 +- **返回值按 `signature.outputs` 的声明给出。** 单输出返回张量，多输出按声明顺序返回 tuple，纯原地写入的 op 返回 `None`。 ### 构造函数只接收编译期参数 -会被编译进生成代码的值 —— tile 尺寸、当作常量的维度、dtype —— 进构造函数，其余留给 `__call__`。 +会被编译进生成代码的值（tile 尺寸、当作常量的维度、dtype）传给构造函数，其余的值留给 `__call__`。 -对 decode 路径这是硬性要求：`seq_len` 逐步递增，batch 随 running set 变化，放进构造函数就是每一步重新编译。 +这条规则对 decode 路径是硬性要求：`seq_len` 逐步递增，batch 随 running set 变化，把它们放进构造函数会导致每一步都重新编译。 ### 形状由 manifest 规定 -算子层不改形状：kernel 收到的就是 manifest 声明的形状，需要哪种 layout 由它在自己的调用包装里处理。 +op 层不改变形状：kernel 收到的就是 manifest 声明的形状，kernel 需要的 layout 由它在自己的调用包装中处理。 -代码与 manifest 说的不一致时以 manifest 为准：输出 dtype、形状规则与参数类型都由它规定，kernel 不得改写。 +代码与 manifest 不一致时以 manifest 为准。输出 dtype、形状规则与参数类型都由 manifest 规定，kernel 不得改写。 -### kernel 服务不了这次调用时怎么报错 +### kernel 无法服务这次调用时如何报错 -kernel 服务不了这次调用就报错，不要降级处理。报错须给出两项信息： +kernel 无法服务这次调用时直接报错，不做降级处理。报错信息须给出两项内容： -- **未满足的是哪一项** —— dtype、形状、arch、无可用实现，还是编译失败。 -- **实际收到的值是什么。** +- **未满足的是哪一项**：dtype、形状、arch、无可用实现，还是编译失败。 +- **实际收到的值。** 只写「不支持」不构成有效的诊断信息。 ## 各阶段允许做什么 {#phase-limits} -decode 路径会被 CUDA graph 捕获，因此各阶段允许执行的操作分别规定如下： +decode 路径会被 CUDA graph 捕获，因此各阶段允许执行的操作规定如下： | 阶段 | 允许 | 不允许 | | --- | --- | --- | -| 查记忆表（键与重建条件见[后文](#memo)） | 一次字典查找 | 其他任何操作 | +| 查记忆表（键与重建条件见 [kernel 的重建条件](#memo)） | 一次字典查找 | 其他任何操作 | | `detect` | 一次谓词判断 | 任何 import，任何加锁 | | 构造 kernel | 选择实现、编译、分配显存、重新 import、建立 handle | 依赖真实张量的调优 | | 调用 kernel | 启动已编译的 kernel，经 torch allocator 分配输出 | 编译、惰性初始化、建立 handle、host 端同步 | -**模块顶层的 import 不得触发编译。** TileOPs 在构造第一个算子时 import 后端模块，编译应当发生在 `build_kernel` 被调用的时候。 +**模块顶层的 import 不得触发编译。** TileOPs 在构造第一个 op 时 import backend 模块，编译应当发生在 `build_kernel` 被调用时。 调用 kernel 还须满足两条与流有关的规则： -- **必须在当前流上启动**，在 CUDA 上即 `torch.cuda.current_stream(device)`，不得改用默认流。自带 launcher 的后端尤其容易违反这一条。 -- **内部分配的生命周期必须跨越异步执行。** 如果只把裸指针传给 launch，对象必须存活到该流执行完成为止。协议不提供 workspace，这一部分安全由后端自己保证。 +- **必须在当前流上启动。** 在 CUDA 上，当前流即 `torch.cuda.current_stream(device)`，不得改用默认流。自带 launcher 的 backend 尤其容易违反这一条。 +- **内部分配的生命周期必须覆盖异步执行。** 如果只把裸指针传给 launch，对应对象必须存活到该流执行完成。协议不提供 workspace，这部分安全由 backend 自己保证。 -调用方需要在捕获之前完成预热，即至少执行一次同形状的非捕获调用，因为构造 kernel 允许编译。捕获期间只允许「查表命中后直接调用」这一条路径。 +因为构造 kernel 时允许编译，调用方需要在捕获之前完成预热，即至少执行一次同形状的非捕获调用。捕获期间只允许「查表命中后直接调用」这一条路径。 ## 安装之后：三种状态 {#three-states} -一旦 `detect` 认领了某一类设备，该设备上的**所有**算子都由这个 target 服务，其中任何一个缺失都会报错，不会改用 TileOPs 自带的实现。唯一的例外是不自己构造 kernel 的复合算子。 +`detect` 认领某一类设备之后，该设备上的**所有** op 都由这个 target 服务，缺少其中任何一个都会报错，不会改用 in-tree 实现。唯一的例外是不自己构造 kernel 的复合 op。 -不回退的理由是：选中一个 target 就意味着这块设备属于另一套硬件，自带的 kernel 在它上面根本启动不了。真去回退，只会把一条清楚的「该 target 未实现此算子」换成一次难以理解的启动失败。 +不回退的理由是：选中一个 target 意味着这块设备属于另一套硬件，in-tree kernel 在它上面无法启动。如果回退，只会把一条清楚的「该 target 未实现此 op」错误换成一次难以理解的启动失败。 -因此安装之后，每个算子处于三种状态之一： +因此安装之后，每个 op 处于以下三种状态之一： | 状态 | 结果 | | --- | --- | -| 该 target 为这个算子注册了 `build_kernel` | 正常执行，整个算子由 target 服务 | -| 没有注册，且算子自己构造 kernel | 报错，指出这个 target 没有为该算子注册 builder，且不会改用自带实现 | -| 没有注册，且算子是复合算子 | 算子照常执行它的组合，每个子算子各自选定 target | +| 该 target 为这个 op 注册了 `build_kernel` | 正常执行，整个 op 由 target 服务 | +| 没有注册，且 op 自己构造 kernel | 报错，指出这个 target 没有为该 op 注册 builder，且不会改用 in-tree 实现 | +| 没有注册，且 op 是复合 op | op 照常执行它的组合，每个子 op 各自选定 target | -覆盖目标模型用到的每一个算子，因此是后端一侧的工作。算子那一侧的前提已经由设计保证：算子层按本次调用的输入给外部路径算记忆键（见[一次调用怎么走到 `build_kernel`](#from-op-layer)）。 +因此，覆盖目标模型用到的每一个 op 是 backend 一侧的工作。op 一侧的前提由设计保证：op 层按本次调用的输入为外部路径计算记忆键（见[一次调用如何走到 `build_kernel`](#from-op-layer)）。 ### 平台无关的前提 -target 定下来之前，算子层不查询与特定硬件绑定的信息 —— 例如 CUDA 的 SM 版本。查了就意味着：在没有该驱动的机器上，调用会在到达 `build_kernel` 之前失败，而失败的原因与这个后端毫无关系。 +target 确定之前，op 层不查询与特定硬件绑定的信息，例如 CUDA 的 SM 版本。如果查询，在没有该驱动的机器上，调用会在到达 `build_kernel` 之前失败，而失败原因与这个 backend 无关。 -万一在自己的硬件上撞到这种失败，调用栈会停在 TileOPs 内部、而不是后端的 `build_kernel` 里。那是 TileOPs 一侧的回退，提 issue 并附上调用栈。 +如果在自己的硬件上遇到这种失败，调用栈会停在 TileOPs 内部，而不是 backend 的 `build_kernel` 中。这属于 TileOPs 一侧的回归，应提交 issue 并附上调用栈。 -在看不到 GPU 的环境里，示例仓库的全部测试照常通过，没有跳过项；这次运行本身就检查了这条前提。 +在看不到 GPU 的环境中，示例仓库的全部测试照常通过，没有跳过项；这次运行本身检查了这条前提。 ## 错误信息与处理 -以下三条均为实测输出，分别对应一种成因和一种处理方式。 +以下三条均为实测输出，每条对应一种成因和一种处理方式。 -**未为该算子注册 builder：** +**未为该 op 注册 builder：** ``` OpNotAvailableError: target 'torch_cpu' registers no kernel builder for SoftmaxFwdOp; @@ -468,7 +439,7 @@ targets that do: []. There is no fall back to the in-tree implementation: those do not run on this target's devices. ``` -为该算子编写并注册一个 builder。 +处理方式是为该 op 编写并注册一个 builder。 **指定了未注册的 target：** @@ -476,17 +447,17 @@ do not run on this target's devices. UnknownTargetError: no backend registered target 'nope'; known targets: ['torch_cpu'] ``` -说明包没有安装成功，或者 target 名拼写有误。可以用 `tileops.backend.registered_targets()` 查看实际注册的内容。 +这说明包没有安装成功，或者 target 名拼写有误。`tileops.backend.registered_targets()` 可以列出实际注册的内容。 -**以 `target=BUILTIN` 强制使用 TileOPs 自带的实现：** +**以 `target=BUILTIN` 强制使用 in-tree 实现：** ``` OpNotAvailableError: RMSNormFwdOp's in-tree kernels do not run on cpu; known targets for this op: ['torch_cpu'] ``` -`BUILTIN` 显式绕过所有后端。自带实现无法在 CPU 张量上运行，这条错误正说明了「不改用自带实现」这条规则所要避免的后果。 +`BUILTIN` 显式绕过所有 backend。in-tree 实现无法在 CPU 张量上运行，这条错误展示的正是「不改用 in-tree 实现」这条规则所要避免的后果。 -**后端包 import 失败**时，TileOPs 会跳过它并发出一条警告，同时把原因收入 `load_failures()`。单个损坏的插件不会导致 TileOPs 无法导入。如果注册过程中途抛出异常，该后端本次注册的内容会**全部回滚**，注册表中不会留下一个只实现了一半的 target。 +**backend 包 import 失败**时，TileOPs 跳过它并发出一条警告，同时把原因收入 `load_failures()`。单个损坏的插件不会导致 TileOPs 无法导入。如果注册过程中途抛出异常，该 backend 本次注册的内容会**全部回滚**，注册表中不会留下只完成一半注册的 target。 ```python from tileops.backend import load_failures @@ -497,39 +468,39 @@ print(load_failures()) | 名字 | 定义 | 在 dispatch 中的位置 | | --- | --- | --- | -| **target** | 一套 kernel 的名字。一个后端发行版带来一套 kernel，为它起一个名字，例如 `"acme"` | 第一层：选中一个 target，本次调用的 kernel 就从它这一套里出 | -| **`detect`** | 后端写的一个函数，一个 target 一个 | 第一层怎么选：接收一块 `torch.device`，回答这类设备是不是自己这套 kernel 的目标设备；不是则返回 `False` | -| **`build_kernel`** | 后端为某个算子写的一个函数，一组 `(算子, target)` 一个 | 第二层：接收本次调用的描述，即各输入张量的 device、dtype、shape 与算子参数，在自己这套 kernel 中选定一个、构造好并返回 | +| **target** | 一套 kernel 的名字。一个 backend 发行版带来一套 kernel，并为它起一个名字，例如 `"acme"` | 第一层：选中一个 target 后，本次调用的 kernel 从它这一套中选出 | +| **`detect`** | backend 编写的一个函数，每个 target 一个 | 第一层的选择依据：接收一个 `torch.device`，回答这类设备是否是自己这套 kernel 的目标设备；不是则返回 `False` | +| **`build_kernel`** | backend 为某个 op 编写的一个函数，每组 `(op, target)` 一个 | 第二层：接收本次调用的描述，即各输入张量的 device、dtype、shape 与 op 参数，从自己这套 kernel 中选定一个，构造好并返回 | -**选择分两层：TileOPs 选 target，target 在自己那套 kernel 里选一个。** 第二层发生在 `build_kernel` 内部，协议不参与：这条路径上不存在 kernel 一级的概念、能力协商与候选筛选。TileOPs 确实执行的候选筛选 —— 可用性、适用范围、优先关系 —— 属于自带实现与范围较小的两种接入方式（[三种接入方式](#three-ways)），target 绕过它们。 +**选择分两层：TileOPs 选 target，target 在自己那套 kernel 中选一个。** 第二层发生在 `build_kernel` 内部，协议不参与，这条路径上没有 kernel 一级的概念、能力协商与候选筛选。TileOPs 实际执行的候选筛选（可用性、适用范围、优先关系）属于 in-tree 实现与范围较小的两种接入方式（见[三种接入方式](#three-ways)），target 绕过这些筛选。 -`detect` 只回答设备的归属，粒度到此为止。**本次调用是否受支持 —— 涉及 dtype、形状与参数组合 —— 由 `build_kernel` 回答**，因为只有它看得到完整的输入描述与参数；不支持时在那里报错。这些判断交给 `detect` 是做不到的，它只拿到一块 `torch.device`。 +`detect` 只回答设备的归属，不回答更细的问题。**本次调用是否受支持（涉及 dtype、形状与参数组合）由 `build_kernel` 回答**，因为只有它看得到完整的输入描述与参数，不支持时也在那里报错。`detect` 只拿到一个 `torch.device`，无法作出这些判断。 -TileOPs 不解析 `torch.device`，而是把它原样传给 `detect`。这样做是因为设备类型与 target 并不一一对应，有三种情形都会让解析出来的字符串失去意义： +TileOPs 把 `torch.device` 原样传给 `detect`，自己不解析它。原因是设备类型与 target 并不一一对应，以下三种情形都会使解析出来的字符串失去意义： - 同一个 device type 可能对应多套 kernel，分属不同厂商。 -- 部分硬件经 `privateuseone` 接入，字符串里不含任何厂商信息。 -- 还有一些后端要读取环境变量、或调用厂商 runtime 才能作出判断。 +- 部分硬件经 `privateuseone` 接入，字符串中不含任何厂商信息。 +- 还有一些 backend 要读取环境变量或调用厂商 runtime 才能作出判断。 -## 算子层的契约 +## op 层的契约 -以下七项是算子层对所有 target 的契约：功能由算子层实现，后端直接重用，无需自行实现。表中按后端作者接触到它们的先后排列： +以下七项是 op 层对所有 target 的契约。这些功能由 op 层实现，backend 直接重用，无需自行实现。表中各项按 backend 作者接触到它们的先后排列： -| # | 算子层提供 | 对后端意味着什么 | +| # | op 层提供 | 对 backend 意味着什么 | | --- | --- | --- | -| 1 | torch 侧的公开 API 与参数语义 | 这个算子如何被调用、参数名与各参数的含义均已确定，后端既不定义也不能改动 | -| 2 | manifest 校验 | dtype 或形状不合规的调用被算子层拒绝，不会到达后端 | -| 3 | 参数按名字传入 | 后端收到的参数名是 manifest `params` 的名字，值是算子实例保存的值。manifest 中默认为 null 的参数，传入的是算子选定的值，可能是 `None`，也可能是一个具体数值 | -| 4 | 输入的连续性归一 | 本次调用不写入的输入都转成连续张量；被写入的输入按调用方传入的原样交给后端，除非 manifest 声明它 `contiguous: true` | -| 5 | kernel 的记忆与重用 | 构造函数按特化调用一次：设备与输入签名相同的后续调用直接使用上一次的返回值。因此构造函数内部可以编译，算子层保证它不会被重复调用 | -| 6 | `torch.compile` 与 CUDA graph 的边界 | 算子层把一次调用包成不透明算子并另配一个 fake，使编译器在不执行的前提下也能推出输出的形状与 dtype。**后端的 kernel 不为编译做任何事**，细节见[接入 torch.compile](torch-compile.md) | -| 7 | roofline、profile 与数值测试 | 算子层已有的测试会用后端的 kernel 跑一遍，与 manifest 的 `ref_api` 比对数值；性能报告照常产出 | +| 1 | torch 侧的公开 API 与参数语义 | 这个 op 如何被调用、参数名与各参数的含义均已确定，backend 既不定义也不能改动 | +| 2 | manifest 校验 | dtype 或形状不合规的调用由 op 层拒绝，不会到达 backend | +| 3 | 参数按名字传入 | backend 收到的参数名是 manifest `params` 的名字，值是 op 实例保存的值。manifest 中默认为 null 的参数，传入的是 op 选定的值，可能是 `None`，也可能是一个具体数值 | +| 4 | 输入的连续性归一 | 本次调用不写入的输入都转成连续张量；被写入的输入按调用方传入的原样传给 backend，除非 manifest 声明它 `contiguous: true` | +| 5 | kernel 的记忆与重用 | 构造函数对每种特化调用一次：设备与输入签名相同的后续调用直接使用上一次的返回值。因此构造函数内部可以编译，op 层保证它不会被重复调用 | +| 6 | `torch.compile` 与 CUDA graph 的边界 | op 层把一次调用包装成不透明 op，并另配一个 fake，使编译器不执行也能推出输出的形状与 dtype。**backend 的 kernel 不为编译做任何事**，细节见[接入 torch.compile](torch-compile.md) | +| 7 | roofline、profile 与数值测试 | op 层已有的测试会用 backend 的 kernel 运行一遍，与 manifest 的 `ref_api` 比对数值；性能报告照常产出 | -七项均与硬件无关，每个 target 得到的完全相同；接入第三方后端不得绕过其中任何一项，也不得另行实现。 +七项均与硬件无关，每个 target 得到的完全相同。接入第三方 backend 时不得绕过其中任何一项，也不得另行实现。 -TileOPs 自带的 kernel（[`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)）是**默认实现**：它没有 target 名，也不进注册表。 +TileOPs 的 in-tree kernel（[`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)）是**默认实现**：它没有 target 名，也不进入注册表。 -**默认走自带实现。** 没有装后端、没有指名 `target=`、也没有设置进程默认值时，调用由自带实现服务。装上一个后端之后，要它的 `detect` 认领了这块设备，或者它被 `target=` 与 `set_default_target` 指名，该算子的 kernel 才换成这个后端的那一套。 +**默认使用 in-tree 实现。** 没有安装 backend、没有通过 `target=` 指名、也没有设置进程默认值时，调用由 in-tree 实现服务。安装一个 backend 之后，只有当它的 `detect` 认领了这块设备，或者它被 `target=` 或 `set_default_target` 指名时，该 op 的 kernel 才换成这个 backend 的那一套。 ## kernel 的重建条件 {#memo} @@ -537,18 +508,18 @@ TileOPs 按**设备加输入签名**记住 `build_kernel` 的返回值。这个 > 本次调用张量所在的设备，加上按 `signature.inputs` 顺序逐项取出的 `(dtype, shape)`；声明为 optional 而本次没有传入的输入，这一项记 `None`。 -也就是说，**设备与输入签名都相同的两次调用，TileOPs 会把同一个 kernel 交回给后端**，不再调用 `build_kernel`。同一个 target 的第二块卡会重新构造一次，因为为一块卡编译出的产物不一定能在另一块卡上启动。参数不进入这个键，它们对一个算子实例而言是固定的。 +因此，**设备与输入签名都相同的两次调用，TileOPs 使用同一个 kernel**，不再调用 `build_kernel`。同一个 target 的第二块卡会重新构造一次，因为为一块卡编译出的产物不一定能在另一块卡上启动。参数不进入这个键，因为它们对一个 op 实例而言是固定的。 -算子层这一侧怎么查表、未命中时怎么调 `build_kernel`，见[一次调用怎么走到 `build_kernel`](#from-op-layer)。 +op 层如何查表、未命中时如何调用 `build_kernel`，见[一次调用如何走到 `build_kernel`](#from-op-layer)。 -两点由此而来： +由此得出两点： -- **条目不一定一直保留。** 一次调用失败时，算子会撤销它的 target 判定，并清空记忆表。后端不得假设自己返回的可调用对象一直存活；它所依赖的资源应由它自己持有引用。 -- **需要更细或更粗的粒度，都在后端一侧解决。** 更细的区分在后端内部处理；希望减少重建次数，可以在 `build_kernel` 内部另加一层缓存。 +- **条目不一定一直保留。** 一次调用失败时，op 会撤销它的 target 判定并清空记忆表。backend 不得假设自己返回的可调用对象一直存活，它所依赖的资源应由它自己持有引用。 +- **需要更细或更粗的粒度时，都在 backend 一侧解决。** 更细的区分在 backend 内部处理；需要减少重建次数时，在 `build_kernel` 内部另加一层缓存。 ## 调用方可用的接口 -以下接口面向使用者，后端作者不需要调用，但调试时用得上： +以下接口面向使用者，backend 作者不需要调用，但调试时会用到： ```python from tileops.backend import ( @@ -561,17 +532,23 @@ set_default_target("torch_cpu") # 进程默认，优先于设备探测 set_default_target(BUILTIN) # 全局关闭替换 ``` -target 的选取顺序是：构造参数 `target=`，其次是进程默认值，最后是设备探测。`BUILTIN` 强制走 TileOPs 自带的实现。指定的 target 没有注册或没有实现该算子即报错，不会改用其他 target。 +target 按以下顺序选取： + +1. 构造参数 `target=`； +2. 进程默认值； +3. 设备探测。 + +`BUILTIN` 强制使用 in-tree 实现。指定的 target 没有注册或没有实现该 op 时直接报错，不会改用其他 target。 ## 协议不支持的情形 -以下情形不在这套协议的范围内，各有其理由： +以下情形不在这套协议的范围内，理由见表： | 不支持 | 理由 | | --- | --- | -| 同一个 target 上存在多个后端 | 一个 target 对应一套 kernel、一个提供者。重复注册同一组 `(算子, target)` 会直接报错，因为这说明安装了两个都声明服务该 target 的包 | -| 跨 target 回退 | 指定的 target 没有实现即报错，不会改用其他 target 执行 | -| 后端改变输入形状，或代替调用方还原输出 | 这是算子层对所有 target 统一提供的服务；要改动就对所有 target 一起改动 | -| 一次调用跨越多个设备 | 所有输入位于同一设备，manifest 声明 `device: cpu` 的张量除外 | -| 调用方提供 workspace 或显式 stream | 后端需要的只是当前流，而 torch 的流本身就是隐式的当前值 | -| 与 autograd 联动 | 这条调用链服务推理，fwd 与 bwd 各自是独立的算子 | +| 同一个 target 上存在多个 backend | 一个 target 对应一套 kernel 与一个提供者。重复注册同一组 `(op, target)` 会直接报错，因为这说明安装了两个都声明服务该 target 的包 | +| 跨 target 回退 | 指定的 target 没有实现时直接报错，不会改用其他 target 执行 | +| backend 改变输入形状，或代替调用方还原输出 | 这是 op 层对所有 target 统一提供的服务；如需改动，就对所有 target 一起改动 | +| 一次调用跨越多个设备 | 所有输入位于同一设备上，manifest 声明 `device: cpu` 的张量除外 | +| 调用方提供 workspace 或显式 stream | backend 需要的只是当前流，而 torch 的流本身就是隐式的当前值 | +| 与 autograd 联动 | 这条调用链服务推理，fwd 与 bwd 各自是独立的 op | diff --git a/docs/index.md b/docs/index.md index 07a2ed68..57f089a4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,20 +1,24 @@ # TileOPs TileOPs is an operator library for large-model inference, built on -[TileLang](https://github.com/tile-ai/tilelang), where one set of operator -interfaces can be implemented by different backends on different hardware. - -What sets it apart from a hand-written library is how it is organised: every -operator is declared as a spec first, and an agent then derives the -implementation from that spec. The spec is the only input to generation and the -standard the result is judged by — correctness against the reference the spec -names, performance against the bound the roofline model gives, neither of them a -judgement call. An implementation can therefore be regenerated from its spec, -while the reverse does not hold. - -To a caller it is simply a set of operators: shapes and dtype come from the call, -the specialized kernel is built and cached on first use and works under CUDA -graphs afterwards, and each op declares whether it supports +[TileLang](https://github.com/tile-ai/tilelang). One set of op interfaces can be +implemented by different backends on different hardware. + +TileOPs differs from a hand-written operator library in how it is organised: every +op is first declared as a spec, and an agent then generates the implementation from +that spec. The spec is the only input to code generation and the standard the result +is accepted against: + +- correctness is judged against the reference implementation the spec names; +- performance is judged against the bound the roofline model gives. + +Neither check depends on human judgement. An implementation can therefore be +regenerated from its spec at any time, while a spec cannot be derived from an +implementation. + +To a caller, TileOPs is a set of ops that can be called directly. Shapes and dtype +are fixed at call time; the specialized kernel is built and cached on first use and +can then be used with CUDA graphs. Each op declares whether it supports `torch.compile(fullgraph=True)`. ## Installation @@ -25,8 +29,8 @@ pip install tileops ## Quick Start -An op commits to nothing at construction. Shapes and dtype come from the inputs -of the call, and the specialized kernel is built and cached on first use. +An op binds no shape at construction. Shapes and dtype come from the tensors passed +to the call, and the specialized kernel is compiled and cached on the first call. ```python import torch @@ -42,12 +46,12 @@ flops, nbytes = op.eval_roofline() # what the call had to do and move ## Where to go next -- [User Guide](user-guide/index.md) — writing a spec, bringing an op into - `torch.compile`, how a benchmark is timed, and adding a hardware backend -- [API Reference](api/index.md) — constructor parameters and call signatures, one - page per op family -- [Benchmarks](benchmarks/index.md) — measured nightly on an H200 against the - alternatives +- [User Guide](user-guide/index.md): reading and writing the manifest, bringing an op + into `torch.compile`, how a benchmark is timed, and adding a hardware backend. +- [API Reference](api/index.md): the constructor parameters and call signatures of + each op family. +- [Benchmarks](benchmarks/index.md): measured nightly on an H200, each workload + against other implementations. ## Links diff --git a/docs/index.zh.md b/docs/index.zh.md index 2803ff7b..1eaf4269 100644 --- a/docs/index.zh.md +++ b/docs/index.zh.md @@ -1,10 +1,15 @@ # TileOPs -TileOPs 是一个面向大模型推理的算子库，构建在 [TileLang](https://github.com/tile-ai/tilelang) 之上，同一套算子接口可以由不同后端在不同硬件上实现。 +TileOPs 是一个面向大模型推理的算子库，构建在 [TileLang](https://github.com/tile-ai/tilelang) 之上。同一套 op 接口可以由不同 backend 在不同硬件上实现。 -它与手写算子库的不同之处在于组织方式：每个算子先以一份 spec 声明，再由 agent 依据这份 spec 生成实现。spec 既是代码生成的唯一依据，也是验收的标准 —— 正确性对照 spec 指定的参考实现，性能对照 roofline 模型给出的上界，两项都不依赖人的判断。因此一个实现可以随时从 spec 重新生成，而反过来做不到。 +TileOPs 与手写算子库的区别在于组织方式：每个 op 先以一份 spec 声明，再由 agent 依据这份 spec 生成实现。spec 是代码生成的唯一依据，也是验收的标准： -对使用者而言，它就是一批可以直接调用的算子：形状与 dtype 在调用时确定，特化后的 kernel 在首次使用时构造并缓存，随后可以与 CUDA graph 配合使用；每个算子各自声明是否支持 `torch.compile(fullgraph=True)`。 +- 正确性以 spec 指定的参考实现为准； +- 性能以 roofline 模型给出的上界为准。 + +两项验收都不依赖人的判断。因此一个实现可以随时从 spec 重新生成，spec 却不能从实现反推。 + +对使用者而言，TileOPs 提供一批可以直接调用的 op。形状与 dtype 在调用时确定；特化后的 kernel 在首次使用时构造并缓存，之后可以与 CUDA graph 配合使用。每个 op 各自声明是否支持 `torch.compile(fullgraph=True)`。 ## 安装 @@ -14,7 +19,7 @@ pip install tileops ## 快速开始 -算子在构造时不绑定任何形状。形状和 dtype 取自调用传入的张量，特化后的 kernel 于首次调用时编译并缓存。 +op 在构造时不绑定任何形状。形状和 dtype 取自调用时传入的张量，特化后的 kernel 在首次调用时编译并缓存。 ```python import torch @@ -28,11 +33,11 @@ d = op(a, b) # -> [M, N] flops, nbytes = op.eval_roofline() # 本次调用所需的计算量与访存量 ``` -## 从这里继续 +## 后续阅读 -- [使用指南](user-guide/index.md) —— 读写 manifest、接入 `torch.compile`、benchmark 怎么计时、接入新硬件后端 -- [API 参考](api/index.md) —— 各算子族的构造参数与调用方式 -- [性能数据](benchmarks/index.md) —— 每晚在 H200 上实测，逐个 workload 与其他实现对比 +- [使用指南](user-guide/index.md)：读写 manifest、接入 `torch.compile`、benchmark 的计时方法、接入新硬件 backend。 +- [API 参考](api/index.md)：各 op family 的构造参数与调用方式。 +- [性能数据](benchmarks/index.md)：每晚在 H200 上实测，逐个 workload 与其他实现对比。 ## 相关链接 diff --git a/docs/manifest.md b/docs/manifest.md deleted file mode 100644 index f5c0b120..00000000 --- a/docs/manifest.md +++ /dev/null @@ -1,291 +0,0 @@ -# Reading and writing an op's spec - -A conventional operator library is organised around its implementations: kernels are -written and tuned one at a time, and what shapes and dtypes each supports, and how fast it -runs, gets described afterwards. - -TileOPs is organised the other way round: an op's specification is declared first, and -the implementation is derived from it. That specification is the op's **spec**, a YAML -entry under -[`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec), -in the file named after its family; those files together are the manifest. - -**A spec makes the op an input to the whole system.** Every stage reads the same -declaration rather than reading the implementation: - -| Consumer | Reads from the spec | Produces | -| --- | --- | --- | -| The op layer | `signature` | the checks around every call, output shape inference, the dtype check, and the operator `torch.compile` sees | -| [The contract tests](https://github.com/tile-ai/TileOPs/tree/main/tests) | `workloads` | one call per workload row and dtype case, run through the op | -| [The nightly benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks) | `workloads` | the device time of each of those calls | -| [Roofline](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/perf) | `roofline` | the FLOPs and bytes one call moves — the denominator of efficiency | -| This site | `signature`, `workloads` | the shapes printed under each row of the Benchmarks pages | -| CI's [spec validator](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py) | every field | the check that declaration and implementation agree — see [the spec validator](#spec-validator) | - -**Every row presupposes a spec**: without one there is no generated validation, no -contract test, no performance data, and nothing in CI holding a regression back. -**Writing a spec is not documenting the op; it is connecting the op to that flow.**{ .keystone } - -This page covers what an entry contains, how to read one, how to write one in five steps, -four forms that recur across the manifest, and what the validator checks. The complete -rules are in [Op Manifest](design/manifest.md); this page does not restate them. - -## What an entry contains - -Each family file is a mapping `op name → entry`; a large family shards into -`_.yaml`. The files merge at load time, and a duplicate op name is an -error. Algebraic data types that several entries share live in `spec/types.yaml`. - -The key is the op's Python class name, `{Name}[{Fwd|Bwd}]Op` — the direction suffix is -required once the other direction also has an entry — and the validator requires -`cls.__name__` to equal it character for character. - -| Field | Required | Contents | -| --- | --- | --- | -| `family` | yes | the public module: the op is importable as `tileops..` | -| `status` | yes | `implemented`, or `spec-only` while no conforming implementation exists | -| `ref_api` | no | the qualified name of the API the op follows, e.g. `torch.nn.functional.rms_norm` | -| `signature` | yes | the op's type, below | -| `workloads` | yes | the calls tests and benchmarks run | -| `roofline` | yes | the cost of one call, specified in [Roofline](design/roofline.md) | -| `composition` | no | for a composite op, its stages in order: the sub-op classes it may hold and its own kernel keys | - -The signature is a function type over named type indices: - -| Sub-field | Contents | -| --- | --- | -| `forall` | every free index with its kind: `Dim` (an axis length), `Shape` (a tuple of axes), `DType[...]` (one of a set of dtypes), `Seq[Int]` (a value list only a generator takes) | -| `params` | the `__init__` parameters: `type`, optional `default` and `kw_only` | -| `inputs` / `outputs` | tensors, each `{dtype, shape}` written in those indices, plus presence and effect flags | -| `types` | type families: a shape chosen by the value of a flag | -| `let` | quantities derived from the indices | -| `shape_rules` | refinements: predicates on index values | -| `dtype_combos` | the supported combinations of several `DType` indices, where not every one works | - -Key order is position — `params` in `__init__`, `inputs` in `forward`, `outputs` in the -returned tuple — so reordering is a breaking change. - -## Reading a spec - -`RMSNormFwdOp`, with two of its eight workload rows: - -```yaml -RMSNormFwdOp: - ref_api: torch.nn.functional.rms_norm - family: norm - status: implemented - signature: - forall: {B: Shape, T: "DType[float16 | bfloat16]"} - params: - normalized_shape: {type: "list[int] | tuple[int, ...]"} - eps: {type: "float | None", default: null} - inputs: - x: {dtype: T, shape: "[*B, *normalized_shape]"} - weight: {dtype: T, shape: "[*normalized_shape]", optional: true} - outputs: - output: {dtype: T, shape: "[*B, *normalized_shape]"} - shape_rules: - - "len(normalized_shape) > 0" - workloads: - - {B: [2048], normalized_shape: [4096], eps: 1.0e-06, some: [weight], - dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-prefill} - - {B: [4, 2048], normalized_shape: [128], dtype_cases: [{T: bfloat16}], label: qk-norm-head} - roofline: - flops: "(4 if present(weight) else 3) * prod(B) * prod(normalized_shape)" -``` - -Read it in five steps: - -1. **`forall`** — what varies between calls. `B` is the leading axes of any rank, and - `T` the dtype. -2. **`inputs` and `outputs`** — a shared name is an equality. `x`, `weight` and `output` - all have dtype `T`, so the three agree; `x` and `output` have one shape term, so they - have one shape. `weight` is `optional: true`, so a call may omit it. -3. **`params`** — `normalized_shape` is a construction argument that also appears in a - shape, spliced with `*`: the trailing axes of `x` must equal it. -4. **`shape_rules`** — what the shapes cannot say. Here, that at least one axis is - normalized. -5. **`workloads`** — each row with each of its `dtype_cases` is one call. A row gives the - indices (`B`), the construction arguments, and in `some` the optional tensors it - passes. Its case id is the label followed by the dtype values, `llama-8b-prefill-float16`: - that is the id a benchmark row carries on the nightly and on this site. - -Reading specs programmatically: - -```python -from tileops.manifest import load_manifest, load_workloads - -ops = load_manifest() # every entry, merged -list(ops["RMSNormFwdOp"]["signature"]["inputs"]) # ['x', 'weight'] -load_workloads("RMSNormFwdOp") # that op's workload rows -``` - -## Writing a spec {#writing-a-spec} - -Five steps, each one checkable immediately. - -1. **Name it and pick the family.** The key is the class name, and the entry goes in - `spec/.yaml`. -2. **Write the signature.** Declare each axis length, shape and dtype that varies in - `forall`, and write every tensor's `dtype` and `shape` in those indices. `params` - is the op's `__init__` parameter list, less the execution-policy parameters the code - owns (`target`, `kernel_map`, `tune`). Optional inputs come after the required ones. - Declare what the reference API supports, not what the current kernel does. -3. **Write the refinements.** `shape_rules` hold predicates on index values, such as - `H % G == 0`; a derived quantity is a `let`; a shape a flag chooses is a type family. - A rule never reads a tensor (`x.shape`, `x is None`): presence is `present(x)`. -4. **Write `workloads`.** Each row gives exactly the indices no generator determines, - every construction parameter without a default, `some` for the optional tensors it - passes, `dtype_cases` where the entry has `DType` indices (a dtype parameter is - written as a parameter), and a `label`. Each optional tensor of an implemented entry is - passed in at least one row and omitted in at least one. The label is part of the case - id, which keys nightly history, so renaming it breaks that history. -5. **Write `roofline`.** Inline `flops` (and `bytes`, where the traffic is not simply - each tensor read or written once) over the same indices, or a `func` that computes - both from the checked call. - -To land an interface before its implementation, write `status: spec-only`: the checks -that need code are skipped. `implemented` turns them on. - -## Four recurring forms - -### A flag chooses a shape - -The layout flags of `GemmFwdOp` decide which axis of `a` carries M. A type family in -`types` states both cases, and each input applies it: - -```yaml - signature: - types: - Mat: - params: {t: Bool, R: Dim, C: Dim} - match: t - cases: - - {when: false, is: "[R, C]"} - - {when: true, is: "[C, R]"} - forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"} - params: - trans_a: {type: bool, default: false} - trans_b: {type: bool, default: true} - inputs: - a: {dtype: T, shape: "Mat[trans_a, M, K]"} - b: {dtype: T, shape: "Mat[trans_b, K, N]"} - outputs: - d: {dtype: T, shape: "[M, N]"} -``` - -### Optional inputs as a switch - -The affine transform of `GroupNormFwdOp` is two optional inputs, and no `affine` flag -beside them: one fact is stated in one place. A workload row passes them through `some`, -and the roofline reads presence with `present`: - -```yaml - signature: - forall: {B: Dim, C: Dim, L: Shape, T: "DType[float32 | float16 | bfloat16]"} - params: - num_groups: {type: int} - eps: {type: float, default: 1.0e-05} - inputs: - x: {dtype: T, shape: "[B, C, *L]"} - weight: {dtype: T, shape: "[C]", optional: true} - bias: {dtype: T, shape: "[C]", optional: true} - outputs: - output: {dtype: T, shape: "[B, C, *L]"} - shape_rules: - - "num_groups > 0 and C % num_groups == 0" - - "B * (C // num_groups) * prod(L) != 1" - workloads: - - {B: 8, C: 128, L: [32, 32], num_groups: 32, dtype_cases: [{T: float16}], label: image} - - {B: 8, C: 128, L: [32, 32], num_groups: 32, some: [weight, bias], - dtype_cases: [{T: float16}], label: image-affine} - roofline: - flops: "(5 + (1 if present(weight) else 0) + (1 if present(bias) else 0)) * B * C * prod(L)" -``` - -(Excerpt: the entry has more rows and dtypes.) An op may dispatch on whether an optional -input was passed; it may not read the tensor's contents to decide. - -### An input the op writes - -`state` in `SSDDecodeFwdOp` is written in place by each decode step, so it declares -`mutated: true` and stays an input; the call returns only `y_out`. The generated operator -names exactly the inputs marked `mutated` as the ones it writes. - -```yaml - inputs: - A: {dtype: float32, shape: "[H, P, N]"} - dt: {dtype: float32, shape: "[B, H, P]"} - x: {dtype: T, shape: "[B, H, P]"} - B_in: {dtype: T, shape: "[B, G, N]"} - C_in: {dtype: T, shape: "[B, G, N]"} - state: {dtype: float32, shape: "[B, H, P, N]", mutated: true, contiguous: true} - outputs: - y_out: {dtype: float32, shape: "[B, H, P]"} -``` - -An output the caller may supply declares `buffer: out` instead: `forward` then takes an -`out` argument, writes it, and returns it. - -### A constraint on a metadata tensor's contents - -A tensor of offsets or lengths gets its values from a generator in `values`, and states -what its contents must satisfy in `requires`. From `GroupedQueryAttentionVarlenFwdOp`, -where `q_lens` is a `Seq[Int]` index and `T_q` the total query length: - -```yaml -cu_seqlens_q: {dtype: int32, shape: "[B + 1]", values: "prefix_sum(q_lens)", - requires: ["prefix_offsets(T_q)"]} -``` - -`B` is solved from the generated tensor, so a row gives `q_lens`, not `B`. - -## Rules at a glance - -**The signature** - -- **Order is position.** `params` order is `__init__` order, `inputs` order is `forward` - order, `outputs` order is return order; reordering is a breaking change. -- **Write against the reference.** Dtypes and parameters follow the authoritative - reference, never the current code; code that disagrees is fixed, with the entry - `spec-only` until it conforms. -- **A shared name is an equality.** Tensors of one shape write one shape term; a - relationship that is not an equality of names is a refinement or a `let`. - -**Rules and presence** - -- **Refinements read indices, never tensors.** No `x.shape`, no `x is None`, no - `isinstance`; presence is `present(x)`. -- **Presence is a switch, contents are not.** An op may select its implementation from - parameters and tensor presence; tensor contents are computation input only. -- **Both sides of an optional input get measured.** Passed and omitted each need a row, - counted per input. - -**Outputs** - -- **Output arity is fixed.** One entry has one set of outputs on every call; an op whose - return changes with a switch is two entries. -- **A written input stays an input.** It declares `mutated: true` and is not listed in - `outputs`. - -## The spec validator {#spec-validator} - -Validation is -[`scripts/validate_manifest.py`](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py), -and a spec can be run through it the moment it is written: - -```bash -python scripts/validate_manifest.py # every entry -python scripts/validate_manifest.py --check-op RMSNormFwdOp # one entry -python scripts/validate_manifest.py --levels schema,signature # skip the benchmark scan -``` - -| Level | Checks | -| --- | --- | -| `schema` | top-level fields, `family`, `ref_api`, `composition`, `roofline`, and `types.yaml` | -| `signature` | the signature on every combination of its discriminants, every workload row instantiated, and effects; for an implemented entry, the class's `__init__`, `forward` and declared kernels and sub-ops | -| `bench` | every benchmark takes its calls from the manifest and its roofline from the op | - -Checks that read code are skipped for a `spec-only` entry. Kernel selection, -multi-kernel ordering, accumulator dtypes, workspaces, tile sizes and autotuning -configuration are not in the manifest, so the validator does not see them. diff --git a/docs/manifest.zh.md b/docs/manifest.zh.md deleted file mode 100644 index ace0ff24..00000000 --- a/docs/manifest.zh.md +++ /dev/null @@ -1,230 +0,0 @@ -# 读写 manifest - -传统算子库以实现为中心：算子逐个写、逐个调优，支持哪些形状、哪些 dtype、跑多快，都由实现事后说明，文档写的是追述。 - -TileOPs 的组织方式相反：算子的规格先声明，实现由规格推导。每个算子的规格称为它的 **spec**，是 [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec) 下以 family 命名的 YAML 文件里的一个条目；这些文件合起来就是 manifest。 - -**一个算子有了 spec，就成为整个系统的数据输入。** 各个环节读同一份声明，而不是各自去读实现： - -| 谁消费 | 读 spec 里的什么 | 产出 | -| --- | --- | --- | -| 算子层 | `signature` | 每次调用前后的检查、输出形状推导、dtype 检查，以及 `torch.compile` 看到的 operator | -| [契约测试](https://github.com/tile-ai/TileOPs/tree/main/tests) | `workloads` | 每个 workload 行的每个 dtype case 对应一次调用，交给算子执行 | -| [每晚的 benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks) | `workloads` | 这些调用各自的 device time | -| [roofline](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/perf) | `roofline` | 一次调用的计算量与访存量，效率的分母 | -| 本文档站 | `signature`、`workloads` | Benchmarks 页每一行下面列出的形状 | -| CI 的 [spec 校验器](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py) | 全部字段 | 检查声明与实现是否一致，见 [spec 校验器](#spec-validator) | - -**每一行都以 spec 为前提**：没有 spec，就没有生成的校验、没有契约测试、没有性能数据，CI 也拦不下任何回退。**为一个算子写 manifest 不是补文档，而是把它接进这条数据流。**{ .keystone } - -本页依次讲：一个条目的构成、怎么读、怎么分五步写、manifest 里反复出现的四种写法，以及校验器查什么。完整规则见 [Op Manifest](design/manifest.md)，本页不重复。 - -## 一个条目的构成 - -每个 family 文件是一个 `算子名 → 条目` 的映射；大的 family 拆成若干个 `_.yaml`。加载时各文件合并，算子名重复即报错。多个条目共用的代数数据类型写在 `spec/types.yaml`。 - -条目的键是算子的 Python 类名 `{Name}[{Fwd|Bwd}]Op`，另一个方向也有条目时方向后缀必须写，校验器要求 `cls.__name__` 与它逐字相同。 - -| 字段 | 必填 | 内容 | -| --- | --- | --- | -| `family` | 是 | 公开模块：算子从 `tileops..` 导入 | -| `status` | 是 | `implemented`；还没有符合 spec 的实现时写 `spec-only` | -| `ref_api` | 否 | 算子语义所依照的 API 的全限定名，如 `torch.nn.functional.rms_norm` | -| `signature` | 是 | 算子的类型，见下表 | -| `workloads` | 是 | 测试与 benchmark 执行的调用 | -| `roofline` | 是 | 一次调用的开销，规范见 [Roofline](design/roofline.md) | -| `composition` | 否 | 复合算子按顺序列出的各阶段：可能持有的子算子类，以及它自己的 kernel key | - -签名是一个以具名类型变量为参数的函数类型： - -| 子字段 | 内容 | -| --- | --- | -| `forall` | 每个自由类型变量及其种类：`Dim`（一个轴的长度）、`Shape`（若干个轴）、`DType[...]`（一组 dtype 中的一个）、`Seq[Int]`（只作生成器参数的整数列表） | -| `params` | `__init__` 的参数：`type`，可选的 `default` 与 `kw_only` | -| `inputs` / `outputs` | 张量，每个写成 `{dtype, shape}`，用上面的类型变量表达，另可带是否可选与副作用的标记 | -| `types` | 类型族：按某个开关的取值选定形状 | -| `let` | 由类型变量算出的量 | -| `shape_rules` | 约束：关于类型变量取值的谓词 | -| `dtype_combos` | 多个 `DType` 变量不是任意组合都支持时，列出支持的组合 | - -键的顺序就是位置：`params` 对应 `__init__`，`inputs` 对应 `forward`，`outputs` 对应返回的 tuple，调换顺序是不兼容的改动。 - -## 读一份 spec - -`RMSNormFwdOp`，八个 workload 行里取两行： - -```yaml -RMSNormFwdOp: - ref_api: torch.nn.functional.rms_norm - family: norm - status: implemented - signature: - forall: {B: Shape, T: "DType[float16 | bfloat16]"} - params: - normalized_shape: {type: "list[int] | tuple[int, ...]"} - eps: {type: "float | None", default: null} - inputs: - x: {dtype: T, shape: "[*B, *normalized_shape]"} - weight: {dtype: T, shape: "[*normalized_shape]", optional: true} - outputs: - output: {dtype: T, shape: "[*B, *normalized_shape]"} - shape_rules: - - "len(normalized_shape) > 0" - workloads: - - {B: [2048], normalized_shape: [4096], eps: 1.0e-06, some: [weight], - dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-prefill} - - {B: [4, 2048], normalized_shape: [128], dtype_cases: [{T: bfloat16}], label: qk-norm-head} - roofline: - flops: "(4 if present(weight) else 3) * prod(B) * prod(normalized_shape)" -``` - -分五步读： - -1. **`forall`**：调用之间哪些量会变。`B` 是任意个数的前导轴，`T` 是 dtype。 -2. **`inputs` 与 `outputs`**：同名即相等。`x`、`weight`、`output` 的 dtype 都是 `T`，三者一致；`x` 与 `output` 写的是同一个形状，所以形状相同。`weight` 标了 `optional: true`，调用时可以不传。 -3. **`params`**：`normalized_shape` 是构造参数，同时用 `*` 展开进形状里：`x` 的末尾几个轴必须等于它。 -4. **`shape_rules`**：形状表达不了的约束。这里是至少归一化一个轴。 -5. **`workloads`**：每一行配上它的每个 `dtype_cases` 就是一次调用。一行给出类型变量（`B`）、构造参数，并在 `some` 里列出这次传入的可选张量。它的 case id 是 label 后接 dtype 取值，如 `llama-8b-prefill-float16`，nightly 与本站的 benchmark 行用的就是这个 id。 - -在代码里读 spec： - -```python -from tileops.manifest import load_manifest, load_workloads - -ops = load_manifest() # every entry, merged -list(ops["RMSNormFwdOp"]["signature"]["inputs"]) # ['x', 'weight'] -load_workloads("RMSNormFwdOp") # that op's workload rows -``` - -## 写一份新 spec {#writing-a-spec} - -五步，每一步写完都能立即检查。 - -1. **起名，定 family。** 键是类名，条目写进 `spec/.yaml`。 -2. **写签名。** 会变的轴长、形状、dtype 都在 `forall` 里声明，每个张量的 `dtype` 与 `shape` 用这些类型变量来写。`params` 是算子 `__init__` 的参数列表，不含代码自己管的执行策略参数（`target`、`kernel_map`、`tune`）。可选输入排在必选输入之后。按参考 API 支持的来声明，不按当前 kernel 支持的来声明。 -3. **写约束。** `shape_rules` 写关于类型变量取值的谓词，如 `H % G == 0`；派生的量写成 `let`；由开关选定的形状写成类型族。规则不读张量（`x.shape`、`x is None`），是否传入写成 `present(x)`。 -4. **写 `workloads`。** 每一行恰好给出没有生成器能确定的类型变量、每个没有默认值的构造参数、`some`（这次传入的可选张量）、`dtype_cases`（条目有 `DType` 类型变量时才写；dtype 参数按参数写）与 `label`。`implemented` 条目的每个可选张量，至少一行传、至少一行不传。label 是 case id 的一部分，而 case id 是 nightly 历史数据的键，改 label 会让历史断开。 -5. **写 `roofline`。** 用同一组类型变量写内联的 `flops`（访存不是「每个张量读或写一次」时再写 `bytes`），或者写一个 `func`，由它从检查过的调用算出两者。 - -接口先于实现落地时写 `status: spec-only`，这时需要读代码的检查都跳过；改成 `implemented` 后这些检查全部打开。 - -## 四种常见写法 - -### 由开关选定形状 - -`GemmFwdOp` 的两个布局开关决定 `a` 的哪个轴是 M。`types` 里的类型族写出两种情况，两个输入各自套用： - -```yaml - signature: - types: - Mat: - params: {t: Bool, R: Dim, C: Dim} - match: t - cases: - - {when: false, is: "[R, C]"} - - {when: true, is: "[C, R]"} - forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"} - params: - trans_a: {type: bool, default: false} - trans_b: {type: bool, default: true} - inputs: - a: {dtype: T, shape: "Mat[trans_a, M, K]"} - b: {dtype: T, shape: "Mat[trans_b, K, N]"} - outputs: - d: {dtype: T, shape: "[M, N]"} -``` - -### 可选输入作开关 - -`GroupNormFwdOp` 的 affine 由两个可选输入表达，旁边不再有 `affine` 开关：一件事只在一处声明。workload 行通过 `some` 传入它们，roofline 用 `present` 读是否传入： - -```yaml - signature: - forall: {B: Dim, C: Dim, L: Shape, T: "DType[float32 | float16 | bfloat16]"} - params: - num_groups: {type: int} - eps: {type: float, default: 1.0e-05} - inputs: - x: {dtype: T, shape: "[B, C, *L]"} - weight: {dtype: T, shape: "[C]", optional: true} - bias: {dtype: T, shape: "[C]", optional: true} - outputs: - output: {dtype: T, shape: "[B, C, *L]"} - shape_rules: - - "num_groups > 0 and C % num_groups == 0" - - "B * (C // num_groups) * prod(L) != 1" - workloads: - - {B: 8, C: 128, L: [32, 32], num_groups: 32, dtype_cases: [{T: float16}], label: image} - - {B: 8, C: 128, L: [32, 32], num_groups: 32, some: [weight, bias], - dtype_cases: [{T: float16}], label: image-affine} - roofline: - flops: "(5 + (1 if present(weight) else 0) + (1 if present(bias) else 0)) * B * C * prod(L)" -``` - -（节选：条目里还有更多行和 dtype。）算子可以按可选输入是否传入来分派，不能读张量内容来决定。 - -### 被写入的输入 - -`SSDDecodeFwdOp` 的 `state` 每个 decode 步都被原地写回，所以它标 `mutated: true`，仍然是输入；调用只返回 `y_out`。生成的 operator 声明写入的参数，恰好是标了 `mutated` 的那些输入。 - -```yaml - inputs: - A: {dtype: float32, shape: "[H, P, N]"} - dt: {dtype: float32, shape: "[B, H, P]"} - x: {dtype: T, shape: "[B, H, P]"} - B_in: {dtype: T, shape: "[B, G, N]"} - C_in: {dtype: T, shape: "[B, G, N]"} - state: {dtype: float32, shape: "[B, H, P, N]", mutated: true, contiguous: true} - outputs: - y_out: {dtype: float32, shape: "[B, H, P]"} -``` - -调用方可以自备的输出改标 `buffer: out`：`forward` 因此多一个 `out` 参数，算子写入它并返回它。 - -### 元数据张量的内容约束 - -偏移或长度这类张量，取值来自 `values` 里的生成器，内容须满足的条件写在 `requires` 里。下例取自 `GroupedQueryAttentionVarlenFwdOp`，其中 `q_lens` 是 `Seq[Int]` 类型变量，`T_q` 是 query 的总长度： - -```yaml -cu_seqlens_q: {dtype: int32, shape: "[B + 1]", values: "prefix_sum(q_lens)", - requires: ["prefix_offsets(T_q)"]} -``` - -`B` 由生成出的张量解出，所以一行给的是 `q_lens`，不是 `B`。 - -## 规则速查 - -**签名** - -- **顺序即位置。** `params` 的顺序是 `__init__` 的参数顺序，`inputs` 是 `forward` 的参数顺序，`outputs` 是返回值顺序；调换顺序是不兼容的改动。 -- **照参考实现写。** dtype 与参数依照权威的参考实现，不照当前代码；代码与之不符时改代码，改好之前条目标 `spec-only`。 -- **同名即相等。** 形状相同的张量写同一个形状；不是「名字相等」的关系写成约束或 `let`。 - -**约束与是否传入** - -- **约束读类型变量，不读张量。** 不写 `x.shape`、`x is None`、`isinstance`；是否传入写成 `present(x)`。 -- **是否传入可以当开关，内容不可以。** 算子可以按参数与张量是否传入选择实现，张量内容只作计算输入。 -- **可选输入两边都要测。** 传与不传各需一行，按输入逐个计，不按组合计。 - -**输出** - -- **输出数量固定。** 一个条目在每次调用上输出都相同；返回值随开关变化的算子拆成两个条目。 -- **被写入的输入仍是输入。** 它标 `mutated: true`，不列进 `outputs`。 - -## spec 校验器 {#spec-validator} - -校验由 [`scripts/validate_manifest.py`](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py) 执行，写完一份 spec 就可以立刻跑： - -```bash -python scripts/validate_manifest.py # every entry -python scripts/validate_manifest.py --check-op RMSNormFwdOp # one entry -python scripts/validate_manifest.py --levels schema,signature # skip the benchmark scan -``` - -| 级别 | 检查什么 | -| --- | --- | -| `schema` | 顶层字段、`family`、`ref_api`、`composition`、`roofline`，以及 `types.yaml` | -| `signature` | 签名在其判别量每种组合上的检查、每个 workload 行能否实例化、副作用声明；`implemented` 条目还要对照类的 `__init__`、`forward` 以及声明的 kernel 与子算子 | -| `bench` | 每个 benchmark 从 manifest 取调用、从算子取 roofline | - -`spec-only` 条目跳过需要读代码的检查。kernel 选择、多 kernel 的执行顺序、累加 dtype、workspace、tile 尺寸与 autotune 配置不在 manifest 里，校验器也就看不到。 diff --git a/docs/new-op.md b/docs/new-op.md index dc24f762..cc384919 100644 --- a/docs/new-op.md +++ b/docs/new-op.md @@ -1,11 +1,11 @@ # Adding a new op -A new op means writing code in the six places below, and the table is in the order to work -through them. +Adding an op means writing code in the six places below. The order of the table is +also the recommended order of writing. -The spec goes first: it decides what the other five files contain, and in the end it is what -they are checked against. **The spec is this pipeline's input, and the other five are written -from it.**{ .keystone } +The spec is written first, because it decides what the other five files contain, and +in the end they are checked against it. **The spec is the input to this workflow, and +the other five places are written from it.**{ .keystone } | # | File | Held to the spec by | Contents | | --- | --- | --- | --- | @@ -13,18 +13,20 @@ from it.**{ .keystone } | 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/…` | the validator, against `__init__`, `forward` and the declared kernels; the checks generated around every call | the op class, subclassing `Op` | | 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/__init__.py` and [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | the validator: the family's `__all__` agrees with the manifest | the op's name, exported by its family and on the public path `tileops..` | | 3 | [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)`/…` | — | the kernel classes, subclassing `Kernel` | -| 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_.py` | the contract tests, which run every workload row | the comparison against the reference `ref_program` | +| 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_.py` | the contract tests, which run every workload row | the numerical comparison against the reference `ref_program` | | 5 | [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops)`bench_.py` | the validator's `bench` level | the benchmark | -`GemmFwdOp` — the plainest matmul there is — runs through all six below. +The steps below take `GemmFwdOp`, the simplest matmul, through all six places. ## Step 1: write the spec -What the fields mean and how to write them is in [writing a spec](manifest.md). A new op -starts at `status: spec-only`: the interface is settled and there is no implementation -yet, so the checks that read code are skipped and do not fail over the missing class. +What each field means and how to write it is in +[Reading and writing the manifest](user-guide/manifest/index.md). A new op starts at +`status: spec-only`: the interface is settled and the implementation is not written +yet. The checks that read code are skipped at this status, so they do not fail on the +missing class. -`GemmFwdOp`'s spec, with one of its workload rows: +`GemmFwdOp`'s spec, with one workload row: ```yaml GemmFwdOp: @@ -55,19 +57,19 @@ GemmFwdOp: flops: "2 * M * N * K" ``` -The spec names no file and no kernel: which kernels serve the op is a fact of the code, -declared on the op class in step 2. +The spec names no file and no kernel. Which kernels serve the op belongs to the code, +and is declared on the op class in step 2. ## Step 2: write the op class {#op-class} The op class subclasses [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/op_base.py) and sits between the spec and the kernel. The checks -around every call — dtypes, shapes, the refinements, output shape inference — are -generated from the signature when the class is defined, so the class writes none of them. -What it writes is how a call reaches a kernel. +around every call (dtypes, shapes, the refinements and output shape inference) are +generated from the signature when the class is defined, and the class writes none of +them. What the class writes is how a call reaches a kernel. ### The class, and its members -`GemmFwdOp`'s skeleton, docstrings elided: +`GemmFwdOp`'s skeleton, with docstrings elided: ```python class GemmFwdOp(Op): @@ -108,49 +110,52 @@ class GemmFwdOp(Op): | --- | --- | --- | | 1 | `__init__` | the names, order and defaults in `signature.params`, then `target`, `kernel_map` and `tune`, closing with `self.dispatch_kernel(kernel_map)` | | 2 | `kernel_types` | the Kernel classes that can serve the op, each under a key; a `kernel_map=` override replaces one by that key | -| 3 | `interfaces` | one entry per kernel call the op makes: the name `kernel_for` uses → the `KernelInterface` class whose implementations serve that call | -| 4 | `forward` | `signature.inputs` — its order, optional inputs last with default `None` | -| 5 | `_eager_forward` | contiguity, the call spec, fetching the kernel and launching it | -| 6 | `compute_roof` | optional: the GPU-profile unit that prices the op's FLOPs, where it is not CUDA-core fp32 | +| 3 | `interfaces` | one entry per kernel call the op makes, mapping the name `kernel_for` uses to the `KernelInterface` class that the implementations serving that call inherit | +| 4 | `forward` | the order of `signature.inputs`, with optional inputs last and defaulting to `None` | +| 5 | `_eager_forward` | making the inputs contiguous, building the call spec, fetching the kernel and calling it | +| 6 | `compute_roof` | optional: the hardware unit whose peak prices the op's FLOPs; the default is fp32 on CUDA cores, and the member is written only when the op uses another unit | -`_infer_output_shapes`, `_validate_dtypes` and `eval_roofline` are generated from the spec -and are not written. +`_infer_output_shapes`, `_validate_dtypes` and `eval_roofline` are generated from the +spec and are not written by hand. -An op without a compile boundary writes the body of `_eager_forward` in `forward` itself; -declaring the boundary moves it behind the generated operator. How that works is in -[bringing an op into torch.compile](torch-compile.md). +An op without a compile boundary writes the body of `_eager_forward` in `forward` +itself. An op that declares the boundary runs that body behind the generated operator. +How that works is in [Bringing an op into torch.compile](torch-compile.md). ### `kernel_for`, and choosing among kernels {#kernel-selection} -A kernel is a compiled artefact, hundreds of milliseconds to seconds to build, while an op -instance is called over and over at different shapes and dtypes. The op layer therefore -keeps a memo table: a kernel this call needs and has built before comes straight back, -and only otherwise is one built and stored. `kernel_for` is that table's only entrance on -the in-tree path; a [target](backends.md) serves the whole op instead and never reaches it. +A kernel is a compiled artefact that takes hundreds of milliseconds to seconds to +build, while an op instance is called many times at different shapes and dtypes. The +op layer therefore keeps a cache table: a kernel this call needs that has been built +before is returned from the table, and otherwise the kernel is built and stored. +`kernel_for` is the only way in-tree implementations reach that table; a +[target](backends.md) serves the whole op and does not go through `kernel_for`. -Its two arguments: +`kernel_for` takes two arguments: -- **`interface`** — a key of `interfaces`, naming one kernel call the op makes. - `GemmFwdOp` makes one, so it declares one, `"gemm"`. A second interface is opened only where the semantics or the call contract changes: +- **`interface`**: a key of `interfaces`, naming one kernel call the op makes. + `GemmFwdOp` makes one kernel call, so it declares one interface, `"gemm"`. A second + kernel interface is added only where the semantics or the call contract changes: `BatchNormFwdOp` has `batch_norm_fwd_train` and `batch_norm_fwd_infer`, which return - different things. A faster kernel for some shape range or some architecture is another - implementation of the interface already there. -- **`call`** — a frozen `CallSpec` subclass carrying the facts needed to select and build - the kernel: shapes, the dtype, the op's semantic parameters, and the device. It - has to be the interface's `request` type. The dispatcher derives the device facts (`arch`, - `sm_count`, `calibration`, `smem_budget`) itself, from - `call.device` on a miss. - -The kernel that comes back is called with the parameters of the interface's abstract -`forward`, in that order. - -An interface is a class in + different things. A kernel that is faster on some shape range or some architecture + is another implementation of the existing kernel interface. +- **`call`**: a frozen `CallSpec` subclass carrying the facts needed to select an + implementation and build the kernel: shapes, the dtype, the op's semantic parameters, + and the device. It must be the type the kernel interface names in `request`. The + dispatcher derives the device facts (`arch`, `sm_count`, `calibration`, + `smem_budget`) from `call.device` on a cache miss; the caller does not fill them in. + +The op calls the returned kernel with the parameters of the kernel interface's +abstract `forward`, in that order. + +A kernel interface is a class in [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), -beside the call spec it names in `request`; a family with one kernel file keeps both in -that file. Its name is `{Name}{Fwd|Bwd}Interface`, with variant words before the -direction, which `interface-names-lint` checks. Its abstract `forward` is the whole contract -an implementation — in-tree or from a backend — is written against, so its docstring -states each tensor's shape, dtype, layout, device and whether it is written in place: +next to the call spec it names in `request`; a family with one kernel file keeps both +in that file. Its name has the form `{Name}{Fwd|Bwd}Interface`, with variant words +before the direction, and `interface-names-lint` checks it. Its abstract `forward` is +the only contract every implementation, in-tree or from a backend, is written against, +so its docstring states each tensor's shape, dtype, memory layout, device, and whether +it is written in place: ```python class GemmFwdInterface(KernelInterface): @@ -174,87 +179,45 @@ class GemmFwdInterface(KernelInterface): """ ``` -An implementation is a class inheriting both `Kernel` and one interface, listed in -`kernel_types` under a key. Which implementation serves a call follows from four -declarations the implementations make about themselves; the op makes none of them: +An implementation is a class that inherits both `Kernel` and one kernel interface, and +is listed in `kernel_types` under a key. Which implementation serves a call follows +from the availability, applicability and precedence each implementation declares about +itself; the op takes no part in the selection. A kernel interface with one +implementation needs no declaration beyond the inheritance. The selection rules, how to +write each declaration, and the common errors are in +[How an op selects a kernel](user-guide/dispatch/index.md) and +[Adding a kernel to an op](user-guide/dispatch/writing.md). + +An op does not write `entry_for` and keeps no kernel cache of its own. An op with no +in-tree implementation, which depends only on an external backend, declares neither +`kernel_types` nor `interfaces`; when no target claims the call's device, the call +raises `OpNotAvailableError`. How a backend adds an implementation or replaces the one +under a key is in [How a backend joins TileOPs](user-guide/dispatch/backends.md). -| # | Declaration | States | Left undeclared | -| --- | --- | --- | --- | -| 1 | `devices`, `supported_archs` | where the implementation runs | CUDA devices, every architecture | -| 2 | `applies(call)`, `refusal(call)` | which calls it serves, stated positively | every call | -| 3 | `general`, `preferred_over` | which implementation wins where two of them serve one call | wins over none | -| 4 | `entry_for(call)` | the build identity, and the builder that runs once per identity | the whole call spec, built by `cls(call)` | - -The dispatcher filters by availability first, then picks the single winner among the -implementations that are left and that apply: `general` loses to every other one, and the -rest are compared by the keys each names in `preferred_over`. Nothing left raises -`no implementation serves this call`, or `OpNotAvailableError` where no key runs on the -call's device type at all; two with no relation between them raise `dispatch is -ambiguous`. Declaration order decides nothing. Where one implementation should give up a -range to another, the one that should win declares `preferred_over`, rather than the -other one excluding that range in its own `applies`. - -`GemmFwdOp`'s three implementations cover the `"gemm"` interface's calls like this: - -| # | Key | Serves | Declares | -| --- | --- | --- | --- | -| 1 | `gemm_tma` | SM90 shapes whose operands TMA can address | `supported_archs = [90]`, and a `refusal` naming the misalignment | -| 2 | `gemv` | shapes of at most two rows contracted over K, where reducing on CUDA cores wins | `supported_archs = [90]`, `applies` through `band_for`, `preferred_over = frozenset({"gemm_tma"})` | -| 3 | `gemm_cp_async` | every shape the other two do not claim, provided a K row spans at least one four-byte load | `supported_archs = [80, 86, 89, 90]`, `general = True`, and a `refusal` for a narrower K row | - -`entry_for(call)` returns the **identity** two calls must share to reuse one kernel, and -the **builder** that runs once per identity. Carry too little in the identity and a second -dtype reuses the first dtype's kernel; carry the whole shape where the kernel depends on -fewer quantities and it compiles once per distinct shape. - -An interface with one implementation needs nothing beyond inheriting it. `RMSNormKernel` -is the only implementation `RMSNormFwdOp` has -([`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)): - -```python -class RMSNormKernel(Kernel, RMSNormFwdInterface): - supported_archs: list[int] = [80, 86, 89, 90] - - @classmethod - def entry_for(cls, call: LayerNormCall) -> Entry: - identity = (call.n, call.eps, call.dtype) - return identity, lambda: cls(*identity) -``` - -An op defines no `entry_for` of its own and keeps no kernel cache of its own — no dict, no -build guarded on an attribute being unset. Holding what `kernel_for` returned in -`self.kernel` is not one. +### Registering -An op with no in-tree implementation, written to depend on a backend, declares neither -`kernel_types` nor `interfaces`; a call on a device no target claims then raises -`OpNotAvailableError`. +Add the op's name to the imports and `__all__` in two places: -A backend adds an implementation to an interface, or replaces the class registered under -one key, -without changing TileOPs; both are in [adding a hardware backend](backends.md). +1. its family's [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops), where the class is implemented; +1. [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops), the public path. -### Registering - -Add the op's name to the imports and `__all__` in two places: its family's -[`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops), -where the class is implemented, and -[`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops), the public path. -Without the second, `from tileops. import ...` will not find the op and the API -reference cannot collect it. +Without the second, `from tileops. import ...` does not find the op, and the +API reference does not include it. ## Step 3: write the kernel -A kernel class subclasses [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) and the interface it implements, lives under [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), is written in -TileLang, compiles at construction and implements `forward`, which the base class's -`__call__` runs. Its constructor is what its own `entry_for` builder calls, and its -`forward` takes the interface's parameters, the `kernel(a, b)` of step 2. +A kernel class subclasses [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) and the kernel interface it implements, lives under [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels), is written in +TileLang, and compiles at construction. It implements `forward`, which the base +class's `__call__` runs. Its constructor is called by the builder its own `entry_for` +returns, and its `forward` takes the kernel interface's parameters, the `kernel(a, b)` +of step 2. -This is the one place of the six the spec does not constrain: a kernel neither reads the -spec nor is checked against it. +The kernel is the only one of the six places the spec does not constrain: a kernel +neither reads the spec nor is checked against it. -How the constructor and the call divide their arguments is a hard requirement: **only -values compiled into the generated code go in the constructor.** `GemmTmaKernel` divides -them like this: +The split of arguments between the constructor and the call is a hard requirement: +**only values compiled into the generated code go in the constructor.** `GemmTmaKernel` +splits them like this: ```python class GemmTmaKernel(Kernel, GemmFwdInterface): @@ -266,13 +229,13 @@ class GemmTmaKernel(Kernel, GemmFwdInterface): ... ``` -`m`, `n`, `k`, the dtype and the two layout flags are constructor arguments because the -generated code treats them as constants: loop bounds, TMA descriptors and the WGMMA shape -all unroll from them, as do the tile sizes. The tensors belong to `forward`, where each -call swaps pointers. +`m`, `n`, `k`, the dtype and the two layout flags are constructor arguments because +the generated code treats them as constants: loop bounds, TMA descriptors and the WGMMA +shape are all unrolled from them, and so are the tile sizes. The tensors are passed to +`forward`, and each call only changes the pointers. -Dividing them wrong costs a recompile. A decode step advances one token at a time, so -`seq_len` grows by one every step and batch changes with the running set: +A wrong split causes recompilation. Decode advances one token per step, so `seq_len` +grows by one every step, and the batch changes with the running set: ```python # wrong: seq_len in the constructor — every step is a new kernel @@ -283,31 +246,36 @@ kernel = AttnKernel(num_heads, head_dim, dtype) out = kernel(q, k, v) # seq_len is read off the tensor shapes ``` -With the first form, `seq_len` ends up in the identity `entry_for` returns, every step -misses, every step compiles, and decode goes nowhere. +With the first form, the build identity `entry_for` returns contains `seq_len`, so +every step misses the cache and compiles once, and decode cannot run at a usable +speed. ## Step 4: write the test -Tests live in [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops) and compare against `ref_program`, the reference the -workload (or the test class) defines, over shapes the test chooses to reach the kernel's -branches. Cases marked `smoke` run on every PR. Cases marked `full` run on any PR that -changes their test file, and on the nightly. Long-running cases are marked `nightly` and -run only on the nightly. The workload rows are not unit-test coverage; the -contract tests already run each of them through the op. +Tests live in [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops) and compare against `ref_program`, the reference +implementation the workload (or the test class) defines. Each test chooses its own +shapes to cover the kernel's branches. Cases fall into three groups by when they run: + +- cases marked `smoke` run on every PR; +- cases marked `full` run on any PR that changes their test file, and in the nightly; +- long-running cases are marked `nightly` and run only in the nightly. + +The workload rows are not part of unit-test coverage, because the contract tests +already run each row through the op. -The scaffolding is `TestBase` and `FixtureBase` from +The test scaffolding is `TestBase` and `FixtureBase` from [`tests/test_base.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_base.py), with the cases in `PARAMS`. -Where the op has an optional input, both sides need a case — passed and not passed often -run different kernels. +When the op has an optional input, it needs at least one case with the input passed +and one without, because the two often run different kernels. ## Step 5: write the benchmark -Benchmarks live in [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops) and time each call through a `ManifestBenchmark` -built around the op and the call's workload. The calls are not written here: -`manifest_calls()` instantiates each workload row with each of its dtype cases and ids -the case by its case id, and the validator's `bench` level fails a benchmark that writes -its own: +Benchmarks live in [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops). Each call is timed by a `ManifestBenchmark` +built from the op and the call's workload. The calls are not written by hand: +`manifest_calls()` instantiates each workload row with each of its dtype cases and +names each case by its case id. A benchmark that writes its own calls fails the +validator's `bench` level: ```python from benchmarks.benchmark_base import ManifestBenchmark, manifest_calls @@ -324,13 +292,14 @@ def test_gemm_bench(call) -> None: bm.compare({"tileops": op, "torch-cublas": workload.ref_program}, a, b) ``` -Record at least one non-TileOPs baseline as well, or the row has nothing to compare -against. Where a baseline needs its input converted, that conversion stays inside its own -timed region. What the reported numbers mean is in [how a benchmark is timed](timing.md). +A benchmark also records at least one non-TileOPs baseline; without one, the row has +nothing to compare against. Where a baseline needs its input converted, the conversion +stays inside that baseline's own timed region. What the reported numbers mean is in +[How a benchmark is timed](timing.md). ## Step 6: flip the status, and let CI take over -With the other five written, check your own work with the three commands below: +With the other five places written, run the three commands below to check the work: ```bash python scripts/validate_manifest.py --check-op GemmFwdOp # spec and code agree @@ -338,15 +307,16 @@ python -m pytest tests/ops/test_gemm.py -v # numerics match ref_p python -m pytest benchmarks/ops/bench_gemm.py # the benchmark produces numbers ``` -With all three passing, flip the spec's `status` from `spec-only` to `implemented`. That -one edit turns on the checks that read code and puts the op inside CI's reach: every later -change is held against the spec by the validator, the tests and the nightly benchmark. +When all three pass, change the spec's `status` from `spec-only` to `implemented`. +That one edit turns on the checks that read code and brings the op under CI: from then +on, the validator, the tests and the nightly benchmark check every change against the +spec. ## Afterwards -Once the op runs, two optional things remain: +Once the op runs, two optional tasks remain: -- Let the op into a user's compiled graph — [bringing an op into - torch.compile](torch-compile.md). -- Let someone else's kernels serve it on other hardware — [adding a hardware - backend](backends.md). +- letting the op into a user's compiled graph, see + [Bringing an op into torch.compile](torch-compile.md); +- letting other kernels serve the op on other hardware, see + [Adding a hardware backend](backends.md). diff --git a/docs/new-op.zh.md b/docs/new-op.zh.md index f838054f..97ce3702 100644 --- a/docs/new-op.zh.md +++ b/docs/new-op.zh.md @@ -1,25 +1,25 @@ -# 添加一个新算子 +# 添加一个新 op -写一个新算子需要在以下六个位置添加实现，表格的顺序也是推荐的动手顺序。 +新增一个 op 需要在以下六个位置添加代码，表格的顺序也是推荐的编写顺序。 -其中 spec 要第一个写：后面五个文件的内容都由它决定，最后也都由它校验。**spec 是这条流程的输入，其余五处都是照它写出来的。**{ .keystone } +spec 最先编写，因为后面五个文件的内容都由 spec 决定，最后也都由 spec 校验。**spec 是这条流程的输入，其余五处都依照它编写。**{ .keystone } | # | 文件 | 由谁对照 spec 检查 | 内容 | | --- | --- | --- | --- | | 1 | [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec)`.yaml` | 校验器的 `schema` 与 `signature` 两级 | spec 本身 | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/…` | 校验器对照 `__init__`、`forward` 与声明的 kernel；每次调用前后生成的检查 | 算子类，继承 `Op` | -| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/__init__.py` 与 [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | 校验器：家族的 `__all__` 与 manifest 一致 | 算子名，由所属家族导出，并出现在公开路径 `tileops..` 上 | +| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/…` | 校验器对照 `__init__`、`forward` 与声明的 kernel；每次调用前后生成的检查 | op 类，继承 `Op` | +| 2 | [`src/tileops/ops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/ops)`/__init__.py` 与 [`src/tileops/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops)`.py` | 校验器：family 的 `__all__` 与 manifest 一致 | op 名，由所属 family 导出，并出现在公开路径 `tileops..` 上 | | 3 | [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)`/…` | —— | kernel 类，继承 `Kernel` | -| 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_<名字>.py` | 契约测试，逐个跑每个 workload 行 | 与参考实现 `ref_program` 的数值比对 | +| 4 | [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)`test_<名字>.py` | 契约测试，逐个运行每个 workload 行 | 与参考实现 `ref_program` 的数值比对 | | 5 | [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops)`bench_<名字>.py` | 校验器的 `bench` 级 | benchmark | -下文以最简单的矩阵乘 `GemmFwdOp` 为例走一遍这六处。 +下文以最简单的矩阵乘 `GemmFwdOp` 为例，依次说明这六处。 ## 第一步：写 spec -spec 各字段的含义与写法见[读写 manifest](manifest.md)。新算子先写成 `status: spec-only`，表示接口已经定下来、实现还没有，这时需要读代码的检查都跳过，不会因为找不到类而报错。 +spec 各字段的含义与写法见[读写 manifest](user-guide/manifest/index.md)。新 op 的 spec 先写成 `status: spec-only`，表示接口已经确定、实现尚未完成。此时所有需要读取代码的检查都会跳过，不会因为找不到类而报错。 -`GemmFwdOp` 的 spec，workload 只列一行： +`GemmFwdOp` 的 spec 只列出一行 workload： ```yaml GemmFwdOp: @@ -50,15 +50,15 @@ GemmFwdOp: flops: "2 * M * N * K" ``` -spec 不写文件路径，也不写 kernel：由哪些 kernel 服务这个算子是代码里的事，在第二步的算子类上声明。 +spec 不写文件路径，也不写 kernel。由哪些 kernel 服务这个 op 属于代码的范围，在第二步的 op 类上声明。 -## 第二步：写算子类 {#op-class} +## 第二步：写 op 类 {#op-class} -算子类继承 [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/op_base.py)，是 spec 与 kernel 之间的一层。每次调用前后的检查 —— dtype、形状、约束、输出形状推导 —— 都在类定义时照签名生成，算子类一条也不写。它要写的是一次调用怎么走到 kernel。 +op 类继承 [`Op`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/op_base.py)，位于 spec 与 kernel 之间。每次调用前后的检查，包括 dtype、形状、约束与输出形状推导，都在类定义时依照签名生成，op 类不手写其中任何一项。op 类需要编写的是一次调用如何到达 kernel。 ### 类的骨架与成员 -`GemmFwdOp` 的骨架，略去 docstring： +`GemmFwdOp` 的骨架如下，略去 docstring： ```python class GemmFwdOp(Op): @@ -95,31 +95,31 @@ class GemmFwdOp(Op): return self.kernel_for("gemm", call)(a, b) ``` -| # | 成员 | 照什么写 | +| # | 成员 | 编写依据 | | --- | --- | --- | -| 1 | `__init__` | `signature.params` 的名字、顺序与默认值，再加 `target`、`kernel_map`、`tune`；结尾调用 `self.dispatch_kernel(kernel_map)` | -| 2 | `kernel_types` | 能服务这个算子的 Kernel 类，各起一个 key；`kernel_map=` 按这个 key 替换其中一个 | -| 3 | `interfaces` | 算子发出的每一个 kernel 调用各占一条：`kernel_for` 用的名字 → 服务这个调用的各实现所继承的 `KernelInterface` 类 | -| 4 | `forward` | `signature.inputs` 的顺序，可选输入排在最后、默认 `None` | -| 5 | `_eager_forward` | 把输入变成连续的，构造 call spec，取出 kernel，再调用它 | -| 6 | `compute_roof` | 可选。算子的 FLOPs 按哪个硬件单元的峰值定价，默认是 CUDA core 上的 fp32，用别的单元时才写 | +| 1 | `__init__` | `signature.params` 的名字、顺序与默认值，再加上 `target`、`kernel_map`、`tune`；末尾调用 `self.dispatch_kernel(kernel_map)` | +| 2 | `kernel_types` | 能服务这个 op 的 Kernel 类，各对应一个 key；`kernel_map=` 按 key 替换其中一个 | +| 3 | `interfaces` | op 发出的每一个 kernel 调用各占一条，从 `kernel_for` 使用的名字映射到服务这个调用的各实现所继承的 `KernelInterface` 类 | +| 4 | `forward` | `signature.inputs` 的顺序，可选输入排在最后，默认值为 `None` | +| 5 | `_eager_forward` | 把输入转为连续张量，构造 call spec，取得 kernel，再调用它 | +| 6 | `compute_roof` | 可选。表示 op 的 FLOPs 按哪个硬件单元的峰值计算，默认是 CUDA core 上的 fp32，只在使用其他单元时编写 | -`_infer_output_shapes`、`_validate_dtypes` 与 `eval_roofline` 都照 spec 生成，不用写。 +`_infer_output_shapes`、`_validate_dtypes` 与 `eval_roofline` 都依照 spec 生成，不需要编写。 -不声明编译边界的算子，把 `_eager_forward` 的内容直接写在 `forward` 里；声明了边界，这些内容挪到生成的 operator 后面。做法见[接入 torch.compile](torch-compile.md)。 +不声明编译边界的 op 把 `_eager_forward` 的内容直接写在 `forward` 中。声明了编译边界的 op 把这些内容放在生成的 operator 之后执行。具体做法见[接入 torch.compile](torch-compile.md)。 ### `kernel_for` 与 kernel 的选择 {#kernel-selection} -kernel 是编译产物，构造一次要几百毫秒到几秒，而一个算子实例会被反复调用，形状与 dtype 各不相同。算子层因此维护一张记忆表：本次调用要的 kernel 已经构造过就取回来，没有才构造并存进去。`kernel_for` 是自带实现走到这张表的唯一入口；[target](backends.md) 服务的是整个算子，不经过它。 +kernel 是编译产物，构造一次需要几百毫秒到几秒，而一个 op 实例会以不同的形状与 dtype 被反复调用。因此 op 层维护一张缓存表：本次调用所需的 kernel 已经构造过时直接取出，否则构造后存入表中。`kernel_for` 是 in-tree 实现访问这张表的唯一入口；[target](backends.md) 服务的是整个 op，不经过 `kernel_for`。 `kernel_for` 接受两个参数： -- **`interface`**：`interfaces` 的一个 key，指算子发出的某一个 kernel 调用。`GemmFwdOp` 只发出一个 kernel 调用，因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新开一个接口：`BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的东西不同。某个形状范围或某个架构上更快的 kernel 是已有接口的另一个实现。 -- **`call`**：一个冻结的 `CallSpec` 子类，带着选择实现和构造 kernel 所需的信息：形状、dtype、算子的语义参数与设备。它必须是这个接口 `request` 指定的类型。设备事实（`arch`、`sm_count`、`calibration`、`smem_budget`）由派发机制在未命中时从 `call.device` 推出，调用方不填。 +- **`interface`**：`interfaces` 的一个 key，指 op 发出的某一个 kernel 调用。`GemmFwdOp` 只发出一个 kernel 调用，因此只声明一个 `"gemm"`。只有语义或调用契约改变时才新增一个 kernel 接口，例如 `BatchNormFwdOp` 的 `batch_norm_fwd_train` 与 `batch_norm_fwd_infer` 返回的内容不同。在某个形状范围或某个架构上更快的 kernel 是已有 kernel 接口的另一个实现。 +- **`call`**：一个冻结的 `CallSpec` 子类，包含选择实现和构造 kernel 所需的信息，即形状、dtype、op 的语义参数与设备。它必须是这个 kernel 接口的 `request` 指定的类型。设备事实（`arch`、`sm_count`、`calibration`、`smem_budget`）由派发机制在缓存未命中时从 `call.device` 推出，调用方不填写。 -算子按接口的抽象 `forward` 方法所声明的参数、以同样的顺序调用取回的 kernel。 +op 按 kernel 接口的抽象 `forward` 方法所声明的参数及其顺序调用取得的 kernel。 -接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 里的一个类，与它在 `request` 中指名的 call spec 放在一起；只有一个 kernel 文件的 family 把两者都写在那个文件里。名字是 `{Name}{Fwd|Bwd}Interface`，变体词写在方向之前，由 `interface-names-lint` 检查。它的抽象 `forward` 是各实现（自带的与后端提供的）唯一依据的契约，docstring 因此写明每个张量的形状、dtype、内存布局、设备，以及是否被原地写入： +kernel 接口是写在 [`src/tileops/kernels//call_spec.py`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 中的一个类，与它在 `request` 中指定的 call spec 放在一起；只有一个 kernel 文件的 family 把两者都写在那个文件中。类名的格式是 `{Name}{Fwd|Bwd}Interface`，表示变体的词写在方向之前，由 `interface-names-lint` 检查。kernel 接口的抽象 `forward` 是所有实现（in-tree 实现与 backend 提供的实现）共同依据的唯一契约，因此它的 docstring 写明每个张量的形状、dtype、内存布局、设备，以及是否被原地写入： ```python class GemmFwdInterface(KernelInterface): @@ -143,56 +143,26 @@ class GemmFwdInterface(KernelInterface): """ ``` -实现是同时继承 `Kernel` 与某一个接口的类，以一个 key 列在 `kernel_types` 中。一次调用由哪个实现服务，取决于各实现自己的四项声明，算子不参与： +实现是同时继承 `Kernel` 与某一个 kernel 接口的类，以一个 key 列在 `kernel_types` 中。一次调用由哪个实现服务，取决于各实现自己声明的可用性、适用范围与优先关系，op 不参与选择。只有一个实现的 kernel 接口，除继承 kernel 接口外不需要其他声明。选择规则、各项声明的写法与常见报错见 [op 如何选择 kernel](user-guide/dispatch/index.md) 与[如何为 op 新增 kernel](user-guide/dispatch/writing.md)。 -| # | 声明 | 说明什么 | 未声明时的默认 | -| --- | --- | --- | --- | -| 1 | `devices`、`supported_archs` | 实现能在哪些设备上运行 | CUDA 设备，全部架构 | -| 2 | `applies(call)`、`refusal(call)` | 实现服务哪些调用，正面写出 | 服务全部调用 | -| 3 | `general`、`preferred_over` | 两个实现都服务同一次调用时谁胜出 | 不胜过任何实现 | -| 4 | `entry_for(call)` | build identity，以及每个 identity 只运行一次的 builder | 以整个 call spec 为 identity，用 `cls(call)` 构造 | - -派发机制先按可用性过滤，再在剩下的、且适用的实现中选出唯一的胜者：`general` 的实现低于其他所有实现，其余按各自 `preferred_over` 列出的 key 比较。一个实现都不剩时报 `no implementation serves this call`；没有任何 key 能在这次调用的设备类型上运行时报 `OpNotAvailableError`；剩下的实现之间互相没有优先关系时报 `dispatch is ambiguous`。实现的声明顺序不影响选择结果。需要让出一段范围时，由应当胜出的一方声明 `preferred_over`，而不是让另一方在自己的 `applies` 里把这段范围排除掉。 - -`GemmFwdOp` 的三个实现这样分担 `"gemm"` 接口的调用： - -| # | key | 服务 | 声明 | -| --- | --- | --- | --- | -| 1 | `gemm_tma` | 操作数能被 TMA 寻址的 SM90 形状 | `supported_archs = [90]`，以及给出未对齐原因的 `refusal` | -| 2 | `gemv` | 沿 K 规约且最多两行的形状，这种形状在 CUDA core 上规约更快 | `supported_archs = [90]`、经 `band_for` 实现的 `applies`、`preferred_over = frozenset({"gemm_tma"})` | -| 3 | `gemm_cp_async` | 其余两个都不认领的全部形状，前提是一行 K 至少占满一次 4 字节读取 | `supported_archs = [80, 86, 89, 90]`、`general = True`，以及拒绝更窄 K 行的 `refusal` | - -`entry_for(call)` 返回两样东西：两次调用共享它才算同一个 kernel 的 **build identity**，以及每个 identity 只运行一次的 **builder**。identity 少带一个量，第二种 dtype 就会复用第一种 dtype 的 kernel；kernel 只依赖其中几个量却把整个形状都带上，就变成每个形状各编译一次。 - -只有一个实现的接口，除继承接口外不需要别的声明。`RMSNormKernel` 是 `RMSNormFwdOp` 唯一的实现（[`src/tileops/kernels/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/norm/rms_norm.py)）： - -```python -class RMSNormKernel(Kernel, RMSNormFwdInterface): - supported_archs: list[int] = [80, 86, 89, 90] - - @classmethod - def entry_for(cls, call: LayerNormCall) -> Entry: - identity = (call.n, call.eps, call.dtype) - return identity, lambda: cls(*identity) -``` - -算子自己不写 `entry_for`，也不自建 kernel 缓存：没有缓存字典，也不以某个属性是否已赋值来决定要不要构造。把 `kernel_for` 的返回值存进 `self.kernel` 不算自建缓存。 +op 不编写 `entry_for`，也不自行维护 kernel 缓存。完全没有 in-tree 实现、只依赖外部 backend 的 op 不写 `kernel_types` 与 `interfaces`；没有 target 认领调用设备时，调用抛出 `OpNotAvailableError`。backend 新增实现或替换某个 key 的方式见 [backend 如何接入](user-guide/dispatch/backends.md)。 -完全没有自带实现、只依赖外部后端的算子，`kernel_types` 与 `interfaces` 都不写；在没有 target 认领设备时，调用会抛 `OpNotAvailableError`。 +### 注册 -后端可以为一个接口新增实现，也可以替换某个 key 登记的类，这两件事都不改动 TileOPs，见[接入一类新硬件](backends.md)。 +op 名需要加入两处的导入与 `__all__`： -### 注册 +1. 所属 family 的 [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops)，即类的实现位置； +1. [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops)，即公开路径。 -把算子名加进两处的导入与 `__all__`：算子所属家族的 [`src/tileops/ops//__init__.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops)（类的实现位置），以及 [`src/tileops/.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops)（公开路径）。缺了后者，`from tileops. import ...` 拿不到这个算子，API 参考也收不到它。 +缺少后者时，`from tileops. import ...` 无法导入这个 op，API 参考也不会收录它。 ## 第三步：写 kernel -kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的那个接口，放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下，用 TileLang 编写，在构造时编译。它实现 `forward`，基类的 `__call__` 会调用它。构造函数由本类 `entry_for` 返回的 builder 调用，`forward` 接受接口规定的参数，就是第二步里的 `kernel(a, b)`。 +kernel 类继承 [`Kernel`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/kernels/kernel_base.py) 与它实现的 kernel 接口，放在 [`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels) 下，用 TileLang 编写，在构造时编译。kernel 类实现 `forward`，由基类的 `__call__` 调用。构造函数由本类 `entry_for` 返回的 builder 调用；`forward` 接受 kernel 接口规定的参数，也就是第二步中的 `kernel(a, b)`。 -它是这六处里唯一不受 spec 约束的一处：kernel 不读 spec，也不对照 spec 检查。 +kernel 是这六处中唯一不受 spec 约束的一处：kernel 不读取 spec，也不对照 spec 检查。 -构造参数与调用参数的划分有一条硬性要求：**只有会被编译进生成代码的值才进构造函数。** `GemmTmaKernel` 是这样分的： +构造参数与调用参数的划分有一条硬性要求：**只有会被编译进生成代码的值才作为构造参数。** `GemmTmaKernel` 的划分如下： ```python class GemmTmaKernel(Kernel, GemmFwdInterface): @@ -204,9 +174,9 @@ class GemmTmaKernel(Kernel, GemmFwdInterface): ... ``` -`m`、`n`、`k`、dtype 与两个布局标志进了构造函数，因为生成的代码里这些值是常量：循环边界、TMA 描述符、WGMMA 的形状都按它们展开，tile 尺寸同理。张量本身留给 `forward`，每次调用只换指针。 +`m`、`n`、`k`、dtype 与两个布局标志是构造参数，因为它们在生成的代码中是常量：循环边界、TMA 描述符与 WGMMA 的形状都按这些值展开，tile 尺寸也是如此。张量本身由 `forward` 接收，每次调用只更换指针。 -分错的代价是重新编译。decode 一步一步往前走，`seq_len` 每步 +1，batch 随 running set 变化： +划分错误的后果是重新编译。decode 逐步推进时，`seq_len` 每一步加 1，batch 随 running set 变化： ```python # 错：seq_len 进了构造函数 —— 每一步都是一个新 kernel @@ -217,19 +187,25 @@ kernel = AttnKernel(num_heads, head_dim, dtype) out = kernel(q, k, v) # seq_len 从张量形状里读 ``` -上一种写法下，`entry_for` 返回的 build identity 里带着 `seq_len`，每步都未命中、每步都编译一次，decode 直接跑不动。 +在前一种写法下，`entry_for` 返回的 build identity 包含 `seq_len`，每一步都无法命中缓存，每一步都编译一次，decode 因此无法正常运行。 ## 第四步：写测试 -测试放在 [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops)，比对对象是 workload（或测试类）定义的参考实现 `ref_program`，形状由测试自己挑，以覆盖 kernel 的各个分支；`smoke` 用例每个 PR 都运行；`full` 用例在改动其测试文件的 PR 和 nightly 中运行；耗时长的用例标 `nightly`，只在 nightly 运行。workload 行不是单元测试的覆盖面，契约测试已经把每一行都交给算子跑过。 +测试放在 [`tests/ops/`](https://github.com/tile-ai/TileOPs/tree/main/tests/ops) 中，比对对象是 workload（或测试类）定义的参考实现 `ref_program`。测试自行选择形状，以覆盖 kernel 的各个分支。用例按运行时机分为三类： + +- `smoke` 用例在每个 PR 中运行； +- `full` 用例在修改了其测试文件的 PR 和 nightly 中运行； +- 耗时长的用例标记为 `nightly`，只在 nightly 中运行。 + +workload 行不属于单元测试的覆盖范围，因为契约测试已经用 op 运行过每一行。 -骨架用 [`tests/test_base.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_base.py) 里的 `TestBase` 与 `FixtureBase`，用例写在 `PARAMS` 里。 +测试骨架使用 [`tests/test_base.py`](https://github.com/tile-ai/TileOPs/blob/main/tests/test_base.py) 中的 `TestBase` 与 `FixtureBase`，用例写在 `PARAMS` 中。 -如果这个算子有可选输入，传与不传各至少要有一条用例 —— 两侧走的往往是不同的 kernel。 +op 有可选输入时，传入与不传入各至少需要一条用例，因为两种情况通常走不同的 kernel。 ## 第五步：写 benchmark -benchmark 放在 [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops)，每个调用交给一个围绕算子与该调用 workload 构造的 `ManifestBenchmark` 计时。调用不自己写：`manifest_calls()` 把每个 workload 行配上它的每个 dtype case 各实例化一次，并以 case id 命名；自己写调用的 benchmark 过不了校验器的 `bench` 级： +benchmark 放在 [`benchmarks/ops/`](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/ops) 中。每个调用由一个基于 op 与该调用的 workload 构造的 `ManifestBenchmark` 计时。调用不手写：`manifest_calls()` 为每个 workload 行的每个 dtype case 各实例化一个调用，并以 case id 命名。手写调用的 benchmark 无法通过校验器的 `bench` 级： ```python from benchmarks.benchmark_base import ManifestBenchmark, manifest_calls @@ -246,11 +222,11 @@ def test_gemm_bench(call) -> None: bm.compare({"tileops": op, "torch-cublas": workload.ref_program}, a, b) ``` -另外至少要记一个非 TileOPs 的基线，否则这一行没有比较对象。基线若需要转换输入，转换的代码留在它自己的计时区间内，不要挪出去。报出来的数字各是什么意思，见[benchmark 怎么计时](timing.md)。 +benchmark 至少还需要记录一个非 TileOPs 的基线，否则这一行没有比较对象。基线需要转换输入时，转换代码保留在基线自己的计时区间内。报出的各个数字的含义见 [benchmark 的计时方法](timing.md)。 -## 第六步：反转实现状态，让算子进入 CI 校验 +## 第六步：反转实现状态，让 op 进入 CI 校验 -上面五处都写完之后，先跑下面三条命令自查一遍： +上述五处都写完之后，先运行以下三条命令自查： ```bash python scripts/validate_manifest.py --check-op GemmFwdOp # spec and code agree @@ -258,11 +234,11 @@ python -m pytest tests/ops/test_gemm.py -v # numerics match ref_p python -m pytest benchmarks/ops/bench_gemm.py # the benchmark produces numbers ``` -三样都过，再把 spec 的 `status` 从 `spec-only` 反转成 `implemented`。这一改动打开所有需要读代码的检查，算子由此进入 CI 的保护范围：往后每次改动，spec 校验器、测试与 nightly benchmark 都会对照 spec 检查一遍。 +三条命令都通过后，再把 spec 的 `status` 从 `spec-only` 改为 `implemented`。这一改动启用所有需要读取代码的检查，op 由此进入 CI 的保护范围：此后每次改动，spec 校验器、测试与 nightly benchmark 都会对照 spec 检查一遍。 -## 接下来 +## 后续步骤 -算子跑起来之后，还有两件可选的事： +op 能够运行之后，还有两项可选的工作： -- 让算子能进使用者的编译图 —— [接入 torch.compile](torch-compile.md)。 -- 让它在别的硬件上由别人的 kernel 服务 —— [接入新硬件后端](backends.md)。 +- 让 op 能够进入使用者的编译图，见[接入 torch.compile](torch-compile.md)。 +- 让 op 在其他硬件上由其他 kernel 服务，见[接入新硬件 backend](backends.md)。 diff --git a/docs/performance-guides/index.md b/docs/performance-guides/index.md index a79f52b6..419cef82 100644 --- a/docs/performance-guides/index.md +++ b/docs/performance-guides/index.md @@ -4,15 +4,15 @@ 1. [Benchmarks](../benchmarks/index.md) — a nightly run on an H200, reporting device time per op per workload against the fastest other implementation of - the same op. How the numbers are taken and how to read the ratio is set out - in "How these numbers are taken" in that section. + the same op. How the numbers are taken and how to read the ratio are + described in "How these numbers are taken" in that section. ## Tools for locating a problem 1. [In-Kernel Timeline Trace](trace-timeline.md) — annotate a kernel body with markers and read back a per-CTA timeline: gaps, stalls, and how far the - producer and consumer overlap. None of that is visible to a per-kernel - profiler. The API is [Trace](../api/trace.md). + producer and consumer overlap. A per-kernel profiler does not show this + information. The API is [Trace](../api/trace.md). ## Tuning practice for TileLang diff --git a/docs/performance-guides/index.zh.md b/docs/performance-guides/index.zh.md index 3761eff4..a8d21365 100644 --- a/docs/performance-guides/index.zh.md +++ b/docs/performance-guides/index.zh.md @@ -2,14 +2,14 @@ ## Nightly 性能数据 -1. [性能数据](../benchmarks/index.md) —— 每晚在 H200 上运行的性能测量，逐算子逐 workload 给出 device time，并与同一算子最快的其他实现对比。数字怎么取、比值怎么读，见该栏的「How these numbers are taken」。 +1. [性能数据](../benchmarks/index.md) —— 每晚在 H200 上运行的性能测量，逐个 op、逐个 workload 给出 device time，并与同一 op 最快的其他实现对比。数字的取法与比值的读法见该栏的「How these numbers are taken」。 ## 调优工具 -1. [核内时间线追踪](trace-timeline.md) —— 在 kernel 体内加标记，运行后得到逐 CTA 的时间线，用于定位空隙、停等以及生产者与消费者的重叠情况。这是 per-kernel profiler 看不到的部分。追踪的 API 见 [Trace](../api/trace.md)。 +1. [核内时间线追踪](trace-timeline.md) —— 在 kernel 体内加标记，运行后得到逐 CTA 的时间线，用于定位空隙、停等以及生产者与消费者的重叠情况。这部分信息是 per-kernel profiler 看不到的。追踪的 API 见 [Trace](../api/trace.md)。 ## TileLang 性能调优最佳实践 -1. [优化访存受限的 kernel](memory-bound/index.md) —— 访存受限在 roofline 上是什么位置，以及 access pattern 决定带宽的两处地方（global memory 与 shared memory）各自的实测结论。 +1. [优化访存受限的 kernel](memory-bound/index.md) —— 说明访存受限在 roofline 上的位置，并给出 access pattern 影响带宽的两处（global memory 与 shared memory）各自的实测结论。 -一次调用的计算量与访存量由 `op.eval_roofline()` 给出，取自 manifest 的 `roofline` 字段，是读测量结果时对照的上限；模型与字段规范见 [Roofline](../design/roofline.md)。 +一次调用的计算量与访存量由 `op.eval_roofline()` 给出，取自 manifest 的 `roofline` 字段，是阅读测量结果时对照的上限。模型与字段规范见 [Roofline](../design/roofline.md)。 diff --git a/docs/performance-guides/memory-bound/global-memory-access.md b/docs/performance-guides/memory-bound/global-memory-access.md index 332650b6..45631d6b 100644 --- a/docs/performance-guides/memory-bound/global-memory-access.md +++ b/docs/performance-guides/memory-bound/global-memory-access.md @@ -1,15 +1,15 @@ # Optimizing Global Memory Access -When a thread reads several elements from a row, the access can be written four -ways. This page measures all four on two workloads and explains how to choose -among them. +When a thread reads several elements from a row, the access can be written in +four ways. This page compares the four on two workloads and gives the basis for +choosing one. ## Checking whether DRAM bandwidth is the current limit {#regime} [Elementwise](https://tile-ai.github.io/TileOPs.github.io/api/elementwise/) and [Reduction](https://tile-ai.github.io/TileOPs.github.io/api/reduction/) are the -typical memory-bound kernels. Each recommendation below states when it applies, -why it applies, and what the wrong and right code look like. +typical memory-bound kernels. Each recommendation on this page states when it +applies, why it applies, and what the wrong and right code look like. Every measurement on this page uses the same conditions: **an input larger than the 60 MiB L2, and enough blocks to fill the whole card** (an H200 has 132 @@ -18,32 +18,33 @@ differences in access pattern show up directly in performance. !!! warning "Where this applies" - Outside these conditions, another factor may set the limit, and some of the - conclusions here can reverse. + Outside these conditions, the main limit may be another factor, and some of + the conclusions on this page can reverse. The table below uses those two conditions to divide the space into three -regimes. It gives the test for each regime and how the conclusions of these two -pages apply there. The named limit is the dominant factor; more complex kernels -usually have more factors active at the same time: +regimes. Each row gives the test for the regime and how the conclusions of the +global memory and shared memory pages apply there. The table lists the dominant +limit; the more complex the kernel, the more factors are active at the same +time. | Regime | Test | Main limit | How to use the conclusions | | --- | --- | --- | --- | | Bandwidth saturated | Input > 60 MiB, blocks at more than twice the SM count | DRAM bandwidth, that is sector utilization | Apply directly | | Small data | Input fits in L2, one call takes tens of microseconds or less | Fixed launch overhead, cache state | Avoid the wrong forms; changing access pattern buys nothing | -| Few blocks | Blocks fewer than twice the SM count | The width of each load instruction, bytes in flight | Keep load width first, and measure every change | +| Few blocks | Blocks fewer than twice the SM count | The width of each load instruction, bytes in flight | Keep load width first, and measure each change separately | - **With small data, launch overhead and cache state dominate.** On 65536 × 4096 (512 MB), the four access patterns of the same row-reduction kernel (fp16, 256 threads, clocks unlocked) measure 4.20 to 4.43 TB/s, within 6% of each other. On 2048 × 4096 (16 MB, which fits in L2), one call takes a dozen or so microseconds, and two measurements of the *same* access pattern can differ by - threefold. Changing the access pattern has no benefit in this regime, because - another factor is setting the pace. + threefold. In this regime another factor limits performance, and changing the + access pattern has no benefit. - **With few blocks, load width matters more than the coalescing rules predict.** - There are too few warps to hide memory latency behind concurrent requests, so - the remaining lever is to make each request wider and keep more bytes in - flight per thread. Any change that trades load width for something else can - reverse here. + With too few warps, memory latency cannot be hidden behind concurrent + requests. The only option is to make each request wider and keep more bytes + in flight per thread. In this regime, any change that gives up load width for + another gain can produce the opposite result. ## Coalescing global memory accesses {#coalescing} @@ -56,12 +57,13 @@ usually have more factors active at the same time: Lookups operate on cache lines, while transfers operate on sectors. On a sector miss, L1 requests only that sector from L2 instead of pulling the whole line. -Therefore **fetching 1 byte costs the same as fetching all 32**. The quality of -a memory instruction is measured by its **sector utilization**: +Therefore **fetching 1 byte costs the same as fetching all 32**, and the +efficiency of a memory instruction is measured by its **sector utilization**: `bytes actually used / (sectors touched × 32)`. The hardware coalesces a warp's 32 accesses into as few 32-byte transactions as -it can. Reaching the minimum requires three things at once: +it can. Reaching the minimum number of transactions requires three conditions +at once: 1. **Contiguous addresses** — the 32 threads of one instruction address a run with no holes in it; @@ -76,7 +78,7 @@ A thread reading $V$ elements ($V$ = elements per row / threads) has four access patterns available. **blocked** — each thread takes one contiguous run. For a fixed `c`, adjacent -threads are $V$ elements apart. This breaks the first requirement, and sector +threads are $V$ elements apart. This breaks the first condition, and sector utilization is $1/V$: ```python @@ -84,16 +86,16 @@ for c in T.serial(V): acc[0] = acc[0] * X[row, tx * V + c] ``` -**striped** — adjacent threads take adjacent elements. The addresses are now +**striped** — adjacent threads take adjacent elements. The addresses are contiguous, but each thread fetches only one element per instruction, which -breaks the third requirement. Reading $V$ elements takes $V$ instructions. +breaks the third condition. Reading $V$ elements takes $V$ instructions. ```python for c in T.serial(V): acc[0] = acc[0] * X[row, c * threads + tx] ``` -**blocked + vectorized** — still one contiguous run per thread, but +**blocked + vectorized** — each thread still takes one contiguous run, but `T.vectorized` reads a full 16 bytes at a time, satisfying all three: ```python @@ -105,8 +107,8 @@ for c in T.serial(V): acc[0] = acc[0] * buf[c] ``` -**staged** — `T.Parallel` performs the copy and consumption reads shared memory, -which also satisfies all three: +**staged** — `T.Parallel` performs the copy, and the consumption phase reads +from shared memory, which also satisfies all three: ```python sh = T.alloc_shared((threads, V + pad), dtype) @@ -125,21 +127,25 @@ With `T.serial`, the loop body runs sequentially on a single thread. The index expression is translated directly into memory instructions, with no coalescing or vectorization. The pattern in the source is the pattern the hardware sees. -`T.vectorized`, `T.Parallel`, and `T.copy` hand that decision to TileLang's -**layout inference**. They differ in how much the programmer still specifies: -`T.vectorized` specifies the access width per thread and infers the thread -mapping; `T.Parallel` specifies neither, so it decides both how loop dimensions -are split across threads and how wide each read is; `T.copy` specifies only a -source region and a destination region, then generates the whole copy -(`coalesced_width` and `loop_layout` are available when the inferred result -needs to be overridden). Layout inference handles vectorization, address -alignment, and avoiding bank conflicts on the shared-memory side. Those are the -hardware-friendly details that are easy to get wrong by hand. - -**Writing indices by hand with `T.serial` means guaranteeing those three -requirements directly; handing the copy to layout inference means specifying -only the copy extent.** The measurements below show which one to choose and what -bandwidth each reaches. +For `T.vectorized`, `T.Parallel`, and `T.copy`, TileLang's **layout inference** +decides the mapping. The three differ in what the programmer writes: + +- `T.vectorized` states the access width per thread, and layout inference + derives the thread mapping; +- `T.Parallel` does not state the width; layout inference decides both how the + loop dimensions are split across threads and how wide each read is; +- `T.copy` states only a source region and a destination region, and generates + the whole copy; `coalesced_width` and `loop_layout` are available when the + inferred result needs to be specified by hand. + +Layout inference handles the rest: vectorization, address alignment, and +avoiding bank conflicts on the shared memory side. Those are the parts that are +hardware-friendly but easy to get wrong by hand. + +**When indices are written by hand with `T.serial`, the programmer guarantees +the three conditions one by one; when the copy is left to layout inference, only +the copy extent is stated.** The measurements below show which form to choose +and the bandwidth each reaches.
@@ -485,9 +491,9 @@ bandwidth each reaches. ## Measurements -Both workloads are measured on an H200. The comparison is the **memory +This section measures two workloads on an H200 and compares the **memory bandwidth** reached by each of the four access patterns: bytes moved divided by -kernel time, in TB/s. The two workloads impose different requirements on the +kernel time, in TB/s. The two workloads place different requirements on the order in which elements are processed, and that requirement determines which access patterns are available. @@ -549,14 +555,16 @@ back. striped is unavailable here because a thread cannot hold a contiguous run. puts it above element-by-element blocked and below vectorized blocked: 3.31 versus 1.83 and 3.81 at $V = 16$. It fits cases where minimizing the code change matters more than extracting the last bit of bandwidth. Because each - thread holds non-contiguous elements, computations that require a contiguous - run per thread, such as a serial prefix, cannot use it. + thread holds non-contiguous elements, it cannot be used by computations that + require each thread to hold a contiguous run, such as a serial prefix + product. The two listings below are complete templates for the recommended access patterns. `M`, `N`, `V`, `threads`, `pad`, and `dtype` are all compile-time -constants. Among the four listings at the top of this page, element-by-element -blocked is the wrong form and should not be copied. striped works, but it is -not the fastest option (see point 4 above). +constants. Among the four listings in +[Coalescing global memory accesses](#coalescing), element-by-element blocked is +the wrong form. striped works, but it is not the fastest option (see point 4 in +Choosing an access pattern). **Recommended at small $V$** — vectorized blocked: @@ -578,8 +586,8 @@ def main(X: T.Tensor((M, N), dtype), Out: T.Tensor((M, threads), "float32")): Out[row, tx] = acc[0] ``` -**Recommended at large $V$** — padded staged (for the crossover, see point 2 -above): +**Recommended at large $V$** — padded staged (for the crossover, see point 2 in +Choosing an access pattern): ```python @T.prim_func diff --git a/docs/performance-guides/memory-bound/global-memory-access.zh.md b/docs/performance-guides/memory-bound/global-memory-access.zh.md index 2aba3372..f10cc19d 100644 --- a/docs/performance-guides/memory-bound/global-memory-access.zh.md +++ b/docs/performance-guides/memory-bound/global-memory-access.zh.md @@ -1,27 +1,27 @@ # 优化 global memory 访问 -一个线程要读一行里的多个元素时，写法有四种。这一页给出它们在两个 workload 上的实测对比，以及怎么挑一种。 +一个线程读取一行中的多个元素时，有四种写法。本页给出它们在两个 workload 上的实测对比，以及选择的依据。 ## 确认 DRAM 带宽是否为当前的瓶颈 {#regime} -[Elementwise](https://tile-ai.github.io/TileOPs.github.io/api/elementwise/) 与 [Reduction](https://tile-ai.github.io/TileOPs.github.io/api/reduction/) 是典型的访存受限 kernel。下面每条建议都写明触发的条件、成因，以及反例与正例代码。 +[Elementwise](https://tile-ai.github.io/TileOPs.github.io/api/elementwise/) 与 [Reduction](https://tile-ai.github.io/TileOPs.github.io/api/reduction/) 是典型的访存受限 kernel。本页每条建议都写明触发条件、成因，以及反例与正例代码。 -这一页的实测都在同一组条件下取得：**输入大于 L2 的 60 MiB，且 block 数足以填满整卡**（H200 有 132 个 SM）。此时 DRAM 带宽是主要瓶颈，访存模式的差别直接反映在性能上。 +本页的实测都在同一组条件下取得：**输入大于 L2 的 60 MiB，且 block 数足以填满整卡**（H200 有 132 个 SM）。此时 DRAM 带宽是主要瓶颈，访存模式的差别直接反映在性能上。 !!! warning "适用范围" - 这组条件之外，主要性能瓶颈可能由别的因素决定，这里给出的一些结论会反转。 + 在这组条件之外，主要瓶颈可能是其他因素，本页的部分结论会反转。 -下表按这两个条件划出三个区间，逐行给出判据，以及这两页的结论在各区间里怎么用。表里的瓶颈是主导因素，kernel 越复杂，同时起作用的因素越多： +下表按这两个条件划分三个区间，逐行给出判据，以及 global memory 与 shared memory 两页的结论在各区间的用法。表中列出的是主导瓶颈；kernel 越复杂，同时起作用的因素越多。 | 区间 | 判据 | 主要瓶颈 | 结论的用法 | | --- | --- | --- | --- | | 带宽饱和 | 输入 > 60 MiB，block 数在 SM 数的两倍以上 | DRAM 带宽，即 sector 利用率 | 直接适用 | | 数据小 | 输入装得进 L2，单次耗时在几十微秒以内 | kernel 发射的固定开销、缓存状态 | 避开反例即可，换 access pattern 没有收益 | -| block 少 | block 数不到 SM 数的两倍 | 每条载入指令的宽度、在飞的字节数 | 保住载入宽度优先，改动逐个实测 | +| block 少 | block 数不到 SM 数的两倍 | 每条载入指令的宽度、在飞的字节数 | 优先保证载入宽度，每项改动单独实测 | -- **数据小时，发射开销与缓存状态占主导。** 同一个行求和 kernel 的四种 access pattern（fp16，256 线程，时钟未锁）在 65536 × 4096（512 MB）上测得的访存带宽是 4.20 到 4.43 TB/s，彼此相差不超过 6%；换成 2048 × 4096（16 MB，装得进 L2）后单次耗时十几微秒，同一个 access pattern 两次测量之间可差三倍。这个区间里换 access pattern 没有收益，制约性能的是别的因素。 -- **block 少时，载入宽度带来的收益大于合并规则算出的差别。** warp 数量不足，靠并发的请求数掩盖访存延迟不再可行，只能让每个请求更宽、每个线程持有更多在飞的字节。以载入宽度换取其他好处的改法，在这个区间都可能反转。 +- **数据小时，发射开销与缓存状态占主导。** 同一个行求和 kernel 的四种 access pattern（fp16，256 线程，时钟未锁）在 65536 × 4096（512 MB）上测得的访存带宽是 4.20 到 4.43 TB/s，彼此相差不超过 6%；换成 2048 × 4096（16 MB，装得进 L2）后单次耗时十几微秒，同一个 access pattern 两次测量之间可差三倍。在这个区间里，制约性能的是其他因素，更换 access pattern 没有收益。 +- **block 少时，载入宽度带来的收益大于合并规则算出的差别。** warp 数量不足时，无法靠并发的请求数掩盖访存延迟，只能让每个请求更宽、每个线程持有更多在飞的字节。牺牲载入宽度来换取其他收益的改法，在这个区间都可能得到相反的结果。 ## 合并 global memory 的访存 {#coalescing} @@ -32,33 +32,33 @@ | cache line | 128 字节 | L1 与 L2 的缓存行，也是缓存查找的单位 | | **sector** | **32 字节** | 一条 cache line 由 4 个 sector 组成，L1 与 L2 之间按 sector 传输 | -缓存查找时以 cache line 为单位，搬运数据时以 sector 为单位：某个 sector 未命中，L1 就只向 L2 请求这一个 sector，不必把整条 cache line 都拉过来。由此得到的结论是**取 1 个字节和取满 32 个字节的代价相同**。于是一条访存指令的好坏由 **sector 利用率**衡量：`真正用到的字节 / (覆盖的 sector 数 × 32)`。 +缓存查找时以 cache line 为单位，搬运数据时以 sector 为单位：某个 sector 未命中，L1 就只向 L2 请求这一个 sector，不必把整条 cache line 都拉过来。因此**取 1 个字节和取满 32 个字节的代价相同**，一条访存指令的效率由 **sector 利用率**衡量：`真正用到的字节 / (覆盖的 sector 数 × 32)`。 -硬件把一个 warp 的 32 个访问合并成尽可能少的 32 字节事务。事务数最少要同时满足三个因素： +硬件把一个 warp 的 32 个访问合并成尽可能少的 32 字节事务。事务数降到最少需要同时满足三个条件： 1. **地址连续** —— 同一条指令里 32 个线程的地址首尾相接，不留空洞； 2. **按 32 字节对齐** —— 起始地址是 32 的倍数，一段数据不会多占一个 sector； 3. **每个线程一次取满 16 字节** —— 一条指令覆盖 $32 \times 16 = 512$ 个连续字节，即 16 个满载的 sector。 -三个因素同时成立时，这条访存指令对硬件最友好。 +三个条件同时成立时，这条访存指令对硬件最友好。 一个线程要读 $V$ 个元素时（$V$ = 一行的元素数 / 线程数），有四种 access pattern。 -**blocked** —— 每个线程负责一段连续的元素。固定 `c` 时相邻线程的地址相隔 $V$ 个元素，违反第一个因素，sector 利用率是 $1/V$： +**blocked** —— 每个线程负责一段连续的元素。固定 `c` 时相邻线程的地址相隔 $V$ 个元素，违反第一个条件，sector 利用率是 $1/V$： ```python for c in T.serial(V): acc[0] = acc[0] * X[row, tx * V + c] ``` -**striped** —— 相邻线程取相邻元素。地址连续了，但每个线程一次只取一个元素，违反第三个因素，$V$ 个元素要发 $V$ 条指令： +**striped** —— 相邻线程取相邻元素。地址连续，但每个线程一次只取一个元素，违反第三个条件，$V$ 个元素需要 $V$ 条指令： ```python for c in T.serial(V): acc[0] = acc[0] * X[row, c * threads + tx] ``` -**blocked + 向量化** —— 仍是每线程一段连续的，但改用 `T.vectorized` 一次读满 16 字节，三个因素全部满足： +**blocked + 向量化** —— 每个线程仍负责一段连续的元素，但改用 `T.vectorized` 一次读满 16 字节，三个条件全部满足： ```python buf = T.alloc_local((V,), dtype) @@ -69,7 +69,7 @@ for c in T.serial(V): acc[0] = acc[0] * buf[c] ``` -**staged** —— 搬运交给 `T.Parallel`，消费改从 shared memory 读，同样满足三个因素： +**staged** —— 由 `T.Parallel` 完成搬运，消费阶段改从 shared memory 读取，同样满足三个条件： ```python sh = T.alloc_shared((threads, V + pad), dtype) @@ -83,15 +83,21 @@ for c in T.serial(V): 四种 access pattern 的差别在于：**「哪个线程读哪些元素、一次读多宽」这个映射由谁决定。** -`T.serial` 的语义是循环体由单个线程顺序执行，索引表达式被逐字翻译成访存指令，不做合并也不做向量化 —— 编程者写出的模式就是硬件看到的模式。 +`T.serial` 的语义是循环体由单个线程顺序执行，索引表达式被逐字翻译成访存指令，不做合并，也不做向量化，因此编程者写出的模式就是硬件看到的模式。 -`T.vectorized`、`T.Parallel`、`T.copy` 则由 TileLang 的 **layout inference** 决定，三者的差别在于编程者还需要写明多少：`T.vectorized` 要写明每线程一次访问的宽度，线程映射由 layout inference 推导；`T.Parallel` 连宽度也不必写，循环维度怎么分给线程、一次读多宽都由它决定；`T.copy` 只写源和目标两个区域，整段搬运由它生成（需要接管推导结果时，另有 `coalesced_width` 与 `loop_layout` 两个参数）。剩下的向量化、地址对齐、以及在 shared memory 一侧避开 bank 冲突，都由 layout inference 负责 —— 这些正是对硬件友好但手写容易出错的部分。 +`T.vectorized`、`T.Parallel`、`T.copy` 的映射则由 TileLang 的 **layout inference** 决定。三者的差别在于编程者需要写明的内容： -**用 `T.serial` 手写下标时，上面三个因素要自己逐一保证；交给 layout inference 时，只需写明搬运的范围。** 三者之间怎么选、各自能跑到多少带宽，见下面的实测。 +- `T.vectorized` 写明每个线程一次访问的宽度，线程映射由 layout inference 推导； +- `T.Parallel` 不必写宽度，循环维度如何分给线程、一次读多宽都由 layout inference 决定； +- `T.copy` 只写源区域与目标区域，整段搬运由它生成；需要自行指定推导结果时，另有 `coalesced_width` 与 `loop_layout` 两个参数。 + +其余的向量化、地址对齐，以及在 shared memory 一侧避开 bank 冲突，都由 layout inference 负责。这些正是对硬件友好、但手写容易出错的部分。 + +**用 `T.serial` 手写下标时，三个条件由编程者逐一保证；交给 layout inference 时，只需写明搬运的范围。** 各种写法的选择与各自达到的带宽见下文的实测。
- + blocked tx * V + c @@ -431,7 +437,7 @@ for c in T.serial(V): ## 实测对比 -我们对两个 workload 在 H200 上进行实测，比较上面四种 access pattern 各自能跑到多少**访存带宽**（搬运的字节数除以 kernel 耗时，单位 TB/s），这两个 workload 的计算对元素的处理顺序有不同要求 —— 这个要求会决定哪几种 access pattern 可用。 +本节在 H200 上实测两个 workload，比较四种 access pattern 各自达到的**访存带宽**（搬运的字节数除以 kernel 耗时，单位 TB/s）。两个 workload 的计算对元素处理顺序的要求不同，这一要求决定了哪几种 access pattern 可用。 测试中 SM 时钟锁在 1830 MHz；输入 bf16 的 $65536 \times 4096$（512 MB，**必须大于 L2 的 60 MiB**，否则测到的是 L2 带宽）；每个配置跑三次，三次的结果一致到 ±0.5%。staged 在表里占两列：一列不加 pad（此时 stride 恰好是 $V$ 个 word，产生 bank 冲突，见[优化 shared memory 访问](shared-memory-access.md)），一列是在若干个 pad 取值中测到的最优值。 @@ -457,15 +463,15 @@ for c in T.serial(V): ## access pattern 的取舍 -1. **逐元素的 blocked 在 $V > 1$ 时总是最差的 access pattern。** 固定 `c` 时相邻线程的地址相隔 $V$ 个元素，sector 利用率是 $1/V$，所以 $V$ 越大越差 —— workload 1 的表里从 $V = 8$ 的 3.02 掉到 $V = 64$ 的 0.48。这个关系由访存合并的规则决定，不随形状改变。 +1. **逐元素的 blocked 在 $V > 1$ 时总是最差的 access pattern。** 固定 `c` 时相邻线程的地址相隔 $V$ 个元素，sector 利用率是 $1/V$，所以 $V$ 越大越差：workload 1 的表中从 $V = 8$ 的 3.02 掉到 $V = 64$ 的 0.48。这个关系由访存合并的规则决定，不随形状改变。 -2. **$V$ 小时用向量化的 blocked；$V$ 大到寄存器压力压低占用率时，改用加了 pad 的 staged。** 向量化把整段留在寄存器里（bf16 是每线程 $V/2$ 个），staged 把它放进 shared memory，用一次同步换回寄存器。翻转点取决于 kernel 里其余部分还剩多少寄存器预算，不是一个固定的 $V$：上面两个 workload 在同一个行宽下就分别落在 $V = 64$ 与 $V = 32$。**这个翻转点要在自己的 kernel 上测。** +2. **$V$ 小时用向量化的 blocked；$V$ 大到寄存器压力压低占用率时，改用加了 pad 的 staged。** 向量化把整段留在寄存器里（bf16 是每线程 $V/2$ 个），staged 把它放进 shared memory，以一次同步为代价节省寄存器。翻转点取决于 kernel 其余部分剩下的寄存器预算，没有固定的 $V$：上面两个 workload 在同一个行宽下分别落在 $V = 64$ 与 $V = 32$。**翻转点需要在具体的 kernel 上实测。** 3. **staged 的 shared 缓冲要避开 bank 冲突。** 声明成 `(threads, V)` 时 stride 恰好是 $V$ 个 word，$V$ 为 2 的幂就一定产生冲突；pad 的算法与候选见[优化 shared memory 访问](shared-memory-access.md#pad-per-chunk)。workload 2 的表里，同一个配置不加 pad 是 0.46，加 pad 是 3.69。 -4. **striped 完全合并，但每个元素要发一条指令。** 所以它好于逐元素的 blocked、差于向量化的 blocked（$V = 16$ 上 3.31 对 1.83 与 3.81），适合改动量比最后一点带宽更重要的场合。它让线程持有的元素不连续，因此要求线程持有连续一段的计算（例如串行前缀）用不了它。 +4. **striped 完全合并，但每个元素需要一条指令。** 因此它好于逐元素的 blocked、差于向量化的 blocked（$V = 16$ 上 3.31 对 1.83 与 3.81），适合改动量比最后一点带宽更重要的场合。它使每个线程持有的元素不连续，因此不能用于要求线程持有连续一段元素的计算，例如串行前缀积。 -下面两段是推荐 access pattern 的完整模板，`M`、`N`、`V`、`threads`、`pad`、`dtype` 都是编译期常量。本页开头的四段代码里，逐元素的 blocked 是反例，不要照抄；striped 可用但不是最快的一种（见上面取舍的第 4 条）。 +下面两段是推荐 access pattern 的完整模板，`M`、`N`、`V`、`threads`、`pad`、`dtype` 都是编译期常量。「合并 global memory 的访存」一节的四段代码中，逐元素的 blocked 是反例；striped 可用，但不是最快的一种（见取舍的第 4 条）。 **推荐的 access pattern，小 $V$** —— 向量化的 blocked： @@ -487,7 +493,7 @@ def main(X: T.Tensor((M, N), dtype), Out: T.Tensor((M, threads), "float32")): Out[row, tx] = acc[0] ``` -**推荐的 access pattern，大 $V$** —— 加了 pad 的 staged（翻转点见上面取舍的第 2 条）： +**推荐的 access pattern，大 $V$** —— 加了 pad 的 staged（翻转点见取舍的第 2 条）： ```python @T.prim_func diff --git a/docs/performance-guides/memory-bound/index.md b/docs/performance-guides/memory-bound/index.md index c6409f1a..3756bb66 100644 --- a/docs/performance-guides/memory-bound/index.md +++ b/docs/performance-guides/memory-bound/index.md @@ -12,8 +12,9 @@ point for tuning. On an H200, the measured values are [**57.27** TFLOP/s for fp32 FMA and **4.50** TB/s of memory bandwidth](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml). The **ridge point** of the roofline is where the bandwidth slope meets the compute ceiling. Dividing compute by bandwidth places it at an arithmetic -intensity of **12.72 flop/byte**: the number of floating-point operations per -byte moved when compute and bandwidth are saturated at the same time: +intensity of **12.72 flop/byte**, the number of floating-point operations per +byte moved when compute and bandwidth are saturated at the same time. The +figure below shows the H200 roofline:
diff --git a/docs/performance-guides/memory-bound/index.zh.md b/docs/performance-guides/memory-bound/index.zh.md index e9d83131..1ad1e589 100644 --- a/docs/performance-guides/memory-bound/index.zh.md +++ b/docs/performance-guides/memory-bound/index.zh.md @@ -2,7 +2,7 @@ ## 什么是访存受限 -在 GPU 上，一个 kernel 跑多快，取决于算力与带宽哪一个先成为瓶颈。TileOPs 用 [macro benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/hardware) 测算出一个**[校准系数](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)**：硬件 spec 给出的理论峰值乘以校准系数，得到实际可达的有效值，以此作为性能优化的指导标准。我们在 H200 上实测出 [fp32 FMA 的算力为 **57.27** TFLOP/s，访存带宽为 **4.50** TB/s](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)。roofline 的**拐点**（ridge point）是带宽斜线与算力上限这两段的交点，两者相除给出它的横坐标，也就是拐点处的算术强度 **12.72 flop/byte** —— 算力与带宽同时用满时，每搬运一个字节对应的浮点运算次数： +在 GPU 上，一个 kernel 的速度取决于算力与带宽哪一个先成为瓶颈。TileOPs 用 [macro benchmark](https://github.com/tile-ai/TileOPs/tree/main/benchmarks/hardware) 测算出一个**[校准系数](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)**：硬件 spec 给出的理论峰值乘以校准系数，得到实际可达的有效值，作为性能优化的参照。在 H200 上实测得到的 [fp32 FMA 的算力为 **57.27** TFLOP/s，访存带宽为 **4.50** TB/s](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/perf/profiles/h200.yaml)。roofline 的**拐点**（ridge point）是带宽斜线与算力上限这两段的交点，算力除以带宽得到它的横坐标，即拐点处的算术强度 **12.72 flop/byte**，也就是算力与带宽同时用满时每搬运一个字节对应的浮点运算次数。H200 的 roofline 如下图所示：
@@ -52,7 +52,7 @@ 算力上限 57.3 TFLOP/s
-
图上这条折线就是 roofline，任何 kernel 的性能点都落在它以下。拐点以左，上限是「算术强度 × 带宽」，可达算力随算术强度线性上升；拐点以右，上限就是算力峰值，不再随算术强度变化。
silu 每搬运 4 个字节做 5 次运算，算术强度 1.25 flop/byte，只有拐点的 1/10，所以即便带宽完全用满，也只能达到算力上限的 10%。 +
图中的折线是 roofline，任何 kernel 的性能点都落在它以下。拐点以左，上限是「算术强度 × 带宽」，可达算力随算术强度线性上升；拐点以右，上限就是算力峰值，不再随算术强度变化。silu 每搬运 4 个字节做 5 次运算，算术强度 1.25 flop/byte，只有拐点的 1/10，所以即便带宽完全用满，也只能达到算力上限的 10%。
diff --git a/docs/performance-guides/memory-bound/shared-memory-access.md b/docs/performance-guides/memory-bound/shared-memory-access.md index 44059b7c..78ef2972 100644 --- a/docs/performance-guides/memory-bound/shared-memory-access.md +++ b/docs/performance-guides/memory-bound/shared-memory-access.md @@ -1,14 +1,14 @@ # Optimizing Shared Memory Access -Routing data through shared memory adds another access pattern to consider. -Writing data from global memory into shared memory and reading it from shared -memory into registers both access shared memory. This page covers bank conflicts -in those two steps: what determines them, how padding removes them, and what the -padding should be computed from. +When data passes through shared memory, the access pattern on shared memory +also matters. Data is written from global memory into shared memory and then +read from shared memory into registers, and both steps access shared memory. +This page describes bank conflicts in those two steps: what determines them, how +padding removes them, and how to compute the padding. Every measurement on this page was taken on an H200 with the SM clock locked at 1830 MHz, an input larger than the 60 MiB L2, and enough blocks to fill the -whole card. Outside those conditions, the conclusions can reverse; the tests are +whole card. Outside those conditions, the conclusions can reverse. The tests are in [Optimizing Global Memory Access](global-memory-access.md#regime). ## The bank structure of shared memory {#bank-conflict} @@ -33,9 +33,8 @@ A shared-memory access pattern should be conflict-free whenever possible. ## What determines the conflict degree -Consider a case that is general enough to reason from. The one-dimensional array -`sh` below lives in shared memory, and each thread reads `chunk` contiguous -elements from it: +Take the following general case as the example. The one-dimensional array `sh` +is in shared memory, and each thread reads `chunk` contiguous elements from it. ```python sh = T.alloc_shared((threads * chunk,), dtype) # one dimension, threads * chunk elements @@ -46,60 +45,61 @@ for c in T.serial(chunk): The 32 threads of a warp execute this loop in lockstep. Within one iteration, `c` has the same value for all of them, while `tx` runs from 0 to 31. The 32 -threads therefore access shared memory at a fixed spacing. At `chunk = 64`, one +threads therefore access shared memory at equal spacing. At `chunk = 64`, one iteration reads elements `c`, `64 + c`, `128 + c`, …, `1984 + c`. Adjacent threads are `chunk` elements apart. That difference is the **stride** of the access, written $S$; in words, $S = \text{chunk} \times E / 4$, where $E$ is the -element size in bytes. The stride is independent of how many dimensions the -array declaration uses and how the index is written. +element size in bytes. The stride is independent of the number of dimensions of +the array and of how the index is written. The other variable is the width of the memory instruction. A vectorized access reads $w$ contiguous words at a time, where $w$ is 1, 2, or 4: 32 bit, the 8 bytes of a `float2`, or the 16 bytes of a `float4`. Thread $t$ then reads words $St$ through $St + w - 1$. -Under two premises, the conflict degree can be counted directly from the -hardware facts above: +When the following two premises hold, the conflict degree can be computed +directly from the bank structure: -1. **$S$ is a whole number of words.** An odd `chunk` in fp16 makes it a half - word, so a common divisor is not defined. The remaining method is to compute - each thread's bank from its byte address and count. +1. **$S$ is a whole number of words.** With an odd `chunk` in fp16, $S$ + contains half a word, and a common divisor cannot be taken. The bank of each + thread then has to be computed from its byte address one by one, and + counted. 2. **The $32w$ words the 32 threads request are distinct**, that is $S \ge w$. - Threads landing on one word are broadcast and cost no extra cycle, so - counting $32w$ words no longer holds; and at $S < w$ the vector ranges of - adjacent threads overlap, which breaks it the same way. + Threads that land on one word are served by broadcast and cost no extra + cycle, so the count of $32w$ words no longer holds. At $S < w$ the vector + ranges of adjacent threads overlap, and the count fails in the same way. The warp fetches $32w$ words in total, and shared memory serves at most 32 words per cycle. The instruction therefore takes at least $w$ cycles. That lower bound comes from width alone and is independent of the addresses; a conflict is any cost above it. $St \bmod 32$ takes only multiples of $g = \gcd(S, 32)$, so the -access reaches $32/g$ banks, each hit $g$ times. The shifts +access reaches only $32/g$ banks, each accessed $g$ times. The shifts $j = 0, \dots, w-1$ are then added on top. Both $g$ and $w$ are powers of two, so one divides the other, giving two cases: | | Where the banks fall | Cycles | Conflict | | --- | --- | --- | --- | | $g \le w$ | All 32 banks requested $w$ times each | $w$, exactly the lower bound | None | -| $g > w$ | Only $(32/g) \cdot w$ banks hit, $g$ times each | $g$ | $g / w$-way | +| $g > w$ | Only $(32/g) \cdot w$ banks accessed, $g$ times each | $g$ | $g / w$-way | $$\text{conflict degree} \ \ge\ \max\left(1,\ \frac{\gcd(S,\ 32)}{w}\right)$$ -For a scalar read ($w = 1$), this is just $\gcd(S, 32)$. When $S$ is coprime +For a scalar read ($w = 1$), the bound is $\gcd(S, 32)$. When $S$ is coprime with 32, the 32 threads cover all 32 banks with no conflict. When $S$ is a multiple of 32, they all land on one bank and serialize 32 ways. fp16 with `chunk = 64` is the latter case: $S$ is 128 bytes, or 32 words. -The result is an inequality because the last step assumes the hardware can fit +The formula is an inequality because the last step assumes the hardware can fit any conflict-free set of words into one cycle. NVIDIA has not published how lanes are actually grouped for 64-bit and 128-bit accesses. ## Changing the stride with padding The stride follows from `chunk`, and `chunk` is usually fixed by the algorithm -rather than free to change. One way to remove an $N$-way conflict is to add -`pad` elements to the end of each run, making the stride `chunk + pad`. That +and often cannot be changed freely. One way to handle an $N$-way conflict is to +add `pad` elements to the end of each run, making the stride `chunk + pad`. That changes the `gcd`, and therefore the conflict degree. For fp16 with -`chunk = 64`, adding just 2 elements changes the stride from 32 words to 33, +`chunk = 64`, adding 2 elements changes the stride from 32 words to 33, which is coprime with 32, and the conflict disappears entirely.
@@ -286,7 +286,8 @@ which is coprime with 32, and the conflict disappears entirely. ## How much padding to use -The formula above gives the test, and **both conditions must hold**: +The conflict formula and the alignment requirement on the start of each run +together give the test, and **both conditions must hold**: | Access width | $w$ | Conflict requirement | Alignment of a run's start | Fastest pad at fp32, `chunk = 64` | | --- | --- | --- | --- | --- | @@ -294,9 +295,9 @@ The formula above gives the test, and **both conditions must hold**: | 64 bit, `float2` | 2 | $\gcd(S, 32) \le 2$ | 8 bytes | `pad = 2` ($S = 66$ words) | | 128 bit, `float4` | 4 | $\gcd(S, 32) \le 4$ | 16 bytes | `pad = 4` ($S = 68$ words) | -The two conditions pull in opposite directions, so neither is sufficient by -itself: **the wider the access, the looser the conflict requirement and the -tighter the alignment.** An odd pad is optimal for a scalar read and the worst +**The wider the access, the looser the conflict requirement and the tighter the +alignment requirement.** The two requirements move in opposite directions, so +neither one is sufficient by itself. An odd pad is optimal for a scalar read and the worst case for a vectorized read, because it shifts the start of each run by 4 bytes. A 128-bit shared read requires natural 16-byte alignment, and its behavior is undefined when the alignment is insufficient. If the compiler can see the @@ -304,9 +305,9 @@ insufficient alignment at compile time, it falls back to a narrower instruction; the $\gcd = 1$ row below shows that case. The measurements below were taken on an H200 with fp32 and `chunk = 64`. The -consumption loop is repeated 32 times so the shared-memory side becomes the +consumption loop is repeated 32 times so the shared memory side becomes the bottleneck, and time scales proportionally with the repeat count. Each entry is -how many times slower that combination is than the fastest row at the same +the time of that combination divided by the time of the fastest row at the same width: | $\gcd(S, 32)$ | 32 bit | 64 bit | 128 bit | Predicted bound (32 / 64 / 128) | @@ -318,8 +319,8 @@ width: | 16 | 12.97 | 7.62 | 7.99 | 16 / 8 / 4 | | 32 | 25.82 | 13.97 | 15.71 | 32 / 16 / 8 | -The three 1.00 entries on the diagonal are the three recommendations in the -previous table. The bound is tight on the $g \le w$ side. On the $g > w$ side, +The three 1.00 entries on the diagonal correspond to the three pads recommended +in the previous table. The bound is tight on the $g \le w$ side. On the $g > w$ side, it is not tight, and the 128-bit column lands at exactly twice the bound (4.04 versus 2, 7.99 versus 4, 15.71 versus 8). Explaining the cause requires lower-level evidence, because NVIDIA has not published how lanes are grouped at @@ -328,8 +329,8 @@ $S = 65$ words is 260 bytes, not a multiple of 16. ## Padding has to be computed from the chunk {#pad-per-chunk} -Padding by a fixed number of bytes falls back into the worst case for some -values of `chunk`. The alignment requirement restricts the candidates to whole +For some values of `chunk`, a pad of a fixed number of bytes brings back the +worst-case conflict. The alignment requirement restricts the candidates to whole multiples of 16 bytes. Write the pad as $k$ multiples of 16 bytes ($k = 1, 2, \dots$): @@ -346,8 +347,8 @@ With $k$ fixed, $\gcd(S, 32)$ varies with `chunk`. For fp16 and bf16 | 72 | 40 | 8 | 44 | **4** | | 128 | 68 | **4** | 72 | 8 | -In the `chunk = 56` row, $S$ is exactly 32 words, so the accesses land on the -same bank as with `pad = 0`: padding was added, but the conflict remains. The +At `chunk = 56` and $k = 1$, $S$ is exactly 32 words, so all accesses land on +the same bank, as with `pad = 0`: the pad is added, but the conflict remains. The fix is to choose whichever of $k = 1$ and $k = 2$ gives the smaller $\gcd(S, 32)$. Whenever $\text{chunk} \times E$ is a multiple of 16, one of the two choices is guaranteed to reduce the $\gcd$ to 4. Writing @@ -370,22 +371,22 @@ differ only in the pad: Those four widths are the hidden size of Qwen2-7B, the hidden size of Llama-3-70B, and the FFN intermediate dimensions of Llama-3-8B and -Llama-3-70B. They are not constructed counterexamples. +Llama-3-70B. All four come from real models. -The shared-memory side of these measurements is not a scalar read. `cuobjdump` +The shared memory side of these measurements is a 64-bit read. `cuobjdump` shows 16 `LDS.64` instructions per thread at $S = 36$ words, so $w = 2$ and the degree is $\gcd(S, 32) / 2$: 16-way and 2-way for the two columns. The degrees differ by a factor of 8, while bandwidth differs by only 2.2, because the 2-way column is limited by DRAM again. When the start of a run is misaligned by -less than 8 bytes, the compiler falls back to `LDS` ($w = 1$). Width is decided -by the compiler, not by the code that declares the pad, so both columns above -must be measured rather than derived from the $\gcd$ alone. +less than 8 bytes, the compiler uses `LDS` ($w = 1$) instead. The width is +decided by the compiler, not by the developer who declares the pad, so both +columns above must be measured rather than derived from the $\gcd$ alone. ## A measured sweep -The measurements below read one element per thread at a time ($w = 1$), use -powers of two for `chunk`, and give each thread its own run. Both premises above -therefore hold. +In the measurements below, each thread reads its own run one element at a time +($w = 1$), and `chunk` is a power of two. Both premises in What determines the +conflict degree therefore hold. H200, SM clock locked at 1830 MHz. fp16, input $65536 \times 4096$ (512 MB, larger than the 60 MiB L2). The kernel stages a whole row into shared memory, @@ -411,13 +412,15 @@ can be used directly to narrow the padding candidates. the table above, 1-way and 2-way both stay above 3.0, 4-way drops to around 2.6, and 8-way and worse fall below 1.6. The formula is useful for narrowing the candidates to those that compute to no more than $w$ ways. Among those - candidates, measure: 2-way edges out 1-way on all four values of `chunk`, but - only by 0.9% to 10%, which is not a rule worth carrying over. + candidates, measure: 2-way is slightly faster than 1-way on all four values + of `chunk`, but the gap ranges from 0.9% to 10%, which does not make a rule + that carries over. 2. **After changing the pad, sweep `chunk` again.** The best entry in the `pad = 0` column is `chunk = 16` (1.63); in the `pad = 4` column, it is `chunk = 64` (3.70). The ordering of the first column is set mostly by - conflict degree: `chunk = 16` hits 8-way, while `chunk = 64` hits 32-way. + conflict degree: `chunk = 16` has an 8-way conflict, while `chunk = 64` has + a 32-way conflict. Once the conflicts are removed, all four cases are 2-way, and the optimum moves. Thread count and `chunk` move together in this table (their product is the row width, 4096), so `chunk` is not the only cause of the move; occupancy @@ -426,7 +429,9 @@ can be used directly to narrow the padding candidates. 3. **After changing `chunk`, recompute the pad.** Writing the pad as a fixed number of bytes assumes $\gcd(S, 32)$ is independent of `chunk`, and - [the table above](#pad-per-chunk) shows that it is not: at `chunk = 56` and + the table in + [Padding has to be computed from the chunk](#pad-per-chunk) shows that it is + not: at `chunk = 56` and $k = 1$, $S$ returns to 32 words. Writing the pad as a function of `chunk`, by taking whichever of $k = 1$ and $k = 2$ gives the smaller $\gcd(S, 32)$, keeps it consistent with the test on this page. @@ -435,10 +440,10 @@ can be used directly to narrow the padding candidates. trade-offs in [Optimizing Global Memory Access](global-memory-access.md#coalescing), small $V$ calls for vectorized blocked. In that pattern, data goes straight into - registers, never touches shared memory, and has no bank conflicts. This page - applies once $V$ grows enough that register pressure cuts occupancy and - staged takes over, and when a whole row has to be shared by every thread in - the block. + registers, never passes through shared memory, and has no bank conflicts. + This page applies in two cases: when $V$ grows enough that register pressure + cuts occupancy and staged is used instead, and when a whole row has to be + shared by every thread in the block. The two listings below declare the shared buffer and differ only in the stride. The wrong form has a stride that is an exact multiple of 32 words: diff --git a/docs/performance-guides/memory-bound/shared-memory-access.zh.md b/docs/performance-guides/memory-bound/shared-memory-access.zh.md index d4f5eb6a..f01506c8 100644 --- a/docs/performance-guides/memory-bound/shared-memory-access.zh.md +++ b/docs/performance-guides/memory-bound/shared-memory-access.zh.md @@ -1,22 +1,22 @@ # 优化 shared memory 访问 -数据经过 shared memory 中转时，多了一处需要考虑 access pattern 的地方：从 global memory 写进来、再从 shared memory 读进寄存器，两步都在访问 shared memory。这一页讲这两步上的 bank conflict —— 它由什么决定、怎么用 pad 消掉、以及 pad 该按什么算。 +数据经过 shared memory 中转时，shared memory 上的 access pattern 也需要考虑：数据从 global memory 写入 shared memory，再从 shared memory 读入寄存器，这两步都访问 shared memory。本页说明这两步上的 bank conflict：它由什么决定，如何用 pad 消除，以及 pad 的取值如何计算。 -这一页的实测都在 H200 上、SM 时钟锁在 1830 MHz、输入大于 L2 的 60 MiB、block 数足以填满整卡 —— 这组条件之外结论可能反转，判据见[优化 global memory 访问](global-memory-access.md#regime)。 +本页的实测条件是：H200，SM 时钟锁在 1830 MHz，输入大于 L2 的 60 MiB，block 数足以填满整卡。在这组条件之外，结论可能反转，判据见[优化 global memory 访问](global-memory-access.md#regime)。 ## shared memory 的 bank 结构 {#bank-conflict} -shared memory 由 32 个 bank 构成，每 bank 宽 4 个字节。将 shared memory 的地址空间按 4 字节划分成 **word**（下文一律按 word 计数），一个地址落在哪个 bank 上，由 `(字节地址 / 4) mod 32` 决定。一个 bank 每周期只能处理一个 word 的访问，多个线程并发访问 shared memory 时，只要访问落在不同的 bank 上，它们就在同一个周期一起完成。多个线程落在同一条 bank 上时，又分三种情形： +shared memory 由 32 个 bank 构成，每个 bank 宽 4 字节。将 shared memory 的地址空间按 4 字节划分成 **word**（下文一律按 word 计数），一个地址落在哪个 bank 上，由 `(字节地址 / 4) mod 32` 决定。一个 bank 每周期只能处理一个 word 的访问。多个线程并发访问 shared memory 时，只要访问落在不同的 bank 上，就在同一个周期内一起完成。多个线程落在同一条 bank 上时，分三种情形： 1. **访问的是不同的 word** —— 硬件把这次请求拆成若干次无冲突的请求依次完成，拆分的次数就是**冲突的路数**。 -2. **读取的是同一个 word** —— 任意两个线程只要落在同一个 word 内（哪怕取的是其中不同的字节），这个 word 会被广播给所有请求它的线程，不产生冲突。分布在不同 bank 上的多个广播还会合并为一次 multicast。 +2. **读取的是同一个 word** —— 任意两个线程只要落在同一个 word 内（即使取的是其中不同的字节），这个 word 会被广播给所有请求它的线程，不产生冲突。分布在不同 bank 上的多个广播还会合并为一次 multicast。 3. **写入的是同一个地址** —— 只有一个写入生效，是哪一个未定义。 访问 shared memory 的 access pattern 应当尽可能做到无 bank conflict。 ## 冲突路数由什么决定 -我们考虑一种不失通用性的情况。下面这个一维数组 `sh` 在 shared memory 上，每个线程读取其中连续的 `chunk` 个元素： +以下面的一般情形为例：一维数组 `sh` 位于 shared memory，每个线程读取其中连续的 `chunk` 个元素。 ```python sh = T.alloc_shared((threads * chunk,), dtype) # 一维数组，threads * chunk 个元素 @@ -25,31 +25,31 @@ for c in T.serial(chunk): acc[0] = acc[0] * sh[tx * chunk + c] # 线程 tx 读自己那一段 ``` -一个 warp 的 32 个线程同步执行这个循环，同一次迭代里 `c` 对它们取同一个值、`tx` 取 0 到 31：这时 32 个线程以等间隔访问 shared memory。例如 `chunk = 64` 时，这 32 个线程在同一次迭代里读的是第 `c`、第 `64 + c`、第 `128 + c`、…… 第 `1984 + c` 个元素。相邻两个线程相差 `chunk` 个元素，这个差值称为这段访问的 **stride**，记作 $S$，换算成 word 是 $S = \text{chunk} \times E / 4$ 个（$E$ 是元素的字节数）。stride 与数组声明成几维、下标怎么写都无关。 +一个 warp 的 32 个线程同步执行这个循环，同一次迭代中 `c` 对它们取同一个值，`tx` 取 0 到 31，因此 32 个线程以相等的间隔访问 shared memory。例如 `chunk = 64` 时，这 32 个线程在同一次迭代里读的是第 `c`、第 `64 + c`、第 `128 + c`、…… 第 `1984 + c` 个元素。相邻两个线程相差 `chunk` 个元素，这个差值称为这段访问的 **stride**，记作 $S$，换算成 word 是 $S = \text{chunk} \times E / 4$ 个（$E$ 是元素的字节数）。stride 与数组的维数、下标的写法都无关。 另一个变量是访存指令的位宽。以向量化的方式一次访问 $w$ 个连续的 word，$w$ 取 1、2、4，对应 32 bit、`float2` 的 8 字节、`float4` 的 16 字节。于是线程 $t$ 读的是第 $St$ 到 $St + w - 1$ 个 word。 -两条前提成立时，冲突路数可以直接从上面的硬件事实数出来： +以下两条前提成立时，冲突路数可以由上述 bank 结构直接算出： -1. **$S$ 是整数个 word。** fp16 的 `chunk` 取奇数时它带半个 word，取公约数无从谈起，只能按字节地址逐个算出线程落在哪条 bank 上再数。 -2. **32 个线程请求的 $32w$ 个 word 互不相同**，即 $S \ge w$。落在同一个 word 上的线程走广播，不占额外的周期，下面按 $32w$ 个 word 计数就不成立；$S < w$ 时相邻线程的向量区间彼此重叠，同样不成立。 +1. **$S$ 是整数个 word。** fp16 的 `chunk` 取奇数时，$S$ 含半个 word，无法取公约数，只能按字节地址逐个算出每个线程落在哪条 bank 上，再计数。 +2. **32 个线程请求的 $32w$ 个 word 互不相同**，即 $S \ge w$。线程落在同一个 word 上时走广播，不占额外的周期，此时按 $32w$ 个 word 计数不成立；$S < w$ 时相邻线程的向量区间彼此重叠，同样不成立。 -整个 warp 要取 $32w$ 个 word，而 shared memory 每周期最多处理 32 个，所以这条指令至少要 $w$ 个周期 —— 这是位宽带来的下界，与地址无关，冲突指的是超出这个下界的部分。$St \bmod 32$ 只取到 $g = \gcd(S, 32)$ 的倍数，也就是 $32/g$ 条 bank，每条被碰到 $g$ 次；再叠加 $j = 0, \dots, w-1$ 的平移。$g$ 与 $w$ 都是 2 的幂，必有一个整除另一个，于是分两种情形： +整个 warp 要取 $32w$ 个 word，而 shared memory 每周期最多处理 32 个，所以这条指令至少需要 $w$ 个周期。这个下界由位宽决定，与地址无关；冲突指超出这个下界的部分。$St \bmod 32$ 只取 $g = \gcd(S, 32)$ 的倍数，即只落在 $32/g$ 条 bank 上，每条被访问 $g$ 次；再叠加 $j = 0, \dots, w-1$ 的平移。$g$ 与 $w$ 都是 2 的幂，必有一个整除另一个，因此分两种情形： | | bank 的落点 | 周期数 | 冲突 | | --- | --- | --- | --- | | $g \le w$ | 32 条 bank 各被请求 $w$ 次 | $w$，正好是下界 | 无 | -| $g > w$ | 只有 $(32/g) \cdot w$ 条被碰到，各 $g$ 次 | $g$ | $g / w$ 路 | +| $g > w$ | 只有 $(32/g) \cdot w$ 条被访问，各 $g$ 次 | $g$ | $g / w$ 路 | $$\text{冲突路数} \ \ge\ \max\left(1,\ \frac{\gcd(S,\ 32)}{w}\right)$$ -标量读（$w = 1$）时它就是 $\gcd(S, 32)$：$S$ 与 32 互素时 32 个线程铺满 32 条 bank，无冲突；$S$ 是 32 的倍数时全部挤在同一条 bank 上，32 路串行 —— fp16、`chunk = 64` 就是后者，$S$ 是 128 字节即 32 个 word。 +标量读（$w = 1$）时它就是 $\gcd(S, 32)$：$S$ 与 32 互素时 32 个线程铺满 32 条 bank，无冲突；$S$ 是 32 的倍数时全部落在同一条 bank 上，32 路串行。fp16、`chunk = 64` 属于后者，$S$ 是 128 字节即 32 个 word。 -写成不等式，是因为最后一步假定硬件能把任意一组无冲突的 word 凑进一个周期，而 NVIDIA 没有公开 64 bit 与 128 bit 访问时 lane 的实际分组方式。 +上式写成不等式，原因是最后一步假定硬件能把任意一组无冲突的 word 放进一个周期，而 NVIDIA 没有公开 64 bit 与 128 bit 访问时 lane 的实际分组方式。 ## 用 pad 改变 stride -stride 由 `chunk` 决定，而 `chunk` 通常由算法决定，往往无法任意改动。当发生$N$路冲突时，一种可行方式是在每段末尾增加 `pad` 个元素，让 stride 变成 `chunk + pad` —— `gcd` 随之改变，冲突路数也就随之改变。例如：fp16、`chunk = 64`时，上只要加 2 个元素，stride 就从 32 个 word 变成 33 个，与 32 互素，冲突完全消失。 +stride 由 `chunk` 决定，而 `chunk` 通常由算法决定，往往不能任意改动。发生 $N$ 路冲突时，一种做法是在每段末尾增加 `pad` 个元素，使 stride 变为 `chunk + pad`，`gcd` 随之改变，冲突路数也随之改变。例如 fp16、`chunk = 64` 时，只需加 2 个元素，stride 就从 32 个 word 变为 33 个，与 32 互素，冲突完全消失。
@@ -228,13 +228,13 @@ stride 由 `chunk` 决定，而 `chunk` 通常由算法决定，往往无法任 fp16、每段 64 个元素、一个 warp 的 32 个线程。格内数字是落在该 bank 上的线程数，空格表示没有线程落上去。 -
四张图是同一个 kernel 在四个 pad 取值下，一个 warp 的 32 个线程在 32 条 bank 上的落点。格内数字是落在这条 bank 上的线程数，最大的那个数就是冲突路数。pad = 0 时 32 个线程全挤在 bank 0；加 2 个元素之后 stride 变成 33 个 word，与 32 互素，32 个线程正好铺满 32 条 bank。
+
四张图是同一个 kernel 在四个 pad 取值下，一个 warp 的 32 个线程在 32 条 bank 上的落点。格内数字是落在这条 bank 上的线程数，最大的那个数就是冲突路数。pad = 0 时 32 个线程全部落在 bank 0；加 2 个元素之后 stride 变成 33 个 word，与 32 互素，32 个线程正好铺满 32 条 bank。
## pad 该取多少 -判据由上面的式子给出，**两条都要满足**： +判据由上面的式子与段起点的对齐要求共同给出，**两条都要满足**： | 访问位宽 | $w$ | 冲突路数要求 | 段起点的对齐要求 | fp32、`chunk = 64` 上最快的 pad | | --- | --- | --- | --- | --- | @@ -242,9 +242,9 @@ stride 由 `chunk` 决定，而 `chunk` 通常由算法决定，往往无法任 | 64 bit，`float2` | 2 | $\gcd(S, 32) \le 2$ | 8 字节 | `pad = 2`（$S = 66$ word） | | 128 bit，`float4` | 4 | $\gcd(S, 32) \le 4$ | 16 字节 | `pad = 4`（$S = 68$ word） | -两条方向相反，所以不能只看一条：**位宽越宽，对冲突路数的要求越松，对对齐的要求越紧**。奇数 pad 对标量读是最优解，对向量化读却是最差的 —— 它让段起点错开 4 字节。128 bit 的 shared 读要求 16 字节自然对齐，对齐不足时这条指令的行为是未定义的；编译期能看出对齐不足时，编译器会退回窄指令，下面实测里 $\gcd = 1$ 那一行就是这种情形。 +**位宽越宽，对冲突路数的要求越松，对对齐的要求越紧**。两条要求方向相反，因此不能只看其中一条。奇数 pad 对标量读最优，对向量化读最差，因为它使段起点错开 4 字节。128 bit 的 shared 读要求 16 字节自然对齐，对齐不足时这条指令的行为是未定义的；编译期能看出对齐不足时，编译器会退回窄指令，下面实测里 $\gcd = 1$ 那一行就是这种情形。 -实测（H200，fp32，`chunk = 64`，消费循环重复 32 遍使 shared 一侧成为瓶颈，耗时随遍数成比例）给出每种组合比同位宽最快的那一行慢多少倍： +下表是实测结果（H200，fp32，`chunk = 64`，消费循环重复 32 遍使 shared 一侧成为瓶颈，耗时随遍数成比例），每个数值是该组合的耗时相对同一位宽最快一行的倍数： | $\gcd(S, 32)$ | 32 bit | 64 bit | 128 bit | 下界预测（32 / 64 / 128） | | --- | --- | --- | --- | --- | @@ -255,15 +255,15 @@ stride 由 `chunk` 决定，而 `chunk` 通常由算法决定，往往无法任 | 16 | 12.97 | 7.62 | 7.99 | 16 / 8 / 4 | | 32 | 25.82 | 13.97 | 15.71 | 32 / 16 / 8 | -对角线上的三个 1.00 就是上表那三行推荐。$g \le w$ 那一侧下界是紧的；$g > w$ 那一侧下界不紧，128 bit 实测恰好是它的 2 倍（4.04 对 2、7.99 对 4、15.71 对 8）；成因要更低层的证据才能定，NVIDIA 未公开这两种位宽下 lane 的分组方式。$\gcd = 1$ 那一行的 2.27 是对齐造成的：$S = 65$ word 即 260 字节，不是 16 的倍数。 +对角线上的三个 1.00 对应上一张表推荐的三个 pad。$g \le w$ 一侧的下界是紧的。$g > w$ 一侧的下界不紧，128 bit 的实测恰好是下界的 2 倍（4.04 对 2、7.99 对 4、15.71 对 8）；成因需要更底层的证据才能确定，NVIDIA 未公开这两种位宽下 lane 的分组方式。$\gcd = 1$ 那一行的 2.27 是对齐造成的：$S = 65$ word 即 260 字节，不是 16 的倍数。 ## pad 要按 chunk 算 {#pad-per-chunk} -固定字节数的 pad 会在某些 `chunk` 上落回最坏情形。段起点的对齐要求把候选限制在 16 字节的整数倍，记 pad 为 $k$ 个 16 字节（$k = 1, 2, \dots$），于是 +在某些 `chunk` 上，固定字节数的 pad 会重新出现最坏情形的冲突。段起点的对齐要求把候选限制在 16 字节的整数倍，记 pad 为 $k$ 个 16 字节（$k = 1, 2, \dots$），于是 $$S = \frac{\text{chunk} \times E}{4} + 4k \ \text{word}$$ -$k$ 固定时 $\gcd(S, 32)$ 随 `chunk` 变。fp16 与 bf16（$E = 2$）的几个 `chunk`： +$k$ 固定时，$\gcd(S, 32)$ 随 `chunk` 变化。下表列出 fp16 与 bf16（$E = 2$）的几个 `chunk`： | chunk | $S$（$k = 1$） | $\gcd(S, 32)$ | $S$（$k = 2$） | $\gcd(S, 32)$ | | --- | --- | --- | --- | --- | @@ -273,9 +273,9 @@ $k$ 固定时 $\gcd(S, 32)$ 随 `chunk` 变。fp16 与 bf16（$E = 2$）的几 | 72 | 40 | 8 | 44 | **4** | | 128 | 68 | **4** | 72 | 8 | -`chunk = 56` 那一行的 $S$ 恰好是 32 个 word，与 `pad = 0` 落在同一条 bank 上 —— pad 加了，冲突没消。做法是在 $k = 1$ 与 $k = 2$ 两个候选里取 $\gcd(S, 32)$ 小的那个。$\text{chunk} \times E$ 是 16 的倍数时，两者必有一个把 $\gcd$ 降到 4：记 $q = \text{chunk} \times E / 16$，则 $S = 4(q + k)$，$\gcd(S, 32) = 4 \gcd(q + k, 8)$，而 $q + 1$ 与 $q + 2$ 一奇一偶。 +`chunk = 56`、$k = 1$ 时 $S$ 恰好是 32 个 word，与 `pad = 0` 一样全部落在同一条 bank 上，加了 pad 也没有消除冲突。做法是在 $k = 1$ 与 $k = 2$ 两个候选里取 $\gcd(S, 32)$ 小的那个。$\text{chunk} \times E$ 是 16 的倍数时，两者必有一个把 $\gcd$ 降到 4：记 $q = \text{chunk} \times E / 16$，则 $S = 4(q + k)$，$\gcd(S, 32) = 4 \gcd(q + k, 8)$，而 $q + 1$ 与 $q + 2$ 一奇一偶。 -实测（H200，SM 时钟锁在 1830 MHz，bf16，CUPTI 设备耗时，每次迭代前清 L2，取 200 次的中位数，镜像 `ghcr.io/tile-ai/tileops-runner:cu132-torch2.13-tl-afcebed1-dev`）。四行的线程数都取到让 `chunk` 等于 56，同一行两列只差 pad： +实测（H200，SM 时钟锁在 1830 MHz，bf16，CUPTI 设备耗时，每次迭代前清 L2，取 200 次的中位数，镜像 `ghcr.io/tile-ai/tileops-runner:cu132-torch2.13-tl-afcebed1-dev`）。四行都选取线程数使 `chunk` 等于 56，同一行的两列只差 pad： | 输入 | 线程数 × chunk | $k = 1$（pad 8 个元素）
TB/s | $k = 2$（pad 16 个元素）
TB/s | | --- | --- | --- | --- | @@ -284,13 +284,13 @@ $k$ 固定时 $\gcd(S, 32)$ 随 `chunk` 变。fp16 与 bf16（$E = 2$）的几 | $1024 \times 14336$ | 256 × 56 | 1.51 | **3.05**{ .win } | | $512 \times 28672$ | 512 × 56 | 1.32 | **2.33**{ .win } | -这四个宽度是 Qwen2-7B 的 hidden size、Llama-3-70B 的 hidden size、Llama-3-8B 与 Llama-3-70B 的 FFN 中间维，不是构造出来的反例。 +这四个宽度分别是 Qwen2-7B 的 hidden size、Llama-3-70B 的 hidden size、Llama-3-8B 与 Llama-3-70B 的 FFN 中间维，都取自实际模型。 -这组测量的 shared 一侧不是标量读：`cuobjdump` 显示 $S = 36$ word 时每线程 16 条 `LDS.64`，即 $w = 2$，所以路数是 $\gcd(S, 32) / 2$ —— 两列分别是 16 路与 2 路。路数差 8 倍而带宽只差 2.2 倍，是因为 2 路那一列的瓶颈已经回到 DRAM。段起点错开 8 字节以下时它退回 `LDS`（$w = 1$）—— 位宽由编译器定，不由声明 pad 的人定，所以上表两列都要实测，不能只算 $\gcd$。 +这组测量中 shared 一侧是 64 bit 读：`cuobjdump` 显示 $S = 36$ word 时每线程 16 条 `LDS.64`，即 $w = 2$，所以路数是 $\gcd(S, 32) / 2$，两列分别是 16 路与 2 路。路数相差 8 倍而带宽只差 2.2 倍，原因是 2 路那一列的瓶颈已经回到 DRAM。段起点错开不足 8 字节时，编译器改用 `LDS`（$w = 1$）。位宽由编译器决定，不由声明 pad 的开发者决定，所以上表两列都要实测，不能只计算 $\gcd$。 ## 实测扫描 -下面的实测每线程逐个元素读（$w = 1$），`chunk` 取 2 的幂，每个线程读自己那一段 —— 上面两个前提都成立。 +下面的实测中，每个线程逐个元素读取自己那一段（$w = 1$），`chunk` 取 2 的幂，因此「冲突路数由什么决定」一节的两条前提都成立。 H200，SM 时钟锁在 1830 MHz。fp16，输入 $65536 \times 4096$（512 MB，大于 60 MiB 的 L2）。kernel 把整行搬进 shared memory，每个线程一段 `chunk + pad` 个元素，逐段做串行前缀积再写回，读加写共 1 GB。括号内是上面公式预测的冲突路数： @@ -305,21 +305,21 @@ H200，SM 时钟锁在 1830 MHz。fp16，输入 $65536 \times 4096$（512 MB， ## 使用时的注意事项 -1. **公式给出候选，最终值靠实测。** 上表里 1 路与 2 路都在 3.0 以上，4 路降到 2.6 附近，8 路及以上跌到 1.6 以下，所以公式的用处是把候选缩到「算出来不超过 $w$ 路」的那几个。这几个之间要实测：四组 `chunk` 上 2 路都略高于 1 路，但差距从 0.9% 到 10% 不等，不构成一条可以照搬的规则。 +1. **公式给出候选，最终取值由实测决定。** 上表里 1 路与 2 路都在 3.0 以上，4 路降到 2.6 附近，8 路及以上跌到 1.6 以下，所以公式的用处是把候选缩到「算出来不超过 $w$ 路」的那几个。这几个候选之间需要实测：四组 `chunk` 上 2 路都略高于 1 路，但差距从 0.9% 到 10% 不等，不构成一条可以照搬的规则。 -2. **改完 pad 要重扫 `chunk`。** `pad = 0` 那一列最优的是 `chunk = 16`（1.63），`pad = 4` 那一列最优的是 `chunk = 64`（3.70）。前一列的排序主要由冲突路数决定 —— `chunk = 16` 撞的是 8 路，`chunk = 64` 撞的是 32 路。冲突消掉之后四组都是 2 路，最优点换了位置。表里线程数与 `chunk` 联动（两者之积恒为行宽 4096），所以换位置的成因不止 `chunk` 一个，占用率与循环长度也跟着在变；能确定的只是改完 pad 之后 `chunk` 的排序会变。 +2. **修改 pad 之后需要重新扫描 `chunk`。** `pad = 0` 那一列最优的是 `chunk = 16`（1.63），`pad = 4` 那一列最优的是 `chunk = 64`（3.70）。前一列的排序主要由冲突路数决定：`chunk = 16` 的冲突是 8 路，`chunk = 64` 的冲突是 32 路。消除冲突之后四组都是 2 路，最优点随之改变。表中线程数与 `chunk` 联动（两者之积恒为行宽 4096），所以最优点改变的原因不只是 `chunk`，占用率与循环长度也在变化；能确定的只是修改 pad 之后 `chunk` 的排序会改变。 -3. **`chunk` 变了要重算 pad。** pad 写成固定字节数，等于假定 $\gcd(S, 32)$ 与 `chunk` 无关，而[上一节](#pad-per-chunk)那张表说明它不是：`chunk = 56`、$k = 1$ 时 $S$ 回到 32 个 word。把 pad 写成 `chunk` 的函数 —— 在 $k = 1, 2$ 里取 $\gcd(S, 32)$ 小的那个 —— 才和这一页的判据自洽。 +3. **`chunk` 改变后需要重新计算 pad。** 把 pad 写成固定字节数，等于假定 $\gcd(S, 32)$ 与 `chunk` 无关，而[pad 要按 chunk 算](#pad-per-chunk)一节的表说明两者相关：`chunk = 56`、$k = 1$ 时 $S$ 回到 32 个 word。把 pad 写成 `chunk` 的函数，即在 $k = 1, 2$ 中取 $\gcd(S, 32)$ 较小的那个，才与本页的判据一致。 -4. **这一页只在数据经过 shared memory 时适用。** 按[优化 global memory 访问](global-memory-access.md#coalescing)里的取舍，$V$ 小时用向量化的 blocked，数据直接进寄存器，不经 shared memory，没有 bank 冲突可言；$V$ 大到寄存器压力压低占用率时才改用 staged，以及整行需要被 block 内所有线程共享时，这一页才适用。 +4. **本页只适用于数据经过 shared memory 的情形。** 按照[优化 global memory 访问](global-memory-access.md#coalescing)中的取舍，$V$ 小时使用向量化的 blocked，数据直接进入寄存器，不经过 shared memory，因此不存在 bank 冲突。本页适用于两种情形：$V$ 大到寄存器压力压低占用率、改用 staged 时；整行需要由 block 内所有线程共享时。 -下面两段是 shared 缓冲的声明，差别只在 stride。反例，stride 恰好是 32 个 word 的倍数： +下面两段是 shared 缓冲的声明，两者只有 stride 不同。反例的 stride 恰好是 32 个 word 的倍数： ```python sh = T.alloc_shared((threads * chunk,), dtype) # stride = chunk 个元素 ``` -正例，pad 按 `chunk` 算，取 16 字节的整数倍里 $\gcd(S, 32)$ 最小的那个： +正例按 `chunk` 计算 pad，在 16 字节的整数倍中取 $\gcd(S, 32)$ 最小的那个： ```python import math diff --git a/docs/timing.md b/docs/timing.md index c9eb2767..b95cabb0 100644 --- a/docs/timing.md +++ b/docs/timing.md @@ -1,31 +1,31 @@ # How a benchmark is timed -The nightly benchmark measures one row per workload per op, reporting `device_busy_ms`: -the union of the execution intervals of every kernel that call produces. CUPTI records each -kernel's device-side start and end, an external correlation id attributes it to an -iteration, L2 is cleared before every iteration, and 25 ms of warm-up plus 100 ms of -measurement give the median. - -**So every number in the table is time the device spent executing kernels: none of the -host cost of issuing the call, and none of the gaps between kernels. For reading the -tables, that is the whole story.**{ .keystone } - -The rest is there when you need it: - -- [How one measurement runs](#how-one-measurement-runs) — the pseudocode, and the five - choices in it: calibration, the iteration count, clearing L2, attribution, failing - closed. -- [What is measured](#what-is-measured) — what `device_busy_ms` is, and why the gaps - between kernels are left out. -- [Why not wall-clock time](#why-not-wall-clock-time) — at decode sizes CUDA events - cannot measure a small kernel. -- [When to change how you measure](#when-to-change-how-you-measure) — only needed when - writing a benchmark yourself; it covers what this method cannot measure. +The nightly benchmark measures one row per workload per op and reports +`device_busy_ms`: the union of the execution intervals on the device of every kernel +that call produces. CUPTI records each kernel's device-side start and end, and an +external correlation id attributes the kernel to an iteration. L2 is cleared before +every iteration, warm-up runs for 25 ms and measurement for 100 ms, and the result is +the median. + +**Every number in the tables is therefore time the device spent executing kernels. It +excludes the host cost of issuing the call and the gaps between kernels. Reading the +tables requires nothing more than this.**{ .keystone } + +The remaining sections are for reference: + +- [How one measurement runs](#how-it-runs): the pseudocode, and its five choices: + calibration, the iteration count, clearing L2, attribution, and failing closed. +- [What is measured](#what-is-measured): the definition of `device_busy_ms`, and why + the gaps between kernels are left out. +- [Why not wall-clock time](#why-not-wall-clock): at decode sizes, CUDA events cannot + measure the execution time of a small kernel. +- [When to change how you measure](#when-to-change): needed only when writing a + benchmark, including the cases this method cannot measure. Every number below was measured on an H200, in the `tileops-runner:cu132-torch2.13` image. -## How one measurement runs +## How one measurement runs {#how-it-runs} ```python from benchmarks.timing import bench_kernel @@ -33,10 +33,10 @@ from benchmarks.timing import bench_kernel samples = bench_kernel(op, args=(x, weight)) # one Sample per iteration ``` -A benchmark rarely calls it directly, going through `ManifestBenchmark.profile()` or -`.compare()`, which take medians over these samples and compute the derived columns. +A benchmark rarely calls `bench_kernel` directly. It calls `ManifestBenchmark.profile()` +or `.compare()`, which take medians over these samples and compute the derived columns. -Inside, `bench_kernel` is three stages — collect, attribute, measure: +Inside, `bench_kernel` has three stages: collect, attribute and measure. ```python # Collect: each call runs under its own iteration number @@ -67,61 +67,66 @@ for i in range(n_repeat): n_kernels.append(len(claimed[i])) ``` -Five choices in it, each for a reason: +The five choices, and the reason for each: -1. **Calibrate.** Three calls estimate what one call costs. -2. **Convert that into an iteration count.** The budgets — 25 ms of warm-up, 100 ms of - measurement — divide by the per-call cost, clamped to `[10, 200]`: a short op gets many - samples, a long one need not run 200 times. +1. **Calibrate.** Three calls estimate the cost of one call. +2. **Convert that into an iteration count.** The budgets, 25 ms of warm-up and 100 ms of + measurement, are divided by the per-call cost, and the result is clamped to + `[10, 200]`. A short op therefore gets more samples, and a long one does not have to + run 200 times. 3. **Clear L2 before every iteration, and drain the device.** Without the clear, the first iteration reads from HBM and every later one from L2, so the median reports the - best case of a full cache hit. Draining keeps the previous iteration out of this one. -4. **Collect and attribute.** The iteration number goes onto CUPTI's external correlation - id stack, and the correlation id a kernel record carries maps back to it. - **Attribution does not look at timestamps** — which iteration a kernel belongs to is + best case of a full cache hit. Draining keeps the previous iteration from overlapping + this one. +4. **Collect and attribute.** Each iteration pushes its iteration number as CUPTI's + external correlation id, so every launch issued inside it carries that id; the + correlation id in a kernel record maps back to the iteration number. + **Attribution does not use timestamps**: which iteration a kernel belongs to is written in its record, independent of when it ran. A kernel shorter than the host - overhead is attributed as reliably as a long one, and a call whose kernel count varies - between iterations still measures. -5. **Fail closed.** Three attribution failures raise three different errors and produce + overhead is therefore attributed as reliably as a long one, and a call whose kernel + count varies between iterations can still be measured. +5. **Fail closed.** Three attribution failures each raise a different error and produce no number: | Case | Raises | Meaning | | --- | --- | --- | -| CUPTI discarded records | `_CUPTIRecordsLostError` | the reading is gone though the iteration did run — the whole phase is measured again, up to 3 attempts in all, asking for a 4× larger buffer each time | -| Nothing discarded, but a kernel carries no iteration number | `_OffThreadLaunchError` | a thread that never pushed an id launched it | -| Nothing discarded, and one iteration has no kernels at all | `_CUPTIAttributionError` | that call never reached the device | +| CUPTI discarded records | `_CUPTIRecordsLostError` | the iteration did run, but the reading is lost; the whole phase is measured again, up to 3 attempts in all, with a 4× larger buffer each time | +| Nothing discarded, but a kernel carries no iteration number | `_OffThreadLaunchError` | the kernel was launched by a thread that never pushed an iteration number | +| Nothing discarded, and one iteration has no kernels at all | `_CUPTIAttributionError` | that call never ran on the device | -## What is measured +## What is measured {#what-is-measured} -**`device_busy_ms`: the union of the execution intervals of every kernel one call -produces.** A CUPTI kernel record gives the device-side execution bounds, with none of the -host cost of issuing the call. Three cases: +**`device_busy_ms` is the length of the union of the execution intervals on the device +of every kernel one call produces.** A CUPTI kernel record gives the device-side start +and end of execution, excluding the host cost of issuing the call. There are three +cases: -- **A single-kernel call** — the kernel's execution time on the device. -- **A multi-kernel call** — the union of the intervals: the total time at least one of - the call's kernels was executing. Two kernels running concurrently are not counted - twice; that would be SM time, not time the device was busy. -- **The gaps between kernels** — not counted. +- **A single-kernel call**: the kernel's execution time on the device. +- **A multi-kernel call**: the union of the intervals, that is, the total time during + which at least one of the call's kernels was executing. Two concurrent kernels are not + counted twice, because the sum of the two would be SM time, not the time the device was + busy. +- **The gaps between kernels**: not counted. -A gap is left out because it cannot be attributed: the device really was idle, but the -cause is either the op's own data dependency or the CPU not having issued the next kernel -yet, and CUPTI's records do not distinguish the two. A quantity whose cause is unknown -cannot judge an implementation. +A gap is left out because its cause cannot be determined. The device really was idle, +but the cause is either the op's own data dependency or the CPU not having issued the +next kernel yet, and CUPTI's records do not distinguish the two. A quantity whose cause +cannot be determined cannot be used to judge an implementation. -`tflops` and `bandwidth_tbs` divide by the same quantity: they describe the throughput -reached while the device was executing, and a denominator that included in-call idleness -would depress them systematically. +`tflops` and `bandwidth_tbs` divide by the same quantity. They describe the throughput +reached while the device was executing, and a denominator that included idle time +within the call would lower them systematically. -Defined this way, the number is immune to how fast the host is. Changing CUPTI's -collection buffer from 256 KB to 32 MB takes the median `latency_ms` of one three-kernel -call from 35 us to 2068 us while `device_busy_ms` stays at 19.1 us — a late host does not -change any kernel's execution time, it only pushes them apart on the timeline, and the -union is the same. +Defined this way, the number does not depend on how fast the host is. Changing CUPTI's +collection buffer from 256 KB to 32 MB raises the median `latency_ms` of one +three-kernel call from 35 us to 2068 us, while `device_busy_ms` stays at 19.1 us. A host +that issues kernels late does not change any kernel's execution time; it only spreads +the kernels apart on the timeline, and the length of the union is unchanged. -## Why not wall-clock time +## Why not wall-clock time {#why-not-wall-clock} -At decode sizes an op can finish faster than the Python call that launched it. Four -methods, one 3 us kernel, four numbers: +At decode sizes, an op can finish faster than the Python call that launched it. Four +methods applied to one 3 us kernel give four readings: | Method | Reading | | --- | --- | @@ -130,41 +135,43 @@ methods, one 3 us kernel, four numbers: | One pair of events around the loop, divided by the iteration count | 6.07 us | | CUDA graph replay | 4.30 us | -The device executed for 1.95 us; the 6 us the event methods read is the rate at which the -CPU issues the next call, not the kernel's execution time. **That is the only reason -TileOPs times with CUPTI**, and it is why a row that fell back to CUDA events cannot be -compared with the others: there `device_busy_ms` and `latency_ms` carry the same number, -and the `timing` field records `cuda-events`. +The device executed for 1.95 us. The 6 us the event methods read is the interval at +which the CPU issues the next call, not the kernel's execution time. **This is the only +reason TileOPs times with CUPTI**, and it is why a row that fell back to CUDA events +cannot be compared with the other rows: in that row `device_busy_ms` and `latency_ms` +hold the same number, and the `timing` field records `cuda-events`. ## Comparing several implementations -Comparing implementations within one case, `compare()` times each twice, in the order -A B C C B A, and takes the median over both passes. +To compare implementations within one case, `compare()` times each implementation +twice, in the order A B C C B A, and takes the median over the samples of both passes. -In a fixed order the implementation that ran first and the one that ran last sit at -different clocks and temperatures, and that difference reads as a difference between the -implementations. A symmetric order puts each implementation's two passes in the first and -second half of the case, cancelling monotonic drift to first order. Two details: +In a fixed order, the implementation that runs first and the one that runs last see +different clocks and temperatures, and that difference reads as a difference between +the implementations. A symmetric order puts each implementation's two passes in the +first and second half of the case, which cancels monotonic drift to first order. Two +details: - **The budget is split, not doubled.** Each pass gets 12.5 ms of warm-up and 50 ms of - measurement, with half the iteration bounds: the point is symmetry, not more samples, - and the sample count matches timing one implementation. -- **Both passes must use the same timing method.** One pass on CUPTI and the other - fallen back to CUDA events raises rather than pooling, which would put one median over - two kinds of measurement. + measurement, with half the iteration bounds. The symmetric order is meant to cancel + drift, not to add samples, so the sample count matches that of timing one + implementation. +- **Both passes must use the same timing method.** When one pass uses CUPTI and the + other falls back to CUDA events, `compare()` raises instead of pooling the results, + since pooling would put two kinds of measurement into one median. -## When to change how you measure +## When to change how you measure {#when-to-change} -The default case needs none of this: one kernel per call, through the Op interface, timed -by `bench_kernel`, no other thread using the GPU — where most ops are today. Seven cases -call for a stop: +The default case needs no change: one kernel per call, called through the Op +interface, timed by `bench_kernel`, with no other thread using the GPU. Most ops are in +this case today. Seven cases need separate handling: | Your case | If you ignore it | What to do | | --- | --- | --- | | The timed closure contains `Tensor.backward` or `torch.autograd.grad` | the backward kernels come from the autograd engine's own thread, carry no iteration number, and the case raises instead of producing a figure | Drive a single fused node with `backward_of(out)`; for a chain, set `torch.autograd.set_multithreading_enabled(False)` | | Another thread in the process uses the GPU, or the timed closure uses CUPTI's `CUSTOM0` external id | those kernels carry no iteration number, or the closure overwrites the one the timer set, and it raises either way | Have the timed call launch its own work; use `CUSTOM1` / `CUSTOM2` instead | -| The op produces its result through `copy_` — in-place elementwise, MoE's write-back | the timer collects the copy but leaves it out of `device_busy_ms` and reports it as `uncounted_copy_ms`, so the reading is short | Pass `count_copies=True` for the case; every tag's reading then includes the copies | -| One call launches several kernels | the gaps between kernels land in `latency_ms`, so comparing by it against a fused implementation charges them to your side | Conclude from `device_busy_ms` only; `latency_ms` compares between rows of equal `n_kernels` | -| One call takes more than 10 ms | the iteration count hits the floor of 10, the wall-clock far exceeds the 100 ms budget, and p10/p90 over 10 samples are coarse | Accept the longer wall-clock, or state an iteration count and the sample size | -| You want a kernel-level benchmark | the op has no spec, so shapes and roofline have to be written by hand and the spec validator cannot see them | Measure through the Op interface and write a [spec](manifest.md) | -| You are adding an external baseline | moving the baseline's input conversion out of its timed region has this repository carry that time instead | Keep the conversion inside the baseline's timed region, and let the import fail where that baseline is the point | +| The op produces its result through `copy_`, as in-place elementwise ops and MoE's write-back do | the timer collects the copy but by default leaves it out of `device_busy_ms` and reports it as `uncounted_copy_ms`, so the reading is too low | Pass `count_copies=True` for the case; every tag's reading then includes the copies | +| One call launches several kernels | the gaps between kernels land in `latency_ms`, so a comparison by it against a fused implementation charges the gaps to the multi-kernel side | Draw conclusions from `device_busy_ms` only; `latency_ms` is comparable only between rows with equal `n_kernels` | +| One call takes more than 10 ms | the iteration count hits the floor of 10, the wall-clock time far exceeds the 100 ms budget, and p10/p90 over 10 samples are coarse | Accept the longer wall-clock time, or state an iteration count and the sample size | +| You want a kernel-level benchmark | the op has no spec, so shapes and roofline have to be written by hand and the spec validator cannot see them | Measure through the Op interface and write a [spec](user-guide/manifest/index.md) | +| You are adding an external baseline | moving the baseline's input conversion out of its timed region makes this repository carry that time instead | Keep the conversion inside the baseline's timed region; where that baseline is the reason the benchmark exists, require its dependency and let the import fail when it is missing | diff --git a/docs/timing.zh.md b/docs/timing.zh.md index dffa3d8b..a919763b 100644 --- a/docs/timing.zh.md +++ b/docs/timing.zh.md @@ -1,15 +1,15 @@ # benchmark 的计时方法 -nightly benchmark 给每个算子的每个 workload 各测一行，报出的 `device_busy_ms` 是这次调用产生的全部 kernel 在设备上执行区间的并集。CUPTI 记下每个 kernel 的执行起止，按 external correlation id 归到某次迭代；每次迭代之前清空 L2，25 ms 预热、100 ms 测量，取中位数。 +nightly benchmark 为每个 op 的每个 workload 各测一行，报出的 `device_busy_ms` 是这次调用产生的全部 kernel 在设备上执行区间的并集。CUPTI 记录每个 kernel 的执行起止时间，并按 external correlation id 把它归到某一次迭代。每次迭代之前清空 L2，预热 25 ms、测量 100 ms，结果取中位数。 -**所以表里的每个数都是设备执行 kernel 的时间：不含 CPU 发起调用的开销，也不含 kernel 之间的空隙。读表到这里就够了。**{ .keystone } +**因此表中的每个数都是设备执行 kernel 的时间，不含 CPU 发起调用的开销，也不含 kernel 之间的空隙。阅读表格只需要知道这一点。**{ .keystone } -余下各节按需要查： +其余各节供需要时查阅： -- [一次测量的流程](#how-it-runs) —— 伪代码，以及校准、迭代次数、清 L2、归属、失败即停五处安排。 -- [被测量的量](#what-is-measured) —— `device_busy_ms` 的定义，以及 kernel 之间的空隙为什么不计入。 -- [为什么不是墙钟时间](#why-not-wall-clock) —— decode 尺度上 CUDA events 测不出小 kernel。 -- [什么时候要改测法](#when-to-change) —— 自己写 benchmark 时才用得上，含这套测法测不到的几种情形。 +- [一次测量的流程](#how-it-runs)：伪代码，以及校准、迭代次数、清空 L2、归属与失败即停五项安排。 +- [被测量的量](#what-is-measured)：`device_busy_ms` 的定义，以及 kernel 之间的空隙不计入的原因。 +- [为什么不是墙钟时间](#why-not-wall-clock)：在 decode 尺度上，CUDA events 测不出小 kernel 的执行时间。 +- [什么时候要改测法](#when-to-change)：只在自行编写 benchmark 时用到，包括这套测法测不到的几种情形。 下文的数字都在 H200 上实测，镜像为 `tileops-runner:cu132-torch2.13`。 @@ -21,9 +21,9 @@ from benchmarks.timing import bench_kernel samples = bench_kernel(op, args=(x, weight)) # 每次迭代一个 Sample ``` -写 benchmark 一般不直接调它，而是走 `ManifestBenchmark.profile()` 或 `.compare()`：那一层在这些样本上取中位数、算派生列。 +编写 benchmark 时一般不直接调用 `bench_kernel`，而是调用 `ManifestBenchmark.profile()` 或 `.compare()`。这一层在样本上取中位数，并计算派生列。 -`bench_kernel` 内部分三段 —— 采集、归属、计量： +`bench_kernel` 内部分为采集、归属与计量三段： ```python # 采集：每次调用在自己的迭代号下执行 @@ -54,37 +54,37 @@ for i in range(n_repeat): n_kernels.append(len(claimed[i])) ``` -五处安排各有原因： +五项安排及其原因如下： -1. **校准。** 先跑 3 次，估出单次调用的耗时。 -2. **换算迭代次数。** 预热 25 ms、测量 100 ms 的预算除以单次耗时，钳在 `[10, 200]`：短算子采样多，长算子不必跑满 200 次。 -3. **每次迭代之前清 L2，并等设备排空。** 不清 L2，第一次迭代从 HBM 读、之后都从 L2 读，中位数报的就是缓存全命中的最好情况；排空是为了上一次迭代不与这一次重叠。 -4. **采集与归属。** 每次迭代把迭代号标记为 CUPTI 的 external correlation id，区间内发出的每次 launch 都带上它，kernel 记录里的 correlation id 再经这层映射回到迭代号。**归属不看时间戳** —— 一个 kernel 属于哪次迭代写在记录里，与它何时执行无关；所以比主机开销还短的 kernel 归属得一样可靠，一次调用的 kernel 数在迭代之间变化也测得出来。 -5. **失败即停。** 三种归属失败各报一种错，不产出数字。 +1. **校准。** 先运行 3 次，估计单次调用的耗时。 +2. **换算迭代次数。** 用预热 25 ms、测量 100 ms 的预算除以单次耗时，结果限制在 `[10, 200]` 之内。耗时短的 op 因此采样更多，耗时长的 op 不必运行满 200 次。 +3. **每次迭代之前清空 L2，并等待设备排空。** 不清空 L2 时，第一次迭代从 HBM 读取，之后的迭代都从 L2 读取，中位数反映的是缓存全部命中的最好情况。等待排空使上一次迭代不与这一次重叠。 +4. **采集与归属。** 每次迭代把迭代号标记为 CUPTI 的 external correlation id，区间内发出的每次 launch 都带上这个 id；kernel 记录中的 correlation id 再经这层映射对应回迭代号。**归属不依据时间戳**：一个 kernel 属于哪次迭代写在记录里，与它何时执行无关。因此比主机开销还短的 kernel 同样能被可靠归属，一次调用的 kernel 数在迭代之间变化时也能测出。 +5. **失败即停。** 三种归属失败各报一种错误，不产出数字。 | 情形 | 报什么 | 含义 | | --- | --- | --- | -| CUPTI 丢了记录 | `_CUPTIRecordsLostError` | 那次迭代跑过，只是读数丢了 —— 整个 phase 重测，总共最多测 3 次，每次把缓冲要大 4 倍 | -| 没有丢弃，但有 kernel 带不上迭代号 | `_OffThreadLaunchError` | 某个没有标记迭代号的线程发起了它 | -| 没有丢弃，某次迭代一个 kernel 都没有 | `_CUPTIAttributionError` | 这次调用没上设备 | +| CUPTI 丢了记录 | `_CUPTIRecordsLostError` | 那次迭代已经执行，只是读数丢失。整个 phase 重新测量，总共最多测 3 次，每次把缓冲扩大 4 倍 | +| 没有丢弃，但有 kernel 无法对应到迭代号 | `_OffThreadLaunchError` | 这个 kernel 由一个没有标记迭代号的线程发起 | +| 没有丢弃，但某次迭代没有任何 kernel | `_CUPTIAttributionError` | 这次调用没有在设备上执行 | ## 被测量的量 {#what-is-measured} -**`device_busy_ms`：一次调用产生的全部 kernel，在设备上执行区间的并集长度。** CUPTI 的 kernel 记录给出设备上的执行起止，不含 CPU 发起这次调用的开销。三种情形： +**`device_busy_ms` 是一次调用产生的全部 kernel 在设备上执行区间的并集长度。** CUPTI 的 kernel 记录给出设备上的执行起止时间，不含 CPU 发起这次调用的开销。分三种情形： -- **单 kernel 的调用** —— 就是这个 kernel 在设备上的执行时长。 -- **多 kernel 的调用** —— 各区间的并集，即设备上至少有一个该调用的 kernel 在执行的总时长。两个 kernel 并发不计成两份，那是 SM 时间，不是设备忙的时间。 -- **kernel 之间的空隙** —— 不计入。 +- **单 kernel 的调用**：这个 kernel 在设备上的执行时长。 +- **多 kernel 的调用**：各执行区间的并集，即设备上至少有一个属于该调用的 kernel 在执行的总时长。两个并发的 kernel 不计为两份，因为两份之和是 SM 时间，不是设备忙碌的时间。 +- **kernel 之间的空隙**：不计入。 -空隙不计入，是因为它归不了因：设备那段时间确实空闲，但成因可能是算子自身的数据依赖，也可能是 CPU 还没发出下一个 kernel，两者在 CUPTI 的记录里没有区别。分不清成因的量，不能用来判断一个实现的好坏。 +空隙不计入，是因为无法确定它的成因。设备在那段时间确实空闲，但成因可能是 op 自身的数据依赖，也可能是 CPU 尚未发出下一个 kernel，两者在 CUPTI 的记录中没有区别。成因无法区分的量不能用来判断一个实现的好坏。 -`tflops` 与 `bandwidth_tbs` 的分母也是这个量：它们描述设备执行期间达到的吞吐，分母含调用内的空闲会把它系统性压低。 +`tflops` 与 `bandwidth_tbs` 的分母也是这个量。它们描述设备执行期间达到的吞吐，分母若包含调用内的空闲时间，会把吞吐系统性地压低。 -这样定义的量对主机的快慢免疫。CUPTI 的采集缓冲从 256 KB 换成 32 MB，同一个三 kernel 调用的 `latency_ms` 中位数从 35 us 涨到 2068 us，`device_busy_ms` 始终 19.1 us —— 主机晚发不改变任何 kernel 的执行时长，只是把它们在时间轴上推远，并集不变。 +这样定义的量不受主机快慢的影响。把 CUPTI 的采集缓冲从 256 KB 换成 32 MB 后，同一个三 kernel 调用的 `latency_ms` 中位数从 35 us 增加到 2068 us，`device_busy_ms` 始终是 19.1 us。主机晚发出 kernel 不改变任何 kernel 的执行时长，只是把它们在时间轴上推远，并集长度不变。 ## 为什么不是墙钟时间 {#why-not-wall-clock} -decode 尺度上，算子的执行时间可能短于发起它的那次 Python 调用。同一个 3 us 的 kernel，四种测法读出四个数： +在 decode 尺度上，op 的执行时间可能短于发起它的那次 Python 调用。对同一个 3 us 的 kernel，四种测法得到四个读数： | 测法 | 读数 | | --- | --- | @@ -93,27 +93,27 @@ decode 尺度上，算子的执行时间可能短于发起它的那次 Python | 整个循环一对 events，再除以迭代数 | 6.07 us | | CUDA graph 重放 | 4.30 us | -设备上实际执行 1.95 us，event 方案读出的 6 us 是 CPU 发起下一次调用的节奏，不是 kernel 的执行时间。**这是 TileOPs 用 CUPTI 计时的唯一理由**，也是为什么退回 CUDA events 之后那一行不能与其余行比较：`device_busy_ms` 与 `latency_ms` 记同一个数，`timing` 字段记为 `cuda-events`。 +设备上的实际执行时间是 1.95 us，event 方案读出的 6 us 是 CPU 发起下一次调用的间隔。**这是 TileOPs 用 CUPTI 计时的唯一理由**，也是退回 CUDA events 的那一行不能与其余行比较的原因：这一行的 `device_busy_ms` 与 `latency_ms` 记录同一个数，`timing` 字段记为 `cuda-events`。 ## 比较多个实现 -同一个用例里比较几个实现，`compare()` 按 A B C C B A 各跑两段，两段样本合并后取中位数。 +在同一个用例中比较多个实现时，`compare()` 按 A B C C B A 的顺序让每个实现各运行两段，两段样本合并后取中位数。 -固定顺序下先跑和后跑的实现处在不同的时钟与温度状态，这个差别会被读成实现之间的差别；对称顺序让每个实现的两段分列全程的前后两半，单调漂移一阶抵消。两个细节： +在固定顺序下，先运行和后运行的实现处在不同的时钟与温度状态，这个差别会被误读为实现之间的差别。对称顺序让每个实现的两段分别位于全程的前半和后半，单调漂移在一阶上相互抵消。有两个细节： -- **预算是拆开的，不是翻倍。** 每段 12.5 ms 预热、50 ms 测量，迭代上下限各取一半 —— 要的是对称，不是更多样本，样本量与单实现计时相当。 -- **两段的计时方法必须一致。** 一段走 CUPTI、另一段退回 CUDA events 时直接报错，不合并，否则一个中位数横跨两种测量。 +- **预算拆分，不翻倍。** 每段预热 12.5 ms、测量 50 ms，迭代次数的上下限各取一半。对称顺序的目的是抵消漂移，不是增加样本，因此样本量与单个实现计时时相当。 +- **两段的计时方法必须一致。** 一段使用 CUPTI、另一段退回 CUDA events 时直接报错，不合并结果，否则一个中位数会混合两种测量方法。 ## 什么时候要改测法 {#when-to-change} -默认情形什么都不用管：一次调用只发一个 kernel、经由 Op 接口、走 `bench_kernel`、进程里没有别的线程用 GPU —— 当前多数算子都是这样。下面七种要停下来处理： +默认情形下不需要任何处理：一次调用只发出一个 kernel，经由 Op 接口，使用 `bench_kernel`，并且进程中没有其他线程使用 GPU。当前多数 op 属于这种情形。以下七种情形需要单独处理： -| 情形 | 不处理会怎样 | 该怎么做 | +| 情形 | 不处理的后果 | 处理方法 | | --- | --- | --- | -| 被测闭包里包了 `Tensor.backward` 或 `torch.autograd.grad` | 反向的 kernel 由 autograd 引擎的线程发出，带不上迭代号，整个用例报错，不产出数字 | 单个融合节点用 `backward_of(out)` 直接驱动；多节点链改用 `torch.autograd.set_multithreading_enabled(False)` | -| 进程里有别的线程在用 GPU，或被测闭包自己用了 CUPTI 的 `CUSTOM0` external id | 那些 kernel 带不上迭代号，或者迭代号被闭包盖掉，同样报错 | 让被计时的调用自己启动它的工作；external id 改用 `CUSTOM1` / `CUSTOM2` | -| 算子靠 `copy_` 回写才产出结果，例如原地 elementwise 与 MoE 的写回 | 计时会采集这次拷贝，但默认不计入 `device_busy_ms`，另记在 `uncounted_copy_ms`，读数因此偏小 | 给这个用例传 `count_copies=True`，所有 tag 的读数都会计入拷贝 | -| 一次调用发多个 kernel | kernel 之间的空隙落在 `latency_ms` 里，用它与融合实现比较，空隙算在多 kernel 这一方 | 结论只用 `device_busy_ms`；`latency_ms` 仅在两行 `n_kernels` 相同时可比 | -| 单次调用超过 10 ms | 迭代次数被下限 10 顶住，墙钟远超 100 ms 的预算，10 个样本给出的 p10/p90 很粗 | 接受更长的墙钟，或显式指定迭代次数并写明样本量 | -| 想加一个 kernel 级的 benchmark | 这个算子没有 spec，形状与 roofline 只能手写，spec 校验器也查不到它 | 经 Op 接口测，并补一份 [spec](manifest.md) | -| 加一个外部基线 | 基线的输入转换若移出它的计时区间，等于本仓库替它承担了这部分时间 | 转换留在基线的计时区间内。这个基线是 benchmark 存在的理由时，要求依赖存在、让 import 失败 | +| 被测闭包中包含 `Tensor.backward` 或 `torch.autograd.grad` | 反向的 kernel 由 autograd 引擎的线程发出，无法对应到迭代号，整个用例报错，不产出数字 | 单个融合节点用 `backward_of(out)` 直接驱动；多节点链改用 `torch.autograd.set_multithreading_enabled(False)` | +| 进程中有其他线程在使用 GPU，或被测闭包自身使用了 CUPTI 的 `CUSTOM0` external id | 那些 kernel 无法对应到迭代号，或迭代号被闭包覆盖，同样报错 | 由被计时的调用自己启动它的工作；external id 改用 `CUSTOM1` / `CUSTOM2` | +| op 依靠 `copy_` 回写才产出结果，例如原地 elementwise 与 MoE 的写回 | 计时会采集这次拷贝，但默认不计入 `device_busy_ms`，而是另记在 `uncounted_copy_ms` 中，读数因此偏小 | 为这个用例传入 `count_copies=True`，所有 tag 的读数都会计入拷贝 | +| 一次调用发出多个 kernel | kernel 之间的空隙计入 `latency_ms`，用它与融合实现比较时，空隙算在多 kernel 的一方 | 结论只依据 `device_busy_ms`；`latency_ms` 只在两行的 `n_kernels` 相同时可比 | +| 单次调用超过 10 ms | 迭代次数停在下限 10，墙钟时间远超 100 ms 的预算，10 个样本给出的 p10/p90 很粗 | 接受更长的墙钟时间，或显式指定迭代次数并写明样本量 | +| 需要新增一个 kernel 级的 benchmark | 这个 op 没有 spec，形状与 roofline 只能手写，spec 校验器也检查不到它 | 经由 Op 接口测量，并补一份 [spec](user-guide/manifest/index.md) | +| 新增一个外部基线 | 如果基线的输入转换被移出它的计时区间，相当于本仓库替它承担了这部分时间 | 转换保留在基线的计时区间内。这个基线是该 benchmark 存在的理由时，要求依赖必须存在，缺失时让 import 失败 | diff --git a/docs/torch-compile.md b/docs/torch-compile.md index 61dbd626..1db1bf27 100644 --- a/docs/torch-compile.md +++ b/docs/torch-compile.md @@ -1,23 +1,31 @@ # Bringing an op into torch.compile -Bringing a TileOPs op into `torch.compile` makes it one node in the user's compiled -graph, and that node's shape does not change with the backend serving it. +A TileOPs op brought into `torch.compile` becomes one node in the user's compiled +graph, and that node does not change with the backend serving it. -It takes one thing: a compile boundary declared at the op layer, with everything outside -it traced by dynamo and everything inside it invisible to the compiler. +Bringing an op in takes one piece of work: declaring a compile boundary at the op +layer. Dynamo traces everything outside the boundary, and everything inside it is +invisible to the compiler. -The body covers the work — checking whether an op is already in, compiling a call, the -five conventions, and the declaration an op that is not in yet needs. The appendix covers -why the boundary can only be drawn there: how dynamo works, where it and the op layer -disagree, and what the boundary costs. +The body covers the work of bringing an op in: + +1. checking whether an op is already in; +1. compiling code that calls it; +1. the five calling conventions; +1. the code an op that is not in yet needs to declare the boundary. + +The appendix explains why the boundary can only be drawn this way: how dynamo works, +where it and the op layer disagree, why the boundary sits at the op layer, and what the +boundary costs and does not provide. ## Calling an op that is already in ### Checking whether an op is in {#supported} -Read the class attribute `compile_op_names`: non-empty means the class declares a -compile boundary (`compile_boundary = True`), so the boundary is at the op layer and -`fullgraph=True` works; an empty tuple means it declares none. +Read the class attribute `compile_op_names`. A non-empty value means the class +declares a compile boundary (`compile_boundary = True`), so the boundary is at the op +layer and `fullgraph=True` works. An empty tuple means the class declares no compile +boundary. ```python >>> from tileops.norm import RMSNormFwdOp @@ -25,13 +33,13 @@ compile boundary (`compile_boundary = True`), so the boundary is at the op layer ('tileops::norm_rms_norm_fwd',) ``` -An op that has not migrated raises under `fullgraph=True`, and breaks the graph under +An op that has not migrated raises under `fullgraph=True` and breaks the graph under the default settings. -### Compiling a call +### Compiling code that calls it -Construct the op instance, hand the function that calls it to `torch.compile`, and -that is all: +Construct the op instance and pass the function that calls it to `torch.compile`. No +other step is needed: ```python import torch @@ -48,53 +56,55 @@ w = torch.randn(4096, device="cuda", dtype=torch.float16) block(x, w) ``` -`TORCH_LOGS=graph_code` prints the captured graph: one node, -`tileops::norm_rms_norm_fwd`, not the calls inside the kernel. +Running with `TORCH_LOGS=graph_code` prints the captured graph. It has one node, +`tileops::norm_rms_norm_fwd`, and the calls inside the kernel do not appear in it. ### The five calling conventions -Each follows from a mechanism at the boundary, and breaking any of them makes the -compiled path behave differently from the eager one. +Each convention follows from a mechanism at the boundary. Breaking any of them makes +the compiled path behave differently from the eager one. - **Construct the op instance once and reuse it.** The instance key is a compile-time - constant and each instance is its own compiled graph, so constructing one inside a loop - recompiles every iteration. -- **Do not rely on strides passing through.** A non-contiguous input the op does not - write is made contiguous inside the node, and an output the op allocates is always - contiguous; convert outside the op if later work needs another layout. An output that - is a written input (`alias`) or a caller's `out` keeps that tensor's storage. -- **Meta tensors cannot warm anything up.** Once an op has a boundary, a call with - meta or fake tensors returns at the fake and never reaches kernel construction. -- **Warm up before a CUDA-graph capture.** Call once with real tensors at the same - shape: building a kernel may compile, while capture allows only a memo hit and the call - that follows it. See [what each phase may do](backends.md#phase-limits). + constant and each instance has its own compiled graph, so constructing an instance + inside a loop recompiles on every iteration. +- **Strides are not passed through.** A non-contiguous input the op does not write is + made contiguous inside the node, and an output the op allocates is always contiguous. + When later work needs another layout, convert outside the op. An output that is a + written input (`alias`) or a caller's `out` keeps that tensor's storage. +- **Meta tensors cannot be used for warm-up.** Once an op declares the boundary, a call + with meta or fake tensors returns at the fake and never reaches kernel construction. +- **Warm up before a CUDA graph capture.** Call the op at least once with real tensors + at the same shape: building a kernel may compile, while a capture allows only a cache + hit followed by the call. What each phase allows is in + [what each phase may do](backends.md#phase-limits). - **A second card may need its own build.** For a call a target serves, the device is - part of the kernel's memo key, so the same instance builds again on a second card. An - in-tree kernel is keyed by the build identity the selected implementation's `entry_for` - returns, which includes the device only when the build depends on it. A `target=` named in the constructor - is honoured on the first compiled call too, and a failed build pins the op to no - target. + part of the kernel's cache key, so the same instance builds the kernel again on a + second card. The cache key of an in-tree kernel is the build identity the selected + implementation's `entry_for` returns, which includes the device only when the build + depends on it. A `target=` given to the constructor also takes effect on the first + compiled call, and a failed build pins the op to no target. ### The three guarantees once an op is in -With the boundary at the op layer, a caller can rely on three things. +With the boundary at the op layer, a caller can rely on three things: - **The graph does not change with the target.** The same code compiles to the same - graph on another backend or another card, so the artefact is independent of who serves - the op. -- **`fullgraph=True` works**, for an op that declares this contract — see [checking - whether an op is in](#supported). -- **Output shape, dtype and stride come from the manifest**, not from how a kernel - tiles or pads internally. An output the op allocates is always contiguous. + graph on another backend or another card, so the compiled artefact does not depend on + the backend. +- **`fullgraph=True` works** for an op that declares this contract; see + [Checking whether an op is in](#supported). +- **Output shape, dtype and stride come from the manifest.** They do not depend on how a + kernel tiles or pads internally. An output the op allocates is always contiguous. ## Declaring the boundary on a new op: `RMSNormFwdOp` -The code bringing one op in takes: how the boundary is declared, how the fake is -written, and why the target is resolved again inside the node. For the tracing, graph -breaks and guards below, see [how dynamo works](#how-dynamo-works). +This section gives the code needed to bring one op in: how the boundary is declared, +how the fake is written, and why the target is resolved again inside the node. The +tracing, graph breaks and guards it refers to are described in +[How dynamo works](#dynamo). -`RMSNormFwdOp` was the first op brought in. Its skeleton, docstrings elided — the full -file is +`RMSNormFwdOp` was the first op brought in. Its skeleton, with docstrings elided; the +full file is [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py): ```python @@ -121,15 +131,19 @@ class RMSNormFwdOp(Op): ``` That is the whole declaration. The operators and their fakes are generated from the -entry, one operator per effect branch: its tensor arguments are `signature.inputs` in -order, what it returns is `signature.outputs`, the arguments it writes are exactly the -inputs marked `mutated`, and each output's shape and dtype come from the signature. Its -name is `tileops::_`, with the family written once where the class -name already opens with it — here `tileops::norm_rms_norm_fwd`; when a branch writes -an input, fills a `buffer: out` or omits an output, its operator's name gains -`_writes_`, `_out` and `_without_` respectively, in that order. -No op chooses its own name, so `compile_op_names` cannot disagree with what was -registered. +manifest entry, one operator per effect branch: + +- its tensor arguments are `signature.inputs`, in order; +- its return value is given by `signature.outputs`; +- the arguments it writes are exactly the inputs marked `mutated`; +- each output's shape and dtype come from the signature. + +The operator's name is `tileops::_`, with the family written only +once when the class name already starts with it; here the name is +`tileops::norm_rms_norm_fwd`. When a branch writes an input, fills a `buffer: out` or +omits an output, its operator's name gains `_writes_`, `_out` and +`_without_` respectively, in that order. No op names its own operator, so +`compile_op_names` cannot disagree with the registered names. The layers one call passes through, and where the boundary falls: @@ -137,68 +151,71 @@ The layers one call passes through, and where the boundary falls:
Op.__call__calls forward, resolves no target
forwardone line, calls the opaque operator
compile boundary
-
the generated operatorrecovers the instance, runs the generated checks, resolves the target, unsettles on failure
-
_eager_forwardcontiguous, kernel, launch
-
The two violet layers are inside dynamo's trace, and that one line of forward is the last thing it reaches; below the boundary the opaque operator takes over, invisible to the compiler.
+
the generated operatorrecovers the instance, runs the generated checks, resolves the target, undoes the resolution on failure
+
_eager_forwardmakes the inputs contiguous, fetches the kernel, launches it
+
The two violet layers are inside dynamo's trace, and the one line of forward is the last thing dynamo traces. Below the boundary the opaque operator runs, invisible to the compiler.
-Three things in it are not free choices. - -**First, the instance is recovered through a string key rather than passed directly.** -The schema's types are a fixed set — `Tensor`, `int`, `float`, `bool`, `str` and a few -more — with no "arbitrary Python object", while what the operator body needs -(`kernel_map`, the settled target, the memo table of built kernels) hangs off the instance -and does not fit a schema argument. Two details of the key are not free either: - -- **A string, not an integer.** A string is a compile-time constant during tracing, - where an integer is generalised to a `SymInt`. -- **Never reused.** Because it is constant, inductor bakes the shape the fake gave - into the artefact, and an op reusing a key would inherit the previous instance's - shape. - -**Second, the fake builds its result with `torch.empty` from the shape and dtype the -signature check infers, not `torch.empty_like(x)`.** What the fake returns has to match real execution in shape, dtype -and stride; a mismatch either fails during tracing or — for a stride — has downstream code -read the wrong layout and go silently wrong. The operator body makes the inputs contiguous -before the kernel writes into a freshly allocated output, so the real output is always -contiguous, while `empty_like` copies the input's strides: a non-contiguous input would -have the fake declare a layout real execution never produces. - -**Third, the target is resolved inside the node, not in `Op.__call__`.** When traced code runs `self.x = ...`, dynamo records a pending -side effect and applies it only after the whole graph has run, while the opaque node runs -before that: a resolution written just outside the node is unreadable inside it. Two -things follow, both inside the node: - -- A resolution made outside the node would have the first compiled call silently run +Three parts of this code are fixed. + +**First, the instance is recovered through a string key, not passed as an object.** +The schema's types are a fixed set, such as `Tensor`, `int`, `float`, `bool` and +`str`, with no "arbitrary Python object". What the operator body needs (`kernel_map`, +the resolved target, and the cache table of built kernels) is stored on the instance and +cannot be split into schema arguments. Two details of the key are also fixed: + +- **The key is a string, not an integer.** A string is a constant during tracing, while + an integer is generalised to a `SymInt`. +- **A key is never reused.** Because the key is a constant, inductor bakes the shape the + fake gave into the compiled artefact, and an op that reused a key would inherit the + previous instance's shape. + +**Second, the fake builds its output with `torch.empty` from the shape and dtype the +signature check infers, not with `torch.empty_like(x)`.** The tensor the fake returns +must match real execution in shape, dtype and stride. A mismatch either fails during +tracing or, for a stride, makes downstream code read memory in the wrong layout and +produce wrong results silently. The operator body makes the inputs contiguous before the +kernel writes into a newly allocated output, so the real output is always contiguous. +`empty_like` copies the input's strides, so for a non-contiguous input the fake would +declare a layout real execution never produces. + +**Third, the target is resolved inside the node, not in `Op.__call__`.** When traced +code runs `self.x = ...`, dynamo records the write as a pending side effect and applies +it only after the whole graph has run, while the opaque node runs before that. A +resolution written just outside the node therefore cannot be read inside it. Two things +follow, and both happen inside the node: + +- A resolution made outside the node would make the first compiled call silently run the wrong implementation. -- Undoing a failed resolution is the job of whichever site made it, since a compiled +- Undoing a failed resolution is the job of the place that made it, because a compiled artefact does not keep the call site's `try/except`. -All three follow from one fact: torch's compilation and declaration mechanisms -work per function, while what needs compiling is one call on an object. +All three have the same cause: torch's compilation and declaration mechanisms work per +function, while what needs compiling is one call on an object. ## Appendix: why the boundary looks like this -### How dynamo works +### How dynamo works {#dynamo} -This section is about how dynamo decides what may enter a graph — which is where -the conditions an op has to satisfy come from. +This section describes how dynamo decides what code may enter a compiled graph. The +conditions an op has to satisfy to be brought into `torch.compile` come from these +rules. -Dynamo is the front end of `torch.compile`, working at CPython's frame evaluation +Dynamo is the front end of `torch.compile`, and works at CPython's frame evaluation layer (PEP 523). -**It has exactly one entry point: `torch.compile`.** `torch.compile(fn)` returns a -wrapper, and tracing happens when that wrapper is called; `nn.Module.compile()` -and the decorator form are two other spellings of the same entry. A call that does -not go through it takes the ordinary Python path and has nothing to do with dynamo -— below, that path is called eager. +**Dynamo has exactly one entry point: `torch.compile`.** `torch.compile(fn)` returns a +wrapper, and tracing happens when the wrapper is called; `nn.Module.compile()` and the +decorator form are two other spellings of the same entry. A call that does not go +through it takes the ordinary Python path and has nothing to do with dynamo. This page +calls that path eager. On the first call, dynamo takes over the frame, symbolically executes the bytecode -instruction by instruction, records the tensor operations as one FX graph, leaves -what cannot enter the graph in Python, and notes a set of guards for the graph — -the premises this trace relied on, such as a tensor's dtype and rank. Later calls -reuse the compiled artefact when every guard holds; if one fails, the new case is -traced again. +instruction by instruction, records the tensor operations as one FX graph, and leaves +what cannot enter the graph to run in Python. It also records a set of guards for the +graph: the premises this trace relied on, such as a tensor's dtype and rank. A later +call reuses the compiled artefact when every guard holds; when one fails, dynamo traces +the new case again. Three terms have fixed meanings on this page: @@ -206,29 +223,29 @@ Three terms have fixed meanings on this page: | --- | --- | | Graph | the FX graph dynamo captured; one trace produces one | | Node | one operator call in the graph, with input edges and the output's shape and dtype | -| Traced | inside dynamo's symbolic execution. Tracing performs no real computation; it records | +| Traced | inside dynamo's symbolic execution. Tracing performs no real computation; it only records | -The graph then goes to a backend (inductor and others) for fusion, memory planning -and code generation. The larger a graph is, the more neighbouring operators can -fuse, so every op in an operator library has to be able to appear as a node in -someone else's graph. +The graph then goes to a compiler backend (inductor and others) for fusion, memory +planning and code generation. The larger a graph is, the more neighbouring operators +can fuse, so every op in an operator library has to be able to appear as a node in a +user's graph. -Two of dynamo's rules matter for bringing an op in: +Two of dynamo's rules decide how an op is brought in: -- **It inlines by default.** A called function is not itself a boundary, and its - body is folded into the same trace. Keeping a stretch of Python out of the trace - takes an explicit declaration. -- **Untraceable code is handled one of two ways.** By default dynamo breaks the graph, falling - back to Python for that stretch, so one graph becomes several; under - `fullgraph=True` it raises instead. Raising surfaces the problem during - development, which is why an operator library treats `fullgraph=True` as its - acceptance criterion. +- **Dynamo inlines by default.** A called function is not itself a boundary, and its + body is folded into the same trace. Keeping a stretch of Python code out of the trace + requires an explicit declaration. +- **Untraceable code is handled in one of two ways.** By default dynamo breaks the + graph: that stretch falls back to Python, and one graph becomes several. Under + `fullgraph=True` dynamo raises instead. Raising surfaces the problem during + development, which is why an operator library uses `fullgraph=True` as its acceptance + criterion. ### Where the op layer and dynamo disagree -Applying those rules to a TileOPs op makes the obstacle plain. Dynamo compiles -frames, that is, functions; a TileOPs op is an object, and one call does four -things, of which only the last belongs in the graph: +Applying those rules to a TileOPs op shows the obstacle. Dynamo compiles frames, that +is, functions, while a TileOPs op is an object. One call does four things, and only the +last belongs in the graph: | What a call does | Should dynamo capture it | | --- | --- | @@ -239,72 +256,79 @@ things, of which only the last belongs in the graph: The table needs three qualifications. -**"Should not be captured" is not "does not run."** All four happen on every call; -the only question is what enters the graph. +**Work that should not be captured still runs.** All four things happen on every call; +the only difference is whether they enter the graph. -**The distinction has to be annotated by hand** — dynamo cannot draw it. Torch -provides two interfaces: `torch.library.custom_op` registers the call as an -operator, so dynamo puts a single node in the graph and does not trace into the -implementation, and `register_fake` tells the compiler what the node outputs, -receiving only the inputs' metadata and never touching real data. +**The distinction has to be annotated by hand**, because dynamo cannot draw it. Torch +provides two interfaces for this: -**Without the annotation, tracing goes in and fails.** With the boundary +- `torch.library.custom_op` registers the call as an operator, so dynamo puts a single + node in the graph and does not trace into the implementation; +- `register_fake` tells the compiler what the node outputs; it receives only the + inputs' metadata and never touches real data. + +**Without the annotation, dynamo traces into the code and fails.** With the boundary undeclared, `RMSNormFwdOp` compiles in neither of its two states: - An instance that has not built a kernel builds one during the call, and dynamo traces into the TileLang JIT inside the constructor. -- An instance that already has one skips the construction, but still re-parses the +- An instance that already has one skips construction, but still re-parses the TileLang program on every call, so dynamo traces into `@tilelang.jit` and stops at `inspect.signature`. -### Why the boundary falls at the op layer +### Why the boundary falls at the op layer {#at-op-layer} + +The boundary could sit at the op layer or lower, at the kernel layer. The difference +shows in the user's compiled graph. + +The node's identity in that graph (its name, arguments, granularity, and the output +its fake declares) is the op the user sees. With the boundary at the kernel layer, +changing backend changes that node: the same op compiles to a different graph under a +different target, and the compiled artefact is tied to the backend. With the boundary +at the op layer, the node's identity is the op's, independent of the backend serving +it. -The boundary could sit at the op layer, or lower down at the kernel layer. The -difference shows up in the user's compiled graph. +That position also decides how the fake is written. The op layer does not know how an +external kernel tiles or pads internally. The only shape rule that holds for every +target is the one in the manifest, so the fake derives its output from the manifest. -The node's identity in that graph — its name, arguments, granularity, and the -output its fake declares — is the operator the user sees. With the boundary at the -kernel layer, changing backend changes that node, the same op compiles to a -different graph under a different target, and the compiled artefact is tied to the -backend. At the op layer it does not: the node's identity is the op's, independent -of who serves it. +The node's interior is invisible to the compiler, but its contract with the outside is +complete: -That position also settles how the fake is written. The op layer does not know how -an external kernel tiles or pads internally; the one shape rule that holds for -every target is the one in the manifest, so the fake derives from it. +- the schema gives the name and argument types; +- the fake gives the output's shape, dtype, device and stride; +- the alias annotations name exactly the inputs the node writes, of which + `RMSNormFwdOp` has none. -The node's interior is invisible to the compiler, but its contract to the outside is -complete: the schema gives the name and argument types, the fake gives the output's -shape, dtype, device and stride, and the alias annotations name exactly the inputs it -writes — none, for `RMSNormFwdOp`. With the contract complete, what is kept and what is -given up separate cleanly: +Because the contract is complete, what the boundary keeps and what it gives up +separate cleanly: -- **Optimisation between nodes proceeds as usual.** Buffer assignment, lifetimes, - reordering against neighbours it does not depend on, deletion when nothing - consumes it. -- **Optimisation inside the node is gone.** Neighbours cannot fuse in, and the - output must be written to memory. +- **Optimisation between nodes proceeds as usual.** This includes buffer assignment, + lifetimes, reordering against neighbours the node does not depend on, and deleting the + node when nothing consumes its output. +- **Optimisation inside the node is lost.** Neighbouring operators cannot fuse into + the node, and its output must be written to memory. -For an operator library the trade is worth it: inside the node is a kernel TileLang -has already compiled, which inductor need not touch. +For an operator library the trade is worth it: inside the node is a kernel TileLang has +already compiled, which inductor does not need to touch. ### What the boundary costs -Measured on an idle H200 at 2048×4096, fp16. Per-call figures are the minimum of -three runs of 2000 iterations × 9 rounds: +Measured on an idle H200 at 2048×4096, fp16. Per-call figures are the minimum of three +runs of 2000 iterations × 9 rounds: | | Boundary at the kernel layer | Boundary at the op layer | | --- | --- | --- | | Kernel time | 0.0119 ms | 0.0117 ms | | Eager, per call | 42.5–45.2 µs | 38.2–42.0 µs | -The kernel itself is unaffected — where the boundary sits has nothing to do with -how the kernel computes. The eager path is 3–5 µs faster because, with the boundary +The kernel itself is unaffected, because where the boundary sits has nothing to do +with how the kernel computes. The eager path is 3–5 µs faster because, with the boundary moved up, a call crosses one operator boundary instead of two. -The cost on the graph side is in [why the boundary falls at the op -layer](#why-the-boundary-falls-at-the-op-layer): fusion does not cross the node boundary, -and the node's output always lands in memory. +The cost on the graph side is described in +[Why the boundary falls at the op layer](#at-op-layer): fusion does not cross the node +boundary, and the node's output is always written to memory. ### What the boundary does not provide @@ -312,5 +336,5 @@ and the node's output always lands in memory. | --- | --- | | Fusion across the node boundary | The node's interior is opaque to the compiler, so the elementwise work on either side stays outside | | autograd through the node | This path serves inference; forward and backward are separate ops | -| Switching target within one compiled artefact | The target belongs to the op instance — another target means another instance, and another graph | -| Building a kernel from meta tensors | A call with meta tensors answers with shapes and dtypes only | +| Switching target within one compiled artefact | The target belongs to the op instance: another target means another instance, and another graph | +| Building a kernel from meta tensors | A call with meta tensors returns only shapes and dtypes | diff --git a/docs/torch-compile.zh.md b/docs/torch-compile.zh.md index b7f338c7..dbdd27cd 100644 --- a/docs/torch-compile.zh.md +++ b/docs/torch-compile.zh.md @@ -1,18 +1,23 @@ # 接入 torch.compile -把一个 TileOPs 算子接入 `torch.compile`，是让它在使用者的编译图里成为一个节点，而这个节点的形态不随服务它的后端变化。 +一个 TileOPs op 接入 `torch.compile` 之后，在使用者的编译图中成为一个节点，这个节点的形态不随服务它的 backend 变化。 -要做的只有一件事：在算子层声明一条编译边界，界外交给 dynamo 追踪，界内对编译器不可见。 +接入只需要一项工作：在 op 层声明一条编译边界。边界之外由 dynamo 追踪，边界之内对编译器不可见。 -正文讲要做的事：判断一个算子是否已接入、编译一段调用它的代码、调用时的五条约定，以及给尚未接入的算子声明这条边界要写什么。 +正文说明接入相关的操作： -附录讲这条边界为什么只能这样划：dynamo 怎么工作、它与算子层在哪里错位、边界为什么落在算子层，以及这条边界的代价与限制。 +1. 判断一个 op 是否已接入； +1. 编译一段调用它的代码； +1. 调用时的五条约定； +1. 为尚未接入的 op 声明这条边界需要编写的代码。 -## 调用一个已接入的算子 +附录说明这条边界为什么只能这样划分：dynamo 的工作方式、它与 op 层的不一致之处、边界为什么位于 op 层，以及这条边界的代价与限制。 -### 判断一个算子是否已接入 {#supported} +## 调用一个已接入的 op -读类属性 `compile_op_names`：非空说明这个类声明了编译边界（`compile_boundary = True`），边界在算子层，`fullgraph=True` 可用；空 tuple 说明它没有声明。 +### 判断一个 op 是否已接入 {#supported} + +读取类属性 `compile_op_names`。它非空时，说明这个类声明了编译边界（`compile_boundary = True`），边界位于 op 层，`fullgraph=True` 可用；它是空 tuple 时，说明这个类没有声明编译边界。 ```python >>> from tileops.norm import RMSNormFwdOp @@ -20,11 +25,11 @@ ('tileops::norm_rms_norm_fwd',) ``` -尚未迁移的算子在 `fullgraph=True` 下报错，默认设置下切图。 +尚未迁移的 op 在 `fullgraph=True` 下报错，在默认设置下切图。 ### 编译一段调用它的代码 -构造算子实例，把调用它的函数交给 `torch.compile`，没有别的步骤： +构造 op 实例，再把调用它的函数传给 `torch.compile`，不需要其他步骤： ```python import torch @@ -41,31 +46,31 @@ w = torch.randn(4096, device="cuda", dtype=torch.float16) block(x, w) ``` -用 `TORCH_LOGS=graph_code` 运行会打印捕获到的图：里面是 `tileops::norm_rms_norm_fwd` 一个节点，不是 kernel 内部的多次调用。 +用 `TORCH_LOGS=graph_code` 运行时会打印捕获到的图，图中只有 `tileops::norm_rms_norm_fwd` 一个节点，kernel 内部的多次调用不出现在图中。 -### 调用时要遵守的五条约定 +### 调用时的五条约定 -五条各对应边界上的一处机制，违反任何一条，编译路径的行为就与 eager 路径不同。 +五条约定各对应边界上的一处机制。违反任何一条，编译路径的行为都会与 eager 路径不同。 -- **算子实例构造一次并反复使用。** 实例键是编译期常量，一个实例对应一张编译图；在循环里新建实例，每次迭代都要重新编译。 -- **不要依赖 stride 原样传递。** 算子不写入的非连续输入在节点内部连续化，算子自己分配的输出恒为连续张量；后续计算需要别的布局，在算子之外自行转换。输出若就是被写入的输入（`alias`）或调用方给的 `out`，沿用那个张量的存储。 -- **不能用 meta 张量预热。** 有了边界，传入 meta 或 fake 张量的调用就在 fake 处返回，走不到构造 kernel 那一步。 -- **CUDA graph 捕获之前先行预热。** 用真实张量、相同形状至少调用一次：构造 kernel 允许编译，捕获期间只允许查表命中后直接调用。各阶段分别允许执行哪些操作，见[各阶段允许做什么](backends.md#phase-limits)。 -- **换一块卡可能要重新构造。** 由 target 服务的调用，设备是 kernel 记忆键的一部分，同一个实例换到另一块卡上会重新构造一次。自带 kernel 的记忆键是选中实现的 `entry_for` 返回的 build identity，只有构建结果与设备有关时才带上设备。构造函数里指名的 `target=` 在首次编译调用中同样生效；构造失败不会把算子固定到任何 target。 +- **op 实例构造一次并反复使用。** 实例键是编译期常量，一个实例对应一张编译图；在循环中新建实例时，每次迭代都要重新编译。 +- **stride 不会原样传递。** op 不写入的非连续输入在节点内部被转为连续张量，op 自己分配的输出总是连续张量；后续计算需要其他布局时，在 op 之外自行转换。输出就是被写入的输入（`alias`）或调用方提供的 `out` 时，沿用那个张量的存储。 +- **meta 张量不能用于预热。** 声明边界之后，传入 meta 或 fake 张量的调用在 fake 函数处返回，不会执行到构造 kernel 的步骤。 +- **CUDA graph 捕获之前需要预热。** 用真实张量以相同形状至少调用一次：构造 kernel 时允许编译，捕获期间只允许在缓存命中后直接调用。各阶段分别允许执行哪些操作，见[各阶段允许做什么](backends.md#phase-limits)。 +- **换到另一块卡时可能重新构造 kernel。** 对由 target 服务的调用，设备是 kernel 缓存键的一部分，同一个实例换到另一块卡上会重新构造一次。in-tree kernel 的缓存键是选中实现的 `entry_for` 返回的 build identity，只有构建结果与设备有关时才包含设备。构造函数中指定的 `target=` 在首次编译调用中同样生效；构造失败不会把 op 固定到任何 target。 ### 接入之后成立的三项保证 -边界落在算子层之后，调用方可以依赖三点。 +边界位于 op 层之后，调用方可以依赖以下三点： -- **编译图不随 target 变化。** 换后端或换硬件，同一段代码编出的图完全相同，编译产物因此与后端无关。 -- **`fullgraph=True` 可用。** 前提是该算子已经声明这条契约，判定见[判断一个算子是否已接入](#supported)。 -- **输出的形状、dtype 与 stride 由 manifest 规定。** 与 kernel 内部怎么分块、怎么 padding 无关。算子自己分配的输出恒为连续张量。 +- **编译图不随 target 变化。** 更换 backend 或硬件后，同一段代码编译出的图完全相同，编译产物因此与 backend 无关。 +- **`fullgraph=True` 可用。** 前提是该 op 已经声明这条契约，判断方法见[判断一个 op 是否已接入](#supported)。 +- **输出的形状、dtype 与 stride 由 manifest 规定。** 它们与 kernel 内部的分块和 padding 方式无关。op 自己分配的输出总是连续张量。 -## 给一个新算子声明编译边界：`RMSNormFwdOp` +## 为新 op 声明编译边界：`RMSNormFwdOp` -接入一个算子要写的代码：边界怎么声明、fake 怎么写、target 判定为什么要在节点内部重做一次。其中的追踪、切图、guard 见[dynamo 是怎么工作的](#dynamo)。 +本节给出接入一个 op 需要编写的代码：边界如何声明、fake 如何编写，以及 target 判定为什么要在节点内部重新执行一次。其中涉及的追踪、切图与 guard 见 [dynamo 是怎么工作的](#dynamo)。 -`RMSNormFwdOp` 是仓内第一个接入的算子。下面是它的骨架，略去 docstring，完整代码见 [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py)： +`RMSNormFwdOp` 是仓库中第一个接入的 op。下面是它的骨架，略去 docstring，完整代码见 [`src/tileops/ops/norm/rms_norm.py`](https://github.com/tile-ai/TileOPs/blob/main/src/tileops/ops/norm/rms_norm.py)： ```python class RMSNormFwdOp(Op): @@ -90,65 +95,72 @@ class RMSNormFwdOp(Op): return self.kernel_for("rms_norm", call)(x, weight) ``` -声明就这么多。operator 与它的 fake 都从条目生成，每个副作用分支一个 operator：张量参数是 `signature.inputs` 的顺序，返回什么看 `signature.outputs`，写哪些参数恰好是标了 `mutated` 的输入，每个输出的形状与 dtype 取自签名。名字是 `tileops::_`（类名本身以 family 开头时只写一次），这里就是 `tileops::norm_rms_norm_fwd`；一个分支的 operator 若写入某个输入、填写 `buffer: out` 或不产出某个输出，名字后面就按这个顺序分别加上 `_writes_<输入>`、`_out`、`_without_<输出>`。没有算子自己起名字，`compile_op_names` 也就不可能和注册的名字对不上。 +声明只有以上内容。operator 与它的 fake 函数都从 manifest 条目生成，每个副作用分支对应一个 operator： + +- 张量参数按 `signature.inputs` 的顺序排列； +- 返回值由 `signature.outputs` 决定； +- 被写入的参数恰好是标记了 `mutated` 的输入； +- 每个输出的形状与 dtype 取自签名。 + +operator 的名字是 `tileops::_`（类名本身以 family 开头时只写一次），这里就是 `tileops::norm_rms_norm_fwd`。一个分支的 operator 若写入某个输入、填写 `buffer: out` 或不产出某个输出，名字后面按这个顺序分别加上 `_writes_<输入>`、`_out`、`_without_<输出>`。op 不自行命名 operator，因此 `compile_op_names` 不可能与注册的名字不一致。 -一次调用经过的各层，以及边界落在哪里： +一次调用经过的各层，以及边界所在的位置：
Op.__call__调用 forward，不判定 target
-
forward一行，调用不透明算子
+
forward一行，调用不透明 operator
编译边界
-
生成出来的算子取回算子实例，跑生成的检查，判定 target，失败则撤销
-
_eager_forward连续化、取 kernel、launch kernel
-
紫色两层在 dynamo 的追踪范围内，forward 那一行是它追到的最后一处；界下由不透明算子接手，编译器看不见。
+
生成的 operator取回 op 实例，执行生成的检查，判定 target，失败时撤销
+
_eager_forward转为连续张量、取得 kernel、launch kernel
+
紫色的两层在 dynamo 的追踪范围内，forward 中的那一行是 dynamo 追踪到的最后一处；边界以下由不透明 operator 执行，编译器看不到。
-其中三处写法不是任选的。 +其中有三处写法是固定的。 -**第一处，算子实例通过字符串键找回，而不是直接传对象。** schema 的类型只有 `Tensor`、`int`、`float`、`bool`、`str` 等固定几种，没有「任意 Python 对象」；而算子体要用的 `kernel_map`、已定下的 target、kernel 记忆表都挂在实例上，摊不成 schema 参数。键还有两个细节不能改： +**第一处，op 实例通过字符串键取回，不直接传递对象。** schema 的类型只有 `Tensor`、`int`、`float`、`bool`、`str` 等固定几种，没有「任意 Python 对象」；而 op 体需要的 `kernel_map`、已确定的 target 与 kernel 缓存表都保存在实例上，无法拆成 schema 参数。键还有两个不能改变的细节： -- **取字符串，不取整数。** 字符串在追踪期是常量，整数会被泛化成 `SymInt`。 -- **从不复用。** 正因为是常量，inductor 会把 fake 给出的形状固化进产物；复用键的算子会继承前一个实例的形状。 +- **键是字符串，不是整数。** 字符串在追踪期是常量，整数会被泛化为 `SymInt`。 +- **键从不重用。** 由于键是常量，inductor 会把 fake 函数给出的形状固定在编译产物中；重用键的 op 会继承前一个实例的形状。 -**第二处，fake 用 `torch.empty` 按签名检查推出的形状与 dtype 构造，而不是 `torch.empty_like(x)`。** fake 返回的张量，形状、dtype 与 stride 三项都必须与真实执行返回的一致；不一致或在追踪期报错，或在运行期按错误布局访问而静默出错。算子体先连续化再写入新分配的输出，真实输出恒为连续，而 `empty_like` 会把入参的 stride 一起复制 —— 非连续输入就让 fake 宣称了一种真实执行不会产出的布局。 +**第二处，fake 函数用 `torch.empty` 按签名检查推出的形状与 dtype 构造输出，不使用 `torch.empty_like(x)`。** fake 函数返回的张量在形状、dtype 与 stride 三项上都必须与真实执行的返回值一致。不一致时，要么在追踪期报错，要么在运行期按错误布局访问内存而静默出错。op 体先把输入转为连续张量，再写入新分配的输出，因此真实输出总是连续的；而 `empty_like` 会复制入参的 stride，输入不连续时，fake 函数声明的布局就是真实执行不会产出的布局。 -**第三处，target 在节点内部判定，不在 `Op.__call__` 里。** 追踪期执行 `self.x = ...`，dynamo 把这次写入记成待办的副作用，等整张图跑完才补上；而不透明节点的执行早于补写，所以节点之外刚写下的判定结果，节点之内读不到。两件事因此都落在节点内部： +**第三处，target 在节点内部判定，不在 `Op.__call__` 中判定。** 追踪期执行 `self.x = ...` 时，dynamo 把这次写入记为待执行的副作用，等整张图执行完才补写；而不透明节点的执行早于补写，所以节点之外刚写下的判定结果在节点之内读不到。因此以下两件事都在节点内部完成： -- 判定若写在节点之外，第一次编译调用会静默用错实现。 -- 判定失败时的撤销由做出判定的那一处负责，因为编译产物不保留调用点的 `try/except`。 +- 判定若写在节点之外，第一次编译调用会静默使用错误的实现。 +- 判定失败时的撤销由做出判定的位置负责，因为编译产物不保留调用点的 `try/except`。 -三处的原因是同一个：torch 的编译与声明机制以函数为单位，而要编译的是一个对象上的一次调用。 +三处写法的原因相同：torch 的编译与声明机制以函数为单位，而需要编译的是一个对象上的一次调用。 ## 附录：这条边界为什么是这样 ### dynamo 是怎么工作的 {#dynamo} -这一节说明 dynamo 如何决定什么能进入编译图 —— 一个算子要接入 `torch.compile`，需要满足的条件由此而来。 +本节说明 dynamo 如何决定哪些代码能进入编译图。一个 op 接入 `torch.compile` 需要满足的条件由此而来。 dynamo 是 `torch.compile` 的前端，工作在 CPython 的帧求值层（PEP 523）。 -**它的触发入口只有一个：`torch.compile`。** `torch.compile(fn)` 返回一个包装体，包装体被调用时才发生追踪；`nn.Module.compile()` 与装饰器写法是同一入口的另外两种形式。不经过这个入口的调用一律走原来的 Python 路径，与 dynamo 无关 —— 下文把这种路径称为 eager。 +**dynamo 只有一个触发入口：`torch.compile`。** `torch.compile(fn)` 返回一个包装对象，包装对象被调用时才发生追踪；`nn.Module.compile()` 与装饰器写法是同一入口的另外两种形式。不经过这个入口的调用都执行原来的 Python 路径，与 dynamo 无关。下文把这种路径称为 eager。 -第一次调用时，dynamo 接管这一帧，逐条符号执行字节码，把其中的张量运算记成一张 FX 图，把无法进入图的部分留在 Python 里，同时为这张图记下一组 guard，也就是本次追踪所依赖的前提，例如某个张量的 dtype 与维数。此后的调用如果 guard 全部成立，直接复用编译产物；只要有一条不成立，就为新的情况重新追踪一次。 +第一次调用时，dynamo 接管这一帧，逐条符号执行字节码，把其中的张量运算记录为一张 FX 图，把无法进入图的部分留在 Python 中执行。同时，dynamo 为这张图记录一组 guard，即本次追踪所依赖的前提，例如某个张量的 dtype 与维数。此后的调用如果所有 guard 都成立，就直接重用编译产物；只要有一条不成立，就为新的情况重新追踪一次。 -本页用到的三个术语，含义固定如下： +本页使用的三个术语含义如下： | 术语 | 含义 | | --- | --- | -| 编译图 | dynamo 捕获下来的那张 FX 图，一次追踪产出一张 | -| 节点 | 图中的一次算子调用，带有输入边以及输出的形状与 dtype | +| 编译图 | dynamo 捕获的 FX 图，一次追踪产出一张 | +| 节点 | 图中的一次 op 调用，带有输入边以及输出的形状与 dtype | | 追踪 | 处在 dynamo 的符号执行范围内。追踪期不执行真实计算，只做记录 | -图随后交给后端（inductor 等）完成融合、内存规划与代码生成。图越大，可融合的相邻算子越多，因此算子库的每个算子都要能作为节点出现在使用者的图里。 +图随后交给编译 backend（inductor 等），由它完成融合、内存规划与代码生成。图越大，可融合的相邻 op 越多，因此算子库中的每个 op 都需要能作为节点出现在使用者的图中。 -对接入而言，dynamo 的两条规则是关键： +dynamo 有两条规则决定了接入的方式： -- **默认一路内联。** 被调用的函数本身不构成边界，函数体会被并入同一次追踪。要让某一段 Python 不被追进去，只能显式声明。 -- **追踪不了的代码有两种处理方式。** 默认设置下切图（graph break），把这一段退回 Python 执行，一张图被切成数张；`fullgraph=True` 下直接报错。后者让问题在开发期暴露，因此算子库以 `fullgraph=True` 作为验收条件。 +- **默认全部内联。** 被调用的函数本身不构成边界，函数体会并入同一次追踪。要让某一段 Python 代码不被追踪，只能显式声明。 +- **无法追踪的代码有两种处理方式。** 默认设置下切图（graph break），这一段退回 Python 执行，一张图被切成多张；`fullgraph=True` 下直接报错。后者让问题在开发期暴露，因此算子库以 `fullgraph=True` 作为验收条件。 -### 算子层与 dynamo 在哪里错位 +### op 层与 dynamo 的不一致之处 -把上面的规则套到 TileOPs 的算子上，接入的障碍就清楚了：dynamo 编译的单位是帧，也就是函数；而 TileOPs 的算子是对象，一次调用要完成四项工作，其中只有最后一项属于图： +把上述规则应用到 TileOPs 的 op 上，接入的障碍是：dynamo 编译的单位是帧，即函数；而 TileOPs 的 op 是对象，一次调用要完成四项工作，其中只有最后一项属于图： | 一次调用完成的工作 | 是否应当被 dynamo 捕获 | | --- | --- | @@ -159,48 +171,54 @@ dynamo 是 `torch.compile` 的前端，工作在 CPython 的帧求值层（PEP 5 这张表需要补充三点。 -**「不应当被捕获」不等于「不执行」。** 四项工作在每次调用中都照常发生，区别只在于是否进入编译图。 +**「不应当被捕获」的工作照常执行。** 四项工作在每次调用中都会执行，区别只在于是否进入编译图。 + +**这个区分需要人工标注**，dynamo 自己无法区分。torch 为此提供两个接口：`torch.library.custom_op` 把这一次调用注册为一个 operator，dynamo 在图中只放一个节点，不追踪其实现；`register_fake` 告诉编译器这个节点的输出是什么，它只接收输入的元信息，不接触真实数据。 + +**不标注时，这些代码会被追踪，并且一定失败。** 以未声明边界的 `RMSNormFwdOp` 为例，实例的两种状态都无法编译： -**这个区分需要人来标注**，dynamo 自己分不出来。torch 为此提供两个接口：`torch.library.custom_op` 把这一次调用注册成一个算子，dynamo 在图里只放一个节点、不追进实现；`register_fake` 告诉编译器这个节点输出什么，它只接收输入的元信息、不接触真实数据。 +- 尚未构造过 kernel 的实例会在本次调用中构造 kernel，dynamo 因此追踪进构造函数中的 TileLang JIT。 +- 已经构造过 kernel 的实例跳过构造，但每次调用仍要重新解析 TileLang program，dynamo 追踪进 `@tilelang.jit`，停在 `inspect.signature`。 -**不标注就会被追进去，而且一定失败。** 以未声明边界的 `RMSNormFwdOp` 为例，实例的两种状态都编不过去： +### 为什么边界位于 op 层 {#at-op-layer} -- 尚未构造过 kernel 的实例，会在本次调用中现场构造，dynamo 于是追进了构造函数里的 TileLang JIT。 -- 已经构造过 kernel 的实例跳过构造，但每次调用仍要重新解析 TileLang program，dynamo 追进 `@tilelang.jit`，停在 `inspect.signature`。 +边界可以划在 op 层，也可以划在更低的 kernel 层。两者的差别体现在使用者的编译图上。 -### 为什么边界落在算子层 {#at-op-layer} +图中那个节点的身份，包括名字、参数、粒度以及 fake 函数给出的输出，就是使用者看到的 op。边界划在 kernel 层时，更换 backend 就更换了这个节点，同一个 op 在不同 target 下会编译出不同的图，编译产物因此与 backend 绑定。边界划在 op 层时，节点的身份由 op 决定，与服务它的 backend 无关。 -边界可以划在算子层，也可以划在更靠下的 kernel 层。两者的差别落在使用者的编译图上。 +这个位置同时决定了 fake 函数的写法。op 层不知道外部 kernel 内部如何分块、如何 padding，对所有 target 都成立的形状规则只有 manifest 中的规则，因此 fake 函数只能依照 manifest 推导。 -图中那个节点的身份 —— 名字、参数、粒度以及 fake 给出的输出 —— 就是使用者看到的算子。如果边界划在 kernel 层，换一个后端就换掉了这个节点，同一个算子在不同 target 下会编出不同的图，编译产物于是与后端绑定。划在算子层则不会：节点的身份由算子决定，与哪个后端在服务它无关。 +节点内部对编译器不可见，但它对外的契约是完整的： -这个位置同时决定了 fake 的写法。算子层并不知道外部 kernel 内部如何分块、如何 padding，唯一对所有 target 共同成立的形状规则写在 manifest 里，因此 fake 只能照 manifest 推导。 +- schema 给出名字与参数类型； +- fake 函数给出输出的形状、dtype、设备与 stride； +- 别名标注恰好列出它写入的入参，`RMSNormFwdOp` 没有这样的入参。 -节点内部对编译器不可见，但它对外的契约是完整的：schema 给出名字与参数类型，fake 给出输出的形状、dtype、设备与 stride，别名标注恰好列出它写入的入参，`RMSNormFwdOp` 一个也没有。契约完整，得失也就分得清楚： +由于契约完整，边界带来的得失可以明确区分： -- **节点之间的优化照常。** 排布 buffer、计算生命周期、与无依赖的相邻节点交换顺序、无人使用时整体删除。 -- **节点内部的优化没有了。** 相邻算子融不进来，输出必须写入显存。 +- **节点之间的优化照常进行。** 包括排布 buffer、计算生命周期、与无依赖的相邻节点交换顺序，以及在输出无人使用时整体删除节点。 +- **节点内部的优化不再进行。** 相邻 op 无法融合进节点，输出必须写入显存。 -对算子库来说这个取舍是值得的：节点内部是 TileLang 已经编译好的 kernel，本来就不需要 inductor 介入。 +对算子库而言，这个取舍是合理的：节点内部是 TileLang 已经编译好的 kernel，本来就不需要 inductor 介入。 ### 边界的代价 -以下数据在空闲的 H200 上实测，形状为 2048×4096、dtype 为 fp16；每次调用的数字取 2000 次迭代 × 9 轮、三次运行中的最小值： +以下数据在空闲的 H200 上实测，形状为 2048×4096，dtype 为 fp16；每次调用的数字取 2000 次迭代 × 9 轮、三次运行中的最小值： -| | 边界位于 kernel 层 | 边界位于算子层 | +| | 边界位于 kernel 层 | 边界位于 op 层 | | --- | --- | --- | | kernel 时间 | 0.0119 ms | 0.0117 ms | | eager 路径每次调用 | 42.5–45.2 µs | 38.2–42.0 µs | -kernel 本身不受影响，边界位于哪一层与 kernel 如何计算无关。eager 路径快 3–5 µs，原因是边界上移之后一次调用只需穿过一层算子边界。 +kernel 本身不受影响，边界位于哪一层与 kernel 如何计算无关。eager 路径快 3–5 µs，原因是边界上移之后，一次调用只需穿过一层 op 边界。 -编译图一侧的代价见[为什么边界落在算子层](#at-op-layer)：融合不跨越节点边界，节点的输出必定写入显存。 +编译图一侧的代价见[为什么边界位于 op 层](#at-op-layer)：融合不跨越节点边界，节点的输出必定写入显存。 ### 编译边界不提供的能力 | 不提供 | 原因 | | --- | --- | | 跨节点边界的融合 | 节点内部对编译器不可见，两侧的 elementwise 计算只能留在节点之外 | -| 穿过节点的 autograd | 这条调用链服务推理，fwd 与 bwd 各自是独立的算子 | -| 在同一份编译产物中切换 target | target 属于算子实例，更换 target 即更换实例，也就更换了编译图 | -| 用 meta 张量构造 kernel | 传入 meta 张量的调用只回答形状与 dtype | +| 穿过节点的 autograd | 这条调用链服务推理，fwd 与 bwd 各自是独立的 op | +| 在同一份编译产物中切换 target | target 属于 op 实例，更换 target 即更换实例，也就更换了编译图 | +| 用 meta 张量构造 kernel | 传入 meta 张量的调用只返回形状与 dtype | diff --git a/docs/user-guide/dispatch/backends.md b/docs/user-guide/dispatch/backends.md new file mode 100644 index 00000000..ad12eee2 --- /dev/null +++ b/docs/user-guide/dispatch/backends.md @@ -0,0 +1,111 @@ +# How a backend joins TileOPs + +A third-party backend has three ways to join, ordered from the smallest range of calls taken over to the largest: replace the implementation behind one key, add an implementation to a kernel interface, or replace a whole op. In the first two, the backend's class and the in-tree implementations follow the same kernel interface; how to write one is in [Adding a kernel to an op](writing.md). + +## 1. Choose the way by the range of calls to take over {#choose} + +![The ranges taken over by the three ways of extending](img/extension.svg) + +**Figure 1** The range each of the three ways of extending takes over. Teal marks the parts provided by the system, purple the parts written by TileOPs developers, and green the parts provided by the backend. + +**Table 1** Comparison of the three ways + +| No. | Item | `kernel_map=` | `register_implementation` | target | +| --- | --- | --- | --- | --- | +| 1 | What changes | the class that runs behind one key; which calls the key serves is unchanged | a new key, with its own applicability and precedence | every call of the op, except calls whose written tensors are all empty | +| 2 | Applies to | the one op instance the caller constructs | every instance of the op constructed after registration | the op instances that select the target | +| 3 | Contract followed | the kernel interface | the kernel interface | the op's signature in the manifest | +| 4 | Calls the new class does not serve | raise an error when the key is selected | still served by the in-tree implementations | none; the target serves every call, except calls whose written tensors are all empty | + +## 2. Replace the implementation behind one key: kernel_map= {#kernel-map} + +`kernel_map=` is a parameter of every op constructor, and its value maps keys to classes. It replaces only the class that runs behind the key in this op instance: + +- Which calls the key serves and its precedence are still decided by the declarations of the originally registered implementation; the replacement's own `applies`, `general`, and `preferred_over` take no part in selection. +- The key is available on devices where either the original implementation or the replacement can run. When the replacement can run on devices where the original cannot, the key also takes part in selection on those devices, and an overlap with other implementations raises an ambiguity error under the same rules. +- When the key is selected, the call is served by the replacement. When the replacement cannot run on the call's device or refuses the call, the call raises an error and does not fall back to the original implementation. +- Calls that select other keys do not consult the replacement. +- The replacement follows the same rule as a registered implementation: it subclasses the kernel interface the key belongs to and is built by its own class method `entry_for(call)`. There is no other way to write it. + +To change which calls a kernel serves, use `register_implementation` in § 3. + +```python +# tests/test_kernel_dispatch.py +class _TorchLayerNorm(Kernel, LayerNormFwdInterface): + """A replacement written against ``LayerNormFwdInterface`` alone.""" + + devices = frozenset({torch.device(run_device()).type}) + + def __init__(self, n: int, eps: float) -> None: + super().__init__() + self.n, self.eps = n, eps + + @classmethod + def entry_for(cls, call: LayerNormCall): + return (call.n, call.eps), lambda: cls(call.n, call.eps) + + def forward(self, x, weight, bias): + return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( + x.dtype + ) + +op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) +``` + +Keys in `kernel_map=` that this op lacks but other ops have are ignored, so a composite op passes one map to all its sub-ops. A key that no op has raises `was given kernel_map keys no op has` at construction. + +## 3. Add an implementation: register_implementation {#register} + +`tileops.backend.register_implementation(op, key, implementation)` adds an implementation to an op. It is the backend form of step 1 in [Adding a kernel to an op](writing.md). `op` is the op's class name, `key` is the new implementation's name, and `implementation` subclasses one of the op's kernel interfaces; the inheritance decides which interface it belongs to. Step 2 is the same as for in-tree implementations. + +```python +# tests/test_kernel_dispatch.py +class _NarrowTorchLayerNorm(_TorchLayerNorm): + """An added implementation for short rows, which wins over the in-tree one there.""" + + preferred_over = frozenset({"layer_norm"}) + + @classmethod + def applies(cls, call: LayerNormCall) -> bool: + return call.n <= 64 + +register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) +``` + +- The new implementation overlaps `LayerNormKernel` on `n <= 64`, and `LayerNormKernel` is not general, so the new implementation declares `preferred_over`. +- Calls the new implementation does not serve, such as `n = 1024`, are still served by the in-tree implementation. +- The new implementation enters only op instances constructed after registration. +- Registering the same key twice under one op raises `BackendError`. A key equal to an in-tree key raises `reuse keys it has` when the op is constructed. + +Registration happens when the backend module is imported. The backend declares an entry point in `pyproject.toml`, and TileOPs imports it when the first op is constructed: + +```toml +[project.entry-points."tileops.backends"] +acme = "tileops_acme" +``` + +When the module fails to import, all its registrations are undone and TileOPs emits a `RuntimeWarning`. The failure records are available through `tileops.backend.load_failures()`. + +## 4. Replace a whole op: target {#target} + +The full target protocol, a runnable template backend, and common errors are in [Adding a hardware backend](../../backends.md). This section only summarizes how a target differs from the other two ways. + +A target is a name a backend gives to a set of kernels. Once a target registers a builder for an op, every call of an op instance that selects the target is served by the target, and the op's own `forward` does not run. The exception is a call whose written tensors are all empty: then neither the target nor the in-tree implementation runs, and the outputs are constructed from the signature. + +```python +from tileops.backend import TensorSpec, register_detector, register_kernel_builder +from .kernels import AcmeRMSNorm + +register_detector(target="acme", detect=lambda device: device.type == "acme") + +def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): + return AcmeRMSNorm(normalized_shape, eps, x.dtype) + +register_kernel_builder(op="RMSNormFwdOp", target="acme", build_kernel=build_rms_norm) +``` + +- `build_kernel` receives a `TensorSpec` for each input in the order of the op's `signature.inputs` (an optional input that is not passed is `None`, such as `RMSNormFwdOp`'s `weight`), and receives `signature.params` as keywords. The kernel it returns receives the input tensors in the same order at call time, and receives the output buffers and execution arguments the caller passes as keywords. Importing the module compiles nothing; building happens when `build_kernel` is called. +- An op instance caches the kernels returned by `build_kernel` by the call's device, and by whether each input is passed, its dtype, and its shape. +- An op instance fixes its target on the first call that needs to run an implementation, and the target does not change afterwards. The construction argument `target=` comes first, then the process default set by `set_default_target`, and finally each target's detector inspects the call's device. When no target claims the device, the in-tree implementation is used; when several targets claim it, the call raises `AmbiguousTargetError`. A call whose written tensors are all empty does not fix the target. When the call that fixes the target fails, the choice is undone and the next call fixes it again. +- A call passed to a target has already passed the checks generated from the signature. Every tensor is on the call's device, except tensors declared `device: cpu`. Every input that is not written is contiguous, and a written input is contiguous when it declares `contiguous: true`. +- When the selected target has registered no builder for this op and the op holds kernels of its own, the call raises `OpNotAvailableError` and does not fall back to the in-tree implementation. diff --git a/docs/user-guide/dispatch/backends.zh.md b/docs/user-guide/dispatch/backends.zh.md new file mode 100644 index 00000000..1f0a626d --- /dev/null +++ b/docs/user-guide/dispatch/backends.zh.md @@ -0,0 +1,111 @@ +# backend 如何接入 TileOPs + +第三方 backend 按要接管的调用范围，从小到大有三种方式：替换一个 key 背后的实现、为一个 kernel 接口新增实现、替换整个 op。前两种方式中 backend 的类与 in-tree 实现遵守同一份 kernel 接口，写法见[如何为 op 新增 kernel](writing.md)。 + +## 1. 按要接管的调用范围选择接入方式 {#choose} + +![三种扩展方式接管的范围](img/extension.svg) + +**图 1** 三种扩展方式各自接管的范围。青色为系统提供的部分，紫色为 TileOPs 开发者编写的部分，绿色为 backend 提供的部分。 + +**表 1** 三种方式的比较 + +| No. | 项目 | `kernel_map=` | `register_implementation` | target | +| --- | --- | --- | --- | --- | +| 1 | 改变什么 | 一个 key 背后运行的类；这个 key 服务哪些调用不变 | 新增一个 key，带有它自己的适用范围与优先关系 | op 的全部调用，所有写入张量都为空的调用除外 | +| 2 | 作用于 | 调用方构造的一个 op 实例 | 注册之后构造的所有该 op 实例 | 选中该 target 的 op 实例 | +| 3 | 依据的契约 | kernel 接口 | kernel 接口 | manifest 中 op 的签名 | +| 4 | 新类不服务的调用 | 该 key 被选中时报错 | 仍由 in-tree 实现服务 | 不存在，target 服务全部调用，所有写入张量都为空的调用除外 | + +## 2. 替换一个 key 背后的实现：kernel_map= {#kernel-map} + +`kernel_map=` 是每个 op 构造函数的参数，值是从 key 到类的映射。它只替换这个 op 实例中该 key 背后运行的类： + +- 这个 key 服务哪些调用、优先关系如何，仍按原来登记的实现的声明决定，替换者自己的 `applies`、`general` 与 `preferred_over` 不参与选择； +- 这个 key 在原来的实现或替换者能运行的设备上可用。替换者能在原实现不能运行的设备上运行时，这个 key 在那些设备上也参与选择，与其他实现重叠时按同样的规则报歧义； +- 该 key 被选中时，这次调用由替换者服务；替换者不能在调用设备上运行或拒绝这次调用时报错，不退回原来的实现； +- 其他 key 被选中的调用不询问替换者； +- 替换者与登记的实现遵守同一条规则：继承这个 key 所属的 kernel 接口，由它自己的类方法 `entry_for(call)` 构建。没有其他写法。 + +要改变一个 kernel 服务哪些调用，用第 3 节的 `register_implementation`。 + +```python +# tests/test_kernel_dispatch.py +class _TorchLayerNorm(Kernel, LayerNormFwdInterface): + """A replacement written against ``LayerNormFwdInterface`` alone.""" + + devices = frozenset({torch.device(run_device()).type}) + + def __init__(self, n: int, eps: float) -> None: + super().__init__() + self.n, self.eps = n, eps + + @classmethod + def entry_for(cls, call: LayerNormCall): + return (call.n, call.eps), lambda: cls(call.n, call.eps) + + def forward(self, x, weight, bias): + return F.layer_norm(x.float(), (self.n,), weight.float(), bias.float(), self.eps).to( + x.dtype + ) + +op = LayerNormFwdOp((32,), kernel_map={"layer_norm": _TorchLayerNorm}, target=BUILTIN) +``` + +`kernel_map=` 中本 op 没有、但其他 op 有的 key 被忽略，复合 op 由此把一份映射传给各个子 op；任何 op 都没有的 key 在构造时报 `was given kernel_map keys no op has`。 + +## 3. 新增实现：register_implementation {#register} + +`tileops.backend.register_implementation(op, key, implementation)` 为 op 新增一个实现，是[如何为 op 新增 kernel](writing.md) 第一步的 backend 形式。`op` 是 op 的类名，`key` 是新实现的名字，`implementation` 继承 op 的某个 kernel 接口，属于哪个接口由继承关系决定。第二步与 in-tree 实现相同。 + +```python +# tests/test_kernel_dispatch.py +class _NarrowTorchLayerNorm(_TorchLayerNorm): + """An added implementation for short rows, which wins over the in-tree one there.""" + + preferred_over = frozenset({"layer_norm"}) + + @classmethod + def applies(cls, call: LayerNormCall) -> bool: + return call.n <= 64 + +register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm) +``` + +- 新实现在 `n <= 64` 上与 `LayerNormKernel` 重叠，后者不是 general，因此新实现声明 `preferred_over`。 +- 新实现不服务的调用仍由 in-tree 实现服务，例如 `n = 1024`。 +- 新实现只进入注册之后构造的 op 实例。 +- 同一个 op 下重复注册同一个 key 报 `BackendError`；key 与 in-tree 的 key 相同时，构造 op 时报 `reuse keys it has`。 + +注册在 backend 模块被导入时完成。backend 在 `pyproject.toml` 中声明 entry point，TileOPs 在构造第一个 op 时导入它： + +```toml +[project.entry-points."tileops.backends"] +acme = "tileops_acme" +``` + +模块导入失败时，它的全部注册被撤销，TileOPs 发出 `RuntimeWarning`，失败记录可以通过 `tileops.backend.load_failures()` 查看。 + +## 4. 替换整个 op：target {#target} + +target 协议的完整写法、一个可运行的模板 backend 与常见报错，见[接入新的硬件 backend](../../backends.md)。本节只概述它与另外两种方式的区别。 + +target 是 backend 为一组 kernel 取的名字。target 为某个 op 注册 builder 之后，选中该 target 的 op 实例的调用都由 target 服务，op 自己的 `forward` 不执行。一次调用写入的所有张量都为空时例外：这时 target 与 in-tree 实现都不执行，输出按签名构造。 + +```python +from tileops.backend import TensorSpec, register_detector, register_kernel_builder +from .kernels import AcmeRMSNorm + +register_detector(target="acme", detect=lambda device: device.type == "acme") + +def build_rms_norm(x: TensorSpec, weight: TensorSpec | None, *, normalized_shape, eps): + return AcmeRMSNorm(normalized_shape, eps, x.dtype) + +register_kernel_builder(op="RMSNormFwdOp", target="acme", build_kernel=build_rms_norm) +``` + +- `build_kernel` 按 op 的 `signature.inputs` 顺序接收每个输入的 `TensorSpec`（未传入的可选输入为 `None`，例如 `RMSNormFwdOp` 的 `weight`），并以关键字接收 `signature.params`。它返回的 kernel 在调用时按同样的顺序接收输入张量，并以关键字接收调用方传入的输出缓冲与执行参数。导入模块时不编译任何东西，构建在 `build_kernel` 被调用时发生。 +- op 实例按调用设备，以及每个输入是否传入、dtype 与 shape，缓存 `build_kernel` 返回的 kernel。 +- op 实例在第一次需要执行实现的调用中确定 target，此后不变：构造参数 `target=` 优先，其次是 `set_default_target` 设置的进程默认值，最后由各 target 的 detector 检测调用设备；没有 target 认领时使用 in-tree 实现，多个 target 认领时报 `AmbiguousTargetError`。写入张量都为空的调用不确定 target；确定 target 的那次调用失败时，选择被撤销，下一次调用重新确定。 +- 传给 target 的调用已经通过由签名生成的检查；除声明了 `device: cpu` 的张量外，所有张量位于调用设备上；所有不被写入的输入都是连续的，被写入的输入在声明了 `contiguous: true` 时是连续的。 +- 选中的 target 没有为这个 op 注册 builder、而 op 自己持有 kernel 时，调用报 `OpNotAvailableError`，不退回 in-tree 实现。 diff --git a/docs/user-guide/dispatch/img/call.svg b/docs/user-guide/dispatch/img/call.svg new file mode 100644 index 00000000..cf7a6bdb --- /dev/null +++ b/docs/user-guide/dispatch/img/call.svg @@ -0,0 +1 @@ +Written by the developerProvided by the systemWritten by the developerOp baseOp classOp baseCallSpecImplementations ofentryOp class_eager_forwardOp class_eager_forwardOp basekernel_forOp basekernel_forCallSpecdevice factsCallSpecdevice factsImplementations ofthe kernel interfaceImplementations ofthe kernel interfaceentryentryOp baseBuild the call specfrom the tensorsand the op's paramskernel_for(interface,call)Look up (interface,call) in the dispatchcachealt[Hit]the resolved entry[Miss]Check the call's typeand hashabilityloop[Each implementation of the interface]unavailable(call),then refusal(call) ifavailableRead arch,sm_count,calibration,smem_budgetResolved fromcall.device on firstreada reason, or NoneTake the one winnerby precedenceamong theavailable,applicable ones; fora key replacedthroughkernel_map=, checkthe replacementserves the callentry_for(call)(build identity,factory)opt[This (implementation class, build identity)is not built yet]Run the factory withcall.device currentRequest tuning onthe entry when theop tunes; store inthe dispatch cacheentryentry(tensors)result diff --git a/docs/user-guide/dispatch/img/call.zh.svg b/docs/user-guide/dispatch/img/call.zh.svg new file mode 100644 index 00000000..519a4b9f --- /dev/null +++ b/docs/user-guide/dispatch/img/call.zh.svg @@ -0,0 +1 @@ +开发者编写系统提供开发者编写Op ..Op .Op ..CallSpeckernel ...entryOp 类_eager_forwardOp 类_eager_forwardOp 基类kernel_forOp 基类kernel_forCallSpec设备事实CallSpec设备事实kernel 接口的实现kernel 接口的实现entryentryOp ..由张量与 op 的参数构造 call speckernel_for(interface,call)以 (interface, call)查派发缓存alt[命中]已解析的 entry[未命中]检查 call 的类型与可哈希性loop[接口的每个实现]unavailable(call)，可用时再问refusal(call)读取 arch、sm_count、calibration、smem_budget首次读取时从call.device 解析理由，或 None在可用且适用的实现中按优先关系取唯一胜者；key 被kernel_map= 替换时，检查替换者能服务这次调用entry_for(call)(build identity, 构建函数)opt[这个 (实现类, build identity)尚未构建]以 call.device 为当前设备运行构建函数op 处于调参模式时对entry 请求调参；写入派发缓存entryentry(张量)结果 diff --git a/docs/user-guide/dispatch/img/extension.svg b/docs/user-guide/dispatch/img/extension.svg new file mode 100644 index 00000000..2489757d --- /dev/null +++ b/docs/user-guide/dispatch/img/extension.svg @@ -0,0 +1 @@ +target: replaces the whole opThe op instance's in-tree implementationImplementations of one kernel interface(installed at construction)build_kernel registered byregister_kernel_builderThe Op class's _eager_forwardClasses in kernel_typesinheriting the interfacekernel_map=replaces the class behind a key,keeping the key's rule(this instance only)register_implementationadds a key(every instance built afterwards)Generated call checksa target registereda builder for the opserved in-treekernel_for selects amongthese by the selection ruleCyan: provided by the systemPurple: written by TileOPs developersGreen: provided by a backend diff --git a/docs/user-guide/dispatch/img/extension.zh.svg b/docs/user-guide/dispatch/img/extension.zh.svg new file mode 100644 index 00000000..5b59e94b --- /dev/null +++ b/docs/user-guide/dispatch/img/extension.zh.svg @@ -0,0 +1 @@ +target：替换整个 opop 实例的 in-tree 实现一个 kernel 接口的实现（构造时安装）register_kernel_builder注册的 build_kernelOp 类的 _eager_forwardkernel_types 中继承该接口的类kernel_map=替换 key 背后运行的类沿用该 key 的规则（只作用于这个实例）register_implementation新增的 key（此后构造的所有实例）生成的调用检查target 为该 op注册了 builder由 in-tree 实现服务kernel_for 按选择规则在这些实现中选择青色：系统提供紫色：TileOPs 开发者编写绿色：backend 提供 diff --git a/docs/user-guide/dispatch/index.md b/docs/user-guide/dispatch/index.md new file mode 100644 index 00000000..f6856504 --- /dev/null +++ b/docs/user-guide/dispatch/index.md @@ -0,0 +1,74 @@ +# How an op selects a kernel + +A TileOPs op can have several kernels, and each kernel serves part of the op's calls. This guide describes how an op selects a kernel on a call, and what a developer writes to add a kernel. It is written for two readers: developers who add a kernel or an op to TileOPs, and authors of third-party backends who replace or add kernels. The reasons behind each design choice are in the design document [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection). + +This page describes the selection process and its terms. The other two pages are: + +- [Adding a kernel to an op](writing.md): the two steps, the default rules, opening a new kernel interface, tests, and common errors; +- [How a backend joins TileOPs](backends.md): `kernel_map=`, `register_implementation`, and targets. + +## 1. Who selects the kernel: the op selects an interface, dispatch selects an implementation {#two-levels} + +The op first selects a **kernel interface** by semantics. Dispatch then selects one of that interface's **implementations** by the call's shape, dtype, and device. + +**Table 1** The op selects an interface, dispatch selects an implementation + +| No. | Level | Decided by | Based on | Written in | +| --- | --- | --- | --- | --- | +| 1 | Select a kernel interface | the op | semantics, that is, where the call contracts differ, such as training versus inference | the op's `_eager_forward`, which names the interface when it calls `kernel_for` | +| 2 | Select an implementation | dispatch, shared by all ops | the availability, applicability, and precedence each implementation declares | the selection algorithm is in `src/tileops/ops/op_base.py`; each implementation's declarations are on its own kernel class | + +Take `BatchNormFwdOp` as an example. The op selects the interface `batch_norm_fwd_train` or `batch_norm_fwd_infer` by its construction argument `training`. The training interface has four implementations, and dispatch selects one of them by batch size, channel count, and spatial size. + +## 2. How a call finds its kernel {#call-path} + +On every call, the op builds a **call spec**, which records the facts of this call that affect selection and building. It then calls `self.kernel_for(interface, call)`: + +```python +# src/tileops/ops/norm/layer_norm.py +call = LayerNormCall(device=x.device, n=math.prod(ns), eps=self.eps, dtype=x.dtype) +return self.kernel_for("layer_norm", call)(x, weight, bias) +``` + +![The path of one kernel_for call](img/call.svg) + +**Figure 1** The hit and miss paths of one `kernel_for` call. Purple marks the parts written by the developer; teal marks the parts provided by the system. + +- **Hit**: one cache lookup by `(interface, call)` returns the resolved entry directly. No implementation is consulted and no device property is read. +- **Miss**: dispatch selects in the following order and writes the result to the cache: + 1. Drop the implementations that cannot run on the call's device, based on `devices` and `supported_archs`. + 2. Drop the implementations that do not serve this call, based on `applies` or `refusal`. + 3. Pick the single winner among the rest: a `general` implementation ranks below every other implementation, and the others are compared by `preferred_over`. + 4. Call the winner's `entry_for(call)` to get a build identity and a factory, then build or reuse the entry. Within one interface, calls share an entry only when they select the same implementation class with the same build identity. An entry can contain one or more kernels. + +When no key can run on the call's device type, the call raises `in-tree kernels do not run on` (`OpNotAvailableError`). When some key supports the device type but no implementation is both available and applicable, it raises `no implementation serves this call`. When several implementations remain with no precedence among them, it raises `dispatch is ambiguous`. Selection does not depend on declaration order, and there are no numeric priorities. + +## 3. Two things to do when adding a kernel {#hooks} + +To add a kernel to an existing kernel interface, TileOPs developers and backend authors both do the following two things: + +**Table 2** The two steps of adding a kernel + +| No. | Step | TileOPs developer | Backend author | Described in | +| --- | --- | --- | --- | --- | +| 1 | Register the implementation | the class subclasses the kernel interface and is added to the op's `kernel_types` | the class subclasses the kernel interface; call `register_implementation(op, key, cls)` | [Adding a kernel to an op § 1](writing.md#register) | +| 2 | Declare which calls the implementation serves | write `applies`; write `preferred_over` only when it overlaps another non-general implementation | same as the TileOPs developer | [Adding a kernel to an op § 2](writing.md#rule) | + +An implementation that declares none of `devices`, `supported_archs`, `applies`, `general`, and `preferred_over` is available on CUDA devices of every architecture, serves every call, and has no precedence relation with other implementations. An interface with a single implementation therefore only needs the class to subclass the interface. + +A backend can also replace the implementation behind an existing key, or replace a whole op; see [How a backend joins TileOPs](backends.md). + +## 4. Terms used in this guide {#terms} + +**Table 3** Terms + +| No. | Term | Meaning | Written in code as | +| --- | --- | --- | --- | +| 1 | kernel interface | a place where an op calls a kernel, which fixes the call contract at that place | a subclass of `KernelInterface`; the op's `interfaces` | +| 2 | implementation | a kernel class that subclasses a kernel interface, registered under a key | a `Kernel` subclass; a key of `kernel_types` | +| 3 | call spec | the immutable facts of one call that affect selection and building | a frozen dataclass subclass of `CallSpec` | +| 4 | device facts | the architecture, the SM count, the calibrated board model the device belongs to, and the shared memory limit per block, read from the call's device only on a miss | `arch`, `sm_count`, `calibration`, `smem_budget` | +| 5 | availability | which devices the implementation can run on | `devices`, `supported_archs` | +| 6 | applicability | which calls the implementation serves | `applies`, `refusal` | +| 7 | precedence | which implementation is selected when several are available and applicable | `general`, `preferred_over` | +| 8 | build identity and entry | the value that decides whether two builds are the same, and the entry built from it, which contains one or more kernels | the pair returned by `entry_for` | diff --git a/docs/user-guide/dispatch/index.zh.md b/docs/user-guide/dispatch/index.zh.md new file mode 100644 index 00000000..0d586a5d --- /dev/null +++ b/docs/user-guide/dispatch/index.zh.md @@ -0,0 +1,74 @@ +# op 如何选择 kernel + +一个 TileOPs op 可以有多个 kernel，每个 kernel 服务一部分调用。本指南说明 op 在一次调用中如何选中 kernel，以及新增一个 kernel 要写什么。本指南面向两类读者：在 TileOPs 中新增 kernel 或 op 的开发者，以及为第三方 backend 替换或新增 kernel 的作者。各项设计的理由见设计文档 [Op Interface Design § Kernel selection](../../design/ops-design.md#kernel-selection)。 + +本页说明选择的过程与术语，另外两页分别是： + +- [如何为 op 新增 kernel](writing.md)：两步的写法、默认规则、新开 kernel 接口、测试与常见报错； +- [backend 如何接入 TileOPs](backends.md)：`kernel_map=`、`register_implementation` 与 target。 + +## 1. 谁来选 kernel：op 选接口，派发机制选实现 {#two-levels} + +op 先按语义选定一个 **kernel 接口**，派发机制再按调用的 shape、dtype 与设备，从这个接口的**实现**中选出一个实现。 + +**表 1** op 选接口与派发机制选实现 + +| No. | 层 | 由谁决定 | 依据 | 写在哪里 | +| --- | --- | --- | --- | --- | +| 1 | 选 kernel 接口 | op | 语义，即调用契约不同之处，例如训练与推理 | op 的 `_eager_forward`，调用 `kernel_for` 时给出接口名 | +| 2 | 选实现 | 派发机制，所有 op 共用 | 每个实现声明的可用性、适用范围与优先关系 | 选择算法在 `src/tileops/ops/op_base.py`；各实现的声明写在自己的 kernel 类上 | + +以 `BatchNormFwdOp` 为例：op 按构造参数 `training` 选择接口 `batch_norm_fwd_train` 或 `batch_norm_fwd_infer`；训练接口有四个实现，派发机制按 batch、channel 数与空间尺寸选中其中一个。 + +## 2. 一次调用如何找到 kernel {#call-path} + +op 每次调用构造一个 **call spec**，它记录这次调用中影响选择与构建的事实，然后调用 `self.kernel_for(interface, call)`： + +```python +# src/tileops/ops/norm/layer_norm.py +call = LayerNormCall(device=x.device, n=math.prod(ns), eps=self.eps, dtype=x.dtype) +return self.kernel_for("layer_norm", call)(x, weight, bias) +``` + +![一次 kernel_for 调用的路径](img/call.svg) + +**图 1** 一次 `kernel_for` 调用的命中与未命中路径。紫色为开发者编写的部分，青色为系统提供的部分。 + +- **命中**：以 `(interface, call)` 查一次缓存，直接返回已解析的 entry，不询问任何实现，也不读取设备属性。 +- **未命中**：派发机制按以下顺序选择，结果写入缓存： + 1. 去掉在调用设备上不能运行的实现，依据 `devices` 与 `supported_archs`； + 2. 去掉不服务这次调用的实现，依据 `applies` 或 `refusal`； + 3. 在剩下的实现中选唯一的胜者：`general` 的实现低于其他所有实现，其余按 `preferred_over` 比较； + 4. 调用胜者的 `entry_for(call)`，得到 build identity 与构建函数，再构建或复用 entry。在同一个 interface 内，只有选中同一实现类且 build identity 相同的调用才共用 entry；一个 entry 可以包含一个或多个 kernel。 + +没有任何 key 能在调用设备类型上运行时，报 `in-tree kernels do not run on`（`OpNotAvailableError`）；有 key 支持这种设备类型、但没有实现同时可用且适用时，报 `no implementation serves this call`；剩下多个互相没有优先关系的实现时，报 `dispatch is ambiguous`。选择不依赖声明顺序，也没有数值优先级。 + +## 3. 新增一个 kernel 要做的两件事 {#hooks} + +为已有的 kernel 接口新增一个 kernel，TileOPs 开发者与 backend 作者都做以下两件事： + +**表 2** 新增一个 kernel 的两步 + +| No. | 步骤 | TileOPs 开发者 | backend 作者 | 说明所在 | +| --- | --- | --- | --- | --- | +| 1 | 登记实现 | 类继承 kernel 接口，加入 op 的 `kernel_types` | 类继承 kernel 接口，调用 `register_implementation(op, key, cls)` | [新增 kernel 1](writing.md#register) | +| 2 | 声明实现服务哪些调用 | 写 `applies`；只在与另一个非 general 实现重叠时写 `preferred_over` | 同左 | [新增 kernel 2](writing.md#rule) | + +实现不声明 `devices`、`supported_archs`、`applies`、`general` 与 `preferred_over` 时，默认在所有架构的 CUDA 设备上可用，服务全部调用，与其他实现没有优先关系。只有一个实现的接口因此只需继承接口。 + +backend 还可以替换一个已有 key 背后的实现，或替换整个 op，见 [backend 如何接入 TileOPs](backends.md)。 + +## 4. 本指南使用的术语 {#terms} + +**表 3** 术语 + +| No. | 术语 | 含义 | 代码中的写法 | +| --- | --- | --- | --- | +| 1 | kernel 接口 | op 调用 kernel 的一个位置，规定这个位置的调用契约 | `KernelInterface` 的子类；op 的 `interfaces` | +| 2 | 实现 | 继承某个 kernel 接口的 kernel 类，以一个 key 登记 | `Kernel` 子类；`kernel_types` 的键 | +| 3 | call spec | 一次调用中影响选择与构建的不可变事实 | `CallSpec` 的冻结 dataclass 子类 | +| 4 | 设备事实 | 架构、SM 数、设备所属的标定板卡型号，以及每个 block 可用的共享内存上限，未命中时才从调用设备读取 | `arch`、`sm_count`、`calibration`、`smem_budget` | +| 5 | 可用性 | 实现能在哪些设备上运行 | `devices`、`supported_archs` | +| 6 | 适用范围 | 实现服务哪些调用 | `applies`、`refusal` | +| 7 | 优先关系 | 多个实现都可用且都适用时谁被选中 | `general`、`preferred_over` | +| 8 | build identity 与 entry | 决定两次构建是否相同的值，以及构建出的 entry，它包含一个或多个 kernel | `entry_for` 返回的二元组 | diff --git a/docs/user-guide/dispatch/writing.md b/docs/user-guide/dispatch/writing.md new file mode 100644 index 00000000..d3772c21 --- /dev/null +++ b/docs/user-guide/dispatch/writing.md @@ -0,0 +1,201 @@ +# Adding a kernel to an op + +This page describes the two steps of adding a kernel to an existing op, the default rules when nothing is declared, when to open a new kernel interface, and how to test the selection. The overall selection process is in [How an op selects a kernel](index.md). + +## 1. Step 1: register the implementation with the op {#register} + +An implementation is a class that subclasses both `Kernel` (or one of its subclasses) and a kernel interface. A TileOPs developer adds it to the op's `kernel_types` under a snake_case key. A backend author calls `register_implementation`; see [How a backend joins TileOPs § 3](backends.md#register). + +```python +# src/tileops/ops/norm/batch_norm.py +class BatchNormFwdOp(Op): + kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { + "fwd_train_whole": BatchNormFwdTrainWholeKernel, + "fwd_train_wide": BatchNormFwdTrainWideKernel, + "fwd_train_split": BatchNormFwdTrainSplitKernel, + "fwd_train_kernel": BatchNormFwdTrainKernel, + "fwd_infer_kernel": BatchNormFwdInferKernel, + } + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = { + "batch_norm_fwd_train": BatchNormTrainFwdInterface, + "batch_norm_fwd_infer": BatchNormInferFwdInterface, + } +``` + +A key belongs to the interface its class subclasses; no separate mapping is written. The implementation class also satisfies the following: + +- `forward` accepts all parameters of the interface's `forward`, positionally; +- the class method `entry_for(call)` returns `(build identity, factory)`. The build identity contains every fact that changes the build result and must be hashable. The factory runs only the first time that identity appears. The default implementation uses the whole call spec as the build identity and builds with `cls(call)`; +- the constructor is the implementation's own choice; the op never constructs an implementation directly; +- tuning does not go through the factory; the op applies it to the built kernel on the call's device. + +## 2. Step 2: declare which calls the implementation serves {#rule} + +An implementation's applicability is written in the class method `applies(call)`. It describes only the calls this implementation serves, not other implementations. When an error message needs to state the reason for a refusal, override `refusal(call)`: it returns `None` when the implementation applies and the reason otherwise, and `applies` then returns `cls.refusal(call) is None`. A check shared by several implementations is written as a property of the family's call spec (for example `AttentionCall.dense_decode_region`); a check shared within one inheritance chain is written as a class method of the base class. + +When a new implementation's applicability overlaps an existing implementation, what it overlaps decides whether precedence is declared: + +**Table 1** Whether to declare precedence + +| No. | The new implementation overlaps | Declaration | +| --- | --- | --- | +| 1 | no implementation | none | +| 2 | only the general implementation | none; the general implementation ranks below every other implementation | +| 3 | another non-general implementation | the side that should win writes `preferred_over = frozenset({""})` | + +- `general = True` marks the fallback: it serves the calls no other implementation serves. Each interface has at most one. +- `preferred_over` states only which side wins when both apply. It does not require one side's applicability to be contained in the other's. +- An implementation does not exclude another implementation's range in its own `applies`. When it should yield, the other implementation declares `preferred_over`. +- An undeclared overlap raises an ambiguity error at call time; it is never decided silently by order. + +The three GQA dense decode implementations are case 3. bs1 serves calls with batch 1, and long-context serves calls with a long KV. The two overlap when batch is 1 and KV is long, and long-context declares that it wins: + +```python +# src/tileops/kernels/attention/gqa_decode.py +class GQADecodeLongContextKernel(GQADecodeKernel): + general: bool = False + preferred_over = frozenset({"gqa_dense_decode_bs1"}) + + @classmethod + def refusal(cls, call: AttentionCall) -> Optional[str]: + served = ( + call.dense_decode_region + and not call.fuse_rope + and call.seqlen_kv >= 1024 + and call.batch == 1 + and call.heads == 32 + and call.heads_kv == 4 + and call.dim == 128 + and call.dtype == torch.float16 + and call.softcap == 0.0 + ) + return None if served else "does not serve this call" +``` + +Two non-general implementations need a precedence declaration as soon as their applicability intersects; one range does not have to contain the other. FP8 decode and the generic FP8 implementation overlap on part of the calls. FP8 decode declares `preferred_over = frozenset({"gqa_dense_fp8"})` and wins inside the intersection; the generic FP8 implementation does not exclude the decode range in its own `applies`. + +`GQADecodeKernel` is general and supports SM80/89/90; bs1 supports only SM90. Availability filters before precedence is compared, so on SM80 bs1 takes no part in selection: + +**Table 2** Selection for batch-1 decode on different architectures + +| No. | Architecture | KV length | Available and applicable implementations | Selected | +| --- | --- | --- | --- | --- | +| 1 | SM90 | 4096 | decode, bs1, long-context | long-context | +| 2 | SM90 | 512 | decode, bs1 | bs1 | +| 3 | SM80 | 4096 | decode, long-context | long-context | +| 4 | SM80 | 512 | decode | decode | + +## 3. Default selection rules when an implementation declares nothing {#default} + +**Table 3** Defaults when nothing is declared + +| No. | Member | Default | Effect | +| --- | --- | --- | --- | +| 1 | `devices` | `frozenset({"cuda"})` | available on CUDA devices | +| 2 | `supported_archs` | `None` | available on every architecture | +| 3 | `applies` | returns `True` | serves every call | +| 4 | `general` | `False` | does not rank below other implementations | +| 5 | `preferred_over` | empty set | wins over no implementation | + +An interface with a single implementation therefore only needs the class to subclass the interface. `LayerNormKernel` declares only its supported architectures and its own `entry_for`: + +```python +# src/tileops/kernels/norm/layer_norm.py +class LayerNormKernel(Kernel, LayerNormFwdInterface): + supported_archs: list[int] = [80, 86, 89, 90] + + @classmethod + def entry_for(cls, call: LayerNormCall) -> Entry: + identity = (call.n, call.eps, call.dtype) + return identity, lambda: cls(*identity) +``` + +## 4. When a new kernel needs a new kernel interface {#interface} + +A new kernel that changes the semantics or the call contract, for example by returning different values, belongs to a new kernel interface. A new kernel that differs only in shape, dtype, architecture, or performance is a new implementation of an existing interface. The same algorithm with a different set of tile sizes or split counts is a parameter inside one implementation, not a new implementation. + +**Table 4** Where common cases belong + +| No. | Case | Belongs to | +| --- | --- | --- | +| 1 | BatchNorm training and inference: the statistics come from different sources, and training also returns `mean` and `rstd` | two interfaces | +| 2 | GLA prefill and decode: the call contract is the same, only the sequence length differs | different implementations of one interface | +| 3 | a faster algorithm on some shape range or some architecture | a new implementation | +| 4 | the same algorithm with a split count of 1 or greater than 1 | a parameter inside the implementation | + +A new interface needs a call spec type and an interface class. When the family has several kernel files, both are written in `src/tileops/kernels//call_spec.py`; when it has a single kernel file, both are written in that kernel file. Interface classes are named `{Name}{Fwd|Bwd}Interface`, with the word for the variant placed before the direction suffix, for example `BatchNormTrainFwdInterface`; the pre-commit hook `interface-names-lint` checks this. Several interfaces of one family can share a call spec type; for example, the six interfaces of BatchNorm and InstanceNorm all use `BatchNormCall`. + +```python +# src/tileops/kernels/norm/call_spec.py +@dataclasses.dataclass(frozen=True) +class LayerNormCall(CallSpec): + """The facts that select an implementation normalizing trailing rows and build it. + + Layer, RMS, fused-add and adaptive layer normalization take it. ``n`` is the row's + element count and ``eps`` the op's epsilon. + """ + + n: int = 0 + eps: float = 1e-5 + dtype: torch.dtype = torch.float16 + +class LayerNormFwdInterface(KernelInterface): + """Layer normalization over the trailing ``call.n`` elements.""" + + request = LayerNormCall + + @abstractmethod + def forward(self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + """Normalize each run of ``call.n`` trailing elements; nothing is written in place. + + Every tensor is contiguous on ``call.device``. + + Args: + x: Any shape whose trailing axes hold ``call.n`` elements, in ``call.dtype``. + weight: ``call.n`` elements of scale in ``call.dtype``. + bias: ``call.n`` elements of shift in ``call.dtype``. + + Returns: + A new output shaped like *x*, in ``call.dtype``. + """ +``` + +- The call spec's fields are the call facts that implementations read in `applies`, `refusal`, and `entry_for`: shapes, dtypes, semantic parameters (including those fixed when the op is constructed), and `device`. Fields hold only immutable values. They do not hold tensor contents, tuning policy, or priorities; device facts are provided by `CallSpec`. +- The interface class's `request` points to the call spec type. The parameter list of the abstract method `forward` is the arguments the op passes when it calls the kernel. Its docstring states, for each tensor, the shape, dtype, memory layout, device, and whether it is written in place, and it states the return value. A backend writes its implementation from this contract alone. +- Methods other than `forward` that the op calls on the entry are also written as abstract methods of the interface. An implementation must implement every abstract method of the interface; otherwise the class cannot be instantiated. +- Values that every implementation must agree on are written as ordinary class methods of the interface and computed by the interface. For example, when the op allocates an output buffer by size before selection, a class method of the interface computes that size from the call spec; the size is not read from the selected implementation. +- The docstring states only what the implementation actually does. If it declares the inputs contiguous, the op calls `.contiguous()` before the call. Behavior the implementation does not guarantee (for example, the result when an input contains NaN) is stated as undefined, not as a guarantee. +- Every op that holds kernels declares `interfaces`, does not override `entry_for`, and does not keep its own kernel cache. + +## 5. How to test the selection, and common errors {#tests} + +Write one test case for each applicability range and for each boundary between adjacent ranges. Check the selected key with `select_implementation`, and give the device facts explicitly so that the test does not depend on the machine that runs it: + +```python +# tests/ops/test_batch_norm.py +def test_each_region_selects_its_one_implementation( + op_cls, interface, n, c, spatial, dtype, key +) -> None: + """Exactly one non-general implementation, or else the general one, serves each shape.""" + op = op_cls() + call = BatchNormCall(arch=90, sm_count=132, n=n, c=c, spatial=spatial, dtype=dtype) + assert op.select_implementation(interface, call) == key +``` + +**Table 5** Common errors + +| No. | When | Error message fragment | Cause | +| --- | --- | --- | --- | +| 1 | constructing the op | `does not implement ` | the class that runs behind a key (including a `kernel_map=` replacement) does not subclass the interface; make it subclass the interface and build through `entry_for(call)` | +| 2 | constructing the op | `forward does not take 's arguments` | `forward`'s parameters do not match the interface | +| 3 | constructing the op | `has more than one general implementation` | an interface has two general implementations | +| 4 | constructing the op | `preferences form a cycle through` | `preferred_over` forms a cycle | +| 5 | constructing the op | `keys implement none of its kernel interfaces` | a class in `kernel_types` subclasses none of the op's kernel interfaces | +| 6 | calling | `in-tree kernels do not run on` | no key can run on the call's device type (`OpNotAvailableError`) | +| 7 | calling | `no implementation serves this call` | some key supports the device type, but no implementation is both available and applicable; the message lists each implementation's reason | +| 8 | calling | `dispatch is ambiguous` | several implementations apply, and none has precedence over the others | +| 9 | calling | `takes a call spec` | the call spec passed to `kernel_for` has the wrong type | +| 10 | calling | `cannot key a dispatch cache` | the call spec has a mutable field, such as a list | +| 11 | calling | `this call spec states ['arch']` | the call spec passed to `kernel_for` gives device facts explicitly; they are derived from `device`, and only `select_implementation` accepts explicit device facts | + +When a new op holds kernels without declaring `interfaces`, the inventory test in `tests/test_kernel_dispatch.py` raises `declare interfaces instead`. diff --git a/docs/user-guide/dispatch/writing.zh.md b/docs/user-guide/dispatch/writing.zh.md new file mode 100644 index 00000000..275dad11 --- /dev/null +++ b/docs/user-guide/dispatch/writing.zh.md @@ -0,0 +1,201 @@ +# 如何为 op 新增 kernel + +本页说明为已有 op 新增一个 kernel 的两步、什么都不声明时的默认规则、何时新开 kernel 接口，以及如何测试选择结果。选择的整体过程见[首页](index.md)。 + +## 1. 第一步：把实现登记到 op {#register} + +实现是同时继承 `Kernel`（或它的子类）与一个 kernel 接口的类。TileOPs 开发者把它加入 op 的 `kernel_types`，key 使用 snake_case；backend 作者调用 `register_implementation`，见 [backend 接入 3](backends.md#register)。 + +```python +# src/tileops/ops/norm/batch_norm.py +class BatchNormFwdOp(Op): + kernel_types: ClassVar[Mapping[str, type[Kernel]]] = { + "fwd_train_whole": BatchNormFwdTrainWholeKernel, + "fwd_train_wide": BatchNormFwdTrainWideKernel, + "fwd_train_split": BatchNormFwdTrainSplitKernel, + "fwd_train_kernel": BatchNormFwdTrainKernel, + "fwd_infer_kernel": BatchNormFwdInferKernel, + } + interfaces: ClassVar[Mapping[str, type[KernelInterface]]] = { + "batch_norm_fwd_train": BatchNormTrainFwdInterface, + "batch_norm_fwd_infer": BatchNormInferFwdInterface, + } +``` + +一个 key 属于它的类所继承的接口，不需要另写映射。实现类还要满足： + +- `forward` 按位置接受接口 `forward` 的全部参数； +- 类方法 `entry_for(call)` 返回 `(build identity, 构建函数)`。build identity 包含一切会改变构建结果的事实，必须可哈希，构建函数只在它第一次出现时运行。默认实现以整个 call spec 为 build identity、以 `cls(call)` 构建； +- 构造函数由实现自行决定，op 从不直接构造实现； +- 调参不经过构建函数，由 op 在调用的设备上作用于构建出的 kernel。 + +## 2. 第二步：声明实现服务哪些调用 {#rule} + +实现的适用范围写在类方法 `applies(call)` 中，它只描述本实现服务的调用，不描述其他实现。需要在报错中说明拒绝的原因时，覆写 `refusal(call)`：适用时返回 `None`，否则返回原因；同时让 `applies` 返回 `cls.refusal(call) is None`。多个实现共用的判断写成 family call spec 的属性（例如 `AttentionCall.dense_decode_region`），同一继承链内共用的判断写成基类的类方法。 + +新实现的适用范围与已有实现重叠时，按重叠的对象决定是否声明优先关系： + +**表 1** 是否声明优先关系 + +| No. | 新实现与谁重叠 | 声明 | +| --- | --- | --- | +| 1 | 不与任何实现重叠 | 无 | +| 2 | 只与 general 实现重叠 | 无，general 实现低于其他所有实现 | +| 3 | 与另一个非 general 实现重叠 | 在应当胜出的一方写 `preferred_over = frozenset({"<另一方的 key>"})` | + +- `general = True` 表示兜底：服务其他实现都不服务的调用，每个接口至多一个。 +- `preferred_over` 只说明两者都适用时谁胜出，不要求一方的适用范围包含在另一方之内。 +- 一个实现不在自己的 `applies` 中排除另一个实现的范围；需要让出时，由另一个实现声明 `preferred_over`。 +- 没有声明的重叠在调用时报歧义，不会被按顺序悄悄决定。 + +GQA dense decode 的三个实现是第 3 种情况。bs1 服务 batch 为 1 的调用，long-context 服务 KV 较长的调用，二者在 batch 为 1 且 KV 较长时重叠，由 long-context 声明它胜出： + +```python +# src/tileops/kernels/attention/gqa_decode.py +class GQADecodeLongContextKernel(GQADecodeKernel): + general: bool = False + preferred_over = frozenset({"gqa_dense_decode_bs1"}) + + @classmethod + def refusal(cls, call: AttentionCall) -> Optional[str]: + served = ( + call.dense_decode_region + and not call.fuse_rope + and call.seqlen_kv >= 1024 + and call.batch == 1 + and call.heads == 32 + and call.heads_kv == 4 + and call.dim == 128 + and call.dtype == torch.float16 + and call.softcap == 0.0 + ) + return None if served else "does not serve this call" +``` + +两个非 general 实现的适用范围只要有交集，就需要声明优先关系，不要求一个范围整体包含另一个。FP8 decode 与通用 FP8 实现在一部分调用上重叠，FP8 decode 声明 `preferred_over = frozenset({"gqa_dense_fp8"})`，在交集内胜出；通用 FP8 实现不在自己的 `applies` 中排除 decode 的范围。 + +`GQADecodeKernel` 是 general，支持 SM80/89/90；bs1 只支持 SM90。可用性在比较优先关系之前过滤，因此 SM80 上 bs1 不参与选择： + +**表 2** batch 为 1 的 decode 在不同架构上的选择 + +| No. | 架构 | KV 长度 | 可用且适用的实现 | 选中 | +| --- | --- | --- | --- | --- | +| 1 | SM90 | 4096 | decode、bs1、long-context | long-context | +| 2 | SM90 | 512 | decode、bs1 | bs1 | +| 3 | SM80 | 4096 | decode、long-context | long-context | +| 4 | SM80 | 512 | decode | decode | + +## 3. 实现什么都不声明时的默认选择规则 {#default} + +**表 3** 未声明时的默认值 + +| No. | 成员 | 默认值 | 效果 | +| --- | --- | --- | --- | +| 1 | `devices` | `frozenset({"cuda"})` | 在 CUDA 设备上可用 | +| 2 | `supported_archs` | `None` | 在所有架构上可用 | +| 3 | `applies` | 返回 `True` | 服务全部调用 | +| 4 | `general` | `False` | 不低于其他实现 | +| 5 | `preferred_over` | 空集 | 不胜过任何实现 | + +只有一个实现的接口因此只需继承接口。`LayerNormKernel` 只声明了支持的架构与自己的 `entry_for`： + +```python +# src/tileops/kernels/norm/layer_norm.py +class LayerNormKernel(Kernel, LayerNormFwdInterface): + supported_archs: list[int] = [80, 86, 89, 90] + + @classmethod + def entry_for(cls, call: LayerNormCall) -> Entry: + identity = (call.n, call.eps, call.dtype) + return identity, lambda: cls(*identity) +``` + +## 4. 新 kernel 何时需要新开 kernel 接口 {#interface} + +新 kernel 改变了语义或调用契约，例如返回值不同，就属于一个新的 kernel 接口；只在 shape、dtype、架构或性能上不同，就是已有接口的一个新实现。同一算法换一组 tile 大小或 split 数，属于实现内部的参数，不新增实现。 + +**表 4** 常见情况的归属 + +| No. | 情况 | 归属 | +| --- | --- | --- | +| 1 | BatchNorm 的训练与推理：统计量来源不同，训练多返回 `mean` 与 `rstd` | 两个接口 | +| 2 | GLA 的 prefill 与 decode：调用契约相同，只是序列长度不同 | 同一接口的不同实现 | +| 3 | 某个 shape 范围或某个架构上有更快的算法 | 新实现 | +| 4 | 同一算法的 split 数为 1 或大于 1 | 实现内部的参数 | + +新接口需要一个 call spec 类型与一个接口类。family 下有多个 kernel 文件时，两者都写在 `src/tileops/kernels//call_spec.py`；只有一个 kernel 文件时，两者都写在这个 kernel 文件中。接口类按 `{Name}{Fwd|Bwd}Interface` 命名，表示变体的词写在方向后缀之前，例如 `BatchNormTrainFwdInterface`，由 pre-commit 的 `interface-names-lint` 检查。同一 family 的多个接口可以共用一个 call spec 类型，例如 BatchNorm 与 InstanceNorm 的六个接口都用 `BatchNormCall`。 + +```python +# src/tileops/kernels/norm/call_spec.py +@dataclasses.dataclass(frozen=True) +class LayerNormCall(CallSpec): + """The facts that select an implementation normalizing trailing rows and build it. + + Layer, RMS, fused-add and adaptive layer normalization take it. ``n`` is the row's + element count and ``eps`` the op's epsilon. + """ + + n: int = 0 + eps: float = 1e-5 + dtype: torch.dtype = torch.float16 + +class LayerNormFwdInterface(KernelInterface): + """Layer normalization over the trailing ``call.n`` elements.""" + + request = LayerNormCall + + @abstractmethod + def forward(self, x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + """Normalize each run of ``call.n`` trailing elements; nothing is written in place. + + Every tensor is contiguous on ``call.device``. + + Args: + x: Any shape whose trailing axes hold ``call.n`` elements, in ``call.dtype``. + weight: ``call.n`` elements of scale in ``call.dtype``. + bias: ``call.n`` elements of shift in ``call.dtype``. + + Returns: + A new output shaped like *x*, in ``call.dtype``. + """ +``` + +- call spec 的字段是各实现在 `applies`、`refusal`、`entry_for` 中读取的调用事实：shape、dtype、语义参数（包括 op 构造时固定的参数）与 `device`。字段只能是不可变的值。不放张量内容、调参策略与优先级；设备事实由 `CallSpec` 提供。 +- 接口类的 `request` 指向 call spec 类型；抽象方法 `forward` 的参数表是 op 调用 kernel 时传入的参数，docstring 写明每个张量的 shape、dtype、内存布局、设备、是否被原地写入，以及返回值。backend 只依据这份契约编写实现。 +- op 在 `forward` 之外还会在 entry 上调用的方法，一并写成接口的抽象方法。实现必须实现接口的全部抽象方法，否则这个类不能实例化。 +- 各实现必须一致的取值，写成接口的普通类方法，由接口算出。例如 op 要在选择之前按大小分配一个输出缓冲区时，这个大小由接口的类方法从 call spec 算出，而不是从已选中的实现上读。 +- docstring 写的必须是实现真正做到的：声明输入连续，op 就要在调用前调用 `.contiguous()`；实现不保证的行为（例如输入含 NaN 时的结果）写明未定义，不写成保证。 +- 持有 kernel 的 op 都声明 `interfaces`，不覆写 `entry_for`，也不自建 kernel 缓存。 + +## 5. 如何测试选择结果，以及常见报错 {#tests} + +为每个适用范围以及相邻范围的边界各写一个用例，用 `select_implementation` 检查选中的 key，并显式给出设备事实，使测试不依赖运行它的机器： + +```python +# tests/ops/test_batch_norm.py +def test_each_region_selects_its_one_implementation( + op_cls, interface, n, c, spatial, dtype, key +) -> None: + """Exactly one non-general implementation, or else the general one, serves each shape.""" + op = op_cls() + call = BatchNormCall(arch=90, sm_count=132, n=n, c=c, spatial=spatial, dtype=dtype) + assert op.select_implementation(interface, call) == key +``` + +**表 5** 常见报错 + +| No. | 时刻 | 报错信息的片段 | 原因 | +| --- | --- | --- | --- | +| 1 | 构造 op | `does not implement ` | key 背后运行的类（包括 `kernel_map=` 的替换者）没有继承接口；让它继承接口，并由 `entry_for(call)` 构建 | +| 2 | 构造 op | `forward does not take 's arguments` | `forward` 的参数与接口不符 | +| 3 | 构造 op | `has more than one general implementation` | 一个接口有两个 general 实现 | +| 4 | 构造 op | `preferences form a cycle through` | `preferred_over` 成环 | +| 5 | 构造 op | `keys implement none of its kernel interfaces` | `kernel_types` 中的类没有继承 op 的任何 kernel 接口 | +| 6 | 调用 | `in-tree kernels do not run on` | 没有任何 key 能在调用设备类型上运行（`OpNotAvailableError`） | +| 7 | 调用 | `no implementation serves this call` | 有 key 支持该设备类型，但没有实现同时可用且适用，信息中列出每个实现的理由 | +| 8 | 调用 | `dispatch is ambiguous` | 多个实现都适用，且互相没有优先关系 | +| 9 | 调用 | `takes a call spec` | 传给 `kernel_for` 的 call spec 类型不对 | +| 10 | 调用 | `cannot key a dispatch cache` | call spec 中有可变的字段，例如 list | +| 11 | 调用 | `this call spec states ['arch']` | 传给 `kernel_for` 的 call spec 显式给出了设备事实；它们应由 `device` 推出，只有 `select_implementation` 接受显式给出的设备事实 | + +新 op 没有声明 `interfaces` 而持有 kernel 时，`tests/test_kernel_dispatch.py` 中的清单测试报 `declare interfaces instead`。 diff --git a/docs/user-guide/index.md b/docs/user-guide/index.md index 66df18b4..8c82f1d5 100644 --- a/docs/user-guide/index.md +++ b/docs/user-guide/index.md @@ -2,8 +2,9 @@ | Page | Contents | | --- | --- | -| [Writing a Spec](../manifest.md) | what an op's spec contains, and how to read and write one | -| [Adding an Op](../new-op.md) | the four steps from a spec to `status: implemented` | -| [Using torch.compile](../torch-compile.md) | what an op looks like inside a compiled graph, and the conventions a caller keeps to | -| [Timing Benchmarks](../timing.md) | how the numbers on the Benchmarks pages are measured | -| [Adding a Backend](../backends.md) | taking over the ops on a class of devices with your own kernels | +| [Reading and writing the manifest](manifest/index.md) | how the system uses a spec, the concepts used to describe a spec, and how to write one | +| [How an op selects a kernel](dispatch/index.md) | how an op selects a kernel, how to add a kernel, and how a backend joins | +| [Adding a new op](../new-op.md) | the six steps from a spec to `status: implemented` | +| [Bringing an op into torch.compile](../torch-compile.md) | what an op looks like inside a compiled graph, and the conventions a caller follows | +| [How a benchmark is timed](../timing.md) | how the numbers on the Benchmarks pages are measured | +| [Adding a hardware backend](../backends.md) | serving the ops on one class of devices with your own kernels | diff --git a/docs/user-guide/index.zh.md b/docs/user-guide/index.zh.md index 208b8d8e..ac2a3fd6 100644 --- a/docs/user-guide/index.zh.md +++ b/docs/user-guide/index.zh.md @@ -2,8 +2,9 @@ | 文档 | 内容 | | --- | --- | -| [读写 manifest](../manifest.md) | 算子的 spec 由什么构成，怎么读、怎么写 | -| [添加新算子](../new-op.md) | 从一份 spec 到 `status: implemented` 的四步 | -| [接入 torch.compile](../torch-compile.md) | 算子在编译图里的形态，以及调用时的约定 | -| [benchmark 怎么计时](../timing.md) | Benchmarks 页上的数字是怎么测出来的 | -| [接入新硬件后端](../backends.md) | 用自己的 kernel 接管某一类设备上的算子 | +| [读写 manifest](manifest/index.md) | 系统如何使用 spec、描述 spec 所用的概念，以及如何写一份 spec | +| [op 如何选择 kernel](dispatch/index.md) | op 如何选中 kernel、如何新增 kernel，以及 backend 如何接入 | +| [添加新 op](../new-op.md) | 从一份 spec 到 `status: implemented` 的六步 | +| [接入 torch.compile](../torch-compile.md) | op 在编译图中的形态，以及调用时的约定 | +| [benchmark 的计时方法](../timing.md) | Benchmarks 页上的数字如何测得 | +| [接入新硬件 backend](../backends.md) | 用自己的 kernel 接管某一类设备上的 op | diff --git a/docs/user-guide/manifest/calls.md b/docs/user-guide/manifest/calls.md new file mode 100644 index 00000000..4036a854 --- /dev/null +++ b/docs/user-guide/manifest/calls.md @@ -0,0 +1,135 @@ +# Calls and validation + +A call has two phases, construction and call, and the checks in both phases are generated from the signature. The validator runs static checks on the signature in CI to ensure that every spec can generate these checks. This page describes the checks run during a call, the kernel selection rules, the validator's checks, and the forms rejected in `shape_rules`. + +## 1. The two phases of a call {#call} + +In the **construction** phase, the op checks parameter values against their `type` (including whether a dtype parameter is in the allowed set), checks that every non-optional construction-time tensor is given, and checks the `invariant` of ADTs. The following are fixed at construction: + +- the values of `Bool`, enum and ADT parameters; +- whether `Maybe` parameters and construction-time tensors are given; +- the indices solved from the shapes and dtypes of construction-time tensors, and the `let` entries computable from these. + +Checks such as non-negative axes, refinements and `invariant` run at construction when they can already be evaluated there, and are deferred to call time otherwise. + +In the **call** phase, the checks generated from the signature around `forward` run in the following order: + +1. determine whether each call-time tensor is passed; +2. check domain restrictions; +3. select type family branches; +4. infer indices, and check the remaining refinements; +5. check the preconditions of output buffers; +6. run the implementation; +7. check the outputs: the number of outputs and whether each is `None`, shape, dtype, device and memory layout; `out` and `alias` outputs must be the corresponding tensor objects, and other outputs must not share storage with inputs. + +When any check fails, the op raises an error that names the unsatisfied declaration. + +If every tensor a call writes (each output and each written input) has no elements, step 6 runs no implementation: new outputs are created on the call device with the checked shapes and dtypes, and `out` and the written inputs are returned unchanged. When the inputs are empty but the outputs are not, the implementation runs as usual. + +## 2. Kernel selection {#selection} + +An op declares in `interfaces` the positions where it calls kernels, and each position corresponds to a kernel interface. An implementation of a kernel interface is a kernel class that inherits it. At call time, the op constructs a call spec and calls `kernel_for(interface, call)` to obtain an entry. Each distinct call spec is looked up only once. When a call spec first appears, the Op base class selects exactly one implementation in the following order: + +1. availability: the implementation's `devices` and `supported_archs`; +2. applicability: the implementation's `applies` and `refusal`; +3. precedence: the implementation's `general` and `preferred_over`. + +The selected implementation's `entry_for(call)` gives the build identity and the factory. The signature declares only the requirements of the algorithm itself, and the restrictions of each implementation are written in that implementation's own declarations. + +The selection rules, how to add a kernel, and how a backend joins are in [How an op selects a kernel](../dispatch/index.md). + +## 3. type inference {#inference} + +The construction phase first solves the indices that construction-time tensors determine, and the call phase then solves the remaining indices from the call-time tensors actually passed. The validator fixes the solving order in advance, independent of the order in which declarations are written. + +**Table 1** Unification rules for input axes + +| No. | Axis form | unification | +| --- | --- | --- | +| 1 | `M`, with `M` unknown | `M :=` actual axis length | +| 2 | `a * M + e`, where `a` is a known positive integer constant, `e` is known, and `M` is the only unknown | `M := (actual axis length - e) / a`, checking that the division is exact and the result is non-negative | +| 3 | `*S`, where `S` is the only unknown part of the shape and the number of other axes is known | `S :=` the tuple of the corresponding axes | +| 4 | `dtype` is an unknown `DType` index `T` | `T :=` actual dtype | +| 5 | any other form | only checks that the equality holds | + +**Table 2** Indices solved from declarations + +| No. | spec | Declaration | Solved index | +| --- | --- | --- | --- | +| 1 | SiluAndMul | `x: "[M, 2 * N]"` | `N`, also checking that the axis length is even | +| 2 | GQA varlen | `cu_seqlens_q: "[B + 1]"` | `B` | +| 3 | GemmW4A16 | `activation: "[M, K]"`, `packed_weight: "[N, K // 2]"` | `K` is solved from `activation`; `K // 2` in `packed_weight` is used only for checking | + +- Every `Dim`, `Shape` and `DType` index used on a branch must be solvable from the inputs or given by a parameter. An index that appears only in outputs must be a parameter or a `let`. +- A generator generates the values of metadata tensors only at workload instantiation, and helps determine the indices a workload row does not write; in an actual call, these indices are likewise solved from the tensors passed. +- If an index cannot be solved, or has more than one solution, the validator rejects the spec. +- For the output buffer `out`, the inference phase only determines whether it is passed; once the output type is fixed, `out` is checked against it. +- If the relation between an axis length and a logical dimension is not affine, a name can stand for the physical axis and a `let` computes the logical dimension. For example, MHCPre writes `b: "[Q]"` and `let: {n: "mhc.expansion(Q)"}`. + +## 4. Call device {#device} + +The call device is determined as shown below. At workload instantiation, tensors are placed by the same rules. + +![Call device](img/device.svg) + +- A tensor that declares `device: cpu` is always on the CPU and does not take part in deciding the call device; among construction-time tensors, too, only those that do not declare `device: cpu` take part. +- When there is no call-time tensor, no `device` parameter and no construction-time tensor that takes part, the call device is the current CUDA device, or `None` when CUDA is unavailable. By design, before this step an explicitly given or process-default target selects a device within the device categories it declares; this step is not implemented yet. +- Once the call device is determined, `out` and each output are checked or allocated on that device. +- A construction-time tensor that does not declare `device: cpu` is copied to the call device at call time and converted to the dtype in the signature. +- For a tensor that declares `contiguous: true`, the generated checks confirm that it is contiguous in memory; tensors generated by workload instantiation are always contiguous. + +## 5. torch.compile and SymInt {#symint} + +Under `torch.compile`, the generated checks are evaluated on SymInt, and each discriminant is already a concrete Python value at that point. + +- An expression that needs to convert a SymBool to a Python boolean is evaluated only at construction. +- An op class can declare a compile boundary, meaning it supports `fullgraph=True`. For such a class, the validator requires every expression to be evaluable on SymInt. +- Expression strings are parsed and checked before code generation, and the generated code does not parse strings at run time. + +The conventions a caller follows under `torch.compile` are in [Bringing an op into torch.compile](../../torch-compile.md). + +## 6. The validator's checks {#validator} + +[`scripts/validate_manifest.py`](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py) checks each spec on every combination of discriminant values. The quantities that take part in the combinations are: + +- the `match` of type families; +- `optional` and `nullable`; +- `mutated`; +- whether output buffers are passed; +- the quantities involved in deciding whether an index is used. + +For combinations excluded by a domain restriction, the validator skips only two checks: type family coverage and the solving order. If a spec has more value combinations to check than the configured limit (256 by default), the validator issues an advisory-level notice but still checks the spec in full. + +On each value combination, the validator checks the following: + +1. The category of each name matches its kind, and the `type` of each parameter is compatible with the kind required at each place it is used. +2. The `cases` of each type family have no gaps and no overlaps over the accepted values, and references between type families have no cycle. +3. Every index can be solved (see [§ 3](#inference)). +4. Dependencies between `let` entries have no cycle. +5. Every expression belongs to the expression language, and every primitive used is built in. +6. On every workload row, the generator results unify with their declarations and every `requires` holds. +7. For a class that declares a compile boundary, every expression can be evaluated on SymInt. +8. Every workload row can be instantiated. +9. On each effect branch, the operator schema, the aliasing and the roofline read and write counts agree. + +For an `implemented` op, the validator also checks that the code agrees with the spec, including `__init__` against `params`, `forward` against the call-time inputs, and composition against the class's `delegate_types` and `kernel_types`. `spec-only` ops skip these code-dependent checks. + +- CI runs the validator on the whole manifest. +- The validator parses a spec field by field. When a field cannot be parsed, the validator reports that field and skips only the checks that depend on it; the other checks run as usual. +- When an op is imported, the manifest is loaded in lenient mode, so the op imports normally even if the manifest is incomplete; strict checks are run only by the validator. + +## 7. Rejected forms {#rejected} + +`shape_rules` contains only refinements. Forms that also act as a shape declaration, a name definition or a presence test are rejected: + +**Table 3** Rejected forms and their rewrites + +| No. | Rejected form | Example | Rewrite as | +| --- | --- | --- | --- | +| 1 | reading a tensor's `shape` | `x.shape == (B, S, H, D)` | declare the shape on the tensor: `x: {shape: "[B, S, H, D]"}` | +| 2 | an equality stating that two tensors have the same shape | `output.shape == input.shape` | the two tensors use the same shape term, such as `[*S]` | +| 3 | an equality defining a new name | `C_in_g == C_in // groups` | `let: {C_in_g: "C_in // groups"}` | +| 4 | `x is None` on a tensor | `bias is None or ...` | `not present(bias) or ...` | +| 5 | `v is None` on a value | `max_seqlen is None` | `not present(max_seqlen)` | +| 6 | `isinstance` | `s[0] if isinstance(s, tuple) else s` | `per_axis(s, 0, 2)` | +| 7 | a set comprehension | `len({d % n for d in dim}) == len(dim)` | `unique_axes(dim, n)` | diff --git a/docs/user-guide/manifest/calls.zh.md b/docs/user-guide/manifest/calls.zh.md new file mode 100644 index 00000000..42377064 --- /dev/null +++ b/docs/user-guide/manifest/calls.zh.md @@ -0,0 +1,135 @@ +# 调用与校验 + +一次调用分为构造与调用两个阶段，两个阶段的检查都由签名生成。validator 在 CI 中对签名做静态检查，确保每个 spec 都能生成这些检查。本页说明调用过程中执行的检查、kernel 的选择规则、validator 的检查项，以及 `shape_rules` 中不被接受的写法。 + +## 1. 一次调用的两个阶段 {#call} + +在**构造**阶段，op 按 `type` 检查参数的取值（包括 dtype 参数是否在允许的集合内），检查非可选的构造期张量都已给出，并检查 ADT 的 `invariant`。以下各量在构造时确定： + +- `Bool`、枚举与 ADT 参数的值； +- `Maybe` 参数与构造期张量是否给出； +- 由构造期张量的形状与 dtype 求得的 index，以及由这些量可以算出的 `let`。 + +轴非负、refinement、`invariant` 等检查，如果在构造时已经可以求值，就在构造时执行，否则推迟到调用时执行。 + +在**调用**阶段，`forward` 外层由签名生成的检查按以下顺序执行： + +1. 确定各调用期张量是否传入； +2. 检查定义域限制； +3. 选择 type family 的分支； +4. 推断 index，并检查其余的 refinement； +5. 检查输出缓冲的前置条件； +6. 执行实现； +7. 检查输出：输出的数量与是否为 `None`、形状、dtype、设备与内存布局；`out` 与 `alias` 输出必须就是对应的张量对象，其他输出不能与输入共享存储。 + +任何一项检查失败时，op 都会报错，并指出未满足的那条声明。 + +如果一次调用写入的所有张量（各输出与被写入的输入）都不含元素，第 6 步不执行任何实现：新的输出按检查过的形状与 dtype 在调用设备上创建，`out` 与被写入的输入原样返回。输入为空而输出不为空时，照常执行实现。 + +## 2. kernel 的选择 {#selection} + +op 在 `interfaces` 中声明它调用 kernel 的位置，每个位置对应一个 kernel 接口；一个 kernel 接口的实现，是继承它的 kernel 类。调用时，op 构造 call spec，再调用 `kernel_for(interface, call)` 取得 entry。同一个 call spec 再次出现时只做一次查找；首次出现时，Op 基类按以下顺序选出唯一一个实现： + +1. 可用性：实现的 `devices` 与 `supported_archs`； +2. 适用性：实现的 `applies` 与 `refusal`； +3. 优先级：实现的 `general` 与 `preferred_over`。 + +选出的实现由 `entry_for(call)` 给出 build identity 与构建函数。签名只声明算法本身的要求，每个实现的限制写在实现自己的声明中。 + +选择规则、新增 kernel 的写法与 backend 的接入方式，见 [op 如何选择 kernel](../dispatch/index.md)。 + +## 3. type inference {#inference} + +构造阶段先求出由构造期张量能确定的 index，调用阶段再从实际传入的调用期张量求出其余的 index。求解的顺序由 validator 预先确定，与声明的书写顺序无关。 + +**表 1** 输入轴的 unification 规则 + +| No. | 轴的写法 | unification | +| --- | --- | --- | +| 1 | `M`，且 `M` 未知 | `M :=` 实际轴长 | +| 2 | `a * M + e`，其中 `a` 是已知的正整数常量，`e` 已知，`M` 是唯一的未知量 | `M := (实际轴长 - e) / a`，并检查能否整除、结果是否非负 | +| 3 | `*S`，其中 `S` 是该形状中唯一未知的部分，其余轴的个数已知 | `S :=` 对应的若干轴组成的元组 | +| 4 | `dtype` 为未知的 `DType` index `T` | `T :=` 实际 dtype | +| 5 | 其他写法 | 只检查等式是否成立 | + +**表 2** 由声明求得的 index + +| No. | spec | 声明 | 求得的 index | +| --- | --- | --- | --- | +| 1 | SiluAndMul | `x: "[M, 2 * N]"` | `N`，同时检查轴长为偶数 | +| 2 | GQA varlen | `cu_seqlens_q: "[B + 1]"` | `B` | +| 3 | GemmW4A16 | `activation: "[M, K]"`、`packed_weight: "[N, K // 2]"` | `K` 由 `activation` 求得，`packed_weight` 的 `K // 2` 只用于检查 | + +- 在某个分支上被用到的每个 `Dim`、`Shape`、`DType` index，都必须能由输入求得，或由参数给出。只出现在输出中的 index 必须是参数或 `let`。 +- generator 只在 workload 实例化时生成 metadata 张量的取值，并帮助确定 workload 行中不写的 index；实际调用时，这些 index 同样从传入的张量求得。 +- 如果某个 index 无法求得，或者有多个解，validator 会拒绝这个 spec。 +- 对于输出缓冲 `out`，推断阶段只确定它是否传入；输出类型确定之后，再按输出类型检查它。 +- 如果轴长与逻辑维度之间不是 affine 关系，可以用一个名字表示物理轴，再用 `let` 计算逻辑维度。例如 MHCPre 写作 `b: "[Q]"` 与 `let: {n: "mhc.expansion(Q)"}`。 + +## 4. 调用设备 {#device} + +调用设备按下图确定。workload 实例化时，张量也按同样的规则放置。 + +![调用设备](img/device.svg) + +- 声明了 `device: cpu` 的张量始终位于 CPU，不参与调用设备的判定；构造期张量也只有未声明 `device: cpu` 的才参与判定。 +- 没有调用期张量、`device` 参数与参与判定的构造期张量时，调用设备取当前的 CUDA 设备；CUDA 不可用时为 `None`。设计上，在这一步之前还会由显式指定或进程默认的 target 在其声明的设备类别中选择设备，这一步尚未实现。 +- 调用设备确定后，`out` 与各输出在该设备上检查或分配。 +- 没有声明 `device: cpu` 的构造期张量，在调用时复制到调用设备，并转换为签名中的 dtype。 +- 声明了 `contiguous: true` 的张量，由生成的检查确认其内存连续；workload 实例化生成的张量总是连续的。 + +## 5. torch.compile 与 SymInt {#symint} + +在 `torch.compile` 下，生成的检查在 SymInt 上求值，此时各 discriminant 已经是确定的 Python 值。 + +- 需要把 SymBool 转为 Python 布尔值的表达式，只在构造时求值。 +- op 类可以声明 compile boundary，表示它支持 `fullgraph=True`。对这样的类，validator 要求其所有表达式都能在 SymInt 上求值。 +- 表达式字符串在代码生成之前完成解析与检查，生成的代码在运行时不再解析字符串。 + +调用方在 `torch.compile` 下需要遵守的约定见[接入 torch.compile](../../torch-compile.md)。 + +## 6. validator 的检查项 {#validator} + +[`scripts/validate_manifest.py`](https://github.com/tile-ai/TileOPs/blob/main/scripts/validate_manifest.py) 在 discriminant 的每一种取值组合上检查每个 spec。参与组合的量包括： + +- type family 的 `match`； +- `optional` 与 `nullable`； +- `mutated`； +- 输出缓冲是否传入； +- 判断 index 是否被用到时涉及的量。 + +对于被定义域限制排除的组合，validator 只跳过 type family 覆盖与求解顺序这两项检查。如果一个 spec 需要检查的取值组合超过配置的上限（默认 256），validator 给出 advisory 级别的提示，但仍然完整地检查这个 spec。 + +在每一种取值组合上，validator 检查以下各项： + +1. 每个名字的类别与其 kind 相符，每个参数的 `type` 与其各使用位置所要求的 kind 相容。 +2. type family 的 `cases` 在接受的取值上既无遗漏也无重叠，type family 之间的引用无环。 +3. 每个 index 都能求得（见[第 3 节](#inference)）。 +4. `let` 之间的依赖无环。 +5. 每个表达式都属于表达式语言，用到的 primitive 都是内建的。 +6. 在每条 workload 行上，generator 的结果都能与声明做 unification，`requires` 都成立。 +7. 对于声明了 compile boundary 的类，其所有表达式都能在 SymInt 上求值。 +8. 每条 workload 行都能实例化。 +9. 在每个 effect 分支上，operator schema、别名关系与 roofline 的读写计数一致。 + +对于 `implemented` 的 op，validator 还检查代码与 spec 是否一致，包括 `__init__` 与 `params`、`forward` 与调用期输入，以及 composition 与类的 `delegate_types`、`kernel_types`。`spec-only` 的 op 跳过这些依赖代码的检查。 + +- CI 对整份 manifest 运行 validator。 +- validator 逐字段解析 spec。某个字段无法解析时，validator 报告这个字段，并只跳过依赖该字段的检查，其余检查照常进行。 +- 导入 op 时，manifest 以宽松模式加载，即使 manifest 不完整，op 也能正常导入；严格的检查只由 validator 执行。 + +## 7. 不被接受的写法 {#rejected} + +`shape_rules` 中只写 refinement。兼做形状声明、名字定义或存在性判断的写法不被接受： + +**表 3** 不被接受的写法及对应的改写 + +| No. | 不被接受的写法 | 例子 | 改写为 | +| --- | --- | --- | --- | +| 1 | 读取张量的 `shape` | `x.shape == (B, S, H, D)` | 在张量上声明形状：`x: {shape: "[B, S, H, D]"}` | +| 2 | 用等式表示两个张量形状相同 | `output.shape == input.shape` | 两个张量使用同一个形状项，例如 `[*S]` | +| 3 | 用等式定义新名字 | `C_in_g == C_in // groups` | `let: {C_in_g: "C_in // groups"}` | +| 4 | 对张量使用 `x is None` | `bias is None or ...` | `not present(bias) or ...` | +| 5 | 对值使用 `v is None` | `max_seqlen is None` | `not present(max_seqlen)` | +| 6 | `isinstance` | `s[0] if isinstance(s, tuple) else s` | `per_axis(s, 0, 2)` | +| 7 | 集合推导式 | `len({d % n for d in dim}) == len(dim)` | `unique_axes(dim, n)` | diff --git a/docs/user-guide/manifest/concepts.md b/docs/user-guide/manifest/concepts.md new file mode 100644 index 00000000..e1cafce6 --- /dev/null +++ b/docs/user-guide/manifest/concepts.md @@ -0,0 +1,154 @@ +# Concepts + +The core of a spec is its signature. Formally, a signature is a polymorphic function type. This page uses GEMM as the example and introduces, in order, the concepts used to describe a signature. How each field is written is in [Spec fields](writing.md) and [Extensions](extensions.md). + +## 1. polymorphic function type {#pft} + +The GEMM signature in mathematical form: + +``` +Mat[t: Bool, R: Dim, C: Dim] = match t { false → [R, C]; true → [C, R] } + +gemm(trans_a, trans_b : Bool) : ∀ (M N K : Dim) (T : DType[float16 | bfloat16]). + Tensor[T, Mat[trans_a, M, K]] → Tensor[T, Mat[trans_b, K, N]] → Tensor[T, [M, N]] +``` + +The corresponding YAML is in [Examples § 1](examples.md#gemm). + +`trans_a` and `trans_b` after the function name are construction parameters; `M`, `N`, `K` and `T` after `∀` are the quantified names. The signature has three layers: + +**Table 1** The three layers of a signature + +| No. | Layer | Mathematical form | Written in the manifest as | +| --- | --- | --- | --- | +| 1 | tensor type | `Tensor[T, s]`, where `T` is the dtype and `s` is the shape | `{dtype: T, shape: "[M, N]"}` | +| 2 | function type | from the types of the input tensors to the types of the output tensors | `inputs`, `outputs` | +| 3 | quantification | `∀` lists every type index and its kind | `forall` | + +A function type with quantification over its type indices is called a polymorphic function type. The runtime checks and `eval_roofline()` are both generated from the signature; for an op that declares a compile boundary, the `torch.library` operator and the fake/meta functions are also generated from it. The validator runs its static checks on the signature. + +## 2. type index and kind {#index} + +A type index is a parameter name in a type, such as `T`, `M`, `N` and `K` in GEMM. Every type index has a kind, and the kind fixes its range of values: + +**Table 2** Common kinds + +| No. | kind | Values | Example | +| --- | --- | --- | --- | +| 1 | `Dim` | a non-negative integer, an axis length | `M`, `N`, `K` in GEMM | +| 2 | `Shape` | a tuple of `Dim`, used for tensors whose rank is not fixed | the shape `[*S]` of an elementwise op | +| 3 | `DType[...]` | one of the listed dtypes | `T` in GEMM | +| 4 | `Bool`, enum, ADT | finitely many values | `trans_a`, `trans_b` in GEMM | + +Every type index in a signature has exactly one source, which is one of the following three: + +- declared in `forall`, such as `M` and `T`; +- a construction parameter, whose kind is determined by the parameter's `type`; for example, the `type` of `trans_a` is `bool`, and the corresponding kind is `Bool`; +- defined by `let`, see [§ 5](#let). + +## 3. type family {#family} + +A type family is a function that gives a shape from the value of a discriminant. A discriminant is a quantity with finitely many values, such as a `Bool` parameter, an enum parameter, an ADT parameter, or whether a tensor is passed. + +`Mat` in GEMM is a type family: when `t` is false the result is `[R, C]`, and when it is true the result is `[C, R]`. The shape of tensor `a` is written `Mat[trans_a, M, K]`, so: + +- at construction the value of `trans_a` is known, `Mat` selects a branch, and the shape of `a` becomes `[M, K]` or `[K, M]`; +- at call time, the values of `M` and `K` are obtained from the actual shape of `a` through unification. + +A type family is defined in the `signature.types` of the spec that uses it. + +## 4. ADT {#adt} + +An ADT (algebraic data type) is a tagged sum type: an ADT value belongs to one of several constructors, and each constructor has its own fields. + +The layout parameter of the MoE staged ops is the ADT `MGroupedLayout`, which has two constructors: + +**Table 3** Constructors of `MGroupedLayout` + +| No. | constructor | Fields | Meaning | +| --- | --- | --- | --- | +| 1 | `contiguous` | `packing`, `metadata_kind`, `alignment` | the rows of each expert are laid out contiguously, and `packing` decides whether they are aligned to `alignment` | +| 2 | `masked` | `max_m` | each expert occupies a fixed `max_m` rows, and extra rows are marked by a mask | + +- A `masked` value has only the field `max_m`, and a `contiguous` value has only the other three fields. +- A type family can pattern-match on the constructor and read `layout.max_m` in the `masked` branch. +- Each constructor corresponds to a Python class; for example, `contiguous` corresponds to `ContiguousLayoutSpec`. +- ADTs shared by several specs are defined in `spec/types.yaml`. + +The definition of `MGroupedLayout` and a spec that uses it are in [Examples § 2](examples.md#moe). + +## 5. `let` {#let} + +`let` names a quantity computed from indices. The output height of MaxPool2d is an example: + +```yaml +let: + kH: "per_axis(kernel_size, 0, 2)" + H_out: "pool.out(H_in, kH, sH, pH, dH, ceil_mode)" +``` + +`H_out` is computed from the input height `H_in` and the construction parameters, and is used in the output shape. The value of a `let` is computed from the signature and is not given in workload rows. + +## 6. refinement {#refinement} + +A refinement is a predicate that constrains the values of indices, written in `shape_rules`. Attention is an example: + +- the types `q: [B, S_q, H, D]` and `k: [B, S_kv, H, D]` by themselves do not constrain how `S_q` relates to `S_kv`; +- the refinement `not is_causal or S_q <= S_kv` requires `S_q <= S_kv` in the causal case. + +## 7. unification and type inference {#unification} + +Unification matches the declared shape against the shape of the actual tensor axis by axis, solves for unknown indices, and checks that each index takes the same value everywhere it appears. Take one GEMM call with `trans_a = trans_b = false`: + +``` +Declared: a: Tensor[T, [M, K]] b: Tensor[T, [K, N]] +Actual: a: bfloat16, (128, 4096) b: bfloat16, (4096, 512) +``` + +**Table 4** Unification for this GEMM call + +| No. | Correspondence | Result | +| --- | --- | --- | +| 1 | dtype of `a` against `T` | `T := bfloat16` | +| 2 | axis 1 of `a` against `M` | `M := 128` | +| 3 | axis 2 of `a` against `K` | `K := 4096` | +| 4 | dtype of `b` against `T` | checked equal to `T` | +| 5 | axis 1 of `b` against `K` | checked equal to `K` | +| 6 | axis 2 of `b` against `N` | `N := 512` | + +Once every index is solved, the output type `Tensor[T, [M, N]]` is fixed as `bfloat16, (128, 512)`. + +Type inference is the process of solving every index at call time; unification is the part of it that handles equalities. Other relations between indices are checked by refinements. Which axis forms unification can solve is in [Calls and validation § 3](calls.md#inference). + +## 8. effect {#effect} + +An effect describes what a call writes to its arguments and how outputs alias inputs. The `inplace` parameter of an activation function is an example: when it is true, the op writes the input tensor and returns that input object as the output: + +```yaml +inputs: {input: {dtype: T, shape: "[*S]", mutated: inplace}} +outputs: {output: {dtype: T, shape: "[*S]", alias: input}} +``` + +Effects determine the read and write counts of the roofline, and the operator schema (its `mutates_args`) generated for an op that declares a compile boundary. All effect declarations are in [Extensions § 6](extensions.md#effects). + +## 9. Glossary {#glossary} + +**Table 5** Meaning of each term and how it is written in the manifest + +| No. | Term | Meaning | Written in the manifest as | Section | +| --- | --- | --- | --- | --- | +| 1 | `Tensor[T, s]` | tensor type parameterized by dtype `T` and shape `s` | a tensor's `dtype`, `shape` | [§ 1](#pft) | +| 2 | polymorphic function type | a function type with quantification over its type indices | `forall` | [§ 1](#pft) | +| 3 | type index | a parameter name in a type | `forall`, construction parameters, `let` | [§ 2](#index) | +| 4 | kind | the sort of a type index, which determines its range of values and the operations available on it | `forall: {M: Dim}` | [§ 2](#index) | +| 5 | discriminant | a quantity with finitely many values, used to select a type family branch or to decide whether a tensor exists | expressions in `match`, `optional`, `nullable`, `mutated` | [§ 3](#family) | +| 6 | type family | a function that gives a shape from the value of a discriminant | `signature.types` | [§ 3](#family) | +| 7 | ADT, constructor | a tagged sum type in which each constructor has its own fields | `adts` in `spec/types.yaml` | [§ 4](#adt) | +| 8 | `let` | a named quantity computed from indices | `let: {H_out: "..."}` | [§ 5](#let) | +| 9 | refinement | a predicate that constrains the values of indices | each entry in `shape_rules` | [§ 6](#refinement) | +| 10 | unification | matching the declared shape against the actual shape axis by axis to solve unknown indices | generated from the signature | [§ 7](#unification) | +| 11 | type inference | the process of solving every index at call time | generated from the signature | [§ 7](#unification) | +| 12 | effect | what a call writes to its arguments, and aliasing | `mutated`, `write_only`, `buffer`, `alias` | [§ 8](#effect) | +| 13 | `Maybe[X]`, `present` | an `X` that may be absent; `present(x)` says whether it is given | optional inputs, nullable outputs, `int \| None` parameters | [Extensions § 1](extensions.md#presence) | +| 14 | primitive | a built-in function callable in expressions, such as `broadcast` and `pool.out` | function calls in expressions | [Extensions § 4](extensions.md#let) | +| 15 | generator | a function that generates the values of a metadata tensor at instantiation, such as `prefix_sum(q_lens)` | a tensor's `values` | [Extensions § 7](extensions.md#generators) | diff --git a/docs/user-guide/manifest/concepts.zh.md b/docs/user-guide/manifest/concepts.zh.md new file mode 100644 index 00000000..0aeda181 --- /dev/null +++ b/docs/user-guide/manifest/concepts.zh.md @@ -0,0 +1,154 @@ +# 概念 + +spec 的核心是签名。签名在形式上是一个 polymorphic function type。本页以 GEMM 为例，依次介绍描述签名所用的概念；各字段的具体写法见[写一个 spec](writing.md)与[扩展写法](extensions.md)。 + +## 1. polymorphic function type {#pft} + +GEMM 的签名写成数学形式如下： + +``` +Mat[t: Bool, R: Dim, C: Dim] = match t { false → [R, C]; true → [C, R] } + +gemm(trans_a, trans_b : Bool) : ∀ (M N K : Dim) (T : DType[float16 | bfloat16]). + Tensor[T, Mat[trans_a, M, K]] → Tensor[T, Mat[trans_b, K, N]] → Tensor[T, [M, N]] +``` + +对应的 YAML 见[示例 1](examples.md#gemm)。 + +函数名之后的 `trans_a`、`trans_b` 是构造参数，`∀` 之后的 `M`、`N`、`K`、`T` 是被量化的名字。这个签名由三层组成： + +**表 1** 签名的三层结构 + +| No. | 层次 | 数学形式 | manifest 中的写法 | +| --- | --- | --- | --- | +| 1 | 张量类型 | `Tensor[T, s]`，其中 `T` 是 dtype，`s` 是形状 | `{dtype: T, shape: "[M, N]"}` | +| 2 | 函数类型 | 从输入张量的类型到输出张量的类型 | `inputs`、`outputs` | +| 3 | quantification | `∀` 列出所有 type index 及其 kind | `forall` | + +对 type index 做了 quantification 的函数类型称为 polymorphic function type。运行时检查与 `eval_roofline()` 都从签名生成；声明了 compile boundary 的 op，其 `torch.library` operator 与 fake/meta 函数也从签名生成。validator 以签名为对象做静态检查。 + +## 2. type index 与 kind {#index} + +type index 是类型中的参数名，例如 GEMM 的 `T`、`M`、`N`、`K`。每个 type index 都有一个 kind，kind 规定了它的取值范围： + +**表 2** 常用的 kind + +| No. | kind | 取值 | 例子 | +| --- | --- | --- | --- | +| 1 | `Dim` | 非负整数，表示轴长 | GEMM 的 `M`、`N`、`K` | +| 2 | `Shape` | 由 `Dim` 组成的元组，用于秩不固定的张量 | 逐元素 op 的形状 `[*S]` | +| 3 | `DType[...]` | 所列 dtype 中的一个 | GEMM 的 `T` | +| 4 | `Bool`、枚举、ADT | 有限个取值 | GEMM 的 `trans_a`、`trans_b` | + +签名中每个 type index 的来源是唯一的，来源有三种： + +- 在 `forall` 中声明，例如 `M`、`T`； +- 来自构造参数，其 kind 由参数的 `type` 决定，例如 `trans_a` 的 `type` 是 `bool`，对应的 kind 是 `Bool`； +- 由 `let` 定义，见[第 5 节](#let)。 + +## 3. type family {#family} + +type family 是一个根据 discriminant 的取值给出形状的函数。discriminant 指取值有限的量，例如 `Bool` 参数、枚举参数、ADT 参数，以及某个张量是否传入。 + +GEMM 中的 `Mat` 就是一个 type family：`t` 为 false 时结果是 `[R, C]`，为 true 时结果是 `[C, R]`。张量 `a` 的形状写作 `Mat[trans_a, M, K]`，于是： + +- 构造时 `trans_a` 的值已经确定，`Mat` 选定分支，`a` 的形状随之确定为 `[M, K]` 或 `[K, M]`； +- 调用时，`M` 与 `K` 的值通过 unification 从 `a` 的实际形状中求得。 + +type family 定义在使用它的 spec 的 `signature.types` 中。 + +## 4. ADT {#adt} + +ADT（algebraic data type）是带 tag 的 sum type：一个 ADT 值属于若干 constructor 中的一个，每个 constructor 有各自的字段。 + +MoE staged 系列 op 的 layout 参数是 ADT `MGroupedLayout`，它有两个 constructor： + +**表 3** `MGroupedLayout` 的 constructor + +| No. | constructor | 字段 | 含义 | +| --- | --- | --- | --- | +| 1 | `contiguous` | `packing`、`metadata_kind`、`alignment` | 各专家的行连续排放，`packing` 决定是否按 `alignment` 对齐 | +| 2 | `masked` | `max_m` | 每个专家固定占 `max_m` 行，多出的行由 mask 标记 | + +- `masked` 值只有 `max_m` 一个字段，`contiguous` 值只有另外三个字段。 +- type family 可以按 constructor 做 pattern matching，并在 `masked` 分支中读取 `layout.max_m`。 +- 每个 constructor 对应一个 Python 类，例如 `contiguous` 对应 `ContiguousLayoutSpec`。 +- 多个 spec 共用的 ADT 定义在 `spec/types.yaml` 中。 + +`MGroupedLayout` 的定义以及使用它的 spec 见[示例 2](examples.md#moe)。 + +## 5. `let` {#let} + +`let` 为由 index 计算得到的量命名。以 MaxPool2d 的输出高度为例： + +```yaml +let: + kH: "per_axis(kernel_size, 0, 2)" + H_out: "pool.out(H_in, kH, sH, pH, dH, ceil_mode)" +``` + +`H_out` 由输入高度 `H_in` 与构造参数计算得到，并用于输出形状。`let` 的值由签名计算，workload 行中不需要给出。 + +## 6. refinement {#refinement} + +refinement 是约束 index 取值的谓词，写在 `shape_rules` 中。以 attention 为例： + +- 类型 `q: [B, S_q, H, D]` 与 `k: [B, S_kv, H, D]` 本身不限制 `S_q` 与 `S_kv` 的大小关系； +- refinement `not is_causal or S_q <= S_kv` 表示在 causal 情形下要求 `S_q <= S_kv`。 + +## 7. unification 与 type inference {#unification} + +unification 将声明的形状与实际张量的形状逐轴对应，求出未知的 index，并检查同一个 index 在各处的取值是否一致。以 `trans_a = trans_b = false` 时的一次 GEMM 调用为例： + +``` +声明： a: Tensor[T, [M, K]] b: Tensor[T, [K, N]] +实际： a: bfloat16, (128, 4096) b: bfloat16, (4096, 512) +``` + +**表 4** 这次 GEMM 调用的 unification 过程 + +| No. | 对应关系 | 结果 | +| --- | --- | --- | +| 1 | `a` 的 dtype 对应 `T` | `T := bfloat16` | +| 2 | `a` 的第 1 轴对应 `M` | `M := 128` | +| 3 | `a` 的第 2 轴对应 `K` | `K := 4096` | +| 4 | `b` 的 dtype 对应 `T` | 检查其与 `T` 相等 | +| 5 | `b` 的第 1 轴对应 `K` | 检查其与 `K` 相等 | +| 6 | `b` 的第 2 轴对应 `N` | `N := 512` | + +所有 index 求出之后，输出类型 `Tensor[T, [M, N]]` 即确定为 `bfloat16, (128, 512)`。 + +type inference 指调用时求出全部 index 的过程，unification 是其中处理等式的部分；index 之间的其他关系由 refinement 检查。哪些轴的写法可以通过 unification 求解，见[调用与校验 3](calls.md#inference)。 + +## 8. effect {#effect} + +effect 描述一次调用对参数的写入以及输出与输入之间的别名关系。以激活函数的 `inplace` 参数为例，当它为 true 时，op 写入输入张量，并将这个输入对象作为输出返回： + +```yaml +inputs: {input: {dtype: T, shape: "[*S]", mutated: inplace}} +outputs: {output: {dtype: T, shape: "[*S]", alias: input}} +``` + +effect 决定了 roofline 的读写计数，以及声明了 compile boundary 的 op 所生成的 operator schema（其中的 `mutates_args`）。所有 effect 声明见[扩展写法 6](extensions.md#effects)。 + +## 9. 术语表 {#glossary} + +**表 5** 术语的含义与 manifest 中的写法 + +| No. | 术语 | 含义 | manifest 中的写法 | 所在小节 | +| --- | --- | --- | --- | --- | +| 1 | `Tensor[T, s]` | 以 dtype `T` 与形状 `s` 为参数的张量类型 | 张量的 `dtype`、`shape` | [1](#pft) | +| 2 | polymorphic function type | 对 type index 做了 quantification 的函数类型 | `forall` | [1](#pft) | +| 3 | type index | 类型中的参数名 | `forall`、构造参数、`let` | [2](#index) | +| 4 | kind | type index 的种类，决定其取值范围与可用的运算 | `forall: {M: Dim}` | [2](#index) | +| 5 | discriminant | 取值有限的量，用于选择 type family 的分支，或决定张量是否存在 | `match`、`optional`、`nullable`、`mutated` 中的表达式 | [3](#family) | +| 6 | type family | 根据 discriminant 的取值给出形状的函数 | `signature.types` | [3](#family) | +| 7 | ADT、constructor | 带 tag 的 sum type，每个 constructor 有各自的字段 | `spec/types.yaml` 中的 `adts` | [4](#adt) | +| 8 | `let` | 由 index 计算得到的具名量 | `let: {H_out: "..."}` | [5](#let) | +| 9 | refinement | 约束 index 取值的谓词 | `shape_rules` 中的每一条 | [6](#refinement) | +| 10 | unification | 将声明的形状与实际形状逐轴对应，求出未知的 index | 由签名生成 | [7](#unification) | +| 11 | type inference | 调用时求出全部 index 的过程 | 由签名生成 | [7](#unification) | +| 12 | effect | 调用对参数的写入与别名关系 | `mutated`、`write_only`、`buffer`、`alias` | [8](#effect) | +| 13 | `Maybe[X]`、`present` | 可能缺省的 `X`；`present(x)` 表示它是否给出 | 可选输入、可空输出、`int \| None` 参数 | [扩展写法 1](extensions.md#presence) | +| 14 | primitive | 表达式中可以调用的内建函数，例如 `broadcast`、`pool.out` | 表达式中的函数调用 | [扩展写法 4](extensions.md#let) | +| 15 | generator | 实例化时为 metadata 张量生成取值的函数，例如 `prefix_sum(q_lens)` | 张量的 `values` | [扩展写法 7](extensions.md#generators) | diff --git a/docs/user-guide/manifest/examples.md b/docs/user-guide/manifest/examples.md new file mode 100644 index 00000000..ce3b1f9a --- /dev/null +++ b/docs/user-guide/manifest/examples.md @@ -0,0 +1,226 @@ +# Examples + +This page collects five real specs from the manifest. Each spec is followed by a table that maps each form used in it to the section that describes it. The specs are taken from `src/tileops/manifest/`, with comments and some workload rows removed. + +## 1. GemmFwdOp: type family {#gemm} + +```yaml +# src/tileops/manifest/spec/gemm.yaml +GemmFwdOp: + ref_api: "torch.matmul" + family: gemm + status: implemented + + signature: + types: + Mat: + params: {t: Bool, R: Dim, C: Dim} + match: t + cases: + - {when: false, is: "[R, C]"} + - {when: true, is: "[C, R]"} + forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"} + params: + trans_a: {type: bool, default: false} + trans_b: {type: bool, default: true} + inputs: + a: {dtype: T, shape: "Mat[trans_a, M, K]"} + b: {dtype: T, shape: "Mat[trans_b, K, N]"} + outputs: + d: {dtype: T, shape: "[M, N]"} + + workloads: + - {M: 1024, N: 1024, K: 1024, trans_a: false, trans_b: false, dtype_cases: [{T: float16}, {T: bfloat16}], label: square-1k} + - {M: 128, N: 2112, K: 7168, trans_a: false, trans_b: true, dtype_cases: [{T: bfloat16}], label: ds-v3-decode-qkv-a} + # remaining rows omitted + + roofline: + flops: "2 * M * N * K" +``` + +**Table 1** The forms used in `GemmFwdOp` + +| No. | Form in the spec | Described in | +| --- | --- | --- | +| 1 | `types.Mat`: `trans_a` and `trans_b` decide the axis order of `a` and `b` | [Extensions § 2](extensions.md#type-family) | +| 2 | `forall` declares three axis lengths and one `DType` index | [Spec fields § 3](writing.md#forall) | +| 3 | `trans_a` and `trans_b` are both construction parameters and discriminants of `Mat` | [Spec fields § 4](writing.md#params) | +| 4 | workload rows give the values of parameters and indices; the first row expands into two cases, with ids `square-1k-float16` and `square-1k-bfloat16` | [Spec fields § 8](writing.md#workloads) | +| 5 | `roofline` gives only `flops`; `bytes` is derived from the signature and equals `(M*K + K*N + M*N)` times the bytes per element | [Spec fields § 9](writing.md#roofline) | + +## 2. MoePrePermuteFwdOp: ADT, let and generator {#moe} + +The type of the `layout` parameter is the ADT `MGroupedLayout` defined in `spec/types.yaml`; its definition is in [Extensions § 3](extensions.md#adt). + +```yaml +# src/tileops/manifest/spec/moe.yaml +MoePrePermuteFwdOp: + family: moe + status: implemented + signature: + types: + LayoutMetadata: + params: {layout: MGroupedLayout, E: Dim, P: Dim} + match: layout + cases: + - {when: {contiguous: {metadata_kind: physical_psum}}, is: "[E]"} + - {when: {contiguous: {metadata_kind: per_row}}, is: "[P]"} + forall: {T: Dim, H: Dim, K: Dim, D: "DType[float16 | bfloat16]"} + params: + layout: {type: MGroupedLayout} + num_local_experts: {type: int} + inputs: + hidden_states: {dtype: D, shape: "[T, H]", contiguous: true} + local_expert_ids: {dtype: int32, shape: "[T, K]", contiguous: true, values: "topk_ids(T, K, num_local_experts)", requires: ["in_range(0, num_local_experts)"]} + let: + P: "moe.capacity(layout, T * K, num_local_experts)" + outputs: + expert_input: {dtype: D, shape: "[P, H]", contiguous: true} + layout_metadata: {dtype: int32, shape: "LayoutMetadata[layout, num_local_experts, P]", contiguous: true} + inverse_indices: {dtype: int32, shape: "[T * K]", contiguous: true} + shape_rules: + - "layout.kind == 'contiguous'" + - "num_local_experts > 0" + - "T > 0" + - "K > 0" + workloads: + - {T: 32, H: 128, K: 2, num_local_experts: 4, layout: {contiguous: {packing: tight, metadata_kind: physical_psum, alignment: 1}}, dtype_cases: [{D: float16}, {D: bfloat16}], label: "decode-t32"} + # remaining rows omitted + roofline: + flops: "0" +``` + +**Table 2** The forms used in `MoePrePermuteFwdOp` + +| No. | Form in the spec | Described in | +| --- | --- | --- | +| 1 | `layout.kind == 'contiguous'` is a domain restriction that excludes the masked layout, so `LayoutMetadata` needs no case for masked | [Extensions § 2](extensions.md#type-family), [§ 3](extensions.md#adt) | +| 2 | `LayoutMetadata` pattern-matches on the constructor and the field `metadata_kind` | [Extensions § 2](extensions.md#type-family) | +| 3 | the `int` parameter `num_local_experts` appears directly in shapes | [Spec fields § 4](writing.md#params) | +| 4 | `P` in `let` is computed by the primitive `moe.capacity` and used in output shapes | [Extensions § 4](extensions.md#let) | +| 5 | the values of `local_expert_ids` are generated by the generator `topk_ids`, and its contents are constrained by `requires` | [Extensions § 7](extensions.md#generators) | +| 6 | each tensor declares `contiguous: true` | [Extensions § 5](extensions.md#placement) | +| 7 | `layout` in the workload row is written as an ADT value | [Extensions § 3](extensions.md#adt) | + +## 3. MeanPoolingFwdOp: optional inputs and metadata tensors {#meanpool} + +```yaml +# src/tileops/manifest/spec/pool.yaml +MeanPoolingFwdOp: + family: pool + status: implemented + + signature: + types: + PoolOut: + params: {p: Bool, B: Dim, NC: Dim, S: Dim, c: Dim, H: Dim, D: Dim} + match: p + cases: + - {when: true, is: "[B, NC, H, D]"} + - {when: false, is: "[B, ceil_div(S, c), H, D]"} + forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[float16 | bfloat16 | float32]", seq_lens: "Seq[Int]"} + params: + chunk_size: {type: int} + accum_dtype: {type: torch.dtype} + inputs: + x: {dtype: T, shape: "[B, S, H, D]"} + offsets: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} + outputs: + output: {dtype: T, shape: "PoolOut[present(offsets), B, NC, S, chunk_size, H, D]"} + shape_rules: + - "chunk_size > 0 and chunk_size % 32 == 0" + - "D <= 128 or D % 128 == 0" + - "category(accum_dtype) == 'float'" + + workloads: + - {B: 1, S: 8192, H: 64, D: 128, chunk_size: 64, accum_dtype: float32, dtype_cases: [{T: float16}, {T: bfloat16}], label: uniform-8k} + - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [offsets], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} + # remaining rows omitted + + roofline: + flops: "B * S * H * D + B * (NC if present(offsets) else ceil_div(S, chunk_size)) * H * D" +``` + +**Table 3** The forms used in `MeanPoolingFwdOp` + +| No. | Form in the spec | Described in | +| --- | --- | --- | +| 1 | `indices` is declared `optional: "present(offsets)"`, so it is passed together with `offsets` or omitted together with it | [Extensions § 1](extensions.md#presence) | +| 2 | the type family `PoolOut` uses `present(offsets)` as its discriminant | [Extensions § 2](extensions.md#type-family) | +| 3 | the first row does not pass `offsets` and does not use `seq_lens`, so the row does not give it; the second row passes `offsets` through `some: [offsets]` and gives `seq_lens`, written here as the value primitive call `"repeat(2048, 4)"` | [Extensions § 1](extensions.md#presence), [Spec fields § 8](writing.md#workloads) | +| 4 | `NS` and `NC` are solved by unification of the generator results and are not given in the row | [Extensions § 7](extensions.md#generators) | +| 5 | `category(accum_dtype) == 'float'` requires the accumulation dtype to be a floating-point type | [Extensions § 9](extensions.md#scalar-dtype) | +| 6 | the roofline formula distinguishes the two cases with `present(offsets)` | [Spec fields § 9](writing.md#roofline) | + +## 4. ReluFwdOp: arbitrary rank and effect {#relu} + +```yaml +# src/tileops/manifest/spec/elementwise_unary_activation.yaml +ReluFwdOp: + ref_api: "torch.nn.functional.relu" + family: elementwise + status: implemented + + signature: + forall: {S: Shape, T: "DType[float16 | bfloat16 | float32]"} + params: + inplace: {type: bool, default: false, kw_only: true} + inputs: + input: {dtype: T, shape: "[*S]", mutated: inplace} + outputs: + output: {dtype: T, shape: "[*S]", alias: input} + + workloads: + - {S: [2048, 4096], dtype_cases: [{T: float16}, {T: bfloat16}], label: "hidden-state-prefill"} + - {S: [1, 4096], dtype_cases: [{T: bfloat16}], label: "hidden-state-decode"} + + roofline: + flops: "prod(S)" +``` + +**Table 4** The forms used in `ReluFwdOp` + +| No. | Form in the spec | Described in | +| --- | --- | --- | +| 1 | the kind of `S` is `Shape`, and the input and output use the same shape term `[*S]` | [Spec fields § 3](writing.md#forall), [§ 5](writing.md#tensors) | +| 2 | `mutated: inplace` and `alias: input` mean that when `inplace` is true the op writes the input and returns it | [Extensions § 6](extensions.md#effects) | +| 3 | `S` in the workload rows is written as a list of integers | [Spec fields § 8](writing.md#workloads) | + +## 5. AlibiFwdOp: no tensor inputs {#alibi} + +```yaml +# src/tileops/manifest/spec/elementwise_generative.yaml +AlibiFwdOp: + family: elementwise + status: implemented + + signature: + params: + seq_len: {type: int, kw_only: true} + num_heads: {type: int, kw_only: true} + out_dtype: {type: "float16 | bfloat16 | float32", default: float32, kw_only: true} + device: {type: "torch.device | str | None", default: null, kw_only: true} + outputs: + output: {dtype: out_dtype, shape: "[num_heads, seq_len, seq_len]"} + shape_rules: + - "seq_len > 0" + - "num_heads > 0" + + workloads: + - {seq_len: 2048, num_heads: 32, out_dtype: float16, label: "mpt-7b-2k"} + - {seq_len: 2048, num_heads: 32, out_dtype: bfloat16, label: "mpt-7b-2k"} + # remaining rows omitted + + roofline: + flops: "3 * num_heads * seq_len * seq_len" +``` + +**Table 5** The forms used in `AlibiFwdOp` + +| No. | Form in the spec | Described in | +| --- | --- | --- | +| 1 | every index comes from a construction parameter, so the spec has no `forall` | [Spec fields § 4](writing.md#params) | +| 2 | the output's `dtype` uses the dtype parameter `out_dtype` directly | [Spec fields § 7](writing.md#dtype) | +| 3 | the op has no call-time tensor inputs, so it declares a `device` parameter | [Extensions § 5](extensions.md#placement), [Calls and validation § 4](calls.md#device) | +| 4 | the two rows have the same `label`, and their case ids are distinguished by the value of the dtype parameter: `mpt-7b-2k-float16` and `mpt-7b-2k-bfloat16` | [Spec fields § 8](writing.md#workloads) | diff --git a/docs/user-guide/manifest/examples.zh.md b/docs/user-guide/manifest/examples.zh.md new file mode 100644 index 00000000..cbc87000 --- /dev/null +++ b/docs/user-guide/manifest/examples.zh.md @@ -0,0 +1,226 @@ +# 示例 + +本页收录 manifest 中的五个真实 spec，每个 spec 之后附有一张表，列出其中各处写法对应的说明。这些 spec 摘自 `src/tileops/manifest/`，删去了注释与部分 workload 行。 + +## 1. GemmFwdOp：type family {#gemm} + +```yaml +# src/tileops/manifest/spec/gemm.yaml +GemmFwdOp: + ref_api: "torch.matmul" + family: gemm + status: implemented + + signature: + types: + Mat: + params: {t: Bool, R: Dim, C: Dim} + match: t + cases: + - {when: false, is: "[R, C]"} + - {when: true, is: "[C, R]"} + forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"} + params: + trans_a: {type: bool, default: false} + trans_b: {type: bool, default: true} + inputs: + a: {dtype: T, shape: "Mat[trans_a, M, K]"} + b: {dtype: T, shape: "Mat[trans_b, K, N]"} + outputs: + d: {dtype: T, shape: "[M, N]"} + + workloads: + - {M: 1024, N: 1024, K: 1024, trans_a: false, trans_b: false, dtype_cases: [{T: float16}, {T: bfloat16}], label: square-1k} + - {M: 128, N: 2112, K: 7168, trans_a: false, trans_b: true, dtype_cases: [{T: bfloat16}], label: ds-v3-decode-qkv-a} + # 其余行省略 + + roofline: + flops: "2 * M * N * K" +``` + +**表 1** `GemmFwdOp` 中各处写法的说明 + +| No. | spec 中的写法 | 说明所在 | +| --- | --- | --- | +| 1 | `types.Mat`：由 `trans_a`、`trans_b` 决定 `a`、`b` 的轴顺序 | [扩展写法 2](extensions.md#type-family) | +| 2 | `forall` 声明三个轴长与一个 `DType` index | [写一个 spec 3](writing.md#forall) | +| 3 | `trans_a`、`trans_b` 既是构造参数，也是 `Mat` 的 discriminant | [写一个 spec 4](writing.md#params) | +| 4 | workload 行给出参数与 index 的取值；第一行展开为两个 case，id 分别是 `square-1k-float16` 与 `square-1k-bfloat16` | [写一个 spec 8](writing.md#workloads) | +| 5 | `roofline` 只给出 `flops`；`bytes` 由签名推导，等于 `(M*K + K*N + M*N)` 乘以每个元素的字节数 | [写一个 spec 9](writing.md#roofline) | + +## 2. MoePrePermuteFwdOp：ADT、let 与 generator {#moe} + +`layout` 参数的类型是 `spec/types.yaml` 中定义的 ADT `MGroupedLayout`，定义见[扩展写法 3](extensions.md#adt)。 + +```yaml +# src/tileops/manifest/spec/moe.yaml +MoePrePermuteFwdOp: + family: moe + status: implemented + signature: + types: + LayoutMetadata: + params: {layout: MGroupedLayout, E: Dim, P: Dim} + match: layout + cases: + - {when: {contiguous: {metadata_kind: physical_psum}}, is: "[E]"} + - {when: {contiguous: {metadata_kind: per_row}}, is: "[P]"} + forall: {T: Dim, H: Dim, K: Dim, D: "DType[float16 | bfloat16]"} + params: + layout: {type: MGroupedLayout} + num_local_experts: {type: int} + inputs: + hidden_states: {dtype: D, shape: "[T, H]", contiguous: true} + local_expert_ids: {dtype: int32, shape: "[T, K]", contiguous: true, values: "topk_ids(T, K, num_local_experts)", requires: ["in_range(0, num_local_experts)"]} + let: + P: "moe.capacity(layout, T * K, num_local_experts)" + outputs: + expert_input: {dtype: D, shape: "[P, H]", contiguous: true} + layout_metadata: {dtype: int32, shape: "LayoutMetadata[layout, num_local_experts, P]", contiguous: true} + inverse_indices: {dtype: int32, shape: "[T * K]", contiguous: true} + shape_rules: + - "layout.kind == 'contiguous'" + - "num_local_experts > 0" + - "T > 0" + - "K > 0" + workloads: + - {T: 32, H: 128, K: 2, num_local_experts: 4, layout: {contiguous: {packing: tight, metadata_kind: physical_psum, alignment: 1}}, dtype_cases: [{D: float16}, {D: bfloat16}], label: "decode-t32"} + # 其余行省略 + roofline: + flops: "0" +``` + +**表 2** `MoePrePermuteFwdOp` 中各处写法的说明 + +| No. | spec 中的写法 | 说明所在 | +| --- | --- | --- | +| 1 | `layout.kind == 'contiguous'` 是一条定义域限制，它排除了 masked layout，因此 `LayoutMetadata` 不需要 masked 对应的 case | [扩展写法 2](extensions.md#type-family)、[3](extensions.md#adt) | +| 2 | `LayoutMetadata` 按 constructor 与字段 `metadata_kind` 做 pattern matching | [扩展写法 2](extensions.md#type-family) | +| 3 | `int` 参数 `num_local_experts` 直接出现在形状中 | [写一个 spec 4](writing.md#params) | +| 4 | `let` 中的 `P` 由 primitive `moe.capacity` 计算得到，并用于输出形状 | [扩展写法 4](extensions.md#let) | +| 5 | `local_expert_ids` 的取值由 generator `topk_ids` 生成，其内容由 `requires` 约束 | [扩展写法 7](extensions.md#generators) | +| 6 | 各张量声明了 `contiguous: true` | [扩展写法 5](extensions.md#placement) | +| 7 | workload 行中的 `layout` 写成 ADT 值 | [扩展写法 3](extensions.md#adt) | + +## 3. MeanPoolingFwdOp：可选输入与 metadata 张量 {#meanpool} + +```yaml +# src/tileops/manifest/spec/pool.yaml +MeanPoolingFwdOp: + family: pool + status: implemented + + signature: + types: + PoolOut: + params: {p: Bool, B: Dim, NC: Dim, S: Dim, c: Dim, H: Dim, D: Dim} + match: p + cases: + - {when: true, is: "[B, NC, H, D]"} + - {when: false, is: "[B, ceil_div(S, c), H, D]"} + forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[float16 | bfloat16 | float32]", seq_lens: "Seq[Int]"} + params: + chunk_size: {type: int} + accum_dtype: {type: torch.dtype} + inputs: + x: {dtype: T, shape: "[B, S, H, D]"} + offsets: {dtype: int32, optional: true, shape: "[NS + 1]", values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", values: "chunk_indices(seq_lens, chunk_size)"} + outputs: + output: {dtype: T, shape: "PoolOut[present(offsets), B, NC, S, chunk_size, H, D]"} + shape_rules: + - "chunk_size > 0 and chunk_size % 32 == 0" + - "D <= 128 or D % 128 == 0" + - "category(accum_dtype) == 'float'" + + workloads: + - {B: 1, S: 8192, H: 64, D: 128, chunk_size: 64, accum_dtype: float32, dtype_cases: [{T: float16}, {T: bfloat16}], label: uniform-8k} + - {B: 1, S: 8192, H: 64, D: 128, seq_lens: "repeat(2048, 4)", chunk_size: 64, accum_dtype: float32, some: [offsets], dtype_cases: [{T: float16}, {T: float32}], label: ragged-even} + # 其余行省略 + + roofline: + flops: "B * S * H * D + B * (NC if present(offsets) else ceil_div(S, chunk_size)) * H * D" +``` + +**表 3** `MeanPoolingFwdOp` 中各处写法的说明 + +| No. | spec 中的写法 | 说明所在 | +| --- | --- | --- | +| 1 | `indices` 声明为 `optional: "present(offsets)"`，因此与 `offsets` 同时传入或同时不传 | [扩展写法 1](extensions.md#presence) | +| 2 | type family `PoolOut` 以 `present(offsets)` 为 discriminant | [扩展写法 2](extensions.md#type-family) | +| 3 | 第一行不传 `offsets`，没有用到 `seq_lens`，因此行中不给出它；第二行通过 `some: [offsets]` 传入 `offsets`，同时给出 `seq_lens`，这里写成值 primitive 调用 `"repeat(2048, 4)"` | [扩展写法 1](extensions.md#presence)、[写一个 spec 8](writing.md#workloads) | +| 4 | `NS` 与 `NC` 由 generator 结果的 unification 求得，不在行中给出 | [扩展写法 7](extensions.md#generators) | +| 5 | `category(accum_dtype) == 'float'` 限制累加 dtype 必须是浮点类型 | [扩展写法 9](extensions.md#scalar-dtype) | +| 6 | roofline 公式通过 `present(offsets)` 区分两种情况 | [写一个 spec 9](writing.md#roofline) | + +## 4. ReluFwdOp：任意秩与 effect {#relu} + +```yaml +# src/tileops/manifest/spec/elementwise_unary_activation.yaml +ReluFwdOp: + ref_api: "torch.nn.functional.relu" + family: elementwise + status: implemented + + signature: + forall: {S: Shape, T: "DType[float16 | bfloat16 | float32]"} + params: + inplace: {type: bool, default: false, kw_only: true} + inputs: + input: {dtype: T, shape: "[*S]", mutated: inplace} + outputs: + output: {dtype: T, shape: "[*S]", alias: input} + + workloads: + - {S: [2048, 4096], dtype_cases: [{T: float16}, {T: bfloat16}], label: "hidden-state-prefill"} + - {S: [1, 4096], dtype_cases: [{T: bfloat16}], label: "hidden-state-decode"} + + roofline: + flops: "prod(S)" +``` + +**表 4** `ReluFwdOp` 中各处写法的说明 + +| No. | spec 中的写法 | 说明所在 | +| --- | --- | --- | +| 1 | `S` 的 kind 是 `Shape`，输入与输出使用同一个形状项 `[*S]` | [写一个 spec 3](writing.md#forall)、[5](writing.md#tensors) | +| 2 | `mutated: inplace` 与 `alias: input` 表示 `inplace` 为 true 时 op 写入输入并将其返回 | [扩展写法 6](extensions.md#effects) | +| 3 | workload 行中的 `S` 写成整数列表 | [写一个 spec 8](writing.md#workloads) | + +## 5. AlibiFwdOp：没有张量输入 {#alibi} + +```yaml +# src/tileops/manifest/spec/elementwise_generative.yaml +AlibiFwdOp: + family: elementwise + status: implemented + + signature: + params: + seq_len: {type: int, kw_only: true} + num_heads: {type: int, kw_only: true} + out_dtype: {type: "float16 | bfloat16 | float32", default: float32, kw_only: true} + device: {type: "torch.device | str | None", default: null, kw_only: true} + outputs: + output: {dtype: out_dtype, shape: "[num_heads, seq_len, seq_len]"} + shape_rules: + - "seq_len > 0" + - "num_heads > 0" + + workloads: + - {seq_len: 2048, num_heads: 32, out_dtype: float16, label: "mpt-7b-2k"} + - {seq_len: 2048, num_heads: 32, out_dtype: bfloat16, label: "mpt-7b-2k"} + # 其余行省略 + + roofline: + flops: "3 * num_heads * seq_len * seq_len" +``` + +**表 5** `AlibiFwdOp` 中各处写法的说明 + +| No. | spec 中的写法 | 说明所在 | +| --- | --- | --- | +| 1 | 所有 index 都来自构造参数，因此 spec 中没有 `forall` | [写一个 spec 4](writing.md#params) | +| 2 | 输出的 `dtype` 直接使用 dtype 参数 `out_dtype` | [写一个 spec 7](writing.md#dtype) | +| 3 | op 没有调用期张量输入，因此声明了 `device` 参数 | [扩展写法 5](extensions.md#placement)、[调用与校验 4](calls.md#device) | +| 4 | 两行的 `label` 相同，case id 以 dtype 参数的值区分，分别是 `mpt-7b-2k-float16`、`mpt-7b-2k-bfloat16` | [写一个 spec 8](writing.md#workloads) | diff --git a/docs/user-guide/manifest/extensions.md b/docs/user-guide/manifest/extensions.md new file mode 100644 index 00000000..6805bc93 --- /dev/null +++ b/docs/user-guide/manifest/extensions.md @@ -0,0 +1,253 @@ +# Extensions + +This page describes forms that only some specs use. Each section covers one case, and each section can be read on its own. The fields most specs need are in [Spec fields](writing.md). + +## 1. Optional inputs, nullable outputs and omittable parameters {#presence} + +Whether a tensor exists is expressed with `present`. FusedMoeSharedExpert is an example: + +```yaml +inputs: + correction_bias: {dtype: float32, shape: "[E]", optional: true} + shared_w_gate_up: {dtype: D, shape: "[2 * S, H]", optional: true} + shared_w_down: {dtype: D, shape: "[H, S]", optional: "present(shared_w_gate_up)"} +outputs: + shared_output: {dtype: D, shape: "[T, H]", nullable: "present(shared_w_gate_up)"} + routed_output: {dtype: D, shape: "[T, H]"} +``` + +- An optional input is declared with `optional: true`. If several tensors must be passed together or omitted together, they share one discriminant, written `optional: ""`, as `shared_w_down` does above. +- An output that may be `None` is declared with `nullable: ""`. In the example above, the op returns `shared_output` when `shared_w_gate_up` is passed, and that output is `None` otherwise. +- The expressions of `optional` and `nullable` may consist only of boolean quantities with finitely many values: `Bool` and enum parameters, ADT tags and fields with finitely many values, and `present(...)`. +- Optional inputs come after required inputs, and `forward` receives them in declaration order with a default of `None`. Omitting an optional input has the same effect as passing `None` explicitly. +- The kind of an `int | None` parameter is `Maybe[Int]`: `present(v)` says whether it is given, and `v.value` is its value. `v.value` may appear only in branches where `present(v)` is true. +- An index is used only on the branches where it appears. In the example above, when `shared_w_gate_up` is not passed, `S` is not used, and the corresponding workload row does not give `S`. +- A workload row lists the optional tensors it passes in `some`. `some` lists only tensors declared `optional: true`; a tensor whose `optional` is an expression follows the value of that expression and is not written in `some`. For example, passing the shared expert is written `some: [shared_w_gate_up]`, and `shared_w_down` is passed along with it. +- For an `implemented` op, each optional tensor must be passed by at least one row and omitted by at least one row. + +## 2. Shapes that vary with parameters: type family {#type-family} + +When a parameter determines the rank of a tensor or the order of its axes, the shape is described with a type family. A type family is defined in the spec's `signature.types`: + +```yaml +signature: + types: + Mat: + params: {t: Bool, R: Dim, C: Dim} + match: t + cases: + - {when: false, is: "[R, C]"} + - {when: true, is: "[C, R]"} + inputs: + a: {dtype: T, shape: "Mat[trans_a, M, K]"} # arguments map to Mat's t, R, C in order +``` + +- A tensor's `shape` is written `[argument, ...]`, and the arguments map to `params` in declaration order. +- The subject of `match` must be a discriminant with finitely many values: `Bool`, an enum, an ADT, `present(...)`, or a tuple of these. When matching a tuple, `when` is written as a list, for example `when: [true, false]`. +- `cases` must have no gaps and no overlaps over all values the spec accepts. +- When several tensors apply the same type family, they always take the same branch. +- References between type families must not form a cycle, and the validator reports an error if a type family is not used by any shape. + +A refinement that reads only discriminants is called a **domain restriction**. A domain restriction is checked before a type family branch is selected, and values it excludes need no corresponding case. Clamp is an example: + +```yaml +types: + ClampOut: + params: {A: Shape, L: Shape, U: Shape, pl: Bool, pu: Bool} + match: [pl, pu] + cases: + - {when: [true, true], is: "[*broadcast(A, L, U)]"} + - {when: [true, false], is: "[*broadcast(A, L)]"} + - {when: [false, true], is: "[*broadcast(A, U)]"} +outputs: + output: {dtype: T, shape: "ClampOut[A, L, U, present(min), present(max)]"} +shape_rules: +- "present(min) or present(max)" +``` + +`present(min) or present(max)` reads only discriminants and excludes the case where neither `min` nor `max` is passed, so the three cases already cover every combination Clamp accepts. Whether a refinement reads only discriminants is decided by all the names that appear in it, independent of the order in which the operands are written. + +## 3. Parameters with fields: ADT {#adt} + +When a parameter has a finite number of forms and each form carries its own fields, the parameter's type is described with an ADT. ADTs are defined in `spec/types.yaml` and shared by several specs: + +```yaml +adts: + MGroupedLayout: + sum: + contiguous: + python: tileops.ops.moe.contracts.ContiguousLayoutSpec + fields: + packing: {type: "'tight' | 'aligned'", python: tileops.ops.moe.contracts.ContiguousPacking} + metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops.ops.moe.contracts.ContiguousMetadata} + alignment: Dim + invariant: "(packing == 'tight') == (alignment == 1) and alignment >= 1" + masked: + python: tileops.ops.moe.contracts.MaskedLayoutSpec + fields: {max_m: Dim} +``` + +- Each constructor corresponds to a Python class, given by `python`. The object's `kind` attribute is the constructor name, each field is an attribute of the same name, and the value of an enum field is the attribute's `.value`. +- An ADT value is written `{constructor: {field: value}}`, and workload rows use the same form. +- `invariant` is an optional refinement on a constructor, checked both at instantiation and at construction. +- A type family can pattern-match on constructors: `{masked: _}` matches any masked value, and `{contiguous: {metadata_kind: per_row}}` also constrains a field with finitely many values. A field specific to a constructor can be read only in a branch that matched that constructor, for example `layout.max_m` in the masked branch. +- ADTs are sealed: constructors and fields are fixed where the ADT is defined. Adding a constructor to a shared ADT is done by editing its definition directly; a spec that does not accept the new constructor can exclude it with a domain restriction, without changing its type families. + +## 4. Computed quantities: let and primitive {#let} + +When a shape or formula needs a quantity computed from indices, the quantity is defined as a `let`. MaxPool2d is an example: + +```yaml +let: + kH: "per_axis(kernel_size, 0, 2)" + sH: "per_axis(stride, 0, 2, fallback=kH)" + H_out: "pool.out(H_in, kH, sH, pH, dH, ceil_mode)" +outputs: + output: {dtype: T, shape: "[N, C, H_out, W_out]"} +``` + +- The value of a `let` is computed from the signature: a `let` that can be evaluated at construction is computed at construction, and the others are computed on every call. Workload rows do not give `let` entries. +- A `let` can refer to other `let` entries, as `sH` refers to `kH`, but dependencies between `let` entries must not form a cycle. + +A primitive is a built-in function callable in expressions, such as `broadcast`, `reduced`, `per_axis` and `ceil_div`. + +- Primitives, generators and the predicates in `requires` are each a fixed set. The full list is in `tileops.manifest.primitives`, and the docstring of each member describes what it computes. +- Members used by only one family carry a family prefix, such as `pool.out` and `moe.capacity`. +- When an argument is outside its domain, the primitive raises an error that points to the declaration that called it. +- Adding a member requires changing `tileops.manifest.primitives`, with tests. +- Every primitive that takes an axis argument handles axes by the same rule: for a zero-dimensional tensor, both `0` and `-1` denote the single scalar axis; otherwise, the axis ranges over `[-rank, rank)`. + +## 5. Construction-time tensors, memory layout and device {#placement} + +- A tensor passed at construction is declared in `params` with `dtype` and `shape`, and can declare `optional: true`. For example, `rescale_factors` in LongRoPE: + + ```yaml + params: + rescale_factors: {dtype: R, shape: "[D // 2]", optional: true} + ``` + +- A tensor that must be contiguous in memory declares `contiguous: true`, for example the inputs and outputs of the MoE staged ops. A tensor without this declaration can have any strides. +- A tensor that must be on the CPU declares `device: cpu`, for example `cu_seqlens_cpu` in GatedDeltaNet. +- An op without call-time tensor inputs declares a `device` parameter, for example Alibi. How the call device is determined is in [Calls and validation § 4](calls.md#device). + +## 6. Writes to arguments: effect {#effects} + +An op without effect declarations only reads its inputs and allocates new tensors for its outputs. If an op writes to an argument, an effect is declared on the corresponding tensor in the signature. Effects determine the read and write counts of the roofline, and the operator schema generated for an op that declares a compile boundary. + +**Table 1** Effect declarations + +| No. | Declaration | Meaning | Example | +| --- | --- | --- | --- | +| 1 | `buffer: out` on an output | `forward` gets a parameter `out` after all inputs. When the caller passes `out`, the op writes the result into it and returns it; otherwise the op allocates a new tensor. `out` has the same shape and dtype as that output | `output` of MoeGroupedGemm | +| 2 | `mutated: true` on an input | the op may write this input, and its contents before the call take part in the computation | | +| 3 | `mutated: true` and `write_only: true` on an input | a result buffer that must be passed: the op overwrites it, and the result depends only on the other inputs; if the op returns `None`, `outputs` is empty | `output` of FusedMoEExperts | +| 4 | `mutated: ""` on an input | the op writes this input only when the expression is true | `mutated: inplace` of activation functions | +| 5 | `alias: ` on an output | when that input is written, this output is that input object itself | `alias: input` of activation functions | + +The `inplace` parameter of activation functions is an example: + +```yaml +params: + inplace: {type: bool, default: false, kw_only: true} +inputs: + input: {dtype: T, shape: "[*S]", mutated: inplace} +outputs: + output: {dtype: T, shape: "[*S]", alias: input} +``` + +Effect declarations must also satisfy the following rules, and the validator rejects specs that violate them: + +- `write_only: true` can be used only together with `mutated: true`; +- `alias` must point to an input that is written; +- an output that declares `alias` cannot also declare `buffer`; +- at most one output declares `buffer: out`. + +For each effect branch, the validator checks that the operator schema, the aliasing and the roofline read and write counts agree. + +## 7. Metadata tensors: generator and requires {#generators} + +Varlen, paged and similar ops use metadata tensors such as `cu_seqlens` and `block_table`. The type of a metadata tensor is written in the signature, and its values are generated at instantiation by a generator written in the tensor's `values` field. MeanPooling is an example: + +```yaml +forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[...]", seq_lens: "Seq[Int]"} +inputs: + offsets: {dtype: int32, optional: true, shape: "[NS + 1]", + values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", + values: "chunk_indices(seq_lens, chunk_size)"} +``` + +The rules for generators: + +- The arguments of a generator are given by the workload row, for example the list of lengths `seq_lens`. A name of kind `Seq[Int]` in `forall` can be used only as a generator argument. A `Seq[Int]` in a workload row can be written as a list of integers or as a call to a value primitive that returns a list, for example `seq_lens: "repeat(512, 64)"`. +- At instantiation, the generator's result is unified with the declared shape, which solves the other indices in the shape, so they are not given in the row. In the example above, `NS` and `NC` are solved from the generated `offsets` and `indices` respectively. In an actual call, these indices are likewise solved from the input tensors. +- A generator is either deterministic or uses private random numbers derived from the workload seed, so the same row generates the same values every time. +- The rank of each generator result is fixed, or determined by its shape arguments. +- A generated tensor declares an integer dtype (`int32` or `int64`). When an argument is outside its domain, or the result is outside the range of the declared dtype, the generator raises an error. +- A generator argument can be a primitive that returns a list, for example `as_tensor(balanced_sizes(M, G))` in GroupedGemm. + +The rules for `requires`: + +- `requires` lists predicates that constrain the contents of a metadata tensor. For example, `prefix_offsets(S)` requires the tensor's first element to be 0, the elements to be non-decreasing, and the last element to be `S`. The contents of the constrained tensor are the predicate's implicit first argument. +- Each predicate reads the constrained tensor at a fixed rank; predicates that give elementwise bounds (such as `in_range`) can apply to tensors of any rank. +- A predicate argument can be another metadata tensor, so that one predicate constrains the relation between two tensors. For example, `batch_offsets` in GroupedGemm declares `requires: ["exclusive_prefix_of(batch_sizes)"]`, which requires it to be the exclusive prefix sum of `batch_sizes`. On every branch where the constrained tensor exists, the tensor used as the argument must also exist. +- The validator checks `requires` at instantiation against the generated values. In an actual call, the caller guarantees these constraints and the op does not check tensor contents; the validator therefore also checks that the predicate is well defined on every branch where the constrained tensor exists. +- A tensor that declares `requires` must also declare `values`. + +## 8. dtype combinations and packed dtypes {#dtype-combos} + +When several `DType` indices allow only specific combinations, the spec lists every allowed combination in `dtype_combos`. Paged GQA is an example: + +```yaml +forall: {..., T: "DType[float16 | bfloat16 | float8_e4m3fn]", KV: "DType[float16 | bfloat16 | float8_e4m3fn]"} +dtype_combos: +- {T: float16, KV: float16} +- {T: bfloat16, KV: bfloat16} +- {T: float16, KV: float8_e4m3fn} +- {T: bfloat16, KV: float8_e4m3fn} +- {T: float8_e4m3fn, KV: float8_e4m3fn} +``` + +- Each row is a mapping from index to dtype. All rows have the same keys, the keys may include dtype parameters, and the rows are all distinct. +- The dtype values of a call must equal one of the rows exactly. +- Every column of `dtype_combos` must be used on every branch the spec accepts. + +Packed dtypes such as fp4 and int4 are stored in a carrier dtype such as `uint8`, and the spec is written in terms of the carrier: + +- `dtype` is the carrier dtype, and `shape` is the carrier shape as PyTorch sees it, for example `packed_weight: "[N, K // 2]"` in GemmW4A16; +- the logical dtype is given by a dtype parameter, or fixed in the spec; +- the roofline counts bytes by the carrier. + +## 9. Scalar parameters whose value range depends on dtype {#scalar-dtype} + +The valid values of some scalar parameters depend on the dtype at call time, for example `alpha` in Add. Such constraints are written as refinements with the primitives `category` and `representable`: + +```yaml +params: + alpha: {type: int | float, default: 1, kw_only: true} +shape_rules: +- "category(alpha) == 'int' or category(alpha) == category(T)" +- "representable(alpha, T)" +``` + +The generated call checks apply the same rule to TileOPs in-tree implementations and to implementations provided by a target. + +## 10. Composite ops: composition {#composition} + +A composite op uses `composition` to record the sub-ops its in-tree implementation may hold, and the positions of its own kernels. FusedMoeSharedExpert is an example: + +```yaml +composition: + kind: composite + stages: + - {name: route_select, op: FusedTopKFwdOp} + - {name: routed_experts, op: FusedMoEExpertsFwdOp} + - {name: shared_expert, op: SharedExpertMLPFwdOp, optional: true} +``` + +- Each stage refers either to an op in the manifest (`op`) or to a key in the op's own `kernel_types` (`kernel`). +- A sub-op that is not held on every call is written as a stage with `optional: true`; whether it is held is decided by code. +- When and how many times sub-ops are constructed, how they are scheduled, and how forward executes are all decided by code. +- The manifest does not prescribe how the roofline of a composite op relates to the rooflines of its stages. +- For an `implemented` op, the validator checks in order that the `op` stages agree with the class's `delegate_types` and the `kernel` stages agree with the class's `kernel_types`. +- `stages` cannot be empty, stage names are all distinct, and `optional` is a boolean. diff --git a/docs/user-guide/manifest/extensions.zh.md b/docs/user-guide/manifest/extensions.zh.md new file mode 100644 index 00000000..fe9e4fb0 --- /dev/null +++ b/docs/user-guide/manifest/extensions.zh.md @@ -0,0 +1,253 @@ +# 扩展写法 + +本页介绍只有部分 spec 才会用到的写法。每一节对应一种情况，各节可以独立阅读。大多数 spec 所需的字段见[写一个 spec](writing.md)。 + +## 1. 可选输入、可空输出与可缺省参数 {#presence} + +张量是否存在通过 `present` 表达。以 FusedMoeSharedExpert 为例： + +```yaml +inputs: + correction_bias: {dtype: float32, shape: "[E]", optional: true} + shared_w_gate_up: {dtype: D, shape: "[2 * S, H]", optional: true} + shared_w_down: {dtype: D, shape: "[H, S]", optional: "present(shared_w_gate_up)"} +outputs: + shared_output: {dtype: D, shape: "[T, H]", nullable: "present(shared_w_gate_up)"} + routed_output: {dtype: D, shape: "[T, H]"} +``` + +- 可选输入使用 `optional: true` 声明。如果几个张量必须同时传入或同时不传，它们共用同一个 discriminant，写作 `optional: "<表达式>"`，例如上例中的 `shared_w_down`。 +- 可能返回 `None` 的输出使用 `nullable: "<表达式>"` 声明。在上例中，传入 `shared_w_gate_up` 时 op 返回 `shared_output`，否则该输出为 `None`。 +- `optional` 与 `nullable` 的表达式只能由取值有限的布尔量组成，包括 `Bool` 与枚举参数、ADT 的 tag 与取值有限的字段，以及 `present(...)`。 +- 可选输入排在必选输入之后，`forward` 按声明顺序接收它们，默认值为 `None`。省略一个可选输入与显式传入 `None` 效果相同。 +- `int | None` 参数的 kind 是 `Maybe[Int]`：`present(v)` 表示它是否给出，`v.value` 是它的值。`v.value` 只能出现在 `present(v)` 为 true 的分支中。 +- 一个 index 只在出现它的分支上被用到。在上例中，不传 `shared_w_gate_up` 时 `S` 没有被用到，对应的 workload 行不需要给出 `S`。 +- workload 行通过 `some` 列出传入的可选张量。`some` 只列出声明为 `optional: true` 的张量；`optional` 为表达式的张量随表达式的取值确定，不写在 `some` 中。例如传入共享专家时写 `some: [shared_w_gate_up]`，`shared_w_down` 随之传入。 +- 对于 `implemented` 的 op，每个可选张量都要至少有一条行传入它，并至少有一条行不传它。 + +## 2. 随参数变化的形状：type family {#type-family} + +当参数决定张量的秩或轴的顺序时，形状使用 type family 描述。type family 定义在 spec 的 `signature.types` 中： + +```yaml +signature: + types: + Mat: + params: {t: Bool, R: Dim, C: Dim} + match: t + cases: + - {when: false, is: "[R, C]"} + - {when: true, is: "[C, R]"} + inputs: + a: {dtype: T, shape: "Mat[trans_a, M, K]"} # 实参依次对应 Mat 的 t、R、C +``` + +- 张量的 `shape` 写作 `[实参, ...]`，实参按 `params` 的声明顺序对应。 +- `match` 的对象必须是取值有限的 discriminant，即 `Bool`、枚举、ADT、`present(...)`，或由它们组成的元组。对元组做 match 时，`when` 写成列表，例如 `when: [true, false]`。 +- `cases` 在 spec 接受的所有取值上必须既无遗漏也无重叠。 +- 多个张量应用同一个 type family 时，它们总是取同一个分支。 +- type family 之间的引用不能成环；如果某个 type family 没有被任何形状使用，validator 会报错。 + +只读取 discriminant 的 refinement 称为**定义域限制**。定义域限制在选择 type family 分支之前检查，被它排除的取值不需要对应的 case。以 Clamp 为例： + +```yaml +types: + ClampOut: + params: {A: Shape, L: Shape, U: Shape, pl: Bool, pu: Bool} + match: [pl, pu] + cases: + - {when: [true, true], is: "[*broadcast(A, L, U)]"} + - {when: [true, false], is: "[*broadcast(A, L)]"} + - {when: [false, true], is: "[*broadcast(A, U)]"} +outputs: + output: {dtype: T, shape: "ClampOut[A, L, U, present(min), present(max)]"} +shape_rules: +- "present(min) or present(max)" +``` + +`present(min) or present(max)` 只读取 discriminant，它排除了 `min` 与 `max` 都不传的情况，因此三个 case 已经覆盖了 Clamp 接受的所有组合。判断一条 refinement 是否只读取 discriminant 时，考察的是它出现的所有名字，与操作数的书写顺序无关。 + +## 3. 带字段的参数：ADT {#adt} + +当一个参数有有限的几种形态，且每种形态带有各自的字段时，参数的类型使用 ADT 描述。ADT 定义在 `spec/types.yaml` 中，供多个 spec 共用： + +```yaml +adts: + MGroupedLayout: + sum: + contiguous: + python: tileops.ops.moe.contracts.ContiguousLayoutSpec + fields: + packing: {type: "'tight' | 'aligned'", python: tileops.ops.moe.contracts.ContiguousPacking} + metadata_kind: {type: "'physical_psum' | 'per_row'", python: tileops.ops.moe.contracts.ContiguousMetadata} + alignment: Dim + invariant: "(packing == 'tight') == (alignment == 1) and alignment >= 1" + masked: + python: tileops.ops.moe.contracts.MaskedLayoutSpec + fields: {max_m: Dim} +``` + +- 每个 constructor 对应一个 Python 类，由 `python` 指定。对象的 `kind` 属性是 constructor 名，各字段是同名属性，枚举字段的取值是属性的 `.value`。 +- ADT 值写作 `{constructor: {字段: 值}}`，workload 行中也使用这种写法。 +- `invariant` 是 constructor 上可选的 refinement，在实例化与构造时都会检查。 +- type family 可以按 constructor 做 pattern matching：`{masked: _}` 匹配任意 masked 值，`{contiguous: {metadata_kind: per_row}}` 同时约束了一个取值有限的字段。constructor 特有的字段只能在匹配了该 constructor 的分支中读取，例如 masked 分支中的 `layout.max_m`。 +- ADT 是 sealed 的，constructor 与字段在定义处确定。为共用的 ADT 增加 constructor 时直接修改其定义；不接受新 constructor 的 spec 可以用一条定义域限制排除它，无需修改 type family。 + +## 4. 计算得到的量：let 与 primitive {#let} + +形状或公式中需要用到由 index 计算得到的量时，这个量定义为 `let`。以 MaxPool2d 为例： + +```yaml +let: + kH: "per_axis(kernel_size, 0, 2)" + sH: "per_axis(stride, 0, 2, fallback=kH)" + H_out: "pool.out(H_in, kH, sH, pH, dH, ceil_mode)" +outputs: + output: {dtype: T, shape: "[N, C, H_out, W_out]"} +``` + +- `let` 的值由签名计算：能在构造时求值的在构造时计算，其余的在每次调用时计算。workload 行中不需要给出 `let`。 +- 一个 `let` 可以引用其他 `let`，例如 `sH` 引用了 `kH`，但 `let` 之间的依赖不能成环。 + +primitive 是表达式中可以调用的内建函数，例如 `broadcast`、`reduced`、`per_axis`、`ceil_div`。 + +- primitive、generator 与 `requires` 中的谓词都是固定的集合，完整列表在 `tileops.manifest.primitives` 中，每个成员的 docstring 说明了它的计算内容。 +- 只供某一个 family 使用的成员带有 family 前缀，例如 `pool.out`、`moe.capacity`。 +- 参数超出定义域时，primitive 会报错，错误信息指向调用它的那条声明。 +- 新增成员需要修改 `tileops.manifest.primitives`，并附带测试。 +- 所有接受轴参数的 primitive 都按同一规则处理轴：对零维张量，`0` 与 `-1` 都表示唯一的标量轴；其他情况下，轴的取值范围是 `[-rank, rank)`。 + +## 5. 构造期张量、内存布局与设备 {#placement} + +- 在构造时传入的张量声明在 `params` 中，带有 `dtype` 与 `shape`，并可以声明 `optional: true`。例如 LongRoPE 的 `rescale_factors`： + + ```yaml + params: + rescale_factors: {dtype: R, shape: "[D // 2]", optional: true} + ``` + +- 要求内存连续的张量声明 `contiguous: true`，例如 MoE staged 系列 op 的输入与输出。没有这项声明的张量可以有任意 stride。 +- 必须位于 CPU 上的张量声明 `device: cpu`，例如 GatedDeltaNet 的 `cu_seqlens_cpu`。 +- 没有调用期张量输入的 op 声明 `device` 参数，例如 Alibi。调用设备的确定方式见[调用与校验 4](calls.md#device)。 + +## 6. 对参数的写入：effect {#effects} + +没有 effect 声明的 op 只读取输入，并为输出分配新的张量。如果 op 会写入参数，需要在签名的相应张量上声明 effect。effect 决定了 roofline 的读写计数，以及声明了 compile boundary 的 op 所生成的 operator schema。 + +**表 1** effect 声明 + +| No. | 声明 | 含义 | 例子 | +| --- | --- | --- | --- | +| 1 | 输出上的 `buffer: out` | `forward` 在所有输入之后增加参数 `out`。调用方传入 `out` 时，op 将结果写入并返回它；未传入时，op 分配新张量。`out` 与该输出的形状和 dtype 相同 | MoeGroupedGemm 的 `output` | +| 2 | 输入上的 `mutated: true` | op 可能写入这个输入，它在调用前的内容参与计算 | | +| 3 | 输入上的 `mutated: true` 与 `write_only: true` | 必须传入的结果缓冲：op 覆盖写入，结果只取决于其他输入；如果 op 返回 `None`，`outputs` 为空 | FusedMoEExperts 的 `output` | +| 4 | 输入上的 `mutated: ""` | 仅当表达式为 true 时，op 才写入这个输入 | 激活函数的 `mutated: inplace` | +| 5 | 输出上的 `alias: <输入名>` | 该输入被写入时，这个输出就是该输入对象本身 | 激活函数的 `alias: input` | + +以激活函数的 `inplace` 为例： + +```yaml +params: + inplace: {type: bool, default: false, kw_only: true} +inputs: + input: {dtype: T, shape: "[*S]", mutated: inplace} +outputs: + output: {dtype: T, shape: "[*S]", alias: input} +``` + +effect 声明还需满足以下规则，validator 会拒绝违反它们的 spec： + +- `write_only: true` 只能与 `mutated: true` 同时使用； +- `alias` 指向的必须是一个会被写入的输入； +- 声明了 `alias` 的输出不能再声明 `buffer`； +- 至多一个输出声明 `buffer: out`。 + +validator 对每个 effect 分支检查 operator schema、别名关系与 roofline 的读写计数是否一致。 + +## 7. metadata 张量：generator 与 requires {#generators} + +varlen、paged 等 op 使用 metadata 张量，例如 `cu_seqlens`、`block_table`。metadata 张量的类型写在签名中，取值则由 generator 在实例化时生成，generator 写在张量的 `values` 字段中。以 MeanPooling 为例： + +```yaml +forall: {B: Dim, S: Dim, H: Dim, D: Dim, NS: Dim, NC: Dim, T: "DType[...]", seq_lens: "Seq[Int]"} +inputs: + offsets: {dtype: int32, optional: true, shape: "[NS + 1]", + values: "prefix_sum(seq_lens)", requires: ["prefix_offsets(S)"]} + indices: {dtype: int32, optional: "present(offsets)", shape: "[NC, 2]", + values: "chunk_indices(seq_lens, chunk_size)"} +``` + +generator 的规则如下： + +- generator 的实参由 workload 行给出，例如长度列表 `seq_lens`。`forall` 中 kind 为 `Seq[Int]` 的名字只能用作 generator 的实参。workload 行中的 `Seq[Int]` 可以写成整数列表，也可以写成返回列表的值 primitive 调用，例如 `seq_lens: "repeat(512, 64)"`。 +- 实例化时，generator 的结果与声明的形状做 unification，形状中的其他 index 由此求得，不需要在行中给出。上例中的 `NS` 与 `NC` 分别由 `offsets` 与 `indices` 的生成结果求得。实际调用时，这些 index 同样由输入张量求得。 +- generator 或者是确定性的，或者使用由 workload 种子派生的私有随机数，因此同一条行每次生成的取值都相同。 +- 每个 generator 结果的秩是固定的，或者由其形状参数决定。 +- 生成的张量声明整数 dtype（`int32` 或 `int64`）；参数超出定义域，或结果超出所声明 dtype 的范围时，generator 会报错。 +- generator 的实参可以是返回列表的 primitive，例如 GroupedGemm 的 `as_tensor(balanced_sizes(M, G))`。 + +`requires` 的规则如下： + +- `requires` 列出约束 metadata 张量内容的谓词。例如 `prefix_offsets(S)` 要求张量首项为 0、单调不减、末项为 `S`。被约束张量的内容是谓词隐含的第一个实参。 +- 每个谓词按固定的秩读取被约束的张量；逐元素给出上下界的谓词（如 `in_range`）可以作用于任意秩的张量。 +- 谓词的实参可以是另一个 metadata 张量，这样一条谓词就能约束两个张量之间的关系，例如 GroupedGemm 的 `batch_offsets` 声明 `requires: ["exclusive_prefix_of(batch_sizes)"]`，要求它是 `batch_sizes` 的 exclusive 前缀和。在被约束的张量存在的每个分支上，作为实参的那个张量也必须存在。 +- validator 在实例化时用生成的取值检查 `requires`。在实际调用中，这些约束由调用方保证，op 不检查张量的内容；因此 validator 还会检查，在被约束的张量存在的每个分支上，谓词都是良定义的。 +- 声明了 `requires` 的张量也必须声明 `values`。 + +## 8. dtype 组合与打包 dtype {#dtype-combos} + +当多个 `DType` index 只允许特定组合时，spec 用 `dtype_combos` 列出所有允许的组合。以 paged GQA 为例： + +```yaml +forall: {..., T: "DType[float16 | bfloat16 | float8_e4m3fn]", KV: "DType[float16 | bfloat16 | float8_e4m3fn]"} +dtype_combos: +- {T: float16, KV: float16} +- {T: bfloat16, KV: bfloat16} +- {T: float16, KV: float8_e4m3fn} +- {T: bfloat16, KV: float8_e4m3fn} +- {T: float8_e4m3fn, KV: float8_e4m3fn} +``` + +- 每一行是一个从 index 到 dtype 的映射。各行的键相同，键中可以包含 dtype 参数；各行互不相同。 +- 一次调用的 dtype 取值必须与其中某一行完全相同。 +- `dtype_combos` 的每一列在 spec 接受的每个分支上都必须被用到。 + +fp4、int4 等打包 dtype 存放在 `uint8` 等载体 dtype 中，spec 按载体书写： + +- `dtype` 写载体 dtype，`shape` 写 PyTorch 所见的载体形状，例如 GemmW4A16 的 `packed_weight: "[N, K // 2]"`； +- 逻辑 dtype 由 dtype 参数给出，或在 spec 中固定； +- roofline 按载体计算字节数。 + +## 9. 取值范围随 dtype 变化的标量参数 {#scalar-dtype} + +有些标量参数的合法取值取决于调用时的 dtype，例如 Add 的 `alpha`。这类约束使用 primitive `category` 与 `representable` 写成 refinement： + +```yaml +params: + alpha: {type: int | float, default: 1, kw_only: true} +shape_rules: +- "category(alpha) == 'int' or category(alpha) == category(T)" +- "representable(alpha, T)" +``` + +生成的调用检查对 TileOPs in-tree 实现与 target 提供的实现应用同一条规则。 + +## 10. 复合 op：composition {#composition} + +复合 op 用 `composition` 记录其 in-tree 实现可能持有的子 op，以及它自身 kernel 的位置。以 FusedMoeSharedExpert 为例： + +```yaml +composition: + kind: composite + stages: + - {name: route_select, op: FusedTopKFwdOp} + - {name: routed_experts, op: FusedMoEExpertsFwdOp} + - {name: shared_expert, op: SharedExpertMLPFwdOp, optional: true} +``` + +- 每个 stage 或者引用一个 manifest 中的 op（`op`），或者引用 op 自身 `kernel_types` 中的一个键（`kernel`）。 +- 并非每次调用都会持有的子 op，写成 `optional: true` 的 stage；是否持有由代码决定。 +- 子 op 的构造时机与次数、调度方式以及 forward 的执行，都由代码决定。 +- manifest 不规定复合 op 的 roofline 与其各 stage 的 roofline 之间的关系。 +- 对于 `implemented` 的 op，validator 按顺序核对：`op` stage 与类的 `delegate_types` 一致，`kernel` stage 与类的 `kernel_types` 一致。 +- `stages` 不能为空，stage 名互不相同，`optional` 为布尔值。 diff --git a/docs/user-guide/manifest/img/call.svg b/docs/user-guide/manifest/img/call.svg new file mode 100644 index 00000000..a2313b8e --- /dev/null +++ b/docs/user-guide/manifest/img/call.svg @@ -0,0 +1 @@ +The path of one op callGeneratedWritten by the developerGeneratedGeneratedOp baseOp classCallerGeneratedOp baseOp classKerneltargetCallerCallerGeneratedchecksGeneratedchecksOp baseOp baseOp classOp classKernelimplementationKernelimplementationtargettargetGeneratedGeneratedOp baseOp classConstructionOp(params)Check params bytype; checkinvariants; run thechecks constructioncan evaluateop instanceCallop(tensors, ...)1. Settle whichoptional    tensors are passed2. Check domainrestrictions3. Choosetype-familybranches4. Infer indices,    check refinements5. Check outputbuffersalt[Every written tensor is empty]Build the outputsfrom the signature;run nothing[A target registered a builder]Call with forward'sinputsoutputs[In-tree implementation]6. forward(tensors,...)kernel_for(interface,call spec)opt[First time this call spec is seen]Select animplementation:availability →applicability →precedenceentry_for(call)build identity andfactoryopt[First time this build identity is seen]Build the entryentryentry(tensors)resultoutputs7. Check theoutputs' count,    shapes, dtypes,device    and layoutoutputsAn op declaring a compile boundary: forward only calls_call_boundary, the same steps run inside the generatedtorch.library operator, and the in-tree implementation is_eager_forward; fake/meta come from the same signature diff --git a/docs/user-guide/manifest/img/call.zh.svg b/docs/user-guide/manifest/img/call.zh.svg new file mode 100644 index 00000000..021d519d --- /dev/null +++ b/docs/user-guide/manifest/img/call.zh.svg @@ -0,0 +1 @@ +一次 op 调用的执行路径系统生成开发者编写......Op ..Op .......Op ..Op .Kernel ..target调用方调用方生成的检查生成的检查Op 基类Op 基类Op 类Op 类Kernel 实现Kernel 实现targettarget......Op ..Op .构造Op(参数)按 type 检查参数；检查invariant；完成构造时可以求值的检查op 实例调用op(张量, ...)1. 确定可选张量是否传入2. 检查定义域限制3. 选择 type family 分支4. 推断 index，    检查 refinement5. 检查输出缓冲alt[所有写入的张量都为空]按签名构造输出，不执行实现[target 为 op 注册了 builder]以 forward 的输入调用输出[in-tree 实现]6. forward(张量, ...)kernel_for(interface, callspec)opt[call spec 首次出现]选择实现：可用性 → 适用性→ 优先级entry_for(call)build identity 与构建函数opt[build identity 首次出现]构造 entryentryentry(张量)结果输出7. 检查输出的数量、    形状、dtype、设备    与内存布局输出声明了 compile boundary 的 op：forward 只调用 _call_boundary，同样的步骤在生成的 torch.library operator 内执行，in-tree 实现写在 _eager_forward 中；fake/meta 由同一份签名生成 diff --git a/docs/user-guide/manifest/img/device.svg b/docs/user-guide/manifest/img/device.svg new file mode 100644 index 00000000..5ae8f62c --- /dev/null +++ b/docs/user-guide/manifest/img/device.svg @@ -0,0 +1 @@ +How the call device is chosen(workload instantiation places tensors by the same rule)A tensor declaring device: cpu stays on the CPUand takes no part belowAny call-time input?device parameter not None?Any construction-time tensor?The current CUDA device;None without CUDATheir device; they must share onedevice; a string is parsedby torch.deviceTheir device; several must share oneOnce the call device is chosen, out and the outputsare checked or allocated on it; construction-time tensorsare copied to it and cast to their signature dtypenononoyesyesyes diff --git a/docs/user-guide/manifest/img/device.zh.svg b/docs/user-guide/manifest/img/device.zh.svg new file mode 100644 index 00000000..4433e959 --- /dev/null +++ b/docs/user-guide/manifest/img/device.zh.svg @@ -0,0 +1 @@ +调用设备的确定方式（workload 实例化按同一规则放置张量）声明 device: cpu 的张量固定在 CPU，不参与下面的判定有调用期输入？参数 device 不为 None？有构造期张量？取当前 CUDA 设备；CUDA 不可用时为 None取这些输入的设备；它们须位于同一设备取 device；字符串按 torch.device 解析取构造期张量的设备；多个时须位于同一设备确定调用设备后，out 与输出在该设备上检查或分配；构造期张量复制到调用设备，并转为签名中的 dtype否否否是是是 diff --git a/docs/user-guide/manifest/img/new-op.svg b/docs/user-guide/manifest/img/new-op.svg new file mode 100644 index 00000000..4de1c8f9 --- /dev/null +++ b/docs/user-guide/manifest/img/new-op.svg @@ -0,0 +1 @@ +Adding an op: the developer's steps and the system's checksWrite the spec with status spec-onlyWrite the call spec, the kernel interfaceand its implementations: override entry_forwhere needed, and declare devices,supported_archs, applies, preferred_overor generalWrite the Op class (__init__, kernel_types,interfaces, forward, docstring), and export itfrom its package and tileops.<family>Write ref_program in workloads/Write correctness tests in tests/ops/:BUILTIN target by default, device fromrun_device(), marked in_tree_kernelsor cuda_only where neededWrite the benchmark function inbenchmarks/ops/: calls from manifest_calls,and the baselines to compareSet status to implementedFix the implementation, not the specThe validator checks the signaturestatically and instantiates everyworkload rowGenerate the call checks, shape inference,dtype checks and eval_roofline from thesignature, and for a compile-boundary classthe operator and fake/metaCI: the validator checks code againstthe spec, then manifest tests andcorrectness tests on a GPUCI checks againnoChecks pass?yesMerge into mainNightly runs the benchmarks by case idand checks every call is measuredThe roofline tool computes efficiency,and the docs site updatesDeveloperSystem diff --git a/docs/user-guide/manifest/img/new-op.zh.svg b/docs/user-guide/manifest/img/new-op.zh.svg new file mode 100644 index 00000000..496dfba3 --- /dev/null +++ b/docs/user-guide/manifest/img/new-op.zh.svg @@ -0,0 +1 @@ +新增一个 op：开发者的步骤与系统的检查写 spec，status 为 spec-only写 call spec、kernel 接口与实现；按需覆写 entry_for，并声明 devices、supported_archs、applies、preferred_over或 general写 Op 类：__init__、kernel_types、interfaces、forward 与 docstring，在实现包与 tileops.<family> 中导出在 workloads/ 写 ref_program在 tests/ops/ 写正确性测试：默认 BUILTIN target，设备取自run_device()，按需标记 in_tree_kernels或 cuda_only在 benchmarks/ops/ 写 benchmark 函数，用 manifest_calls 取得调用，选择对比基线status 改为 implemented修改实现，不修改 specvalidator 静态检查签名，实例化每条 workload 行由签名生成调用检查、形状推导、dtype 检查与 eval_roofline；compile boundary 类另有 operator 与 fake/metaCI：validator 核对代码与 spec，manifest 测试，GPU 上的正确性测试CI 重新检查否检查通过？是合入 mainnightly 按 case id 运行 benchmark，检查每个调用都被测到roofline 工具计算效率，文档站更新开发者系统 diff --git a/docs/user-guide/manifest/img/overview.svg b/docs/user-guide/manifest/img/overview.svg new file mode 100644 index 00000000..f5907793 --- /dev/null +++ b/docs/user-guide/manifest/img/overview.svg @@ -0,0 +1 @@ +TileOPs components and how they call each otherspecsrc/tileops/manifest/spec/Written by the developerProvided by the systemPublicationImplementationCode generationcall checks, shape inferencedtype checks, eval_rooflineOp class__init__, forward, docstringkernel_types, interfacesOp basekernel_for: selectionentry cache, target dispatchKernel implementationinherits a kernel interfaceapplies, entry_for as neededValidation and measurementReferenceref_program in workloads/validatorvalidate_manifest.pyWorkload instantiationmanifest_callsCorrectness teststests/ops/Benchmark functionbenchmarks/ops/Manifest testsmeta calls, target conformanceManifestBenchmarktiming, FLOPs and bytesPublicationnightlyruns every benchmarkkeeps history by case idRoofline toolSOL efficiencyand boundDocs sitereads specs, docstringsand benchmark resultssignature, rooflineworkload rows, all fieldsinstalls methodskernel_forbuilds and cachestests and benchmarks call the opchecks the interfacereferencetorch baselinecallscallsbaselinesdevice time diff --git a/docs/user-guide/manifest/img/overview.zh.svg b/docs/user-guide/manifest/img/overview.zh.svg new file mode 100644 index 00000000..fa2d3c8d --- /dev/null +++ b/docs/user-guide/manifest/img/overview.zh.svg @@ -0,0 +1 @@ +TileOPs 的组件与调用关系specsrc/tileops/manifest/spec/开发者编写系统提供，开发者不修改发布实现代码生成调用检查、形状推导dtype 检查、eval_rooflineOp 类__init__、forward、docstringkernel_types、interfacesOp 基类kernel_for：选择实现entry 缓存、target 派发Kernel 实现继承 kernel 接口applies、entry_for 按需声明验证与测量参考实现workloads/ 的 ref_programvalidatorvalidate_manifest.pyworkload 实例化manifest_calls正确性测试tests/ops/benchmark 函数benchmarks/ops/manifest 测试meta 调用、target conformanceManifestBenchmark计时、FLOPs 与字节数发布nightly运行全部 benchmark按 case id 记录历史roofline 工具SOL 效率与瓶颈判定文档站读取 spec、docstring与 benchmark 结果签名、rooflineworkload 行、全部字段安装生成的方法kernel_for构造并缓存测试与 benchmark 执行调用核对接口数值参考torch 基线调用调用对比基线device time diff --git a/docs/user-guide/manifest/index.md b/docs/user-guide/manifest/index.md new file mode 100644 index 00000000..b49fe5df --- /dev/null +++ b/docs/user-guide/manifest/index.md @@ -0,0 +1,152 @@ +# Reading and writing the manifest + +The manifest describes the external contract of every public op in TileOPs. It consists of the YAML files under [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec). Each op has one entry in the manifest, and that entry is called the op's **spec**. The op's implementation follows the spec, and the validator checks in CI that the implementation agrees with the spec. When the two disagree, the implementation is what changes. + +Once the developer has written the spec for an op, the system generates the following from it: + +- the construction and call checks, output shape inference, and dtype checks; +- the calls and inputs used by the manifest tests and the benchmarks; +- the FLOPs and byte count of each call; +- the content of the docs site that is taken from the spec, such as the support matrix. + +The developer writes: + +- the kernel; +- the op class, including `__init__`, `forward` and the docstring; +- the numerical reference and the correctness tests; +- the comparison baselines for the benchmark. + +The first three sections of this page describe, in order, the components of the system, how one call executes, and the full workflow for adding an op. + +## 1. System overview {#overview} + +The components of TileOPs fall into four layers: + +1. Declaration layer: the spec. +1. Implementation layer: the code generated from the spec, and the op and kernel written by the developer. +1. Verification and measurement layer: tests, benchmarks and the validator. +1. Publishing layer: nightly, the roofline tools and the docs site. + +In the figure below, purple components are written by the developer, cyan components are provided by the system, and light green components belong to the publishing layer. + +![Components of TileOPs and how they call each other](img/overview.svg) + +**Table 1** Parts written by the developer + +| No. | Part | Location | Contents | +| --- | --- | --- | --- | +| 1 | spec | `src/tileops/manifest/spec/.yaml` | signature, workload rows, roofline formulas | +| 2 | kernel interface and implementations | `src/tileops/kernels/`; the kernel interface and the call spec usually live in the family's `call_spec.py`, and a family with a single kernel file writes them in that file | The kernel interface fixes the type of the call spec and the parameters of `forward`; an implementation is a subclass of `Kernel` that also inherits the corresponding kernel interface. The default `entry_for(call)` constructs the implementation directly from the call spec; it is overridden only when that does not express the construction arguments correctly | +| 3 | Op class | `src/tileops/ops/`, exported by `tileops.` | an `__init__` that matches `params`; `kernel_types` (key to implementation class) and `interfaces` (call site to kernel interface); in the in-tree implementation (usually `_eager_forward`, see compile boundary below), constructing the call spec and obtaining the entry to call through `kernel_for`; a Google-style docstring, from which the docs site generates the API reference | +| 4 | reference implementation | `workloads/` | `ref_program` on the workload class named after the op, or on a parameterized workload class shared by the family, plus the input construction that the workload rows cannot determine | +| 5 | correctness tests | `tests/ops/` | numerical tolerances, and the shapes needed to cover every branch of the kernel | +| 6 | benchmark function | the benchmark file in `benchmarks/ops/` for the module the op belongs to; a new file is added only when no suitable one exists | a test function parameterized with `manifest_calls(Op)`, and the choice of comparison baselines, such as the torch reference implementation and kernels from other libraries | + +The following are written only when needed: + +- when an inline formula cannot express the cost, a roofline `func` in `tileops.perf.formulas`; +- when the op's FLOPs are a matrix-multiply contraction and the best implementation should use tensor cores, an override of `compute_roof()`. It states the compute unit the best implementation should use, independent of the unit the current kernel actually uses; +- when one kernel interface has several implementations, a declaration on each implementation of which calls it serves (`applies`) and which devices it runs on (`devices`, `supported_archs`); when the applicability of implementations overlaps, `preferred_over` declares which one takes precedence, and the fallback implementation declares `general = True`, see [Adding a kernel to an op](../dispatch/writing.md#rule); +- an external backend does not modify the op. Before an op instance is constructed, it calls `tileops.backend.register_implementation("", "", implementation class)` to add an implementation for a kernel interface. The first argument is the op name in the manifest, and `key` must be a key the op does not use yet. The new implementation enters only instances constructed afterwards, and it is checked against the kernel interface when the instance is constructed, see [How a backend joins TileOPs](../dispatch/backends.md#register); +- a composite op declares `delegate_types` and `kernel_types` on the class, matching the spec's `composition`; +- an op that supports `fullgraph=True` declares `compile_boundary = True`: `forward` only calls `_call_boundary`, the in-tree implementation is written in `_eager_forward`, and a cold-compile test is registered in `tests/compile_contract.py`. + +**Table 2** Parts provided by the system + +| No. | Component | Input | Work done for the developer | +| --- | --- | --- | --- | +| 1 | code generation | signature, `roofline` | construction checks, call checks wrapped around `forward`, `_infer_output_shapes`, `_validate_dtypes` and `eval_roofline()`; a class that declares a compile boundary also gets a `torch.library` operator and fake/meta functions | +| 2 | Op base class | `kernel_types`, `interfaces`, call spec | at construction, checking that every implementation conforms to its kernel interface; at call time, selecting the implementation and constructing and caching the entry; holding sub-ops (`delegate_for`), target dispatch and autotune | +| 3 | workload instantiation | workload rows | expanding each row into calls according to `dtype_cases`, and generating input tensors that fit the call, including the values of shapes, dtypes, parameters and metadata | +| 4 | `ManifestBenchmark` | calls and the op | timing, generating the case id, taking FLOPs and byte count from `eval_roofline()`, and recording results under the op name | +| 5 | manifest tests | all specs | running every call on meta tensors; checking target conformance; checking that the public API agrees with the manifest; checking that the roofline byte count agrees with what the signature derives | +| 6 | validator | all fields | static checks of the signature; for `implemented` ops, checking `__init__`, `forward` and `composition` | +| 7 | nightly | benchmarks | running every benchmark each night, checking that every call is measured, and recording history by case id | +| 8 | roofline tools | device time, GPU profile | computing SOL efficiency and identifying the bottleneck | +| 9 | docs site | spec, docstring, benchmark and roofline results | generating the API reference, the support matrix and the performance pages | + +The calls used by the manifest tests and the benchmarks, including shapes, dtypes and input tensors, are all generated from the workload rows; the benchmark function only decides which implementations to compare against. The correctness tests choose their own shapes that cover every branch of the kernel. When ordinary random tensors do not satisfy the value range of an input, the workload class overrides `gen_inputs()`. + +## 2. The execution path of one call {#call-path} + +An op call has two phases: construction and call. In the figure below, cyan participants are generated from the signature or provided by the Op base class, and purple participants are written by the developer. + +![The execution path of one op call](img/call.svg) + +- In the construction phase, the generated checks check the parameters against their `type`, and run the checks that can already be evaluated at construction. +- In the call phase, the generated checks first determine the inputs, select the branch and infer the indices, and only then call the implementation written by the developer. After the implementation returns, they check the number, shape, dtype, device and memory layout of the outputs, and verify that `out` and `alias` outputs are the corresponding tensor objects. The implementation therefore does not repeat checks the spec already declares. +- The implementation obtains the entry to call through `kernel_for(interface, call)`; the entry is usually a single kernel. The Op base class looks up each distinct call spec only once. When a call spec first appears, the Op base class selects exactly one of the kernel interface's implementations, whose `entry_for(call)` returns the build identity and the factory. Within one kernel interface, an entry with the same implementation class and the same build identity is constructed only once. The selection rules are in [Calls and validation § 2](calls.md#selection). +- For an op that declares a compile boundary, `forward` only calls the generated `_call_boundary`, and the in-tree implementation is written in `_eager_forward`. +- When a target has registered a builder for the op through `register_kernel_builder`, the kernel returned by the target is called after the generated checks, and the op's `forward` does not run. Implementations added through `register_implementation` belong to the in-tree path and are still selected by `forward` through `kernel_for`. +- If every tensor a call writes (each output and each written input) has no elements, neither the in-tree implementation nor the target runs: new outputs are created on the call device with the checked shapes and dtypes, and `out` and the written inputs are returned unchanged. When the inputs are empty but the outputs are not, the implementation runs as usual. + +The details of each check are in [Calls and validation § 1](calls.md#call). + +## 3. Workflow for adding an op {#new-op} + +The left side of the figure below shows the developer's steps; the right side shows the generation and checks the system runs after each step. + +![Workflow for adding an op](img/new-op.svg) + +1. The spec is committed with `status: spec-only`. At this point the validator runs only static checks and does not require the code to exist. +2. The developer writes the kernel and the op class, and exports the op from the package that holds the implementation (for example `src/tileops/ops/reduction/__init__.py`) and from the public module `src/tileops/.py`; the `__all__` of both must include it. As soon as the op class exists, the methods generated from the signature are added to it. +3. The developer writes the reference implementation, the correctness tests and the benchmark function. When the family or module already has the corresponding file, they are added to that file. Correctness tests use the `BUILTIN` target by default, and the test device comes from `workloads.device.run_device()`. Tests that depend on in-tree kernel state are marked `pytest.mark.in_tree_kernels` or pass `target=BUILTIN` explicitly. Tests that need CUDA regardless of the target are marked `pytest.mark.cuda_only`. Device availability is checked with `workloads.device.run_device_available()`, not `torch.cuda.is_available()`. Benchmarks also accept the `--tileops-target` and `--tileops-device` options and use `BUILTIN` by default; benchmarks time with CUDA events and CUPTI and run only on CUDA devices. See the design document [Testing](../../design/testing.md). +4. After `status` changes to `implemented`, the validator starts checking the code against the spec, and CI runs the manifest tests and the correctness tests on GPU. When a check fails, the implementation is changed, not the spec. +5. After the change is merged into main, nightly runs the benchmarks by case id, the roofline tools compute efficiency, and the docs site is updated accordingly. + +## 4. The spec and the implementation {#authority} + +The spec is the authority for the op's external contract, and the implementation follows the spec. + +- The spec is written from an authoritative reference, such as the PyTorch API the op follows semantically, not reverse-engineered from existing TileOPs code. +- `status` is `implemented` only when the implementation conforms to the spec. When an op that is already `implemented` is found to disagree with its spec, `status` goes back to `spec-only` and the implementation is changed, not the spec. +- The runtime checks are generated from the signature, and implementing an op does not require writing them by hand. When a generated check is wrong, the code generation or the validator is fixed; the op does not work around the check. +- Checks that depend on code are skipped only for `spec-only` ops. There is no per-op switch to turn off a check. + +## 5. Contents and scope of the spec {#scope} + +The spec describes the op's external contract. Its contents fall into five groups of fields: + +**Table 3** The five groups of fields in a spec + +| No. | Content | Fields | Described in | +| --- | --- | --- | --- | +| 1 | type signature | `signature` | [Spec fields](writing.md) | +| 2 | side effects | `mutated`, `write_only`, `buffer`, `alias` on tensors | [Extensions § 6](extensions.md#effects) | +| 3 | test cases | `workloads`, and `values`, `requires` on tensors | [Spec fields § 8](writing.md#workloads), [Extensions § 7](extensions.md#generators) | +| 4 | cost model | `roofline` | [Spec fields § 9](writing.md#roofline) | +| 5 | internal structure of a composite op | `composition` | [Extensions § 10](extensions.md#composition) | + +The last four groups all build on the signature: + +- side effects are annotated on tensors in the signature; +- workload rows are concrete values of the signature; +- cost formulas use the names in the signature. + +The following belong to the implementation, are decided by code, and do not appear in the spec: + +- source paths; +- the choice of kernels, and the order in which several kernels are called; +- accumulation dtype, workspace, tile sizes and autotune configuration. + +## 6. File layout {#layout} + +- The YAML files are in `src/tileops/manifest/spec/`. Each family has one file `.yaml`; a large family is split into several `_.yaml` files. +- Each file is a non-empty mapping from op name to spec, and the `family` of every spec in a file equals the family the file name denotes. +- At load time all files are merged into one manifest. A duplicated op name, or a file that does not follow the naming rules above, is an error. +- ADTs shared by several specs are defined in `spec/types.yaml`, see [Extensions § 3](extensions.md#adt). +- The key of a spec is the op's Python class name, and the validator requires `cls.__name__` to equal the key exactly. Keys end in `FwdOp` or `BwdOp`, and words that name a variant go before the direction suffix, for example `GroupedQueryAttentionPagedFwdOp`. + +## 7. Pages of this guide {#pages} + +**Table 4** Contents of each page in this guide + +| No. | Page | Contents | +| --- | --- | --- | +| 1 | [Concepts](concepts.md) | the concepts used to describe an op's type | +| 2 | [Spec fields](writing.md) | how to write the fields most specs need | +| 3 | [Extensions](extensions.md) | how to write optional inputs, shapes that vary with parameters, side effects, metadata tensors and similar cases | +| 4 | [Calls and validation](calls.md) | the signature-based checks at call time, the validator's checks, and forms that are rejected | +| 5 | [Examples](examples.md) | real specs from the manifest | + +Reference tables such as the full set of field values, the kind mapping and the expression language are in the design document [Manifest](../../design/manifest.md#reference-tables), and this guide does not repeat them. diff --git a/docs/user-guide/manifest/index.zh.md b/docs/user-guide/manifest/index.zh.md new file mode 100644 index 00000000..c2a80400 --- /dev/null +++ b/docs/user-guide/manifest/index.zh.md @@ -0,0 +1,152 @@ +# 读写 manifest + +manifest 描述 TileOPs 每个 public op 的外部契约，由 [`src/tileops/manifest/spec/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/manifest/spec) 下的 YAML 文件组成。每个 op 在 manifest 中占一个条目，这个条目称为该 op 的 **spec**。op 的实现以 spec 为依据，validator 在 CI 中检查实现与 spec 是否一致；两者不一致时，修改的是实现。 + +开发者为一个 op 写好 spec 之后，以下内容都由系统根据 spec 生成： + +- 构造与调用时的检查、输出形状推导与 dtype 检查； +- manifest 测试与 benchmark 所用的调用及其输入； +- 每次调用的 FLOPs 与字节数； +- 文档站中取自 spec 的内容，例如支持矩阵。 + +开发者需要编写的是： + +- kernel； +- op 类，包括 `__init__`、`forward` 与 docstring； +- 数值参考与正确性测试； +- benchmark 的对比基线。 + +本页前三节依次说明系统由哪些组件组成、一次调用如何执行，以及新增一个 op 的完整流程。 + +## 1. 系统概览 {#overview} + +TileOPs 的组件分为四层： + +1. 声明层：spec。 +1. 实现层：由 spec 生成的代码，以及开发者编写的 op 与 kernel。 +1. 验证与测量层：测试、benchmark 与 validator。 +1. 发布层：nightly、roofline 工具与文档站。 + +下图中紫色的组件由开发者编写，青色的组件由系统提供，浅绿色的组件属于发布层。 + +![TileOPs 的组件与调用关系](img/overview.svg) + +**表 1** 开发者编写的部分 + +| No. | 部分 | 位置 | 内容 | +| --- | --- | --- | --- | +| 1 | spec | `src/tileops/manifest/spec/.yaml` | 签名、workload 行、roofline 公式 | +| 2 | kernel 接口与实现 | `src/tileops/kernels/`；kernel 接口与 call spec 通常放在 family 的 `call_spec.py` 中，只有一个 kernel 文件的 family 写在该文件里 | kernel 接口规定 call spec 的类型与 `forward` 的参数；实现是 `Kernel` 的子类，并继承相应的 kernel 接口。默认的 `entry_for(call)` 直接用 call spec 构造实现，不能准确表示构造参数时再覆写它 | +| 3 | Op 类 | `src/tileops/ops/`，并由 `tileops.` 导出 | 与 `params` 一致的 `__init__`；`kernel_types`（键到实现类）与 `interfaces`（调用位置到 kernel 接口）；在 in-tree 实现中（通常是 `_eager_forward`，见下文 compile boundary）构造 call spec，并通过 `kernel_for` 取得要调用的 entry；Google 风格的 docstring，文档站的 API 参考由它生成 | +| 4 | 参考实现 | `workloads/` | 以该 op 命名的 workload 类，或 family 共用的参数化 workload 类上的 `ref_program`，以及 workload 行无法确定的输入构造 | +| 5 | 正确性测试 | `tests/ops/` | 数值容差，以及覆盖 kernel 各分支所需的形状 | +| 6 | benchmark 函数 | `benchmarks/ops/` 中该 op 所属模块的 benchmark 文件，没有合适的文件时再新增 | 以 `manifest_calls(Op)` 参数化测试函数，并选择对比基线，例如 torch 参考实现与其他库的 kernel | + +以下几项只在需要时编写： + +- 内联公式无法表达代价时，在 `tileops.perf.formulas` 中编写 roofline `func`； +- op 的 FLOPs 属于矩阵乘收缩、最优实现应使用 tensor core 时，覆写 `compute_roof()`。它表示最优实现应当使用的计算单元，与当前 kernel 实际使用的单元无关； +- 同一个 kernel 接口有多个实现时，在各实现上声明它服务哪些调用（`applies`）与能在哪些设备上运行（`devices`、`supported_archs`）；实现之间的适用范围有重叠时，用 `preferred_over` 声明哪个优先，兜底的实现声明 `general = True`，见[如何为 op 新增 kernel](../dispatch/writing.md#rule)； +- 外部 backend 不修改 op，而是在构造 op 实例之前调用 `tileops.backend.register_implementation("", "", 实现类)`，为某个 kernel 接口增加实现。第一个参数是 manifest 中的 op 名，`key` 必须是该 op 尚未使用的新键；新增的实现只进入之后构造的实例，并在实例构造时接受 kernel 接口检查，见 [backend 如何接入 TileOPs](../dispatch/backends.md#register)； +- 复合 op 在类上声明 `delegate_types` 与 `kernel_types`，与 spec 的 `composition` 对应； +- 支持 `fullgraph=True` 的 op 声明 `compile_boundary = True`：`forward` 只调用 `_call_boundary`，in-tree 实现写在 `_eager_forward` 中，并在 `tests/compile_contract.py` 中登记冷编译测试。 + +**表 2** 系统提供的部分 + +| No. | 组件 | 输入 | 为开发者完成的工作 | +| --- | --- | --- | --- | +| 1 | 代码生成 | 签名、`roofline` | 构造检查、包裹 `forward` 的调用检查、`_infer_output_shapes`、`_validate_dtypes` 与 `eval_roofline()`；声明了 compile boundary 的类还会得到 `torch.library` operator 与 fake/meta 函数 | +| 2 | Op 基类 | `kernel_types`、`interfaces`、call spec | 构造时检查每个实现是否符合其 kernel 接口；调用时选择实现，构造并缓存 entry；子 op 的持有（`delegate_for`）、target 派发与 autotune | +| 3 | workload 实例化 | workload 行 | 将每条行按 `dtype_cases` 展开为调用，并生成符合调用的输入张量，包括形状、dtype、参数与 metadata 的取值 | +| 4 | `ManifestBenchmark` | 调用与 op | 计时、生成 case id、从 `eval_roofline()` 取得 FLOPs 与字节数，并以 op 名记录结果 | +| 5 | manifest 测试 | 全部 spec | 在 meta 张量上执行每个调用；检查 target conformance；检查公开 API 与 manifest 一致；检查 roofline 字节数与签名的推导一致 | +| 6 | validator | 全部字段 | 静态检查签名；对 `implemented` 的 op 核对 `__init__`、`forward` 与 `composition` | +| 7 | nightly | benchmark | 每晚运行全部 benchmark，检查每个调用都被测到，并按 case id 记录历史数据 | +| 8 | roofline 工具 | device time、GPU profile | 计算 SOL 效率，判定瓶颈 | +| 9 | 文档站 | spec、docstring、benchmark 与 roofline 结果 | 生成 API 参考、支持矩阵与性能页面 | + +manifest 测试与 benchmark 所用的调用，包括形状、dtype 与输入张量，都由 workload 行生成，benchmark 函数只决定与哪些实现对比。正确性测试另选能覆盖 kernel 各分支的形状；普通随机张量不能满足输入的取值范围时，在 workload 类中覆写 `gen_inputs()`。 + +## 2. 一次调用的执行路径 {#call-path} + +一次 op 调用分为构造与调用两个阶段。下图中青色的参与者由签名生成或由 Op 基类提供，紫色的参与者由开发者编写。 + +![一次 op 调用的执行路径](img/call.svg) + +- 构造阶段，生成的检查按 `type` 检查参数，并完成构造时已经可以求值的检查。 +- 调用阶段，生成的检查先确定输入、选择分支并推断 index，然后才调用开发者编写的实现；实现返回之后，再检查输出的数量、形状、dtype、设备与内存布局，并核对 `out` 与 `alias` 输出是否就是对应的张量对象。因此实现中不需要重复 spec 已经声明的检查。 +- 实现通过 `kernel_for(interface, call)` 取得要调用的 entry，通常就是一个 kernel。同一个 call spec 再次出现时，Op 基类只做一次查找；首次出现时，Op 基类在该 kernel 接口的实现中选出唯一一个，由它的 `entry_for(call)` 返回 build identity 与构建函数，在同一个 kernel 接口内，同一实现类、相同 build identity 的 entry 只构造一次。选择规则见[调用与校验 2](calls.md#selection)。 +- 声明了 compile boundary 的 op，`forward` 只调用生成的 `_call_boundary`，in-tree 实现写在 `_eager_forward` 中。 +- 某个 target 通过 `register_kernel_builder` 为 op 注册了 builder 时，生成的检查之后调用 target 返回的 kernel，op 的 `forward` 不执行。通过 `register_implementation` 增加的实现属于 in-tree 路径，仍由 `forward` 经 `kernel_for` 选中。 +- 如果一次调用写入的所有张量（各输出与被写入的输入）都不含元素，in-tree 实现与 target 都不执行：新的输出按检查过的形状与 dtype 在调用设备上创建，`out` 与被写入的输入原样返回。输入为空而输出不为空时，照常执行实现。 + +各检查的细节见[调用与校验 1](calls.md#call)。 + +## 3. 新增一个 op 的流程 {#new-op} + +下图左侧是开发者的步骤，右侧是系统在各步骤之后执行的生成与检查。 + +![新增一个 op 的流程](img/new-op.svg) + +1. spec 以 `status: spec-only` 提交。此时 validator 只做静态检查，不要求代码存在。 +2. 开发者编写 kernel 与 op 类，并在实现所在的包（例如 `src/tileops/ops/reduction/__init__.py`）与公开模块 `src/tileops/.py` 中导出 op，两处的 `__all__` 都要包含它。op 类一旦存在，由签名生成的方法就会加入这个类。 +3. 开发者编写参考实现、正确性测试与 benchmark 函数。family 或模块已有对应文件时，在已有文件中添加。正确性测试默认使用 `BUILTIN` target，测试设备取自 `workloads.device.run_device()`；依赖 in-tree kernel 状态的测试标记 `pytest.mark.in_tree_kernels` 或显式传入 `target=BUILTIN`；无论 target 为何都需要 CUDA 的测试标记 `pytest.mark.cuda_only`；判断设备是否可用时调用 `workloads.device.run_device_available()`，而不是 `torch.cuda.is_available()`。benchmark 同样接受 `--tileops-target` 与 `--tileops-device` 选项，默认使用 `BUILTIN`；benchmark 用 CUDA events 与 CUPTI 计时，只在 CUDA 设备上运行。详见设计文档 [Testing](../../design/testing.md)。 +4. `status` 改为 `implemented` 后，validator 开始核对代码与 spec，CI 运行 manifest 测试与 GPU 上的正确性测试。检查失败时修改实现，不修改 spec。 +5. 合入 main 之后，nightly 按 case id 运行 benchmark，roofline 工具计算效率，文档站随之更新。 + +## 4. spec 与实现的关系 {#authority} + +spec 是 op 外部契约的依据，实现以 spec 为准。 + +- spec 依据权威参考编写，例如 op 在语义上所参照的 PyTorch API，而不是从 TileOPs 现有的代码反推。 +- 只有实现符合 spec 时，`status` 才是 `implemented`。已经 `implemented` 的 op 被发现与 spec 不符时，`status` 改回 `spec-only`，并修改实现，而不是修改 spec。 +- 运行时检查由签名生成，实现 op 时不需要手写这些检查。生成的检查有误时，应修复代码生成或 validator，而不是在 op 中绕开它们。 +- 依赖代码的检查只对 `spec-only` 的 op 跳过，没有按 op 关闭某项检查的开关。 + +## 5. spec 的内容与范围 {#scope} + +spec 描述 op 的外部契约，其内容分为五组字段： + +**表 3** spec 的五组字段 + +| No. | 内容 | 字段 | 说明所在 | +| --- | --- | --- | --- | +| 1 | 类型签名 | `signature` | [写一个 spec](writing.md) | +| 2 | 副作用 | 张量上的 `mutated`、`write_only`、`buffer`、`alias` | [扩展写法 6](extensions.md#effects) | +| 3 | 测试用例 | `workloads`，以及张量上的 `values`、`requires` | [写一个 spec 8](writing.md#workloads)、[扩展写法 7](extensions.md#generators) | +| 4 | 代价模型 | `roofline` | [写一个 spec 9](writing.md#roofline) | +| 5 | 复合 op 的内部结构 | `composition` | [扩展写法 10](extensions.md#composition) | + +后四组字段都以签名为基础： + +- 副作用标注在签名中的张量上； +- workload 行是签名的具体取值； +- 代价公式使用签名中的名字。 + +以下内容属于实现，由代码决定，不出现在 spec 中： + +- 源码路径； +- kernel 的选择，以及多个 kernel 的调用顺序； +- 累加 dtype、workspace、tile 大小与 autotune 配置。 + +## 6. 文件组织 {#layout} + +- YAML 文件位于 `src/tileops/manifest/spec/`。每个 family 对应一个文件 `.yaml`；规模较大的 family 拆分为若干个 `_.yaml`。 +- 每个文件是从 op 名到 spec 的非空映射，文件中每个 spec 的 `family` 都与文件名所表示的 family 相同。 +- 加载时所有文件合并为一份 manifest。同一个 op 名重复出现，或者文件不符合上述命名规则，都会报错。 +- 多个 spec 共用的 ADT 定义在 `spec/types.yaml` 中，见[扩展写法 3](extensions.md#adt)。 +- spec 的键是 op 的 Python 类名，validator 要求 `cls.__name__` 与键完全相同。键以 `FwdOp` 或 `BwdOp` 结尾，表示变体的词写在方向后缀之前，例如 `GroupedQueryAttentionPagedFwdOp`。 + +## 7. 本指南的内容 {#pages} + +**表 4** 本指南各页的内容 + +| No. | 页面 | 内容 | +| --- | --- | --- | +| 1 | [概念](concepts.md) | 描述 op 类型所用的概念 | +| 2 | [写一个 spec](writing.md) | 大多数 spec 所需字段的写法 | +| 3 | [扩展写法](extensions.md) | 可选输入、随参数变化的形状、副作用、metadata 张量等情况的写法 | +| 4 | [调用与校验](calls.md) | 调用时基于签名的检查、validator 的检查项，以及不被接受的写法 | +| 5 | [示例](examples.md) | manifest 中的真实 spec | + +字段的完整取值、kind 映射、表达式语言等参考表收录在设计文档 [Manifest](../../design/manifest.md#reference-tables) 中，本指南不再重复。 diff --git a/docs/user-guide/manifest/writing.md b/docs/user-guide/manifest/writing.md new file mode 100644 index 00000000..66c641c5 --- /dev/null +++ b/docs/user-guide/manifest/writing.md @@ -0,0 +1,220 @@ +# Spec fields + +Most specs use only the following fields, and this page describes how to write them: + +1. `forall`; +1. the `shape` and `dtype` of tensors; +1. construction parameters; +1. simple `shape_rules`; +1. workload rows; +1. inline roofline formulas. + +How to write optional inputs, shapes that vary with parameters, side effects, metadata tensors and similar cases is in [Extensions](extensions.md). + +## 1. Structure of a spec {#anatomy} + +The SiluAndMul spec uses only the fields above: + +```yaml +SiluAndMulFwdOp: + family: elementwise + status: implemented + signature: + forall: {M: Dim, N: Dim, T: "DType[float16 | bfloat16 | float32]"} + inputs: + x: {dtype: T, shape: "[M, 2 * N]"} + outputs: + output: {dtype: T, shape: "[M, N]"} + workloads: + - {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill} + - {M: 1, N: 14336, dtype_cases: [{T: bfloat16}], label: llama-8b-ffn-decode} + roofline: + flops: "6 * M * N" +``` + +The following lists every top-level field a spec can contain. Each comment gives the section that describes the field: + +```yaml +: + family: # 2 + status: implemented | spec-only # 2 + ref_api: # 2, optional + signature: + forall: {: } # 3 + params: {

: {type, default, kw_only}} # 4 + inputs: {: {dtype, shape, optional, mutated, values, requires, ...}} # 5 + outputs: {: {dtype, shape, nullable, buffer, alias, ...}} # 5 + types: {: {params, match, cases}} # Extensions 2 + let: {: ""} # Extensions 4 + shape_rules: [""] # 6 + dtype_combos: [{: }] # Extensions 8 + workloads: [{, some, dtype_cases, label}] # 8 + roofline: {flops, bytes} | {func} # 9 + composition: {kind: composite, stages} # Extensions 10 +``` + +## 2. Top-level fields {#top} + +`family`, `status`, `signature`, `workloads` and `roofline` are required, and `workloads` contains at least one row. + +- `family` is the name of the public module the op belongs to. The op can be imported by class name from `tileops.`, and the `__all__` of that module agrees with the manifest. The name of the YAML file that holds the spec is also determined by `family`, see [Reading and writing the manifest § 6](index.md#layout). +- `status` takes one of two values: + - `implemented` means an implementation that conforms to the spec exists; + - `spec-only` means no conforming implementation exists yet; the code may be absent or only partly done. Checks that depend on code are skipped only for `spec-only` ops. +- `status` decides only which code-dependent checks run, and does not affect the methods generated from the signature: any class that has a spec gets all the generated methods, including `eval_roofline()`. +- `ref_api` is optional and records the fully qualified name of the API the op follows semantically, for example `torch.matmul`. The validator checks its format, and when the corresponding module can be imported, it also checks that the name exists. + +## 3. forall and kind {#forall} + +`forall` declares every free type index in the signature and its kind, for example `forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"}`. + +**Table 1** Kinds available in `forall` + +| No. | kind | Values | How it is determined at call time | How it is written in a workload row | +| --- | --- | --- | --- | --- | +| 1 | `Dim` | a non-negative integer, an axis length | solved by unification of the inputs | an integer | +| 2 | `Shape` | a tuple of `Dim` | solved by unification of the inputs | a list of integers | +| 3 | `DType[a \| b]` | one of the listed dtypes | solved by unification of the inputs | `dtype_cases` | +| 4 | `Seq[Int]` | a list of integers, such as `q_lens` | exists only at instantiation, as an argument to a generator (see [Extensions § 7](extensions.md#generators)) | a list of integers, or a call to a value primitive that returns a list, such as `"repeat(512, 64)"` | + +- A `Dim` can be used where an `Int` is required, and a `Shape` can be used where a `Seq[Int]` is required. +- Every axis in a shape is an integer expression. An axis whose kind is not `Dim`, and every element of a sequence expanded with `*p`, must be non-negative; the generated checks confirm this at construction or at call time. +- Values in YAML are converted to Python values by the declared kind or `type`: `Shape` and `tuple[...]` become tuples, `Seq[Int]` and `list[...]` stay lists, dtype names become `torch.dtype`, and ADT values become the corresponding Python objects. +- Within one spec, the names of indices, `let` entries and tensors are all distinct. `out` is a reserved name and cannot be used as the name of a tensor, parameter or index. + +## 4. params {#params} + +`signature.params` corresponds one to one with the parameter list of `__init__`. An ordinary construction parameter declares `type`, and optionally `default` and `kw_only` (keyword-only). A tensor passed at construction declares `dtype` and `shape` instead, see [Extensions § 5](extensions.md#placement). + +- The spec is the authority. For an `implemented` op, the validator compares `params` with `__init__` item by item, and requires the same set of parameters, order, `default` and `kw_only`. +- The only extra parameters the code may have are execution-policy parameters: `kernel_map`, `tune` and `target`, which every op has; implementation objects injected by the caller; and the reserved parameter `config`, which is passed only to the kernel. +- The call-time inputs and output buffers in the signature form, in order, the leading part of the `forward` parameter list. `forward` may append code-defined execution parameters after them; those parameters are not part of the signature. + +Construction parameters can appear directly in types without extra annotation: + +- an `int` parameter can be written in a shape, for example `num_local_experts` in MoE; +- a `list[int]` or `tuple[int, ...]` parameter can be expanded in a shape as `*p`, for example `normalized_shape` in RMSNorm; +- a dtype parameter (whose `type` is a union of dtype names) can be used directly as a tensor's `dtype`, for example `out_dtype` in Alibi; +- how to write an `int | None` parameter is in [Extensions § 1](extensions.md#presence). + +The full mapping from `type` to kind is in the design document [Manifest Table 3](../../design/manifest.md#t-types). Parameters of types such as `float` or an unconstrained `str` do not take part in type inference, but they can appear in `shape_rules` and roofline formulas; the validator checks their `type`, `default`, and how they are used in those expressions. + +## 5. inputs and outputs {#tensors} + +Each tensor is written `{dtype: ..., shape: "..."}`, which corresponds to the type `Tensor[T, s]`. + +- A shape is written `"[" axis, ... "]"`, or as an application of a type family. Each axis can be an expression, `*S` or `*primitive(...)`; `"[]"` denotes a zero-dimensional tensor. +- Tensors with the same shape use the same shape term; for example, the input and output of an elementwise op are both written `[*S]`. +- Each spec has exactly one signature, so every call returns outputs with the same names and the same count. If the number of outputs of an op varies with some switch, or the same parameter can be either a scalar or a tensor, the op is split into several specs. +- An op selects its implementation from its parameters and from which tensors are passed. The contents of a tensor are only input to the computation and do not affect implementation selection. + +## 6. shape_rules {#refinement} + +Each entry in `shape_rules` is a refinement, a predicate that constrains the values of indices, checked after unification. + +```yaml +shape_rules: +- "not is_causal or S_q <= S_kv" +- "D > 0 and D % 2 == 0" +``` + +- Shapes, `let`, type families, refinements and inline roofline formulas use the same closed expression language, with the same operator precedence as Python. The parts of the expression language are in the design document [Manifest Table 10](../../design/manifest.md#t-lang). +- A refinement can depend only on quantities available at run time. A refinement that can be evaluated at construction is checked at construction; the others are checked on every call. +- Constraints on the contents of a metadata tensor are written in the tensor's `requires`, not in `shape_rules`, see [Extensions § 7](extensions.md#generators). +- A condition narrows kinds within the branch it selects: in the branch where `present(v)` is true, `Maybe[X]` narrows to `X`; in the branch where `x == 'a'` or `x in ('a', 'b')` holds, the values of `x` narrow to those literals. If a comparison uses a literal that is not a member of the corresponding enum or dtype set, the validator rejects the refinement. +- Whether a refinement can be satisfied is the responsibility of the spec's author. +- Forms such as `x.shape == (...)`, `x is None` and `isinstance` are rejected; the corresponding rewrites are in [Calls and validation § 7](calls.md#rejected). + +## 7. dtype {#dtype} + +A tensor's `dtype` is a dtype expression, which takes one of four forms: + +**Table 2** Forms of a dtype expression + +| No. | Form | Meaning | Example | +| --- | --- | --- | --- | +| 1 | a `DType` index in `forall` | takes a value in the declared set, solved by unification of the inputs | `a: {dtype: T}` | +| 2 | a dtype parameter | takes the value of the construction parameter | `output: {dtype: out_dtype}` | +| 3 | a constant | a fixed dtype | `cu_seqlens_q: {dtype: int32}` | +| 4 | a dtype primitive | computed from other dtypes | `promote_int_to_float(T)`, `coalesce_dtype(out_dtype, D)` | + +If the spec has no `dtype_combos`, each `DType` index takes values in its own set independently. How to write a spec in which several `DType` indices allow only specific combinations is in [Extensions § 8](extensions.md#dtype-combos). + +## 8. Workload rows {#workloads} + +Each workload row determines one call. The benchmarks, nightly and the manifest tests all generate their calls from workload rows. + +```yaml +workloads: +- {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill} +``` + +**Table 3** Keys of a workload row + +| No. | Key | Value | +| --- | --- | --- | +| 1 | a construction parameter name | the parameter's value; parameters without a default must be given | +| 2 | `some` | the optional tensors passed in this call, see [Extensions § 1](extensions.md#presence) | +| 3 | a `Dim`, `Shape` or `Seq[Int]` index in `forall` | the indices used on the current branch that are not solved by a generator | +| 4 | `dtype_cases` | a non-empty list, each item of which is one assignment of the `DType` indices used on the current branch, for example `[{T: float16}, {T: bfloat16}]`; written only when such indices exist; dtype parameters are given as construction parameters | +| 5 | `label` | the name of the row | + +- An index is used on a branch when it appears in one of the following places on that branch: a shape, a dtype, a refinement, a generator argument, `requires`, or an inline roofline formula. A workload row must give every index used on the current branch, and only those indices; `let` entries and indices solved by a generator are not given in the row. + - For this decision, each expression is first simplified for the current branch. If the condition of a refinement is always true on that branch, the indices in it do not count as used. + - When a `let` is used, the indices in its expression also count as used. + - A roofline in `func` form makes no index used. + - A discriminant that selects a type family branch, decides whether a tensor exists, or decides whether an output is `None` always counts as used. +- Each row expands into several cases according to `dtype_cases`. The case id joins the following three parts in order with `-`: + + 1. `label`; + 1. each dtype value in `dtype_cases`, in the declaration order of `forall`; + 1. the values of dtype parameters, in the declaration order of `params`. + + For example, the id of the first case in the example above is `llama-8b-ffn-prefill-float16`. +- Nightly history is keyed by case id, so changing a `label` breaks the history of that row. +- A `label` is non-empty, at most 24 characters long, and contains only characters in `[A-Za-z0-9._-]`. A `label` describes only the scenario the row models, that is, the model and its use or the synthetic purpose, plus the qualifiers needed to tell similar rows apart, for example `llama-8b-ffn-prefill`. The op name and dtype already appear in the case id and are not repeated. Rows that share a `label` may differ only in dtype. +- Case ids within one spec are all distinct. +- At instantiation, the workload row determines the shapes, dtypes, parameter values, which optional tensors are passed, and the values of metadata tensors. The rest is generated by fixed rules: the device is determined as described in [Calls and validation § 4](calls.md#device), strides are contiguous, tensors do not alias each other, and ordinary data is random. +- The validator re-infers the call from the instantiated inputs and requires the result to agree with the workload row. +- Workload rows are not responsible for unit-test coverage; the shapes needed to cover every branch of the kernel are chosen separately by each op's tests. + +## 9. roofline {#roofline} + +`roofline` gives the FLOPs and byte count of one call. It is written either as inline formulas or as a `func`: + +```yaml +roofline: + flops: "2 * M * N * K" # inline formula; when bytes is omitted it is derived from the signature +# or +roofline: + func: "tileops.perf.formulas.gqa_fwd_roofline" +``` + +- Inline formulas are written in the expression language and can refer to the signature's indices, construction parameters and `let` entries, as well as `present(t)`, `bytes(t)` (the byte count of tensor `t`) and built-in primitives. When the cost depends on whether a tensor is passed, the formula can distinguish the cases with `present(...)`. +- When `bytes` is omitted, the byte count is derived by reading or writing each tensor in full once: + - an input that is not written counts as one read, an output counts as one write, a `mutated` input counts as one read and one write, and a `write_only` input counts as one write only; + - each tensor name in the signature counts as separate storage, even if the caller passes the same tensor to two parameters; only aliases declared in the signature (`buffer`, `alias`) count as the same storage; + - the byte count of each tensor is `prod(shape) * bits(dtype) / 8`, and packed dtypes are counted by their carrier. +- The byte count is the algorithm's minimum memory traffic: intermediate results are not counted, and an input that the algorithm reads only in part counts the distinct elements actually read. When the two differ by less than 1% on every workload row, the whole tensor may be counted. +- If the derived byte count does not match this minimum traffic, the spec writes `bytes` explicitly, with a corresponding test. +- `flops` is the minimum amount of arithmetic the algorithm needs for the call, not the number of instructions the hardware executes: + - values already computed count as reused; + - recurrent ops such as linear attention and state-space scans count by the chunked algorithm at the chunk size when the signature has a chunk size, and by per-token recurrence otherwise; + - elementwise operations, attention and MoE are counted by the conventions in the design document [Roofline § 1.3](../../design/roofline.md#13-convention); + - a path that is an identity map on the call's dtype counts as 0. +- `func` is used when the formula needs Python logic. It points to a module-level function `f(call) -> tuple[int, int]` in `tileops.perf.formulas`. The function reads only its argument `call` and does not read the op instance. `call` is one checked call and provides the following interface: + + **Table 4** The argument `call` of a roofline `func` + + | No. | Interface | Contents | + | --- | --- | --- | + | 1 | `call.ix` | the parameters, the indices and dtype indices solved for this call, and the `let` entries used, that is, the names an inline formula can refer to | + | 2 | `call.present(t)` | whether tensor `t` is passed, held or returned; `call.present("out")` says whether the caller passed `out` | + | 3 | `call.tensors[t]` | the `(shape, dtype name)` of tensor `t` | + | 4 | `call.bytes(t)` | the byte count of tensor `t` | + | 5 | `call.values(t)` | the contents of metadata tensor `t`; calling it on meta tensors raises an error, because meta tensors have no values | + | 6 | `call.stages` | the calls completed by each sub-op of a composite op in this call, indexed by stage name | + +- Every spec generates an `eval_roofline()` method, which computes FLOPs and byte count from the op's most recently completed call. The benchmarks obtain the numbers through this method and write them into the results; the roofline tools read the benchmark results and do not call the op directly. + +The full roofline rules are in the design document [Roofline](../../design/roofline.md). diff --git a/docs/user-guide/manifest/writing.zh.md b/docs/user-guide/manifest/writing.zh.md new file mode 100644 index 00000000..2ce74a44 --- /dev/null +++ b/docs/user-guide/manifest/writing.zh.md @@ -0,0 +1,216 @@ +# 写一个 spec + +大多数 spec 只用到以下字段，本页说明它们的写法： + +1. `forall`； +1. 张量的 `shape` 与 `dtype`； +1. 构造参数； +1. 简单的 `shape_rules`； +1. workload 行； +1. 内联的 roofline 公式。 + +可选输入、随参数变化的形状、副作用、metadata 张量等情况的写法见[扩展写法](extensions.md)。 + +## 1. spec 的结构 {#anatomy} + +SiluAndMul 的 spec 只用到了上述字段： + +```yaml +SiluAndMulFwdOp: + family: elementwise + status: implemented + signature: + forall: {M: Dim, N: Dim, T: "DType[float16 | bfloat16 | float32]"} + inputs: + x: {dtype: T, shape: "[M, 2 * N]"} + outputs: + output: {dtype: T, shape: "[M, N]"} + workloads: + - {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill} + - {M: 1, N: 14336, dtype_cases: [{T: bfloat16}], label: llama-8b-ffn-decode} + roofline: + flops: "6 * M * N" +``` + +下面是 spec 可以包含的全部顶层字段，注释中标出了说明该字段的小节： + +```yaml +: + family: <模块名> # 2 + status: implemented | spec-only # 2 + ref_api: <完整限定名> # 2，可选 + signature: + forall: {: } # 3 + params: {

: {type, default, kw_only}} # 4 + inputs: {: {dtype, shape, optional, mutated, values, requires, ...}} # 5 + outputs: {: {dtype, shape, nullable, buffer, alias, ...}} # 5 + types: {: {params, match, cases}} # 扩展写法 2 + let: {: "<表达式>"} # 扩展写法 4 + shape_rules: [""] # 6 + dtype_combos: [{: }] # 扩展写法 8 + workloads: [{<参数与 index>, some, dtype_cases, label}] # 8 + roofline: {flops, bytes} | {func} # 9 + composition: {kind: composite, stages} # 扩展写法 10 +``` + +## 2. 顶层字段 {#top} + +`family`、`status`、`signature`、`workloads` 与 `roofline` 是必填字段，`workloads` 至少包含一条行。 + +- `family` 是 op 所在的公开模块名。op 可以通过 `tileops.` 按类名导入，该模块的 `__all__` 与 manifest 保持一致。spec 所在的 YAML 文件名也由 `family` 决定，见[读写 manifest 6](index.md#layout)。 +- `status` 的取值有两种： + - `implemented` 表示已有符合 spec 的实现； + - `spec-only` 表示尚无符合 spec 的实现，代码可能不存在，也可能只完成了一部分。依赖代码的检查只对 `spec-only` 的 op 跳过。 +- `status` 只决定哪些依赖代码的检查会运行，不影响由签名生成的方法：只要一个类有对应的 spec，它就会得到全部生成的方法，包括 `eval_roofline()`。 +- `ref_api` 为可选字段，记录 op 在语义上所参照的 API 的完整限定名，例如 `torch.matmul`。validator 检查它的格式；在对应模块可以导入时，还会检查该名字确实存在。 + +## 3. forall 与 kind {#forall} + +`forall` 声明签名中所有自由的 type index 及其 kind，例如 `forall: {M: Dim, N: Dim, K: Dim, T: "DType[float16 | bfloat16]"}`。 + +**表 1** `forall` 中可用的 kind + +| No. | kind | 取值 | 调用时如何确定 | workload 行中的写法 | +| --- | --- | --- | --- | --- | +| 1 | `Dim` | 非负整数，表示轴长 | 由输入的 unification 求得 | 整数 | +| 2 | `Shape` | 由 `Dim` 组成的元组 | 由输入的 unification 求得 | 整数列表 | +| 3 | `DType[a \| b]` | 所列 dtype 中的一个 | 由输入的 unification 求得 | `dtype_cases` | +| 4 | `Seq[Int]` | 整数列表，例如 `q_lens` | 只在实例化时存在，作为 generator 的实参（见[扩展写法 7](extensions.md#generators)） | 整数列表，或返回列表的值 primitive 调用，例如 `"repeat(512, 64)"` | + +- 要求 `Int` 的位置可以使用 `Dim`，要求 `Seq[Int]` 的位置可以使用 `Shape`。 +- 形状中的每个轴都是一个整数表达式。kind 不是 `Dim` 的轴，以及以 `*p` 展开的序列中的每个元素，都必须是非负数；生成的检查会在构造时或调用时确认这一点。 +- YAML 中的值按声明的 kind 或 `type` 转换为 Python 值：`Shape` 与 `tuple[...]` 转为 tuple，`Seq[Int]` 与 `list[...]` 保持为 list，dtype 名转为 `torch.dtype`，ADT 值转为对应的 Python 对象。 +- 在同一个 spec 中，index、`let` 与张量的名字互不相同。`out` 是保留名，不能用作张量、参数或 index 的名字。 + +## 4. params {#params} + +`signature.params` 与 `__init__` 的参数表一一对应。普通的构造参数声明 `type`，以及可选的 `default` 与 `kw_only`（表示只能以关键字方式传入）；构造时传入的张量改为声明 `dtype` 与 `shape`，见[扩展写法 5](extensions.md#placement)。 + +- 以 spec 为准。对于 `implemented` 的 op，validator 逐项比较 `params` 与 `__init__`，要求参数集合、顺序、`default` 与 `kw_only` 都相同。 +- 代码中可以额外出现的只有执行策略参数：所有 op 共有的 `kernel_map`、`tune`、`target`，由调用方注入的实现对象，以及只传给 kernel 的保留参数 `config`。 +- 签名中的调用期输入与输出缓冲，按顺序构成 `forward` 参数表的开头部分。`forward` 可以在其后追加由代码定义的执行参数，这些参数不属于签名。 + +构造参数可以直接出现在类型中，不需要额外标注： + +- `int` 参数可以写在形状中，例如 MoE 的 `num_local_experts`； +- `list[int]` 或 `tuple[int, ...]` 参数可以以 `*p` 的形式展开在形状中，例如 RMSNorm 的 `normalized_shape`； +- dtype 参数（`type` 为若干 dtype 名的并集）可以直接作为张量的 `dtype`，例如 Alibi 的 `out_dtype`； +- `int | None` 参数的写法见[扩展写法 1](extensions.md#presence)。 + +`type` 到 kind 的完整映射见设计文档 [Manifest 表 3](../../design/manifest.md#t-types)。`float`、取值不受限的 `str` 等类型的参数不参与类型推导，但可以出现在 `shape_rules` 与 roofline 公式中，validator 检查它们的 `type`、`default` 以及在这些表达式中的用法。 + +## 5. inputs 与 outputs {#tensors} + +每个张量写作 `{dtype: ..., shape: "..."}`，对应类型 `Tensor[T, s]`。 + +- 形状写作 `"[" 轴, ... "]"`，或者写成 type family 的应用。每个轴可以是表达式、`*S` 或 `*primitive(...)`；`"[]"` 表示零维张量。 +- 形状相同的张量使用同一个形状项，例如逐元素 op 的输入与输出都写作 `[*S]`。 +- 每个 spec 只有一份签名，因此每次调用返回的输出名字与数量都相同。如果一个 op 的输出数量随某个开关变化，或者同一个参数既可以是标量也可以是张量，这个 op 应拆分为多个 spec。 +- op 根据参数以及张量是否传入来选择实现；张量的内容只作为计算的输入，不影响实现的选择。 + +## 6. shape_rules {#refinement} + +`shape_rules` 中的每一条都是一个 refinement，即约束 index 取值的谓词，在 unification 之后检查。 + +```yaml +shape_rules: +- "not is_causal or S_q <= S_kv" +- "D > 0 and D % 2 == 0" +``` + +- 形状、`let`、type family、refinement 与内联 roofline 公式使用同一套封闭的表达式语言，运算优先级与 Python 相同。表达式语言的组成见设计文档 [Manifest 表 10](../../design/manifest.md#t-lang)。 +- refinement 只能依赖运行时可以得到的量。能在构造时求值的 refinement 在构造时检查，其余的在每次调用时检查。 +- 对 metadata 张量内容的约束写在张量的 `requires` 中，而非 `shape_rules` 中，见[扩展写法 7](extensions.md#generators)。 +- 条件判断会在其所选分支内收窄 kind：`present(v)` 为 true 的分支中，`Maybe[X]` 收窄为 `X`；`x == 'a'` 或 `x in ('a', 'b')` 成立的分支中，`x` 的取值收窄为这些字面量。如果一个比较所用的字面量不是对应枚举或 dtype 集合中的任何成员，validator 会拒绝这条 refinement。 +- refinement 是否可以满足，由 spec 的作者负责。 +- `x.shape == (...)`、`x is None`、`isinstance` 等写法不被接受，对应的改写方式见[调用与校验 7](calls.md#rejected)。 + +## 7. dtype {#dtype} + +张量的 `dtype` 是一个 dtype 表达式，有以下四种形式： + +**表 2** dtype 表达式的形式 + +| No. | 形式 | 含义 | 例子 | +| --- | --- | --- | --- | +| 1 | `forall` 中的 `DType` index | 在声明的集合内取值，由输入的 unification 求得 | `a: {dtype: T}` | +| 2 | dtype 参数 | 取构造参数的值 | `output: {dtype: out_dtype}` | +| 3 | 常量 | 固定的 dtype | `cu_seqlens_q: {dtype: int32}` | +| 4 | dtype primitive | 由其他 dtype 计算得到 | `promote_int_to_float(T)`、`coalesce_dtype(out_dtype, D)` | + +如果 spec 没有 `dtype_combos`，各个 `DType` index 在各自的集合内独立取值。多个 `DType` index 只允许特定组合时的写法见[扩展写法 8](extensions.md#dtype-combos)。 + +## 8. workload 行 {#workloads} + +每条 workload 行确定一次调用。benchmark、nightly 与 manifest 相关的测试都从 workload 行生成调用。 + +```yaml +workloads: +- {M: 2048, N: 14336, dtype_cases: [{T: float16}, {T: bfloat16}], label: llama-8b-ffn-prefill} +``` + +**表 3** workload 行的键 + +| No. | 键 | 值 | +| --- | --- | --- | +| 1 | 构造参数名 | 参数的值；没有默认值的参数必须给出 | +| 2 | `some` | 本次调用传入的可选张量，见[扩展写法 1](extensions.md#presence) | +| 3 | `forall` 中的 `Dim`、`Shape`、`Seq[Int]` index | 当前分支用到、且不由 generator 求得的 index | +| 4 | `dtype_cases` | 一个非空列表，每一项是当前分支用到的 `DType` index 的一组取值，例如 `[{T: float16}, {T: bfloat16}]`；只在有这类 index 时写；dtype 参数按构造参数给出 | +| 5 | `label` | 这条行的名字 | + +- 一个 index 在某个分支上被用到，是指它出现在该分支的以下位置之一：形状、dtype、refinement、generator 实参、`requires`、内联 roofline 公式。workload 行必须给出当前分支用到的每一个 index，且只给出这些 index；`let` 与由 generator 求得的 index 不在行中给出。 + - 判断时，各处表达式先按当前分支化简。如果一条 refinement 的条件在该分支上恒为 true，其中出现的 index 不算被用到。 + - 一个 `let` 被用到时，它的表达式中出现的 index 也算被用到。 + - `func` 形式的 roofline 不会使任何 index 被用到。 + - 选择 type family 分支、决定张量是否存在或决定输出是否为 `None` 的 discriminant，总是算作被用到。 +- 每条行按 `dtype_cases` 展开为若干个 case。case id 由以下三部分依次以 `-` 连接而成： + + 1. `label`； + 1. `dtype_cases` 中的各 dtype 值，按 `forall` 的声明顺序； + 1. dtype 参数的值，按 `params` 的声明顺序。 + + 例如上例第一个 case 的 id 是 `llama-8b-ffn-prefill-float16`。 +- nightly 的历史数据以 case id 为键，因此修改 `label` 会使这条行的历史记录中断。 +- `label` 不能为空，长度不超过 24 个字符，只能包含 `[A-Za-z0-9._-]` 中的字符。`label` 只描述这条行所建模的场景，即模型及其用途或合成目的，再加上区分同类行所需的限定词，例如 `llama-8b-ffn-prefill`；op 名与 dtype 已经出现在 case id 中，不必重复。共用同一个 `label` 的行只能在 dtype 上不同。 +- 同一个 spec 内的 case id 互不相同。 +- 实例化时，workload 行确定形状、dtype、参数取值、可选张量是否传入，以及 metadata 张量的取值。其余部分按固定规则生成：设备的确定方式见[调用与校验 4](calls.md#device)，stride 连续，张量之间没有别名，普通数据随机生成。 +- validator 会根据实例化后的输入重新推断这次调用，并要求结果与 workload 行一致。 +- workload 行不承担单元测试的覆盖职责，覆盖 kernel 各分支所需的形状由各 op 的测试另行选取。 + +## 9. roofline {#roofline} + +`roofline` 给出一次调用的 FLOPs 与字节数，有内联公式与 `func` 两种写法： + +```yaml +roofline: + flops: "2 * M * N * K" # 内联公式；省略 bytes 时由签名推导 +# 或 +roofline: + func: "tileops.perf.formulas.gqa_fwd_roofline" +``` + +- 内联公式使用表达式语言书写，可以引用签名中的 index、构造参数与 `let`，以及 `present(t)`、`bytes(t)`（张量 `t` 的字节数）和内建 primitive。如果代价取决于某个张量是否传入，公式可以使用 `present(...)` 区分。 +- 省略 `bytes` 时，字节数按每个张量完整读写一次推导： + - 未被写入的输入计一次读，输出计一次写，`mutated` 的输入读写各计一次，`write_only` 的输入只计一次写； + - 签名中的每个张量名按一块独立的存储计算，即使调用方把同一个张量传给两个参数；只有签名中声明的别名（`buffer`、`alias`）视为同一块存储； + - 每个张量的字节数为 `prod(shape) * bits(dtype) / 8`，打包 dtype 按载体计算。 +- 字节数表示算法的最少访存量：中间结果不计入，算法只读取一部分的输入按实际读取的不同元素计数；两者在每条 workload 行上相差不到 1% 时，可以按整个张量计。 +- 如果推导出的字节数与上述最少访存量不符，spec 需要显式写出 `bytes`，并附带相应的测试。 +- `flops` 是算法完成本次调用所需的最少算术量，不是硬件实际执行的指令数：已经算出的值按复用计；linear attention、状态空间扫描等循环 op，签名中带 chunk size 时按该 chunk size 的分块算法计，否则按逐 token 递推计；逐元素运算、attention 与 MoE 按设计文档 [Roofline 1.3](../../design/roofline.md#13-convention) 的约定计数；在本次调用的 dtype 上等同于恒等映射的路径计 0。 +- 公式需要 Python 逻辑时使用 `func`。它指向 `tileops.perf.formulas` 中的一个模块级函数 `f(call) -> tuple[int, int]`。函数只读取参数 `call`，不读取 op 实例。`call` 是一次经过检查的调用，提供以下接口： + + **表 4** roofline `func` 的参数 `call` + + | No. | 接口 | 内容 | + | --- | --- | --- | + | 1 | `call.ix` | 参数、本次调用求得的 index 与 dtype index、用到的 `let`，即内联公式可以引用的名字 | + | 2 | `call.present(t)` | 张量 `t` 是否传入、持有或返回；`call.present("out")` 表示调用方是否传入了 `out` | + | 3 | `call.tensors[t]` | 张量 `t` 的 `(形状, dtype 名)` | + | 4 | `call.bytes(t)` | 张量 `t` 的字节数 | + | 5 | `call.values(t)` | metadata 张量 `t` 的内容；在 meta 张量上调用时报错，因为 meta 张量没有取值 | + | 6 | `call.stages` | 复合 op 的各子 op 在本次调用中完成的调用，按 stage 名索引 | + +- 每个 spec 都会生成一个 `eval_roofline()` 方法，它基于 op 最近一次完成的调用计算 FLOPs 与字节数。benchmark 通过这个方法取得数值并写入结果，roofline 工具读取 benchmark 的结果，不直接调用 op。 + +roofline 的完整规则见设计文档 [Roofline](../../design/roofline.md)。 diff --git a/figures/user-guide/dispatch/call.en.puml b/figures/user-guide/dispatch/call.en.puml new file mode 100644 index 00000000..372b02a2 --- /dev/null +++ b/figures/user-guide/dispatch/call.en.puml @@ -0,0 +1,65 @@ +@startuml call +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam sequenceMessageAlign left +skinparam maxMessageSize 125 +skinparam BoxPadding 6 +skinparam sequence { + LifeLineBorderColor #CFC6E6 + ParticipantBorderColor #8E6CCF + DividerBackgroundColor #F8F5FD + DividerBorderColor #CFC6E6 + GroupBorderColor #8E6CCF + GroupHeaderFontColor #4A2C8F +} +hide stereotype +skinparam participant<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam participant<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +box "Written by the developer" #F8F5FD + participant "Op class\n_eager_forward" as op <> +end box +box "Provided by the system" #F3FBFC + participant "Op base\nkernel_for" as base <> + participant "CallSpec\ndevice facts" as facts <> +end box +box "Written by the developer" #F8F5FD + participant "Implementations of\nthe kernel interface" as impl <> + participant "entry" as entry <> +end box + +op -> op : Build the call spec from the tensors and the op's params +op -> base : kernel_for(interface, call) +activate base +base -> base : Look up (interface, call) in the dispatch cache +alt Hit + base --> op : the resolved entry +else Miss + base -> base : Check the call's type and hashability + loop Each implementation of the interface + base -> impl : unavailable(call), then refusal(call) if available + impl -> facts : Read arch, sm_count, calibration, smem_budget + facts -> facts : Resolved from call.device on first read + impl --> base : a reason, or None + end + base -> base : Take the one winner by precedence among the available, applicable ones; for a key replaced through kernel_map=, check the replacement serves the call + base -> impl : entry_for(call) + impl --> base : (build identity, factory) + opt This (implementation class, build identity)\nis not built yet + base -> entry : Run the factory with call.device current + end + base -> base : Request tuning on the entry when the op tunes; store in the dispatch cache + base --> op : entry +end +deactivate base +op -> entry : entry(tensors) +entry --> op : result +@enduml diff --git a/figures/user-guide/dispatch/call.zh.puml b/figures/user-guide/dispatch/call.zh.puml new file mode 100644 index 00000000..503659f9 --- /dev/null +++ b/figures/user-guide/dispatch/call.zh.puml @@ -0,0 +1,65 @@ +@startuml call +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam sequenceMessageAlign left +skinparam maxMessageSize 125 +skinparam BoxPadding 6 +skinparam sequence { + LifeLineBorderColor #CFC6E6 + ParticipantBorderColor #8E6CCF + DividerBackgroundColor #F8F5FD + DividerBorderColor #CFC6E6 + GroupBorderColor #8E6CCF + GroupHeaderFontColor #4A2C8F +} +hide stereotype +skinparam participant<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam participant<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +box "开发者编写" #F8F5FD + participant "Op 类\n_eager_forward" as op <> +end box +box "系统提供" #F3FBFC + participant "Op 基类\nkernel_for" as base <> + participant "CallSpec\n设备事实" as facts <> +end box +box "开发者编写" #F8F5FD + participant "kernel 接口的\n实现" as impl <> + participant "entry" as entry <> +end box + +op -> op : 由张量与 op 的参数构造 call spec +op -> base : kernel_for(interface, call) +activate base +base -> base : 以 (interface, call)查派发缓存 +alt 命中 + base --> op : 已解析的 entry +else 未命中 + base -> base : 检查 call 的类型与可哈希性 + loop 接口的每个实现 + base -> impl : unavailable(call)，可用时再问 refusal(call) + impl -> facts : 读取 arch、sm_count、calibration、smem_budget + facts -> facts : 首次读取时从 call.device 解析 + impl --> base : 理由，或 None + end + base -> base : 在可用且适用的实现中按优先关系取唯一胜者；key 被 kernel_map= 替换时，检查替换者能服务这次调用 + base -> impl : entry_for(call) + impl --> base : (build identity, 构建函数) + opt 这个 (实现类, build identity)\n尚未构建 + base -> entry : 以 call.device 为当前设备运行构建函数 + end + base -> base : op 处于调参模式时对 entry 请求调参；写入派发缓存 + base --> op : entry +end +deactivate base +op -> entry : entry(张量) +entry --> op : 结果 +@enduml diff --git a/figures/user-guide/dispatch/extension.en.puml b/figures/user-guide/dispatch/extension.en.puml new file mode 100644 index 00000000..ed88526d --- /dev/null +++ b/figures/user-guide/dispatch/extension.en.puml @@ -0,0 +1,54 @@ +@startuml extension +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam rectangle { + RoundCorner 8 +} +skinparam package { + BorderColor #CFC6E6 + BackgroundColor #FBFAFE + FontColor #4A2C8F +} +hide stereotype +skinparam rectangle<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam rectangle<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +skinparam rectangle<> { + BackgroundColor #E8F6EA + BorderColor #6DBA7A +} +rectangle "Generated call checks" as chk <> + +package "target: replaces the whole op" { + rectangle "build_kernel registered by\nregister_kernel_builder" as tgt <> +} + +package "The op instance's in-tree implementation" { + rectangle "The Op class's _eager_forward" as body <> + package "Implementations of one kernel interface\n(installed at construction)" as snap { + rectangle "Classes in kernel_types\ninheriting the interface" as intree <> + rectangle "kernel_map=\nreplaces the class behind a key,\nkeeping the key's rule\n(this instance only)" as repl <> + rectangle "register_implementation\nadds a key\n(every instance built afterwards)" as added <> + intree -[hidden]down- repl + repl -[hidden]down- added + } +} + +chk --> tgt : a target registered\na builder for the op +chk --> body : served in-tree +body --> snap : kernel_for selects among\nthese by the selection rule + +legend bottom + Cyan: provided by the system + Purple: written by TileOPs developers + Green: provided by a backend +endlegend +@enduml diff --git a/figures/user-guide/dispatch/extension.zh.puml b/figures/user-guide/dispatch/extension.zh.puml new file mode 100644 index 00000000..43f98c26 --- /dev/null +++ b/figures/user-guide/dispatch/extension.zh.puml @@ -0,0 +1,54 @@ +@startuml extension +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam rectangle { + RoundCorner 8 +} +skinparam package { + BorderColor #CFC6E6 + BackgroundColor #FBFAFE + FontColor #4A2C8F +} +hide stereotype +skinparam rectangle<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam rectangle<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +skinparam rectangle<> { + BackgroundColor #E8F6EA + BorderColor #6DBA7A +} +rectangle "生成的调用检查" as chk <> + +package "target：替换整个 op" { + rectangle "register_kernel_builder\n注册的 build_kernel" as tgt <> +} + +package "op 实例的 in-tree 实现" { + rectangle "Op 类的 _eager_forward" as body <> + package "一个 kernel 接口的实现（构造时安装）" as snap { + rectangle "kernel_types 中\n继承该接口的类" as intree <> + rectangle "kernel_map=\n替换 key 背后运行的类\n沿用该 key 的规则\n（只作用于这个实例）" as repl <> + rectangle "register_implementation\n新增的 key\n（此后构造的所有实例）" as added <> + intree -[hidden]down- repl + repl -[hidden]down- added + } +} + +chk --> tgt : target 为该 op\n注册了 builder +chk --> body : 由 in-tree 实现服务 +body --> snap : kernel_for 按选择规则\n在这些实现中选择 + +legend bottom + 青色：系统提供 + 紫色：TileOPs 开发者编写 + 绿色：backend 提供 +endlegend +@enduml diff --git a/figures/user-guide/manifest/call.en.puml b/figures/user-guide/manifest/call.en.puml new file mode 100644 index 00000000..23d20d18 --- /dev/null +++ b/figures/user-guide/manifest/call.en.puml @@ -0,0 +1,92 @@ +@startuml call +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam sequenceMessageAlign left +skinparam maxMessageSize 125 +skinparam BoxPadding 6 +skinparam sequence { + LifeLineBorderColor #CFC6E6 + ParticipantBorderColor #8E6CCF + ActorBorderColor #6B5FA0 + ActorBackgroundColor #EFE9FA + DividerBackgroundColor #F8F5FD + DividerBorderColor #CFC6E6 + GroupBorderColor #8E6CCF + GroupHeaderFontColor #4A2C8F +} +hide stereotype +skinparam participant<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam participant<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +skinparam participant<> { + BackgroundColor #F4F4F6 + BorderColor #9A98A8 +} +title The path of one op call + +actor "Caller" as caller +box "Generated" #F3FBFC + participant "Generated\nchecks" as chk <> + participant "Op base" as base <> +end box +box "Written by the developer" #F8F5FD + participant "Op class" as op <> + participant "Kernel\nimplementation" as kern <> +end box +participant "target" as tgt <> + +== Construction == +caller -> chk : Op(params) +activate chk +chk -> chk : Check params by type; check invariants; run the checks construction can evaluate +chk --> caller : op instance +deactivate chk + +== Call == +caller -> chk : op(tensors, ...) +activate chk +chk -> chk : 1. Settle which optional\n tensors are passed\n2. Check domain restrictions\n3. Choose type-family branches\n4. Infer indices,\n check refinements\n5. Check output buffers +alt Every written tensor is empty + chk -> chk : Build the outputs from the signature; run nothing +else A target registered a builder + chk -> tgt : Call with forward's inputs + tgt --> chk : outputs +else In-tree implementation + chk -> op : 6. forward(tensors, ...) + activate op + op -> base : kernel_for(interface, call spec) + activate base + opt First time this call spec is seen + base -> base : Select an implementation: availability → applicability → precedence + base -> kern : entry_for(call) + kern --> base : build identity and factory + opt First time this build identity is seen + base -> kern : Build the entry + end + end + base --> op : entry + deactivate base + op -> kern : entry(tensors) + kern --> op : result + op --> chk : outputs + deactivate op +end +chk -> chk : 7. Check the outputs' count,\n shapes, dtypes, device\n and layout +chk --> caller : outputs +deactivate chk + +note over chk, kern #FBFAFE + An op declaring a compile boundary: forward only calls + _call_boundary, the same steps run inside the generated + torch.library operator, and the in-tree implementation is + _eager_forward; fake/meta come from the same signature +end note +@enduml diff --git a/figures/user-guide/manifest/call.zh.puml b/figures/user-guide/manifest/call.zh.puml new file mode 100644 index 00000000..35014e25 --- /dev/null +++ b/figures/user-guide/manifest/call.zh.puml @@ -0,0 +1,91 @@ +@startuml call +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam sequenceMessageAlign left +skinparam maxMessageSize 150 +skinparam BoxPadding 6 +skinparam sequence { + LifeLineBorderColor #CFC6E6 + ParticipantBorderColor #8E6CCF + ActorBorderColor #6B5FA0 + ActorBackgroundColor #EFE9FA + DividerBackgroundColor #F8F5FD + DividerBorderColor #CFC6E6 + GroupBorderColor #8E6CCF + GroupHeaderFontColor #4A2C8F +} +hide stereotype +skinparam participant<> { + BackgroundColor #E2F6F8 + BorderColor #3AA9B8 +} +skinparam participant<> { + BackgroundColor #EFE9FA + BorderColor #8E6CCF +} +skinparam participant<> { + BackgroundColor #F4F4F6 + BorderColor #9A98A8 +} +title 一次 op 调用的执行路径 + +actor "调用方" as caller +box "系统生成" #F3FBFC + participant "生成的\n检查" as chk <> + participant "Op 基类" as base <> +end box +box "开发者编写" #F8F5FD + participant "Op 类" as op <> + participant "Kernel 实现" as kern <> +end box +participant "target" as tgt <> + +== 构造 == +caller -> chk : Op(参数) +activate chk +chk -> chk : 按 type 检查参数；检查 invariant；完成构造时可以求值的检查 +chk --> caller : op 实例 +deactivate chk + +== 调用 == +caller -> chk : op(张量, ...) +activate chk +chk -> chk : 1. 确定可选张量是否传入\n2. 检查定义域限制\n3. 选择 type family 分支\n4. 推断 index，\n 检查 refinement\n5. 检查输出缓冲 +alt 所有写入的张量都为空 + chk -> chk : 按签名构造输出，不执行实现 +else target 为 op 注册了 builder + chk -> tgt : 以 forward 的输入调用 + tgt --> chk : 输出 +else in-tree 实现 + chk -> op : 6. forward(张量, ...) + activate op + op -> base : kernel_for(interface, call spec) + activate base + opt call spec 首次出现 + base -> base : 选择实现：可用性 → 适用性→ 优先级 + base -> kern : entry_for(call) + kern --> base : build identity 与构建函数 + opt build identity 首次出现 + base -> kern : 构造 entry + end + end + base --> op : entry + deactivate base + op -> kern : entry(张量) + kern --> op : 结果 + op --> chk : 输出 + deactivate op +end +chk -> chk : 7. 检查输出的数量、\n 形状、dtype、设备\n 与内存布局 +chk --> caller : 输出 +deactivate chk + +note over chk, kern #FBFAFE + 声明了 compile boundary 的 op：forward 只调用 _call_boundary， + 同样的步骤在生成的 torch.library operator 内执行， + in-tree 实现写在 _eager_forward 中；fake/meta 由同一份签名生成 +end note +@enduml diff --git a/figures/user-guide/manifest/device.en.puml b/figures/user-guide/manifest/device.en.puml new file mode 100644 index 00000000..404209d2 --- /dev/null +++ b/figures/user-guide/manifest/device.en.puml @@ -0,0 +1,56 @@ +@startuml device +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam linetype ortho +skinparam nodesep 60 +skinparam ranksep 28 +skinparam rectangle { + BorderColor #3AA9B8 + BackgroundColor #E2F6F8 + RoundCorner 12 +} +skinparam card { + BorderColor #8E6CCF + BackgroundColor #F8F5FD +} + +title How the call device is chosen\n(workload instantiation places tensors by the same rule) + +card "A tensor declaring device: cpu stays on the CPU\nand takes no part below" as pre +card "Any call-time input?" as c1 +card "device parameter not None?" as c2 +card "Any construction-time tensor?" as c3 +card "The current CUDA device;\nNone without CUDA" as r5 #E2F6F8;line:3AA9B8 + +rectangle "Their device; they must share one" as r1 +rectangle "device; a string is parsed\nby torch.device" as r2 +rectangle "Their device; several must share one" as r3 + +pre -down-> c1 +c1 -down-> c2 : no +c2 -down-> c3 : no +c3 -down-> r5 : no +c1 -right-> r1 : yes +c2 -right-> r2 : yes +c3 -right-> r3 : yes + +r1 -[hidden]down- r2 +r2 -[hidden]down- r3 + +note bottom of r5 #FBFAFE + Once the call device is chosen, out and the outputs + are checked or allocated on it; construction-time tensors + are copied to it and cast to their signature dtype +end note +@enduml diff --git a/figures/user-guide/manifest/device.zh.puml b/figures/user-guide/manifest/device.zh.puml new file mode 100644 index 00000000..c4bef693 --- /dev/null +++ b/figures/user-guide/manifest/device.zh.puml @@ -0,0 +1,55 @@ +@startuml device +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam linetype ortho +skinparam nodesep 60 +skinparam ranksep 28 +skinparam rectangle { + BorderColor #3AA9B8 + BackgroundColor #E2F6F8 + RoundCorner 12 +} +skinparam card { + BorderColor #8E6CCF + BackgroundColor #F8F5FD +} + +title 调用设备的确定方式\n（workload 实例化按同一规则放置张量） + +card "声明 device: cpu 的张量固定在 CPU，不参与下面的判定" as pre +card "有调用期输入？" as c1 +card "参数 device 不为 None？" as c2 +card "有构造期张量？" as c3 +card "取当前 CUDA 设备；\nCUDA 不可用时为 None" as r5 #E2F6F8;line:3AA9B8 + +rectangle "取这些输入的设备；它们须位于同一设备" as r1 +rectangle "取 device；字符串按 torch.device 解析" as r2 +rectangle "取构造期张量的设备；多个时须位于同一设备" as r3 + +pre -down-> c1 +c1 -down-> c2 : 否 +c2 -down-> c3 : 否 +c3 -down-> r5 : 否 +c1 -right-> r1 : 是 +c2 -right-> r2 : 是 +c3 -right-> r3 : 是 + +r1 -[hidden]down- r2 +r2 -[hidden]down- r3 + +note bottom of r5 #FBFAFE + 确定调用设备后，out 与输出在该设备上检查或分配； + 构造期张量复制到调用设备，并转为签名中的 dtype +end note +@enduml diff --git a/figures/user-guide/manifest/new-op.en.puml b/figures/user-guide/manifest/new-op.en.puml new file mode 100644 index 00000000..975e96b9 --- /dev/null +++ b/figures/user-guide/manifest/new-op.en.puml @@ -0,0 +1,68 @@ +@startuml new-op +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam activity { + BorderColor #8E6CCF + BackgroundColor #FFFFFF + DiamondBorderColor #3AA9B8 + DiamondBackgroundColor #E2F6F8 + StartColor #6B5FA0 + EndColor #6B5FA0 +} +skinparam swimlane { + BorderColor #CFC6E6 + TitleFontStyle bold +} +title Adding an op: the developer's steps and the system's checks + +|#F8F5FD|Developer| +start +:Write the spec with status spec-only; +|#F3FBFC|System| +:The validator checks the signature +statically and instantiates every +workload row; +|Developer| +:Write the call spec, the kernel interface +and its implementations: override entry_for +where needed, and declare devices, +supported_archs, applies, preferred_over +or general; +:Write the Op class (~__init~__, kernel_types, +interfaces, forward, docstring), and export it +from its package and tileops.; +|System| +:Generate the call checks, shape inference, +dtype checks and eval_roofline from the +signature, and for a compile-boundary class +the operator and fake/meta; +|Developer| +:Write ref_program in workloads/; +:Write correctness tests in tests/ops/: +BUILTIN target by default, device from +run_device(), marked in_tree_kernels +or cuda_only where needed; +:Write the benchmark function in +benchmarks/ops/: calls from manifest_calls, +and the baselines to compare; +:Set status to implemented; +|System| +:CI: the validator checks code against +the spec, then manifest tests and +correctness tests on a GPU; +while (Checks pass?) is (no) + |Developer| + :Fix the implementation, not the spec; + |System| + :CI checks again; +endwhile (yes) +:Merge into main; +:Nightly runs the benchmarks by case id +and checks every call is measured; +:The roofline tool computes efficiency, +and the docs site updates; +stop +@enduml diff --git a/figures/user-guide/manifest/new-op.zh.puml b/figures/user-guide/manifest/new-op.zh.puml new file mode 100644 index 00000000..14ac51da --- /dev/null +++ b/figures/user-guide/manifest/new-op.zh.puml @@ -0,0 +1,62 @@ +@startuml new-op +skinparam defaultFontName "Noto Sans SC" +skinparam defaultFontSize 13 +skinparam shadowing false +skinparam ArrowColor #6B5FA0 +skinparam TitleFontColor #4A2C8F +skinparam activity { + BorderColor #8E6CCF + BackgroundColor #FFFFFF + DiamondBorderColor #3AA9B8 + DiamondBackgroundColor #E2F6F8 + StartColor #6B5FA0 + EndColor #6B5FA0 +} +skinparam swimlane { + BorderColor #CFC6E6 + TitleFontStyle bold +} +title 新增一个 op：开发者的步骤与系统的检查 + +|#F8F5FD|开发者| +start +:写 spec，status 为 spec-only; +|#F3FBFC|系统| +:validator 静态检查签名， +实例化每条 workload 行; +|开发者| +:写 call spec、kernel 接口与实现； +按需覆写 entry_for，并声明 devices、 +supported_archs、applies、preferred_over +或 general; +:写 Op 类：~__init~__、kernel_types、 +interfaces、forward 与 docstring， +在实现包与 tileops. 中导出; +|系统| +:由签名生成调用检查、形状推导、 +dtype 检查与 eval_roofline； +compile boundary 类另有 operator 与 fake/meta; +|开发者| +:在 workloads/ 写 ref_program; +:在 tests/ops/ 写正确性测试： +默认 BUILTIN target，设备取自 +run_device()，按需标记 in_tree_kernels +或 cuda_only; +:在 benchmarks/ops/ 写 benchmark 函数， +用 manifest_calls 取得调用，选择对比基线; +:status 改为 implemented; +|系统| +:CI：validator 核对代码与 spec， +manifest 测试，GPU 上的正确性测试; +while (检查通过？) is (否) + |开发者| + :修改实现，不修改 spec; + |系统| + :CI 重新检查; +endwhile (是) +:合入 main; +:nightly 按 case id 运行 benchmark， +检查每个调用都被测到; +:roofline 工具计算效率，文档站更新; +stop +@enduml diff --git a/figures/user-guide/manifest/overview.py b/figures/user-guide/manifest/overview.py new file mode 100644 index 00000000..deb56a39 --- /dev/null +++ b/figures/user-guide/manifest/overview.py @@ -0,0 +1,257 @@ +"""Draw overview.svg: the TileOPs components as a layered component diagram on a fixed grid. + +Graphviz reorders the components within a layer once edges cross layers, so this diagram is +laid out by hand: three layers of equal width, three columns, and orthogonal edges that run in +the gaps between columns. Run it as ``overview.py ``; render.sh does both. +""" + +import sys +from pathlib import Path + +FONT = "'Noto Sans SC','PingFang SC','Microsoft YaHei',sans-serif" +INK = "#2A2540" +SUB = "#5B5670" +LINE = "#6B5FA0" +BAND_FILL = "#FBFAFE" +BAND_LINE = "#CFC6E6" +BAND_INK = "#4A2C8F" +# Fill of each category, and the border and title colour its boxes take. +DEV = "#EFE9FA" +SYS = "#E2F6F8" +PUB = "#E8F6EA" +EDGE = {DEV: ("#8E6CCF", "#4A2C8F"), SYS: ("#3AA9B8", "#0B5E6A"), PUB: ("#6DBA7A", "#2E6B3A")} + +WIDTH = 700 +COL_W = 184 +GAP = 28 +BAND_X = 20 +BAND_W = WIDTH - 2 * BAND_X +COL_X = [BAND_X + 24 + i * (COL_W + GAP) for i in range(3)] + + +def text_width(s: str, size: float) -> float: + return sum(size if ord(c) > 0x2E7F else size * 0.56 for c in s) + + +def esc(s: str) -> str: + return s.replace("&", "&").replace("<", "<").replace(">", ">") + + +class Svg: + def __init__(self) -> None: + self.parts: list[str] = [] + self.boxes: dict[str, tuple[float, float, float, float]] = {} + + def text(self, x, y, s, size=13, weight="normal", fill=INK, anchor="start"): + self.parts.append( + f'{esc(s)}' + ) + + def band(self, y, h, label): + tab_w = text_width(label, 13) * (1.0 if LANG == "zh" else 1.12) + 22 + self.parts.append( + f'' + ) + self.parts.append( + f'' + ) + self.text(BAND_X + 10, y + 15, label, weight="bold", fill=BAND_INK) + + def box(self, key, col, y, title, subs=(), fill=SYS, h=None): + x = COL_X[col] + h = h or 30 + 16 * len(subs) + stroke, ink = EDGE[fill] + self.parts.append( + f'' + ) + self.text(x + 10, y + 20, title, weight="bold", fill=ink) + for i, s in enumerate(subs): + self.text(x + 10, y + 38 + 16 * i, s, size=11, fill=SUB) + self.boxes[key] = (x, y, COL_W, h) + + def edge(self, points, label=None, at=0, dashed=False, dx=6, dy=-4, anchor="start"): + """A polyline through *points*, arrowhead at the end, label beside segment *at*.""" + d = "M" + " L".join(f"{x:.1f},{y:.1f}" for x, y in points) + dash = ' stroke-dasharray="6,4"' if dashed else "" + self.parts.append( + f'' + ) + if label: + (x1, y1), (x2, y2) = points[at], points[at + 1] + mx, my = (x1 + x2) / 2 + dx, (y1 + y2) / 2 + dy + w = text_width(label, 12) + bx = mx if anchor == "start" else mx - w / 2 if anchor == "middle" else mx - w + self.parts.append( + f'' + ) + self.text(mx, my, label, size=12, fill=INK, anchor=anchor) + + # Anchor points of a box. + def top(self, k, f=0.5): + x, y, w, _ = self.boxes[k] + return (x + w * f, y) + + def bottom(self, k, f=0.5): + x, y, w, h = self.boxes[k] + return (x + w * f, y + h) + + def left(self, k): + x, y, _, h = self.boxes[k] + return (x, y + h / 2) + + def right(self, k): + x, y, w, h = self.boxes[k] + return (x + w, y + h / 2) + + def render(self, height) -> str: + head = ( + f'' + '' + f'' + ) + return head + "".join(self.parts) + "\n" + + +def L(zh: str, en: str) -> str: + """The label in the language being drawn.""" + return zh if LANG == "zh" else en + + +def draw() -> str: + s = Svg() + s.text(WIDTH / 2, 28, L("TileOPs 的组件与调用关系", "TileOPs components and how they call each other"), size=15, weight="bold", fill=BAND_INK, anchor="middle") + + # 顶行:图例在左侧,spec 居中。 + s.box("spec", 1, 48, "spec", ["src/tileops/manifest/spec/"], fill=DEV) + lx, ly = COL_X[0], 44 + s.parts.append( + f'' + ) + for i, (color, name) in enumerate([(DEV, L("开发者编写", "Written by the developer")), (SYS, L("系统提供，开发者不修改", "Provided by the system")), (PUB, L("发布", "Publication"))]): + yy = ly + 12 + 21 * i + s.parts.append( + f'' + ) + s.text(lx + 36, yy + 11, name, size=12) + + # 实现层。 + y1 = 140 + s.band(y1, 228, L("实现", "Implementation")) + r1 = y1 + 48 + s.box("gen", 0, r1, L("代码生成", "Code generation"), [L("调用检查、形状推导", "call checks, shape inference"), L("dtype 检查、eval_roofline", "dtype checks, eval_roofline")]) + s.box("op", 1, r1, L("Op 类", "Op class"), [L("__init__、forward、docstring", "__init__, forward, docstring"), L("kernel_types、interfaces", "kernel_types, interfaces")], fill=DEV) + s.box("base", 2, r1, L("Op 基类", "Op base"), [L("kernel_for：选择实现", "kernel_for: selection"), L("entry 缓存、target 派发", "entry cache, target dispatch")]) + s.box("kernel", 2, r1 + 104, L("Kernel 实现", "Kernel implementation"), [L("继承 kernel 接口", "inherits a kernel interface"), L("applies、entry_for 按需声明", "applies, entry_for as needed")], fill=DEV) + + # 验证与测量层。 + y2 = y1 + 228 + 40 + s.band(y2, 294, L("验证与测量", "Validation and measurement")) + q1 = y2 + 36 + q2 = q1 + 96 + q3 = q2 + 82 + s.box("ref", 0, q1, L("参考实现", "Reference"), [L("workloads/ 的 ref_program", "ref_program in workloads/")], fill=DEV) + s.box("val", 1, q1, "validator", ["validate_manifest.py"]) + s.box("inst", 2, q1, L("workload 实例化", "Workload instantiation"), ["manifest_calls"]) + s.box("test", 0, q2, L("正确性测试", "Correctness tests"), ["tests/ops/"], fill=DEV) + s.box("bench", 1, q2, L("benchmark 函数", "Benchmark function"), ["benchmarks/ops/"], fill=DEV) + s.box("mtest", 2, q2, L("manifest 测试", "Manifest tests"), [L("meta 调用、target conformance", "meta calls, target conformance")]) + s.box("mb", 1, q3, "ManifestBenchmark", [L("计时、FLOPs 与字节数", "timing, FLOPs and bytes")]) + + # 发布层。 + y3 = y2 + 294 + 40 + s.band(y3, 112, L("发布", "Publication")) + p1 = y3 + 36 + s.box("nightly", 0, p1, "nightly", [L("运行全部 benchmark", "runs every benchmark"), L("按 case id 记录历史", "keeps history by case id")], fill=PUB) + s.box("m5", 1, p1, L("roofline 工具", "Roofline tool"), [L("SOL 效率", "SOL efficiency"), L("与瓶颈判定", "and bound")], fill=PUB) + s.box("site", 2, p1, L("文档站", "Docs site"), [L("读取 spec、docstring", "reads specs, docstrings"), L("与 benchmark 结果", "and benchmark results")], fill=PUB) + height = y3 + 112 + 20 + + # spec 到两层的输入。 + sb = s.bottom("spec", 0.3) + turn = (sb[1] + y1) / 2 + 8 + gx0 = s.top("gen")[0] + s.edge([sb, (sb[0], turn), (gx0, turn), s.top("gen")], L("签名、roofline", "signature, roofline"), at=0, dx=6, dy=4) + sx, sy = s.right("spec") + rail = BAND_X + BAND_W + 10 + drop = COL_X[2] + COL_W / 2 + s.edge( + [(sx, sy), (rail, sy), (rail, y2 - 18), (drop, y2 - 18), (drop, y2 + 20)], + L("workload 行、全部字段", "workload rows, all fields"), + at=0, + dx=0, + dy=-5, + anchor="middle", + ) + + # 实现层内部。 + ay = r1 + 20 + s.edge([(s.right("gen")[0], ay), (s.left("op")[0], ay)]) + s.edge([(s.right("op")[0], ay), (s.left("base")[0], ay)]) + s.text(s.right("gen")[0] + GAP / 2, r1 - 7, L("安装生成的方法", "installs methods"), size=12, anchor="middle") + s.text(s.right("op")[0] + GAP / 2, r1 - 7, "kernel_for", size=12, anchor="middle") + s.edge([s.bottom("base"), s.top("kernel")], L("构造并缓存", "builds and caches")) + + # 验证与测量层对实现层的调用与核对。 + # Clear of the band's tab, whose width follows the language of its label. + tab_end = BAND_X + text_width(L("验证与测量", "Validation and measurement"), 13) * (1.0 if LANG == "zh" else 1.12) + 22 + tx = max(COL_X[0] + COL_W / 2, tab_end + 24) + s.edge( + [(tx, y2 + 20), (tx, y1 + 228)], + L("测试与 benchmark 执行调用", "tests and benchmarks call the op"), + dashed=True, + dx=6 if LANG == "zh" else -6, + anchor="start" if LANG == "zh" else "end", + ) + vx = s.top("val")[0] + s.edge([(vx, q1), (vx, s.bottom("op")[1])], L("核对接口", "checks the interface"), at=0, dy=30, dashed=True) + + # 验证与测量层内部。 + s.edge([s.bottom("ref", 0.3), s.top("test", 0.3)], L("数值参考", "reference")) + gap_y1 = (q1 + 46 + q2) / 2 + s.edge( + [s.bottom("ref", 0.8), (s.bottom("ref", 0.8)[0], gap_y1), (s.top("bench", 0.2)[0], gap_y1), s.top("bench", 0.2)], + L("torch 基线", "torch baseline"), + at=1, + dx=0, + dy=-4, + anchor="middle", + ) + s.edge([s.bottom("inst", 0.7), s.top("mtest", 0.7)], L("调用", "calls")) + gx = COL_X[2] - GAP / 2 + ix = s.bottom("inst", 0.2)[0] + mby = s.right("mb")[1] + s.edge( + [s.bottom("inst", 0.2), (ix, gap_y1), (gx, gap_y1), (gx, mby), s.right("mb")], + L("调用", "calls"), + at=3, + dx=6, + dy=-4, + ) + s.edge([s.bottom("bench"), s.top("mb")], L("对比基线", "baselines")) + + # 验证与测量层到发布层。 + mbx, mbb = s.bottom("mb") + nx = s.top("nightly")[0] + mid = (mbb + p1) / 2 + 6 + s.edge([(mbx, mbb), (mbx, mid), (nx, mid), (nx, p1)], "device time", at=1, dx=0, dy=-4, anchor="middle") + s.edge([s.right("nightly"), s.left("m5")]) + s.edge([s.right("m5"), s.left("site")]) + + return s.render(height) + + +LANG = "zh" + +if __name__ == "__main__": + LANG = sys.argv[1] + Path(sys.argv[2]).write_text(draw(), encoding="utf-8") diff --git a/figures/user-guide/render.sh b/figures/user-guide/render.sh new file mode 100755 index 00000000..7213e254 --- /dev/null +++ b/figures/user-guide/render.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# Render the user-guide figures into docs/user-guide//img/. +# +# figures/user-guide/render.sh +# +# .zh.puml becomes img/.zh.svg, the Chinese page's figure, and .en.puml +# becomes img/.svg, the English one; manifest/overview.py draws both languages. The +# PlantUML release is pinned because releases lay diagrams out differently. +set -euo pipefail +cd "$(dirname "$0")" + +VERSION="1.2026.6" +JAR="${PLANTUML_JAR:-$HOME/.local/lib/plantuml.jar}" +if ! java -jar "$JAR" -version 2>/dev/null | grep -q "PlantUML version ${VERSION}"; then + echo "PlantUML ${VERSION} not found at ${JAR}; set PLANTUML_JAR. It also needs graphviz and Noto Sans SC." >&2 + exit 1 +fi + +DOCS=../../docs/user-guide +tmp="$(mktemp -d)" +trap 'rm -rf "$tmp"' EXIT + +for src in */*.puml; do + guide="$(dirname "$src")" + base="$(basename "$src" .puml)" # .zh or .en + name="${base%.*}" + lang="${base##*.}" + suffix="" + if [ "$lang" = zh ]; then suffix=".zh"; fi + out="$DOCS/$guide/img/$name$suffix.svg" + rm -rf "$tmp/out" && mkdir -p "$tmp/out" + java -jar "$JAR" -tsvg -failfast2 -o "$tmp/out" "$src" + # PlantUML writes one font-family and no trailing newline. + sed "s/font-family=\"'Noto Sans SC'\"/font-family=\"'Noto Sans SC','PingFang SC','Microsoft YaHei',sans-serif\"/g" \ + "$tmp"/out/*.svg >"$out" + if [[ -n "$(tail -c1 "$out")" ]]; then printf '\n' >>"$out"; fi + echo "$out" +done + +python3 manifest/overview.py zh "$DOCS/manifest/img/overview.zh.svg" +python3 manifest/overview.py en "$DOCS/manifest/img/overview.svg" +echo "$DOCS/manifest/img/overview.{zh.,}svg" diff --git a/mkdocs.yml b/mkdocs.yml index 28e20258..97ac87ba 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -31,7 +31,17 @@ nav: - Home: index.md - User Guide: - user-guide/index.md - - Writing a Spec: manifest.md + - Reading and Writing Specs: + - user-guide/manifest/index.md + - Concepts: user-guide/manifest/concepts.md + - Spec Fields: user-guide/manifest/writing.md + - Extensions: user-guide/manifest/extensions.md + - Calls and Validation: user-guide/manifest/calls.md + - Examples: user-guide/manifest/examples.md + - Kernel Dispatch: + - user-guide/dispatch/index.md + - Adding a Kernel: user-guide/dispatch/writing.md + - Backend Integration: user-guide/dispatch/backends.md - Adding an Op: new-op.md - Timing Benchmarks: timing.md - Using torch.compile: torch-compile.md @@ -152,14 +162,22 @@ plugins: Home: 首页 Design: 设计文档 User Guide: 使用指南 - Writing a Spec: 读写 manifest - Adding an Op: 添加新算子 - Timing Benchmarks: benchmark 怎么计时 + Reading and Writing Specs: 读写 manifest + Concepts: 概念 + Spec Fields: 写一个 spec + Extensions: 扩展写法 + Calls and Validation: 调用与校验 + Examples: 示例 + Kernel Dispatch: op 如何选择 kernel + Adding a Kernel: 如何为 op 新增 kernel + Backend Integration: backend 如何接入 + Adding an Op: 添加新 op + Timing Benchmarks: benchmark 的计时方法 Using torch.compile: 接入 torch.compile - Adding a Backend: 接入新硬件后端 + Adding a Backend: 接入新硬件 backend Architecture: 架构 - Manifest: 算子清单 - Op Interfaces: 算子接口 + Manifest: manifest 规范 + Op Interfaces: op 接口 Slot Rules: Slot 规则 Roofline: Roofline Testing: 测试