|
| 1 | +# Declarative graph-rewrite for the tinygrad port |
| 2 | + |
| 3 | +## Context / problem |
| 4 | + |
| 5 | +The ML port (tinygrad-in-RSScript) reimplements tinygrad's **declarative** |
| 6 | +scheduler/codegen as **bespoke imperative recursion**. tinygrad expresses every |
| 7 | +rewrite as data — `PatternMatcher([(UPat(...), lambda node: ...)])` applied by one |
| 8 | +generic `graph_rewrite(sink, pm, ...)` driver with built-in fixpoint + |
| 9 | +memoization. The port hand-rolls each pass (`indexing_rangeify_rewrite_node`, |
| 10 | +`graph_rewrite_memo_d`, …) with manual `olds/news` memo arrays. |
| 11 | + |
| 12 | +That paraphrase-of-a-declarative-engine is the root cause of the port's worst |
| 13 | +bugs. Canonical example (chained-reduce): tinygrad's rule is |
| 14 | +`(UPat(BUFFERIZE), bufferize_to_store)` — it matches a *bare* BUFFERIZE node and |
| 15 | +writes using **that node's own ranges**. The port re-expressed it as "match |
| 16 | +`INDEX(STAGE)` adjacency, use the **consumer's** ranges," which (a) conflates |
| 17 | +write-ranges with read-ranges and (b) silently breaks when another bespoke pass |
| 18 | +(`flatten_bufferize`) perturbs the graph shape. A real PatternMatcher makes the |
| 19 | +scheduler a near-mechanical **transliteration** of `rangeify.py` — rules match the |
| 20 | +same node shapes tinygrad matches, so this whole class of divergence can't arise. |
| 21 | +Paraphrase drifts; transliteration doesn't. |
| 22 | + |
| 23 | +## What RSScript already has (premises that are NOT gaps) |
| 24 | + |
| 25 | +Grounded against the implementation — the original ask over-stated the gap: |
| 26 | +- **Closures are first-class runtime values**: `VmValue::Closure(Rc<VmClosure>)`, |
| 27 | + `MakeClosure`/`CallClosure` in the reg-VM; closures capture (incl. owned) |
| 28 | + values; `List.fold`'s `|acc, x|` is one in use. |
| 29 | +- **An `Fn(...)` type exists** in the checker (e.g. `noescape Fn()`), with |
| 30 | + function-value desugaring (`syntax/function_value_desugar.rs`). |
| 31 | +- **Tuple destructuring is done**: `let (a, b, c) = expr` |
| 32 | + (`syntax/desugar.rs::expand_tuple_destructuring`). No work needed. |
| 33 | + |
| 34 | +## The actual keystone gap: escaping / storable closures |
| 35 | + |
| 36 | +A `PatternMatcher` is a **stored list of closures that outlives its definition |
| 37 | +site** and is called repeatedly by the rewrite driver. RSScript today restricts |
| 38 | +closures to **`noescape`** — they may be *forwarded down* into a call but not |
| 39 | +*stored up* (the diagnostic: "Forwarding a local closure is only allowed when the |
| 40 | +target parameter is `noescape Fn()`"). This is a deliberate review-first |
| 41 | +restriction (an escaping closure that captures is a retention/aliasing concern). |
| 42 | + |
| 43 | +So the precise need is **owned, escaping closures storable in a collection** |
| 44 | +(`List<Fn(UOp) -> UOp>` of rules). The VM already represents them (a `Closure` is |
| 45 | +an `Rc`, trivially storable); the gap is the **checker's escape rule** plus a |
| 46 | +**storable `Fn` value type** in signatures. Relaxing this is the one real language |
| 47 | +decision — and it must clear RSScript's admission bar: an escaping+capturing |
| 48 | +closure has to *phrase as a reviewer question* (what it captured, whether it can |
| 49 | +retain), e.g. via an explicit `own`/move capture annotation and the `Fn` type |
| 50 | +surfacing in the signature, so nothing is implicit. |
| 51 | + |
| 52 | +## Layered plan (priority order) |
| 53 | + |
| 54 | +1. **L1 — language (keystone): owned/escaping storable closures.** Allow a |
| 55 | + closure that captures owned values to escape into a value of a first-class |
| 56 | + `Fn(args) -> ret` type and be stored in containers. Make the capture + escape |
| 57 | + explicit in the signature (review-first). Everything else rides on this. |
| 58 | +2. **L2 — library/runtime: PatternMatcher + graph_rewrite.** With L1, build a |
| 59 | + `UOp`/`UPat` type and a native `graph_rewrite(sink, rules, bottom_up, name)` |
| 60 | + driver (fixpoint + memoization) plus a gated `toposort(gate: Fn(UOp)->Bool)` — |
| 61 | + as a runtime/library facility (the Tensor-kernel pattern), NOT more language |
| 62 | + surface. The port then transliterates `rangeify.py` rule-for-rule. Also |
| 63 | + retroactively simplifies divmod/symbolic (already faking a fixpoint via |
| 64 | + `graph_rewrite_memo_d`) and codegen. |
| 65 | +3. **L3 — ergonomic sugar (independent): iteration.** No `for`-in / comprehension |
| 66 | + exists (the `for` keyword is `protocol … for Type` impls). `while i < |
| 67 | + List.len { List.get(...) }` is ~3× the Python and a transcription-error source. |
| 68 | + `for x in xs` / `map`/`filter` is a separate, independently-weighable nicety. |
| 69 | +- **Done already:** tuple destructuring (`let (a,b,c) = …`). |
| 70 | + |
| 71 | +## The decision |
| 72 | + |
| 73 | +The win is declarative *fidelity* — transliteration over paraphrase eliminates a |
| 74 | +whole bug class, and it's the highest-leverage thing the port could get. The only |
| 75 | +deep change is **L1 (escaping storable closures)**; L2 is then a library, L3 is |
| 76 | +optional sugar. L1 is exactly the kind of expressivity RSScript has deliberately |
| 77 | +withheld, so the gate is: can an escaping+capturing closure be made *explicit |
| 78 | +enough in the signature* to satisfy the review-first contract? If yes, this is the |
| 79 | +next language-roadmap item. If the answer is "only with implicit retention rules," |
| 80 | +it stays out and the port keeps a (smaller, well-tested) hand-rolled engine. |
0 commit comments