[Docs] Describe op to kernel dispatch as kernel interfaces and implementations

docs/backends.md

@@ -11,16 +11,98 @@ devices run your kernels.
 **A backend supplies one thing: something callable that computes this call.**
 Everything else is the op layer's.
 
-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 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.
 
 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.
 
+## 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.
+
+**`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
+
+
+register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm)
+```
+
+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.
+
+`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.
+
 ## Four things to write
 
 | # | What to do |
@@ -222,9 +304,9 @@ signature. Whatever it returns, the op layer stores and launches.
 
 Four things follow from that:
 
-- **`kernel_for` and `entry_for` are the op author's, not a backend's.** They serve the
-  in-tree path only: which in-tree kernel to fetch, 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.
+- **`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.
 - **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.
@@ -233,8 +315,9 @@ Four things follow from that:
   `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 with no in-tree implementation may leave `entry_for` out, and then
-  a call with no target claiming the device raises `OpNotAvailableError`.
+  `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`.
 
 ## Writing a backend that runs {#runnable}
 
@@ -553,9 +636,11 @@ print(load_failures())
 | **`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;
-there is no kernel-level concept, no capability negotiation and no candidate
-filtering.
+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 —
@@ -592,10 +677,10 @@ third-party backend neither bypasses one 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 default state is no substitution.** With no backend claiming a device, calls on
-it run the shipped implementation; only once a backend is installed and has claimed
-that device does the op's kernel become the backend's. The protocol has no notion of a
-"default target".
+**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`.
 
 ## When a kernel is rebuilt {#memo}
 

docs/backends.zh.md

@@ -6,10 +6,68 @@ TileLang 是多后端 DSL，每种硬件各有一套独立的 kernel，由各自
 
 **后端只提供一件事：一个能算这次调用的可调用对象。** 其余都由算子层负责。
 
-前半按动手顺序列出要做的事：要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么，以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。
+本页讲的是 target，接管整个算子的那一种接入方式。前半部分按实现顺序列出要做的事：要写的四样东西、协议中的四个函数、一次调用怎么走到它们、一个可直接安装运行的后端、怎么从模板改造为面向真实硬件的后端、编写 kernel 的四条规则、各阶段允许做什么，以及装好之后每个算子处于哪种状态、各条错误信息对应什么原因。
 
 后半说明协议何以如此设计：两层选择、算子层的契约、kernel 的重建条件、调用方可用的接口，以及刻意不支持的情形。
 
+## 三种接入方式 {#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
+
+
+register_implementation("LayerNormFwdOp", "torch_short_rows", _NarrowTorchLayerNorm)
+```
+
+这样 `n <= 64` 的调用由 `_NarrowTorchLayerNorm` 服务，`n = 1024` 仍由自带的 `LayerNormKernel` 服务。新增的实现只对注册之后构造的算子实例生效。同一个算子下重复注册同一个 key 报 `BackendError`；key 与自带实现的 key 相同时，构造实例时报 `reuse keys it has`。
+
+`register_implementation` 在后端模块被导入时执行，通过 target 所用的同一个 entry point 触发；`kernel_map=` 则由调用方在构造算子时传入，不经过注册。以下各节讲 target。
+
 ## 写一个后端要做的四件事
 
 | # | 做什么 |
@@ -184,10 +242,10 @@ register_kernel_builder(op="GemmFwdOp", target="acme", build_kernel=build_gemm)
 
 四点对应关系值得记住：
 
-- **`kernel_for` 与 `entry_for` 由算子作者写，与后端无关。** 它们只服务自带实现：取哪个自带 kernel、按什么查表、又怎么构造。target 选中后端时整个算子由 target 服务，这几处都不会执行。
+- **`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` 内部另加一层缓存。算子完全没有自带实现时 `entry_for` 可以不写，那时没有 target 认领设备，调用直接抛 `OpNotAvailableError`。
+- **不必自己缓存构建结果。** 设备与输入签名相同时，算子层不会再调用一次 `build_kernel`；需要更细的区分或更少的重建，在 `build_kernel` 内部另加一层缓存。专为外部后端而写的算子，`kernel_types` 与 `interfaces` 都不声明，没有 target 认领设备时调用直接抛 `OpNotAvailableError`。
 
 ## 实现一个可运行的后端 {#runnable}
 
@@ -443,7 +501,7 @@ print(load_failures())
 | **`detect`** | 后端写的一个函数，一个 target 一个 | 第一层怎么选：接收一块 `torch.device`，回答这类设备是不是自己这套 kernel 的目标设备；不是则返回 `False` |
 | **`build_kernel`** | 后端为某个算子写的一个函数，一组 `(算子, target)` 一个 | 第二层：接收本次调用的描述，即各输入张量的 device、dtype、shape 与算子参数，在自己这套 kernel 中选定一个、构造好并返回 |
 
-**选择分两层：TileOPs 选 target，target 在自己那套 kernel 里选一个。** 第二层发生在 `build_kernel` 内部，协议不参与，也不存在 kernel 一级的概念、能力协商与候选筛选。
+**选择分两层：TileOPs 选 target，target 在自己那套 kernel 里选一个。** 第二层发生在 `build_kernel` 内部，协议不参与：这条路径上不存在 kernel 一级的概念、能力协商与候选筛选。TileOPs 确实执行的候选筛选 —— 可用性、适用范围、优先关系 —— 属于自带实现与范围较小的两种接入方式（[三种接入方式](#three-ways)），target 绕过它们。
 
 `detect` 只回答设备的归属，粒度到此为止。**本次调用是否受支持 —— 涉及 dtype、形状与参数组合 —— 由 `build_kernel` 回答**，因为只有它看得到完整的输入描述与参数；不支持时在那里报错。这些判断交给 `detect` 是做不到的，它只拿到一块 `torch.device`。
 
@@ -471,7 +529,7 @@ TileOPs 不解析 `torch.device`，而是把它原样传给 `detect`。这样做
 
 TileOPs 自带的 kernel（[`src/tileops/kernels/`](https://github.com/tile-ai/TileOPs/tree/main/src/tileops/kernels)）是**默认实现**：它没有 target 名，也不进注册表。
 
-**默认状态是不替换。** 没有后端认领某块设备时，这台设备上的调用走自带实现；装上一个后端、并且它认领了这块设备，该算子的 kernel 才换成后端的那一套。协议里因此不存在「默认 target」这个概念。
+**默认走自带实现。** 没有装后端、没有指名 `target=`、也没有设置进程默认值时，调用由自带实现服务。装上一个后端之后，要它的 `detect` 认领了这块设备，或者它被 `target=` 与 `set_default_target` 指名，该算子的 kernel 才换成这个后端的那一套。
 
 ## kernel 的重建条件 {#memo}
 

docs/manifest.md

@@ -48,7 +48,7 @@ required once the other direction also has an entry — and the validator requir
 | `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 roles |
+| `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:
 

docs/manifest.zh.md

@@ -33,7 +33,7 @@ TileOPs 的组织方式相反：算子的规格先声明，实现由规格推导
 | `signature` | 是 | 算子的类型，见下表 |
 | `workloads` | 是 | 测试与 benchmark 执行的调用 |
 | `roofline` | 是 | 一次调用的开销，规范见 [Roofline](design/roofline.md) |
-| `composition` | 否 | 复合算子按顺序列出的各阶段：可能持有的子算子类，以及它自己的 kernel 角色 |
+| `composition` | 否 | 复合算子按顺序列出的各阶段：可能持有的子算子类，以及它自己的 kernel key |
 
 签名是一个以具名类型变量为参数的函数类型：