From dc446cc54ca2257f06a4d631d285e13d77226932 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:09 -0400 Subject: [PATCH 1/9] Attribute HBM traffic per tensor in the latency trace MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `dram_bytes` answers how much traffic a kernel moved; nothing answered which tensor moved it. That makes a claim about composition — "this weight tensor's traffic is negligible" — uncheckable, because it is a statement about one term of a total rather than about the total. `_estimate` already computes each op's bytes and FLOPs and `record` already receives them, but only the per-category sums were kept. Retain them per operation instead, alongside the memory the operation addressed, and add `LatencyReport.traffic_by_target()` to aggregate on it. The target is the enclosing `MemRef.base_ptr` — an element index, so it compares directly against the pointer `execute_function` binds to an argument, which `KTIRInterpreter.arg_ptrs` now retains. Bytes whose origin is not resolvable are kept under a `None` key rather than dropped: a breakdown that silently omits a row adds up to less than the total while looking complete. A distributed memory view appears as one key per partition, since `distributed_load` charges each surviving partition separately; folding those into one row per argument needs each argument's element extent, which the report does not know. `tests/test_latency.py` covers both halves of the new method. The attribution is hand-counted off the IR for `matmul_small` — one row per argument at the pointer it was bound to, four times as many bytes against B as against A, and the rows summing to `dram_bytes` — because a total of 45,056 says nothing about the split and a sentence calling one of the two negligible survives the total but not the breakdown. Without the trace the method raises rather than returning an empty mapping, and it distinguishes the two ways there is nothing to report: a report covering no core at all, and one covering cores that were never traced. Returning an empty mapping for either would make a missing constructor argument look like a finding about the kernel. The same file gains a chip-wide hand count of bytes and FLOPs for `matmul_small`, the companion to the existing single-core count for the simplest kernel, at exact equality rather than a tolerance — both sides are integer counts, and a mis-charged operation is exactly the small disagreement a tolerance hides. The stick-rule arithmetic behind the two `matmul_small` hand counts — this one and the attribution above — is written once and shared, so the two cannot drift into disagreeing about the same kernel. No cost formula changes. The trace only carries figures it was already given. Signed-off-by: WarningRan --- docs/latency.md | 40 ++++++++ ktir_cpu/interpreter.py | 12 +++ ktir_cpu/latency.py | 200 ++++++++++++++++++++++++++++++++++++---- tests/test_latency.py | 143 ++++++++++++++++++++++++++++ 4 files changed, 375 insertions(+), 20 deletions(-) diff --git a/docs/latency.md b/docs/latency.md index 2d8f11d..44ae998 100644 --- a/docs/latency.md +++ b/docs/latency.md @@ -123,6 +123,46 @@ Total: 40,960 bytes **Open question**: the model applies `hbm_bandwidth_tb_s` uniformly across all dtypes. Whether f16 data tensors and i32 index tensors achieve the same effective HBM bandwidth on Spyre is unconfirmed — if they differ, a separate bandwidth parameter would be needed. +## Per-tensor traffic attribution + +`dram_bytes` says how much HBM traffic a kernel moved; `traffic_by_target()` says +*which tensor* moved it — which is what makes a claim about composition ("this +weight tensor's traffic is negligible") checkable at all. Both read the same +per-operation figures: `record()` folds each op's `nbytes` and `flops` into the +per-category aggregates and, with `trace_latency=True`, also keeps them per op +alongside the memory that op addressed. + +```python +interp = KTIRInterpreter(latency_config=HardwareConfig(), trace_latency=True) +interp.load(mlir); interp.execute_function("add_kernel", **args) +report = interp.get_latency_report() + +ptr_to_arg = {p: a for a, p in interp.arg_ptrs.items() if isinstance(p, int)} +for target, row in report.traffic_by_target().items(): + print(ptr_to_arg.get(target, target), row["nbytes"], row["ops"]) +``` + +For `examples/triton-ktir/vector_add_ktir.mlir` at 4096 elements, `grid = [32, 1]`: + +``` +x_ptr 8192 {'ktdp.load': 32} +y_ptr 8192 {'ktdp.load': 32} +output_ptr 8192 {'ktdp.store': 32} +``` + +Keys are element indices — the origin of the view an operation addressed, which is +the same value `execute_function` binds to a pointer argument, hence the `arg_ptrs` +lookup. Two properties matter when reading the result: + +- **The rows always sum to the category total.** Traffic whose origin could not be + resolved is kept under the key `None` rather than dropped, so a breakdown never + adds up to less than `dram_bytes` while looking complete. +- **A distributed memory view appears as one key per partition**, not one per + tensor: `distributed_load` charges each surviving partition separately, and each + partition is a distinct region of memory. The report cannot fold them into + arguments — it does not know each argument's element extent; fold key `t` into + argument `a` when `ptr[a] <= t < ptr[a] + numel(a)`. + ## Hardware parameters Default `HardwareConfig()` values: diff --git a/ktir_cpu/interpreter.py b/ktir_cpu/interpreter.py index e3b8d1d..7c6fd41 100644 --- a/ktir_cpu/interpreter.py +++ b/ktir_cpu/interpreter.py @@ -71,6 +71,9 @@ def __init__( # Keeping the old name for now; the type is TransferBackend. self.ring_backend: Optional[TransferBackend] = None self._env: Optional[ExecutionEnv] = None + # arg name -> pointer value bound by the last execute_function call. + # Element index for tensor arguments, verbatim for scalars. + self.arg_ptrs: Dict[str, Any] = {} self._parser: Optional[KTIRParserBase] = parser self._latency_config: Optional[HardwareConfig] = latency_config self._latency_tracker: Optional[LatencyTracker] = ( @@ -141,6 +144,10 @@ def execute_function(self, func_name: str, **kwargs) -> Dict[str, np.ndarray]: if not self.module: raise RuntimeError("No module loaded. Call load() first.") + # Cleared before anything can raise, so a failed run leaves no pointer map + # behind: cost attribution reading the previous call's bindings would name + # the wrong tensors and look like a correct breakdown while doing it. + self.arg_ptrs = {} func = self.module.get_function(func_name) self._prepare_execution(func.grid) @@ -171,6 +178,11 @@ def execute_function(self, func_name: str, **kwargs) -> Dict[str, np.ndarray]: # Scalar argument (like n) input_ptrs[arg_name] = tensor + # Retained for cost attribution: LatencyReport.traffic_by_target keys on + # the element index a view starts at, which is exactly the value bound + # here, so this is what turns those keys back into argument names. + self.arg_ptrs = dict(input_ptrs) + for core in self.grid_executor.cores: core._use_counts = func.use_counts diff --git a/ktir_cpu/latency.py b/ktir_cpu/latency.py index 1509121..9ab6f26 100644 --- a/ktir_cpu/latency.py +++ b/ktir_cpu/latency.py @@ -184,12 +184,41 @@ def ring_bytes_per_cycle(self) -> float: # Per-core latency counters # --------------------------------------------------------------------------- +def _sole_partition_origin(ref: DistributedTileRef) -> Optional[int]: + """Origin of the one partition *ref* survived into, or ``None`` if not one. + + A distributed access resolves to the partitions it actually touches. When + that is a single partition its origin identifies the memory; when it is + several there is no single origin, and naming one of them would report a + guess as a fact. + """ + if len(ref.partitions) != 1: + return None + return ref.partitions[0].memref.base_ptr + + @dataclass class _TraceEntry: - """Single operation trace entry.""" + """Single operation trace entry. + + ``nbytes`` and ``flops`` are the same figures ``record`` folds into the + per-category aggregates. Keeping them per-op as well is what lets a caller + attribute a kernel's HBM traffic to individual tensors instead of reading a + single total: a claim that "the weight tensor is negligible" is checkable + against ``traffic_by_target`` and not against ``dram_bytes``. + + ``target`` identifies the memory the op addressed, as the **element index** + of the origin of the enclosing view (``MemRef.base_ptr``), or ``None`` when + the op touches no memory or the origin is not derivable. It is an element + index rather than a byte address so it can be compared directly with the + pointer values ``KTIRInterpreter.execute_function`` binds to arguments. + """ op_type: str cycles: float category: str + nbytes: int = 0 + flops: float = 0.0 + target: Optional[int] = None @dataclass @@ -242,7 +271,8 @@ def dram_bytes(self) -> int: return self.bytes_by_category.get("memory", 0) def record(self, category: str, cycles: float, op_type: str = "", - flops: float = 0.0, nbytes: int = 0): + flops: float = 0.0, nbytes: int = 0, + target: Optional[int] = None): if category.startswith("compute_"): self.cycles_by_category[category] = self.cycles_by_category.get(category, 0.0) + cycles self.flops_by_category[category] = self.flops_by_category.get(category, 0.0) + flops @@ -257,7 +287,10 @@ def record(self, category: str, cycles: float, op_type: str = "", self.bytes_by_category[category] = self.bytes_by_category.get(category, 0) + nbytes if self.trace is not None: - self.trace.append(_TraceEntry(op_type=op_type, cycles=cycles, category=category)) + self.trace.append(_TraceEntry( + op_type=op_type, cycles=cycles, category=category, + nbytes=nbytes, flops=flops, target=target, + )) # --------------------------------------------------------------------------- @@ -300,7 +333,10 @@ def record_op(self, core_id: int, op_type: str, result: Any, operands: List[Any] trace=[] if self._trace else None ) category, cycles, flops, nbytes = self._estimate(op_type, result, operands) - self.counters[core_id].record(category, cycles, op_type, flops=flops, nbytes=nbytes) + self.counters[core_id].record( + category, cycles, op_type, flops=flops, nbytes=nbytes, + target=self._target(operands) if nbytes else None, + ) def report(self) -> "LatencyReport": """Build a LatencyReport from accumulated counters.""" @@ -392,6 +428,74 @@ def _estimate(self, op_type: str, result: Any, operands: List[Any]) -> Tuple[str # Unknown category raise NotImplementedError(f"Unknown category {category}") + @staticmethod + def _first_ref(operands: List[Any], *types: type) -> Optional[Any]: + """The first operand that is one of *types*, or ``None``. + + The walk is shared because the two questions asked over it are not: what + tensor an op addressed and what memory space it landed in are answered by + different bodies, but by the same operand. Two hand-written copies of the + ladder had already drifted — ``_target`` answers for a bare + ``DistributedTileRef`` operand and ``_memory_space`` does not — and the + remaining difference is now a visible argument list rather than a + discrepancy between two ladders nobody reads side by side. + + The five ref types are unrelated classes, so the order of *types* selects + nothing; it is each caller's own dispatch that depends on order. + """ + for v in operands: + if isinstance(v, types): + return v + return None + + @staticmethod + def _target(operands: List[Any]) -> Optional[int]: + """Element index of the origin of the view this op addressed. + + Mirrors the operand walk in :meth:`_memory_space`, but answers *which + tensor* rather than *which memory space*. Always the enclosing + ``MemRef.base_ptr`` (an element index), never ``TileRef.base_ptr`` + (a byte address), so the value is directly comparable with the pointer + an argument was bound to. + + A ``DistributedTileRef`` carries the partitions *this access* survived + into, not every partition of the view, so its origin is that partition's + and not the logical tensor's: a tensor sharded two ways yields two + distinct targets. Folding those back into one row per tensor is the + caller's job, and needs each argument's element extent. + + An access surviving into more than one partition has no single origin, so + it returns ``None`` and its bytes are reported unattributed. Charging + them to the lowest origin would read as a fact about that partition; no + kernel available here exercises the case, which is the other reason not + to encode a guess about it. + + Returns ``None`` when no origin is derivable. Callers must surface an + unattributed bucket rather than dropping it: silently omitting bytes + from an attribution table makes the remaining rows add up to something + that looks complete and is not. + """ + v = LatencyTracker._first_ref( + operands, MemRef, TileRef, DistributedTileRef, AccessTile, + IndirectAccessTile) + if isinstance(v, MemRef): + return v.base_ptr + if isinstance(v, TileRef): + return v.memref.base_ptr + if isinstance(v, DistributedTileRef): + return _sole_partition_origin(v) + if isinstance(v, AccessTile): + parent = v.parent_ref + if isinstance(parent, DistributedTileRef): + return _sole_partition_origin(parent) + return parent.memref.base_ptr + if isinstance(v, IndirectAccessTile): + # An indirect load charges the gathered data and the index lookups + # as one figure (see _data_size), so this row carries both. The + # parent is the larger of the two by construction. + return v.parent_ref.base_ptr + return None + @staticmethod def _memory_space(operands: List[Any]) -> str: """Return the memory space of the memory op's TileRef target. @@ -403,22 +507,28 @@ def _memory_space(operands: List[Any]) -> str: Returns "HBM" when no TileRef is found (e.g. tt.load which always reads from HBM via pointer arithmetic). """ - for v in operands: - if isinstance(v, MemRef): - return v.memory_space - if isinstance(v, TileRef): - return v.memref.memory_space - if isinstance(v, AccessTile): - if isinstance(v.parent_ref, DistributedTileRef): - if any(p.memref.memory_space == "HBM" - for p in v.parent_ref.partitions): - return "HBM" - return v.parent_ref.partitions[0].memref.memory_space - return v.parent_ref.memref.memory_space - if isinstance(v, IndirectAccessTile): - all_lx = (v.parent_ref.memory_space == "LX" and - all(iv.memory_space == "LX" for iv in v.index_views)) - return "LX" if all_lx else "HBM" + # No DistributedTileRef in the list, and that is the pre-existing + # asymmetry with _target rather than a new decision: a bare one reaches + # here only through an op whose operand is the whole distributed view, + # which nothing under examples/ emits, and "HBM" is the safe default for + # a space this cannot resolve. Adding it would change a reported figure. + v = LatencyTracker._first_ref( + operands, MemRef, TileRef, AccessTile, IndirectAccessTile) + if isinstance(v, MemRef): + return v.memory_space + if isinstance(v, TileRef): + return v.memref.memory_space + if isinstance(v, AccessTile): + if isinstance(v.parent_ref, DistributedTileRef): + if any(p.memref.memory_space == "HBM" + for p in v.parent_ref.partitions): + return "HBM" + return v.parent_ref.partitions[0].memref.memory_space + return v.parent_ref.memref.memory_space + if isinstance(v, IndirectAccessTile): + all_lx = (v.parent_ref.memory_space == "LX" and + all(iv.memory_space == "LX" for iv in v.index_views)) + return "LX" if all_lx else "HBM" return "HBM" @staticmethod @@ -622,6 +732,56 @@ def per_core_summary(self) -> List[Dict[str, Any]]: }) return summaries + def traffic_by_target(self, category: str = "memory") -> Dict[Any, Dict[str, Any]]: + """Chip-wide bytes of one transport, attributed per tensor. + + ``dram_bytes`` answers *how much* traffic a kernel moved; this answers + *which tensor moved it*. The distinction is what makes a claim like + "the weight tensor's traffic is negligible" checkable — against a + breakdown, not against a total. + + Keys are ``_TraceEntry.target`` (the element index a view starts at, so + comparable with the pointer an argument was bound to). Bytes whose + origin was not derivable are kept under the key ``None`` instead of + being dropped, so the rows always sum to the category total. + + A key is the origin of the view the op addressed, which for a + distributed memory view is **one partition**, not the whole tensor: + ``distributed_load`` charges each surviving partition separately, so a + tensor sharded two ways appears as two keys at its two partition + origins. That is the accurate answer to "which memory moved" and not + the answer to "which argument moved"; folding partitions back into + arguments needs each argument's element extent, which this class does + not know. Callers that want per-argument rows should fold a key ``t`` + into argument ``a`` when ``ptr[a] <= t < ptr[a] + numel(a)``. + + Requires the tracker to have been created with ``trace_latency=True``: + the per-op figures are only retained then. Raises otherwise rather + than returning an empty mapping, which would read as "no traffic". + The two ways of having no trace are reported separately, because + "nothing ran" and "tracing was off" call for different fixes and one + condition covering both names the wrong one half the time. + """ + if not self.counters: + raise RuntimeError( + "traffic_by_target has nothing to attribute: this report " + "covers no core, so the kernel has not been run through it." + ) + if all(c.trace is None for c in self.counters.values()): + raise RuntimeError( + "traffic_by_target needs the per-op trace: construct the " + "interpreter with trace_latency=True." + ) + out: Dict[Any, Dict[str, Any]] = {} + for counters in self.counters.values(): + for entry in counters.trace or (): + if entry.category != category or not entry.nbytes: + continue + row = out.setdefault(entry.target, {"nbytes": 0, "ops": {}}) + row["nbytes"] += entry.nbytes + row["ops"][entry.op_type] = row["ops"].get(entry.op_type, 0) + 1 + return out + # ------------------------------------------------------------------ # Roofline — unified formulation # ------------------------------------------------------------------ diff --git a/tests/test_latency.py b/tests/test_latency.py index 8853b91..35ebdcf 100644 --- a/tests/test_latency.py +++ b/tests/test_latency.py @@ -23,6 +23,7 @@ from ktir_cpu import KTIRInterpreter, HardwareConfig, LatencyReport from ktir_cpu.dtypes import stick_to_elem_idx +from ktir_cpu.memory import HBMSimulator from conftest import EXAMPLES_DIR, get_test_params, parse_example @@ -93,6 +94,30 @@ def _run_matmul(path, func_name, entry, cfg, trace=False): return interp.get_latency_report() +def _matmul_hand_count(entry): + """A matmul entry's execute kwargs, and the stick rule that prices one tile. + + Shared by the two tests that hand-count this kernel, because the stick rule is + exactly the part that has to stay identical between them: rounding a tile row + up to a whole stick is where a hand count and the model most easily disagree, + and two copies of it can only be corrected in one place. + + The rule is specific to these shapes -- every tile row here is a contiguous run + of at most one stick and every view's row stride is a whole number of sticks, + so a row costs one stick and a tile costs one stick per row. It does not + generalise, which is why it is written out here rather than read back out of + the model the tests are checking. + """ + kwargs = {k: v for k, v in entry["execute_kwargs"].items() if v is not None} + f16 = 2 + + def tile_bytes(rows, cols): + stick = HBMSimulator.STICK_BYTES + return rows * -(-cols * f16 // stick) * stick # ceil to whole sticks + + return kwargs, tile_bytes + + def _run_vector_reduce(path, func_name, entry, cfg, trace=False): """Run a vector reduce (per-core tile) and return report.""" interp = KTIRInterpreter(latency_config=cfg, trace_latency=trace) @@ -650,6 +675,124 @@ def test_matmul_flops(self, path, func_name, entry): # Each iteration does one linalg.matmul of shape (bm × bk) × (bk × bn) assert core0.total_flops >= 2.0 * bm * bn * bk * n_iters + @pytest.mark.parametrize("path,func_name,entry", get_test_params("matmul_kernel_small")) + def test_matmul_chip_bytes_and_flops_match_hand_count(self, path, func_name, entry): + """Chip-wide traffic and FLOPs, counted off the IR by hand, exactly. + + The companion to ``test_vector_add_flops_and_bytes``: that one pins one + core of the simplest kernel, this one pins the whole chip of a tiled one, + where the count has to get the grid, the loop trip count and stick + granularity right rather than just the element count. + + Exact equality and not a tolerance. Both sides are integer counts, so the + only thing a tolerance would buy is silence about a real disagreement — + and a small one is what a mis-charged operation looks like. Charging an + integer compare as a float compute, which is what the model did until + recently, moves a total by parts in 100,000. + + The stick rule the count rests on is written out in + ``_matmul_hand_count``, which the other hand-counting test shares. + """ + kwargs, tile_bytes = _matmul_hand_count(entry) + M, N, K = kwargs["M"], kwargs["N"], kwargs["K"] + bm, bn, bk = (kwargs["BLOCK_SIZE_M"], kwargs["BLOCK_SIZE_N"], + kwargs["BLOCK_SIZE_K"]) + cores = (M // bm) * (N // bn) # grid [M/bm, N/bn], one C tile each + k_iters = K // bk # scf.for trip count per core + # Per core: one A tile and one B tile per iteration, then one C store. + per_core = k_iters * (tile_bytes(bm, bk) + tile_bytes(bk, bn)) + per_core += tile_bytes(bm, bn) + # Per core: one linalg.matmul per iteration, plus the accumulate that folds + # it into the running C tile. + flops = cores * k_iters * (2 * bm * bn * bk + bm * bn) + + report = _run_matmul(path, func_name, entry, HardwareConfig()) + assert sum(c.dram_bytes for c in report.counters.values()) == cores * per_core + assert sum(c.total_flops for c in report.counters.values()) == flops + assert sum(c.comm_bytes for c in report.counters.values()) == 0, ( + "a matmul on independent C tiles has nothing to exchange" + ) + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("matmul_kernel_small")) + def test_traffic_by_target_attributes_every_byte_to_an_argument( + self, path, func_name, entry + ): + """Per-tensor attribution, hand-counted, with nothing left over. + + This is the method that makes a prose claim about one tensor checkable. + The chip total here is 45,056 bytes and says nothing about the split; the + breakdown says B moves four times what A moves, because every core reads + a (bk x bn) B tile against an (bm x bk) A tile. A sentence calling one of + the two negligible survives the total and does not survive this. + + What discriminates is the key set together with the per-row counts: + attribution that resolved every byte but keyed it on something + uncomparable -- a tile's byte address rather than a view's element index + -- would still add up. The closing sum against ``dram_bytes`` is not + independent of those on this kernel, where every byte is attributable and + no ``None`` row exists; it is here as the invariant the method's docstring + promises, so a later change that starts dropping underivable bytes rather + than parking them under ``None`` fails on a kernel that has some. + """ + kwargs, tile_bytes = _matmul_hand_count(entry) + M, N, K = kwargs["M"], kwargs["N"], kwargs["K"] + bm, bn, bk = (kwargs["BLOCK_SIZE_M"], kwargs["BLOCK_SIZE_N"], + kwargs["BLOCK_SIZE_K"]) + cores = (M // bm) * (N // bn) + k_iters = K // bk + + interp = KTIRInterpreter(latency_config=HardwareConfig(), trace_latency=True) + interp.load(path) + rng = np.random.default_rng(42) + A = rng.standard_normal((M, K)).astype(np.float16) + B = rng.standard_normal((K, N)).astype(np.float16) + C = np.zeros((M, N), dtype=np.float16) + interp.execute_function(func_name, a_ptr=A, b_ptr=B, c_ptr=C, **kwargs) + report = interp.get_latency_report() + + rows = report.traffic_by_target() + # One key per argument, at the element index that argument was bound to. + assert set(rows) == {interp.arg_ptrs[n] for n in ("a_ptr", "b_ptr", "c_ptr")} + # Every core loads one A tile and one B tile per iteration, and stores its + # C tile once at the end. + expected = { + interp.arg_ptrs["a_ptr"]: (cores * k_iters, tile_bytes(bm, bk), "ktdp.load"), + interp.arg_ptrs["b_ptr"]: (cores * k_iters, tile_bytes(bk, bn), "ktdp.load"), + interp.arg_ptrs["c_ptr"]: (cores, tile_bytes(bm, bn), "ktdp.store"), + } + for target, (n_ops, per_op, op_type) in expected.items(): + assert rows[target]["nbytes"] == n_ops * per_op + assert rows[target]["ops"] == {op_type: n_ops} + assert sum(r["nbytes"] for r in rows.values()) == sum( + c.dram_bytes for c in report.counters.values() + ), "attribution dropped bytes instead of parking them under None" + + @pytest.mark.parametrize("path,func_name,entry", get_test_params("matmul_kernel_small")) + def test_traffic_by_target_without_trace_raises(self, path, func_name, entry): + """No trace is not zero traffic, so it raises instead of returning {}. + + The per-op figures only exist under ``trace_latency=True``. An empty + mapping is a valid answer -- a kernel that moved nothing has one -- so + returning it here would make a missing constructor argument look like a + finding about the kernel. + """ + report = _run_matmul(path, func_name, entry, HardwareConfig(), trace=False) + assert report.counters, "the run itself should still have been counted" + with pytest.raises(RuntimeError, match="trace_latency=True"): + report.traffic_by_target() + + def test_traffic_by_target_on_an_empty_report_says_nothing_ran(self): + """An empty report is not a tracing mistake, and must not be reported as one. + + Both states have no trace to attribute, and the fixes are opposite: pass + ``trace_latency=True``, or run the kernel at all. One message covering both + sends half its readers to the wrong one. + """ + from ktir_cpu.latency import LatencyReport + report = LatencyReport(config=HardwareConfig(), counters={}) + with pytest.raises(RuntimeError, match="covers no core"): + report.traffic_by_target() + def test_empty_report_roofline(self): """roofline() on empty report returns empty dict.""" from ktir_cpu.latency import LatencyReport From 252e451c663341a2e496fde810d296e5a5a18637 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:15 -0400 Subject: [PATCH 2/9] Add the kernelentry declaration, argument specs and a zero-price audit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The three parts of kernelentry that stand on their own, ahead of the engine that reads them: what a kernel declares, how its arguments are built, and the one check that is asked of the repository rather than of any kernel. A declaration is a row — the kernel's file, the parameters the gate runs it at, its arguments, a reference, and any excuse. Six claim states, and two of the distinctions are load-bearing. `undetermined` is separate from `closed` because a check the engine could not evaluate must not read as one that passed — and without the state such a claim gets omitted instead, which is indistinguishable from clean in any summary. `deferred` is separate from `waived` because only one of them comes back: a deferral carries an issue reference and reads as a known gap rather than a failure, following `tests/test_spec_gaps.py`. Either excuse reports itself as unnecessary once its check starts passing, which is what fails the build until the declaration is updated — the xfail marker cannot do that on its own, since it is applied from the claim's current state and is simply absent once the gap closes. That is what lets a kernel arrive over two pull requests without the first pretending the cost leg is done. Both excuses are validated where they are written rather than where they are read: a deferral naming no issue, and a tolerance naming an argument the kernel does not have, are rejected at declaration. An argument is a spec resolved against the row's parameters — `normal(("M", "K"))`, `zeros(("M", "N"))`, a bare string forwarding a parameter — so declaring a kernel writes a row and not a file, and a callable `(params, rng)` remains for the arguments no spec covers. Specs are per argument rather than per kernel, which is what lets a kernel declare one tensor the long way and leave the rest as rows. Every draw is seeded from the argument's own name: one generator drawn twice gives the second tensor the first one's values, and a kernel that swapped its operands would then still agree with its reference. The pricing check is repository-wide rather than a per-kernel claim, which a probe-only prototype over six real artifacts settled. Asked per kernel it produced 66 of that prototype's 74 open claims, about fifteen of them the same structural ops in every kernel — `arith.constant`, `scf.*`, `construct_memory_view`, `return` — which really are free; per kernel it is a chore each contributor waives their way through, and a waiver mapping that is mostly noise stops carrying information. Asked once per repository it is `pricing.py`: every op the registry prices `zero` — 46 of 106 today — has to appear in exactly one of two mappings, each entry carrying a written reason. `ZERO_COST_OPS` is free by decision, 21 ops; `UNJUDGED_ZERO_OPS` is priced zero with nobody having decided that, 25 ops, each reason naming the issue that would settle it, because an open question with no issue behind it is how such a list becomes permanent. An op in neither, or in both, is a finding. So registering an op without naming a category no longer passes silently — and it does not become priced either, which is the point: the gate asks for a decision, and recording "not yet judged" against an issue is a legitimate answer. What the audit cannot see is a category that is simply wrong. An integer compare billed to the float pipe is priced, so a comparison against `zero` passes it, and that defect survived in this repository for as long as it existed. Only a reader who knows the op's semantics catches that one. `conformance.py` records which kernels under `examples/` the MLIR frontend rejects, and why. That has to be written down rather than measured at render time: the check needs an optional dependency, so on a machine without the MLIR bindings the claim reads `skip`, and a report carrying whatever this machine saw would be a document about the machine. It is a record and not a verdict — it does not say whether the file or the parser should change. Same shape and same reason as `tests/mlir_frontend/test_registry_consistency.py::FRONTEND_UNSUPPORTED`. Signed-off-by: WarningRan --- ktir_cpu/kernelentry/__init__.py | 244 +++++++++++++++++++++ ktir_cpu/kernelentry/conformance.py | 73 +++++++ ktir_cpu/kernelentry/pricing.py | 315 ++++++++++++++++++++++++++++ ktir_cpu/kernelentry/tensorspec.py | 278 ++++++++++++++++++++++++ 4 files changed, 910 insertions(+) create mode 100644 ktir_cpu/kernelentry/__init__.py create mode 100644 ktir_cpu/kernelentry/conformance.py create mode 100644 ktir_cpu/kernelentry/pricing.py create mode 100644 ktir_cpu/kernelentry/tensorspec.py diff --git a/ktir_cpu/kernelentry/__init__.py b/ktir_cpu/kernelentry/__init__.py new file mode 100644 index 0000000..55f0a36 --- /dev/null +++ b/ktir_cpu/kernelentry/__init__.py @@ -0,0 +1,244 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""What it takes for this simulator to fully support one kernel, as a computed set. + +"Fully supported" is otherwise an adjective, and an adjective cannot be checked. +Here it is a **ledger**: a set of claims derived from the kernel itself, each one +either closed, or open with the file that would close it named. The set is not a +fixed checklist — every distinct op in the kernel contributes claims, every output +tensor contributes claims — so it grows with the kernel rather than with this +module. + +A contributor's loop is:: + + python -m ktir_cpu.kernelentry probe examples/latency/my_kernel.py + python -m ktir_cpu.kernelentry adopt examples/latency/my_kernel.py + python -m ktir_cpu.kernelentry verify --all + +``probe`` is read-only and answers "what does the simulator not support about my +kernel" before any work starts. ``adopt`` writes only the files it owns and never +edits the interpreter: when a new handler or a repriced op is needed it prints the +edit for a human to apply. ``verify`` is the gate, shared with +``tests/test_kernelentry.py`` so CI enforces the same engine. + +Two idioms here are the repository's own: a waiver mapping whose every entry +carries a reason (``tests/mlir_frontend/test_registry_consistency.py``), and +``xfail(strict=True)`` so a gap that closes fails the build until the excuse for it +is removed (``tests/test_spec_gaps.py``). +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Dict, Optional + +from .tensorspec import build_tensors, validate_specs + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent +EXAMPLES_DIR = REPO_ROOT / "examples" + +# --------------------------------------------------------------------------- +# Claim states +# --------------------------------------------------------------------------- + +CLOSED = "closed" +OPEN = "open" +UNDETERMINED = "undetermined" +DEFERRED = "deferred" +WAIVED = "waived" +SKIP = "skip" + +#: States that keep the ledger from being clean, i.e. that fail ``verify``. +BLOCKING = (OPEN, UNDETERMINED) + +# ``undetermined`` exists because its absence is invisible. A claim the engine +# cannot evaluate — an output tensor it failed to identify, a figure it could not +# resolve — must not be omitted, because an omitted claim reads exactly like a +# closed one in any summary. It is distinct from ``waived`` (never applies to +# this kernel) and from ``skip`` (this machine lacks a dependency; CI decides). + +_STATE_RANK = {CLOSED: 0, WAIVED: 1, SKIP: 2, DEFERRED: 3, UNDETERMINED: 4, OPEN: 5} + +FUNCTION, COST = "function", "cost" + + +@dataclass(frozen=True) +class Claim: + """One checkable assertion about one kernel. + + ``closer`` names what would move the claim to ``closed`` — a file, a registry, + a missing declaration field. It is the difference between a report that says + a kernel is unsupported and a report that says what to do about it. + """ + + id: str + leg: str + state: str + detail: str = "" + closer: str = "" + #: True when the check needs an optional dependency, so its state is a fact + #: about the machine as much as about the kernel. The committed support report + #: renders these to a fixed value: a document compared verbatim cannot depend + #: on which environment generated it, or CI and a laptop disagree forever. + env_dependent: bool = False + + @property + def blocking(self) -> bool: + return self.state in BLOCKING + + def sort_key(self) -> tuple: + return (-_STATE_RANK[self.state], self.id) + + +# --------------------------------------------------------------------------- +# Entry declaration +# --------------------------------------------------------------------------- + +@dataclass +class KernelEntry: + """The single declaration a kernel needs in order to be gated. + + *path* is the kernel's ``.mlir``, relative to ``examples/``, and it is the + kernel's source — hand-written IR, or captured compiler output such as + ``examples/triton-ktir/``, which is "kernels as the Triton -> KTIR path emits + them". + + *gate_params* is deliberately a reduced shape: gate cost is driven by shape and + not by the number of entries, which is why ``examples/latency/`` holds reduced + sizes in the first place. A ``cost.*`` claim evaluated there checks the + *composition* of a kernel's cost — which tensor dominates — and not the + absolute figure at full size. + + ``waived`` and ``deferred`` both excuse a claim; the difference is whether + anything will ever come back for it. A waived claim never applies here and its + value is the reason. A deferred claim does apply and is not met yet: its value + is an issue reference, the claim runs as ``xfail(strict=True)``, and closing the + gap fails the build until the deferral is removed. That is what lets the second + of a kernel's two PRs carry the cost leg without the first one having to lie. + """ + + name: str + func: str + path: str + + gate_params: Dict[str, Any] = field(default_factory=dict) + + #: ``{arg_name: spec}`` — the arguments to call the kernel with, as data. + #: See ``ktir_cpu/kernelentry/tensorspec.py`` for the vocabulary and for the + #: per-argument escape hatch. Empty means the kernel cannot be driven, which + #: ``exec.runs`` reports rather than skipping. + tensors: Dict[str, Any] = field(default_factory=dict) + #: ``reference(params, tensors) -> {arg_name: ndarray}``. Must compute in + #: f32 or wider: a reference evaluated in f16 reproduces the very overflow it + #: is supposed to catch, and then agrees with the kernel about a wrong answer. + #: *tensors* is a pristine rebuild, not what the run left behind. + reference: Optional[Callable[..., Dict[str, Any]]] = None + + waived: Dict[str, str] = field(default_factory=dict) + deferred: Dict[str, str] = field(default_factory=dict) + #: ``{arg_name: (rtol, atol)}`` — a wider pair than the ledger's default for one + #: output, because f16 error is set by the magnitude of a dot product's *terms* + #: and not of its result. Loosening a tolerance weakens the claim, so it belongs + #: in the row where a reviewer reads the reason beside it, not in the engine. + tolerance: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not str(self.path).strip(): + raise ValueError( + f"{self.name}: path= is the kernel's .mlir, relative to " + "examples/, and there is nothing to check without it" + ) + for mapping, what in ((self.waived, "waived"), (self.deferred, "deferred")): + for claim_id, reason in mapping.items(): + if not str(reason).strip(): + raise ValueError( + f"{self.name}: {what}[{claim_id!r}] needs a reason. An " + "excuse with no reason is indistinguishable from an " + "oversight." + ) + for claim_id, reason in self.deferred.items(): + # The issue is the whole of a deferral's promise: the report groups + # deferrals by issue, so a closed issue with an open gap stays visible. + if not re.search(r"#\d+", str(reason)): + raise ValueError( + f"{self.name}: deferred[{claim_id!r}] does not name an issue " + f"as #N: {reason!r}. Use waived= for a check that will never " + "apply here; a deferral has to say where it is tracked." + ) + for name, pair in self.tolerance.items(): + # The ledger reads this mapping with .get(), so a misspelled argument + # name is indistinguishable from a considered widening that took effect. + if name not in self.tensors: + raise ValueError( + f"{self.name}: tolerance[{name!r}] names no declared tensor " + f"(have {sorted(self.tensors)}). A tolerance for an argument " + "that does not exist is silently never applied." + ) + try: + rtol, atol = (float(v) for v in pair) + except (TypeError, ValueError): + raise ValueError( + f"{self.name}: tolerance[{name!r}] must be an (rtol, atol) " + f"pair of numbers, not {pair!r}" + ) from None + if rtol < 0 or atol < 0: + raise ValueError( + f"{self.name}: tolerance[{name!r}] = {pair!r} is negative, " + "which no comparison can satisfy" + ) + validate_specs(self.tensors, self.gate_params, self.name) + + def build_tensors(self, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: + """The keyword arguments for one run, rebuilt from the declared specs.""" + return build_tensors(self.tensors, + self.gate_params if params is None else params) + + @property + def mlir_path(self) -> Path: + """Absolute path of this kernel's ``.mlir``.""" + return EXAMPLES_DIR / self.path + + def mlir_text(self) -> str: + """The kernel's MLIR, as committed.""" + return self.mlir_path.read_text() + + +# --------------------------------------------------------------------------- +# Registry +# --------------------------------------------------------------------------- + +_ENTRIES: Dict[str, KernelEntry] = {} + + +def register_entry(entry: KernelEntry) -> KernelEntry: + """Add *entry* to the registry that ``probe --all`` and the gate iterate.""" + if entry.name in _ENTRIES: + raise ValueError(f"duplicate kernelentry name {entry.name!r}") + _ENTRIES[entry.name] = entry + return entry + + +def registered() -> Dict[str, KernelEntry]: + """Every declaration discovered so far, by name.""" + return dict(_ENTRIES) + + +__all__ = [ + "BLOCKING", "CLOSED", "COST", "Claim", "DEFERRED", "EXAMPLES_DIR", "FUNCTION", + "KernelEntry", "OPEN", "REPO_ROOT", "SKIP", "UNDETERMINED", "WAIVED", + "register_entry", "registered", +] diff --git a/ktir_cpu/kernelentry/conformance.py b/ktir_cpu/kernelentry/conformance.py new file mode 100644 index 0000000..b248b3a --- /dev/null +++ b/ktir_cpu/kernelentry/conformance.py @@ -0,0 +1,73 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Kernels under ``examples/`` that the MLIR frontend does not accept. + +This is a record, not a verdict. ``parse.frontend`` can only be *decided* where the +MLIR bindings are installed, so on most machines it reads ``skip`` — and a committed +support report that printed the outcome it saw would be a document about the machine +that generated it. Writing the outcome down here instead makes the report a function +of the repository, and ``tests/mlir_frontend/test_kernelentry_adapt.py`` holds this +mapping to the truth in the one environment that can check it. Same shape and same +reason as ``FRONTEND_UNSUPPORTED`` in +``tests/mlir_frontend/test_registry_consistency.py``: a committed mapping of known +gaps, each with a reason, verified rather than trusted. + +What this mapping deliberately does not say is which side is wrong. Unlike +``FRONTEND_UNSUPPORTED``, which covers ops the dialect does not define at all, these +are defined ops, and the disagreement is a typing one rather than a spelling one. +Gap row 2a in ``docs/gap_analysis.md`` carries it; ``inter_tile_produce`` / +``inter_tile_reduce`` are not in RFC 0682, so the specification does not adjudicate +it either. +""" + +from __future__ import annotations + +from typing import Dict + +# One gap, reached through three kernels, so the reason is named once rather than +# repeated per path: all three write the inter-tile ops in the form only the regex +# parser reads +# : T -> !ktdp.tile_future +# rather than the dialect's own +# -> <(T), groups = S> (produce) +# : <(T), groups = S> -> R (reduce) +# so the frontend stops at the first inter-tile op with `expected '->'`. Rewriting +# only that spelling moves the error rather than removing it: all three reduce +# `tensor<1x128xf16>` to `tensor<128xf16>`, and the dialect verifies that a reduce +# result matches the future's partial type, so it then fails with `result types must +# match future partial types`. Closing that means deciding whether the reduce should +# reshape at all, which is gap row 2a's question, not this mapping's. +_RESHAPING_REDUCE = ( + "inter-tile ops in the form only the regex parser reads, and a reduce that " + "reshapes its result, which the dialect's type relation does not express " + "(gap row 2a)" +) + +#: Repository-relative kernel path -> why the frontend rejects it. +#: +#: A kernel belongs here only if it has no declaration. A declared kernel records +#: the same fact as a ``deferred`` or ``waived`` ``parse.frontend`` claim, and two +#: places recording one fact is how they come to disagree. +FRONTEND_REJECTS: Dict[str, str] = { + "examples/ktir/ring_reduce.mlir": _RESHAPING_REDUCE, + "examples/ktir/ring_reduce_inner_loop.mlir": _RESHAPING_REDUCE, + "examples/latency/ring_reduce_multi_group.mlir": _RESHAPING_REDUCE, + + # A different defect class, and it does not overlap with the one above: this + # kernel has no `coordinate_set` on `construct_memory_view` at all, which the + # dialect requires and the regex parser does not. + "examples/ktir/nested_yield.ktir": + "construct_memory_view has no coordinate_set, which the dialect requires", +} diff --git a/ktir_cpu/kernelentry/pricing.py b/ktir_cpu/kernelentry/pricing.py new file mode 100644 index 0000000..87f6618 --- /dev/null +++ b/ktir_cpu/kernelentry/pricing.py @@ -0,0 +1,315 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Splitting "this op is free" from "nobody priced this op". + +``@register()`` defaults ``latency_category`` to ``"zero"``, so the registry holds +one state where there are two facts: an op the hardware really does no measurable +work for, and an op whose cost nobody has decided yet. A kernel leaning on the +second reports a lower cost than the hardware would, and every figure the cost leg +prints is read out of that registry — which makes a *clean* cost report +untrustworthy rather than merely incomplete. + +This module is the committed record that splits them. Every op priced ``zero`` +must appear in exactly one of two mappings: + +``ZERO_COST_OPS`` + Free by decision, with the reason the hardware does no measurable work. + +``UNJUDGED_ZERO_OPS`` + Not judged yet, with what is unresolved about it and the issue tracking the + decision. :func:`audit` accepts these — an open question that is written + down is not the failure mode being guarded against. + +:func:`audit` fails on an op in **neither**, which is what stops a newly +registered op from being silently free. It also fails on the three ways the +record can rot: an op in both mappings, an entry whose op has since been priced, +and an entry naming an op that is no longer registered. All four are one-line +edits to this file; none of them touches a cost formula. + +**Repository-wide, not per kernel** — measured, not assumed. A per-kernel +prototype of this check over six real artifacts produced 74 false positives, 66 of +them here: about fifteen of the ops it flagged are the same structural ones in +every kernel (``arith.constant``, ``ktdp.construct_memory_view``, ``scf.for``), +genuinely free. Per kernel it is a chore each contributor waives their way +through, and a waiver mapping that is mostly noise stops carrying information. +So it is asked once of the registry, and no kernel declaration mentions it. + +**What this check cannot see: a wrong category.** It compares against ``zero``, +so an op priced in the wrong non-zero class passes it — an integer compare billed +to the float pipe is invisible here. Only a reader who knows the op's semantics +catches that, which is why the reasons below are written for a reader rather than +being generated. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List + +import ktir_cpu.dialects # noqa: F401 — import triggers @register side effects +from ktir_cpu.dialects import registry + +# --------------------------------------------------------------------------- +# Free by decision +# --------------------------------------------------------------------------- + +#: Ops that cost nothing because of what they *are*, with the reason each one +#: does no measurable work. Four kinds appear here and it is worth reading them +#: as four: a value that exists at compile time, a terminator that only names +#: values, addressing metadata that computes where data is without moving it, and +#: an orchestrator whose body is charged op by op. +ZERO_COST_OPS: Dict[str, str] = { + # -- compile-time values ------------------------------------------------ + "arith.constant": + "a compile-time literal, and registered `no_lx_charge=True` for the same " + "reason: the scratchpad is charged when a consumer materializes the " + "value into a working tile, not here", + + # -- terminators: they name values leaving a region, and issue nothing --- + "func.return": "a terminator; it names the values leaving the function", + "return": "the regex parser's spelling of `func.return`; same reason", + "linalg.yield": "a region terminator; it names the combiner's result", + "scf.yield": "a region terminator", + "tensor.yield": "a region terminator, for a `tensor.generate` body", + "ktdp.yield_partial": "a region terminator; it names this core's partial", + "ktdp.yield_reduced": "a region terminator; it names the reduced value", + "region.bb0_args": + "not an op at all — the parser's record of a region's block arguments, " + "which the enclosing op's handler binds", + + # -- addressing: where the data is, not the data ------------------------- + "ktdp.construct_memory_view": + "computes an address, moves nothing; the `ktdp.load` or `ktdp.store` " + "that reads through the view is what carries the bytes", + "ktdp.construct_distributed_memory_view": + "the same address computation, per partition; the traffic is charged on " + "the loads and stores that address it", + "ktdp.construct_access_tile": + "narrows an existing view to a tile's worth of it — index arithmetic on " + "the view, with no access performed", + "ktdp.construct_indirect_access_tile": + "the same narrowing with a gathered index set; the gather itself is " + "charged on the indirect load, whose figure covers both the data and " + "the index lookups", + "ktdp.get_compute_tile_id": + "reads the core's own coordinate in the grid, which is available to it " + "without a memory access", + + # -- the produce half of a cross-core reduce ----------------------------- + "ktdp.inter_tile_produce": + "publishes this core's partial to the scheduler's mailbox; the wire time " + "for the whole exchange is charged once, as `comm`, on the matching " + "`ktdp.inter_tile_reduce`. Pricing both would count one transfer twice", + + # -- orchestrators: the body is charged, op by op ------------------------ + "linalg.reduce": + "executes its combiner region rather than mapping to a fixed reduction, " + "so the arithmetic inside it is charged individually: one core's trace " + "for the RMSNorm generator carries 256 `arith.addf` entries totalling " + "1,280 cycles from inside a reduce. The orchestrator is free; the " + "arithmetic is not", + "tensor.generate": + "evaluates its region body per index, so the ops in the body are charged " + "individually — the same split as `linalg.reduce`", + "linalg.index": + "produces the iteration index inside a region body; the arithmetic that " + "consumes it is charged, again the `linalg.reduce` split", + "scf.for": + "loop control; the body's ops are charged once per iteration, so a cost " + "for the loop itself would be on top of the work it drives", + "scf.if": + "branch control; the ops of the taken branch are charged", + + # -- allocation without initialization ----------------------------------- + "tensor.empty": + "names an uninitialized buffer. No data moves, and a consumer that " + "writes into it pays for the write", +} + +# --------------------------------------------------------------------------- +# Not judged yet +# --------------------------------------------------------------------------- + +#: Ops priced ``zero`` where zero has not been decided, with what is unresolved. +#: Every entry names the issue that would settle it. ``audit`` accepts these: +#: the state being guarded against is an op nobody looked at, not a question +#: somebody wrote down. +#: +#: Most of these share one shape. ``_unary`` in ``ktir_cpu/dialects/_helpers.py`` +#: applies its function to a whole ``Tile`` when given one, so a cast that reads +#: like a scalar conversion in the IR is an elementwise pass over a tile at run +#: time — and an elementwise pass is what every op in ``compute_float`` is priced +#: for. Whether Spyre does measurable work for them cannot be settled from +#: RFC 0682: a conversion folded into the consumer's read is free, and a +#: materialized one is not. That is a hardware question, so these are decisions +#: to be made rather than a fix to be applied. +UNJUDGED_ZERO_OPS: Dict[str, str] = { + # -- tile-wide work, and it shows in kernels already on main (#211) ------ + "linalg.fill": + "writes a scalar across the whole `outs` tile. Eight kernels under " + "`examples/` use it. #211", + "linalg.broadcast": + "expands a tile along new dimensions — free if the consumer reads it " + "strided, not free if it is materialized. #211", + "arith.sitofp": + "integer to float across the whole tile, via `_unary`. #211", + + # -- the rest of the cast cluster: same question, cheaper to be wrong ---- + "arith.extf": "widens every element of a tile, via `_unary`. #211", + "arith.truncf": "narrows every element of a tile, via `_unary`. #211", + "arith.convertf": "converts every element of a tile, via `_unary`. #211", + "arith.extsi": "sign-extends every element of a tile. #211", + "arith.extui": "zero-extends every element of a tile, via `_unary`. #211", + "arith.trunci": "truncates every element of a tile, via `_unary`. #211", + "arith.fptosi": "float to signed integer across the tile, via `_unary`. #211", + "arith.fptoui": "float to unsigned integer across the tile, via `_unary`. #211", + "arith.uitofp": "unsigned integer to float across the tile, via `_unary`. #211", + "arith.bitcast": + "reinterprets a tile's bits under another type. Free if it is a type " + "relabel and not free if the data is copied; the handler uses " + "`ndarray.view`, which is the free reading. #211", + "arith.index_cast": + "zero is defensible — the handler returns a Python `int`, so this is one " + "scalar conversion rather than a tile's worth — but it has not been " + "decided. #211", + "arith.index_castui": + "the same one-scalar conversion as `arith.index_cast`, undecided for the " + "same reason. #211", + + # -- shape metadata, or a copy? ------------------------------------------ + "linalg.transpose": + "free if it is a stride permutation, not free if the data moves; the " + "handler calls `np.transpose(...).copy()`, which is the second reading. " + "#211", + "tensor.reshape": + "reinterprets the same elements under a new shape. Free as metadata, not " + "free if the layout is rebuilt. #211", + "tensor.expand_shape": "the same question as `tensor.reshape`. #211", + "tensor.collapse_shape": "the same question as `tensor.reshape`. #211", + "tensor.extract_slice": + "reads a strided sub-tensor. Free if the consumer reads the parent " + "strided, not free if the slice is materialized. #211", + "tensor.insert_slice": + "writes a sub-tensor into a destination, which is a copy of the slice's " + "worth of elements unless it folds into the producer. #211", + "tensor.splat": + "broadcasts one scalar across a whole tile — the same question as " + "`linalg.fill`, and priced the same way. #211", + "tensor.from_elements": + "builds a small tensor from N scalar operands, so its cost is N element " + "writes rather than a tile's worth. #211", + "tensor.extract": + "reads one element out of a tile. One scalar read, so zero is " + "defensible, but undecided. #211", + + # -- not a pricing question at all --------------------------------------- + "ktdp.coreid": + "not an op to price: it is not in the authoritative `ktdp` dialect and " + "survives only on the regex path, so the resolution is to reconcile or " + "remove it rather than to give it a category. #88", +} + + +# --------------------------------------------------------------------------- +# The audit +# --------------------------------------------------------------------------- + +#: An op priced ``zero`` and listed in neither mapping. The state AC7 exists to +#: reject: registering an op without a category makes it free, and nothing said so. +UNLISTED = "unlisted" +#: An op in both mappings — the record contradicts itself about whether the +#: question is settled. +BOTH = "both" +#: An op with a real category that is still listed here. Someone priced it and +#: the entry outlived the decision; the fix is to delete the line. +PRICED = "priced" +#: An entry naming an op no longer in the registry, i.e. a rename or a removal +#: that left the record behind. +UNREGISTERED = "unregistered" + + +@dataclass(frozen=True) +class Finding: + """One way the pricing record and the registry disagree. + + ``fix`` names the edit rather than describing the problem, for the same + reason a ``Claim`` carries a ``closer``: the difference between a report that + says something is wrong and a report that says what to do about it. + """ + + op: str + kind: str + detail: str + fix: str + + def __str__(self) -> str: + return f"{self.op}: {self.detail}\n {self.fix}" + + +def zero_priced_ops() -> List[str]: + """Every registered op whose ``latency_category`` is ``zero``.""" + return sorted(op for op in registry._REGISTRY + if registry.get_latency_category(op) == "zero") + + +def audit() -> List[Finding]: + """Compare the record above against the registry. Empty means agreement. + + Both directions are checked, because only one of them is the interesting + one and the other is the one that rots. A zero-priced op missing from the + record is the defect this exists for; an entry that outlived its op, or its + ``zero``, is how the record stops meaning anything. + """ + findings: List[Finding] = [] + zero = set(zero_priced_ops()) + + for op in sorted(zero - set(ZERO_COST_OPS) - set(UNJUDGED_ZERO_OPS)): + findings.append(Finding( + op, UNLISTED, + "priced `zero` and listed in neither mapping, so its cost is the " + "`@register()` default rather than a decision", + "add it to ZERO_COST_OPS with the reason the hardware does no " + "measurable work for it, or to UNJUDGED_ZERO_OPS with what is " + "unresolved and the issue tracking it — both in " + "ktir_cpu/kernelentry/pricing.py", + )) + + for op in sorted(set(ZERO_COST_OPS) & set(UNJUDGED_ZERO_OPS)): + findings.append(Finding( + op, BOTH, + "listed as free by decision and as not yet judged at the same time", + "keep one of the two entries in ktir_cpu/kernelentry/pricing.py", + )) + + for mapping, which in ((ZERO_COST_OPS, "ZERO_COST_OPS"), + (UNJUDGED_ZERO_OPS, "UNJUDGED_ZERO_OPS")): + for op in sorted(mapping): + if op not in registry._REGISTRY: + findings.append(Finding( + op, UNREGISTERED, + f"listed in {which} but no longer has a `@register` handler", + f"remove it from {which} in ktir_cpu/kernelentry/pricing.py, " + "or restore the handler", + )) + elif op not in zero: + category = registry.get_latency_category(op) + findings.append(Finding( + op, PRICED, + f"listed in {which} but priced `{category}`, so the entry " + "outlived the decision it records", + f"remove it from {which} in ktir_cpu/kernelentry/pricing.py", + )) + + return findings diff --git a/ktir_cpu/kernelentry/tensorspec.py b/ktir_cpu/kernelentry/tensorspec.py new file mode 100644 index 0000000..3979cb3 --- /dev/null +++ b/ktir_cpu/kernelentry/tensorspec.py @@ -0,0 +1,278 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The arguments a kernel is called with, written as data rather than as code. + +Every kernel in this repository is driven by the same few tensor shapes: a normal +draw, a zeroed output, a broadcast constant, a tiled row, an index vector, and the +scalars the signature declares. Written as a function per kernel, that vocabulary +was re-spelled eighteen times and each spelling could differ in ways nothing +checked — the seed, the dtype, whether the output was zeroed. Written as a table, +declaring a kernel is a row, and the ways two kernels can differ are the ways the +vocabulary allows. + +A value in an entry's ``tensors`` mapping is one of four things: + +* a **spec** from this module — ``normal``, ``zeros``, ``full``, ``tile``, + ``arange``, ``integers``, ``asarray``, ``param``; +* a **string**, naming a parameter to forward to the kernel unchanged; +* any other **literal**, forwarded as it stands; +* a **callable** ``(params, rng) -> value``, for input a spec cannot express. + +The last is the escape hatch, and it is per argument rather than per kernel: RoPE +needs its cos/sin tables built from angles, and gets to declare those two the long +way while its other two arguments stay rows in the table. + +Shapes are resolved against the entry's parameters: an ``int`` is itself, a +``str`` is a parameter name, and a callable is evaluated on the parameters for a +dimension that is an expression of them. + +**Determinism is a requirement, not a convenience.** The reference comparison and +the committed cost derivation each rebuild the arguments, and neither may depend on +which ran first, so every random draw comes from a generator seeded by the argument +name. Seeding per argument rather than per kernel is also what keeps two arguments +of the same kernel independent: one generator drawn twice gives the second tensor +the first one's values wherever their shapes overlap, and a kernel that swapped its +two operands would then still agree with its reference. +""" + +from __future__ import annotations + +import math +import zlib +from dataclasses import dataclass +from typing import Any, Callable, Dict, Mapping, Sequence, Tuple, Union + +import numpy as np + +#: One seed for the whole ledger. Which value it is does not matter; that it is +#: written down once, and that no declaration gets to choose its own, does. +SEED = 42 + +_DTYPES = { + "f16": np.float16, "f32": np.float32, "f64": np.float64, + "i32": np.int32, "i64": np.int64, +} + +Dim = Union[int, str, Callable[[Mapping[str, Any]], int]] +Shape = Union[Dim, Sequence[Dim]] + + +def _rng(name: str) -> np.random.Generator: + """A generator for one argument, stable across runs, hosts and orderings.""" + return np.random.default_rng([SEED, zlib.crc32(name.encode("utf-8"))]) + + +def dtype_of(name: str) -> np.dtype: + """The numpy dtype a spec's ``dtype=`` string names.""" + try: + return np.dtype(_DTYPES[name]) + except KeyError: + raise ValueError( + f"unknown dtype {name!r}; use one of {', '.join(sorted(_DTYPES))}" + ) from None + + +def _scalar(value: Any, params: Mapping[str, Any]) -> Any: + """A spec's own scalar argument: a parameter name, or a literal.""" + return params[value] if isinstance(value, str) else value + + +def _shape(shape: Shape, params: Mapping[str, Any]) -> Tuple[int, ...]: + dims = shape if isinstance(shape, (tuple, list)) else (shape,) + out = [] + for dim in dims: + if isinstance(dim, str): + out.append(int(params[dim])) + elif callable(dim): + out.append(int(dim(params))) + else: + out.append(int(dim)) + return tuple(out) + + +#: ``scale=`` for a weight matrix initialised the way the layer would be, which +#: keeps every f16 intermediate downstream of it inside range. Named rather than +#: written out because the reason for it is the same wherever it appears. +FAN_IN: Callable[[Tuple[int, ...]], float] = lambda shape: 1.0 / math.sqrt(shape[0]) + + +@dataclass(frozen=True) +class normal: + """A standard normal draw, optionally scaled.""" + + shape: Shape + dtype: str = "f16" + scale: Union[float, Callable[[Tuple[int, ...]], float]] = 1.0 + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + shape = _shape(self.shape, params) + scale = self.scale(shape) if callable(self.scale) else self.scale + return (rng.standard_normal(shape) * scale).astype(dtype_of(self.dtype)) + + +@dataclass(frozen=True) +class zeros: + """A zeroed output tensor.""" + + shape: Shape + dtype: str = "f16" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + return np.zeros(_shape(self.shape, params), dtype=dtype_of(self.dtype)) + + +@dataclass(frozen=True) +class full: + """One value everywhere. *value* may name a parameter.""" + + shape: Shape + value: Any + dtype: str = "f16" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + return np.full(_shape(self.shape, params), _scalar(self.value, params), + dtype=dtype_of(self.dtype)) + + +@dataclass(frozen=True) +class tile: + """*row*, repeated — a vector a kernel views at matrix shape. + + The repetition is what makes the reference able to say the tensor is + position-dependent along one axis and constant along the other. + """ + + row: Any + reps: Shape + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + return np.tile(self.row(params, rng), _shape(self.reps, params)) + + +@dataclass(frozen=True) +class arange: + """Consecutive values over *shape*, so a fold's result is hand-checkable.""" + + shape: Shape + start: Any = 0 + dtype: str = "f16" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + shape = _shape(self.shape, params) + start = int(_scalar(self.start, params)) + size = int(np.prod(shape)) if shape else 0 + return np.arange(start, start + size, + dtype=dtype_of(self.dtype)).reshape(shape) + + +@dataclass(frozen=True) +class integers: + """A draw from ``[low, high)`` — an index tensor, not a value tensor. + + *low* and *high* may name parameters. Drawn over the whole range rather than + set to identity on purpose: indices that happen to equal their position make + an indirect access indistinguishable from a direct one. + """ + + shape: Shape + high: Any + low: Any = 0 + dtype: str = "i32" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + return rng.integers(int(_scalar(self.low, params)), + int(_scalar(self.high, params)), + size=_shape(self.shape, params), + dtype=dtype_of(self.dtype)) + + +@dataclass(frozen=True) +class asarray: + """A literal sequence, or a parameter holding one, as a tensor.""" + + value: Any + dtype: str = "i64" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> np.ndarray: + return np.asarray(_scalar(self.value, params), dtype=dtype_of(self.dtype)) + + +@dataclass(frozen=True) +class param: + """A scalar argument, at a width the kernel reads it as. + + A bare string in the mapping forwards a parameter as the Python value it is; + this is for the arguments where the width is part of the signature — an ``i32`` + extent is not an index constant, and the interpreter reads the two differently. + """ + + name: str + dtype: str = "" + + def __call__(self, params: Mapping[str, Any], + rng: np.random.Generator) -> Any: + value = params[self.name] + return dtype_of(self.dtype).type(value) if self.dtype else value + + +def build_tensors(specs: Mapping[str, Any], + params: Mapping[str, Any]) -> Dict[str, Any]: + """The keyword arguments to call the kernel with. + + Called once per run and again for whatever needs a pristine copy: a kernel + whose output argument aliases its input writes over what it was given, and a + reference reading that would be comparing the result against itself. + """ + built: Dict[str, Any] = {} + for name, spec in specs.items(): + if isinstance(spec, str): + built[name] = params[spec] + elif callable(spec): + built[name] = spec(params, _rng(name)) + else: + built[name] = spec + return built + + +def validate_specs(specs: Mapping[str, Any], params: Mapping[str, Any], + where: str) -> None: + """Reject a mapping that cannot be built, at declaration time. + + A parameter name misspelled in a spec would otherwise surface as a ``KeyError`` + from inside the engine, on the one kernel, at the moment it ran — which reads + as a fault in the tool rather than a typo in the row. + """ + for name, spec in specs.items(): + wanted = [spec] if isinstance(spec, str) else [] + if isinstance(spec, param): + wanted = [spec.name] + for key in wanted: + if key not in params: + raise ValueError( + f"{where}: tensors[{name!r}] names parameter {key!r}, which " + f"is not in gate_params ({', '.join(sorted(params))})" + ) + + +__all__ = ["FAN_IN", "SEED", "arange", "asarray", "build_tensors", "dtype_of", + "full", "integers", "normal", "param", "tile", "validate_specs", + "zeros"] From 9e577f4726585c091498ac89b3ea18b4e54a3f07 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:15 -0400 Subject: [PATCH 3/9] Decide kernel support as a computed ledger MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit "Fully supported" was an adjective, and an adjective cannot be checked. Three kernels have now gone through this repository, each as a kernel plus a notebook section, and what they delivered under `tests/` varies sevenfold across the commits that carry them: 471, 363, 220, 84 and 66 lines. The 84 asserts an arithmetic-intensity formula, per-core cycle structure and three input guards, and compares no output value. Two simulator gaps were found by walking into them mid-change rather than by asking first. Closes #209. kernelentry makes support a ledger: a set of claims derived from the kernel itself, each closed or open with the file that would close it named. The set is not a checklist — every distinct op contributes a claim, every output tensor contributes two — so it grows with the kernel rather than with this module. python -m ktir_cpu.kernelentry probe examples/latency/my_kernel.mlir python -m ktir_cpu.kernelentry adopt my_kernel python -m ktir_cpu.kernelentry probe --all --write-report python -m ktir_cpu.kernelentry verify --all `probe` takes the bare `.mlir` and writes nothing, so the first question — what does the simulator not support about this kernel — can be asked before there is a declaration to ask it with. `adopt` writes only the files it owns and never edits the interpreter: when a handler or a repriced op is needed it prints the change and the file, because a tool that quietly changed what the simulator charges would be changing the answer it was asked for. `verify` shares one engine with the gate, so the local loop and CI cannot disagree. Support is five questions asked in order — read by the regex parser, read by the MLIR frontend, runs, output right, cost pinned — and a kernel with no declaration answers a prefix of them. That ordering is why the generated report has two tables rather than one with holes: an unasked question has no cell, not an empty one. Output tensors are identified from the trace rather than by walking SSA names back from a store. The walk does not survive `construct_distributed_memory_view`, and when it fails it fails silently: the prototype produced a kernel with no output claims at all and a summary line reading 60/61 closed. Deriving the set from the trace also means a declaration can name an output the trace does not, so a reference entry no store wrote is reported `undetermined` at the id its comparison would have had. Left unasked it is the same defect one step over: a comparison that vanishes rather than fails. There is deliberately no per-op frontend-reachability claim. `tests/mlir_frontend/test_registry_consistency.py` already asserts that every executor op is frontend-installed or allow-listed, and an op appearing in a kernel is necessarily registered, so an unreachable op already breaks the build repository-wide; `parse.frontend` catches it for one kernel specifically. A third check would restate them and would need the allow-list moved out of that test. `cost.derivation` commits a generated attribution of the kernel's cost, so a change arrives as a diff naming the term that moved and a reviewer reads a breakdown instead of a number. It cannot catch the cost model being wrong, since it is produced by the same code as the measurement, and it says so in its own footer. That question is not asked per kernel: there is one cost model behind all thirty-three kernels, and `tests/test_latency.py` already asks it by mechanism, with hand-counted bytes, FLOPs and cycles per latency category and across the hardware parameters that scale them. A per-declaration hand count would have each kernel re-answer one question about one model, so support is five ordered questions and not six. Reading `examples/` for discovery, and the frontend allow-list out of the test that enforces it, makes this package a development tool rather than part of the library, so `pyproject.toml` excludes it from the wheel. Packaged it would resolve to a `site-packages` directory carrying neither path, and fail where it read them rather than at import. Editable installs, which is how the repository is worked in, are unaffected. `docs/kernelentry.md` is the contributor entry point — the five questions, the six states, how to declare a kernel, how to excuse a claim, and which kernels cannot be declared at all — and `CONTRIBUTING.md` links it from the pull-request section. Signed-off-by: WarningRan --- CONTRIBUTING.md | 11 + docs/kernelentry.md | 309 +++++++++++ ktir_cpu/kernelentry/__main__.py | 22 + ktir_cpu/kernelentry/cli.py | 803 +++++++++++++++++++++++++++++ ktir_cpu/kernelentry/derivation.py | 308 +++++++++++ ktir_cpu/kernelentry/ledger.py | 607 ++++++++++++++++++++++ ktir_cpu/kernelentry/ops.py | 364 +++++++++++++ pyproject.toml | 4 + 8 files changed, 2428 insertions(+) create mode 100644 docs/kernelentry.md create mode 100644 ktir_cpu/kernelentry/__main__.py create mode 100644 ktir_cpu/kernelentry/cli.py create mode 100644 ktir_cpu/kernelentry/derivation.py create mode 100644 ktir_cpu/kernelentry/ledger.py create mode 100644 ktir_cpu/kernelentry/ops.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index dfedda4..7dcb6e9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -21,6 +21,17 @@ once wheels are available. uv run pytest -v ``` +## Adding a kernel + +If your change adds a KTIR kernel, see [docs/kernelentry.md](docs/kernelentry.md) — +the path from "here is a kernel" to "the simulator supports it", and what the machine +checks along it versus what only a person can. Start before you write anything, by +pointing the tool at the IR: + +```bash +uv run python -m ktir_cpu.kernelentry probe examples/latency/my_kernel.mlir +``` + ## Pull Requests 1. Fork the repository and create a feature branch. diff --git a/docs/kernelentry.md b/docs/kernelentry.md new file mode 100644 index 0000000..b8c4b47 --- /dev/null +++ b/docs/kernelentry.md @@ -0,0 +1,309 @@ +# Adding a kernel + +You have a KTIR kernel and you want `ktir_cpu` to support it — to run it, agree +that it computes the right thing, and report a cost you can trust. "Fully +supported" is otherwise an adjective, and an adjective cannot be checked. Here it +is a **ledger**: claims derived from your kernel, each closed or open with the +file that would close it named. + +Sections 1 to 5 are five commands, one each. What follows them is reference. + +```bash +uv run python -m ktir_cpu.kernelentry probe examples/latency/my_kernel.mlir # 1 +$EDITOR examples/entries.py # 2 +uv run python -m ktir_cpu.kernelentry adopt my_kernel # 3 +uv run python -m ktir_cpu.kernelentry probe --all --write-report # 4 +uv run python -m ktir_cpu.kernelentry verify --all # 5 +``` + +## 1. Probe, before you write anything + +```bash +uv run python -m ktir_cpu.kernelentry probe examples/latency/my_kernel.mlir +``` + +Point it at the `.mlir` itself; add `--func` if the file declares more than one +function. A path is read cold — no declaration is consulted even if one exists — so +this step is available before step 2 rather than after it. Once the kernel is +declared, name it instead: `probe my_kernel`. + +It answers the questions that need no declaration: which of your kernel's ops have +no execution handler, and whether it survives both parse paths. Anything that needs +to *run* the kernel reports that there is no declaration yet. `probe` writes +nothing. + +## 2. Add a row + +Declaring a kernel is one entry in the `ENTRIES` table at the bottom of +`examples/entries.py`. There is no new file to create, nothing to register and no +test to write — the file is found by name at any depth under `examples/`, so a +vendored kernel set carries its own `entries.py`: + +```python +KernelEntry( + name="my_kernel", + func="my_kernel", # the func.func name, not the file name + path="latency/my_kernel.mlir", # relative to examples/ + gate_params={"M": 16, "N": 64, "K": 64}, + tensors={ + "a_ptr": normal(("M", "K")), # the arguments, as data + "b_ptr": normal(("K", "N")), + "c_ptr": zeros(("M", "N")), + "K": "K", # a bare string forwards a parameter + }, + reference=matmul_reference, # (params, tensors) -> {arg: ndarray} +), +``` + +`matmul_small` in that table is this row filled in. Four of the fields have a rule +behind them. + +**`tensors` is a mapping, not a function.** The specs — `normal`, `zeros`, `full`, +`tile`, `arange`, `integers`, `asarray`, `param` — are declared in +`ktir_cpu/kernelentry/tensorspec.py` and resolved against `gate_params`: an `int` +dimension is itself, a `str` names a parameter, and a callable is evaluated on the +parameters for a dimension that is an expression of them. A bare string as a *value* +forwards that parameter to the kernel unchanged. Where no spec fits, a callable +`(params, rng) -> value` is the escape hatch, per argument rather than per kernel: +RoPE declares its cos/sin tables the long way and leaves its other two arguments as +rows. Every draw is seeded from the argument's own name, so two arguments are +independent — which is what makes a swapped operand observable — and a rebuild is +bit-identical, so the reference is handed a pristine rebuild rather than the arrays +the kernel has by then written over. + +**`path=` is the kernel, not a copy of it.** The `.mlir` under `examples/` *is* the +source — hand-written IR whose point is its exact shape (`examples/ktir/`, +`examples/rfc/`) or captured compiler output (`examples/triton-ktir/`). + +**`gate_params` small.** CI runs it on every push, so a `cost.*` claim is evaluated +at reduced size: it checks the *composition* of your kernel's cost, not the absolute +figure at full size. + +**`reference` computed in f32 or wider,** and written against the short form of the +answer rather than the kernel's own decomposition. Every value in a KTIR kernel is +f16, so an f16 reference reproduces the kernel's own rounding and then agrees with it +about a wrong answer; the ledger rejects one rather than comparing against it. A +reference that sums shards the way the kernel sums them checks the arithmetic while +assuming the decomposition, which is usually the part under test. No reference at all +leaves `out..reference` **open**, not skipped. `out..nontrivial` covers the +same f16 range from the other side: an input scale large enough to overflow can zero +an output outright, and no cost report would flag it. + +Two fields are for the cases the rules above do not fit: + +- `waived=` / `deferred=` — usually omitted, and no row declares either today. See + *When a claim will not close*. +- `tolerance={"c_ptr": (rtol, atol)}` — a wider pair than the ledger's default, for + one output. `ffn_swiglu` is the only row that declares one and its comment says + why. This is the one input to a claim that can be adjusted until it passes, so it + belongs in the row where the reason is read beside it, and the claim states the + pair even when it closes. + +## 3. Adopt + +```bash +uv run python -m ktir_cpu.kernelentry adopt my_kernel +``` + +It writes your kernel's section of `docs/kernel_cost.md` and leaves the others as +committed, so adopting one kernel does not depend on having run the rest. You commit +the result. It never edits the interpreter: a new op handler or a repriced op is +printed for you to apply, because a tool that quietly changed what the simulator +charges would be changing the answer you asked it for. + +You are not asked to write that section, only to read it — and so does the reviewer, +in the diff. Read it there rather than here: a copy of a figure in this document is a +figure nothing regenerates. Four things to look at: + +- **The composition**, not the totals: which tensor dominates, what the + `traffic_ratio` is, which unit the cycles land on. "The weight operand's traffic is + negligible" does not survive next to a row attributing most of the traffic to it. +- **What it cannot catch is the cost model being wrong.** The same code produces the + breakdown and the figure it is checked against, so a mis-charged op moves both. +- **``** is traffic whose origin could not be resolved — printed rather + than dropped, so the rows always sum to the total. An access landing in more than + one partition of a distributed view produces it. +- **`traffic_ratio` below 1 is not an error.** The denominator is the whole footprint + of the tensors you declared, so a kernel that indexes instead of sweeping moves + less: `examples/triton-ktir/indexed_add.mlir` reaches two of 128 slices, and its + ratio says so. + +## 4. Regenerate the reports + +```bash +uv run python -m ktir_cpu.kernelentry probe --all --write-report +``` + +Writes `docs/kernel_support.md` (per kernel) and `docs/supported_ops.md` (per op), +both committed on the same discipline as a lock file. Not `adopt`'s job: `adopt` acts +on one kernel, these are whole-repository views, and a partial one would claim the +kernels it omits are absent — so `--write-report` requires `--all`. + +Adding any kernel under `examples/` makes them stale whether or not you declared it, +so the omission appears in your own diff instead of being an absence. `verify` names +the command when it finds a document stale. + +## 5. The gate + +```bash +uv run python -m ktir_cpu.kernelentry verify --all +uv run pytest -q --ignore=tests/mlir_frontend +uv run pytest tests/mlir_frontend/ -q # skips without mlir_ktdp; CI runs it +``` + +`verify` shares one engine with `tests/test_kernelentry.py`, so the local loop and CI +cannot disagree. Two of the things it checks belong to the repository rather than to +any kernel — that the two generated reports are current, and that the op registry's +zero prices are all accounted for — so `verify --all` can fail with every declared +kernel clean. It says which of the two it is. + +The third command is the one a green local run hides. The regex parser validates +against no dialect, so a form only it accepts is one the dialect does not define, and +only the MLIR frontend notices. Without `mlir_ktdp` installed that suite skips at +module level and `parse.frontend` reads `skip`, so CI decides it (see README for +installing the bindings); `verify` counts those claims in its summary line instead of +calling the kernels fully supported. + +What it catches, it does not adjudicate: it reports that a form is one the dialect +does not define, not which side should change. Four kernels under `examples/` are in +that position, all undeclared and recorded with their reason in +`ktir_cpu/kernelentry/conformance.py`. If yours lands there, what to do splits on the +kind of disagreement. Where only the spelling differs, rewriting it is a rewrite and +nothing else — the two cross-core kernels declared here were two lines each, both +parsers accepting the result and the cost figures identical to the digit. Where the +dialect disagrees about *types*, closing the gap means deciding something about the +op, which `docs/gap_analysis.md` row 2a carries for the three `ring_reduce*` kernels; +for a declared kernel there, `deferred` against an issue is the honest state. + +And per `CLAUDE.md`: CC the maintainers in a comment on the pull request, plus the +issue author if it is linked to an issue. + +--- + +## What is checked, and what you supply + +Grouped by the question each claim answers, in the order the tool asks them. Those +questions are the columns of `docs/kernel_support.md`; reading splits in two there, +because the two parsers fail for different reasons and are fixed by different people. +Everything you write is in the last column. + +| question | claim | what it asserts | you supply | +|---|---|---|---| +| **is it read** | `parse.regex` | `KTIRInterpreter.load` accepts the kernel | — | +| | `parse.frontend` | the MLIR frontend accepts it and MLIR's own verifier passes | — | +| **does it run** | `op..handler` | every distinct op has an execution handler | — | +| | `exec.runs` | it executes on the declared grid without overflowing LX | `gate_params`, `tensors` | +| | `out.identified` | the engine could work out which tensors the kernel writes | — | +| **is the output right** | `out..reference` | each output matches an independent reference | `reference` | +| | `out..nontrivial` | no output is entirely zero or NaN | — | +| **is the cost trustworthy** | `cost.derivation` | the kernel's section of `docs/kernel_cost.md` matches what the model reports now | `adopt` writes it; you read it | + +## When a claim will not close + +| state | meaning | when the gap closes | +|---|---|---| +| `closed` | checked, passed | — | +| `open` | should hold, does not | — | +| `undetermined` | applies, but the engine could not evaluate it | — | +| `deferred` | known gap, tracked against an issue | reported as unnecessary, which **fails the build** | +| `waived` | never applies to this kernel | reported as unnecessary | +| `skip` | this environment lacks a dependency; CI decides | — | + +`open` and `undetermined` both block the gate. They are separate because a check that +*could not run* must not read as one that passed — with no such state an unevaluable +claim gets quietly omitted, which looks identical to a clean result. + +`deferred` and `waived` both excuse a claim; the difference is whether anything ever +comes back for it. Deferral is what lets a kernel arrive over two pull requests +without the first one having to pretend the cost leg is done. + +```python +deferred={"cost.derivation": "# — the cost leg lands in the follow-up"}, +waived={"out.y.nontrivial": "fully masked at this shape, so all-zero is correct"}, +``` + +- A deferral needs an issue reference, `#` followed by digits; only the format is + checked, because verifying more would make a local gate depend on the network. So a + deferral against a closed issue passes, and that is how this state rots — a gap that + outlives the issue tracking it is caught by a reader of `docs/kernel_support.md`, + which prints every deferral with its issue, and by nothing else. A waiver needs a + reason, and an empty one is rejected. +- Either excuse overrides any state but `closed`, a missing optional dependency + included, so an excused claim reads the same on every machine. Both are printed in + `docs/kernel_support.md`, so an accumulation is visible in the diff rather than only + in the code. +- **The claim no longer exists** → `waived.stale.`; usually a typo in the id, or a + leftover from a kernel that has since changed. +- **The claim now passes on its own** → `waived.unnecessary.` / + `deferred.unnecessary.`, and the check itself reports `closed`. An excuse that + outlives what made it necessary is misinformation. This is what makes a deferral + expire — not the `xfail(strict=True)` marker, which is applied from the claim's + current state and so is simply absent once the gap closes. + +## If your kernel needs a new op + +Price it, and pin the price in `tests/test_latency.py` rather than on your +declaration: it holds hand-counted bytes, FLOPs and cycles per latency category and +per hardware parameter that scales them, which is one question about the one cost +model every kernel here shares. + +`@register()` defaults `latency_category` to `"zero"`, so "this op is free" and +"nobody priced this op" are the same state in the registry, and a kernel leaning on +an unpriced op reports a lower cost than the hardware would with nothing saying so. +`ktir_cpu/kernelentry/pricing.py` splits the two apart. Every op the registry prices +`zero` has to appear in exactly one of two mappings, each entry carrying a written +reason: + +- **`ZERO_COST_OPS`** — free by decision. A terminator, a compile-time constant, an + address computation, an orchestrator whose body is priced op by op. +- **`UNJUDGED_ZERO_OPS`** — priced zero with nobody having decided that. Each reason + names the issue that would settle it, because an open question with no issue behind + it is how the list becomes permanent. + +An op in neither, or in both, is a finding: `verify` fails and `probe --all` prints it +with the fix. So registering an op without naming a category no longer passes +silently — but it does not become priced either, and that is the point. The gate asks +you to decide, and recording "not yet judged" against an issue is a legitimate answer. + +The check compares against `zero`, so it **cannot see a category that is simply +wrong** — an integer compare billed to the float pipe passed it for as long as it +existed, and only a reader who knows the op's semantics catches that. +`docs/supported_ops.md` carries the result as a **cost** column, printing the decision +rather than the registry's default. + +## What no claim can tell you + +The ledger decides whether the simulator supports your kernel, not whether your +kernel is the kernel you meant: whether the IR expresses the algorithm, whether the +grid is the one you want to model, whether `access_tile_order` is lexicographic with +the rightmost dimension innermost, whether overlapping coordinate sets in a +distributed view are constrained enough to have defined behaviour. Those are read by a +person, against RFC 0682 and the checklist in `CLAUDE.md`. + +## Where this repository stands + +18 of the 33 kernels under `examples/` are declared, every claim they raise is closed, +and none of the eighteen defers or waives one. `docs/kernel_support.md` is the current +state per kernel, and its second table is the other 15: what *has* been asked of them, +with the columns nobody asked absent rather than empty. + +None of those 15 is waiting on somebody to transcribe a reference. Thirteen take no +tensor arguments the `tensors=` mapping can express — the eight `examples/rfc/*` files +address memrefs at absolute HBM bases, and five more (the three ring-reduce kernels +and the two `rmsnorm_4core_*`) take raw HBM element indices with no shape attached, +which their tests supply by replacing `KTIRInterpreter._prepare_execution`. That is a +property of the argument convention rather than of those files, and it is the live one: +the two rmsnorm kernels are the most recent to arrive. The last two are deliberate: +`softmax_wide.mlir` overflows LX on purpose and its test asserts the exception, so +`exec.runs` failing is the kernel behaving; `nested_yield.ktir` is a reproducer written +down to the one op under test, so being minimal IR rather than dialect-valid IR is what +it is for. + +What the 18 do not buy is cost coverage at real shapes for free. Probing all of them +takes about 21 s and five kernels are 97% of it — `layernorm_fwd_ktir` alone is 9.4 s, +then `paged_attention`, `softmax_fwd_ktir`, `matmul_fwd_ktir` and `rope_fwd_4x2` — +while the remaining 13 come to well under a second together. Shape drives that, which +is the trade `gate_params` exists to make; for these five there is nothing to trade, +because their shapes are baked into the IR. They are here anyway, on the grounds that a +cost model checked only at reduced size is checked where the padding, the tail core and +the page table are not real yet. diff --git a/ktir_cpu/kernelentry/__main__.py b/ktir_cpu/kernelentry/__main__.py new file mode 100644 index 0000000..2d51664 --- /dev/null +++ b/ktir_cpu/kernelentry/__main__.py @@ -0,0 +1,22 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``python -m ktir_cpu.kernelentry``.""" + +import sys + +from .cli import main + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ktir_cpu/kernelentry/cli.py b/ktir_cpu/kernelentry/cli.py new file mode 100644 index 0000000..fd41c60 --- /dev/null +++ b/ktir_cpu/kernelentry/cli.py @@ -0,0 +1,803 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``probe`` / ``adopt`` / ``verify``. + + python -m ktir_cpu.kernelentry probe examples/latency/matmul_small.mlir + python -m ktir_cpu.kernelentry probe matmul_small + python -m ktir_cpu.kernelentry probe --all + python -m ktir_cpu.kernelentry adopt matmul_small + python -m ktir_cpu.kernelentry verify --all + +A kernel is addressed either way, and the difference is the point of the first +line: a ``.mlir`` path is read cold, with no declaration consulted even if one +exists, which is what makes the first question answerable before any paperwork. A +bare name is the declared entry in ``examples/entries.py``, with all five questions +in reach. + +``probe`` writes nothing. ``adopt`` writes only what it owns — one kernel's section +of the committed cost document — and never edits the interpreter: when a kernel +needs a new handler or an op repriced, it prints the edit for a person to make. A +tool that silently changes what the simulator charges would be changing the answer +to the question it is being asked. + +``verify`` shares :func:`ktir_cpu.kernelentry.ledger.probe` with +``tests/test_kernelentry.py``, so the local loop and the gate cannot disagree. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import re +from pathlib import Path +from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple + +from . import ( + BLOCKING, CLOSED, COST, DEFERRED, EXAMPLES_DIR, FUNCTION, OPEN, REPO_ROOT, + SKIP, UNDETERMINED, WAIVED, KernelEntry, registered, +) +from .conformance import FRONTEND_REJECTS +from .ledger import Ledger, probe +from .pricing import audit, zero_priced_ops + +_STATE_LABEL = { + CLOSED: "closed", OPEN: "OPEN", UNDETERMINED: "UNDETERMINED", + DEFERRED: "deferred", WAIVED: "waived", SKIP: "skip", +} + + +# --------------------------------------------------------------------------- +# Discovery +# --------------------------------------------------------------------------- + +#: Declarations already executed, so that loading is idempotent. These modules are +#: not put in ``sys.modules`` — they are loaded by file location, not by name — so +#: nothing else would stop a second call from executing the same file again. +_LOADED: set = set() + + +def load_declaration(path: Path) -> None: + """Execute one declaration file, registering whatever it declares. + + ``examples/`` is not an importable package — ``examples/triton-ktir`` is not + even a legal identifier for an ``import`` statement — so the file is loaded by + location. That constraint is also why the declarations are one table rather + than a module per kernel: a declaration under ``examples/`` cannot import a + sibling, so anything shared between two kernels has to live with them. + + Loading the same file twice is a no-op rather than an error. Two test modules + call :func:`discover_all` at import time, and ``pytest tests/`` collects both + into one session — without this, the second one re-executes every declaration + and ``register_entry`` rejects the duplicate name, taking down collection for + the whole run. The duplicate-name guard stays: two *different* files claiming + one name is the hazard it is there for. + """ + resolved = path.resolve() + if resolved in _LOADED: + return + spec = importlib.util.spec_from_file_location(f"_kernelentry_{path.stem}", path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load {path}") + module = importlib.util.module_from_spec(spec) + _LOADED.add(resolved) + spec.loader.exec_module(module) + + +#: Declarations are found by this name, at any depth under ``examples/``. +DECLARATION_FILE = "entries.py" + + +def discover_all() -> None: + """Register every declared kernel — today, the one ``examples/entries.py``. + + A glob for one fixed name, which is neither of the two extremes. A single + hard-coded path would make a second table (a vendored kernel set, a fork's + own) silently do nothing, and silence is the failure mode this whole module + exists to remove — ``undeclared_kernels`` already walks the same tree for the + same reason. Importing *any* ``.py`` under ``examples/`` is the other + extreme, and too much: it makes every helper file a declaration table, so a + module put there for one kernel's reference implementation gets executed as + one. Naming the file is the whole convention, and it is what + ``docs/kernelentry.md`` tells a contributor to open. + """ + for path in sorted(EXAMPLES_DIR.rglob(DECLARATION_FILE)): + load_declaration(path) + + +def undeclared_kernels() -> List[str]: + """Kernels under ``examples/`` that no entry declares. + + ``probe --all`` can only report on entries that exist, so without this the + report would describe a handful of kernels and read as though it described + the repository. Coverage is only honest when what is *not* covered is on the + page too. + """ + declared = {entry.mlir_path.resolve() for entry in registered().values()} + found = [] + for path in sorted(EXAMPLES_DIR.rglob("*")): + if path.suffix not in (".mlir", ".ktir") or path.resolve() in declared: + continue + found.append(str(path.relative_to(REPO_ROOT))) + return found + + +def function_names(path: Path) -> Optional[List[str]]: + """The functions the regex parser reads out of *path*, ``None`` if it rejects it. + + Separated from ``entry_for_bare_kernel`` because two callers want different + things from the same read. Somebody probing one file wants the parser's own + exception, which says where in the file it gave up; the repository report wants + that rejection as a cell, because a walk that stops at the first unreadable file + reports nothing about the other thirty. + """ + from ktir_cpu import KTIRInterpreter + + interp = KTIRInterpreter() + try: + interp.load(path.read_text()) + except Exception: # noqa: BLE001 — any rejection is the answer, not an error + return None + return sorted(interp.module.functions) + + +def entry_for_bare_kernel(path: Path, func: Optional[str] = None) -> KernelEntry: + """A throwaway entry for a ``.mlir`` that has no declaration yet. + + The first question a kernel raises — which of its ops the simulator has no + handler for, and whether it survives both parse paths — needs no declaration to + answer. Requiring one first would mean writing the paperwork before finding out + whether the kernel can run at all, which is the opposite of the order that + helps. The claims that do need a declaration report that as their reason. + """ + from ktir_cpu import KTIRInterpreter + + if func is None: + # Read through the interpreter directly rather than through + # ``function_names``, so a file this parser rejects raises the parser's own + # error here. Somebody probing a kernel they have just written needs to + # know where the parse gave up, which a swallowed exception cannot say. + interp = KTIRInterpreter() + interp.load(path.read_text()) + names = sorted(interp.module.functions) + if not names: + # Distinct from the several-functions case below: telling somebody to + # name a function in a file that declares none sends them to look for + # something that is not there. + raise SystemExit( + f"{path.name} is read by the parser but declares no function, so " + "there is no kernel in it to probe" + ) + if len(names) > 1: + raise SystemExit( + f"{path.name} declares {len(names)} functions ({', '.join(names)}); " + "name one with --func" + ) + func = names[0] + try: + rel = str(path.resolve().relative_to(EXAMPLES_DIR)) + except ValueError: + raise SystemExit(f"{path} is not under examples/; move it there first") + return KernelEntry(name=f"{path.stem} (no declaration)", func=func, path=rel) + + +def _resolve(target: Optional[str], use_all: bool, + func: Optional[str] = None) -> List[KernelEntry]: + """The entries one invocation is about. + + Two ways to name one kernel, and the difference between them is deliberate. A + ``.mlir`` path is read cold — no declaration is consulted even if one exists, + which is what makes the first question answerable before any paperwork does. A + bare name is the declared entry, with all five questions in reach. + """ + if use_all: + discover_all() + return [registered()[name] for name in sorted(registered())] + if not target: + raise SystemExit("name a declared kernel, or give a .mlir path, or --all") + path = Path(target) + if not path.exists(): + path = REPO_ROOT / target + if path.suffix in (".mlir", ".ktir"): + return [entry_for_bare_kernel(path, func)] + discover_all() + entries = registered() + if target not in entries: + raise SystemExit( + f"no declared kernel named {target!r}. Declared: " + f"{', '.join(sorted(entries))}. An undeclared kernel is addressed by " + f"its path instead." + ) + return [entries[target]] + + +# --------------------------------------------------------------------------- +# Rendering +# --------------------------------------------------------------------------- + +def _tally_text(ledger: Ledger, leg: str) -> str: + """Counts per state, for the console. + + The committed report does not use this: it is compared verbatim, and + ``parse.frontend`` reads ``closed`` where the MLIR bindings are installed and + ``skip`` where they are not, so any tally of it would depend on the machine that + generated the document. The report answers that claim from a committed record + instead — see ``_support_rows``. + """ + tally = ledger.tally(leg) + total = sum(tally.values()) + if not total: + return "no claims" + parts = [f"{tally.get(CLOSED, 0)}/{total} closed"] + for state in (OPEN, UNDETERMINED, DEFERRED, WAIVED, SKIP): + if tally.get(state): + parts.append(f"{tally[state]} {_STATE_LABEL[state]}") + return ", ".join(parts) + + +def render_ledger(ledger: Ledger, *, verbose: bool) -> str: + entry = ledger.entry + lines = [ + f"{entry.name} function {entry.func}", + f" function leg {_tally_text(ledger, FUNCTION)}", + f" cost leg {_tally_text(ledger, COST)}", + ] + shown = ledger.claims if verbose else [ + c for c in ledger.claims if c.state in BLOCKING or c.state == DEFERRED] + for claim in sorted(shown, key=lambda c: c.sort_key()): + lines.append(f" {_STATE_LABEL[claim.state]:<13} {claim.id}") + if claim.detail: + lines.append(f" {claim.detail}") + if claim.closer and claim.state in BLOCKING: + lines.append(f" -> {claim.closer}") + return "\n".join(lines) + + +#: How many rows of "Read, not declared" have an empty ``execute_kwargs``. The +#: figure belongs to ``tests/conftest.py::EXAMPLE_PARAMS``, and a document +#: renderer that imports the test tree to compute one integer would point the +#: dependency the wrong way round. Instead +#: ``tests/test_kernelentry.py::test_support_report_covers_conftest_examples`` +#: pins it from the side that imports both, so it cannot drift unnoticed. +EMPTY_EXECUTE_KWARGS_ROWS = 8 + +_COLUMNS = ("read: regex", "read: frontend", "runs", "output", + "cost: derivation") + + +def _issue(detail: str) -> str: + """The issue an excuse names, for a cell too narrow to carry the whole reason. + + ``KernelEntry`` will not accept a deferral that names no issue, so there is no + fallback here on purpose: a ``?`` in this column would put the report's own + inability to read a reason where a reader expects a fact about the kernel. + """ + match = re.search(r"#\d+", detail or "") + if match is None: + raise ValueError( + f"deferral reason names no issue as #N: {detail!r}. This is " + "supposed to be unreachable — KernelEntry.__post_init__ rejects it." + ) + return match.group(0) + + +def _cell(claims: Sequence) -> str: + """One column of one kernel, from the claims that column covers. + + Deferred outranks open in the same cell: a gap somebody wrote down and tracked + is a different fact from one nobody has looked at, and the sections below carry + the reason either way. A waiver is reported rather than folded into ``yes``, + because "this check does not apply here" is a decision a reader may want to + disagree with. + """ + if not claims: + return "—" + deferred = [c for c in claims if c.state == DEFERRED] + if deferred: + return f"deferred {_issue(deferred[0].detail)}" + if any(c.state in BLOCKING for c in claims): + return "**no**" + waived = sum(1 for c in claims if c.state == WAIVED) + if waived == len(claims): + return "waived" + return "yes" if not waived else f"yes ({waived} waived)" + + +def _frontend_cell(rel: str, ledger=None) -> str: + """Whether the MLIR frontend accepts this file, from a committed record. + + Never from the generating run. The check needs the optional MLIR bindings, so a + machine without them would print ``skip`` for all thirty kernels and the document + would describe the machine. ``conformance.py`` carries the rejections and a + declaration carries its own excuse; both are in the repository, and + ``tests/mlir_frontend/test_kernelentry_adapt.py`` holds the first to the real + frontend in both directions. + """ + if ledger is not None: + excused = [c for c in ledger.claims + if c.id == "parse.frontend" and c.state in (DEFERRED, WAIVED)] + if excused: + return _cell(excused) + return "**no**" if rel in FRONTEND_REJECTS else "yes" + + +def _row_without_one_kernel(rel: str) -> Optional[Tuple[str, List[str]]]: + """The undeclared row for a file no one kernel can be read out of. + + ``None`` when one can, and the caller probes it as usual. Two ways this + happens and they are different answers, which is why the branch is not one: the + regex parser rejects the file, and the first column *is* that answer; or it + reads the file and finds no single function to probe, and both parse answers + stand while the op count does not. Folding the two together and printing + ``yes`` puts a file nothing could read in the column that says it was read, + which is the one outcome this ledger exists to make impossible. + """ + names = function_names(REPO_ROOT / rel) + if names is None: + return ("?", ["**no**", _frontend_cell(rel)]) + if len(names) != 1: + return ("?", ["yes", _frontend_cell(rel)]) + return None + + +class Survey(NamedTuple): + """Everything one walk over ``examples/`` knows, for both generated documents.""" + + declared: List[Tuple[str, str, List[str]]] + undeclared: List[Tuple[str, str, List[str]]] + kernels_by_op: Dict[str, List[str]] + handlers: int + unhandled: int + + +def survey(ledgers: Sequence[Ledger]) -> Survey: + """One row per kernel under ``examples/``, declared or not, plus the op inversion. + + A row is ``(path, ops, cells)``. Declared and undeclared kernels are separate + lists because the five questions are asked *in order*: the last three need input + that only a declaration supplies, so an undeclared kernel answers a **prefix** of + the five rather than a row with holes in it. An earlier version printed one + table and filled the unasked cells with em-dashes; at 29 of 31 rows that was most + of the document by area, and it read as a repository that supports almost nothing + — the opposite of what those cells actually say. Two tables, each with only the + columns whose question was asked, cannot be misread that way. + + ``kernels_by_op`` is the same walk inverted, and it is what + ``render_ops_report`` publishes. Deriving it here rather than traversing + ``examples/`` a second time keeps one answer to "which ops does this file use" — + the parser's, read back off the claims — instead of two that can disagree. + + ``handlers`` and ``unhandled`` are repo-wide and come free with the walk: the + ``ops`` column is how many distinct ops a kernel uses, and asking whether the + interpreter has a handler for each is the same question. + """ + by_path = {} + for ledger in ledgers: + rel = str(ledger.entry.mlir_path.resolve().relative_to(REPO_ROOT.resolve())) + by_path[rel] = ledger + + declared: List[Tuple[str, str, List[str]]] = [] + undeclared: List[Tuple[str, str, List[str]]] = [] + kernels_by_op: Dict[str, List[str]] = {} + handlers = unhandled = 0 + + def _ops(claims) -> List[str]: + """The op names behind this kernel's handler claims.""" + return [c.id[len("op."):-len(".handler")] + for c in claims if c.id.endswith(".handler")] + + for rel in sorted(set(by_path) | set(undeclared_kernels())): + ledger = by_path.get(rel) + if ledger is None: + row = _row_without_one_kernel(rel) + if row is not None: + undeclared.append((rel, row[0], row[1])) + continue + ledger = probe(entry_for_bare_kernel(REPO_ROOT / rel)) + + ops = _ops(ledger.claims) + handlers += len(ops) + unhandled += sum(1 for c in ledger.claims + if c.id.endswith(".handler") and c.state != CLOSED) + for op in ops: + kernels_by_op.setdefault(op, []).append(rel) + + pick = lambda pred: [c for c in ledger.claims if pred(c.id)] + regex = _cell(pick(lambda i: i == "parse.regex")) + if rel in by_path: + declared.append((rel, str(len(ops)), [ + regex, + _frontend_cell(rel, ledger), + _cell(pick(lambda i: i == "exec.runs")), + _cell(pick(lambda i: i.startswith("out."))), + _cell(pick(lambda i: i == "cost.derivation")), + ])) + else: + undeclared.append((rel, str(len(ops)), + [regex, _frontend_cell(rel)])) + + return Survey(declared, undeclared, kernels_by_op, handlers, unhandled) + + +def _count(n: int, total: int) -> str: + """``All 31 are`` or ``29 of 31 are`` — never ``All 29`` out of thirty-one. + + The first phrasing is only available when the count is the whole set; used + otherwise it reads as though nothing were missing, in the one sentence a reader + takes the repository's state from. + """ + return f"All {total} are" if n == total else f"{n} of {total} are" + + +def _status_line(sv: Survey, green, blocking) -> str: + """Where the repository stands, in one sentence, before any vocabulary. + + The mass goes first — kernels, ops behind them, how many run — because a reader + who meets the narrow tables first counts rows rather than reading them. + """ + rows = sv.declared + sv.undeclared + claims = "claim is" if len(blocking) == 1 else "claims are" + # ``handlers`` sums the per-kernel op counts, so it is what the ``ops`` columns + # add up to rather than a count of distinct ops in the registry. + ops = ("has a handler for every one" if not sv.unhandled + else f"has a handler for {sv.handlers - sv.unhandled} of them") + return ( + f"{len(rows)} kernels under `examples/`, using {sv.handlers} ops between " + f"them; the interpreter {ops}. " + f"{_count(sum(1 for r in rows if r[2][0] == 'yes'), len(rows))} read by " + f"the regex parser and " + f"{sum(1 for r in rows if r[2][1] == 'yes')} are also accepted by the MLIR " + f"frontend and MLIR's own verifier. " + f"{len(sv.declared)} are declared to this ledger, " + f"{'all' if len(green) == len(sv.declared) else len(green)} of which " + f"answer all five questions below; {len(blocking)} {claims} open across " + f"the repository." + ) + + +def render_report(ledgers: Sequence[Ledger], + sv: Optional[Survey] = None) -> str: + """The committed support report: how far this repository supports what. + + Written to ``docs/kernel_support.md`` and compared there on lock-file + discipline, so "what is supported" has a history rather than being re-derived. + ``sv`` is accepted so ``generated_docs`` can walk ``examples/`` once for both + documents; passing nothing walks it here. + """ + sv = sv if sv is not None else survey(ledgers) + blocking = [(led, c) for led in ledgers for c in led.blocking + if not c.env_dependent] + green = [r for r in sv.declared if not any(c == "**no**" for c in r[2])] + + lines = [ + "# Kernel support", + "", + "Generated by `python -m ktir_cpu.kernelentry probe --all --write-report`;", + "`tests/test_kernelentry.py` fails while it is stale.", + "", + # The state of the repository, before the vocabulary for describing it. + _status_line(sv, green, blocking), + "", + "Support is not one property. It is five questions asked in order, and the " + "columns are them: " + "**read: regex** — `KTIRInterpreter.load` accepts the file; " + "**read: frontend** — the MLIR frontend accepts it and MLIR's own " + "verifier passes; **runs** — it executes on the declared grid without " + "overflowing LX; **output** — every tensor it writes matches a " + "reference computed in f32, and none is silently all-zero; " + "**cost: derivation** — the committed per-tensor cost breakdown still " + "matches what the model reports.", + "", + "A cell reads `yes` when the check ran and passed, **`no`** when it ran " + "and failed or could not be evaluated, `deferred #N` against a tracked " + "issue, and `waived` where the check does not apply to that kernel. A " + "kernel is fully supported when no cell in its row is **`no`**.", + "", + "`docs/kernelentry.md` is how to declare a kernel and what a declaration " + "has to supply. `docs/supported_ops.md` is the same repository seen per " + "op rather than per kernel.", + "", + f"## Declared to this ledger ({len(sv.declared)})", + "", + "| kernel | ops | " + " | ".join(_COLUMNS) + " |", + "|" + "---|" * (len(_COLUMNS) + 2), + ] + for rel, ops, cells in sv.declared: + lines.append(f"| `{rel}` | {ops} | " + " | ".join(cells) + " |") + + lines += [ + "", + f"## Read, not declared ({len(sv.undeclared)})", + "", + "Both reading questions are answered for these, and every op they use has " + "an execution handler. The other three have not been *asked*, which is not " + "the same as answered no: they need what only a declaration supplies — " + "tensors to drive the kernel with, and a reference independent of the " + "simulator — so their columns are absent here rather than empty.", + "", + "`tests/` already drives every file below, from " + "`tests/conftest.py::EXAMPLE_PARAMS`, but what that listing supplies is not " + "a declaration waiting to be copied: `execute_kwargs` is empty for " + f"{EMPTY_EXECUTE_KWARGS_ROWS} of them and carries raw HBM element indices, or a " + "scalar size, for the rest. `docs/kernelentry.md` groups these files by " + "which reason applies.", + "", + "| kernel | ops | read: regex | read: frontend |", + "|---|---|---|---|", + ] + for rel, ops, cells in sv.undeclared: + lines.append(f"| `{rel}` | {ops} | " + " | ".join(cells) + " |") + + if any(ops == "?" for _, ops, _ in sv.undeclared): + lines += ["", + "`?` — no one kernel to count ops for, either because the file " + "holds several functions or because the regex parser does not " + "read it at all. The **read: regex** cell says which."] + + def _listed(state: str): + """Detail rows for one excuse state. + + An applied excuse clears ``env_dependent`` — it overrides any state but + ``closed`` — so this is the same filter the blocking list uses, not a + special case. + """ + return [(led, c) for led in ledgers for c in led.claims + if c.state == state and not c.env_dependent] + + deferred = _listed(DEFERRED) + waived = _listed(WAIVED) + + lines += ["", f"## Open and undetermined ({len(blocking)})", ""] + if not blocking: + lines.append("None.") + for led, claim in sorted(blocking, key=lambda kv: kv[1].sort_key()): + lines.append(f"- `{led.entry.name}` — **{claim.id}** " + f"({_STATE_LABEL[claim.state]}): {claim.detail}") + + lines += ["", f"## Deferred ({len(deferred)})", ""] + if not deferred: + lines.append("None.") + for led, claim in deferred: + lines.append(f"- `{led.entry.name}` — **{claim.id}**: {claim.detail}") + + # Grouped by reason: the same waiver on several kernels is one decision, and + # listing it once per kernel pads the document without adding anything. + by_reason: Dict[str, List[str]] = {} + for led, claim in waived: + by_reason.setdefault(claim.detail, []).append(f"{led.entry.name}:{claim.id}") + lines += ["", f"## Waived ({len(waived)})", ""] + if not waived: + lines.append("None.") + for reason in sorted(by_reason): + where = ", ".join(f"`{w}`" for w in sorted(by_reason[reason])) + lines.append(f"- {where} — {reason}") + + rejected = sorted(rel for rel, _, cells in sv.declared + sv.undeclared + if cells[1] == "**no**" and rel in FRONTEND_REJECTS) + # Split by state: a declared kernel's rejection is a deferred claim and is + # listed above, which a reader counting kernels cannot tell unless it is said. + also_deferred = sorted(rel for rel, _, cells in sv.declared + if cells[1].startswith("deferred")) + lines += ["", f"## Not accepted by the MLIR frontend ({len(rejected)})", ""] + lines.append("A record, not a verdict: it does not say whether the file or " + "the parser reading it should change.") + if also_deferred: + lines.append("") + lines.append("No kernel is listed twice: a declared kernel's rejection is a " + "deferred `parse.frontend` claim, listed under **Deferred** " + f"above rather than here ({len(also_deferred)} of them).") + lines.append("") + if not rejected: + lines.append("None.") + # Grouped by reason, as the waived list is: one gap reached through several + # kernels is one gap, and repeating its sentence per kernel says nothing more. + rejects_by_reason: Dict[str, List[str]] = {} + for rel in rejected: + rejects_by_reason.setdefault(FRONTEND_REJECTS[rel], []).append(rel) + for reason in sorted(rejects_by_reason, key=lambda r: rejects_by_reason[r][0]): + where = ", ".join(f"`{w}`" for w in rejects_by_reason[reason]) + lines.append(f"- {where} — {reason}") + + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Verbs +# --------------------------------------------------------------------------- + +def generated_docs(ledgers: Sequence[Ledger]) -> Dict[Path, str]: + """Every document this tool owns, as ``path -> contents``. + + Both are derived from one walk over ``examples/``: they are two views of the + same repository — per kernel and per op — and a second walk could disagree + with the first about which ops a kernel uses. + + Returning the pair rather than writing it is what lets ``verify`` and + ``tests/test_kernelentry.py`` check staleness with the code that writes them, + so a document cannot be checked against a different renderer than the one + that generated it. + """ + from .ops import render_ops_report + + sv = survey(ledgers) + docs = REPO_ROOT / "docs" + return { + docs / "kernel_support.md": render_report(ledgers, sv), + docs / "supported_ops.md": render_ops_report(sv.kernels_by_op), + } + + +def cmd_probe(args: argparse.Namespace) -> int: + entries = _resolve(args.target, args.all, args.func) + ledgers = [probe(entry) for entry in entries] + for ledger in ledgers: + print(render_ledger(ledger, verbose=args.verbose)) + print() + + if args.write_report: + if not args.all: + raise SystemExit("--write-report needs --all: a partial report would " + "claim the kernels it omits are absent") + for path, text in generated_docs(ledgers).items(): + path.write_text(text) + print(f"wrote {path.relative_to(REPO_ROOT)}") + + blocking = sum(len(led.blocking) for led in ledgers) + print(f"{len(ledgers)} kernel(s), {blocking} open or undetermined claim(s)" + f"{_skip_note(ledgers)}") + if args.all: + # Repository-wide, so it is asked once and only when the walk was whole. + findings = audit() + print(f"{len(findings)} pricing finding(s) against " + f"{len(zero_priced_ops())} zero-priced op(s)") + for finding in findings: + print(f" {finding}") + return 0 + + +def cmd_adopt(args: argparse.Namespace) -> int: + """Write the files this tool owns, and print the edits it will not make.""" + from .derivation import DOCUMENT, render_derivation, update_document + + entries = _resolve(args.target, args.all, args.func) + path = REPO_ROOT / DOCUMENT + for entry in entries: + print(f"{entry.name}:") + + ledger = probe(entry) + if ledger.report is None: + print(" the kernel did not run, so no derivation was written:") + for claim in ledger.blocking: + print(f" {claim.id}: {claim.detail}") + else: + body = render_derivation(entry, ledger, dict(entry.gate_params)) + if update_document(path, entry.name, body): + print(f" wrote the {entry.name} section of {DOCUMENT} — read it " + "before committing; it is the attribution a reviewer confirms") + else: + print(f" the {entry.name} section of {DOCUMENT} is unchanged") + + _print_manual_edits(ledger) + return 0 + + +def _print_manual_edits(ledger: Ledger) -> None: + """Say what a human has to change, with the file, and change nothing.""" + remaining = [c for c in ledger.blocking + if not c.id.startswith("cost.derivation")] + if not remaining: + return + print(" edits for you to make (this tool does not touch the interpreter):") + for claim in sorted(remaining, key=lambda c: c.sort_key()): + print(f" {claim.id}") + print(f" {claim.detail}") + if claim.closer: + print(f" in: {claim.closer}") + + +def _skip_note(ledgers: Sequence[Ledger]) -> str: + """What a clean run did not actually check, for the line that says it is clean. + + ``skip`` is not blocking — the claim needs an optional dependency, and CI has + it — but a summary that reports only the kernels leaves "fully supported" + standing on checks this machine never ran. Naming them is the difference + between a gate that passed and a gate that was not fully evaluated here. + """ + skipped = [c for led in ledgers for c in led.claims if c.state == SKIP] + if not skipped: + return "" + ids = sorted({c.id for c in skipped}) + return (f" — {len(skipped)} claim(s) skipped on this machine " + f"({', '.join(ids)}); CI is where that layer is checked") + + +def _repo_wide_problems(ledgers: Sequence[Ledger]) -> List[str]: + """Checks that belong to the repository rather than to any one kernel. + + Both are here for the same reason: they are true or false once, not once per + kernel. The pricing audit was measured to be noise per kernel — 66 of a + prototype's 74 false positives — and a stale generated document is a fact + about the document. + """ + problems: List[str] = [] + for path, expected in generated_docs(ledgers).items(): + rel = path.relative_to(REPO_ROOT) + if not path.exists(): + problems.append(f"{rel} is missing — regenerate with " + "`probe --all --write-report`") + elif path.read_text() != expected: + problems.append(f"{rel} is stale — regenerate with " + "`probe --all --write-report`") + problems += [str(finding) for finding in audit()] + return problems + + +def cmd_verify(args: argparse.Namespace) -> int: + entries = _resolve(args.target, args.all, args.func) + ledgers = [probe(entry) for entry in entries] + bad_kernels = 0 + for ledger in ledgers: + if ledger.clean: + print(f"ok {ledger.entry.name}") + continue + bad_kernels += 1 + print(render_ledger(ledger, verbose=False)) + + problems = _repo_wide_problems(ledgers) if args.all else [] + for problem in problems: + print(problem) + + if bad_kernels or problems: + # Counted separately because they are different failures with different + # fixes: a kernel with an open claim needs work on the kernel or the + # interpreter, and a repository-wide problem is one edit in one file. + parts = [] + if bad_kernels: + parts.append(f"{bad_kernels} kernel(s) with open claims") + if problems: + parts.append(f"{len(problems)} repository-wide problem(s)") + print(f"\n{', '.join(parts)}") + return 1 + print(f"\n{len(ledgers)} kernel(s) fully supported{_skip_note(ledgers)}") + return 0 + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m ktir_cpu.kernelentry", + description="What it takes for this simulator to fully support a kernel.") + sub = parser.add_subparsers(dest="verb", required=True) + + for name, handler, helptext in ( + ("probe", cmd_probe, "compute the ledger; writes nothing"), + ("adopt", cmd_adopt, "write the files this tool owns; print the rest"), + ("verify", cmd_verify, "gate: fail while any claim is open"), + ): + child = sub.add_parser(name, help=helptext) + child.add_argument("target", nargs="?", + help="a declared kernel's name, or the path of a " + ".mlir under examples/ to read cold") + child.add_argument("--func", default=None, + help="function to probe, when a bare .mlir declares " + "more than one") + child.add_argument("--all", action="store_true", + help="every kernel declared in examples/entries.py") + child.add_argument("-v", "--verbose", action="store_true", + help="show closed claims too") + child.add_argument("--write-report", action="store_true", + help="write docs/kernel_support.md (probe --all only)") + child.set_defaults(handler=handler) + + args = parser.parse_args(argv) + return args.handler(args) diff --git a/ktir_cpu/kernelentry/derivation.py b/ktir_cpu/kernelentry/derivation.py new file mode 100644 index 0000000..6dd92e4 --- /dev/null +++ b/ktir_cpu/kernelentry/derivation.py @@ -0,0 +1,308 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The cost report a reviewer confirms, rendered so a diff names what moved. + +Asking an author to write down what a kernel should cost is the highest-value +check available and the one most likely to be waived away, because deriving it is +work. Asking a reviewer to *confirm* a derivation is nearly free. So the tool +derives it and commits the result: the figures arrive in the pull request as a +readable attribution, and a claim in prose that disagrees with the table beside +it does not survive review. + +That inversion is why this file exists at all. It is also the whole of what this +mechanism can do: the derivation is computed by the same code that produces the +measurement, so a mis-charged operation moves the total and its breakdown together +and the page stays self-consistent. Whether the model itself charges correctly is +a question about the model rather than about any one kernel, and is asked in +``tests/test_latency.py``, against hand-counted bytes, FLOPs and cycles. + +All of them land in one generated document, ``docs/kernel_cost.md``, a section per +kernel. One file rather than one beside each ``.mlir`` because a derivation is not +a fact about a kernel a reader goes looking for individually — it is the cost of the +repository, read down the page and diffed as a whole, and eighteen files repeated +this module's caveat eighteen times. + +Two rendering rules, both there so the committed file is a useful diff: + +* Every figure is formatted to four significant digits. A verbatim comparison on + full-precision floats would fail on last-bit noise and teach people to ignore + it. +* Nothing varies between runs at fixed parameters — no timestamps, no dictionary + iteration order, no host detail. A diff must mean the cost changed. +""" + +from __future__ import annotations + +import math +import re +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +#: The generated document every derivation is a section of. +DOCUMENT = "docs/kernel_cost.md" + +TITLE = "# Kernel cost derivations" +HEADER = ("") +PREAMBLE = """\ +One section per declared kernel, at the parameters its declaration gates on. Each +is what the cost model reports, **not an independent check of it**: the figures and +the totals they sum to come from the same code, so a mis-charged op moves both +together and the section stays self-consistent. Whether the model charges +correctly is asked of the model rather than of any one kernel, in +`tests/test_latency.py`, against hand-counted bytes, FLOPs and cycles. + +What a section is for is the composition — which tensor dominates the traffic, what +the ratio to the kernel's declared footprint is, which unit the cycles land on. +Confirm that against what the kernel is supposed to do; a sentence in prose that +disagrees with the table below it does not survive review. Every figure is +formatted to four significant digits, and nothing here varies between runs at fixed +parameters, so a diff means the cost changed.""" + +_SECTION = re.compile( + r"^## (?P[^\n]+)\n\n```text\n(?P.*?)\n```\n", re.M | re.S) + + +def _sig(value: float, digits: int = 4) -> str: + """Format to *digits* significant digits, without exponent noise for ints.""" + if value == 0: + return "0" + if isinstance(value, (int, np.integer)) or float(value).is_integer(): + return f"{int(value):,}" + magnitude = math.floor(math.log10(abs(value))) + decimals = max(0, digits - 1 - magnitude) + return f"{value:,.{decimals}f}" + + +class ArgExtents: + """Which declared argument a view origin belongs to. + + A distributed memory view is charged per partition, so one tensor arrives as + several origins, and an origin belongs to the argument whose element extent + contains it. The output claims and the cost derivation both need that + mapping and had a copy each — same construction, same tie-break, in two files + whose drift would surface as a kernel writing to a tensor in one report and + not in the other. + + An argument with no array behind it has no extent and is therefore never + matched: the ring-reduce kernels take raw element indices rather than + tensors, and a pointer with no length cannot contain anything. Callers get + ``None`` for those and must report them, not drop them. + """ + + def __init__(self, arg_ptrs: Dict[str, Any], tensors: Dict[str, Any]) -> None: + self._extents: List[Tuple[int, int, str]] = [] + self._numel: Dict[str, int] = {} + for name, ptr in arg_ptrs.items(): + value = tensors.get(name) + n = int(value.size) if isinstance(value, np.ndarray) else 0 + if isinstance(ptr, int) and n: + self._extents.append((ptr, ptr + n, name)) + self._numel[name] = n + + def of(self, target: Any) -> Optional[str]: + """The argument *target* lies within, or ``None`` if no extent holds it.""" + if not isinstance(target, int): + return None + hit = [name for lo, hi, name in self._extents if lo <= target < hi] + if not hit: + return None + # Narrowest extent wins: a tensor laid out inside another's address range + # — which the hand-placed multi-tensor examples do — would match both. + return min(hit, key=lambda nm: self._numel[nm]) + + +def _fold_targets(report, arg_ptrs: Dict[str, Any], + tensors: Dict[str, Any], category: str + ) -> Tuple[List[Tuple[str, int, Dict[str, int]]], int]: + """Attribute one transport's bytes to argument names. + + A distributed memory view is charged per partition, so one tensor arrives as + several origins; an origin belongs to the argument whose element extent + contains it. Bytes that match no argument are returned separately and + printed, never folded into a neighbour: a breakdown that silently drops a row + adds up to less than the total while looking complete. + """ + extents = ArgExtents(arg_ptrs, tensors) + per_arg: Dict[str, Dict[str, Any]] = {} + unattributed = 0 + for target, row in report.traffic_by_target(category).items(): + name = extents.of(target) + if name is None: + unattributed += row["nbytes"] + continue + slot = per_arg.setdefault(name, {"nbytes": 0, "ops": {}}) + slot["nbytes"] += row["nbytes"] + for op, count in row["ops"].items(): + slot["ops"][op] = slot["ops"].get(op, 0) + count + + ordered = sorted(per_arg.items(), key=lambda kv: (-kv[1]["nbytes"], kv[0])) + return [(name, slot["nbytes"], slot["ops"]) for name, slot in ordered], unattributed + + +def unique_bytes(tensors: Dict[str, Any]) -> int: + """Bytes a kernel would move if it read and wrote each tensor exactly once. + + The denominator of ``traffic_ratio``. Derived from the declared tensors + rather than from the trace, so it is independent of how the kernel tiles its + access — which is the point of comparing the two. That independence is also + why the ratio can land below 1: a kernel indexing into a tensor rather than + sweeping it moves less than the tensor's whole footprint, and the ratio says + so rather than being clamped. + """ + total = 0 + for value in tensors.values(): + if isinstance(value, np.ndarray): + total += int(value.nbytes) + return total + + +def render_derivation(entry, ledger, params: Dict[str, Any]) -> str: + """The committed text for one entry at one parameter set.""" + report = ledger.report + if report is None: + raise ValueError("no latency report — the kernel did not run") + + tensors = entry.build_tensors(params) if entry.tensors else {} + # Re-deriving the tensors gives shapes and extents; the values are irrelevant + # here and the factory is required to be deterministic at fixed parameters. + arg_ptrs = ledger.arg_ptrs + + lines: List[str] = [f"kernel: {entry.name} function: {entry.func}"] + lines.append("parameters: " + ", ".join( + f"{k}={v}" for k, v in sorted(params.items()))) + lines.append(f"grid cores active: {len(report.counters)}") + lines.append("") + + hbm_total = sum(c.dram_bytes for c in report.counters.values()) + rows, unattributed = _fold_targets(report, arg_ptrs, tensors, "memory") + lines.append(f"hbm_bytes = {hbm_total:,}") + for name, nbytes, ops in rows: + share = (nbytes / hbm_total * 100) if hbm_total else 0.0 + op_text = ", ".join(f"{op} x{count}" for op, count in sorted(ops.items())) + lines.append(f" {name:<18} {nbytes:>14,} {share:5.1f}% {op_text}") + if unattributed: + share = unattributed / hbm_total * 100 if hbm_total else 0.0 + lines.append(f" {'':<18} {unattributed:>14,} {share:5.1f}% " + "origin did not fall inside any argument's extent") + + logical = unique_bytes(tensors) + if logical: + ratio = hbm_total / logical + line = (f"unique_bytes = {logical:,}" + f" traffic_ratio = {_sig(ratio)}x") + if ratio < 1: + line += " (below 1: the kernel reaches only part of a declared tensor)" + lines.append(line) + lines.append("") + + comm_total = sum(c.comm_bytes for c in report.counters.values()) + if comm_total: + # Broken down by op rather than by argument. Cross-core traffic moves a + # tile between cores; it has no origin in an HBM tensor, so asking which + # argument moved it is the wrong question and answering "unattributed" + # for all of it would imply a missing feature instead of a category error. + lines.append(f"comm_bytes = {comm_total:,}") + for op_type, count, nbytes in _comm_by_op(report): + lines.append(f" {op_type:<34} n={count:<6} {nbytes:>12,}") + lines.append("") + + flops = sum(c.total_flops for c in report.counters.values()) + lines.append(f"flops = {flops:,}") + if hbm_total: + lines.append(f"arithmetic_intensity = flops / hbm_bytes = " + f"{flops:,} / {hbm_total:,} = {_sig(flops / hbm_total)}") + lines.append("") + + critical = report.critical_core + lines.append(f"kernel_cycles = {_sig(report.kernel_cycles)}" + f" bottleneck = {report.bottleneck}") + lines.append(f" critical core compute {_sig(critical.compute_cycles)}" + f" memory {_sig(critical.memory_cycles)}" + f" comm {_sig(critical.comm_cycles)}") + for category, cycles in sorted(critical.cycles_by_category.items()): + lines.append(f" {category:<28} {_sig(cycles)}") + lines.append("") + + lines.append("cycles charged per op type, critical core:") + for op_type, count, cycles in _per_op(report, critical): + lines.append(f" {op_type:<34} n={count:<6} {_sig(cycles)}") + lines.append("") + + while lines and not lines[-1]: + lines.pop() + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# The generated document +# --------------------------------------------------------------------------- + +def read_sections(text: str) -> Dict[str, str]: + """Every kernel's derivation in *text*, by kernel name.""" + return {m.group("name"): m.group("body") for m in _SECTION.finditer(text)} + + +def render_document(sections: Dict[str, str]) -> str: + """The whole document, from a name -> derivation mapping. + + Ordered by name rather than by discovery, so adopting one kernel rewrites one + section and a reordering never shows up as a diff. + """ + parts = [TITLE, "", HEADER, "", PREAMBLE, ""] + for name in sorted(sections): + parts += [f"## {name}", "", "```text", sections[name], "```", ""] + return "\n".join(parts).rstrip("\n") + "\n" + + +def update_document(path, name: str, body: str) -> bool: + """Put *body* in as *name*'s section, leaving the others as committed. + + Returns whether anything changed. Adopting one kernel must not depend on + every other kernel having been run in the same invocation: the alternative is + that a single-kernel ``adopt`` silently deletes seventeen sections. + """ + existing = read_sections(path.read_text()) if path.exists() else {} + if existing.get(name) == body: + return False + existing[name] = body + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(render_document(existing)) + return True + + +def _comm_by_op(report) -> List[Tuple[str, int, int]]: + """Chip-wide cross-core bytes, by the op that moved them.""" + counts: Dict[str, int] = {} + nbytes: Dict[str, int] = {} + for counters in report.counters.values(): + for tentry in counters.trace or (): + if tentry.category != "comm" or not tentry.nbytes: + continue + counts[tentry.op_type] = counts.get(tentry.op_type, 0) + 1 + nbytes[tentry.op_type] = nbytes.get(tentry.op_type, 0) + tentry.nbytes + return sorted(((op, counts[op], nbytes[op]) for op in counts), + key=lambda row: (-row[2], row[0])) + + +def _per_op(report, core) -> List[Tuple[str, int, float]]: + counts: Dict[str, int] = {} + cycles: Dict[str, float] = {} + for tentry in core.trace or (): + counts[tentry.op_type] = counts.get(tentry.op_type, 0) + 1 + cycles[tentry.op_type] = cycles.get(tentry.op_type, 0.0) + tentry.cycles + return sorted(((op, counts[op], cycles[op]) for op in counts), + key=lambda row: (-row[2], row[0])) diff --git a/ktir_cpu/kernelentry/ledger.py b/ktir_cpu/kernelentry/ledger.py new file mode 100644 index 0000000..104a452 --- /dev/null +++ b/ktir_cpu/kernelentry/ledger.py @@ -0,0 +1,607 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The one engine that decides a kernel's claims. + +``probe``, ``verify`` and ``tests/test_kernelentry.py`` all call :func:`probe`; +none of them re-implements a check. A second implementation would be a second +opinion, and the two would eventually disagree about whether a kernel is +supported. + +Every check here delegates to machinery that already exists — the regex parser, +the MLIR frontend's own ``verify()``, the dialect handler registry, the frontend +adapter table, ``LatencyReport`` — so a claim closing means the library accepted +the kernel, not that this module was satisfied. + +The hazard this module is built against is **false positives**. A ledger that +flags correct kernels gets waived into silence, and then the waiver mapping +carries no information. A probe-only prototype run over six real artifacts -- +the example files and notebook generators of the three kernels most recently +added to this repository -- produced 74 unexpected open claims, of which 66 came +from a single check that fired on every structural op in every kernel; the +remaining rows were the informative ones. That check is not here: whether an op is priced is a question +about the one cost model every kernel shares, and asking it per kernel made every +kernel answer it. When adding a claim, the question to answer first is not +"what does this catch" but "what correct kernel does this flag". +""" + +from __future__ import annotations + +import contextlib +import io +from dataclasses import dataclass, field +from typing import Any, Dict, Iterable, List, Optional, Tuple + +import numpy as np + +import ktir_cpu.dialects # noqa: F401 — import triggers @register side effects +from ktir_cpu import KTIRInterpreter +from ktir_cpu.dialects import registry +from ktir_cpu.ir_types import _iter_ops +from ktir_cpu.latency import HardwareConfig + +from .derivation import ArgExtents, DOCUMENT, read_sections, render_derivation + +from . import ( + CLOSED, COST, DEFERRED, FUNCTION, OPEN, REPO_ROOT, SKIP, UNDETERMINED, WAIVED, + Claim, KernelEntry, +) + +#: Default tolerances for `out..reference`. Wide because every value in a +#: KTIR kernel is f16 and cross-core folds round their running sum at each step; +#: `tests/test_examples.py` uses the same pair for the split-K matmul and the +#: decode SDPA fold. An entry may widen it for one output through +#: ``KernelEntry.tolerance``, which the claim then reports so the looser pair is +#: not invisible in a green run. +REF_RTOL, REF_ATOL = 2e-2, 2e-1 + + +@dataclass +class Ledger: + """Every claim about one kernel, plus what the engine learned while probing.""" + + entry: KernelEntry + claims: List[Claim] = field(default_factory=list) + #: Populated when the kernel executed, for callers that want the figures + #: rather than the verdicts (``adopt`` writes the derivation from these). + report: Any = None + outputs: Dict[str, Any] = field(default_factory=dict) + output_args: Tuple[str, ...] = () + #: arg name -> pointer value from the run, needed to turn the attribution's + #: element-index keys back into argument names. + arg_ptrs: Dict[str, Any] = field(default_factory=dict) + + @property + def blocking(self) -> List[Claim]: + return [c for c in self.claims if c.blocking] + + @property + def clean(self) -> bool: + return not self.blocking + + def by_leg(self, leg: str) -> List[Claim]: + return [c for c in self.claims if c.leg == leg] + + def tally(self, leg: str) -> Dict[str, int]: + out: Dict[str, int] = {} + for claim in self.by_leg(leg): + out[claim.state] = out.get(claim.state, 0) + 1 + return out + + +class _Builder: + """Accumulates claims, applying the entry's waivers and deferrals uniformly. + + Every claim goes through :meth:`add`, so ``waived``/``deferred`` cannot be + honoured for some checks and forgotten for others — and a waiver naming a + claim the kernel never raises is reported instead of being ignored, since a + stale waiver reads as a considered decision about a check that no longer runs. + """ + + def __init__(self, entry: KernelEntry) -> None: + self.entry = entry + self.claims: List[Claim] = [] + self._seen: set[str] = set() + # Excuses whose check now passes on its own, waived and deferred alike. + self._unnecessary: set[str] = set() + #: Of those, the ones whose check needs an optional dependency, so that + #: "this excuse is no longer needed" is itself only decidable where that + #: dependency is installed — which keeps it out of the committed report. + self._unnecessary_env: set[str] = set() + + def add(self, claim_id: str, leg: str, state: str, detail: str = "", + closer: str = "", env_dependent: bool = False) -> None: + self._seen.add(claim_id) + # One loop over the two excuse mappings, in the order that decides which + # one wins, so a guard cannot be written for one and forgotten for the + # other. ``finish()`` walks the same pair. + for mapping, excused_state in ((self.entry.waived, WAIVED), + (self.entry.deferred, DEFERRED)): + if claim_id not in mapping: + continue + if state == CLOSED: + # Excusing a check that passes is misinformation, not caution: + # the excuse outlives whatever made it necessary and nothing + # says so. ``finish()`` turns this into a claim of its own. + self._unnecessary.add(claim_id) + if env_dependent: + self._unnecessary_env.add(claim_id) + break + # The author's reason replaces the engine's: an excuse says why this + # is knowingly not met and, for a deferral, where it is tracked — + # strictly more informative than repeating what the check looked for. + state, detail = excused_state, mapping[claim_id] + # The declaration decided this one, and it decided it the same way on + # every machine: the excuse applies to any state but CLOSED, and a + # missing optional dependency is one of those. So the outcome stops + # being a fact about the host, and leaving it out of the committed + # report would hide a failure the author already wrote down. + env_dependent = False + break + self.claims.append( + Claim(claim_id, leg, state, detail, closer, env_dependent)) + + def finish(self) -> List[Claim]: + for mapping, what in ((self.entry.waived, "waived"), + (self.entry.deferred, "deferred")): + for claim_id in mapping: + if claim_id not in self._seen: + self.claims.append(Claim( + f"{what}.stale.{claim_id}", FUNCTION, OPEN, + f"{what} names {claim_id!r}, which this kernel does not " + "raise — remove it or fix the id", + closer=f"the entry's {what}= mapping", + )) + for claim_id in sorted(self._unnecessary): + # A deferral is meant to expire on its own once the gap closes. It + # cannot do that through xfail(strict=True) alone: the marker is + # applied from the claim's *current* state, so a closed claim gets no + # marker and the test simply passes. The expiry has to be a claim. + what = "waived" if claim_id in self.entry.waived else "deferred" + self.claims.append(Claim( + f"{what}.unnecessary.{claim_id}", FUNCTION, OPEN, + f"{what} names {claim_id!r}, which now passes on its own — " + f"remove it from {what}=", + closer=f"the entry's {what}= mapping", + env_dependent=claim_id in self._unnecessary_env, + )) + return self.claims + + +# --------------------------------------------------------------------------- +# Derivations from the IR and from a run +# --------------------------------------------------------------------------- + +def distinct_ops(func) -> List[str]: + """Every op type appearing in *func*, regions included, sorted.""" + return sorted({op.op_type for op in _iter_ops(func.operations)}) + + +def store_targets(report) -> Tuple[Dict[Any, int], int]: + """Bytes written by store ops, keyed by the memory they addressed. + + Returns ``(targets, unattributed)``. Reading the trace rather than walking + SSA names is what makes this work for a distributed memory view: the walk + from a store back through ``construct_access_tile`` to a pointer argument + does not survive ``construct_distributed_memory_view``, and when it fails it + fails silently — the prototype produced a kernel with no ``out.*`` claims at + all and a summary line reading ``60/61 closed``. The trace carries the + resolved memory instead, so the failure mode becomes a nonzero + *unattributed* count that the caller must report. + + "Wrote" means wrote to HBM, so the filter names the ``memory`` category + explicitly rather than inheriting it. A store whose view is entirely in LX + is already excluded — ``_estimate`` charges it zero bytes because the tile + never leaves the chip — but that is a coincidence between two files, and this + one is about what the kernel produced. An op priced under any other category + is not output traffic no matter what it is called. + """ + targets: Dict[Any, int] = {} + unattributed = 0 + for counters in report.counters.values(): + for tentry in counters.trace or (): + # Substring, because "store" is how the one storing op is spelled + # today (``ktdp.store``) and a dialect may add a qualified spelling. + if tentry.category != "memory" or "store" not in tentry.op_type: + continue + if not tentry.nbytes: + continue + if tentry.target is None: + unattributed += tentry.nbytes + else: + targets[tentry.target] = targets.get(tentry.target, 0) + tentry.nbytes + return targets, unattributed + + +def fold_to_args(targets: Iterable[Any], arg_ptrs: Dict[str, Any], + tensors: Dict[str, Any]) -> Tuple[set, set]: + """Attribute view origins to the arguments whose element extent contains them. + + The mapping itself is ``derivation.ArgExtents``, shared with the cost + derivation: both answer "which argument is this origin" and two copies of + that answer would eventually disagree about what a kernel wrote to. Returns + ``(arg_names, unresolved_targets)`` — the second is never dropped, because a + target that matched no argument means the ledger does not know what the + kernel wrote to. + """ + extents = ArgExtents(arg_ptrs, tensors) + args, unresolved = set(), set() + for target in targets: + name = extents.of(target) + if name is None: + unresolved.add(target) + else: + args.add(name) + return args, unresolved + + +# --------------------------------------------------------------------------- +# The probe +# --------------------------------------------------------------------------- + +def probe(entry: KernelEntry, *, params: Optional[Dict[str, Any]] = None, + hardware: Optional[HardwareConfig] = None) -> Ledger: + """Compute *entry*'s ledger against the simulator as it stands right now. + + Read-only: nothing is written, and the kernel is executed only in the + simulator's own memory model. *params* defaults to the entry's + ``gate_params``. + """ + params = dict(entry.gate_params if params is None else params) + hardware = hardware or HardwareConfig() + build = _Builder(entry) + ledger = Ledger(entry=entry) + + text = _mlir_text(build, entry) + if text is None: + ledger.claims = build.finish() + return ledger + + interp = _claim_parse(build, entry, text, hardware) + if interp is None: + ledger.claims = build.finish() + return ledger + + func = interp.module.get_function(entry.func) + _claim_ops(build, func) + _claim_frontend_parse(build, text) + _claim_execution(build, ledger, entry, interp, params) + _claim_cost(build, ledger, entry, params) + + ledger.claims = build.finish() + return ledger + + +def _mlir_text(build: _Builder, entry: KernelEntry) -> Optional[str]: + try: + return entry.mlir_text() + except Exception as exc: + build.add("parse.regex", FUNCTION, OPEN, + f"could not obtain MLIR: {type(exc).__name__}: {exc}", + closer=str(entry.mlir_path)) + return None + + +def _claim_parse(build: _Builder, entry: KernelEntry, text: str, + hardware: HardwareConfig) -> Optional[KTIRInterpreter]: + interp = KTIRInterpreter(latency_config=hardware, trace_latency=True) + try: + interp.load(text) + interp.module.get_function(entry.func) + except Exception as exc: + build.add("parse.regex", FUNCTION, OPEN, + f"{type(exc).__name__}: {exc}", + closer="the kernel, or ktir_cpu/parser.py") + return None + build.add("parse.regex", FUNCTION, CLOSED) + return interp + + +def _claim_ops(build: _Builder, func) -> None: + """One handler claim per distinct op. + + Deliberately *not* also a per-op frontend-reachability claim. + ``tests/mlir_frontend/test_registry_consistency.py`` already asserts that + every executor op is either frontend-installed or in that file's + ``FRONTEND_UNSUPPORTED`` allow-list — and an op appearing in a kernel is + necessarily registered, or the claim below would open. So an unreachable op + already breaks the build repository-wide, and ``parse.frontend`` catches it + for this kernel specifically. A third check would only restate them, at the + cost of reading a list that lives in a test module. + """ + for op in distinct_ops(func): + handled = registry.dispatch(op) is not None + build.add(f"op.{op}.handler", FUNCTION, CLOSED if handled else OPEN, + "" if handled else "no execution handler is registered", + closer="ktir_cpu/dialects/ — add a @register handler") + + +def _claim_frontend_parse(build: _Builder, text: str) -> None: + """Does the real MLIR frontend accept the kernel, and does verify() pass? + + Absent ``mlir_ktdp`` this is ``skip``, never ``closed``: the frontend has no + catch-all where the regex parser has one, so a kernel that only works on the + tolerant path is exactly what this claim exists to catch, and reporting it as + closed locally would hide that until CI. + """ + try: + # MLIRFrontendParser raises ImportError from its *constructor*, not on + # import, so the guard has to wrap construction. Getting this wrong makes + # a missing local dependency look like a kernel the frontend rejected — + # the loudest possible false positive, on every kernel at once. + from ktir_cpu.mlir_frontend.parser import MLIRFrontendParser + parser = MLIRFrontendParser() + except ImportError as exc: + build.add("parse.frontend", FUNCTION, SKIP, + f"{exc} — this layer is verified by CI only", + closer="CI", env_dependent=True) + return + try: + parser.parse_module(text) + except Exception as exc: + build.add("parse.frontend", FUNCTION, OPEN, + f"{type(exc).__name__}: {str(exc)[:200]}", + closer="the kernel, or ktir_cpu/mlir_frontend/parser.py", + env_dependent=True) + return + build.add("parse.frontend", FUNCTION, CLOSED, env_dependent=True) + + +def _claim_execution(build: _Builder, ledger: Ledger, entry: KernelEntry, + interp: KTIRInterpreter, params: Dict[str, Any]) -> None: + """exec.runs, then reference / nontrivial for each output it identifies.""" + if not entry.tensors: + detail = ("no tensors= on the entry, so the kernel cannot be driven") + build.add("exec.runs", FUNCTION, OPEN, detail, + closer="the entry's tensors= mapping") + return + + try: + tensors = entry.build_tensors(params) + except Exception as exc: + build.add("exec.runs", FUNCTION, OPEN, + f"tensors= raised {type(exc).__name__}: {exc}", + closer="the entry's tensors= mapping") + return + + try: + with contextlib.redirect_stdout(io.StringIO()): + outputs = interp.execute_function(entry.func, **tensors) + except Exception as exc: + build.add("exec.runs", FUNCTION, OPEN, + f"{type(exc).__name__}: {str(exc)[:200]}", + closer="the kernel, or the interpreter") + return + build.add("exec.runs", FUNCTION, CLOSED) + ledger.outputs = outputs or {} + ledger.report = interp.get_latency_report() + ledger.arg_ptrs = dict(interp.arg_ptrs) + + targets, unattributed = store_targets(ledger.report) + out_args, unresolved = fold_to_args(targets, interp.arg_ptrs, tensors) + ledger.output_args = tuple(sorted(out_args)) + + identified = not (unattributed or unresolved) and bool(out_args) + if unattributed or unresolved: + build.add("out.identified", FUNCTION, UNDETERMINED, + f"{unattributed} bytes of store traffic had no resolvable " + f"origin and {len(unresolved)} origin(s) matched no argument, " + "so the set of output tensors is not fully known", + closer="ktir_cpu/latency.py::LatencyReport.traffic_by_target") + elif not out_args: + build.add("out.identified", FUNCTION, UNDETERMINED, + "no store traffic was recorded, so no output tensor could be " + "identified — a kernel that writes nothing is either wrong or " + "not driven by the arguments tensors= supplies", + closer="the kernel, or the entry's tensors= mapping") + else: + build.add("out.identified", FUNCTION, CLOSED, + f"outputs: {', '.join(ledger.output_args)}") + + expected: Optional[Dict[str, Any]] = None + reference_error: Optional[str] = None + if entry.reference is not None: + try: + # A pristine rebuild, not the mapping the run was given: a kernel + # whose output argument aliases its input has overwritten that array + # by now, and a reference reading it would compare the result with + # itself and agree. The specs are deterministic, so the two builds + # are the same arrays until the kernel touches one of them. + expected = entry.reference(params=params, + tensors=entry.build_tensors(params)) + except Exception as exc: + # Reported through each output's own claim rather than as a claim of + # its own. A claim that exists only when something fails is absent + # from the ledger the rest of the time, so it never appears in the + # support report and a waiver naming it reads as stale. + reference_error = f"reference= raised {type(exc).__name__}: {exc}" + + for name in ledger.output_args: + _claim_nontrivial(build, name, outputs.get(name)) + _claim_reference(build, name, outputs.get(name), expected, reference_error) + if identified and expected is not None: + _claim_unwritten_references(build, ledger.output_args, expected) + + +def _claim_nontrivial(build: _Builder, name: str, value: Any) -> None: + """An all-zero or all-NaN output, which no cost report would ever flag. + + The decode-SDPA notebook helper already computes an fp32 reference and counts + zero rows on every run, and its docstring gives the reason: every value in + that kernel is f16 including the cross-core fold's running sum, so an input + scale large enough to overflow shows up here and nowhere else in the report. + This generalises that check. A legitimately all-zero output (a fully masked + kernel) waives the claim with that as the reason. + """ + claim_id = f"out.{name}.nontrivial" + if value is None: + # The store trace named this argument as written, and the run did not + # return it. Reporting the dtype of a None instead — which is what + # asarray produces — would name the symptom and hide the cause. + build.add(claim_id, FUNCTION, UNDETERMINED, + f"the kernel wrote to {name!r} but the run returned no such " + "output, so there is nothing to inspect", + closer="the entry's tensors= mapping") + return + got = np.asarray(value) + if got.dtype.kind not in "fc": + build.add(claim_id, FUNCTION, WAIVED, + f"dtype {got.dtype} is not floating point") + return + as_f32 = got.astype(np.float32) + if not np.any(got): + build.add(claim_id, FUNCTION, OPEN, "output is entirely zero", + closer="the kernel, or a waiver if zero is correct here") + return + if np.any(np.isnan(as_f32)): + build.add(claim_id, FUNCTION, OPEN, "output contains NaN", + closer="the kernel") + return + flat = as_f32.reshape(-1, as_f32.shape[-1]) if as_f32.ndim else as_f32 + zero_rows = int(np.sum(~np.any(flat, axis=1))) if as_f32.ndim else 0 + if zero_rows: + build.add(claim_id, FUNCTION, OPEN, + f"{zero_rows}/{flat.shape[0]} rows are entirely zero", + closer="the kernel, or a waiver if that is correct here") + return + build.add(claim_id, FUNCTION, CLOSED) + + +def _claim_reference(build: _Builder, name: str, value: Any, + expected: Optional[Dict[str, Any]], + reference_error: Optional[str] = None) -> None: + claim_id = f"out.{name}.reference" + if reference_error is not None: + build.add(claim_id, FUNCTION, OPEN, reference_error, + closer="the entry's reference= callable") + return + if value is None: + build.add(claim_id, FUNCTION, UNDETERMINED, + f"the run returned no output named {name!r}", + closer="the entry's tensors= mapping") + return + got = np.asarray(value) + if expected is None: + build.add(claim_id, FUNCTION, OPEN, + "no reference declared, so correctness is unknown — this is " + "open rather than skipped because a kernel that writes " + "plausible nonsense produces a perfectly consistent report", + closer="the entry's reference= callable") + return + if name not in expected: + build.add(claim_id, FUNCTION, UNDETERMINED, + f"reference= returned no entry for {name!r}", + closer="the entry's reference= callable") + return + want = np.asarray(expected[name]) + if want.dtype.kind == "f" and want.dtype.itemsize < 4: + build.add(claim_id, FUNCTION, OPEN, + f"reference for {name!r} is {want.dtype}: a reference computed " + "at the kernel's own precision reproduces the kernel's own " + "overflow and then agrees with it — compute in f32 or wider", + closer="the entry's reference= callable") + return + if want.shape != got.shape: + build.add(claim_id, FUNCTION, OPEN, + f"shape {got.shape} != reference {want.shape}", + closer="the kernel, or the reference") + return + rtol, atol = build.entry.tolerance.get(name, (REF_RTOL, REF_ATOL)) + close = np.allclose(got.astype(np.float32), want.astype(np.float32), + rtol=rtol, atol=atol) + diff = np.abs(got.astype(np.float32) - want.astype(np.float32)) + if close: + # A widened pair is stated even when the claim closes. A tolerance is the + # one part of this claim that can be adjusted until it passes, so the run + # that passes has to say which pair it passed against. + build.add(claim_id, FUNCTION, CLOSED, + f"max abs diff {float(diff.max()):.4g} against a declared " + f"rtol={rtol}, atol={atol}" + if (rtol, atol) != (REF_RTOL, REF_ATOL) else "") + return + build.add(claim_id, FUNCTION, OPEN, + f"max abs diff {float(diff.max()):.4g} at output magnitude " + f"{float(np.abs(want).max()):.4g} " + f"(rtol={rtol}, atol={atol})", + closer="the kernel, or the reference") + + +def _claim_unwritten_references(build: _Builder, output_args: Tuple[str, ...], + expected: Dict[str, Any]) -> None: + """A reference entry for a tensor no store in the trace wrote. + + The mirror of the ``name not in expected`` branch above, and both directions are + needed because the claim set is derived from the store trace rather than from + the declaration: a reference key the trace never names raises no claim at all, + so an output that stops being written loses its comparison instead of failing + it — which is one of the defects this ledger exists to catch. Reported at the + id the comparison would have had, so the vanished claim reappears under its own + name rather than as a remark on a different one. + + Asked only where ``out.identified`` closed. Where it did not, which tensors the + kernel wrote is exactly what is in doubt, "no store wrote it" is not established + for any key, and that claim already blocks. + """ + for name in sorted(set(expected) - set(output_args)): + build.add(f"out.{name}.reference", FUNCTION, UNDETERMINED, + f"reference= returned an entry for {name!r} and no store in the " + "trace wrote it, so nothing was compared against it", + closer="the entry's reference= callable, or the kernel") + + +def _claim_cost(build: _Builder, ledger: Ledger, entry: KernelEntry, + params: Dict[str, Any]) -> None: + """cost.derivation — the committed attribution, compared verbatim. + + The derivation is generated, committed, and compared verbatim, so a change in + what the kernel costs arrives as a reviewable diff naming the term that moved. + That is what a sentence in prose disagreeing with the model does not survive. + + What it structurally cannot do is catch the cost model itself being wrong: the + derivation is produced by the same code as the measurement, so a mis-charged op + moves the total and its breakdown together and the page stays self-consistent. + That question is asked of the model rather than of any one kernel, and lives in + ``tests/test_latency.py``, which holds hand-counted bytes, FLOPs and cycles + against every latency category and every hardware parameter that scales them. + Re-asking it here, once per declaration, would make each kernel re-answer one + question about the single cost model they all share. + """ + if ledger.report is None: + build.add("cost.derivation", COST, UNDETERMINED, + "the kernel did not run, so there is nothing to derive", + closer="exec.runs") + return + + path = REPO_ROOT / DOCUMENT + try: + rendered = render_derivation(entry, ledger, params) + except Exception as exc: + build.add("cost.derivation", COST, UNDETERMINED, + f"could not render: {type(exc).__name__}: {exc}", + closer="ktir_cpu/kernelentry/derivation.py") + return + + committed = read_sections(path.read_text()).get(entry.name) if path.exists() else None + if committed is None: + build.add("cost.derivation", COST, OPEN, + f"{DOCUMENT} has no section for {entry.name} — run `adopt` to " + "write it, then read it: it is the attribution a reviewer confirms", + closer=DOCUMENT) + elif committed == rendered: + build.add("cost.derivation", COST, CLOSED) + else: + build.add("cost.derivation", COST, OPEN, + f"the {entry.name} section of {DOCUMENT} is stale — regenerate " + "with `adopt` and read the diff, which names the term that moved", + closer=DOCUMENT) diff --git a/ktir_cpu/kernelentry/ops.py b/ktir_cpu/kernelentry/ops.py new file mode 100644 index 0000000..0e32f2a --- /dev/null +++ b/ktir_cpu/kernelentry/ops.py @@ -0,0 +1,364 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""``docs/supported_ops.md`` — the same repository seen per op, not per kernel. + +``docs/kernel_support.md`` answers five questions about each *kernel*, and the +last three of them need a declaration, so it is mostly a report about how far +that ledger reaches. This document answers a narrower question that needs no +declaration at all — *does this repository have a handler for this op, and can +the MLIR frontend reach it* — and it answers it for every op in the registry, +including the ones no example exercises. + +Every column is read from something the repository already maintains: + +============================ ======================================== +column source +============================ ======================================== +executor handler ``registry._REGISTRY`` +MLIR frontend ``MLIRTypeAdapter._adapt_handlers`` +reason, where not reachable ``FRONTEND_UNSUPPORTED`` in + ``tests/mlir_frontend/test_registry_consistency.py`` +kernels the walk in ``cli.survey``, inverted +pricing ``ZERO_COST_OPS`` / ``UNJUDGED_ZERO_OPS`` in + ``ktir_cpu/kernelentry/pricing.py`` +============================ ======================================== + +The first two are the same two sets ``test_registry_consistency.py`` already +asserts about; this module does not invent a second criterion, it publishes the +one that is enforced. The last comes from the ledger's own walk rather than a +second traversal of ``examples/``, so "which ops does this file use" has one +answer here and in ``docs/kernel_support.md``. + +**The pricing column does not print the registry's answer.** ``@register()`` +defaults ``latency_category`` to ``"zero"``, so "this op is free" and "nobody +priced this op" are the same state there, and a column that printed it would +report a default as a decision. What it prints instead is +``ktir_cpu/kernelentry/pricing.py``, which splits that state in two and is +audited against the registry by ``tests/test_kernelentry.py``. The audit is +repository-wide rather than per kernel because the ops it would flag are the same +structural ones in every kernel — measured: 66 of a per-kernel prototype's 74 +false positives came from asking it once per kernel. + +One thing it still cannot see: a **wrong** category. It compares against +``zero``, so an op billed to the wrong non-zero class passes it. That needs a +reader who knows the op's semantics, which is why the reasons in ``pricing.py`` +are written rather than generated. +""" + +from __future__ import annotations + +import importlib.util +from typing import Dict, List + +import ktir_cpu.dialects # noqa: F401 — import triggers @register side effects +from ktir_cpu.dialects import registry + +from . import REPO_ROOT +from .pricing import UNJUDGED_ZERO_OPS, ZERO_COST_OPS + +CONSISTENCY_TEST = (REPO_ROOT / "tests" / "mlir_frontend" + / "test_registry_consistency.py") + +# Above this many files, "Where each op appears" summarises instead of naming +# every path: an op that every example uses is one fact, not one fact per file. +_MAX_NAMED_FILES = 3 + +# What an op's appearance in one example directory is worth as evidence. The +# distinction matters for the ``kernels`` column: ``rfc/`` examples are +# specification cases the suite expects to fail execution, so an op seen only +# there has been parsed, not run. +CATEGORY_NOTE = { + "ktir": "hand-written dialect cases; executed by the test suite", + "latency": "small kernels for the latency tests; executed by the test suite", + "rfc": "RFC 0682 specification examples; **expected to fail execution** " + "(they carry absolute addresses rather than arguments), so an op " + "seen only here is parsed, not run", + "sdsc": "kernels from the SuperDSC lowering path; executed by the test suite", + "triton-ktir": "kernels as the Triton → KTIR path emits them, i.e. captured " + "compiler output; executed by the test suite", +} + + +def _frontend_allowlist() -> Dict[str, str]: + """Read ``FRONTEND_UNSUPPORTED`` out of the test that enforces it. + + It lives in the test rather than the package because it is a statement about + a known divergence, not a fact the library needs at run time. That makes it + un-importable as a module path, so it is loaded by location — the same way + this package already reads ``examples/`` by path. Reading a checkout from + library code is sound only because this package is a development tool and is + not packaged: ``pyproject.toml`` excludes ``ktir_cpu.kernelentry`` from the + wheel, so there is no install in which these paths are absent. + """ + spec = importlib.util.spec_from_file_location( + "_registry_consistency", CONSISTENCY_TEST) + if spec is None or spec.loader is None: # pragma: no cover — path is fixed + raise ValueError(f"cannot load {CONSISTENCY_TEST}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + try: + return dict(module.FRONTEND_UNSUPPORTED) + except AttributeError: # pragma: no cover — guarded so a rename is loud + raise ValueError( + f"{CONSISTENCY_TEST.relative_to(REPO_ROOT)} no longer defines " + "FRONTEND_UNSUPPORTED; update ktir_cpu/kernelentry/ops.py to read " + "whatever replaced it.") + + +def _frontend_handlers() -> set: + """Ops the MLIR frontend can adapt. + + Host-independent, which this document needs and the kernel report has to work + for: ``_adapt_handlers`` is filled by ``@MLIRTypeAdapter.install`` decorators + in this package, not by anything the MLIR bindings provide, and the module + guards its own ``mlir_ktdp`` import. So this table reads the same on a + machine without the bindings, and a committed document generated there does + not contradict one generated in CI. Imported plainly rather than guarded for + the same reason: if that ever stops being true, it should fail loudly instead + of publishing "no op is reachable". + """ + from ktir_cpu.mlir_frontend.parser import MLIRTypeAdapter + + return set(MLIRTypeAdapter._adapt_handlers) + + +def _pricing_cell(op: str) -> str: + """What the cost model charges for *op*, with ``zero`` split into two answers. + + A real category is printed as itself. ``zero`` is printed as the decision + behind it, never as the registry's word for it, because the registry cannot + tell a decision from its own default. ``**unlisted**`` should not occur: + ``tests/test_kernelentry.py`` fails on it, the same way a bare ``**no**`` in + the frontend column is a test failure rather than a row. + """ + if op not in registry._REGISTRY: + return "—" + category = registry.get_latency_category(op) + if category != "zero": + return f"`{category}`" + if op in ZERO_COST_OPS: + return "free, by decision" + if op in UNJUDGED_ZERO_OPS: + return "**not judged**" + return "**unlisted**" + + +def render_ops_report(kernels_by_op: Dict[str, List[str]]) -> str: + """The committed op-level report. A pure function of the registries. + + ``kernels_by_op`` maps an op name to the kernels under ``examples/`` that use + it, as ``cli.survey`` observed them. It is passed in rather than recomputed + so that both generated documents agree about what an op is: the parser's + answer, read back off the ledger's claims. + """ + executor = dict(registry._REGISTRY) + frontend = _frontend_handlers() + allowlist = _frontend_allowlist() + + all_ops = sorted(set(executor) | frontend | set(kernels_by_op)) + unhandled = [op for op in sorted(kernels_by_op) if op not in executor] + unexercised = sorted(set(executor) - set(kernels_by_op)) + zero = [op for op in sorted(executor) + if registry.get_latency_category(op) == "zero"] + free = [op for op in zero if op in ZERO_COST_OPS] + unjudged = [op for op in zero if op in UNJUDGED_ZERO_OPS] + + out = [ + "", + "", + "# Supported operations", + "", + "Op-level truth about this interpreter, read from the registries " + "themselves. `tests/test_kernelentry.py` fails while this file and the " + "registries disagree, so it cannot go stale silently.", + "", + f"{len(executor)} ops have an execution handler. " + f"{len(set(executor) & frontend)} of them are also reachable through " + f"the MLIR frontend; {len(allowlist)} are deliberately not " + "([why](#ops-not-reachable-through-the-mlir-frontend)). " + f"{len(kernels_by_op)} are used by a kernel under `examples/`, leaving " + f"{len(unexercised)} that no example exercises " + "([which](#ops-no-example-exercises)).", + "", + f"{len(zero)} of them cost nothing. That is two facts, not one: " + f"{len(free)} are [free by decision](#ops-that-are-free-by-decision) and " + f"{len(unjudged)} are [priced zero with nobody having decided]" + "(#ops-priced-zero-without-a-decision). `@register()` defaults " + "`latency_category` to `zero`, so the registry itself cannot tell those " + "apart — `ktir_cpu/kernelentry/pricing.py` is where they are split, and " + "an op priced zero that appears in neither list fails " + "`tests/test_kernelentry.py`.", + "", + "**Division of labour.** `docs/kernel_support.md` asks five questions " + "about each *kernel*, and three of them need a declared entry — " + "so it reports how far that ledger reaches. This file asks one question " + "about each *op*, needs no declaration, and therefore covers the whole " + "registry. `docs/gap_analysis.md` is the third: conformance against " + "RFC 0682, judged by a reader rather than generated.", + "", + "## Matrix", + "", + "| op | executor handler | MLIR frontend | cost | kernels |", + "|---|---|---|---|---|", + ] + for op in all_ops: + has_exec = "yes" if op in executor else "**no**" + if op in frontend: + fe = "yes" + elif op in allowlist: + fe = "no, by design" + else: + fe = "**no**" + n = len(kernels_by_op.get(op, ())) + out.append(f"| `{op}` | {has_exec} | {fe} | {_pricing_cell(op)} | " + f"{n or 'none'} |") + + out += [ + "", + "### How to read the columns", + "", + "- **executor handler** — a `@register` handler exists, so the " + "interpreter can execute the op. `no` means it cannot; such an op is " + "listed here only because the MLIR frontend or an example mentions it.", + "- **MLIR frontend** — an `@MLIRTypeAdapter.install` handler exists, so " + "the op survives the real MLIR parser and its `verify()`. `no, by " + "design` is an entry in the allowlist below. A bare **no** should not " + "occur: `tests/mlir_frontend/test_registry_consistency.py` fails on it.", + "- **cost** — what the latency model charges. A named category is the " + "one `@register` gave it. `free, by decision` and `**not judged**` are " + "both `zero` in the registry, split apart by " + "`ktir_cpu/kernelentry/pricing.py`: the first has a written reason why " + "the hardware does no measurable work, the second has a written " + "statement of what is unresolved and the issue tracking it. A cost " + "column can only catch a *missing* price — an op billed to the wrong " + "non-zero unit passes it, and needs a reader.", + "- **kernels** — how many files under `examples/` use the op; " + "[the files themselves](#where-each-op-appears) are listed below, with " + "what each directory is worth as evidence. `none` means the op is " + "registered but no example exercises it, so only unit tests, if any, " + "cover it.", + "", + "What this file does not say: whether a handler is **correct**, and " + "whether a particular attribute or type of the op is supported. The unit " + "here is the op name — `ktdp.construct_access_tile` having a handler " + "says nothing about a particular `base_map` or `coordinate_set` reaching " + "the conclusion the specification does. `docs/gap_analysis.md` tracks " + "that.", + "", + "## Ops not reachable through the MLIR frontend", + "", + "Executor ops with no MLIR frontend handler, and the reason each is " + "allowed to stay that way. This is the allowlist " + "`tests/mlir_frontend/test_registry_consistency.py` enforces — an op " + "missing from both the frontend and this list fails that test.", + "", + "| op | reason |", + "|---|---|", + ] + for op, reason in sorted(allowlist.items()): + out.append(f"| `{op}` | {' '.join(reason.split())} |") + + if unhandled: + out += [ + "", + "## Ops used by an example with no execution handler", + "", + "These appear in a file under `examples/` that the interpreter " + "cannot execute as written:", + "", + ] + out += [f"- `{op}` — {', '.join(kernels_by_op[op])}" for op in unhandled] + + out += [ + "", + "## Ops that are free by decision", + "", + f"{len(free)} ops the cost model charges nothing for, and the reason each " + "one does no measurable work. Four kinds: a value that exists at compile " + "time, a terminator that only names values, addressing metadata that " + "computes where data is without moving it, and an orchestrator whose " + "region is charged op by op.", + "", + "| op | why it is free |", + "|---|---|", + ] + out += [f"| `{op}` | {' '.join(ZERO_COST_OPS[op].split())} |" for op in free] + + out += [ + "", + "## Ops priced zero without a decision", + "", + f"{len(unjudged)} ops that cost nothing today because " + "`latency_category` defaults to `zero`, not because anyone decided they " + "are free. A kernel using one of these reports a lower cost than the " + "hardware would. They are listed rather than fixed because each is a " + "hardware question RFC 0682 does not settle — a conversion folded into " + "the consumer's read is free and a materialized one is not — and " + "repricing one changes the committed derivations of every kernel using " + "it.", + "", + "| op | what is unresolved | kernels |", + "|---|---|---|", + ] + out += [f"| `{op}` | {' '.join(UNJUDGED_ZERO_OPS[op].split())} | " + f"{len(kernels_by_op.get(op, ())) or 'none'} |" for op in unjudged] + + out += [ + "", + "## Ops no example exercises", + "", + f"{len(unexercised)} registered ops that no file under `examples/` uses. " + "Not a defect on its own — an op can be covered by a unit test — but it " + "is where a handler nothing has ever run would hide:", + "", + ] + out += [f"- `{op}`" for op in unexercised] or ["None."] + + out += [ + "", + "## Where each op appears", + "", + f"Up to {_MAX_NAMED_FILES} files are named; above that the directories " + "and a count, " + "because the identity of the file is what matters when there are few and " + "the coverage class is what matters when there are many. Paths are " + "relative to `examples/`, and what each directory is worth as evidence:", + "", + ] + # Keys arrive repo-relative (``examples/ktir/x.mlir``); the category is the + # segment after ``examples/``, so strip that first rather than splitting on + # the leading slash and getting ``examples`` for every one of them. + def _category(rel: str) -> str: + return rel.split("/")[1] if rel.startswith("examples/") else rel + + seen_dirs = sorted({_category(k) for ks in kernels_by_op.values() for k in ks}) + for name in seen_dirs: + note = CATEGORY_NOTE.get( + name, "**no note** — add one to `CATEGORY_NOTE` in " + "`ktir_cpu/kernelentry/ops.py`") + out.append(f"- `{name}/` — {note}") + out.append("") + for op in sorted(kernels_by_op): + files = sorted(kernels_by_op[op]) + if len(files) <= _MAX_NAMED_FILES: + where = ", ".join(k.replace("examples/", "") for k in files) + else: + dirs = ", ".join(f"{d}/" for d in sorted({_category(k) + for k in files})) + where = f"{len(files)} files in {dirs}" + out.append(f"- `{op}` — {where}") + out.append("") + return "\n".join(out) diff --git a/pyproject.toml b/pyproject.toml index d4b6e35..6e41ed1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,10 @@ mlir-frontend = [ [tool.setuptools.packages.find] include = ["ktir_cpu*"] +# ktir_cpu.kernelentry is a development tool rather than part of the library: it +# reads examples/ and tests/ by path, and a wheel carries neither. Excluded for +# the same reason examples/ itself is not packaged. +exclude = ["ktir_cpu.kernelentry*"] [tool.setuptools.package-data] "*" = ["*.mlir"] From 4f368fe82473fd537221b2e120919c6551b66516 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:15 -0400 Subject: [PATCH 4/9] Write two cross-core kernels in the dialect's own assembly form MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ledger the previous commit adds asks the MLIR frontend whether it accepts each kernel, and these two were the only ones it refused that anything here declares. Both write `inter_tile_produce` / `inter_tile_reduce` as : T -> !ktdp.tile_future which the regex parser reads and the built dialect does not print, so the frontend stops at the first of the two ops with `expected '->'`. The dialect's own custom assembly is -> <(T), groups = S> (produce) : <(T), groups = S> -> R (reduce) and the regex parser reads that form as well, which settles which of the two spellings is authoritative: both parsers accept the second, only one accepts the first. Two lines per file. The authors' line breaks are left alone, so the diff is the type syntax and nothing besides it, and nothing about either kernel moves: both parsers accept the result, and the cost derivations are identical to the digit, because the spelling of a type is not an input to the cost model. This precedes the generated reports so that no committed document records an acceptance the tree does not have at that commit. The three `ring_reduce*` kernels are deliberately left in the old form. Rewriting their spelling does not make them parse: all three reduce `tensor<1x128xf16>` to `tensor<128xf16>`, and the dialect verifies that a reduce's result matches the future's partial type, so the error becomes `result types must match future partial types` instead of disappearing. The reduce here takes its result type as a reshape target (`_result_shape` -> `attach_reshape` in `ktir_cpu/dialects/ktdp_ops.py`), so closing that gap means deciding whether it should reshape at all — a decision about the op rather than about how those three files are written down. They stay recorded as rejected, with that distinction in the reason and in gap row 2a. Signed-off-by: WarningRan --- examples/ktir/ffn_swiglu_4core.mlir | 4 ++-- examples/sdsc/sdpa_pv_ksplit.mlir | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/ktir/ffn_swiglu_4core.mlir b/examples/ktir/ffn_swiglu_4core.mlir index c96feaa..879ac22 100644 --- a/examples/ktir/ffn_swiglu_4core.mlir +++ b/examples/ktir/ffn_swiglu_4core.mlir @@ -272,7 +272,7 @@ module { %fut = ktdp.inter_tile_produce producer_tiles_per_group = #all_tiles - : tensor<1024xf16> -> !ktdp.tile_future, groups = affine_set<(g) : (g == 0)>> + -> <(tensor<1024xf16>), groups = affine_set<(g) : (g == 0)>> { ^bb0(%gid: index): ktdp.yield_partial %out_partial_flat : tensor<1024xf16> @@ -285,7 +285,7 @@ module { %out_reduced_flat = ktdp.inter_tile_reduce(%fut) consumer_tiles_per_group = #all_tiles, identity(%add_id : tensor<1024xf16>) - : !ktdp.tile_future, groups = affine_set<(g) : (g == 0)>> -> tensor<1024xf16> + : <(tensor<1024xf16>), groups = affine_set<(g) : (g == 0)>> -> tensor<1024xf16> { ^bb0(%lhs: tensor<1024xf16>, %rhs: tensor<1024xf16>): %init = tensor.empty() : tensor<1024xf16> diff --git a/examples/sdsc/sdpa_pv_ksplit.mlir b/examples/sdsc/sdpa_pv_ksplit.mlir index cb8907f..b41d0cd 100644 --- a/examples/sdsc/sdpa_pv_ksplit.mlir +++ b/examples/sdsc/sdpa_pv_ksplit.mlir @@ -91,7 +91,7 @@ module { %prod_16 = linalg.matmul ins(%a_13, %b_14 : tensor<1x512xf16>, tensor<512x64xf16>) outs(%init_15 : tensor<1x64xf16>) -> tensor<1x64xf16> %fut_17 = ktdp.inter_tile_produce producer_tiles_per_group = #red_tiles - : tensor<1x64xf16> -> !ktdp.tile_future, groups = affine_set<(g) : (g >= 0, -g + 1 >= 0)>> + -> <(tensor<1x64xf16>), groups = affine_set<(g) : (g >= 0, -g + 1 >= 0)>> { ^bb0(%gid: index): ktdp.yield_partial %prod_16 : tensor<1x64xf16> @@ -102,7 +102,7 @@ module { %reduced_21 = ktdp.inter_tile_reduce(%fut_17) consumer_tiles_per_group = #red_tiles, identity(%add_id_20 : tensor<1x64xf16>) - : !ktdp.tile_future, groups = affine_set<(g) : (g >= 0, -g + 1 >= 0)>> -> tensor<1x64xf16> + : <(tensor<1x64xf16>), groups = affine_set<(g) : (g >= 0, -g + 1 >= 0)>> -> tensor<1x64xf16> { ^bb0(%lhs: tensor<1x64xf16>, %rhs: tensor<1x64xf16>): %r_init_22 = tensor.empty() : tensor<1x64xf16> From d9e929418c436906f7eddbfca868c2fed7f462b7 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:16 -0400 Subject: [PATCH 5/9] Declare 18 kernels to the ledger MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every kernel under `examples/` that can be driven through the `tensors=` seam is declared, and every claim the eighteen raise is closed. Declaring one is a row in the `ENTRIES` table in `examples/entries.py` — no new file, nothing to register, no test to write. The blocker looked larger than it was: seventeen of these kernels already had a reference inside `tests/`, computed in f32 and then cast back to f16 before comparison, and the cast is what made it a reference to the kernel's own rounding rather than an independent one. Dropping it is a one-line change per kernel, at the tolerances the ledger already used. Three needed more than a transcription: - `ffn_swiglu_4core` overflows f16 in `exp` at unit-variance weights and writes an all-zero output, which no cost report would flag. Its inputs are scaled by 1/sqrt(fan-in) — standard initialisation, not a number chosen to make a check pass — and `out.*.nontrivial` is what noticed. - `layernorm_fwd_ktir` writes three tensors and the end-to-end test checks two. `Rstd` now has a reference, so the reciprocal standard deviation of a fold over 8192 f16 values is compared against something for the first time. - `paged_attention` is referenced at all thirty-two grid positions rather than the first, and with the softmax taken flat over the masked row rather than folded tile by tile. Both matter: a mask indexed by query position is exactly what can be right at `pid0 = 0` and wrong everywhere else, and the tiled fold is the kernel's own algorithm, so a reference written that way checks the arithmetic and not the decomposition. The fifteen that remain undeclared are not a backlog. Eight `examples/rfc/*` files take no tensor arguments — their tensors are memrefs at absolute HBM bases, so there is nothing to bind, and the two that take anything take one scalar size. Five take raw HBM element indices with no shape attached — the three ring-reduce kernels and the two `rmsnorm_4core_*` — and the tests that run them seed memory by replacing `_prepare_execution`, which a declaration cannot express; that is a property of the argument convention rather than of those files, and it is the live one, since the two rmsnorm kernels are the most recent to arrive and their tests seed memory the same way. `softmax_wide.mlir` overflows LX on purpose and its test asserts the exception, so `exec.runs` failing is the kernel behaving. `nested_yield.ktir` is rejected by the frontend for a `construct_memory_view` with no `coordinate_set`, and it is a reproducer written down to the one op under test, so being minimal rather than dialect-valid is what it is for. No declaration defers or waives a claim. That is worth stating rather than assuming, because the two that would have — the cross-core kernels of the previous commit — are the reason `deferred` exists in the engine at all, and the state ends up exercised only by the gate. The gap this branch found and did not close is the one it cannot: the three `ring_reduce*` kernels are on the page, in the record of what the frontend refuses, with the decision their author has to make written beside them. The gate's cost is shape, not entry count: probing all eighteen takes about 16 s, and timed one kernel at a time, five kernels whose shapes are baked into their IR are 94% of it. They are here anyway, because a cost model checked only at reduced size is checked where the padding, the tail core and the page table are not real yet. Signed-off-by: WarningRan --- examples/entries.py | 672 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 672 insertions(+) create mode 100644 examples/entries.py diff --git a/examples/entries.py b/examples/entries.py new file mode 100644 index 0000000..ec013f8 --- /dev/null +++ b/examples/entries.py @@ -0,0 +1,672 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Every kernel under ``examples/`` that this repository's simulator gates on. + +Declaring a kernel is adding a row to ``ENTRIES`` at the bottom of this file: the +kernel's ``.mlir``, the function in it, the parameters to gate at, the arguments to +call it with, and a reference to compare its output against. Nothing else — no new +file, no registration call, no test. + +Two things are worth knowing before adding a row. + +**The arguments are data.** ``tensors`` is a mapping from the kernel's argument +names to specs — ``normal``, ``zeros``, ``full``, ``tile``, ``arange``, +``integers``, ``asarray``, ``param`` — resolved against ``gate_params``; a bare +string forwards a parameter unchanged, and a callable ``(params, rng)`` is the +escape hatch for input no spec expresses. ``ktir_cpu/kernelentry/tensorspec.py`` +holds the vocabulary and the reason each draw is seeded the way it is. + +**A reference computes in f32 or wider, and does not fold the way the kernel +folds.** Both halves matter. Every kernel here holds its intermediates in f16, +so a reference evaluated in f16 reproduces the same overflow and then agrees with +the kernel about a wrong answer — which is what the end-to-end tests in +``tests/test_examples.py`` do when they cast a result back to f16 before comparing. +And a reference written in the kernel's own decomposition — summing shards the way +the shards are summed, carrying a running softmax denominator tile by tile — +checks the arithmetic while assuming the decomposition, which is usually the part +under test. Write the short form of the answer. + +``docs/kernelentry.md`` is the contributor's entry point; this file is where the +kernels are. +""" + +from __future__ import annotations + +import math +from typing import Any, Dict + +import numpy as np + +from ktir_cpu.kernelentry import KernelEntry, register_entry +from ktir_cpu.kernelentry.tensorspec import ( + FAN_IN, arange, asarray, full, integers, normal, param, tile, zeros, +) + +# --------------------------------------------------------------------------- +# Argument specs no vocabulary covers +# --------------------------------------------------------------------------- + +def rope_cos(params: Dict[str, Any], rng) -> np.ndarray: + """The RoPE cosine table, built in f64 before rounding to the kernel's f16. + + The angles are input, not part of what is checked: computed at f16 precision + the reference and the kernel would rotate by slightly different amounts, and + the difference would be charged to the kernel. + """ + return np.cos(_rope_angles(params)).astype(np.float16) + + +def rope_sin(params: Dict[str, Any], rng) -> np.ndarray: + """The RoPE sine table; see :func:`rope_cos`.""" + return np.sin(_rope_angles(params)).astype(np.float16) + + +def _rope_angles(params: Dict[str, Any]) -> np.ndarray: + half = params["D"] // 2 + freqs = 10000.0 ** (-np.arange(half, dtype=np.float64) * 2.0 / params["D"]) + return np.outer(np.arange(params["S"], dtype=np.float64), freqs) + + +def ridge_row(params: Dict[str, Any], rng) -> np.ndarray: + """A row whose maximum is in the middle, repeated down the rows. + + Ordered this way on purpose: on a monotonic row, a fold that returned its + final input rather than the running maximum would agree with the reference. + """ + cols = params["cols"] + row = np.linspace(-2.0, 2.0, cols, dtype=np.float16) + row[cols // 2:] = row[cols // 2:][::-1] + return np.broadcast_to(row, (params["rows"], cols)).copy() + + +def padded_rows(params: Dict[str, Any], rng) -> np.ndarray: + """Real values up to ``n_real_cols``, then ``-inf`` to the block width. + + The padding is the point of the full-size softmax entry: the lowering pads a + row to the block width, and a kernel that included the padding in its + denominator would still produce a plausible distribution. + """ + rows, cols = params["n_rows"], params["n_cols"] + real = params["n_real_cols"] + out = np.full((rows, cols), -np.inf, dtype=np.float16) + out[:, :real] = rng.standard_normal((rows, real)).astype(np.float16) + return out + + +#: ``indexed_add``'s views, which are baked into its MLIR rather than derived +#: from its parameters. +INDEXED_ADD_X = (128, 64, 8, 128) +INDEXED_ADD_Y = (2, 32, 8, 128) + + +# --------------------------------------------------------------------------- +# References +# --------------------------------------------------------------------------- + +def matmul_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """``C = A @ B``, accumulated in f32 across all of K at once. + + The kernel folds K in f16 — sixteen accumulation steps at + ``matmul_fwd_ktir``'s shape — and a reference folding the same way would + reproduce the same rounding and then agree about a wrong answer. + """ + a = np.asarray(tensors["a_ptr"], dtype=np.float32) + b = np.asarray(tensors["b_ptr"], dtype=np.float32) + return {"c_ptr": a @ b} + + +def softmax_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """Row-wise softmax with the shift, the exponential and the sum all in f32. + + The kernel subtracts the row maximum before exponentiating, so what needs + widening is not the exponential but the denominator: an f16 sum of values near + 1 loses the low bits of every term after the first few. Any ``-inf`` padding + is left in rather than sliced off — ``exp`` of it is exactly zero, so both + sides are asked the same question about the padded columns. + """ + x = np.asarray(tensors["input_ptr"], dtype=np.float32) + e = np.exp(x - x.max(axis=1, keepdims=True)) + return {"output_ptr": e / e.sum(axis=1, keepdims=True)} + + +def vector_add_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + x = np.asarray(tensors["x_ptr"], dtype=np.float64) + y = np.asarray(tensors["y_ptr"], dtype=np.float64) + return {"output_ptr": x + y} + + +def swiglu_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """``x + (silu(x @ W_gate) * (x @ W_up)) @ W_down``, the whole block in f32. + + Written against the unsharded expression, which is what makes it a check of + ``ffn_swiglu_4core``'s sharding: a reference summing four 256-wide partials + the way that kernel does would agree with it about a wrong split. + """ + x = np.asarray(tensors["x_ptr"], dtype=np.float32) + w_gate = np.asarray(tensors["w_gate_ptr"], dtype=np.float32) + w_up = np.asarray(tensors["w_up_ptr"], dtype=np.float32) + w_down = np.asarray(tensors["w_down_ptr"], dtype=np.float32) + gate = x @ w_gate + silu = gate / (1.0 + np.exp(-gate)) + return {"out_ptr": x + (silu * (x @ w_up)) @ w_down} + + +def row_sum_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """The row sum, broadcast back across the row, accumulated in f32.""" + data = np.asarray(tensors["arg0"], dtype=np.float32) + return {"arg0": np.broadcast_to( + data.sum(axis=1, keepdims=True), data.shape).copy()} + + +def row_max_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """The row maximum, broadcast back across the row.""" + data = np.asarray(tensors["arg0"], dtype=np.float32) + return {"arg0": np.broadcast_to( + data.max(axis=1, keepdims=True), data.shape).copy()} + + +def scalar_broadcast_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """Exact rather than approximate: broadcasting a value is the one output + where any difference at all is a defect.""" + return {"out_ptr": np.full((params["out_rows"], params["out_cols"]), + params["value"], dtype=np.float32)} + + +def rope_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """The half-layout rotation in f32, over the same tables the kernel was given. + + y[:, :D/2] = x[:, :D/2] * cos - x[:, D/2:] * sin + y[:, D/2:] = x[:, :D/2] * sin + x[:, D/2:] * cos + """ + H, S, D = params["H"], params["S"], params["D"] + half = D // 2 + x = np.asarray(tensors["x_ptr"], dtype=np.float32).reshape(H, S, D) + cos = np.asarray(tensors["cos_ptr"], dtype=np.float32)[np.newaxis] + sin = np.asarray(tensors["sin_ptr"], dtype=np.float32)[np.newaxis] + first, second = x[:, :, :half], x[:, :, half:] + y = np.concatenate([first * cos - second * sin, + first * sin + second * cos], axis=-1) + return {"out_ptr": y.reshape(H * S, D)} + + +def layernorm_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """Mean, rstd and the normalised rows, all in f32. + + Widening matters more here than in a kernel with one output: the kernel folds + both the sum and the sum of squares over 8192 f16 values, and a reference + folding them the same way would agree with whatever those two reductions + drifted to. ``eps`` is read from the parameters rather than repeated, so + changing it in the row changes both sides. + """ + x = np.asarray(tensors["X"], dtype=np.float32) + w = np.asarray(tensors["W"], dtype=np.float32) + b = np.asarray(tensors["B"], dtype=np.float32) + mean = x.mean(axis=1) + rstd = 1.0 / np.sqrt(x.var(axis=1) + params["eps"]) + return {"Y": (x - mean[:, None]) * rstd[:, None] * w + b, + "Mean": mean, "Rstd": rstd} + + +def indexed_add_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """The gather done in numpy, with the add in f32. + + The gather itself is exact — it moves f16 values without arithmetic — so the + widening is only for the add. Getting the *indices* wrong is the failure this + reference is really for, and no tolerance hides it: two different rows of + ``x`` are uncorrelated. + """ + x = np.asarray(tensors["x_ptr"], dtype=np.float32) + y = np.asarray(tensors["y_ptr"], dtype=np.float32) + index = np.asarray(tensors["index_ptr"], dtype=np.intp) + start = params["dim1_start"] + rows = x[index][:, start:start + INDEXED_ADD_Y[1], :, :] + return {"output_ptr": rows + y} + + +def sdpa_2d_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """Attention with every intermediate in f32. + + Widening all of them rather than only the matmuls is deliberate: the softmax + denominator sums 32 terms, and an f32 reference that folded that sum in f16 + would reproduce the kernel's rounding in the one place the two are most + likely to differ. + """ + q = np.asarray(tensors["q_ptr"], dtype=np.float32) + k = np.asarray(tensors["k_ptr"], dtype=np.float32) + v = np.asarray(tensors["v_ptr"], dtype=np.float32) + scores = (q @ k.T) * np.float32(params["scale"]) + e = np.exp(scores - scores.max(axis=1, keepdims=True)) + return {"output_ptr": (e / e.sum(axis=1, keepdims=True)) @ v} + + +def paged_attention_reference(params: Dict[str, Any], + tensors: Dict[str, Any]) -> Dict[str, Any]: + """Causal attention at every grid position, in f32 and untiled. + + Two differences from the reference in ``tests/test_examples.py``, both load + bearing. It covers all 32 grid positions rather than the first, because a + mask indexed by query position is exactly what can be right at ``pid0 = 0`` + and wrong everywhere else. And it reduces over the whole masked row at once + instead of carrying a running maximum and denominator tile by tile: the tiled + form is the kernel's own algorithm, so a reference written that way checks the + arithmetic and assumes the decomposition. + """ + q = np.asarray(tensors["query_ptr"], dtype=np.float32) + k_cache = np.asarray(tensors["key_cache_ptr"], dtype=np.float32) + v_cache = np.asarray(tensors["value_cache_ptr"], dtype=np.float32) + table = np.asarray(tensors["block_tables_ptr"]) + scale = float(tensors["scale"]) + context_len = int(tensors["context_len"]) + num_tiles = int(tensors["num_tiles"]) + block_q = params["block_q"] + per_kv = params["num_query_heads"] // params["num_kv_heads"] + + out = np.zeros_like(q) + pages = table[0, :num_tiles] + for pid1 in range(params["num_kv_heads"]): + # (num_tiles * blk_size, head_size): the pages this KV head reads, in the + # order block_tables gives them. + keys = np.concatenate([k_cache[p, :, pid1, :] for p in pages]) + values = np.concatenate([v_cache[p, :, pid1, :] for p in pages]) + for pid0 in range(params["num_tokens"] // block_q): + rows = slice(pid0 * block_q, (pid0 + 1) * block_q) + heads = slice(pid1 * per_kv, (pid1 + 1) * per_kv) + tile_q = q[rows, heads, :].reshape(block_q * per_kv, -1) + scores = tile_q @ keys.T * scale + # Query row r sits at absolute position context_len + pid0*block_q + # + r // per_kv, and may not see past it. + positions = context_len + pid0 * block_q + np.arange( + block_q * per_kv) // per_kv + scores[np.arange(keys.shape[0])[None, :] > positions[:, None]] = -np.inf + p = np.exp(scores - scores.max(axis=1, keepdims=True)) + attended = (p @ values) / p.sum(axis=1, keepdims=True) + out[rows, heads, :] = attended.reshape(block_q, per_kv, -1) + return {"output_ptr": out} + + +# --------------------------------------------------------------------------- +# The kernels +# --------------------------------------------------------------------------- + +ENTRIES = [ + # --- examples/ktir: hand-written IR ----------------------------------- + KernelEntry( + # Three matmuls and a sigmoid in one kernel: the longest chain of f16 + # intermediates here at a gate-sized shape. The weights are unit variance + # and unscaled, which saturates the sigmoid on purpose — `gate` spans + # [-24.3, +21.2], and 14 of its 128 f16 values leave f16 range in the + # exponential: 11 overflow to inf, where 1 / inf makes sigmoid 0, and 3 + # underflow to exactly 0, where sigmoid is exactly 1. Both are the value + # the f32 reference converges to there, so the output still agrees with + # it. It is the only entry that exercises saturation, and + # `ktir_cpu/ops/_helpers.py` reports the overflowing 11 as a + # RuntimeWarning while it runs. + name="ffn_swiglu", func="ffn_swiglu", path="ktir/ffn_swiglu.mlir", + gate_params={"seq": 1, "d_model": 64, "d_ffn": 128}, + tensors={ + "x_ptr": normal(("seq", "d_model")), + "w_gate_ptr": normal(("d_model", "d_ffn")), + "w_up_ptr": normal(("d_model", "d_ffn")), + "w_down_ptr": normal(("d_ffn", "d_model")), + "out_ptr": zeros(("seq", "d_model")), + }, + reference=swiglu_reference, + # The one entry that declares its own tolerance, and the unscaled weights + # are why. `silu * up` reaches 331 and the down projection sums 128 of + # those, so every output element accumulates through a peak far above the + # value it ends on — 183 at the smallest of the 64, 2095 at the largest. + # A relative tolerance is charged against the value, not against the peak + # the sum passed through, so an element that cancels escapes it: element + # 30 comes down from a peak of 338 to -1.85, and two f16 rounding steps at + # 338 put it 0.34 away, 19% off. atol has to cover that at whichever + # element cancels, so it is set from the largest peak in the tensor rather + # than fitted to the one that cancels here: 2 * 2095 * 2^-11 = 2.05. + # Only 2 of the 64 elements need it at all and the worst needs 0.31, so + # the declared pair sits ~6x above what this seed asks for — deliberately, + # because a shape change re-rolls the tensors and moves which element + # cancels. It is a correctness-leg number only: the cost leg reads the + # same run, and cycle counts there are structural rather than + # data-dependent, so no tolerance of any width moves them. + tolerance={"out_ptr": (2e-2, 2.0)}, + ), + KernelEntry( + # The same block with the hidden dimension sharded over 4 cores: `x` is + # replicated, each core owns a 256-wide slice of W_gate / W_up and the + # matching 256-row slice of W_down, and the four [4, 256] partials fold + # through inter_tile_produce / inter_tile_reduce before the residual. + # Weights scaled by fan-in, unlike the single-core entry: `gate` widens + # with sqrt(d_model), so unscaled at 256 it would have std 16.8 against + # the single-core entry's 8.2 and 2069 of its 4096 values would saturate + # instead of 22 of 128. A comparison mostly between saturated values + # would stop being about whether the shards folded. + name="ffn_swiglu_4core", func="ffn_swiglu_4core", + path="ktir/ffn_swiglu_4core.mlir", + gate_params={"seq": 4, "d_model": 256, "d_ffn": 1024}, + tensors={ + "x_ptr": normal(("seq", "d_model")), + "w_gate_ptr": normal(("d_model", "d_ffn"), scale=FAN_IN), + "w_up_ptr": normal(("d_model", "d_ffn"), scale=FAN_IN), + "w_down_ptr": normal(("d_ffn", "d_model"), scale=FAN_IN), + "out_ptr": zeros(("seq", "d_model")), + }, + reference=swiglu_reference, + ), + KernelEntry( + # A linalg.reduce combiner written as an explicit (%in, %out) region + # rather than the { arith.addf } shorthand, which is the form the Triton + # Spyre ConvertTTReduce pass emits. One argument, input and output both, + # so the sum is written back over the data it came from — which is why the + # reference is handed a pristine rebuild rather than what the run left. + name="reduce_generic", func="reduce_explicit_region", + path="ktir/reduce_generic.mlir", + gate_params={"rows": 1, "cols": 4}, + tensors={"arg0": arange(("rows", "cols"), start=1)}, + reference=row_sum_reference, + ), + KernelEntry( + # `max` written as arith.cmpf ogt + arith.select rather than a single + # arith.maximumf, so the result cannot come from recognising a combiner op + # by name: every op in the region has to run. + name="reduce_multiop", func="reduce_multiop", + path="ktir/reduce_multiop.mlir", + gate_params={"rows": 1, "cols": 8}, + tensors={"arg0": ridge_row}, + reference=row_max_reference, + ), + KernelEntry( + # Collapses a tensor<1x1xf16> to a scalar tensor and broadcasts it, + # which is the shape the Triton -> KTIR lowering emits for a 1x1 + # broadcast. + name="scalar_broadcast", func="scalar_broadcast", + path="ktir/scalar_broadcast.mlir", + gate_params={"value": 2.5, "out_rows": 4, "out_cols": 64}, + tensors={ + "in_ptr": full((1, 1), "value"), + "out_ptr": zeros(("out_rows", "out_cols")), + }, + reference=scalar_broadcast_reference, + ), + + # --- examples/latency: reduced shapes, which is why they are the gate --- + KernelEntry( + # A tiled matmul on a [2, 2] grid. Every parameter here is baked into the + # MLIR — view sizes, grid, tile constants — so they are parameters of this + # row rather than of the kernel: changing one does not change the kernel, + # it makes this row disagree with it, which out.c_ptr.reference reports. + name="matmul_small", func="matmul_kernel_small", + path="latency/matmul_small.mlir", + gate_params={"M": 16, "N": 64, "K": 64, + "BLOCK_SIZE_M": 8, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + tensors={ + "a_ptr": normal(("M", "K")), + "b_ptr": normal(("K", "N")), + "c_ptr": zeros(("M", "N")), + "K": "K", + "BLOCK_SIZE_M": "BLOCK_SIZE_M", + "BLOCK_SIZE_N": "BLOCK_SIZE_N", + "BLOCK_SIZE_K": "BLOCK_SIZE_K", + }, + reference=matmul_reference, + ), + KernelEntry( + # Half-layout RoPE on a [4, 2] grid at LLaMA-8B / Granite-8B shapes: 40 + # heads, 4096 positions, D=128, flattened to [H*S, D]. The cos/sin tables + # are inputs rather than constants, so the kernel is checked against the + # rotation it was asked for and not against whatever tables it read. + name="rope_fwd_4x2", func="rope_fwd_kernel", + path="latency/rope_fwd_4x2.mlir", + gate_params={"H": 40, "S": 4096, "D": 128}, + tensors={ + "x_ptr": normal((lambda p: p["H"] * p["S"], "D")), + "cos_ptr": rope_cos, + "sin_ptr": rope_sin, + "out_ptr": zeros((lambda p: p["H"] * p["S"], "D")), + }, + reference=rope_reference, + ), + KernelEntry( + # softmax_fwd_ktir.mlir at reduced size, which is what makes it a gate + # kernel: the same row-wise max / exp / sum / divide chain over 64 rows + # instead of 4096. tests/test_latency.py already drives it for its cost + # breakdown; this row is what says those figures came from a kernel + # computing the right thing. + name="softmax_small", func="softmax_kernel_small", + path="latency/softmax_small.mlir", + gate_params={"n_rows": 64, "n_cols": 64}, + tensors={ + "output_ptr": zeros(("n_rows", "n_cols")), + "input_ptr": normal(("n_rows", "n_cols")), + "n_rows": "n_rows", + }, + reference=softmax_reference, + ), + KernelEntry( + # The same kernel with each linalg.reduce combiner written as an explicit + # region instead of the shorthand. Same inputs and same reference as the + # row above, which is the whole point of the second file: two spellings + # fed different inputs could be reported as agreeing without ever + # computing the same thing. + name="softmax_small_explicit", func="softmax_kernel_small_explicit", + path="latency/softmax_small_explicit.mlir", + gate_params={"n_rows": 64, "n_cols": 64}, + tensors={ + "output_ptr": zeros(("n_rows", "n_cols")), + "input_ptr": normal(("n_rows", "n_cols")), + "n_rows": "n_rows", + }, + reference=softmax_reference, + ), + + # --- examples/sdsc: decode attention, split across cores --------------- + KernelEntry( + # C = A @ B for (1 x 8192) @ (8192 x 128): the output split x2 and the KV + # contraction split x16, so partial sums fold across cores in two strided + # reduce groups. The entry that exercises what a single-core kernel cannot + # — comm bytes in the derivation, and a reference that has to tolerate an + # f16 fold, since nothing in the kernel is wider than f16 and each of the + # 16 fold steps rounds its running sum back. + name="sdpa_pv_ksplit", func="sdpa_pv_ksplit", + path="sdsc/sdpa_pv_ksplit.mlir", + gate_params={"M": 1, "N": 128, "K": 8192}, + tensors={ + "a_ptr": normal(("M", "K")), + "b_ptr": normal(("K", "N")), + "c_ptr": zeros(("M", "N")), + }, + reference=matmul_reference, + ), + + # --- examples/triton-ktir: captured Triton -> KTIR output -------------- + KernelEntry( + # An indirect gather on the leading axis: the row of `x` each core reads + # comes out of an i64 index tensor at run time, so this is the one + # declared kernel whose addresses are not a function of its grid + # position. It reaches ktdp.construct_indirect_access_tile. + name="indexed_add", func="indexed_add_kernel", + path="triton-ktir/indexed_add.mlir", + gate_params={"dim1_start": 0, "index": (3, 7)}, + tensors={ + "x_ptr": normal(INDEXED_ADD_X), + "y_ptr": normal(INDEXED_ADD_Y), + "index_ptr": asarray("index", "i64"), + "output_ptr": zeros(INDEXED_ADD_Y), + "dim1_start": "dim1_start", + }, + reference=indexed_add_reference, + ), + KernelEntry( + # Layer norm at full size, and the one kernel here with three outputs: + # the normalised rows, plus the per-row mean and reciprocal standard + # deviation the backward pass would consume. All three are referenced, + # which tests/test_examples.py does not do for Rstd — and Rstd is the + # interesting one, because it is where a reduction over 8192 f16 elements + # either survives a division and a square root or does not. + # + # W and B are a weight and bias vector, but the MLIR views them at + # [n_rows, n_cols] — the same row read once per row of X rather than + # broadcast — so they are declared at that shape, built by tiling one row. + # The row count is deliberately not a multiple of the 32-core grid: + # 1151 = 35*32 + 31, so the last core takes a short tail. + name="layernorm_fwd_ktir", func="_layer_norm_fwd_fused", + path="triton-ktir/layernorm_fwd_ktir.mlir", + gate_params={"n_rows": 1151, "n_cols": 8192, "eps": 1e-5, + "BLOCK_SIZE": 1024}, + tensors={ + "X": normal(("n_rows", "n_cols")), + "Y": zeros(("n_rows", "n_cols")), + "W": tile(normal("n_cols"), ("n_rows", 1)), + "B": tile(normal("n_cols"), ("n_rows", 1)), + "Mean": zeros("n_rows"), + "Rstd": zeros("n_rows"), + "N": "n_cols", + "eps": "eps", + "BLOCK_SIZE": "BLOCK_SIZE", + }, + reference=layernorm_reference, + ), + KernelEntry( + # Split-K matmul at full size. matmul_small is the same kernel shape at + # [16, 64, 64]; what only appears here is the depth of the accumulation — + # K=2048 over BLOCK_SIZE_K=128 is 16 f16 accumulation steps, and the + # rounding that produces is proportional to the output magnitude rather + # than bounded by it. + name="matmul_fwd_ktir", func="matmul_kernel", + path="triton-ktir/matmul_fwd_ktir.mlir", + gate_params={"M": 64, "N": 8192, "K": 2048, + "BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 512, + "BLOCK_SIZE_K": 128}, + tensors={ + "a_ptr": normal(("M", "K")), + "b_ptr": normal(("K", "N")), + "c_ptr": zeros(("M", "N")), + "K": "K", + "BLOCK_SIZE_M": "BLOCK_SIZE_M", + "BLOCK_SIZE_N": "BLOCK_SIZE_N", + "BLOCK_SIZE_K": "BLOCK_SIZE_K", + }, + reference=matmul_reference, + ), + KernelEntry( + # Unified attention over a paged KV cache: block_tables names which of the + # 64 cache pages each step reads, so the kernel's traffic is decided at + # run time by data rather than by its loop bounds. The page table is drawn + # over the whole cache rather than set to identity, so the pages a step + # reads are not the pages a direct view would have reached — which is what + # makes the indirect path observable in the output at all. + name="paged_attention", func="kernel_unified_attention_spyre_2d", + path="triton-ktir/paged_attention.mlir", + gate_params={ + "num_tokens": 8, "num_query_heads": 32, "num_kv_heads": 8, + "head_size": 128, "num_blks": 64, "blk_size": 16, + "max_num_blocks_per_seq": 16, "block_q": 2, "num_tiles": 8, + "context_len": 120, + }, + tensors={ + "output_ptr": zeros(("num_tokens", "num_query_heads", "head_size")), + "query_ptr": normal(("num_tokens", "num_query_heads", "head_size")), + "key_cache_ptr": normal( + ("num_blks", "blk_size", "num_kv_heads", "head_size")), + "value_cache_ptr": normal( + ("num_blks", "blk_size", "num_kv_heads", "head_size")), + "block_tables_ptr": integers((1, "max_num_blocks_per_seq"), + high="num_blks"), + "cur_batch_start_index": 0, + "block_table_offset": 0, + "num_tiles": "num_tiles", + "context_len": "context_len", + "scale": lambda p, rng: 1.0 / math.sqrt(p["head_size"]), + }, + reference=paged_attention_reference, + ), + KernelEntry( + # softmax(Q @ K^T * scale) @ V on one core, so the whole attention chain — + # two matmuls with a row-wise max, exp and sum between them — runs with no + # cross-core fold. The scale is a constant in the MLIR, which is why it is + # a parameter of this row: changing it here makes the row disagree with + # the kernel. + name="sdpa_2d", func="sdpa_kernel_2d", + path="triton-ktir/sdpa_2d.mlir", + gate_params={"n_rows": 32, "head_dim": 64, "scale": 0.125}, + tensors={ + "q_ptr": normal(("n_rows", "head_dim")), + "k_ptr": normal(("n_rows", "head_dim")), + "v_ptr": normal(("n_rows", "head_dim")), + "output_ptr": zeros(("n_rows", "head_dim")), + }, + reference=sdpa_2d_reference, + ), + KernelEntry( + # Row-wise softmax at full size. softmax_small is the same chain at + # [64, 64] and is the one the gate leans on; this row exists because the + # padding is only real here — see padded_rows. + name="softmax_fwd_ktir", func="softmax_kernel", + path="triton-ktir/softmax_fwd_ktir.mlir", + gate_params={"n_rows": 4096, "n_cols": 1024, "n_real_cols": 778}, + tensors={ + "output_ptr": zeros(("n_rows", "n_cols")), + "input_ptr": padded_rows, + "n_rows": "n_rows", + }, + reference=softmax_reference, + ), + KernelEntry( + # A symbolic extent: the views are memref and the access tile's + # coordinate set covers d0 in [0, 1023], so n_elements masks the tail at + # run time rather than being baked in. Gated at 1024, the value at which + # the mask is a no-op — masking is what the smaller extents in + # tests/test_examples.py exercise, and this ledger's question is whether + # the kernel computes the right thing, not how many extents it does so at. + name="vector_add_dynamic_ktir", func="add_kernel_dynamic", + path="triton-ktir/vector_add_dynamic_ktir.mlir", + gate_params={"n_elements": 1024}, + tensors={ + "x_ptr": normal("n_elements", "f32"), + "y_ptr": normal("n_elements", "f32"), + "output_ptr": zeros("n_elements", "f32"), + # Read as an i32 function argument, not as an index constant. + "n_elements": param("n_elements", "i32"), + }, + reference=vector_add_reference, + ), + KernelEntry( + # The smallest kernel on this path: one load per input, one linalg.add, + # one store. Its value here is as the floor — anything this ledger reports + # about a larger kernel it also reports about this one. + name="vector_add_ktir", func="add_kernel", + path="triton-ktir/vector_add_ktir.mlir", + gate_params={"n_elements": 4096, "BLOCK_SIZE": 128}, + tensors={ + "x_ptr": normal("n_elements"), + "y_ptr": normal("n_elements"), + "output_ptr": zeros("n_elements"), + "BLOCK_SIZE": "BLOCK_SIZE", + }, + reference=vector_add_reference, + ), +] + +for _entry in ENTRIES: + register_entry(_entry) From aba47bf6c94ed3784758f7ba4c7d819d152e47d7 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:16 -0400 Subject: [PATCH 6/9] Commit the generated support, op and cost reports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generated and committed on the same discipline as a lock file: regenerating them is one command, and the gate compares them verbatim, so adding a kernel without regenerating them is a red build rather than a document that quietly goes stale. `docs/kernel_support.md` and `docs/supported_ops.md` are the same repository seen per kernel and per op; `docs/kernel_cost.md` is the committed cost attribution the `cost.derivation` claim compares against. Between them they carry figures nothing here stated before: 490 op occurrences across thirty-three kernels with a handler for every one, 106 handlers of which 102 are frontend-reachable and four are by design not, and 52 of those 106 exercised by an example. Four kernels are not accepted by the MLIR frontend — the three `ring_reduce*` for the typing gap recorded earlier in this branch, and `nested_yield.ktir` for a `construct_memory_view` with no `coordinate_set` — which nothing said before either: the frontend suite covers eleven kernels by inheriting regex test classes and none of these four is among them, so this is a coverage hole rather than a break. The op document's cost column prints the pricing decision rather than the registry's default, so an op priced zero with nobody having decided that reads differently from one that is free. `docs/kernel_cost.md` is worth reading rather than skimming. 72.7% of `matmul_small`'s HBM traffic is the B operand, at a traffic ratio of 3.667 over what its tensors logically require; `indexed_add` reports 0.02320 for the same figure, which is not an error — it reads two of 128 slices, and clamping the ratio at 1 would erase exactly the fact worth seeing. The spread across the eighteen is 0.02320 to 32. Signed-off-by: WarningRan --- docs/kernel_cost.md | 746 +++++++++++++++++++++++++++++++++++++++++ docs/kernel_support.md | 78 +++++ docs/supported_ops.md | 325 ++++++++++++++++++ 3 files changed, 1149 insertions(+) create mode 100644 docs/kernel_cost.md create mode 100644 docs/kernel_support.md create mode 100644 docs/supported_ops.md diff --git a/docs/kernel_cost.md b/docs/kernel_cost.md new file mode 100644 index 0000000..4b9159f --- /dev/null +++ b/docs/kernel_cost.md @@ -0,0 +1,746 @@ +# Kernel cost derivations + + + +One section per declared kernel, at the parameters its declaration gates on. Each +is what the cost model reports, **not an independent check of it**: the figures and +the totals they sum to come from the same code, so a mis-charged op moves both +together and the section stays self-consistent. Whether the model charges +correctly is asked of the model rather than of any one kernel, in +`tests/test_latency.py`, against hand-counted bytes, FLOPs and cycles. + +What a section is for is the composition — which tensor dominates the traffic, what +the ratio to the kernel's declared footprint is, which unit the cycles land on. +Confirm that against what the kernel is supposed to do; a sentence in prose that +disagrees with the table below it does not survive review. Every figure is +formatted to four significant digits, and nothing here varies between runs at fixed +parameters, so a diff means the cost changed. + +## ffn_swiglu + +```text +kernel: ffn_swiglu function: ffn_swiglu +parameters: d_ffn=128, d_model=64, seq=1 +grid cores active: 1 + +hbm_bytes = 49,408 + w_down_ptr 16,384 33.2% ktdp.load x1 + w_gate_ptr 16,384 33.2% ktdp.load x1 + w_up_ptr 16,384 33.2% ktdp.load x1 + out_ptr 128 0.3% ktdp.store x1 + x_ptr 128 0.3% ktdp.load x1 +unique_bytes = 49,408 traffic_ratio = 1x + +flops = 49,984.0 +arithmetic_intensity = flops / hbm_bytes = 49,984.0 / 49,408 = 1.012 + +kernel_cycles = 453 bottleneck = memory + critical core compute 67 memory 386 comm 0 + compute_float 11 + compute_matmul 48 + compute_transcendental 8 + +cycles charged per op type, critical core: + ktdp.load n=4 385 + linalg.matmul n=3 48 + math.exp n=1 8 + arith.mulf n=2 4 + arith.addf n=2 3 + arith.divf n=1 2 + arith.negf n=1 2 + ktdp.store n=1 1 + arith.constant n=2 0 + ktdp.construct_access_tile n=5 0 + ktdp.construct_memory_view n=5 0 + return n=1 0 + tensor.empty n=3 0 +``` + +## ffn_swiglu_4core + +```text +kernel: ffn_swiglu_4core function: ffn_swiglu_4core +parameters: d_ffn=1024, d_model=256, seq=4 +grid cores active: 4 + +hbm_bytes = 1,583,104 + w_down_ptr 524,288 33.1% ktdp.load x4 + w_gate_ptr 524,288 33.1% ktdp.load x4 + w_up_ptr 524,288 33.1% ktdp.load x4 + x_ptr 8,192 0.5% ktdp.load x4 + out_ptr 2,048 0.1% ktdp.store x1 +unique_bytes = 1,576,960 traffic_ratio = 1.004x + +comm_bytes = 24,576 + ktdp.inter_tile_reduce n=4 24,576 + +flops = 6,332,416.0 +arithmetic_intensity = flops / hbm_bytes = 6,332,416.0 / 1,583,104 = 4 + +kernel_cycles = 14,256 bottleneck = memory + critical core compute 1,744 memory 12,416 comm 96 + compute_float 144 + compute_int 0 + compute_matmul 1,536 + compute_transcendental 64 + +cycles charged per op type, critical core: + ktdp.load n=4 12,352 + linalg.matmul n=3 1,536 + ktdp.inter_tile_reduce n=1 96 + ktdp.store n=1 64 + math.exp n=1 64 + linalg.add n=3 48 + arith.addf n=2 32 + arith.mulf n=2 32 + arith.divf n=1 16 + arith.negf n=1 16 + arith.addi n=9 0 + arith.cmpi n=1 0 + arith.constant n=9 0 + arith.muli n=1 0 + ktdp.construct_access_tile n=5 0 + ktdp.construct_distributed_memory_view n=3 0 + ktdp.construct_memory_view n=14 0 + ktdp.get_compute_tile_id n=1 0 + ktdp.inter_tile_produce n=1 0 + ktdp.yield_partial n=1 0 + ktdp.yield_reduced n=3 0 + linalg.fill n=1 0 + return n=1 0 + scf.if n=1 0 + tensor.collapse_shape n=1 0 + tensor.empty n=7 0 + tensor.expand_shape n=1 0 +``` + +## indexed_add + +```text +kernel: indexed_add function: indexed_add_kernel +parameters: dim1_start=0, index=(3, 7) +grid cores active: 16 + +hbm_bytes = 395,264 + x_ptr 133,120 33.7% ktdp.load x16 + output_ptr 131,072 33.2% ktdp.store x16 + y_ptr 131,072 33.2% ktdp.load x16 +unique_bytes = 17,039,376 traffic_ratio = 0.02320x (below 1: the kernel reaches only part of a declared tensor) + +flops = 65,536.0 +arithmetic_intensity = flops / hbm_bytes = 65,536.0 / 395,264 = 0.1658 + +kernel_cycles = 3,152 bottleneck = memory + critical core compute 64 memory 3,088 comm 0 + compute_float 64 + +cycles charged per op type, critical core: + ktdp.load n=2 2,064 + ktdp.store n=1 1,024 + arith.addf n=1 64 + arith.constant n=1 0 + ktdp.construct_access_tile n=2 0 + ktdp.construct_indirect_access_tile n=1 0 + ktdp.construct_memory_view n=4 0 + ktdp.get_compute_tile_id n=1 0 + return n=1 0 +``` + +## layernorm_fwd_ktir + +```text +kernel: layernorm_fwd_ktir function: _layer_norm_fwd_fused +parameters: BLOCK_SIZE=1024, eps=1e-05, n_cols=8192, n_rows=1151 +grid cores active: 32 + +hbm_bytes = 113,442,560 + X 56,573,952 49.9% ktdp.load x27624 + B 18,857,984 16.6% ktdp.load x9208 + W 18,857,984 16.6% ktdp.load x9208 + Y 18,857,984 16.6% ktdp.store x9208 + Mean 147,328 0.1% ktdp.store x1151 + Rstd 147,328 0.1% ktdp.store x1151 +unique_bytes = 75,436,540 traffic_ratio = 1.504x + +flops = 77,794,939.0 +arithmetic_intensity = flops / hbm_bytes = 77,794,939.0 / 113,442,560 = 0.6858 + +kernel_cycles = 925,060 bottleneck = memory + critical core compute 38,020 memory 887,040 comm 0 + compute_float 38,018 + compute_transcendental 2.250 + +cycles charged per op type, critical core: + ktdp.load n=1440 737,280 + ktdp.store n=360 149,760 + arith.addf n=1692 14,977 + arith.mulf n=864 13,824 + arith.subf n=576 9,216 + math.sqrt n=36 2.250 + arith.divf n=108 1.688 + arith.constant n=112 0 + arith.index_cast n=36 0 + arith.sitofp n=36 0 + ktdp.construct_access_tile n=1800 0 + ktdp.construct_memory_view n=6 0 + ktdp.get_compute_tile_id n=1 0 + linalg.reduce n=72 0 + linalg.yield n=792 0 + return n=1 0 + scf.for n=109 0 + scf.yield n=864 0 + tensor.extract n=72 0 + tensor.splat n=252 0 +``` + +## matmul_fwd_ktir + +```text +kernel: matmul_fwd_ktir function: matmul_kernel +parameters: BLOCK_SIZE_K=128, BLOCK_SIZE_M=32, BLOCK_SIZE_N=512, K=2048, M=64, N=8192 +grid cores active: 32 + +hbm_bytes = 72,351,744 + b_ptr 67,108,864 92.8% ktdp.load x512 + a_ptr 4,194,304 5.8% ktdp.load x512 + c_ptr 1,048,576 1.4% ktdp.store x32 +unique_bytes = 34,865,152 traffic_ratio = 2.075x + +flops = 2,155,872,256.0 +arithmetic_intensity = flops / hbm_bytes = 2,155,872,256.0 / 72,351,744 = 29.80 + +kernel_cycles = 634,880 bottleneck = memory + critical core compute 69,632 memory 565,248 comm 0 + compute_float 4,096 + compute_int 0 + compute_matmul 65,536 + +cycles charged per op type, critical core: + ktdp.load n=32 557,056 + linalg.matmul n=16 65,536 + ktdp.store n=1 8,192 + arith.addf n=16 4,096 + arith.constant n=3 0 + arith.muli n=2 0 + ktdp.construct_access_tile n=33 0 + ktdp.construct_memory_view n=3 0 + ktdp.get_compute_tile_id n=1 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=16 0 + tensor.empty n=16 0 +``` + +## matmul_small + +```text +kernel: matmul_small function: matmul_kernel_small +parameters: BLOCK_SIZE_K=32, BLOCK_SIZE_M=8, BLOCK_SIZE_N=32, K=64, M=16, N=64 +grid cores active: 4 + +hbm_bytes = 45,056 + b_ptr 32,768 72.7% ktdp.load x8 + a_ptr 8,192 18.2% ktdp.load x8 + c_ptr 4,096 9.1% ktdp.store x4 +unique_bytes = 12,288 traffic_ratio = 3.667x + +flops = 133,120.0 +arithmetic_intensity = flops / hbm_bytes = 133,120.0 / 45,056 = 2.955 + +kernel_cycles = 392 bottleneck = memory + critical core compute 40 memory 352 comm 0 + compute_float 8 + compute_int 0 + compute_matmul 32 + +cycles charged per op type, critical core: + ktdp.load n=4 320 + ktdp.store n=1 32 + linalg.matmul n=2 32 + arith.addf n=2 8 + arith.constant n=3 0 + arith.muli n=2 0 + ktdp.construct_access_tile n=5 0 + ktdp.construct_memory_view n=3 0 + ktdp.get_compute_tile_id n=1 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=2 0 + tensor.empty n=2 0 +``` + +## paged_attention + +```text +kernel: paged_attention function: kernel_unified_attention_spyre_2d +parameters: blk_size=16, block_q=2, context_len=120, head_size=128, max_num_blocks_per_seq=16, num_blks=64, num_kv_heads=8, num_query_heads=32, num_tiles=8, num_tokens=8 +grid cores active: 32 + +hbm_bytes = 2,293,760 + key_cache_ptr 1,081,344 47.1% ktdp.load x256 + value_cache_ptr 1,081,344 47.1% ktdp.load x256 + output_ptr 65,536 2.9% ktdp.store x32 + query_ptr 65,536 2.9% ktdp.load x32 +unique_bytes = 4,325,440 traffic_ratio = 0.5303x (below 1: the kernel reaches only part of a declared tensor) + +flops = 17,688,576.0 +arithmetic_intensity = flops / hbm_bytes = 17,688,576.0 / 2,293,760 = 7.712 + +kernel_cycles = 18,928 bottleneck = memory + critical core compute 1,008 memory 17,920 comm 0 + compute_float 407 + compute_int 21 + compute_matmul 512 + compute_transcendental 68 + +cycles charged per op type, critical core: + ktdp.load n=17 17,408 + ktdp.store n=1 512 + linalg.matmul n=16 512 + linalg.generic n=49 179 + arith.mulf n=24 145 + math.exp n=16 68 + arith.addf n=48 17 + arith.maximumf n=48 17 + arith.subf n=16 17 + arith.cmpi n=8 16 + arith.divf n=1 16 + arith.select n=8 16 + arith.addi n=33 4 + arith.divui n=8 1 + arith.bitcast n=1 0 + arith.constant n=8 0 + arith.extf n=17 0 + arith.muli n=10 0 + arith.truncf n=1 0 + ktdp.construct_access_tile n=2 0 + ktdp.construct_indirect_access_tile n=16 0 + ktdp.construct_memory_view n=5 0 + ktdp.get_compute_tile_id n=1 0 + linalg.index n=16 0 + linalg.reduce n=16 0 + linalg.transpose n=8 0 + linalg.yield n=129 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=8 0 + tensor.collapse_shape n=17 0 + tensor.empty n=57 0 + tensor.expand_shape n=1 0 + tensor.splat n=35 0 +``` + +## reduce_generic + +```text +kernel: reduce_generic function: reduce_explicit_region +parameters: cols=4, rows=1 +grid cores active: 1 + +hbm_bytes = 256 + arg0 256 100.0% ktdp.load x1, ktdp.store x1 +unique_bytes = 8 traffic_ratio = 32x + +flops = 4.0 +arithmetic_intensity = flops / hbm_bytes = 4.0 / 256 = 0.01562 + +kernel_cycles = 2.062 bottleneck = memory + critical core compute 0.06250 memory 2 comm 0 + compute_float 0.06250 + +cycles charged per op type, critical core: + ktdp.load n=1 1 + ktdp.store n=1 1 + arith.addf n=3 0.06250 + arith.constant n=2 0 + ktdp.construct_access_tile n=2 0 + ktdp.construct_memory_view n=2 0 + linalg.fill n=1 0 + linalg.reduce n=1 0 + linalg.yield n=3 0 + return n=1 0 + tensor.empty n=1 0 + tensor.extract n=1 0 + tensor.splat n=1 0 +``` + +## reduce_multiop + +```text +kernel: reduce_multiop function: reduce_multiop +parameters: cols=8, rows=1 +grid cores active: 1 + +hbm_bytes = 256 + arg0 256 100.0% ktdp.load x1, ktdp.store x1 +unique_bytes = 16 traffic_ratio = 16x + +flops = 16.0 +arithmetic_intensity = flops / hbm_bytes = 16.0 / 256 = 0.06250 + +kernel_cycles = 2.250 bottleneck = memory + critical core compute 0.2500 memory 2 comm 0 + compute_float 0.2500 + +cycles charged per op type, critical core: + ktdp.load n=1 1 + ktdp.store n=1 1 + arith.cmpf n=4 0.1250 + arith.select n=4 0.1250 + arith.constant n=2 0 + ktdp.construct_access_tile n=2 0 + ktdp.construct_memory_view n=2 0 + linalg.fill n=1 0 + linalg.reduce n=1 0 + linalg.yield n=4 0 + return n=1 0 + tensor.empty n=1 0 + tensor.extract n=1 0 + tensor.splat n=1 0 +``` + +## rope_fwd_4x2 + +```text +kernel: rope_fwd_4x2 function: rope_fwd_kernel +parameters: D=128, H=40, S=4096 +grid cores active: 8 + +hbm_bytes = 85,983,232 + out_ptr 41,943,040 48.8% ktdp.store x1280 + x_ptr 41,943,040 48.8% ktdp.load x1280 + cos_ptr 1,048,576 1.2% ktdp.load x32 + sin_ptr 1,048,576 1.2% ktdp.load x32 +unique_bytes = 84,934,656 traffic_ratio = 1.012x + +flops = 62,914,560.0 +arithmetic_intensity = flops / hbm_bytes = 62,914,560.0 / 85,983,232 = 0.7317 + +kernel_cycles = 794,624 bottleneck = memory + critical core compute 122,880 memory 671,744 comm 0 + compute_float 122,880 + compute_int 0 + +cycles charged per op type, critical core: + ktdp.load n=168 344,064 + ktdp.store n=160 327,680 + arith.mulf n=320 81,920 + arith.addf n=80 20,480 + arith.subf n=80 20,480 + arith.addi n=164 0 + arith.constant n=8 0 + arith.muli n=86 0 + ktdp.construct_access_tile n=328 0 + ktdp.construct_memory_view n=4 0 + ktdp.get_compute_tile_id n=1 0 + scf.for n=5 0 + scf.yield n=80 0 +``` + +## scalar_broadcast + +```text +kernel: scalar_broadcast function: scalar_broadcast +parameters: out_cols=64, out_rows=4, value=2.5 +grid cores active: 1 + +hbm_bytes = 640 + out_ptr 512 80.0% ktdp.store x1 + in_ptr 128 20.0% ktdp.load x1 +unique_bytes = 514 traffic_ratio = 1.245x + +flops = 0 +arithmetic_intensity = flops / hbm_bytes = 0 / 640 = 0 + +kernel_cycles = 5 bottleneck = memory + critical core compute 0 memory 5 comm 0 + +cycles charged per op type, critical core: + ktdp.store n=1 4 + ktdp.load n=1 1 + arith.constant n=1 0 + ktdp.construct_access_tile n=2 0 + ktdp.construct_memory_view n=2 0 + linalg.broadcast n=1 0 + return n=1 0 + tensor.collapse_shape n=1 0 + tensor.empty n=1 0 +``` + +## sdpa_2d + +```text +kernel: sdpa_2d function: sdpa_kernel_2d +parameters: head_dim=64, n_rows=32, scale=0.125 +grid cores active: 1 + +hbm_bytes = 16,384 + k_ptr 4,096 25.0% ktdp.load x1 + output_ptr 4,096 25.0% ktdp.store x1 + q_ptr 4,096 25.0% ktdp.load x1 + v_ptr 4,096 25.0% ktdp.load x1 +unique_bytes = 16,384 traffic_ratio = 1x + +flops = 268,288.0 +arithmetic_intensity = flops / hbm_bytes = 268,288.0 / 16,384 = 16.38 + +kernel_cycles = 528 bottleneck = compute + critical core compute 400 memory 128 comm 0 + compute_float 80 + compute_matmul 256 + compute_transcendental 64 + +cycles charged per op type, critical core: + linalg.matmul n=2 256 + ktdp.load n=3 96 + math.exp n=1 64 + ktdp.store n=1 32 + arith.addf n=6 16 + arith.divf n=1 16 + arith.maximumf n=6 16 + arith.mulf n=1 16 + arith.subf n=1 16 + arith.constant n=4 0 + ktdp.construct_access_tile n=4 0 + ktdp.construct_memory_view n=4 0 + ktdp.get_compute_tile_id n=1 0 + linalg.broadcast n=2 0 + linalg.fill n=3 0 + linalg.reduce n=2 0 + linalg.transpose n=1 0 + linalg.yield n=12 0 + return n=1 0 + tensor.empty n=7 0 + tensor.splat n=1 0 +``` + +## sdpa_pv_ksplit + +```text +kernel: sdpa_pv_ksplit function: sdpa_pv_ksplit +parameters: K=8192, M=1, N=128 +grid cores active: 32 + +hbm_bytes = 2,130,176 + b_ptr 2,097,152 98.4% ktdp.load x32 + a_ptr 32,768 1.5% ktdp.load x32 + c_ptr 256 0.0% ktdp.store x2 +unique_bytes = 2,113,792 traffic_ratio = 1.008x + +comm_bytes = 126,976 + ktdp.inter_tile_reduce n=32 126,976 + +flops = 2,127,872.0 +arithmetic_intensity = flops / hbm_bytes = 2,127,872.0 / 2,130,176 = 0.9989 + +kernel_cycles = 16,813 bottleneck = memory + critical core compute 79 memory 16,672 comm 62 + compute_float 15 + compute_int 0 + compute_matmul 64 + +cycles charged per op type, critical core: + ktdp.load n=2 16,640 + linalg.matmul n=1 64 + ktdp.inter_tile_reduce n=1 62 + ktdp.store n=1 32 + linalg.add n=15 15 + arith.cmpi n=1 0 + arith.constant n=6 0 + arith.muli n=3 0 + ktdp.construct_access_tile n=3 0 + ktdp.construct_memory_view n=3 0 + ktdp.get_compute_tile_id n=1 0 + ktdp.inter_tile_produce n=1 0 + ktdp.yield_partial n=1 0 + ktdp.yield_reduced n=15 0 + linalg.fill n=1 0 + scf.if n=1 0 + tensor.empty n=16 0 +``` + +## softmax_fwd_ktir + +```text +kernel: softmax_fwd_ktir function: softmax_kernel +parameters: n_cols=1024, n_real_cols=778, n_rows=4096 +grid cores active: 32 + +hbm_bytes = 16,777,216 + input_ptr 8,388,608 50.0% ktdp.load x4096 + output_ptr 8,388,608 50.0% ktdp.store x4096 +unique_bytes = 16,777,216 traffic_ratio = 1x + +flops = 20,971,520.0 +arithmetic_intensity = flops / hbm_bytes = 20,971,520.0 / 16,777,216 = 1.250 + +kernel_cycles = 147,456 bottleneck = memory + critical core compute 16,384 memory 131,072 comm 0 + compute_float 8,192 + compute_transcendental 8,192 + +cycles charged per op type, critical core: + ktdp.load n=128 65,536 + ktdp.store n=128 65,536 + math.exp n=128 8,192 + arith.addf n=1408 2,048 + arith.divf n=128 2,048 + arith.maxnumf n=1408 2,048 + arith.subf n=128 2,048 + arith.constant n=386 0 + ktdp.construct_access_tile n=256 0 + ktdp.construct_memory_view n=2 0 + ktdp.get_compute_tile_id n=1 0 + linalg.reduce n=256 0 + linalg.yield n=2816 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=128 0 + tensor.extract n=256 0 + tensor.splat n=512 0 +``` + +## softmax_small + +```text +kernel: softmax_small function: softmax_kernel_small +parameters: n_cols=64, n_rows=64 +grid cores active: 32 + +hbm_bytes = 16,384 + input_ptr 8,192 50.0% ktdp.load x64 + output_ptr 8,192 50.0% ktdp.store x64 +unique_bytes = 16,384 traffic_ratio = 1x + +flops = 20,480.0 +arithmetic_intensity = flops / hbm_bytes = 20,480.0 / 16,384 = 1.250 + +kernel_cycles = 144 bottleneck = memory + critical core compute 16 memory 128 comm 0 + compute_float 8 + compute_transcendental 8 + +cycles charged per op type, critical core: + ktdp.load n=2 64 + ktdp.store n=2 64 + math.exp n=2 8 + arith.addf n=14 2 + arith.divf n=2 2 + arith.maximumf n=14 2 + arith.subf n=2 2 + arith.constant n=8 0 + ktdp.construct_access_tile n=4 0 + ktdp.construct_memory_view n=2 0 + ktdp.get_compute_tile_id n=1 0 + linalg.reduce n=4 0 + linalg.yield n=28 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=2 0 + tensor.extract n=4 0 + tensor.splat n=8 0 +``` + +## softmax_small_explicit + +```text +kernel: softmax_small_explicit function: softmax_kernel_small_explicit +parameters: n_cols=64, n_rows=64 +grid cores active: 32 + +hbm_bytes = 16,384 + input_ptr 8,192 50.0% ktdp.load x64 + output_ptr 8,192 50.0% ktdp.store x64 +unique_bytes = 16,384 traffic_ratio = 1x + +flops = 20,480.0 +arithmetic_intensity = flops / hbm_bytes = 20,480.0 / 16,384 = 1.250 + +kernel_cycles = 144 bottleneck = memory + critical core compute 16 memory 128 comm 0 + compute_float 8 + compute_transcendental 8 + +cycles charged per op type, critical core: + ktdp.load n=2 64 + ktdp.store n=2 64 + math.exp n=2 8 + arith.addf n=14 2 + arith.divf n=2 2 + arith.maximumf n=14 2 + arith.subf n=2 2 + arith.constant n=8 0 + ktdp.construct_access_tile n=4 0 + ktdp.construct_memory_view n=2 0 + ktdp.get_compute_tile_id n=1 0 + linalg.reduce n=4 0 + linalg.yield n=28 0 + return n=1 0 + scf.for n=1 0 + scf.yield n=2 0 + tensor.extract n=4 0 + tensor.splat n=8 0 +``` + +## vector_add_dynamic_ktir + +```text +kernel: vector_add_dynamic_ktir function: add_kernel_dynamic +parameters: n_elements=1024 +grid cores active: 1 + +hbm_bytes = 12,288 + output_ptr 4,096 33.3% ktdp.store x1 + x_ptr 4,096 33.3% ktdp.load x1 + y_ptr 4,096 33.3% ktdp.load x1 +unique_bytes = 12,288 traffic_ratio = 1x + +flops = 1,024.0 +arithmetic_intensity = flops / hbm_bytes = 1,024.0 / 12,288 = 0.08333 + +kernel_cycles = 112 bottleneck = memory + critical core compute 16 memory 96 comm 0 + compute_float 16 + +cycles charged per op type, critical core: + ktdp.load n=2 64 + ktdp.store n=1 32 + arith.addf n=1 16 + arith.constant n=1 0 + arith.index_cast n=1 0 + ktdp.construct_access_tile n=3 0 + ktdp.construct_memory_view n=3 0 + return n=1 0 +``` + +## vector_add_ktir + +```text +kernel: vector_add_ktir function: add_kernel +parameters: BLOCK_SIZE=128, n_elements=4096 +grid cores active: 32 + +hbm_bytes = 24,576 + output_ptr 8,192 33.3% ktdp.store x32 + x_ptr 8,192 33.3% ktdp.load x32 + y_ptr 8,192 33.3% ktdp.load x32 +unique_bytes = 24,576 traffic_ratio = 1x + +flops = 4,096.0 +arithmetic_intensity = flops / hbm_bytes = 4,096.0 / 24,576 = 0.1667 + +kernel_cycles = 194 bottleneck = memory + critical core compute 2 memory 192 comm 0 + compute_float 2 + compute_int 0 + +cycles charged per op type, critical core: + ktdp.load n=2 128 + ktdp.store n=1 64 + arith.addf n=1 2 + arith.muli n=1 0 + ktdp.construct_access_tile n=3 0 + ktdp.construct_memory_view n=3 0 + ktdp.get_compute_tile_id n=1 0 + return n=1 0 +``` diff --git a/docs/kernel_support.md b/docs/kernel_support.md new file mode 100644 index 0000000..07ee85f --- /dev/null +++ b/docs/kernel_support.md @@ -0,0 +1,78 @@ +# Kernel support + +Generated by `python -m ktir_cpu.kernelentry probe --all --write-report`; +`tests/test_kernelentry.py` fails while it is stale. + +33 kernels under `examples/`, using 490 ops between them; the interpreter has a handler for every one. All 33 are read by the regex parser and 29 are also accepted by the MLIR frontend and MLIR's own verifier. 18 are declared to this ledger, all of which answer all five questions below; 0 claims are open across the repository. + +Support is not one property. It is five questions asked in order, and the columns are them: **read: regex** — `KTIRInterpreter.load` accepts the file; **read: frontend** — the MLIR frontend accepts it and MLIR's own verifier passes; **runs** — it executes on the declared grid without overflowing LX; **output** — every tensor it writes matches a reference computed in f32, and none is silently all-zero; **cost: derivation** — the committed per-tensor cost breakdown still matches what the model reports. + +A cell reads `yes` when the check ran and passed, **`no`** when it ran and failed or could not be evaluated, `deferred #N` against a tracked issue, and `waived` where the check does not apply to that kernel. A kernel is fully supported when no cell in its row is **`no`**. + +`docs/kernelentry.md` is how to declare a kernel and what a declaration has to supply. `docs/supported_ops.md` is the same repository seen per op rather than per kernel. + +## Declared to this ledger (18) + +| kernel | ops | read: regex | read: frontend | runs | output | cost: derivation | +|---|---|---|---|---|---|---| +| `examples/ktir/ffn_swiglu.mlir` | 13 | yes | yes | yes | yes | yes | +| `examples/ktir/ffn_swiglu_4core.mlir` | 28 | yes | yes | yes | yes | yes | +| `examples/ktir/reduce_generic.mlir` | 13 | yes | yes | yes | yes | yes | +| `examples/ktir/reduce_multiop.mlir` | 14 | yes | yes | yes | yes | yes | +| `examples/ktir/scalar_broadcast.mlir` | 9 | yes | yes | yes | yes | yes | +| `examples/latency/matmul_small.mlir` | 13 | yes | yes | yes | yes | yes | +| `examples/latency/rope_fwd_4x2.mlir` | 13 | yes | yes | yes | yes | yes | +| `examples/latency/softmax_small.mlir` | 15 | yes | yes | yes | yes | yes | +| `examples/latency/softmax_small_explicit.mlir` | 18 | yes | yes | yes | yes | yes | +| `examples/sdsc/sdpa_pv_ksplit.mlir` | 18 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/indexed_add.mlir` | 9 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/layernorm_fwd_ktir.mlir` | 19 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/matmul_fwd_ktir.mlir` | 13 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/paged_attention.mlir` | 35 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/sdpa_2d.mlir` | 18 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/softmax_fwd_ktir.mlir` | 15 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/vector_add_dynamic_ktir.mlir` | 8 | yes | yes | yes | yes | yes | +| `examples/triton-ktir/vector_add_ktir.mlir` | 8 | yes | yes | yes | yes | yes | + +## Read, not declared (15) + +Both reading questions are answered for these, and every op they use has an execution handler. The other three have not been *asked*, which is not the same as answered no: they need what only a declaration supplies — tensors to drive the kernel with, and a reference independent of the simulator — so their columns are absent here rather than empty. + +`tests/` already drives every file below, from `tests/conftest.py::EXAMPLE_PARAMS`, but what that listing supplies is not a declaration waiting to be copied: `execute_kwargs` is empty for 8 of them and carries raw HBM element indices, or a scalar size, for the rest. `docs/kernelentry.md` groups these files by which reason applies. + +| kernel | ops | read: regex | read: frontend | +|---|---|---|---| +| `examples/ktir/nested_yield.ktir` | 8 | yes | **no** | +| `examples/ktir/ring_reduce.mlir` | 20 | yes | **no** | +| `examples/ktir/ring_reduce_inner_loop.mlir` | 22 | yes | **no** | +| `examples/ktir/softmax_wide.mlir` | 19 | yes | yes | +| `examples/latency/ring_reduce_multi_group.mlir` | 22 | yes | **no** | +| `examples/latency/rmsnorm_4core_2x2.mlir` | 29 | yes | yes | +| `examples/latency/rmsnorm_4core_4x1.mlir` | 20 | yes | yes | +| `examples/rfc/add-with-control-flow.mlir` | 13 | yes | yes | +| `examples/rfc/distributed-view-copy-dynamic.mlir` | 10 | yes | yes | +| `examples/rfc/distributed-view-copy-rowmerge-dynamic.mlir` | 9 | yes | yes | +| `examples/rfc/distributed-view-copy.mlir` | 7 | yes | yes | +| `examples/rfc/indirect-access-copy.mlir` | 7 | yes | yes | +| `examples/rfc/indirect-scatter.mlir` | 7 | yes | yes | +| `examples/rfc/paged-tensor-copy.mlir` | 9 | yes | yes | +| `examples/rfc/paged-tensor-write.mlir` | 9 | yes | yes | + +## Open and undetermined (0) + +None. + +## Deferred (0) + +None. + +## Waived (0) + +None. + +## Not accepted by the MLIR frontend (4) + +A record, not a verdict: it does not say whether the file or the parser reading it should change. + +- `examples/ktir/nested_yield.ktir` — construct_memory_view has no coordinate_set, which the dialect requires +- `examples/ktir/ring_reduce.mlir`, `examples/ktir/ring_reduce_inner_loop.mlir`, `examples/latency/ring_reduce_multi_group.mlir` — inter-tile ops in the form only the regex parser reads, and a reduce that reshapes its result, which the dialect's type relation does not express (gap row 2a) \ No newline at end of file diff --git a/docs/supported_ops.md b/docs/supported_ops.md new file mode 100644 index 0000000..1373ed3 --- /dev/null +++ b/docs/supported_ops.md @@ -0,0 +1,325 @@ + + +# Supported operations + +Op-level truth about this interpreter, read from the registries themselves. `tests/test_kernelentry.py` fails while this file and the registries disagree, so it cannot go stale silently. + +106 ops have an execution handler. 102 of them are also reachable through the MLIR frontend; 4 are deliberately not ([why](#ops-not-reachable-through-the-mlir-frontend)). 52 are used by a kernel under `examples/`, leaving 54 that no example exercises ([which](#ops-no-example-exercises)). + +46 of them cost nothing. That is two facts, not one: 21 are [free by decision](#ops-that-are-free-by-decision) and 25 are [priced zero with nobody having decided](#ops-priced-zero-without-a-decision). `@register()` defaults `latency_category` to `zero`, so the registry itself cannot tell those apart — `ktir_cpu/kernelentry/pricing.py` is where they are split, and an op priced zero that appears in neither list fails `tests/test_kernelentry.py`. + +**Division of labour.** `docs/kernel_support.md` asks five questions about each *kernel*, and three of them need a declared entry — so it reports how far that ledger reaches. This file asks one question about each *op*, needs no declaration, and therefore covers the whole registry. `docs/gap_analysis.md` is the third: conformance against RFC 0682, judged by a reader rather than generated. + +## Matrix + +| op | executor handler | MLIR frontend | cost | kernels | +|---|---|---|---|---| +| `arith.absf` | yes | yes | `compute_float` | none | +| `arith.addf` | yes | yes | `compute_float` | 15 | +| `arith.addi` | yes | yes | `compute_int` | 10 | +| `arith.andi` | yes | yes | `compute_int` | none | +| `arith.bitcast` | yes | yes | **not judged** | 1 | +| `arith.ceildivsi` | yes | yes | `compute_int` | none | +| `arith.ceildivui` | yes | yes | `compute_int` | none | +| `arith.cmpf` | yes | yes | `compute_float` | 1 | +| `arith.cmpi` | yes | yes | `compute_int` | 7 | +| `arith.constant` | yes | yes | free, by decision | 32 | +| `arith.convertf` | yes | no, by design | **not judged** | none | +| `arith.divf` | yes | yes | `compute_float` | 11 | +| `arith.divsi` | yes | yes | `compute_int` | none | +| `arith.divui` | yes | yes | `compute_int` | 4 | +| `arith.extf` | yes | yes | **not judged** | 1 | +| `arith.extsi` | yes | yes | **not judged** | none | +| `arith.extui` | yes | yes | **not judged** | none | +| `arith.floordivsi` | yes | yes | `compute_int` | none | +| `arith.fptosi` | yes | yes | **not judged** | none | +| `arith.fptoui` | yes | yes | **not judged** | none | +| `arith.index_cast` | yes | yes | **not judged** | 6 | +| `arith.index_castui` | yes | yes | **not judged** | none | +| `arith.maxf` | yes | yes | `compute_float` | none | +| `arith.maximumf` | yes | yes | `compute_float` | 2 | +| `arith.maxnumf` | yes | yes | `compute_float` | none | +| `arith.maxsi` | yes | yes | `compute_int` | none | +| `arith.maxui` | yes | yes | `compute_int` | none | +| `arith.minf` | yes | yes | `compute_float` | none | +| `arith.minimumf` | yes | yes | `compute_float` | none | +| `arith.minnumf` | yes | yes | `compute_float` | none | +| `arith.minsi` | yes | yes | `compute_int` | none | +| `arith.minui` | yes | yes | `compute_int` | none | +| `arith.mulf` | yes | yes | `compute_float` | 8 | +| `arith.muli` | yes | yes | `compute_int` | 17 | +| `arith.negf` | yes | yes | `compute_float` | 2 | +| `arith.ori` | yes | yes | `compute_int` | none | +| `arith.remf` | yes | yes | `compute_float` | none | +| `arith.remsi` | yes | yes | `compute_int` | none | +| `arith.remui` | yes | yes | `compute_int` | 1 | +| `arith.select` | yes | yes | `compute_float` | 3 | +| `arith.shli` | yes | yes | `compute_int` | none | +| `arith.shrsi` | yes | yes | `compute_int` | none | +| `arith.shrui` | yes | yes | `compute_int` | none | +| `arith.sitofp` | yes | yes | **not judged** | 3 | +| `arith.subf` | yes | yes | `compute_float` | 8 | +| `arith.subi` | yes | yes | `compute_int` | none | +| `arith.truncf` | yes | yes | **not judged** | 1 | +| `arith.trunci` | yes | yes | **not judged** | none | +| `arith.uitofp` | yes | yes | **not judged** | none | +| `arith.xori` | yes | yes | `compute_int` | none | +| `func.return` | yes | yes | free, by decision | none | +| `ktdp.construct_access_tile` | yes | yes | free, by decision | 33 | +| `ktdp.construct_distributed_memory_view` | yes | yes | free, by decision | 5 | +| `ktdp.construct_indirect_access_tile` | yes | yes | free, by decision | 6 | +| `ktdp.construct_memory_view` | yes | yes | free, by decision | 33 | +| `ktdp.coreid` | yes | no, by design | **not judged** | none | +| `ktdp.get_compute_tile_id` | yes | yes | free, by decision | 20 | +| `ktdp.inter_tile_produce` | yes | yes | free, by decision | 6 | +| `ktdp.inter_tile_reduce` | yes | yes | `comm` | 6 | +| `ktdp.load` | yes | yes | `memory` | 32 | +| `ktdp.region_terminator` | **no** | yes | — | none | +| `ktdp.store` | yes | yes | `memory` | 33 | +| `ktdp.yield_partial` | yes | yes | free, by decision | 6 | +| `ktdp.yield_reduced` | yes | yes | free, by decision | 6 | +| `linalg.add` | yes | yes | `compute_float` | 7 | +| `linalg.batch_matmul` | yes | yes | `compute_matmul` | none | +| `linalg.broadcast` | yes | yes | **not judged** | 2 | +| `linalg.fill` | yes | yes | **not judged** | 9 | +| `linalg.generic` | yes | yes | `compute_float` | 1 | +| `linalg.index` | yes | yes | free, by decision | 1 | +| `linalg.matmul` | yes | yes | `compute_matmul` | 7 | +| `linalg.max` | yes | yes | `compute_float` | none | +| `linalg.reduce` | yes | yes | free, by decision | 11 | +| `linalg.transpose` | yes | yes | **not judged** | 2 | +| `linalg.yield` | yes | yes | free, by decision | 4 | +| `math.absf` | yes | yes | `compute_float` | none | +| `math.absi` | yes | yes | `compute_float` | none | +| `math.ceil` | yes | yes | `compute_float` | none | +| `math.cos` | yes | yes | `compute_transcendental` | none | +| `math.erf` | yes | yes | `compute_transcendental` | none | +| `math.exp` | yes | yes | `compute_transcendental` | 8 | +| `math.floor` | yes | yes | `compute_float` | none | +| `math.fma` | yes | yes | `compute_float` | none | +| `math.log` | yes | yes | `compute_transcendental` | none | +| `math.log1p` | yes | yes | `compute_transcendental` | none | +| `math.log2` | yes | yes | `compute_transcendental` | none | +| `math.powf` | yes | yes | `compute_transcendental` | none | +| `math.rsqrt` | yes | yes | `compute_transcendental` | 2 | +| `math.sin` | yes | yes | `compute_transcendental` | none | +| `math.sqrt` | yes | yes | `compute_transcendental` | 1 | +| `math.tanh` | yes | yes | `compute_transcendental` | none | +| `region.bb0_args` | yes | no, by design | free, by decision | 7 | +| `return` | yes | no, by design | free, by decision | 31 | +| `scf.for` | yes | yes | free, by decision | 16 | +| `scf.if` | yes | yes | free, by decision | 5 | +| `scf.yield` | yes | yes | free, by decision | 14 | +| `tensor.collapse_shape` | yes | yes | **not judged** | 3 | +| `tensor.empty` | yes | yes | free, by decision | 15 | +| `tensor.expand_shape` | yes | yes | **not judged** | 5 | +| `tensor.extract` | yes | yes | **not judged** | 9 | +| `tensor.extract_slice` | yes | yes | **not judged** | none | +| `tensor.from_elements` | yes | yes | **not judged** | none | +| `tensor.generate` | yes | yes | free, by decision | none | +| `tensor.insert_slice` | yes | yes | **not judged** | none | +| `tensor.reshape` | yes | yes | **not judged** | none | +| `tensor.splat` | yes | yes | **not judged** | 11 | +| `tensor.yield` | yes | yes | free, by decision | none | + +### How to read the columns + +- **executor handler** — a `@register` handler exists, so the interpreter can execute the op. `no` means it cannot; such an op is listed here only because the MLIR frontend or an example mentions it. +- **MLIR frontend** — an `@MLIRTypeAdapter.install` handler exists, so the op survives the real MLIR parser and its `verify()`. `no, by design` is an entry in the allowlist below. A bare **no** should not occur: `tests/mlir_frontend/test_registry_consistency.py` fails on it. +- **cost** — what the latency model charges. A named category is the one `@register` gave it. `free, by decision` and `**not judged**` are both `zero` in the registry, split apart by `ktir_cpu/kernelentry/pricing.py`: the first has a written reason why the hardware does no measurable work, the second has a written statement of what is unresolved and the issue tracking it. A cost column can only catch a *missing* price — an op billed to the wrong non-zero unit passes it, and needs a reader. +- **kernels** — how many files under `examples/` use the op; [the files themselves](#where-each-op-appears) are listed below, with what each directory is worth as evidence. `none` means the op is registered but no example exercises it, so only unit tests, if any, cover it. + +What this file does not say: whether a handler is **correct**, and whether a particular attribute or type of the op is supported. The unit here is the op name — `ktdp.construct_access_tile` having a handler says nothing about a particular `base_map` or `coordinate_set` reaching the conclusion the specification does. `docs/gap_analysis.md` tracks that. + +## Ops not reachable through the MLIR frontend + +Executor ops with no MLIR frontend handler, and the reason each is allowed to stay that way. This is the allowlist `tests/mlir_frontend/test_registry_consistency.py` enforces — an op missing from both the frontend and this list fails that test. + +| op | reason | +|---|---| +| `arith.convertf` | not a real upstream arith op — unknown to the MLIR parser; reconcile to a real cast op or remove. | +| `ktdp.coreid` | non-spec op (not in the ktdp dialect). Reconcile to the real core-identity form or remove. See issue #88. | +| `region.bb0_args` | regex-only synthetic op; the frontend carries bb0 names on linalg.generic via the bb0_names attribute. | +| `return` | alias of func.return for the bare-`return` regex form; the bindings always emit func.return. | + +## Ops that are free by decision + +21 ops the cost model charges nothing for, and the reason each one does no measurable work. Four kinds: a value that exists at compile time, a terminator that only names values, addressing metadata that computes where data is without moving it, and an orchestrator whose region is charged op by op. + +| op | why it is free | +|---|---| +| `arith.constant` | a compile-time literal, and registered `no_lx_charge=True` for the same reason: the scratchpad is charged when a consumer materializes the value into a working tile, not here | +| `func.return` | a terminator; it names the values leaving the function | +| `ktdp.construct_access_tile` | narrows an existing view to a tile's worth of it — index arithmetic on the view, with no access performed | +| `ktdp.construct_distributed_memory_view` | the same address computation, per partition; the traffic is charged on the loads and stores that address it | +| `ktdp.construct_indirect_access_tile` | the same narrowing with a gathered index set; the gather itself is charged on the indirect load, whose figure covers both the data and the index lookups | +| `ktdp.construct_memory_view` | computes an address, moves nothing; the `ktdp.load` or `ktdp.store` that reads through the view is what carries the bytes | +| `ktdp.get_compute_tile_id` | reads the core's own coordinate in the grid, which is available to it without a memory access | +| `ktdp.inter_tile_produce` | publishes this core's partial to the scheduler's mailbox; the wire time for the whole exchange is charged once, as `comm`, on the matching `ktdp.inter_tile_reduce`. Pricing both would count one transfer twice | +| `ktdp.yield_partial` | a region terminator; it names this core's partial | +| `ktdp.yield_reduced` | a region terminator; it names the reduced value | +| `linalg.index` | produces the iteration index inside a region body; the arithmetic that consumes it is charged, again the `linalg.reduce` split | +| `linalg.reduce` | executes its combiner region rather than mapping to a fixed reduction, so the arithmetic inside it is charged individually: one core's trace for the RMSNorm generator carries 256 `arith.addf` entries totalling 1,280 cycles from inside a reduce. The orchestrator is free; the arithmetic is not | +| `linalg.yield` | a region terminator; it names the combiner's result | +| `region.bb0_args` | not an op at all — the parser's record of a region's block arguments, which the enclosing op's handler binds | +| `return` | the regex parser's spelling of `func.return`; same reason | +| `scf.for` | loop control; the body's ops are charged once per iteration, so a cost for the loop itself would be on top of the work it drives | +| `scf.if` | branch control; the ops of the taken branch are charged | +| `scf.yield` | a region terminator | +| `tensor.empty` | names an uninitialized buffer. No data moves, and a consumer that writes into it pays for the write | +| `tensor.generate` | evaluates its region body per index, so the ops in the body are charged individually — the same split as `linalg.reduce` | +| `tensor.yield` | a region terminator, for a `tensor.generate` body | + +## Ops priced zero without a decision + +25 ops that cost nothing today because `latency_category` defaults to `zero`, not because anyone decided they are free. A kernel using one of these reports a lower cost than the hardware would. They are listed rather than fixed because each is a hardware question RFC 0682 does not settle — a conversion folded into the consumer's read is free and a materialized one is not — and repricing one changes the committed derivations of every kernel using it. + +| op | what is unresolved | kernels | +|---|---|---| +| `arith.bitcast` | reinterprets a tile's bits under another type. Free if it is a type relabel and not free if the data is copied; the handler uses `ndarray.view`, which is the free reading. #211 | 1 | +| `arith.convertf` | converts every element of a tile, via `_unary`. #211 | none | +| `arith.extf` | widens every element of a tile, via `_unary`. #211 | 1 | +| `arith.extsi` | sign-extends every element of a tile. #211 | none | +| `arith.extui` | zero-extends every element of a tile, via `_unary`. #211 | none | +| `arith.fptosi` | float to signed integer across the tile, via `_unary`. #211 | none | +| `arith.fptoui` | float to unsigned integer across the tile, via `_unary`. #211 | none | +| `arith.index_cast` | zero is defensible — the handler returns a Python `int`, so this is one scalar conversion rather than a tile's worth — but it has not been decided. #211 | 6 | +| `arith.index_castui` | the same one-scalar conversion as `arith.index_cast`, undecided for the same reason. #211 | none | +| `arith.sitofp` | integer to float across the whole tile, via `_unary`. #211 | 3 | +| `arith.truncf` | narrows every element of a tile, via `_unary`. #211 | 1 | +| `arith.trunci` | truncates every element of a tile, via `_unary`. #211 | none | +| `arith.uitofp` | unsigned integer to float across the tile, via `_unary`. #211 | none | +| `ktdp.coreid` | not an op to price: it is not in the authoritative `ktdp` dialect and survives only on the regex path, so the resolution is to reconcile or remove it rather than to give it a category. #88 | none | +| `linalg.broadcast` | expands a tile along new dimensions — free if the consumer reads it strided, not free if it is materialized. #211 | 2 | +| `linalg.fill` | writes a scalar across the whole `outs` tile. Eight kernels under `examples/` use it. #211 | 9 | +| `linalg.transpose` | free if it is a stride permutation, not free if the data moves; the handler calls `np.transpose(...).copy()`, which is the second reading. #211 | 2 | +| `tensor.collapse_shape` | the same question as `tensor.reshape`. #211 | 3 | +| `tensor.expand_shape` | the same question as `tensor.reshape`. #211 | 5 | +| `tensor.extract` | reads one element out of a tile. One scalar read, so zero is defensible, but undecided. #211 | 9 | +| `tensor.extract_slice` | reads a strided sub-tensor. Free if the consumer reads the parent strided, not free if the slice is materialized. #211 | none | +| `tensor.from_elements` | builds a small tensor from N scalar operands, so its cost is N element writes rather than a tile's worth. #211 | none | +| `tensor.insert_slice` | writes a sub-tensor into a destination, which is a copy of the slice's worth of elements unless it folds into the producer. #211 | none | +| `tensor.reshape` | reinterprets the same elements under a new shape. Free as metadata, not free if the layout is rebuilt. #211 | none | +| `tensor.splat` | broadcasts one scalar across a whole tile — the same question as `linalg.fill`, and priced the same way. #211 | 11 | + +## Ops no example exercises + +54 registered ops that no file under `examples/` uses. Not a defect on its own — an op can be covered by a unit test — but it is where a handler nothing has ever run would hide: + +- `arith.absf` +- `arith.andi` +- `arith.ceildivsi` +- `arith.ceildivui` +- `arith.convertf` +- `arith.divsi` +- `arith.extsi` +- `arith.extui` +- `arith.floordivsi` +- `arith.fptosi` +- `arith.fptoui` +- `arith.index_castui` +- `arith.maxf` +- `arith.maxnumf` +- `arith.maxsi` +- `arith.maxui` +- `arith.minf` +- `arith.minimumf` +- `arith.minnumf` +- `arith.minsi` +- `arith.minui` +- `arith.ori` +- `arith.remf` +- `arith.remsi` +- `arith.shli` +- `arith.shrsi` +- `arith.shrui` +- `arith.subi` +- `arith.trunci` +- `arith.uitofp` +- `arith.xori` +- `func.return` +- `ktdp.coreid` +- `linalg.batch_matmul` +- `linalg.max` +- `math.absf` +- `math.absi` +- `math.ceil` +- `math.cos` +- `math.erf` +- `math.floor` +- `math.fma` +- `math.log` +- `math.log1p` +- `math.log2` +- `math.powf` +- `math.sin` +- `math.tanh` +- `tensor.extract_slice` +- `tensor.from_elements` +- `tensor.generate` +- `tensor.insert_slice` +- `tensor.reshape` +- `tensor.yield` + +## Where each op appears + +Up to 3 files are named; above that the directories and a count, because the identity of the file is what matters when there are few and the coverage class is what matters when there are many. Paths are relative to `examples/`, and what each directory is worth as evidence: + +- `ktir/` — hand-written dialect cases; executed by the test suite +- `latency/` — small kernels for the latency tests; executed by the test suite +- `rfc/` — RFC 0682 specification examples; **expected to fail execution** (they carry absolute addresses rather than arguments), so an op seen only here is parsed, not run +- `sdsc/` — kernels from the SuperDSC lowering path; executed by the test suite +- `triton-ktir/` — kernels as the Triton → KTIR path emits them, i.e. captured compiler output; executed by the test suite + +- `arith.addf` — 15 files in ktir/, latency/, triton-ktir/ +- `arith.addi` — 10 files in ktir/, latency/, rfc/, triton-ktir/ +- `arith.bitcast` — triton-ktir/paged_attention.mlir +- `arith.cmpf` — ktir/reduce_multiop.mlir +- `arith.cmpi` — 7 files in ktir/, latency/, sdsc/, triton-ktir/ +- `arith.constant` — 32 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `arith.divf` — 11 files in ktir/, latency/, triton-ktir/ +- `arith.divui` — 4 files in latency/, rfc/, triton-ktir/ +- `arith.extf` — triton-ktir/paged_attention.mlir +- `arith.index_cast` — 6 files in latency/, rfc/, triton-ktir/ +- `arith.maximumf` — latency/softmax_small_explicit.mlir, triton-ktir/paged_attention.mlir +- `arith.mulf` — 8 files in ktir/, latency/, triton-ktir/ +- `arith.muli` — 17 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `arith.negf` — ktir/ffn_swiglu.mlir, ktir/ffn_swiglu_4core.mlir +- `arith.remui` — latency/ring_reduce_multi_group.mlir +- `arith.select` — ktir/reduce_multiop.mlir, ktir/softmax_wide.mlir, triton-ktir/paged_attention.mlir +- `arith.sitofp` — latency/rmsnorm_4core_2x2.mlir, latency/rmsnorm_4core_4x1.mlir, triton-ktir/layernorm_fwd_ktir.mlir +- `arith.subf` — 8 files in ktir/, latency/, triton-ktir/ +- `arith.truncf` — triton-ktir/paged_attention.mlir +- `ktdp.construct_access_tile` — 33 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `ktdp.construct_distributed_memory_view` — 5 files in ktir/, latency/, rfc/ +- `ktdp.construct_indirect_access_tile` — 6 files in rfc/, triton-ktir/ +- `ktdp.construct_memory_view` — 33 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `ktdp.get_compute_tile_id` — 20 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `ktdp.inter_tile_produce` — 6 files in ktir/, latency/, sdsc/ +- `ktdp.inter_tile_reduce` — 6 files in ktir/, latency/, sdsc/ +- `ktdp.load` — 32 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `ktdp.store` — 33 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `ktdp.yield_partial` — 6 files in ktir/, latency/, sdsc/ +- `ktdp.yield_reduced` — 6 files in ktir/, latency/, sdsc/ +- `linalg.add` — 7 files in ktir/, latency/, rfc/, sdsc/ +- `linalg.broadcast` — ktir/scalar_broadcast.mlir, triton-ktir/sdpa_2d.mlir +- `linalg.fill` — 9 files in ktir/, latency/, sdsc/, triton-ktir/ +- `linalg.generic` — triton-ktir/paged_attention.mlir +- `linalg.index` — triton-ktir/paged_attention.mlir +- `linalg.matmul` — 7 files in ktir/, latency/, sdsc/, triton-ktir/ +- `linalg.reduce` — 11 files in ktir/, latency/, triton-ktir/ +- `linalg.transpose` — triton-ktir/paged_attention.mlir, triton-ktir/sdpa_2d.mlir +- `linalg.yield` — 4 files in ktir/, latency/, triton-ktir/ +- `math.exp` — 8 files in ktir/, latency/, triton-ktir/ +- `math.rsqrt` — latency/rmsnorm_4core_2x2.mlir, latency/rmsnorm_4core_4x1.mlir +- `math.sqrt` — triton-ktir/layernorm_fwd_ktir.mlir +- `region.bb0_args` — 7 files in ktir/, latency/, sdsc/, triton-ktir/ +- `return` — 31 files in ktir/, latency/, rfc/, triton-ktir/ +- `scf.for` — 16 files in ktir/, latency/, rfc/, triton-ktir/ +- `scf.if` — 5 files in ktir/, latency/, sdsc/ +- `scf.yield` — 14 files in ktir/, latency/, rfc/, triton-ktir/ +- `tensor.collapse_shape` — ktir/ffn_swiglu_4core.mlir, ktir/scalar_broadcast.mlir, triton-ktir/paged_attention.mlir +- `tensor.empty` — 15 files in ktir/, latency/, rfc/, sdsc/, triton-ktir/ +- `tensor.expand_shape` — 5 files in ktir/, latency/, triton-ktir/ +- `tensor.extract` — 9 files in ktir/, latency/, triton-ktir/ +- `tensor.splat` — 11 files in ktir/, latency/, triton-ktir/ From 33e5165d36388852f0c1e1b2989955526d991b59 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:16 -0400 Subject: [PATCH 7/9] Gate the ledger in CI, and test that every claim can fail MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The gate for everything above: one parameterized case per claim per declared kernel, over the same engine `verify` runs, so the local loop and CI cannot come to different conclusions about the same tree. A claim's state decides its marker, so a deferral is an `xfail(strict=True)` and closing the gap it names fails the build until the declaration stops claiming it. A gate can also be green because nothing it asks is hard, and that is what the second suite is for. It removes a reference, narrows one to a slice of the output, swaps in a reference computed in f16, adds a reference key no store wrote, drops a declared tolerance, staleness-checks a waiver and a deferral by closing the gap each excuses — and asserts the claim turns each time. A claim that cannot be made to fail there is decoration, and the suite is where that gets established rather than assumed. Three more things are asserted rather than trusted: - The generated documents are functions of the repository and not of the host, checked by rendering them twice with the MLIR bindings stubbed out. Without that they would differ between CI and a contributor's laptop, and each would call the other's stale forever. - The support report covers every example `tests/conftest.py` already drives, so a kernel cannot be added to the test parameters and stay invisible to the ledger. The same test pins the one figure the renderer cannot compute — how many of those listings pass no arguments — because the package cannot import the test suite, and a hand-written count in a generated document is the rot this whole change is against. - The pricing audit's blind spot is written down as a test: an op billed to the wrong pipe is still priced, so no comparison against `zero` can catch it. Left unstated it is a gap someone rediscovers; stated, it is the reason the follow-up asks for a semantic audit. `tests/mlir_frontend/test_kernelentry_adapt.py` verifies `conformance.py` against the real frontend in both directions, so an unrecorded rejection and a rejection since fixed each fail. It also asserts that `parse.frontend` comes back as a real answer rather than `skip` where the bindings are installed, because a leg that silently skips is a leg that reports green without being asked. Signed-off-by: WarningRan --- tests/mlir_frontend/test_kernelentry_adapt.py | 133 +++ tests/test_kernelentry.py | 778 ++++++++++++++++++ 2 files changed, 911 insertions(+) create mode 100644 tests/mlir_frontend/test_kernelentry_adapt.py create mode 100644 tests/test_kernelentry.py diff --git a/tests/mlir_frontend/test_kernelentry_adapt.py b/tests/mlir_frontend/test_kernelentry_adapt.py new file mode 100644 index 0000000..ed86913 --- /dev/null +++ b/tests/mlir_frontend/test_kernelentry_adapt.py @@ -0,0 +1,133 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The frontend leg of the gate: ``parse.frontend`` must actually be decided. + +Without ``mlir_ktdp`` the ledger reports ``parse.frontend`` as ``skip``, which is +correct — but a claim that is only ever skipped is a claim that never runs. This +module is where it runs. It lives under ``tests/mlir_frontend/`` so it inherits +that package's module-level skip, and therefore only executes where the bindings +exist: in CI. + +Why the claim matters more than the other parse claim: the regex parser accepts +unregistered ops through a catch-all, and the frontend has none — it runs MLIR's +own verifier. A kernel that works only on the tolerant path is exactly what a +green local run hides. +""" + +from __future__ import annotations + +import pytest + +from ktir_cpu.kernelentry import BLOCKING, DEFERRED, SKIP, registered +from ktir_cpu.kernelentry.cli import discover_all +from ktir_cpu.kernelentry.ledger import probe + +discover_all() +ENTRIES = registered() + + +@pytest.mark.parametrize("name", sorted(ENTRIES)) +def test_kernel_parses_on_the_mlir_frontend(name: str): + """Every declared kernel is accepted by the real frontend, verifier included. + + A deferral passes here. The declaration having recorded a known frontend gap + against an issue is a different situation from nobody having looked, and it is + the gate in ``tests/test_kernelentry.py`` that holds the deferral to account — + duplicating that here would report one gap as two failures. + """ + ledger = probe(ENTRIES[name]) + claim = next(c for c in ledger.claims if c.id == "parse.frontend") + assert claim.state != SKIP, ( + "parse.frontend was skipped in an environment that has mlir_ktdp — the " + "ledger's import guard is too broad, and the claim will never be decided" + ) + assert claim.state not in BLOCKING, f"{claim.state}: {claim.detail}" + if claim.state == DEFERRED: + pytest.skip(f"deferred: {claim.detail}") + + +class TestTheCommittedRecordOfFrontendRejections: + """``conformance.FRONTEND_REJECTS`` must say what the frontend actually says. + + ``docs/kernel_support.md`` prints that mapping instead of what the generating + machine saw, because the check needs the optional MLIR bindings and the report + is compared verbatim. A committed record only earns that role if something + checks it, and this is the one environment that can: here the frontend is real. + + Both directions, because each failure mode is its own kind of wrong. A missing + entry means the report calls a rejected kernel accepted — the exact outcome + listing undeclared kernels at all is meant to prevent. A stale entry means it + reports a gap that somebody has since closed, which is how a record stops being + read. + """ + + @staticmethod + def _rejected() -> dict: + """Every kernel under ``examples/`` the frontend will not accept.""" + from ktir_cpu.kernelentry import EXAMPLES_DIR, REPO_ROOT + from ktir_cpu.mlir_frontend.parser import MLIRFrontendParser + + out = {} + for path in sorted(EXAMPLES_DIR.rglob("*")): + if path.suffix not in (".mlir", ".ktir"): + continue + rel = str(path.relative_to(REPO_ROOT)) + try: + MLIRFrontendParser().parse_module(path.read_text()) + except Exception as exc: # noqa: BLE001 — any rejection counts + out[rel] = str(exc) + return out + + def test_every_rejected_kernel_is_recorded_or_declared(self): + """Recorded in the mapping, or declared with the claim excused. Not neither.""" + from ktir_cpu.kernelentry import EXAMPLES_DIR, REPO_ROOT + from ktir_cpu.kernelentry.conformance import FRONTEND_REJECTS + + excused = set() + for entry in ENTRIES.values(): + if "parse.frontend" in entry.waived or "parse.frontend" in entry.deferred: + excused.add(str(entry.mlir_path.resolve() + .relative_to(REPO_ROOT.resolve()))) + + unrecorded = { + rel: err for rel, err in self._rejected().items() + if rel not in FRONTEND_REJECTS and rel not in excused + } + assert not unrecorded, ( + "the MLIR frontend rejects these kernels and nothing in the repository " + "says so, so docs/kernel_support.md reports them as accepted. Add each " + "to FRONTEND_REJECTS in ktir_cpu/kernelentry/conformance.py with a " + "reason, or excuse parse.frontend in its declaration:\n " + + "\n ".join(f"{rel}: {err.splitlines()[0]}" + for rel, err in sorted(unrecorded.items())) + ) + + def test_the_record_has_no_stale_entries(self): + """A recorded gap that has been closed must come out of the record.""" + from ktir_cpu.kernelentry import REPO_ROOT + from ktir_cpu.kernelentry.conformance import FRONTEND_REJECTS + + rejected = self._rejected() + missing = [rel for rel in FRONTEND_REJECTS if not (REPO_ROOT / rel).exists()] + assert not missing, ( + "FRONTEND_REJECTS names kernels that no longer exist (remove them): " + f"{sorted(missing)}" + ) + now_accepted = [rel for rel in FRONTEND_REJECTS if rel not in rejected] + assert not now_accepted, ( + "the MLIR frontend now accepts these, so FRONTEND_REJECTS is reporting " + "a gap that is closed — remove the entries and regenerate " + f"docs/kernel_support.md: {sorted(now_accepted)}" + ) diff --git a/tests/test_kernelentry.py b/tests/test_kernelentry.py new file mode 100644 index 0000000..04da405 --- /dev/null +++ b/tests/test_kernelentry.py @@ -0,0 +1,778 @@ +# Copyright 2025 The Torch-Spyre Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The gate: one test per claim, over every kernel that declares an entry. + +This module contains no checks of its own. It calls +``ktir_cpu.kernelentry.ledger.probe`` — the same function the CLI calls — and +turns each claim into a test result. A check implemented here as well as there +would be a second opinion, and the two would eventually disagree about whether a +kernel is supported. + +A ``deferred`` claim becomes ``xfail(strict=True)``, following +``tests/test_spec_gaps.py``, so a deferral reads as a known gap rather than a +failure. What makes the deferral *expire* is not that marker, though: the marker +is applied from the claim's current state, so once the gap closes the claim is no +longer deferred, no marker is applied, and this test simply passes. The engine +raises ``deferred.unnecessary.`` for that case, which is what actually fails +the build until the declaration is updated. + +Cost: each entry is probed once, at ``gate_params``, and the result is cached for +every claim derived from it — without that cache the full-size kernels would be +re-executed once per claim, and ``layernorm_fwd_ktir`` alone raises 30. Gate cost +is driven by shape rather than by how many kernels are declared, which is the +trade ``gate_params`` makes; ``docs/kernelentry.md`` carries the measurement. +""" + +from __future__ import annotations + +import re +from typing import Dict, List, Tuple + +import pytest + +from ktir_cpu.dialects import registry +from ktir_cpu.kernelentry import ( + BLOCKING, CLOSED, DEFERRED, REPO_ROOT, KernelEntry, registered, +) +from ktir_cpu.kernelentry.cli import discover_all, render_report +from ktir_cpu.kernelentry.ledger import Ledger, probe +from ktir_cpu.kernelentry.pricing import ( + BOTH, PRICED, UNJUDGED_ZERO_OPS, UNLISTED, UNREGISTERED, ZERO_COST_OPS, audit, +) + +discover_all() +ENTRIES: Dict[str, KernelEntry] = registered() + +_CACHE: Dict[str, Ledger] = {} + + +def ledger_for(name: str) -> Ledger: + if name not in _CACHE: + _CACHE[name] = probe(ENTRIES[name]) + return _CACHE[name] + + +def _claim_params() -> List[pytest.param]: + """One parameter per (kernel, claim), with deferred claims marked xfail. + + The ledger has to be computed to know the claim set, which is the point — the + set is derived from each kernel rather than listed here, so a kernel with more + ops contributes more tests without this file changing. + """ + params = [] + for name in sorted(ENTRIES): + for claim in ledger_for(name).claims: + marks = [] + if claim.state == DEFERRED: + marks.append(pytest.mark.xfail( + strict=True, reason=claim.detail or "deferred")) + params.append(pytest.param(name, claim.id, + marks=marks, id=f"{name}-{claim.id}")) + return params + + +def test_at_least_one_entry_is_declared(): + """A gate over an empty set passes, and would do so silently forever.""" + assert ENTRIES, ( + "no kernel entries found under examples/ — the gate would pass " + "vacuously. See docs/kernelentry.md." + ) + + +@pytest.mark.parametrize("name,claim_id", _claim_params()) +def test_claim(name: str, claim_id: str): + """One claim about one kernel. + + ``undetermined`` fails alongside ``open`` on purpose: a check the engine could + not evaluate must not read as one that passed, which is exactly what happens + when such a claim is quietly omitted instead. + """ + claim = next(c for c in ledger_for(name).claims if c.id == claim_id) + assert claim.state not in BLOCKING, ( + f"{claim.state}: {claim.id}\n {claim.detail}\n" + f" fix in: {claim.closer or 'unknown'}" + ) + if claim.state == DEFERRED: + # Reached only while the deferral holds; the xfail mark above turns a pass + # here into the failure that forces the declaration to be updated. + pytest.fail(f"deferred: {claim.detail}") + + +class TestClaimsDetectViolations: + """A gate nobody has watched fail is a gate nobody knows the state of. + + ``test_claim`` above asserts that every claim is closed, which passes just as + well for an engine that closes everything. These are the other direction: each + one breaks something a claim is supposed to notice, and asserts it opens. + + Each case mutates a *copy* of the declaration, so the entry the rest of the + suite reads is untouched. + """ + + @staticmethod + def _copy(name: str, **overrides) -> KernelEntry: + import dataclasses + + return dataclasses.replace(ENTRIES[name], **overrides) + + @staticmethod + def _state(entry: KernelEntry, claim_id: str) -> str: + return next(c.state for c in probe(entry).claims if c.id == claim_id) + + def test_a_missing_reference_opens_the_reference_claim(self): + """Absent, not skipped: a kernel writing plausible nonsense reports cleanly.""" + entry = self._copy("matmul_small", reference=None) + assert self._state(entry, "out.c_ptr.reference") == "open" + + def test_a_narrow_reference_opens_the_reference_claim(self): + """A reference at the kernel's own precision reproduces its own overflow.""" + import numpy as np + + def narrow(params, tensors): + a = np.asarray(tensors["a_ptr"], dtype=np.float32) + b = np.asarray(tensors["b_ptr"], dtype=np.float32) + return {"c_ptr": (a @ b).astype(np.float16)} + + entry = self._copy("matmul_small", reference=narrow) + assert self._state(entry, "out.c_ptr.reference") == "open" + + def test_a_reference_for_an_unwritten_tensor_is_undetermined(self): + """The claim set is derived from the store trace, not from the reference. + + So a reference entry the trace never names is the direction that raises no + claim at all rather than a failing one, and the support report carries no + per-kernel claim count in which a vanished claim would show up. The defect + it hides is the one this ledger exists to catch: a kernel that stops writing + a declared output loses that output's comparison instead of failing it. + """ + import numpy as np + + def with_ghost(params, tensors): + got = ENTRIES["matmul_small"].reference(params=params, tensors=tensors) + return {**got, "ghost_ptr": np.zeros(4, dtype=np.float32)} + + entry = self._copy("matmul_small", reference=with_ghost) + assert self._state(entry, "out.ghost_ptr.reference") == "undetermined" + assert not probe(entry).clean + + def test_dropping_a_declared_tolerance_opens_the_reference_claim(self): + """The one widened tolerance in the repository has to be load-bearing. + + A tolerance is the one input to this claim that can be adjusted until it + passes, so the entry that declares its own has to be the entry that needs + it. If this test starts failing, the override is decoration and the row + should lose it rather than keep it. + """ + entry = self._copy("ffn_swiglu", tolerance={}) + assert self._state(entry, "out.out_ptr.reference") == "open" + + def test_a_declared_tolerance_is_stated_by_the_claim_that_passes(self): + """A green run has to say which pair it was green against.""" + claim = next(c for c in probe(ENTRIES["ffn_swiglu"]).claims + if c.id == "out.out_ptr.reference") + assert claim.state == "closed" + assert "declared rtol=" in claim.detail + + def test_a_waiver_for_a_claim_the_kernel_never_raises_is_reported(self): + """A stale waiver reads as a decision about a check that no longer runs.""" + entry = self._copy( + "matmul_small", + waived={"op.linalg.batch_matmul.handler": "this kernel has no batched matmul"}, + ) + states = {c.id: c.state for c in probe(entry).claims} + assert "waived.stale.op.linalg.batch_matmul.handler" in states + assert states["waived.stale.op.linalg.batch_matmul.handler"] == "open" + + def test_an_excuse_with_no_reason_is_rejected_at_declaration(self): + with pytest.raises(ValueError, match="needs a reason"): + self._copy("matmul_small", waived={"parse.regex": " "}) + + def test_a_deferral_naming_no_issue_is_rejected_at_declaration(self): + """The issue is the whole of what a deferral promises. + + Without it the report cannot group the gap, nobody comes back for it, and + what was declared as temporary is a waiver written in the column that is + supposed to expire. + """ + with pytest.raises(ValueError, match="does not name an issue"): + self._copy("matmul_small", + deferred={"parse.regex": "will look at this later"}) + + def test_a_tolerance_for_an_argument_the_kernel_does_not_have_is_rejected(self): + """The ledger reads this mapping with ``.get()``, so a typo is invisible. + + An entry that is never read looks exactly like a widening that took + effect, and the claim it was meant to loosen passes or fails for reasons + the reviewer is no longer looking at. + """ + with pytest.raises(ValueError, match="names no declared tensor"): + self._copy("ffn_swiglu", tolerance={"out_ptrr": (2e-2, 2.0)}) + + def test_a_malformed_tolerance_is_rejected_at_declaration(self): + """Not an (rtol, atol) pair, caught where the row is read rather than deep + in a comparison whose message would be about the kernel.""" + with pytest.raises(ValueError, match="pair of numbers"): + self._copy("ffn_swiglu", tolerance={"out_ptr": 2e-2}) + with pytest.raises(ValueError, match="is negative"): + self._copy("ffn_swiglu", tolerance={"out_ptr": (-1e-2, 2.0)}) + + def test_a_waiver_for_a_claim_that_now_passes_is_reported(self): + """An excuse outliving what made it necessary is misinformation. + + The stale-waiver check above only catches a waiver naming a check the + kernel never raises. This is the other case: the check runs and passes, + and the waiver would otherwise keep reporting it as excused. + """ + entry = self._copy( + "matmul_small", + waived={"parse.regex": "this kernel is known not to parse"}, + ) + states = {c.id: c.state for c in probe(entry).claims} + assert states["parse.regex"] == "closed", ( + "a passing check must report as closed, not as waived" + ) + assert states["waived.unnecessary.parse.regex"] == "open" + + def test_a_deferral_for_a_claim_that_now_passes_is_reported(self): + """A deferral has to expire on its own, and xfail alone does not do it. + + The xfail(strict=True) marker in the gate is applied from the claim's + *current* state, so once the gap closes the claim is no longer deferred, no + marker is applied, and the test simply passes — leaving the deferral in the + declaration with nothing pointing at it. The expiry has to be a claim. + """ + entry = self._copy( + "matmul_small", + deferred={"parse.regex": "#1 — assumed not to be readable yet"}, + ) + states = {c.id: c.state for c in probe(entry).claims} + assert states["parse.regex"] == "closed", ( + "a passing check must report as closed, not as deferred" + ) + assert states["deferred.unnecessary.parse.regex"] == "open" + + +class TestFrontendClaimWithoutBindings: + """``parse.frontend``'s three outcomes, checked without the MLIR bindings. + + The claim itself can only be *decided* where ``mlir_ktdp`` is installed, which + on most machines is CI only. Its three-way branch is this repository's code + though, and getting it wrong is expensive in a specific way: treating a missing + local dependency as a rejection would open the claim on every kernel at once, + and treating a rejection as a skip would hide the one thing this claim exists + to catch. So the branch is exercised here against stand-ins. + """ + + ENTRY = "matmul_small" + + def _state_with(self, monkeypatch, stub) -> str: + import ktir_cpu.mlir_frontend.parser as frontend + + monkeypatch.setattr(frontend, "MLIRFrontendParser", stub) + ledger = probe(ENTRIES[self.ENTRY]) + return next(c.state for c in ledger.claims if c.id == "parse.frontend") + + def test_a_parser_that_accepts_closes_the_claim(self, monkeypatch): + class Accepts: + def parse_module(self, text): + return object() + + assert self._state_with(monkeypatch, Accepts) == "closed" + + def test_a_parser_that_rejects_opens_the_claim(self, monkeypatch): + class Rejects: + def parse_module(self, text): + raise ValueError("op 'ktdp.load' verification failed") + + assert self._state_with(monkeypatch, Rejects) == "open" + + def test_absent_bindings_skip_rather_than_open_the_claim(self, monkeypatch): + """The constructor raises ImportError, not the import — hence the guard.""" + class NoBindings: + def __init__(self): + raise ImportError( + "mlir_ktdp not installed; " + "MLIRFrontendParser is unavailable." + ) + + assert self._state_with(monkeypatch, NoBindings) == "skip" + + +class TestProbingAKernelWithNoDeclaration: + """``probe`` on a bare ``.mlir``, which is the first thing anyone does. + + Requiring a declaration before the tool will say anything inverts the useful + order: the questions worth asking first — which ops have no handler, does it + survive both parse paths — do not depend on one. + """ + + KERNEL = "triton-ktir/layernorm_fwd_ktir.mlir" + + def _ledger(self): + from ktir_cpu.kernelentry import EXAMPLES_DIR + from ktir_cpu.kernelentry.cli import entry_for_bare_kernel + + return probe(entry_for_bare_kernel(EXAMPLES_DIR / self.KERNEL)) + + def test_the_op_and_parse_claims_are_answered(self): + states = {c.id: c.state for c in self._ledger().claims} + handlers = [v for k, v in states.items() if k.endswith(".handler")] + assert handlers and all(v == CLOSED for v in handlers) + assert states["parse.regex"] == CLOSED + + def test_claims_needing_a_declaration_say_so(self): + claims = {c.id: c for c in self._ledger().claims} + assert claims["exec.runs"].state == "open" + assert "tensors=" in claims["exec.runs"].detail + assert claims["cost.derivation"].state == "undetermined" + + def test_a_multi_function_module_asks_which_one(self, tmp_path): + """Guessing would probe an arbitrary function and report on the wrong kernel. + + Written on a file this test makes rather than one under ``examples/``: every + kernel there holds exactly one function today, so a test pointed at one of + them would pass without the branch ever running. + """ + from ktir_cpu.kernelentry.cli import entry_for_bare_kernel + + multi = tmp_path / "two_functions.mlir" + multi.write_text("module {\n" + " func.func @first() { return }\n" + " func.func @second() { return }\n" + "}\n") + with pytest.raises(SystemExit, match="--func"): + entry_for_bare_kernel(multi) + + def test_a_module_with_no_function_says_that_instead(self, tmp_path): + """Not the same message: there is nothing to name with ``--func``. + + The regex parser has a catch-all fallback and validates against no dialect, + so text that is not KTIR at all reaches this point *accepted*, with no + function in it. Sending somebody to name one sends them looking for + something that is not there. + """ + from ktir_cpu.kernelentry.cli import entry_for_bare_kernel + + empty = tmp_path / "not_ktir.mlir" + empty.write_text("module {\n this is not KTIR at all\n}\n") + with pytest.raises(SystemExit, match="declares no function"): + entry_for_bare_kernel(empty) + + def test_a_file_the_parser_rejects_is_a_cell_not_a_crash(self, tmp_path, + monkeypatch): + """One unreadable file must not take the report for the other thirty with it. + + And the cell it gets is the answer, not ``yes``: an earlier version caught + every failure to read out a single kernel as "several functions" and printed + ``yes`` in the column that says the parser read the file — a claim nothing + had checked, in the one document whose job is to not do that. + """ + from ktir_cpu.kernelentry import cli + + unreadable = tmp_path / "undecodable.mlir" + unreadable.write_bytes(b"module {\n \xff\xfe not text\n}\n") + monkeypatch.setattr(cli, "undeclared_kernels", + lambda: [str(unreadable)]) + + assert cli.function_names(unreadable) is None + sv = cli.survey([]) + assert sv.undeclared == [(str(unreadable), "?", ["**no**", "yes"])] + # And the headline counts it out rather than saying "All 1": the sentence a + # reader takes the repository's state from cannot read as complete while a + # row below it says otherwise. + assert "0 of 1 are read by the regex parser" in cli.render_report([], sv) + + +def test_the_generated_documents_do_not_depend_on_this_machine(): + """Both committed documents must be functions of the repository, not the host. + + It is compared verbatim, and `parse.frontend` reads `closed` where the MLIR + bindings are installed and `skip` where they are not — so a report that showed + what the machine writing it happened to see would differ between CI and a + contributor's laptop, and each would call the other's stale, forever. Two things keep it + host-independent: the frontend column is read from the committed record in + `conformance.py` and from a declaration's own excuse, and claims still marked + environment-dependent are left out of the lists below the table. + + `docs/supported_ops.md` is host-independent for a different reason — both its + registries are decorator tables in this package rather than anything the + bindings provide — and it is checked here so that a later column reading a + real parser fails rather than committing what one machine could see. + """ + import ktir_cpu.mlir_frontend.parser as frontend + + from ktir_cpu.kernelentry.cli import generated_docs + + def rendered(ledgers): + return {p.name: t for p, t in generated_docs(ledgers).items()} + + # The first pass reuses the session cache; the second cannot, because the + # stub has to be in place while the kernels are probed. So this test costs one + # extra full probe of every declared kernel — most of this module's runtime, + # and the reason it is one test rather than one per document. + as_is = rendered([ledger_for(n) for n in sorted(ENTRIES)]) + + class NoBindings: + def __init__(self): + raise ImportError( + "mlir_ktdp not installed; " + "MLIRFrontendParser is unavailable." + ) + + original = frontend.MLIRFrontendParser + try: + frontend.MLIRFrontendParser = NoBindings + without = rendered([probe(ENTRIES[n]) for n in sorted(ENTRIES)]) + finally: + frontend.MLIRFrontendParser = original + + differ = sorted(n for n in as_is if as_is[n] != without[n]) + assert not differ, ( + f"{', '.join(differ)} differs between a machine with the MLIR bindings " + "and one without: it is rendering a state it observed rather than one " + "that is written down" + ) + + +@pytest.mark.parametrize("name", ["kernel_support.md", "supported_ops.md"]) +def test_generated_document_is_current(name: str): + """Each generated document must match what the engine reports now. + + Same discipline as a lock file: the documents are generated, and this is what + keeps generation from being optional. Regenerate with:: + + python -m ktir_cpu.kernelentry probe --all --write-report + + Parametrized over the names rather than over ``generated_docs`` itself so that + a document dropped from the writer fails here instead of quietly reducing this + test to the ones that remain. + """ + from ktir_cpu.kernelentry.cli import generated_docs + + ledgers = [ledger_for(name_) for name_ in sorted(ENTRIES)] + docs = {path.name: text for path, text in generated_docs(ledgers).items()} + assert name in docs, f"{name} is no longer generated by generated_docs()" + + path = REPO_ROOT / "docs" / name + assert path.exists(), ( + f"{name} is missing — run " + "`python -m ktir_cpu.kernelentry probe --all --write-report`" + ) + assert path.read_text() == docs[name], ( + f"{name} is stale — regenerate with " + "`python -m ktir_cpu.kernelentry probe --all --write-report`" + ) + + +def test_support_report_covers_conftest_examples(): + """Every kernel in the report is already driven from ``EXAMPLE_PARAMS``. + + The report says so in prose, and the sentence is load-bearing: it is what + separates "this ledger has not asked" from "the simulator cannot". A reader + who does not believe it has to go count 31 files by hand, and the claim rots + silently the moment somebody adds an example without arguments — so assert it + here instead. Failing means the report's opening paragraph has become wrong + and the new file needs either an ``EXAMPLE_PARAMS`` entry or a different + sentence. + + The report cannot check this itself, and neither can it count how many of + those listings are empty: ``ktir_cpu`` ships as a package and + ``tests/conftest.py`` is not part of it. A test is the only place both sides + are importable, so the figure the report prints is pinned here too. + """ + from conftest import EXAMPLE_PARAMS + + from ktir_cpu.kernelentry.cli import (EMPTY_EXECUTE_KWARGS_ROWS, + undeclared_kernels) + + kwargs_by_path: Dict[str, List[dict]] = {} + for entries in EXAMPLE_PARAMS.values(): + for entry in entries: + kwargs_by_path.setdefault( + f"examples/{entry['path']}", []).append(entry["execute_kwargs"]) + undeclared = set(undeclared_kernels()) + rows = undeclared | { + str(entry.mlir_path.resolve().relative_to(REPO_ROOT.resolve())) + for entry in registered().values() + } + missing = sorted(rows - set(kwargs_by_path)) + assert not missing, ( + "docs/kernel_support.md says tests/ already drives every file in its " + "table, from tests/conftest.py::EXAMPLE_PARAMS, and these are not: " + f"{missing}" + ) + + # A path listed twice with different arguments would make "empty" ambiguous, + # so require it of every listing for that path rather than of one of them. + empty = sorted(p for p in undeclared + if all(not kw for kw in kwargs_by_path[p])) + assert len(empty) == EMPTY_EXECUTE_KWARGS_ROWS, ( + "docs/kernel_support.md says execute_kwargs is empty for " + f"{EMPTY_EXECUTE_KWARGS_ROWS} of the files it lists as read-not-declared, " + f"but {len(empty)} are: {empty}. Update " + "ktir_cpu.kernelentry.cli.EMPTY_EXECUTE_KWARGS_ROWS." + ) + + # The same sentence's other half — what the non-empty listings carry — would + # otherwise be prose nothing checks. An index or a size is a scalar; a listing + # supplying an array would be the declaration the sentence says it is not. + nonscalar = sorted( + (path, key) for path in undeclared for kw in kwargs_by_path[path] + for key, value in kw.items() if not isinstance(value, (int, float)) + ) + assert not nonscalar, ( + "docs/kernel_support.md says the non-empty execute_kwargs listings carry raw " + f"HBM element indices or a scalar size, and these carry neither: {nonscalar}" + ) + + +class TestAnExcusedEnvironmentDependentClaim: + """A declared excuse on ``parse.frontend`` must be reported, not filtered out. + + A claim whose outcome depends on the machine is left out of the report's lists, + because a document that named it would say something different depending on who + generated it. An excuse in the declaration removes that dependence: it applies + to every state but ``closed``, and a missing dependency is one of those, so both + machines report the same thing. Keeping such a claim out of the report anyway + would hide a known, written-down gap behind "your machine could not tell" — + the one outcome this report is supposed to make impossible. + + Both cases stub the parser rather than reading whichever answer this machine + happens to give, so they assert the same thing in CI and on a laptop. + """ + + @staticmethod + def _probe(monkeypatch, stub, **overrides): + import dataclasses + + import ktir_cpu.mlir_frontend.parser as frontend + + monkeypatch.setattr(frontend, "MLIRFrontendParser", stub) + return probe(dataclasses.replace(ENTRIES["matmul_small"], **overrides)) + + class _NoBindings: + def __init__(self): + raise ImportError( + "mlir_ktdp not installed; " + "MLIRFrontendParser is unavailable." + ) + + class _Accepts: + def parse_module(self, text): + return object() + + def test_the_deferral_survives_into_the_committed_report(self, monkeypatch): + """Without the bindings the claim is ``skip``, which the deferral covers.""" + ledger = self._probe( + monkeypatch, self._NoBindings, + deferred={"parse.frontend": "#1 — known not to parse there"}) + claim = next(c for c in ledger.claims if c.id == "parse.frontend") + assert claim.state == DEFERRED + assert not claim.env_dependent, ( + "an applied excuse pins the outcome on every machine, so the claim is " + "no longer environment-dependent" + ) + assert "#1 — known not to parse there" in render_report([ledger]) + + def test_an_unnecessary_excuse_for_it_stays_out_of_the_report(self, monkeypatch): + """The other direction: only a machine with the bindings sees it is stale. + + Where the bindings are absent the claim is ``skip``, the excuse applies, and + nothing is unnecessary. Where they are present the check may pass and the + excuse become stale — so the expiry claim's own outcome depends on the + machine, and leaving it out of the document is what keeps the document + host-independent. The gate reads raw states, so CI still fails on it. + """ + ledger = self._probe( + monkeypatch, self._Accepts, + waived={"parse.frontend": "the MLIR frontend cannot see this kernel"}) + claims = {c.id: c for c in ledger.claims} + assert claims["parse.frontend"].state == CLOSED + expiry = claims["waived.unnecessary.parse.frontend"] + assert expiry.state == "open" + assert expiry.env_dependent + assert "waived.unnecessary" not in render_report([ledger]) + + +class TestArgumentSpecs: + """The properties the rest of the ledger reads the declared arguments through. + + ``tensors`` is a table rather than a callable, and two claims depend on what + that table guarantees: the reference is handed a *rebuild* rather than what the + run was given, and the committed derivation rebuilds again to recover shapes. + Neither is sound unless a rebuild is the same tensors. + """ + + def test_rebuilding_an_entry_s_arguments_gives_the_same_values(self): + """Why the reference can be handed a pristine copy at all. + + A kernel whose output argument aliases its input has overwritten that array + by the time the reference runs; handing the reference a rebuild is what + keeps it from comparing the result against itself. That substitution is only + valid if the rebuild is bit-identical to what the kernel was given. + """ + import numpy as np + + for name, entry in sorted(ENTRIES.items()): + first, second = entry.build_tensors(), entry.build_tensors() + assert sorted(first) == sorted(second), name + for arg in first: + np.testing.assert_array_equal( + np.asarray(first[arg]), np.asarray(second[arg]), + err_msg=f"{name}.{arg} is not reproducible") + + def test_two_arguments_of_one_kernel_are_drawn_independently(self): + """Seeded per argument, not per kernel. + + One generator drawn twice gives the second tensor the first one's values + wherever their shapes overlap — and a kernel that swapped its two operands + would then still agree with its reference. This is the property that makes + the swap observable. + """ + import numpy as np + + tensors = ENTRIES["matmul_small"].build_tensors() + a = np.asarray(tensors["a_ptr"], dtype=np.float32).ravel() + b = np.asarray(tensors["b_ptr"], dtype=np.float32).ravel() + overlap = min(a.size, b.size) + assert not np.array_equal(a[:overlap], b[:overlap]) + + def test_a_spec_naming_a_parameter_that_does_not_exist_is_rejected(self): + """At declaration, not on the one kernel at the moment it ran. + + A misspelled parameter name would otherwise surface as a ``KeyError`` from + inside the engine, which reads as a fault in the tool rather than a typo in + the row. + """ + from ktir_cpu.kernelentry.tensorspec import param, zeros + + with pytest.raises(ValueError, match="not in gate_params"): + KernelEntry(name="_typo", func="f", path="ktir/ffn_swiglu.mlir", + gate_params={"rows": 1}, + tensors={"out": zeros("rows"), "n": "colums"}) + with pytest.raises(ValueError, match="not in gate_params"): + KernelEntry(name="_typo", func="f", path="ktir/ffn_swiglu.mlir", + gate_params={"rows": 1}, + tensors={"n": param("colums", "i32")}) + + +def test_discovering_declarations_twice_is_a_no_op(): + """Two test modules discover at import time, and one session collects both. + + Declarations are loaded by file location rather than by module name, so they + never enter ``sys.modules`` and nothing stops a second call from executing the + same file again — at which point ``register_entry`` rejects the duplicate name + and takes down collection for the entire run, not just this module. That failure + only appears when both modules are collected together, which is what + ``uv run pytest tests/`` does and what running either file alone does not. + """ + before = dict(registered()) + discover_all() + discover_all() + assert dict(registered()) == before + + +class TestPricingAudit: + """`zero` in the registry is two facts, and this is the test that splits them. + + ``@register()`` defaults ``latency_category`` to ``"zero"``, so registering an + op without naming a category makes it free and says nothing. The record in + ``ktir_cpu/kernelentry/pricing.py`` is what makes that a decision; these tests + hold the record to the registry in both directions, because the direction that + rots is not the one the check exists for. + + Repository-wide rather than per kernel, and that came from measurement: asked + once per kernel, this check produced 66 of a prototype's 74 false positives, + since the ops it flags are the same structural ones in every kernel. + """ + + def test_every_zero_priced_op_is_free_by_decision_or_recorded_as_unjudged(self): + findings = audit() + assert not findings, ( + "the pricing record and the op registry disagree:\n" + + "\n".join(f" {finding}" for finding in findings) + ) + + def test_every_reason_is_written(self): + """An entry with no reason is a waiver, not a decision. + + The same rule ``KernelEntry`` applies to ``waived`` and ``deferred``: an + excuse with no reason is indistinguishable from an oversight. + """ + for mapping, which in ((ZERO_COST_OPS, "ZERO_COST_OPS"), + (UNJUDGED_ZERO_OPS, "UNJUDGED_ZERO_OPS")): + for op, reason in mapping.items(): + assert reason.strip(), f"{which}[{op!r}] needs a reason" + + def test_every_unjudged_op_names_the_issue_that_would_settle_it(self): + """The half of the record that expires has to say what would expire it. + + ``ZERO_COST_OPS`` is a decision and stands on its reason alone. An + unjudged op is an open question, and an open question with no issue behind + it is how the list becomes permanent. + """ + for op, reason in UNJUDGED_ZERO_OPS.items(): + assert re.search(r"#\d+", reason), ( + f"UNJUDGED_ZERO_OPS[{op!r}] does not name an issue. Without one " + "the entry is a decision to leave it unpriced, spelled as a " + "question." + ) + + def test_a_newly_registered_op_is_not_silently_free(self): + """The state AC7 of #209 exists to reject.""" + with registry.temp_registry(): + registry._REGISTRY["ktdp.invented_op"] = lambda op, ctx, env: None + registry._LATENCY_CATEGORIES["ktdp.invented_op"] = "zero" + findings = audit() + assert [f.op for f in findings] == ["ktdp.invented_op"] + assert findings[0].kind == UNLISTED + assert "pricing.py" in findings[0].fix + + def test_an_entry_that_outlived_its_zero_is_reported(self): + """Deciding an op's price has to take its entry with it. + + This is the direction that rots: #211 will price three of these, and a + record still calling them unjudged afterwards would describe a repository + that no longer exists. + """ + with registry.temp_registry(): + registry._LATENCY_CATEGORIES["linalg.fill"] = "compute_float" + findings = audit() + assert [(f.op, f.kind) for f in findings] == [("linalg.fill", PRICED)] + + def test_an_entry_whose_op_is_gone_is_reported(self): + """A rename or a removal must not leave the record behind.""" + with registry.temp_registry(): + del registry._REGISTRY["ktdp.coreid"] + findings = audit() + assert [(f.op, f.kind) for f in findings] == [("ktdp.coreid", UNREGISTERED)] + + def test_claiming_both_answers_for_one_op_is_reported(self, monkeypatch): + """Free by decision and not yet judged are mutually exclusive.""" + monkeypatch.setitem(UNJUDGED_ZERO_OPS, "scf.for", "undecided #1") + findings = audit() + assert [(f.op, f.kind) for f in findings] == [("scf.for", BOTH)] + + def test_the_audit_cannot_see_a_wrong_category(self): + """Stated as a test so the limitation is not mistaken for coverage. + + The check compares against ``zero``, so an op billed to the wrong non-zero + unit passes it untouched — an integer compare charged to the float pipe + was a real defect in this repository, and this audit would not have found + it. Only a reader who knows the op's semantics does. + """ + with registry.temp_registry(): + registry._LATENCY_CATEGORIES["arith.addi"] = "compute_matmul" + assert audit() == [] From 7b646abdbadfe10a984a9154c44c6d7f66d1e066 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Tue, 25 Aug 2026 10:35:16 -0400 Subject: [PATCH 8/9] Correct the op inventories docs/supported_ops.md contradicts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generating an op matrix makes two hand-maintained inventories checkable for the first time, and both were wrong in the same direction — reporting as missing what has been implemented for months. `docs/gap_analysis.md` rows 13–19 (`arith.cmpf`, `negf`, `absf`, `minf`, `minnumf`, the float casts, the signed and bitwise int ops) and rows 28–29 (`tensor.extract_slice`, `tensor.insert_slice`) were ❌ or 🟡 with handlers in the registry; `cmpf` and `negf` are used by an example. They are now ✅ with the handler cited, and the two summary lists at the end agree with them. The Arith and Tensor preambles no longer carry an inventory at all — they point at the generated document, because a list copied into prose is what let those rows sit wrong, and re-copying it correctly would only restart the clock. The coverage-backlog paragraph named three RFC-mentioned ops as absent; `memref.subview` is the only one still absent, and it now names just that one. README's Supported Subset had the same problem from the other side: it described rectangular coordinate sets only, and omitted `base_map`, `access_tile_order`, distributed and indirect views, and cross-core reduction. Its Not-Supported list named `ktdp.transfer` / `ktdp.reduce`, which no longer exist, and `tensor.insert_slice`, which does; and it called `coordinate_set` unenforced, where what is actually unchecked is overlap between a distributed view's partitions — `find_partition` returns the first match, which RFC 0682 §3.3 leaves unspecified. Both lists are rewritten and now link the two generated documents rather than restating them. New gap row 36a, found by the same matrix: `ktdp.region_terminator` has a frontend adapter and no execution handler. The MLIR bindings' region walk is documented in `mlir_frontend/parser.py` as emitting an implicit terminator that never appears in text IR; an adapter is installed so the walk survives, `adapt_block` keeps it in the region's op list, and `_execute_op` would raise `Unknown operation` for it there. No kernel under `examples/` reaches it, and the reason is not the one it looks like. Three inter-tile kernels do get through the frontend — `rmsnorm_4core_2x2.mlir` and the two rewritten earlier in this branch — and walking every op the frontend produces for each, into the `inter_tile_produce` and `inter_tile_reduce` regions, finds `ktdp.yield_partial` and `ktdp.yield_reduced` terminating them and no `ktdp.region_terminator` in any of the three. So the row records an asymmetry that is real and a construct whose trigger is not established, rather than a kernel that would abort. Row 2a is corrected by the same leg. It described the four-op design's upstream status rather than this repository, and the status it described is no longer the case; it now states what is true of the code here — `ktdp.inter_tile_reduce` takes its result type as a reshape target, the dialect verifies the result against the future's partial type instead, and the three `ring_reduce*` kernels are the only non-identity reduce reshape and the only inter-tile kernels the frontend refuses. Rows 2b-2d are left alone: their implementation status is unmeasured here. §I and §J are left alone; they are dated 2026-05-30 and are a different kind of claim. Signed-off-by: WarningRan --- README.md | 23 +++++++++++++++-------- docs/gap_analysis.md | 42 +++++++++++++++++++++++------------------- 2 files changed, 38 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index b02974d..0a1fac7 100644 --- a/README.md +++ b/README.md @@ -184,20 +184,27 @@ This keeps operation knowledge (behavior + cost classification) co-located in on This interpreter covers a **subset** of KTIR ([RFC 0682](https://github.com/torch-spyre/RFCs/blob/main/0682-KtirSpec/0682-KtirSpecRFC.md)). The following are supported: -- Embarrassingly parallel kernels (no inter-core communication required) -- `ktdp.load` / `ktdp.store` with rectangular-slice semantics -- `ktdp.construct_access_tile` (rectangular tiles only) -- Arithmetic, math, and linalg dialect ops (see `ktir_cpu/dialects/`) +- Embarrassingly parallel kernels, plus cross-core reduction through the experimental `ktdp.inter_tile_produce` / `ktdp.inter_tile_reduce` pair +- `ktdp.load` / `ktdp.store` over general affine coordinate sets, not only rectangular slices +- `ktdp.construct_access_tile` with `base_map`, `access_tile_set` and `access_tile_order` evaluated +- `ktdp.construct_memory_view`, `construct_distributed_memory_view` and `construct_indirect_access_tile` +- Arithmetic, math, linalg and tensor dialect ops (see `ktir_cpu/dialects/`) - `scf.for` / `scf.if` control flow - Multi-core grid execution - Cycle-approximate latency estimation +Per-op status is generated from the registries rather than listed here: [`docs/supported_ops.md`](docs/supported_ops.md). Per-kernel status: [`docs/kernel_support.md`](docs/kernel_support.md). + **Not yet supported or unreliable:** -- `ktdp.construct_distributed_memory_view` — not implemented -- `ktdp.construct_indirect_access_tile` — not implemented -- `ktdp.transfer` / `ktdp.reduce` (communication ops) — present but **unreliable**: the multi-round communication model re-executes the entire function per round, causing incorrect latency accumulation and potential correctness issues with cyclic communication patterns. See `docs/gap_analysis.md` for details. -- `tensor.extract_slice` / `memref.subview` +- `memref.subview`, and the `memref` dialect generally — no module exists +- `linalg.map` +- `scf.parallel` / `scf.forall` / `scf.reduce` +- `ktdp.inter_tile_consume` (broadcast) and `ktdp.inter_tile_reduce_scatter` — the reduce half of the four-op inter-tile design is implemented, these two are not +- `ktdp.region_terminator` — has an MLIR frontend adapter but no execution handler; no kernel under `examples/` reaches it +- `coordinate_set` overlap between the partitions of a distributed view is unchecked — `find_partition` returns the first match, which RFC 0682 §3.3 leaves unspecified + +See [`docs/gap_analysis.md`](docs/gap_analysis.md) for the full conformance picture. `ktdp.transfer` / `ktdp.reduce` no longer exist — the experimental `ktdp.inter_tile_produce` / `ktdp.inter_tile_reduce` pair replaced them. ## RFC conformance diff --git a/docs/gap_analysis.md b/docs/gap_analysis.md index 3e2cc4f..baf01ab 100644 --- a/docs/gap_analysis.md +++ b/docs/gap_analysis.md @@ -13,7 +13,7 @@ |---|---------------|--------|-------| | 1 | `ktdp.construct_distributed_memory_view` | ✅ | Handler and parser implemented in `ktir_cpu/dialects/ktdp_ops.py`; produces `DistributedMemRef` (composition of N per-partition `MemRef`s). Per-partition routing at access time via `MemoryOps.distributed_tile_access` → `DistributedTileRef`. Tests in `tests/test_distributed_view.py`. | | 2 | `ktdp.construct_indirect_access_tile` | ✅ | Handler and parser implemented in `ktir_cpu/dialects/ktdp_ops.py`; tests passing in `tests/test_indirect_access.py`. Both `ktdp.load` (gather, via `MemoryOps.indirect_load`) and `ktdp.store` (scatter, via `MemoryOps.indirect_store`) accept `IndirectAccessTile` (#44 closed). | -| 2a | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` (four-op design, reduce path) | 🧪 experimental | **Spec status: not yet merged** — the four-op design lives in [ktir-mlir-frontend PR #23](https://github.com/torch-spyre/ktir-mlir-frontend/pull/23). Implemented here ahead of merge so the simulator can exercise it. Op names, attribute keys, and `!ktdp.tile_future<...>` type may change to track the upstream PR. Implementation: handlers and parsers in `ktir_cpu/dialects/ktdp_ops.py`; `TileFuture` per-core handle binds the local partial and a not-yet-running `RingReduceBackend`; reduce delivery triggers the backend with the combiner region as `reduce_fn`. Per-core wire bytes flow to the latency tracker via `Tile.comm_bytes` (mirrors `unique_sticks`). End-to-end tests: `tests/test_examples.py::TestRingReduceExecution` (contiguous reduce groups) and `tests/test_examples.py::TestMulticoreSdpaExecution` (strided reduce groups, 32 cores). Replaces the legacy `ktdp.reduce` / `ktdp.transfer` ops. See `docs/cross_core_scheduling.md`. | +| 2a | `ktdp.inter_tile_produce` + `ktdp.inter_tile_reduce` (four-op design, reduce path) | 🧪 experimental | The four-op design lives in [ktir-mlir-frontend PR #23](https://github.com/torch-spyre/ktir-mlir-frontend/pull/23). **One divergence stands:** `ktdp.inter_tile_reduce` here takes its result type as a reshape target (`_result_shape` → `attach_reshape`), while the dialect verifies `result types must match future partial types`, so a reduce that reshapes cannot be written in the dialect's own assembly. `examples/ktir/ring_reduce.mlir`, `examples/ktir/ring_reduce_inner_loop.mlir` and `examples/latency/ring_reduce_multi_group.mlir` reduce `tensor<1x128xf16>` to `tensor<128xf16>` — the only non-identity reduce reshape here — and are read by the regex parser only; the other three inter-tile kernels reduce shape-to-shape and the frontend accepts them. Implementation: handlers and parsers in `ktir_cpu/dialects/ktdp_ops.py`; `TileFuture` per-core handle binds the local partial and a not-yet-running `RingReduceBackend`; reduce delivery triggers the backend with the combiner region as `reduce_fn`. Per-core wire bytes flow to the latency tracker via `Tile.comm_bytes` (mirrors `unique_sticks`). End-to-end tests: `tests/test_examples.py::TestRingReduceExecution` (contiguous reduce groups) and `tests/test_examples.py::TestMulticoreSdpaExecution` (strided reduce groups, 32 cores). Replaces the legacy `ktdp.reduce` / `ktdp.transfer` ops. See `docs/cross_core_scheduling.md`. | | 2b | `ktdp.inter_tile_consume` (broadcast delivery) | ❌ experimental | Same upstream PR. Not implemented. | | 2c | `ktdp.inter_tile_reduce_scatter` | ❌ experimental | Same upstream PR. Not implemented. | | 2d | `producer_dependency_per_consumer` (per-tile sync) | 🟡 experimental | Same upstream PR. Parsed and stored on the delivery op AST; runtime is full-barrier mode (waits for all producers). Per-tile mode unimplemented. | @@ -54,24 +54,25 @@ Currently implemented: `scf.for`, `scf.if`, `scf.yield`. This section is best read as a **coverage backlog**, not a list of equally strong RFC violations. The RFC explicitly defines the `ktdp` surface and -explicitly calls out only a small subset of non-`ktdp` ops. Missing -`linalg.add`, `tensor.extract_slice`, and `memref.subview` are more directly -grounded in the RFC text than every absent op from the broader Arith/Math -dialects. +explicitly calls out only a small subset of non-`ktdp` ops, so an absent op the +RFC names — `memref.subview` is the one still missing — is a stronger finding +than an absent op from the broader Arith/Math dialects. ### Arith dialect -The spec references the [full Arith dialect](https://mlir.llvm.org/docs/Dialects/ArithOps/). Currently implemented: `addf`, `subf`, `mulf`, `divf`, `addi`, `subi`, `muli`, `divui`, `remui`, `constant`, `maxf`, `maxnumf`, `extf`, `truncf`, `index_cast`, `sitofp`, `cmpi`, `select`. +The spec references the [full Arith dialect](https://mlir.llvm.org/docs/Dialects/ArithOps/). + +Which arith ops have a handler is no longer listed here. It is generated per op in [`docs/supported_ops.md`](supported_ops.md), read from the registry itself and held current by `tests/test_kernelentry.py` — an inventory copied into prose is what let rows 13–19 below sit at ❌ for months after they were implemented. | # | Operation | Status | Notes | |---|-----------|--------|-------| -| 13 | `arith.cmpf` | ❌ | float compare — only `arith.cmpi` (int compare) exists | -| 14 | `arith.negf` | ❌ | | -| 15 | `arith.absf` | ❌ | | -| 16 | `arith.minf` | ❌ | only `maxf` / `maxnumf` exist | -| 17 | `arith.minnumf` | ❌ | | -| 18 | `arith.fptosi`, `arith.fptoui`, `arith.uitofp` | 🟡 | only `sitofp` exists | -| 19 | `arith.divsi`, `arith.remsi`, `arith.andi`, `arith.ori`, `arith.xori`, `arith.ceildivsi`, `arith.floordivsi` | ❌ | only unsigned variants `divui`/`remui` exist | +| 13 | `arith.cmpf` | ✅ | `ktir_cpu/dialects/arith_ops.py:170`, priced `COMPUTE_FLOAT`; shares a parser with `arith.cmpi`. Used by `examples/ktir/reduce_multiop.mlir` to build a max fold. | +| 14 | `arith.negf` | ✅ | Registered in `ktir_cpu/dialects/arith_ops.py`; used by `examples/ktir/ffn_swiglu.mlir`. | +| 15 | `arith.absf` | ✅ | Registered in `ktir_cpu/dialects/arith_ops.py`. | +| 16 | `arith.minf` | ✅ | Registered alongside `maxf` / `maxnumf` / `minimumf` / `minnumf` in `ktir_cpu/dialects/arith_ops.py`. | +| 17 | `arith.minnumf` | ✅ | Same registration as row 16. | +| 18 | `arith.fptosi`, `arith.fptoui`, `arith.uitofp` | ✅ | All three registered in `ktir_cpu/dialects/arith_ops.py` alongside `sitofp`. | +| 19 | `arith.divsi`, `arith.remsi`, `arith.andi`, `arith.ori`, `arith.xori`, `arith.ceildivsi`, `arith.floordivsi` | ✅ | All registered in `ktir_cpu/dialects/arith_ops.py` alongside the unsigned variants. | ### Math dialect @@ -97,12 +98,12 @@ The spec references the [full Linalg dialect](https://mlir.llvm.org/docs/Dialect ### Tensor dialect -Currently implemented: `tensor.splat`, `tensor.extract`, `tensor.expand_shape`, `tensor.collapse_shape`. +Which tensor ops have a handler is generated in [`docs/supported_ops.md`](supported_ops.md), for the reason given under Arith above. | # | Operation | Status | Notes | |---|-----------|--------|-------| -| 28 | `tensor.extract_slice` | ❌ | Spec explicitly calls this out for tensor-level slicing | -| 29 | `tensor.insert_slice`, `tensor.collapse_shape` | 🟡 | `collapse_shape` implemented; `insert_slice` still missing | +| 28 | `tensor.extract_slice` | ✅ | Handler at `ktir_cpu/dialects/tensor_ops.py:217`, frontend adapter at `ktir_cpu/mlir_frontend/parser.py:772`. No example under `examples/` uses it, so only unit tests cover it. | +| 29 | `tensor.insert_slice`, `tensor.collapse_shape` | ✅ | Both implemented. `insert_slice` handler at `ktir_cpu/dialects/tensor_ops.py:258`, regex parser at `:714`, frontend adapter at `ktir_cpu/mlir_frontend/parser.py:786`. No example under `examples/` uses it; `tests/test_dialects_exec.py` and `tests/test_dialects_parse.py` cover it. | ### MemRef dialect @@ -122,6 +123,7 @@ The spec explicitly mentions `memref.subview` for view-based transformations. ** | 34 | `ktdp.load` only implements rectangular slice semantics | ✅ | Now enumerates coordinates from `access_tile_set` and applies `access_tile_order`; supports general polyhedral regions. | | 35 | `ktdp.store` only implements rectangular slice semantics | ✅ | Same coordinate-set enumeration as load. | | 36 | `module { }` is tolerated, but module-level structure is not modeled | 🟡 | The parser can find `func.func` inside a `module { ... }` wrapper, but it does not model module-level attributes, declarations, or non-function top-level constructs. | +| 36a | `ktdp.region_terminator` has a frontend adapter and no execution handler | ❌ | The MLIR bindings' region walk emits an implicit terminator that never appears in text IR; `ktir_cpu/mlir_frontend/parser.py:238` installs an adapter for it so the walk survives, and `adapt_block` keeps it in the region's op list. The interpreter has no handler, so `_execute_op` would raise `Unknown operation` if a kernel ever reached one. None does: walking every op the frontend produces for the three inter-tile kernels it accepts — `examples/latency/rmsnorm_4core_2x2.mlir`, `examples/ktir/ffn_swiglu_4core.mlir` and `examples/sdsc/sdpa_pv_ksplit.mlir`, regions included — finds `ktdp.yield_partial` / `ktdp.yield_reduced` terminating them and no `ktdp.region_terminator` anywhere. So the asymmetry is real and the adapter unreached, but which construct makes the bindings emit an implicit terminator is not established. `docs/supported_ops.md` shows it as the one op with a frontend adapter and no handler. | ## G. Parser Limitations @@ -142,16 +144,16 @@ Blocks running spec-compliant KTIR programs: Limits dialect coverage for real-world kernels: - **#9–12**: ❌ SCF parallel/reduce operations -- **#13–19**: ❌/🟡 Many standard arith ops (cmpf, negf, absf, minf, signed int ops) - **#20–24**: ✅ All math ops now implemented (log2, log1p, tanh, sin, cos, rsqrt, absf, ceil, floor, erf, powf, fma) -- **#28, #30–31**: ❌ `tensor.extract_slice`, entire `memref` dialect +- **#30–31**: ❌ Entire `memref` dialect +- **#36a**: ❌ `ktdp.region_terminator` has a frontend adapter and no execution handler — an asymmetry no kernel under `examples/` reaches, on a construct whose trigger is not established - **#32**: 🟡 Dynamic sizes/strides not supported ### Lower Priority Extensibility and completeness: - **#3, #4**: ❌/🟡 Dynamic access tile dimensions, generic `MemorySpaceAttr` -- **#27, #29**: 🟡 Remaining linalg/tensor ops (`linalg.map`, `tensor.insert_slice`) +- **#27**: 🟡 `linalg.map` still missing (`broadcast` and `transpose` implemented) - **#36, #39**: 🟡 Module-level handling, full function signatures ### Resolved @@ -159,7 +161,9 @@ Extensibility and completeness: - **#6, #7, #8**: ✅ `access_tile_set`, `access_tile_order`, `base_map` - **#20–24**: ✅ All math ops (rsqrt, log2, log1p, tanh, sin, cos, absf, ceil, floor, erf, powf, fma) - **#25**: ✅ `linalg.add` +- **#13–19**: ✅ The standard arith ops (cmpf, negf, absf, minf/minnumf, the float casts, the signed and bitwise int ops) - **#26**: ✅ `linalg.generic` +- **#28, #29**: ✅ `tensor.extract_slice`, `tensor.insert_slice` - **#33, #34, #35**: ✅ Access tile coordinate semantics - **#37, #38**: ✅ Affine expression evaluation and alias support From 5893d20b3165313739b23b2274d07fc5e2b7f069 Mon Sep 17 00:00:00 2001 From: WarningRan Date: Thu, 27 Aug 2026 10:20:07 -0400 Subject: [PATCH 9/9] Let a cold probe read a kernel outside this repository MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `probe ` refused any `.mlir` not under `examples/`, telling the caller to move it there first. That inverts what the tool is for. The kernels this interpreter is asked to support are versioned elsewhere; `examples/` holds fixtures, and the two questions a cold probe answers — whether every op in the file has an execution handler, and whether both parse paths accept it — are about the simulator rather than about where the file sits. A cold probe writes nothing, so a path outside the tree leaves nothing behind to go stale. The entry keeps the relative form when the file is under `examples/` and the absolute path when it is not, so a cold probe and a declared entry name the same kernel the same way. Declarations are deliberately not relaxed with it: a row in the committed report names a repository-relative path, and a row pointing at one person's own disk is a row nobody else can regenerate. Signed-off-by: WarningRan --- docs/kernelentry.md | 5 +++++ ktir_cpu/kernelentry/__init__.py | 5 ++++- ktir_cpu/kernelentry/cli.py | 18 +++++++++++------- tests/test_kernelentry.py | 20 ++++++++++++++++++++ 4 files changed, 40 insertions(+), 8 deletions(-) diff --git a/docs/kernelentry.md b/docs/kernelentry.md index b8c4b47..91771a2 100644 --- a/docs/kernelentry.md +++ b/docs/kernelentry.md @@ -27,6 +27,11 @@ function. A path is read cold — no declaration is consulted even if one exists this step is available before step 2 rather than after it. Once the kernel is declared, name it instead: `probe my_kernel`. +The file does not have to be in this repository. `examples/` holds fixtures, while +the kernels this interpreter is asked to support are versioned elsewhere, so a cold +probe takes any path — a checkout beside this one, a scratch file you are still +editing. Only the steps that write anything need the kernel to live here. + It answers the questions that need no declaration: which of your kernel's ops have no execution handler, and whether it survives both parse paths. Anything that needs to *run* the kernel reports that there is no declaration yet. `probe` writes diff --git a/ktir_cpu/kernelentry/__init__.py b/ktir_cpu/kernelentry/__init__.py index 55f0a36..2c8cd77 100644 --- a/ktir_cpu/kernelentry/__init__.py +++ b/ktir_cpu/kernelentry/__init__.py @@ -115,7 +115,10 @@ class KernelEntry: *path* is the kernel's ``.mlir``, relative to ``examples/``, and it is the kernel's source — hand-written IR, or captured compiler output such as ``examples/triton-ktir/``, which is "kernels as the Triton -> KTIR path emits - them". + them". A declaration's path is relative because the report names it that way, + and a row pointing at somebody's own disk is a row nobody else can regenerate. + An absolute path is what a kernel read cold from outside this repository carries + instead: it is answering the questions that need no declaration. *gate_params* is deliberately a reduced shape: gate cost is driven by shape and not by the number of entries, which is why ``examples/latency/`` holds reduced diff --git a/ktir_cpu/kernelentry/cli.py b/ktir_cpu/kernelentry/cli.py index fd41c60..bb80d86 100644 --- a/ktir_cpu/kernelentry/cli.py +++ b/ktir_cpu/kernelentry/cli.py @@ -185,11 +185,14 @@ def entry_for_bare_kernel(path: Path, func: Optional[str] = None) -> KernelEntry "name one with --func" ) func = names[0] - try: - rel = str(path.resolve().relative_to(EXAMPLES_DIR)) - except ValueError: - raise SystemExit(f"{path} is not under examples/; move it there first") - return KernelEntry(name=f"{path.stem} (no declaration)", func=func, path=rel) + # Relative under ``examples/``, so a cold probe and a declared entry name the + # same kernel the same way; absolute for a kernel versioned outside this + # repository, which a cold read takes because it writes nothing. + resolved = path.resolve() + where = (resolved.relative_to(EXAMPLES_DIR) + if resolved.is_relative_to(EXAMPLES_DIR) else resolved) + return KernelEntry(name=f"{path.stem} (no declaration)", func=func, + path=str(where)) def _resolve(target: Optional[str], use_all: bool, @@ -786,8 +789,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ): child = sub.add_parser(name, help=helptext) child.add_argument("target", nargs="?", - help="a declared kernel's name, or the path of a " - ".mlir under examples/ to read cold") + help="a declared kernel's name, or the path of any " + ".mlir to read cold, inside this repository " + "or not") child.add_argument("--func", default=None, help="function to probe, when a bare .mlir declares " "more than one") diff --git a/tests/test_kernelentry.py b/tests/test_kernelentry.py index 04da405..1dce879 100644 --- a/tests/test_kernelentry.py +++ b/tests/test_kernelentry.py @@ -338,6 +338,26 @@ def test_claims_needing_a_declaration_say_so(self): assert "tensors=" in claims["exec.runs"].detail assert claims["cost.derivation"].state == "undetermined" + def test_a_kernel_outside_this_repository_is_probed_where_it_lies(self, tmp_path): + """The case the tool is for: a kernel somebody brings, versioned elsewhere. + + The entry keeps the absolute path, there being no relative form for it to + have, and every claim that needs no declaration is answered anyway. + """ + from ktir_cpu.kernelentry import EXAMPLES_DIR + from ktir_cpu.kernelentry.cli import entry_for_bare_kernel + + outside = tmp_path / "brought_from_elsewhere.mlir" + outside.write_text((EXAMPLES_DIR / self.KERNEL).read_text()) + + entry = entry_for_bare_kernel(outside) + assert entry.mlir_path.resolve() == outside.resolve() + + states = {c.id: c.state for c in probe(entry).claims} + handlers = [v for k, v in states.items() if k.endswith(".handler")] + assert handlers and all(v == CLOSED for v in handlers) + assert states["parse.regex"] == CLOSED + def test_a_multi_function_module_asks_which_one(self, tmp_path): """Guessing would probe an arbitrary function and report on the wrong kernel.