Experimental compiler approach for Tensor library

tensorcery_compiler.nim

## Tattletale
## Copyright (c) 2026 Mamy André-Ratsimbazafy
## Licensed and distributed under either of
##   * MIT license (license terms in the root directory or at http://opensource.org/licenses/MIT).
##   * Apache v2 license (license terms in the root directory or at http://opensource.org/licenses/LICENSE-2.0).
## at your option. This file may not be copied, modified, or distributed except according to those terms.

## tensorcery_compiler, an exprgrad-style kernel DSL over the tensorcery engine.
## One loop-free algorithm statement plus a schedule block lower to a Metal
## kernel through a pass pipeline over a small nest IR.
##
##   xposeNaive[j, i] ++= a[i, j] | (i, j)
##
## A scheduled kernel takes a do block.
##
## Lowering is derivation-first. Every kernel emits:
## - a device function `proc <name>Fn[T](...)` with the views, threadgroup
##   scratch pointers and launch geometry as parameters
## - a thin entry template owning the geometry lets, the threadgroup-memory
##   declarations and the call
##
## Each function body is a tile program over TileAlgebra primitives
## (`gd` views, `loadTileBounded`, elementwise tile ops, `row_sum`, `mma_AB`, `storeTileMasked`).
##
## Index math, the atom configuration (`getTileConfig`) and the element
## dtype derive from the IR, the schedule and the tensorcery engine.
## No emitter literals enter the emitted code.
##
## Run with `engine.eval(mm(m1, m2))` from `proc main()`.
## The engine allocates the output `o`.
## The module exports nothing, `main` is the only consumer.
##
## Compile-time echo prints every derived function and entry, framed by separators.
##
## Section banners below carry the pass pipeline, the lowering rules and the einsum section 2 recipes (Rocktaschel 2018) with their checks.

import
  std/[macros, macrocache, strutils, tables, typetraits],
  workspace/data_structures/src/small_seqs,
  workspace/ceramic/src/int_tuples,
  workspace/ceramic/src/layouts,
  workspace/ceramic/src/layout_constructors,
  workspace/ceramic/src/layout_indexing,
  workspace/ceramic/src/tensors,
  workspace/ceramic/src/ptr_arithmetic,
  workspace/ceramic/src/layout_algebra,
  workspace/ceramic/src/tile_algebra,
  workspace/ceramic/src/tile_algebra/tile_io_bounded,
  workspace/experimental/tensorcery,
  workspace/crucible

# Element dtype every kernel derives from, the buffer element type
# tensorcery's marshalling (`operandArg`) hands the device.
type HostElem = typeof(operandArg(default(tensorcery.Tensor)).buf[][0])
static:
  doAssert HostElem is float32, "substrate elem type should be float32"

proc toMetadata(vals: openArray[int32]): Metadata =
  ## Output-shape builder for `eval` expansions, resolved at this module's scope.
  ## Spills past the inline capacity, rank-5 shapes included.
  result = Metadata.init()
  for v in vals:
    result.add v

proc evalNotImplemented(): tensorcery.Tensor =
  ## Backend-dispatch placeholder for unimplemented backends.
  raiseAssert "eval: backend not implemented"

func rows(t: tensorcery.Tensor): int32 =
  ## Row count of a host tensor.
  int32(t.shape[0])

func cols(t: tensorcery.Tensor): int32 =
  ## Column count of a host tensor.
  int32(t.shape[1])

func tdim(t: tensorcery.Tensor, k: int): int32 =
  ## Size of axis `k` of a host tensor at any rank, grid clauses resolve the name at this module's scope.
  int32(t.shape[k])

func cdiv(a, b: int32): int32 =
  ## Ceil division for launch-geometry expressions.
  (a + b - 1) div b

const
  KernelTrace = CacheTable"tc.trace"
    ## Per-kernel pass trace recorded at macro expansion, one line per
    ## executed pass. Compile-time only, read through `kernelTraceStr`.
  SchedInfo = CacheTable"tc.sched"
    ## Per-kernel provenance, schedule values and the tile types they chose.
    ## Compile-time only, read through `schedInfoStr`.
  DefaultTile = 16
  DefaultThreads = 32

macro kernelTraceStr(name: static string): string =
  ## Pass trace of the kernel `name` (see `KernelTrace`).
  newLit(KernelTrace[name].strVal)

macro schedInfoStr(name: static string): string =
  ## Schedule provenance of the kernel `name` (see `SchedInfo`).
  newLit(SchedInfo[name].strVal)

# ─────────────────────────────────────────────────────────────────────────
# The nest IR shared by every pass
# ─────────────────────────────────────────────────────────────────────────
type
  IdxKind = enum ikIter, ikConst
    ## One read index slot, an iterator name or a constant.

  IndexRef = object
    case kind: IdxKind
    of ikIter: name: string
    of ikConst: value: int32

  ReadIR = object
    tensor: string
    indices: seq[IndexRef]
      ## One index per read slot, slot k addressing tensor axis k, fewer slots add no offset (vector reads address axis 0).

  ExprKind = enum xeRead, xeAdd, xeMul, xeLit
  ExprIR = ref object
    ## Expression tree over reads, `+` and `*`, float literals.
    case kind: ExprKind
    of xeRead: read: ReadIR
    of xeAdd, xeMul: lhs, rhs: ExprIR
    of xeLit: value: float32

  LoopRole = enum lrOut, lrReduce
  Extent = object
    ## An iterator's extent, the tensor and axis providing the loop bound.
    tensor: string ## Tensor whose shape provides the extent.
    axis: int      ## Axis of that shape.

  LoopIR = object
    ## One loop of the logical nest, an out dim or a reduction iterator.
    name: string    ## Logical iterator name, tiling keeps it outer, the inner loop takes `name & "t"`.
    extent: Extent  ## For an outer loop, the logical extent it iterates.
    tile: int       ## 0 = untiled.
    innerOf: string ## "" for outer loops, the parent loop for a tile split.
    role: LoopRole

  CachePlan = object
    tensor: string
    buffer: seq[tuple[iter: string, tile: int]]
      ## Staged buffer dims, one per iterator appearing in the tensor's reads.

  AlgoIR = object
    name: string        ## Kernel name, the LHS name of the `++=` statement.
    target: string      ## Output tensor name, always `o` in the kernel.
    dtype: string       ## Element dtype of every view, tensorcery's host element type name.
    outDims: seq[string]
    reduceIters: seq[string] ## Iterators not indexing any output dim.
    iters: seq[string]  ## Logical iterator names, parse order.
    loops: seq[LoopIR]  ## Ordered outer -> inner loop nest.
    expr: ExprIR
    reads: seq[ReadIR]  ## Reads in expression order.
    cached: seq[string] ## Tensors marked by `cacheAnnotatePass`.
    cachePlans: seq[CachePlan] ## Built by `cachePlanPass`.
    threads: int
    parallel: seq[string] ## `parallel(...)` annotations, default mapping.

# ─────────────────────────────────────────────────────────────────────────
# The pass framework over the nest IR
# ─────────────────────────────────────────────────────────────────────────
# A schedule is a `seq[Pass]` composed in written order, mirroring
# crucible's pass system in miniature.
type
  PassKind = enum pkTransform, pkAnalysis, pkLowering
  Pass = ref object
    name: string
    description: string
    kind: PassKind
    run: proc(ir: var AlgoIR) {.closure.}
    emit: proc(ir: var AlgoIR, pass: Pass): NimNode {.closure.}
      ## Lowering passes only, records the chosen form in `pass.description`.

proc describe(ir: AlgoIR): string =
  ## One-line IR summary for the pass trace.
  var parts: seq[string] = @[]
  for lp in ir.loops:
    var s = lp.name & "["
    if lp.extent.tensor.len > 0:
      s &= "ax" & $lp.extent.axis & "(" & lp.extent.tensor & ")"
    if lp.innerOf.len > 0:
      s &= " inner-of " & lp.innerOf
    if lp.tile > 0:
      s &= " tile=" & $lp.tile
    s &= (if lp.role == lrOut: " out]" else: " reduce]")
    parts.add(s)
  var rs: seq[string] = @[]
  for rd in ir.reads:
    var idx: seq[string] = @[]
    for ix in rd.indices:
      idx.add(case ix.kind
        of ikIter: ix.name
        of ikConst: $ix.value)
    rs.add(rd.tensor & "[" & idx.join(", ") & "]")
  var cache: seq[string] = @[]
  for cp in ir.cachePlans:
    var dims: seq[string] = @[]
    for d in cp.buffer:
      dims.add($d.tile)
    cache.add(cp.tensor & " -> smem[" & dims.join(", ") & "]")
  result = "out " & ir.outDims.join(",") &
    " | loops: " & parts.join(" ") &
    " | reads: " & rs.join(", ") &
    " | cache: " & cache.join(", ") &
    " | threads: " & $ir.threads

proc runPipeline(ir: var AlgoIR, passes: seq[Pass]): NimNode =
  ## Runs the passes in order, records the IR after each pass in `KernelTrace[ir.name]`, returns the lowering pass's emitted node.
  var trace = ir.name & ":\n"
  for p in passes:
    when defined(schedTrace):
      echo "[sched-pass] ", p.name, " (", p.kind, "): ", p.description
    case p.kind
    of pkLowering:
      result = p.emit(ir, p)
      trace &= "  lower: " & p.description & "\n"
    else:
      p.run(ir)
      trace &= "  after " & p.name & ": " & describe(ir) & "\n"
  KernelTrace[ir.name] = newLit(trace)

# ─────────────────────────────────────────────────────────────────────────
# Schedule micro-passes over the schedule AST
# ─────────────────────────────────────────────────────────────────────────
proc tilePass(loop: string, size: int): Pass =
  ## Splits `loop` into an outer step-`size` loop and an inner extent-`size` loop, read indices stay logical (lowering rewrites them).
  Pass(name: "tile(" & loop & "," & $size & ")",
       description: "split loop '" & loop & "' into outer step " & $size &
                    " and inner extent " & $size,
       kind: pkTransform,
       run: proc(ir: var AlgoIR) =
         for i in 0 ..< ir.loops.len:
           if ir.loops[i].name == loop and ir.loops[i].innerOf.len == 0:
             let inner = LoopIR(name: loop & "t",
                                extent: ir.loops[i].extent,
                                tile: 0, innerOf: loop,
                                role: ir.loops[i].role)
             ir.loops[i].tile = size
             ir.loops.insert(inner, i + 1)
             return
         raiseAssert("tilePass: no outer loop named '" & loop & "'"))

proc cacheAnnotatePass(tensors: seq[string]): Pass =
  ## Records the tensors to stage, each must be read by the algorithm.
  Pass(name: "cache(" & tensors.join(",") & ")",
       description: "mark tensors for threadgroup-memory staging",
       kind: pkTransform,
       run: proc(ir: var AlgoIR) =
         for t in tensors:
           var found = false
           for rd in ir.reads:
             if rd.tensor == t: found = true
           doAssert found, "cache: tensor '" & t & "' is not read"
           if t notin ir.cached:
             ir.cached.add(t))

proc cachePlanPass(): Pass =
  ## Sizes each staged buffer from the tile sizes of the iterators the tensor's reads depend on, one plan per cached tensor.
  Pass(name: "cachePlan",
       description: "size the staged buffers from the schedule tiles",
       kind: pkAnalysis,
       run: proc(ir: var AlgoIR) =
         for t in ir.cached:
           var planned = false
           for cp in ir.cachePlans:
             if cp.tensor == t: planned = true
           if planned: continue
           var plan = CachePlan(tensor: t)
           for rd in ir.reads:
             if rd.tensor != t: continue
             for ix in rd.indices:
               if ix.kind != ikIter: continue
               var matched = false
               for lp in ir.loops:
                 if lp.name == ix.name and lp.innerOf.len == 0:
                   doAssert lp.tile > 0,
                     "cache: iterator '" & ix.name &
                     "' staging '" & t & "' must be tiled"
                   plan.buffer.add((ix.name, lp.tile))
                   matched = true
               doAssert matched,
                 "cache: iterator '" & ix.name & "' not found in the nest"
           ir.cachePlans.add(plan))

proc threadsPass(n: int): Pass =
  ## Sets the threadgroup width.
  Pass(name: "threads(" & $n & ")",
       description: "threadgroup width " & $n,
       kind: pkTransform,
       run: proc(ir: var AlgoIR) = ir.threads = n)

proc parallelPass(loop: string): Pass =
  ## Annotates a loop for parallel placement, the default grid mapping
  ## already parallelizes the out dims, so this pass records only.
  Pass(name: "parallel(" & loop & ")",
       description: "annotate loop '" & loop & "' for parallel placement",
       kind: pkAnalysis,
       run: proc(ir: var AlgoIR) = ir.parallel.add(loop))

# ─────────────────────────────────────────────────────────────────────────
# algorithm statement parsing
# ─────────────────────────────────────────────────────────────────────────
proc parseTarget(t: NimNode, ir: var AlgoIR) =
  ## Accepts `name[o0, o1]` and the historical `name*[o0, o1]` infix form.
  ## - the module exports nothing, so the target carries no `*` marker
  ## - an empty bracket `name[]` is the all-reduce (scalar) target
  ## - dim counts 0 to 5 work
  var nameNode: NimNode
  var dims: NimNode
  if t.kind == nnkBracketExpr:
    nameNode = t[0]
    dims = t
  elif t.kind == nnkInfix and t[0].strVal == "*" and t[2].kind == nnkBracket:
    nameNode = t[1]
    dims = t[2]
  else:
    error("target must be written name[o0, o1]", t)
  ir.name = nameNode.strVal
  ir.target = nameNode.strVal
  for i in 1 ..< dims.len:
    ir.outDims.add(dims[i].strVal)
  doAssert ir.outDims.len <= 5,
    "at most five output dimensions"

proc parseIndex(n: NimNode): IndexRef =
  ## One read index, an iterator name or an integer constant (a parenthesized index arrives wrapped in a Par node).
  var n = n
  if n.kind == nnkPar and n.len == 1:
    n = n[0]
  case n.kind
  of nnkIdent, nnkSym: IndexRef(kind: ikIter, name: n.strVal)
  of nnkIntLit: IndexRef(kind: ikConst, value: int32(n.intVal))
  else:
    error("index must be an iterator name or an integer literal", n)
    IndexRef(kind: ikConst)

proc parseRead(n: NimNode, ir: var AlgoIR): ExprIR =
  ## `t[i0, i1]` parses as BracketExpr(t, i0, i1), possibly with the index
  ## pair wrapped in a Par node, vectors read as `t[i]` (one slot).
  doAssert n.kind == nnkBracketExpr,
    "reads are written tensor[i0, i1, ...]"
  var rd = ReadIR(tensor: n[0].strVal)
  if n.len == 2 and n[1].kind == nnkPar:
    for i in 0 ..< n[1].len:
      rd.indices.add(parseIndex(n[1][i]))
  else:
    for i in 1 ..< n.len:
      rd.indices.add(parseIndex(n[i]))
  doAssert rd.indices.len in 1 .. 5,
    "reads carry 1 to 5 index slots"
  ir.reads.add(rd)
  ExprIR(kind: xeRead, read: rd)

proc parseExprNode(n: NimNode, ir: var AlgoIR): ExprIR =
  ## Expression tree over reads, `+`, `*`, tensor reads, float literals.
  case n.kind
  of nnkInfix:
    let op = n[0].strVal
    doAssert op in ["+", "*"],
      "only + and * are supported in the reduction expression"
    var e = ExprIR(kind: (if op == "+": xeAdd else: xeMul))
    e.lhs = parseExprNode(n[1], ir)
    e.rhs = parseExprNode(n[2], ir)
    e
  of nnkBracketExpr: parseRead(n, ir)
  of nnkPar:
    doAssert n.len == 1, "unsupported parenthesized expression " & $n
    parseExprNode(n[0], ir)
  of nnkIntLit, nnkFloatLit:
    ExprIR(kind: xeLit, value: float32(n.floatVal))
  else:
    error("unsupported expression node " & $n.kind, n)
    ExprIR(kind: xeLit)

proc parseValue(v: NimNode, ir: var AlgoIR) =
  ## Parses `expr | (i1, i2, ...)`, the `|` separates the expression from an iterator binding tuple of any length.
  doAssert v.kind == nnkInfix and v[0].strVal == "|" and v.len == 3,
    "value must be written expr | (i1, i2, ...)"
  ir.expr = parseExprNode(v[1], ir)
  let iters = v[2]
  doAssert iters.kind in {nnkPar, nnkBracket, nnkTupleConstr},
    "the iterator binding must be a tuple"
  for it in iters:
    ir.iters.add(it.strVal)
  for d in ir.outDims:
    doAssert d in ir.iters,
      "target dim '" & d & "' is not bound in the iterator tuple"

proc iterExtent(ir: AlgoIR, iter: string): Extent =
  ## An iterator's extent is the (tensor, axis) of the first read slot using it, later reads are assumed shape-aligned.
  for rd in ir.reads:
    for k, ix in rd.indices:
      if ix.kind == ikIter and ix.name == iter:
        return Extent(tensor: rd.tensor, axis: k)
  raiseAssert("no read binds iterator '" & iter & "'; its extent is unknown")

proc buildLoops(ir: var AlgoIR) =
  ## Ordered logical loop nest, out dims first, then the reduction
  ## iterators in the order the statement binds them.
  for it in ir.iters:
    if it notin ir.outDims and it notin ir.reduceIters:
      ir.reduceIters.add(it)
  var names = ir.outDims
  names.add(ir.reduceIters)
  for nm in names:
    ir.loops.add(LoopIR(name: nm, extent: ir.iterExtent(nm),
                        role: (if nm in ir.reduceIters: lrReduce else: lrOut)))

# ─────────────────────────────────────────────────────────────────────────
# Lowering matcher over IR facts
# ─────────────────────────────────────────────────────────────────────────
# Form selection from IR facts only. Every form emits a tile program:
# - blocked copy, no reduction and one read with two out dims, its two slots a permutation of the out dims, both dims tiled.
# - mma, product of two reads, slot classes (out0, reduce) and (reduce, out1), one reduction iterator, tiles all in [16, 32],
#   staged threadgroup tiles when every read is cache-planned, register otherwise.
# - everything else lowers to the derived elementwise, reduction or all-reduce tile program.
type
  LowerForm = enum lfScalar, lfBlockedCopy, lfMmaStaged, lfMmaRegister

proc outerTile(ir: AlgoIR, iter: string): int =
  for lp in ir.loops:
    if lp.name == iter and lp.innerOf.len == 0:
      return lp.tile
  0

proc iterClass(ir: AlgoIR, iter: string): string =
  ## Classifies an iterator, "o0", "o1", ... by output-dim position, "r" for a reduction iterator, empty when unknown.
  for d, dim in ir.outDims:
    if dim == iter: return "o" & $d
  if iter in ir.reduceIters: return "r"
  ""

proc readSlotClass(ir: AlgoIR, rd: ReadIR, slot: int): string =
  ## Iterator class of one read slot, a slot past the read's rank is "".
  if slot >= rd.indices.len:
    return ""
  (case rd.indices[slot].kind
    of ikIter: ir.iterClass(rd.indices[slot].name)
    of ikConst: "")

proc chooseForm(ir: AlgoIR): tuple[form: LowerForm, tiles: tuple[m, n, k: int]] =
  result.form = lfScalar
  if ir.reduceIters.len == 0 and ir.expr.kind == xeRead and
      ir.outDims.len == 2:
    let rd = ir.reads[0]
    let c0 = ir.readSlotClass(rd, 0)
    let c1 = ir.readSlotClass(rd, 1)
    if c0 != c1 and c0 in ["o0", "o1"] and c1 in ["o0", "o1"]:
      let t0 = ir.outerTile(ir.outDims[0])
      let t1 = ir.outerTile(ir.outDims[1])
      if t0 > 0 and t1 > 0:
        result.form = lfBlockedCopy
        result.tiles = (t0, t1, 0)
        return
  if ir.outDims.len == 2 and ir.reduceIters.len == 1 and
      ir.expr.kind == xeMul and
      ir.expr.lhs.kind == xeRead and ir.expr.rhs.kind == xeRead:
    let (ra, rb) = (block:
      var a = -1
      var b = -1
      for i, rd in ir.reads:
        if ir.readSlotClass(rd, 0) == "o0" and ir.readSlotClass(rd, 1) == "r":
          a = i
        elif ir.readSlotClass(rd, 0) == "r" and ir.readSlotClass(rd, 1) == "o1":
          b = i
      (a, b))
    if ra >= 0 and rb >= 0 and ra != rb:
      let tiles = (ir.outerTile(ir.outDims[0]), ir.outerTile(ir.outDims[1]),
                   ir.outerTile(ir.reduceIters[0]))
      if tiles[0] in [16, 32] and tiles[1] in [16, 32] and
          tiles[2] in [16, 32]:
        var allCached = true
        for rd in ir.reads:
          if rd.tensor notin ir.cached: allCached = false
        result.form = (if allCached: lfMmaStaged else: lfMmaRegister)
        result.tiles = (tiles[0], tiles[1], tiles[2])

# ─────────────────────────────────────────────────────────────────────────
# Node construction helpers for the lowering
# ─────────────────────────────────────────────────────────────────────────
# Plain idents, not quote-do gensyms, for any name a sibling fragment or wrapper references by name (`o`, `a`, `bx`, `acc`, `sm_a`, ...).
# A quote-do `let` gensyms its name and leaves such references unresolved at template instantiation, intra-fragment temporaries may gensym.
proc id(s: string): NimNode = newIdentNode(s)
proc mkGdView[T](buf: ptr UncheckedArray[T], sb, sd, sr, sc: int32): GlView[T] =
  ## Returns a view over a global buffer of dynamic (-1) extents
  ## with the given strides, tile_algebra's `gd` in proc form.
  ##
  ## Plain calls resolve inside macro-emitted procedure bodies on this Nim.
  ## Method-call syntax does not, which is the proc form's reason.
  block:
    let sh = (-1, -1, -1, -1)
    let st = (int(sb), int(sd), int(sr), int(sc))
    GlView[T](data: buf,
              layout: make_layout((int(sh[0]), int(sh[1]), int(sh[2]),
                                   int(sh[3])),
                                  (st[0], st[1], st[2], st[3])))
template offsetBuf*[T](p: ptr UncheckedArray[T], off: int32): ptr UncheckedArray[T] =
  ## Advance a buffer pointer by `off` elements, ptr_arithmetic's `+%`
  ## operator under a callable name.
  ##
  ## Imported operator templates do not resolve inside macro-emitted
  ## procedure bodies, named templates do.
  ##
  ## Returns the offset pointer.
  cast[ptr UncheckedArray[T]](cast[uint64](p) + uint64(off) * uint64(sizeof(T)))
proc callI32(n: NimNode): NimNode = nnkCall.newTree(id"int32", n)
proc i32lit(v: int): NimNode = newLit(v)
proc infixF(op: string, a, b: NimNode): NimNode =
  nnkInfix.newTree(id(op), a, b)
proc dot(a: NimNode, field: string): NimNode =
  nnkDotExpr.newTree(a, id(field))

proc layoutField(t: string, f: string, k: int): NimNode =
  ## Builds `t.layout.<f>[k]`, shape or stride access on a device view, the emitter knows the axis at compile time.
  nnkBracketExpr.newTree(
    nnkDotExpr.newTree(nnkDotExpr.newTree(id(t), ident"layout"), ident(f)),
    newLit(k))

proc extentExpr(e: Extent, device: bool): NimNode =
  ## Loop extent, `tdim(t, k)` in host clause position, `int32(t.layout.shape[k])` in kernel-body position.
  if device:
    callI32(layoutField(e.tensor, "shape", e.axis))
  else:
    nnkCall.newTree(id"tdim", id(e.tensor), newLit(e.axis))

# ─────────────────────────────────────────────────────────────────────────
# Tile-program derivation from the IR + schedule
# ─────────────────────────────────────────────────────────────────────────
# Tile plan, mapping the out dims onto a register-tile plane plus
# batch/depth view slots and, for reductions, chunked or serial reduce loops.
#
# Per read, each axis is either VARYING or BROADCAST.
#
# - VARYING maps a slot to the axis dim, the view stride being the slot's stride,
#   the tile origin carrying the grid or chunk coordinate and the load
#   limit equal to the axis extent
# - BROADCAST has stride 0, origin 0 and the load limit equal to the tile extent, since loadTileBounded's limits are absolute over the view
type
  TilePlan = object
    rowDim, colDim: string ## Out dims mapped to the tile plane, for an all-reduce the row/col chunk iters.
    tileR, tileC: int      ## Tile plane extents, schedule tiles or DefaultTile.
    prefixDims: seq[string]## Out dims mapped to the gd view's batch/depth slots (0..2).
    chunkIter: string      ## Reduce iter chunked along the tile cols, "" without one.
    serialIters: seq[string] ## Reduce iters stepping per tile (folded into the gather base).
    allReduce: bool        ## Scalar target, cross-lane combine instead of a tile store.
    flatRect: bool         ## All-reduce with one-slot reads, reduce space unranked over the atom rectangle.

proc slotBinds(rd: ReadIR, name: string): bool =
  for ix in rd.indices:
    if ix.kind == ikIter and ix.name == name:
      return true
  false

proc readKey(rd: ReadIR): string =
  ## Identity of a read, tensor name plus index list, the gathered-tile
  ## table key when one tensor is read several times.
  result = rd.tensor & "("
  for i, ix in rd.indices:
    if i > 0: result &= ","
    result &= (case ix.kind
      of ikIter: ix.name
      of ikConst: $ix.value)
  result &= ")"

proc makeTilePlan(ir: AlgoIR): TilePlan =
  let L = ir.outDims.len
  if L == 0:
    result.allReduce = true
    result.rowDim = ir.reduceIters[0]
    result.tileR = ir.outerTile(ir.reduceIters[0])
    if result.tileR == 0: result.tileR = DefaultTile
    if ir.reduceIters.len >= 2:
      # rows chunk the first reduce iter, cols the second, the rest serial
      result.colDim = ir.reduceIters[1]
      let tc = ir.outerTile(result.colDim)
      result.tileC = (if tc == 0: DefaultTile else: tc)
      result.serialIters = ir.reduceIters[2 ..^ 1]
    else:
      var singleSlot = true
      for rd in ir.reads:
        var iterSlots = 0
        for ix in rd.indices:
          if ix.kind == ikIter: inc iterSlots
        if iterSlots != 1: singleSlot = false
      if singleSlot:
        result.flatRect = true
      else:
        raiseAssert(
          "reported gap: scalar all-reduce over a multi-slot read with a " &
          "single reduce iterator needs a ragged flat unrank that " &
          "loadTileBounded's per-axis limits cannot express; reproduce with " &
          "`s[] ++= a[i, 0] | (i)`")
  else:
    result.rowDim = (if L >= 2: ir.outDims[^2] else: ir.outDims[0])
    result.colDim = (if L >= 2: ir.outDims[^1] else: "")
    result.tileR = ir.outerTile(result.rowDim)
    if result.tileR == 0: result.tileR = DefaultTile
    if result.colDim.len > 0:
      let t = ir.outerTile(result.colDim)
      result.tileC = (if t == 0: DefaultTile else: t)
    else:
      result.tileC = DefaultTile
    result.prefixDims = (if L >= 3: ir.outDims[0 ..< L - 2] else: @[])
    if result.prefixDims.len > 2:
      raiseAssert(
        "reported gap: output rank above 4 has no grid axis for the " &
        "prefix dims, the gd view's batch and depth slots cover two. " &
        "Reproduce with `x[a, b, c, i, j] ++= w[a, b, c, i, j] | (a, b, c, i, j)`")
    if result.colDim.len == 0:
      # no col dim, chunk the reduce space along the tile cols, valid only
      # when no read references a col dim (elementwise tiles must align)
      for it in countdown(ir.reduceIters.len - 1, 0):
        var bound = false
        for rd in ir.reads:
          if rd.slotBinds(ir.reduceIters[it]): bound = true
        if bound:
          result.chunkIter = ir.reduceIters[it]
          break
      let t = ir.outerTile(result.chunkIter)
      if result.chunkIter.len > 0 and t > 0: result.tileC = t
    for it in ir.reduceIters:
      if it != result.chunkIter:
        result.serialIters.add(it)

proc slotStride(rd: ReadIR, name: string): NimNode =
  ## int32 view-stride sum of the slots bound to iterator `name`, 0 when unbound (a broadcast axis).
  var acc: NimNode = nil
  for k, ix in rd.indices:
    if ix.kind == ikIter and ix.name == name:
      let term = callI32(layoutField(rd.tensor, "stride", k))
      acc = (if acc == nil: term else: infixF("+", acc, term))
  (if acc == nil: i32lit(0) else: acc)

proc readBase(rd: ReadIR, resolve: Table[string, NimNode]): NimNode =
  ## Returns the int32 offset folded into the gather's buffer pointer.
  ## Constant slots and resolved serial reduce iters contribute.
  ## Out dims and the chunk iter are covered by the gd strides instead.
  var acc: NimNode = nil
  for k, ix in rd.indices:
    case ix.kind
    of ikConst:
      let term = infixF("*", i32lit(ix.value.int),
                        callI32(layoutField(rd.tensor, "stride", k)))
      acc = (if acc == nil: term else: infixF("+", acc, term))
    of ikIter:
      if ix.name in resolve:
        let term = infixF("*", resolve[ix.name],
                          callI32(layoutField(rd.tensor, "stride", k)))
        acc = (if acc == nil: term else: infixF("+", acc, term))
  (if acc == nil: i32lit(0) else: acc)

proc gdViewNode(viewName, tensorName: string, bStr, dStr, rStr, cStr, base: NimNode): NimNode =
  ## Emits `let viewName = mkGdView(offsetBuf(t.data, base), s...)`, the read's gd view.
  ## A zero base skips the pointer arithmetic.
  let buf = (if base.kind == nnkIntLit and base.intVal == 0:
               dot(id(tensorName), "data")
             else:
               nnkCall.newTree(id"offsetBuf", dot(id(tensorName), "data"), base))
  nnkLetSection.newTree(nnkIdentDefs.newTree(id(viewName), newEmptyNode(),
    nnkCall.newTree(id"mkGdView", buf, bStr, dStr, rStr, cStr)))

proc emitGather(ir: AlgoIR, plan: TilePlan, rdIdx: int, resolve: Table[string, NimNode]): tuple[gd: string, tile: string, stmts: NimNode] =
  ## One read's per-threadgroup register tile, a gd view over the read's
  ## tensor with derived strides and `loadTileBounded` at the derived origin.
  ##
  ## Origins and limits reference the fn-local `by`, `bx`, `rchunk`,
  ## `cchunk` and the loop variables in `resolve`.
  let rd = ir.reads[rdIdx]
  let t = rd.tensor
  let gdName = "gd" & $rdIdx
  let tileName = "gt" & $rdIdx

  var bStr, dStr: NimNode = i32lit(0)
  if plan.prefixDims.len >= 1: bStr = rd.slotStride(plan.prefixDims[0])
  if plan.prefixDims.len >= 2: dStr = rd.slotStride(plan.prefixDims[1])
  var rStr, cStr: NimNode
  var rOrig, cOrig: NimNode
  var limR, limC: NimNode
  if plan.flatRect:
    # the reduce space unranked over the atom rectangle, cell (r, c) is
    # rchunk*atomM*atomN + r*atomN + c, page-tail zeros past the extent
    rStr = id"atomN"
    cStr = i32lit(1)
    rOrig = id"rchunk"
    cOrig = i32lit(0)
    limR = id"flatRows"
    limC = id"atomN"
  else:
    let rVaries = rd.slotBinds(plan.rowDim)
    rStr = (if rVaries: rd.slotStride(plan.rowDim) else: i32lit(0))
    rOrig = (if rVaries: (if plan.allReduce: id"rchunk" else: id"by")
             else: i32lit(0))
    var cVaries: bool
    if plan.colDim.len > 0:
      cVaries = rd.slotBinds(plan.colDim)
      cStr = (if cVaries: rd.slotStride(plan.colDim) else: i32lit(0))
      cOrig = (if cVaries: (if plan.allReduce: id"cchunk" else: id"bx")
               else: i32lit(0))
    elif plan.chunkIter.len > 0:
      cVaries = rd.slotBinds(plan.chunkIter)
      cStr = (if cVaries: rd.slotStride(plan.chunkIter) else: i32lit(0))
      cOrig = (if cVaries: id"chunk" else: i32lit(0))
    else:
      cVaries = false
      cStr = i32lit(0)
      cOrig = i32lit(0)
    # Varying axes take the real absolute extent as the load limit.
    # Broadcast axes take the tile extent, loadTileBounded's limits being
    # absolute over the view.
    limR = (if rVaries:
              callI32(extentExpr(ir.iterExtent(plan.rowDim), device = true))
            else:
              id"tileRc")
    if plan.colDim.len > 0 and cVaries:
      limC = callI32(extentExpr(ir.iterExtent(plan.colDim), device = true))
    elif plan.chunkIter.len > 0 and cVaries:
      limC = callI32(extentExpr(ir.iterExtent(plan.chunkIter), device = true))
    else:
      limC = id"tileCc"

  var stmts = nnkStmtList.newTree()
  stmts.add(gdViewNode(gdName, t, bStr, dStr, rStr, cStr,
                       rd.readBase(resolve)))
  stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tileName),
    nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
    newEmptyNode())))
  stmts.add(nnkCall.newTree(nnkDotExpr.newTree(id(tileName),
    ident"loadTileBounded"), id(gdName),
    nnkPar.newTree(
      (if plan.prefixDims.len >= 1: id"bCoord" else: i32lit(0)),
      (if plan.prefixDims.len >= 2: id"dCoord" else: i32lit(0)),
      rOrig, cOrig),
    limR, limC))
  (gd: gdName, tile: tileName, stmts: stmts)

proc applyScalar(tile: string, scalar: NimNode, stmts: NimNode, counter: var int): string =
  ## Multiplies a tile by a scalar through `mul(dst, src, s)`.
  ## Returns the input tile unchanged without a scalar.
  if scalar == nil:
    return tile
  let tmp = "st" & $counter
  inc counter
  stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tmp),
    nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
    newEmptyNode())))
  stmts.add(nnkCall.newTree(dot(id(tmp), "mul"), id(tile), scalar))
  tmp

proc emitTileValue(e: ExprIR, tileOf: Table[string, string], stmts: NimNode, counter: var int): tuple[tile: string, scalar: NimNode] =
  ## Combines gathered tiles into one rt_l tile by walking the IR expression
  ## tree:
  ## - tile mul through `mul(dst, src, src2)`
  ## - tile add through `convert` plus `addScaled(dst, src, 1.0)`
  ## - float literals as scalar multipliers carried on the side
  ##
  ## A literal-only tree is a reported gap, since no scalar-broadcast
  ## fill primitive exists for a tile.
  case e.kind
  of xeLit:
    (tile: "", scalar: newLit(e.value))
  of xeRead:
    (tile: tileOf[readKey(e.read)], scalar: nil)
  of xeMul:
    let (lt, lsc) = emitTileValue(e.lhs, tileOf, stmts, counter)
    let (rt, rsc) = emitTileValue(e.rhs, tileOf, stmts, counter)
    var scalar: NimNode
    if lsc != nil and rsc != nil:
      scalar = infixF("*", lsc, rsc)
    elif lsc != nil: scalar = lsc
    else: scalar = rsc
    if lt.len == 0 and rt.len == 0:
      raiseAssert(
        "reported gap: a literal without a read has no tile to combine " &
        "into (TileAlgebra has no scalar-broadcast fill); reproduce with " &
        "`o[i, j] ++= 2.0 | (i, j)`")
    if lt.len == 0:
      (tile: rt, scalar: scalar)
    elif rt.len == 0:
      (tile: lt, scalar: scalar)
    else:
      let tmp = "mt" & $counter
      inc counter
      stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tmp),
        nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
        newEmptyNode())))
      stmts.add(nnkCall.newTree(dot(id(tmp), "mul"), id(lt), id(rt)))
      (tile: tmp, scalar: scalar)
  of xeAdd:
    let (lt, lsc) = emitTileValue(e.lhs, tileOf, stmts, counter)
    let (rt, rsc) = emitTileValue(e.rhs, tileOf, stmts, counter)
    if lt.len == 0 or rt.len == 0:
      raiseAssert(
        "reported gap: a literal summand has no broadcast fill; " &
        "reproduce with `o[i, j] ++= a[i, j] + 1.0 | (i, j)`")
    doAssert lsc == nil and rsc == nil,
      "reported gap: a scaled summand inside + has no fused form; " &
      "reproduce with `o[i, j] ++= 2.0 * a[i, j] + b[i, j] | (i, j)`"
    let tmp = "at" & $counter
    inc counter
    stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tmp),
      nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
      newEmptyNode())))
    stmts.add(nnkCall.newTree(dot(id(tmp), "convert"), id(lt)))
    stmts.add(nnkCall.newTree(dot(id(tmp), "addScaled"), id(rt), newLit(1.0'f32)))
    (tile: tmp, scalar: nil)

# ─────────────────────────────────────────────────────────────────────────
# threadgroup-memory staging, gemm expansion, blocked tile copy
# ─────────────────────────────────────────────────────────────────────────
proc smemDecl(name: string, rows, cols: int): NimNode =
  ## Builds a `var name: array[rows * cols, HostElem] {.smem.}` declaration,
  ## hoisted into the entry via the launch spec and passed to the device fn
  ## as a threadgroup-pointer parameter (the Phase A crucible capability).
  nnkVarSection.newTree(nnkIdentDefs.newTree(
    nnkPragmaExpr.newTree(id(name), nnkPragma.newTree(id"smem")),
    nnkBracketExpr.newTree(id"array", newLit(rows * cols), id"HostElem"),
    newEmptyNode()))

proc emitCopyFromGlobal(outBody: var NimNode, dst: string, dims: tuple[r, c: int], region: NimNode) =
  ## Strided flat staging loop, the threadgroup copying a (r, c) row-major
  ## tile of the global view into the threadgroup-memory buffer.
  ## - the tile origin carries the region index expressions
  ## - the destination is the device fn's threadgroup-pointer parameter
  let base = region[0]
  let rE = region[1]
  let cE = region[2]
  let i = genSym(nskVar, "stageI")
  let dstN = id(dst)
  let rL = newLit(dims.r)
  let cL = newLit(dims.c)
  outBody.add quote do:
    var `i` = int32(thread_index_in_threadgroup)
    while `i` < int32(`rL` * `cL`):
      `dstN`[`i`] = `base`.data[(`rE` + `i` div int32(`cL`)) *
                                `base`.rowStride + `cE` + `i` mod int32(`cL`)]
      `i` += int32(threads_per_threadgroup.x)

proc emitGemm(outBody: var NimNode, asName, bsName, cfName: string, asDims, bsDims, fragDims: tuple[r, c: int]) =
  ## Staged mma accumulate, per-atom-slab threadgroup-to-register loads plus
  ## `mma_AB` into the accumulator.
  ## - `asDims`/`bsDims` name the staged buffers' (rows, cols)
  ## - `fragDims` names the accumulator tile's (rows, cols)
  ##
  ## - the atom record and every loop bound are consts in the emitted function, derived from `getTileConfig` at codegen time
  ## - the emitter writes no atom value
  let bm = fragDims.r
  let bn = fragDims.c
  let bk = asDims.c
  let asN = id(asName)
  let bsN = id(bsName)
  let cfN = id(cfName)
  let bmL = newLit(bm)
  let bnL = newLit(bn)
  let bkL = newLit(bk)
  outBody.add quote do:
    const atom = getTileConfig(HostElem, HostElem)
    const atomM = atom.getM()
    const atomN = atom.getN()
    const atomK = atom.getK()
    const atomVpt = atom.getVpt()
    const kSlabs = int32(`bkL`) div atomK
    let gemmLane = int32(thread_index_in_threadgroup)
    let gemmCell = crd2idx(atom.getLayoutA(), (gemmLane, 0)).toIntVal()
    let gemmLr = gemmCell mod atomM
    let gemmLc = gemmCell div atomM
    for gemmKk in 0'i32 ..< kSlabs:
      var ta: rt_l(HostElem, `bmL`, atomK)
      var tb: rt_r(HostElem, atomK, `bnL`)
      for n in 0'i32 ..< int32(`bmL` div atomM):
        for m in 0'i32 ..< int32(atomK div atomN):
          for v in 0'i32 ..< int32(atomVpt):
            ta.frags[n][m].frag[v] =
              `asN`[(gemmLr + n * int32(atomM)) * int32(`bkL`) +
                    gemmKk * int32(atomK) + gemmLc + m * int32(atomN) + v]
      for m in 0'i32 ..< int32(`bnL` div atomN):
        for n in 0'i32 ..< int32(atomK div atomM):
          for v in 0'i32 ..< int32(atomVpt):
            tb.frags[m][n].frag[v] =
              `bsN`[(gemmKk * int32(atomK) + gemmLr + n * int32(atomM)) *
                    int32(`bnL`) + gemmLc + m * int32(atomN) + v]
      `cfN`.mma_AB(ta, tb)

proc emitCopyToGlobal(outBody: var NimNode, region: NimNode, frag: string, fragDims: tuple[r, c: int]) =
  ## `storeTile` of the accumulator tile into the global output, tile index
  ## derived from the region indices (multiples of the frag dims).
  let base = region[0]
  let rE = region[1]
  let cE = region[2]
  let gd = genSym(nskLet, "gdStoreView")
  let fragN = id(frag)
  let rT = genSym(nskLet, "tileR")
  let cT = genSym(nskLet, "tileC")
  let frL = newLit(fragDims.r)
  let fcL = newLit(fragDims.c)
  outBody.add quote do:
    let `gd` = mkGdView(`base`.data, `base`.rowStride, 0'i32,
      `base`.rowStride, `base`.colStride)
    let `rT` = `rE` div int32(`frL`)
    let `cT` = `cE` div int32(`fcL`)
    storeTile(`gd`, `fragN`, (0, 0, `rT`, `cT`))

proc emitMmaStaged(ir: AlgoIR, tiles: tuple[m, n, k: int]): tuple[body, smemDecls: NimNode, scratch: seq[string]] =
  ## Staged mma kernel body, one threadgroup tile per cached operand staged
  ## from global memory per reduce chunk with `syncthreads()` barriers.
  ## - the staged buffers live in the entry and arrive as threadgroup pointer parameters (threadgroup memory cannot be declared in a device fn)
  ## - tile origins come from the out-dim grid coordinates and the reduce chunk
  let (bm, bn, bk) = tiles
  var ra, rb: int = -1
  for i, rd in ir.reads:
    if ir.readSlotClass(rd, 0) == "o0" and ir.readSlotClass(rd, 1) == "r":
      ra = i
    elif ir.readSlotClass(rd, 0) == "r" and ir.readSlotClass(rd, 1) == "o1":
      rb = i
  let tenA = ir.reads[ra].tensor
  let tenB = ir.reads[rb].tensor
  let bufA = "sm_" & tenA
  let bufB = "sm_" & tenB
  let frag = "acc"
  let reduceVar = "it"
  let yOrig = infixF("*", id"by", i32lit(bm))
  let xOrig = infixF("*", id"bx", i32lit(bn))

  var smem = nnkStmtList.newTree(smemDecl(bufA, bm, bk),
                                 smemDecl(bufB, bk, bn))
  var stmts = nnkStmtList.newTree()
  stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(
    id(frag), nnkCall.newTree(id"rt_l", id"T", i32lit(bm), i32lit(bn)),
    newEmptyNode())))
  stmts.add(nnkCall.newTree(dot(id(frag), "zero")))
  var chunkBody = nnkStmtList.newTree()
  emitCopyFromGlobal(chunkBody, bufA, (bm, bk),
    nnkBracketExpr.newTree(id(tenA), yOrig,
      infixF("*", id(reduceVar), i32lit(bk))))
  emitCopyFromGlobal(chunkBody, bufB, (bk, bn),
    nnkBracketExpr.newTree(id(tenB),
      infixF("*", id(reduceVar), i32lit(bk)), xOrig))
  chunkBody.add(nnkCall.newTree(id"syncthreads"))
  emitGemm(chunkBody, bufA, bufB, frag, (bm, bk), (bk, bn), (bm, bn))
  chunkBody.add(nnkCall.newTree(id"syncthreads"))
  stmts.add(nnkForStmt.newTree(id(reduceVar),
    nnkInfix.newTree(id"..<", newLit(0'i32),
      callI32(infixF("div",
        extentExpr(ir.iterExtent(ir.reduceIters[0]), device = true),
        i32lit(bk)))),
    chunkBody))
  emitCopyToGlobal(stmts, nnkBracketExpr.newTree(id"o", yOrig, xOrig),
                   frag, (bm, bn))
  (body: stmts, smemDecls: smem, scratch: @[bufA, bufB])

proc gdView(gdName, tensor, stride2, stride3: string): NimNode =
  ## `let gdX = t.data.gd(...)` register-tile view, dims 2 and 3 of the view
  ## being the tile's row and column directions.
  nnkLetSection.newTree(nnkIdentDefs.newTree(id(gdName), newEmptyNode(),
    nnkCall.newTree(id"mkGdView", dot(id(tensor), "data"),
      dot(id(tensor), "rowStride"), newLit(0),
      nnkDotExpr.newTree(id(tensor), id(stride2)),
      nnkDotExpr.newTree(id(tensor), id(stride3)))))

proc emitMmaRegister(ir: AlgoIR, tiles: tuple[m, n, k: int]): NimNode =
  ## Register-only mma kernel body, gd views over each operand, full-chunk
  ## register tiles per chunk. The right operand's gd axes swap into a transposed `rt_r` tile.
  ##
  ## Loads and stores are unguarded, the mma form's schedule assuming shape
  ## divisibility by the tiles. This is a reported limitation, not guarded.
  
  let (bm, bn, bk) = tiles
  var ra, rb: int = -1
  for i, rd in ir.reads:
    if ir.readSlotClass(rd, 0) == "o0" and ir.readSlotClass(rd, 1) == "r":
      ra = i
    elif ir.readSlotClass(rd, 0) == "r" and ir.readSlotClass(rd, 1) == "o1":
      rb = i
  let tenA = ir.reads[ra].tensor
  let tenB = ir.reads[rb].tensor
  let gdA = "gd_" & tenA
  let gdB = "gd_" & tenB
  let ta = "ta_" & tenA
  let tb = "tb_" & tenB
  let acc = "acc"
  let reduceVar = "it"

  var stmts = nnkStmtList.newTree()
  stmts.add(gdView(gdA, tenA, "rowStride", "colStride"))
  stmts.add(gdView(gdB, tenB, "colStride", "rowStride"))
  stmts.add(gdView("gd_o", "o", "rowStride", "colStride"))
  stmts.add(nnkVarSection.newTree(
    nnkIdentDefs.newTree(id(acc),
      nnkCall.newTree(id"rt_l", id"T", i32lit(bm), i32lit(bn),
        nnkCall.newTree(dot(id"T", "getTileConfig"), id"T")),
      newEmptyNode()),
    nnkIdentDefs.newTree(id(ta),
      nnkCall.newTree(id"rt_l", id"T", i32lit(bm), i32lit(bk)),
      newEmptyNode()),
    nnkIdentDefs.newTree(id(tb),
      nnkCall.newTree(id"rt_r", id"T", i32lit(bk), i32lit(bn)),
      newEmptyNode())))
  stmts.add(nnkCall.newTree(dot(id(acc), "zero")))
  var chunkBody = nnkStmtList.newTree()
  chunkBody.add(nnkCall.newTree(dot(id(ta), "loadTile"), id(gdA),
    nnkPar.newTree(newLit(0), newLit(0), id"by", id(reduceVar))))
  chunkBody.add(nnkCall.newTree(dot(id(tb), "loadTile"), id(gdB),
    nnkPar.newTree(newLit(0), newLit(0), id"bx", id(reduceVar))))
  chunkBody.add(nnkCall.newTree(dot(id(acc), "mma_AB"), id(ta), id(tb)))
  stmts.add(nnkForStmt.newTree(id(reduceVar),
    nnkInfix.newTree(id"..<", newLit(0'i32),
      callI32(infixF("div",
        extentExpr(ir.iterExtent(ir.reduceIters[0]), device = true),
        i32lit(bk)))),
    chunkBody))
  stmts.add(nnkCall.newTree(dot(id"gd_o", "storeTile"), id(acc),
    nnkPar.newTree(newLit(0), newLit(0), id"by", id"bx")))
  result = stmts

# ─────────────────────────────────────────────────────────────────────────
# kernel body, the derived tile program
# ─────────────────────────────────────────────────────────────────────────
proc emitTileKernel(ir: AlgoIR, plan: TilePlan):
    tuple[body: NimNode, smemDecls: NimNode, scratch: seq[string]] =
  ## Derived tile program for the scalar family, shared by the blocked copy (an elementwise permutation gather at the schedule tiles).
  ##
  ## Fn-local contract for the emitted fragment, group 1:
  ## - `tileRc`/`tileCc` as tile plane extents, atom dims for the flat rectangle, asserted atom-multiples at codegen time
  ## - `by`/`bx`/`bz` as grid coordinates
  ## - `rchunk`/`cchunk`/`chunk`/serial iterators as loop variables
  ##
  ## Fn-local contract for the emitted fragment, group 2:
  ## - `bCoord`/`dCoord` as the bz split over the prefix dims
  ## - `rowExtent`/`colExtent` as real plane extents and `flatRows` as the flat rectangle's row count
  ##
  ## Fn-local contract for the emitted fragment, group 3:
  ## - `os1` as the rank-1 output's element stride
  ## - `laneRow`/`laneCol` as the calling lane's cell in the atom record
  var body = nnkStmtList.newTree()
  # direct nodes, not quote, for the later fragments reference these consts
  # by name and a quote-do `const` gensyms its names
  body.add(nnkConstSection.newTree(
    nnkConstDef.newTree(id"atom", newEmptyNode(),
      nnkCall.newTree(id"getTileConfig", id"T", id"T")),
    nnkConstDef.newTree(id"atomM", newEmptyNode(),
      nnkCall.newTree(dot(id"atom", "getM"))),
    nnkConstDef.newTree(id"atomN", newEmptyNode(),
      nnkCall.newTree(dot(id"atom", "getN"))),
    nnkConstDef.newTree(id"atomVpt", newEmptyNode(),
      nnkCall.newTree(dot(id"atom", "getVpt")))))
  body.add(nnkStaticStmt.newTree(nnkStmtList.newTree(
    nnkCall.newTree(id"doAssert",
      infixF("==", nnkCall.newTree(dot(id"atom", "getThreadCount")),
             newLit(32)),
      newLit("tile programs need one atom's lanes per threadgroup")))))
  if plan.flatRect:
    body.add(nnkConstSection.newTree(
      nnkConstDef.newTree(id"tileRc", newEmptyNode(), id"atomM"),
      nnkConstDef.newTree(id"tileCc", newEmptyNode(), id"atomN")))
  else:
    body.add(nnkConstSection.newTree(
      nnkConstDef.newTree(id"tileRc", newEmptyNode(), newLit(plan.tileR)),
      nnkConstDef.newTree(id"tileCc", newEmptyNode(), newLit(plan.tileC))))
  body.add(quote do:
    static:
      doAssert tileRc mod atomM == 0 and tileCc mod atomN == 0,
        "schedule tile sizes must be multiples of the atom's M and N")

  # prefix dims, decomposed from bz over batch and depth
  case plan.prefixDims.len
  of 1:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"bCoord", newEmptyNode(), id"bz")))
  of 2:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"bExtent", newEmptyNode(),
      extentExpr(ir.iterExtent(plan.prefixDims[1]), device = true))))
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"bCoord", newEmptyNode(), infixF("div", id"bz", id"bExtent"))))
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"dCoord", newEmptyNode(), infixF("mod", id"bz", id"bExtent"))))
  else: discard

  let L = ir.outDims.len
  if L >= 1:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"rowExtent", newEmptyNode(),
      extentExpr(ir.iterExtent(plan.rowDim), device = true))))
  if L >= 2:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"colExtent", newEmptyNode(),
      extentExpr(ir.iterExtent(plan.colDim), device = true))))
  if plan.flatRect:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"flatRows", newEmptyNode(),
      callI32(infixF("div",
        infixF("+",
          extentExpr(ir.iterExtent(ir.reduceIters[0]), device = true),
          infixF("-", id"atomN", newLit(1))),
        id"atomN")))))

  let hasReduce = ir.reduceIters.len > 0
  let isReduceStore = hasReduce and L >= 1
  if plan.flatRect or (hasReduce and L == 1) or plan.allReduce:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"lane", newEmptyNode(), nnkCall.newTree(id"int", id"tid"))))
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"cell", newEmptyNode(),
      nnkCall.newTree(id"toIntVal",
        nnkCall.newTree(id"crd2idx",
          nnkCall.newTree(dot(id"atom", "getLayoutA")),
          nnkPar.newTree(id"lane", newLit(0)))))))
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"laneRow", newEmptyNode(),
      infixF("mod", id"cell", id"atomM"))))
  if plan.allReduce:
    body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      id"laneCol", newEmptyNode(), infixF("div", id"cell", id"atomM"))))

  var accName = ""
  if isReduceStore or plan.allReduce:
    accName = "acc"
    body.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(accName),
      nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
      newEmptyNode())))
    body.add(nnkCall.newTree(dot(id(accName), "zero")))

  # loop nest with the rows chunk (all-reduce) outer, the col chunk next
  # and then the serial iters. The innermost body gathers every read,
  # combines the tiles and adds the product into the accumulator
  var resolve = initTable[string, NimNode]()
  for it in plan.serialIters:
    resolve[it] = id(it)

  var inner = nnkStmtList.newTree()
  var tileOf = initTable[string, string]()
  for i in 0 ..< ir.reads.len:
    let g = ir.emitGather(plan, i, resolve)
    tileOf[readKey(ir.reads[i])] = g.tile
    for st in g.stmts:
      inner.add(st)
  let accI = id(accName)
  var counter = 0
  let (valueTile, valueScalar) = emitTileValue(ir.expr, tileOf, inner,
                                               counter)
  let value = applyScalar(valueTile, valueScalar, inner, counter)
  doAssert value.len > 0,
    "reported gap: a literal-only statement has no tile (see emitTileValue)"
  if accName.len > 0:
    inner.add(nnkCall.newTree(dot(id(accName), "addScaled"),
      id(value), newLit(1.0'f32)))
  else:
    # elementwise form, the combined tile stores directly
    var oBat, oDep: NimNode = i32lit(0)
    if L >= 3: oBat = callI32(layoutField("o", "stride", 0))
    if L == 4: oDep = callI32(layoutField("o", "stride", 1))
    let oRow = (if L == 1: callI32(layoutField("o", "stride", 1))
              else: callI32(layoutField("o", "stride", L - 2)))
    let oCol = (if L >= 2: callI32(layoutField("o", "stride", L - 1))
              else: i32lit(0))
    inner.add(gdViewNode("gdO", "o", oBat, oDep, oRow, oCol, i32lit(0)))
    inner.add(nnkCall.newTree(dot(id"gdO", "storeTileMasked"), id(value),
      nnkPar.newTree(
        (if plan.prefixDims.len >= 1: id"bCoord" else: i32lit(0)),
        (if plan.prefixDims.len >= 2: id"dCoord" else: i32lit(0)),
        (if L >= 2: id"by" else: i32lit(0)),
        (if L >= 2: id"bx" else: i32lit(0))),
      (if L == 1: infixF("-", id"rowExtent",
                          infixF("*", id"by", id"tileRc"))
       else: id"rowExtent"),
      (if L >= 2: infixF("-", id"colExtent",
                          infixF("*", id"bx", id"tileCc"))
       else: id"tileCc")))

  # wrap the chunks and serial iters, outermost first
  var stmts = inner
  proc wrapFor(varName: string, bound: NimNode, s: NimNode): NimNode =
    nnkForStmt.newTree(id(varName),
      nnkInfix.newTree(id"..<", newLit(0'i32), callI32(bound)),
      nnkStmtList.newTree(s))
  if plan.allReduce:
    stmts = wrapFor("rchunk",
      extentExpr(ir.iterExtent(plan.rowDim), device = true), stmts)
    if plan.colDim.len > 0:
      stmts = wrapFor("cchunk",
        extentExpr(ir.iterExtent(plan.colDim), device = true), stmts)
  elif plan.chunkIter.len > 0:
    stmts = wrapFor("chunk",
      infixF("div",
        infixF("+",
          extentExpr(ir.iterExtent(plan.chunkIter), device = true),
          infixF("-", id"tileCc", newLit(1))),
        id"tileCc"),
      stmts)
  for it in countdown(plan.serialIters.len - 1, 0):
    stmts = wrapFor(plan.serialIters[it],
      extentExpr(ir.iterExtent(plan.serialIters[it]), device = true), stmts)
  body.add(stmts)

  # epilogues, one per store form
  if plan.allReduce:
    let combineBound = (if plan.flatRect: id"atomM" else: id"tileRc")
    let rvI = id"rvAcc"
    body.add(quote do:
      var `rvI`: rv(HostElem, tileRc, tileCc)
      `rvI`.row_sum(`accI`))
    if plan.flatRect:
      body.add(quote do:
        if laneCol == 0:
          part[laneRow] = `rvI`.data[0])
    else:
      body.add(quote do:
        if laneCol == 0:
          for n in 0'i32 ..< int32(tileRc div atomM):
            part[laneRow + n * atomM] = `rvI`.data[n * atomVpt])
    body.add(nnkCall.newTree(id"threadgroup_barrier"))
    body.add(quote do:
      if tid == 0:
        var total = 0.0'f32
        for i in 0'i32 ..< int32(`combineBound`):
          total += part[i]
        o.data[0] = total)
    # the all-reduce partial buffer is threadgroup memory owned by the entry
    var partSize = ir.threads
    if not plan.flatRect and plan.tileR > partSize:
      partSize = plan.tileR
    result.smemDecls = nnkStmtList.newTree(smemDecl("part", partSize, 1))
    result.scratch = @["part"]
  elif L == 1 and hasReduce:
    let rvI = id"rvAcc"
    body.add(quote do:
      var `rvI`: rv(HostElem, tileRc, tileCc)
      `rvI`.row_sum(`accI`)
      let os1 = int32(o.layout.stride[1])
      let rvRows = int32(tileRc div atomM)
      for n in 0'i32 ..< rvRows:
        let row = by * tileRc + laneRow + n * atomM
        if row < rowExtent:
          o.data[row * os1] = `rvI`.data[n * atomVpt])
  elif L >= 2 and hasReduce:
    var oBat, oDep: NimNode = i32lit(0)
    if L >= 3: oBat = callI32(layoutField("o", "stride", 0))
    if L == 4: oDep = callI32(layoutField("o", "stride", 1))
    body.add(gdViewNode("gdO", "o", oBat, oDep,
      callI32(layoutField("o", "stride", L - 2)),
      callI32(layoutField("o", "stride", L - 1)), i32lit(0)))
    body.add(nnkCall.newTree(dot(id"gdO", "storeTileMasked"), id(accName),
      nnkPar.newTree(
        (if plan.prefixDims.len >= 1: id"bCoord" else: i32lit(0)),
        (if plan.prefixDims.len >= 2: id"dCoord" else: i32lit(0)),
        id"by", id"bx"),
      infixF("-", id"rowExtent", infixF("*", id"by", id"tileRc")),
      infixF("-", id"colExtent", infixF("*", id"bx", id"tileCc"))))
  result.body = body

# ─────────────────────────────────────────────────────────────────────────
# kernel emission, device fn plus thin entry and compile-time AST echo
# ─────────────────────────────────────────────────────────────────────────
proc tensorViewNode(elem: NimNode, rank: int): NimNode =
  ## Rank-N runtime-shaped `TensorView[Elem, dims, dims]` formal type.
  var dims = nnkPar.newTree()
  for i in 0 ..< rank:
    dims.add(ident"int")
  nnkBracketExpr.newTree(ident"TensorView", elem, dims, dims)

proc tensorRank(ir: AlgoIR, tensorName: string): int =
  ## Device view rank of a tensor, its widest read, at least 2 (vectors live in rank-2 (n, 1) views).
  result = 2
  for rd in ir.reads:
    if rd.tensor == tensorName:
      result = max(result, rd.indices.len)

proc tensorParams(ir: AlgoIR): seq[string] =
  for rd in ir.reads:
    if rd.tensor notin result:
      result.add(rd.tensor)

proc emitDeviceFunc(ir: AlgoIR, body: NimNode, scratch: seq[string]): NimNode =
  ## Device function, generic over the element type, taking the views, the threadgroup scratch pointers and the launch geometry as parameters.
  ##
  ## Crucible attaches device semantics itself when it pulls the called
  ## proc (nim_to_gpu registerGenericInstOrExternalProc), so no pragma is needed.
  ##
  ## The compile-time target is preset in genKernel. On this Nim the body
  ## sems eagerly at module scope and getTileConfig needs that target.
  let fnName = id(ir.name & "Fn")
  var formals = nnkFormalParams.newTree(newEmptyNode())
  formals.add(nnkIdentDefs.newTree(ident"o",
    tensorViewNode(id"T", max(2, ir.outDims.len)), newEmptyNode()))
  for t in ir.tensorParams():
    formals.add(nnkIdentDefs.newTree(ident(t),
      tensorViewNode(id"T", ir.tensorRank(t)), newEmptyNode()))
  for sc in scratch:
    formals.add(nnkIdentDefs.newTree(
      nnkPragmaExpr.newTree(ident(sc), nnkPragma.newTree(ident"smem")),
      nnkPtrTy.newTree(
        nnkBracketExpr.newTree(ident"UncheckedArray", id"T")),
      newEmptyNode()))
  for g in ["tid", "nthr", "bx", "by", "bz"]:
    formals.add(nnkIdentDefs.newTree(ident(g), ident"int32", newEmptyNode()))
  # No export marker, the module exporting nothing.
  # Shape constraints from this Nim's macro emission:
  # - the 7-child procdef shape (name, pattern, genericParams, formals, pragmas, exceptions, body) sems
  #   the 8-child parser shape with a leading export marker does not
  # - the `{.device.}` pragma fails to instantiate on a macro-emitted procdef, so the marker is omitted
  #   and crucible's pull path attaches device semantics (nim_to_gpu registerGenericInstOrExternalProc)
  let genParams = nnkGenericParams.newTree(
    nnkIdentDefs.newTree(ident"T", newEmptyNode(), newEmptyNode()))
  result = nnkProcDef.newTree(
    fnName, newEmptyNode(), genParams, formals,
    newEmptyNode(), newEmptyNode(), body)

proc emitEntry(ir: AlgoIR, grid, outShape, smemDecls: NimNode, scratch: seq[string]): NimNode =
  ## Thin entry holding the geometry lets, the device fn call and the scratch call arguments.
  ##
  ## - the threadgroup-memory declarations hoist into the wrapper via the spec
  ## - scratch arguments use `cast[ptr UncheckedArray[HostElem]](addr <smem local>)`, the only `{.smem.}` source surviving the pull path
  ##
  ## Registers the launch spec for `engine.eval(name(...))`.
  let spec = nnkBracket.newTree(
    nnkExprEqExpr.newTree(ident"output", ident"o"),
    nnkExprEqExpr.newTree(ident"outShape", outShape),
    nnkExprEqExpr.newTree(ident"grid", grid),
    nnkExprEqExpr.newTree(ident"blk",
      nnkPar.newTree(newLit(ir.threads), newLit(1), newLit(1))))
  if smemDecls.len > 0:
    spec.add(nnkExprEqExpr.newTree(ident"smemDecls", smemDecls))
  DescSpecs[ir.name] = spec
  var formals = nnkFormalParams.newTree(newEmptyNode())
  formals.add(nnkIdentDefs.newTree(ident"o",
    tensorViewNode(id"HostElem", max(2, ir.outDims.len)), newEmptyNode()))
  for t in ir.tensorParams():
    formals.add(nnkIdentDefs.newTree(ident(t),
      tensorViewNode(id"HostElem", ir.tensorRank(t)), newEmptyNode()))

  var lowered = nnkStmtList.newTree()
  let geometry = @["tid", "nthr", "bx", "by", "bz"]
  let geometryRhs = @[
    nnkCall.newTree(ident"int32", ident"thread_index_in_threadgroup"),
    nnkCall.newTree(ident"int32",
      nnkDotExpr.newTree(ident"threads_per_threadgroup", ident"x")),
    nnkCall.newTree(ident"int32",
      nnkDotExpr.newTree(ident"threadgroup_position_in_grid", ident"x")),
    nnkCall.newTree(ident"int32",
      nnkDotExpr.newTree(ident"threadgroup_position_in_grid", ident"y")),
    nnkCall.newTree(ident"int32",
      nnkDotExpr.newTree(ident"threadgroup_position_in_grid", ident"z"))]
  for gi in 0 ..< geometry.len:
    lowered.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
      newIdentNode(geometry[gi]), newEmptyNode(), geometryRhs[gi])))

  var call = nnkCall.newTree(id(ir.name & "Fn"), ident"o")
  for t in ir.tensorParams():
    call.add(ident(t))
  for s in scratch:
    call.add(nnkCast.newTree(
      nnkPtrTy.newTree(
        nnkBracketExpr.newTree(ident"UncheckedArray", id"HostElem")),
      nnkPrefix.newTree(ident"addr", ident(s))))
  for g in geometry:
    call.add(ident(g))
  lowered.add(call)

  result = nnkTemplateDef.newTree(
    newIdentNode(ir.name), newEmptyNode(), newEmptyNode(),
    formals, newEmptyNode(), newEmptyNode(),
    lowered)

  echo "####################[ derived kernel: " & ir.name & " ]####################"
  echo result.repr
  echo "####################[ end derived kernel: " & ir.name & " ]####################"

proc formName(ir: AlgoIR, form: LowerForm): string =
  case form
  of lfScalar:
    if ir.outDims.len == 0:
      "all-reduce tile program (atom rectangle or chunked tiles, cross-lane partial combine)"
    elif ir.reduceIters.len == 0:
      "elementwise tile program (derived gather, combine, masked store)"
    else:
      "reduction tile program (derived gather, combine, row_sum + rv or tile store)"
  of lfBlockedCopy: "blocked tile copy (schedule-tiled transposed gather and store)"
  of lfMmaStaged: "mma kernel, smem-staged operands + gemm accumulation"
  of lfMmaRegister: "mma kernel, register tiles (loadTile + mma_AB)"

proc lowerEmit(ir: var AlgoIR, pass: Pass): NimNode =
  let choice = ir.chooseForm()
  pass.description = formName(ir, choice.form)
  var body: NimNode
  var smemDecls = nnkStmtList.newTree()
  var scratch: seq[string] = @[]
  let L = ir.outDims.len

  var gridX, gridY, gridZ: NimNode
  if choice.form in {lfScalar, lfBlockedCopy}:
    let plan = ir.makeTilePlan()
    (body, smemDecls, scratch) = ir.emitTileKernel(plan)
    if L == 0:
      gridX = newLit(1)
      gridY = newLit(1)
      gridZ = newLit(1)
    else:
      if L >= 2:
        gridX = nnkCall.newTree(ident"cdiv",
          extentExpr(ir.iterExtent(plan.colDim), device = false),
          i32lit(plan.tileC))
      else:
        gridX = newLit(1)
      gridY = nnkCall.newTree(ident"cdiv",
        extentExpr(ir.iterExtent(plan.rowDim), device = false),
        i32lit(plan.tileR))
      case plan.prefixDims.len
      of 1:
        gridZ = extentExpr(ir.iterExtent(plan.prefixDims[0]), device = false)
      of 2:
        gridZ = infixF("*",
          extentExpr(ir.iterExtent(plan.prefixDims[0]), device = false),
          extentExpr(ir.iterExtent(plan.prefixDims[1]), device = false))
      else:
        gridZ = newLit(1)
  else:
    let tiles = choice.tiles
    if choice.form == lfMmaStaged:
      (body, smemDecls, scratch) = ir.emitMmaStaged(tiles)
    else:
      body = ir.emitMmaRegister(tiles)
    gridX = nnkCall.newTree(ident"cdiv",
      extentExpr(ir.iterExtent(ir.outDims[1]), device = false),
      i32lit(tiles[1]))
    gridY = nnkCall.newTree(ident"cdiv",
      extentExpr(ir.iterExtent(ir.outDims[0]), device = false),
      i32lit(tiles[0]))
    gridZ = newLit(1)

  let grid = nnkPar.newTree(gridX, gridY, gridZ)
  var outDimsHost: seq[NimNode] = @[]
  if L == 0:
    outDimsHost = @[callI32(i32lit(1)), callI32(i32lit(1))]
  elif L == 1:
    outDimsHost = @[callI32(i32lit(1)),
      extentExpr(ir.iterExtent(ir.outDims[0]), device = false)]
  else:
    for d in 0 ..< L:
      outDimsHost.add(extentExpr(ir.iterExtent(ir.outDims[d]),
                                 device = false))
  let outShape = nnkPrefix.newTree(ident"@", nnkBracket.newTree(outDimsHost))

  # schedule provenance, which schedule values produced which tile types
  let (bm, bn, bk) = choice.tiles
  var provenance: string
  case choice.form
  of lfScalar:
    if L == 0:
      provenance = (if ir.makeTilePlan().flatRect:
        "all-reduce, one threadgroup, atom-rectangle unrank + partial combine"
      else:
        "all-reduce, one threadgroup, chunked tiles + partial combine")
    else:
      provenance = "tile program " & $ir.makeTilePlan().tileR & "x" &
        $ir.makeTilePlan().tileC &
        ", derived gather/combine, atom from getTileConfig(HostElem, HostElem)"
  of lfBlockedCopy:
    provenance = "blocked tile copy " & $bm & "x" & $bn &
      ", transposed gather and masked store"
  of lfMmaStaged:
    provenance = "tiles y=" & $bm & " x=" & $bn & " k=" & $bk &
      ", cache: smem staged tiles, acc rt_l(HostElem, " & $bm & ", " &
      $bn & "), gemm per-atom-slab"
  of lfMmaRegister:
    provenance = "tiles y=" & $bm & " x=" & $bn & " k=" & $bk &
      ", register-only: ta rt_l(HostElem, " & $bm & ", " & $bk & "), tb rt_r(" &
      $bk & ", " & $bn & "), acc rt_l(HostElem, " & $bm & ", " & $bn & ")"
  SchedInfo[ir.name] = newLit(provenance)

  let fnDef = ir.emitDeviceFunc(body, scratch)
  let entry = ir.emitEntry(grid, outShape, smemDecls, scratch)
  echo "=== derived function for '", ir.name, "' ===\n", fnDef.repr
  echo "=== thin entry for '", ir.name, "' ===\n", entry.repr
  when defined(onlyFn):
    result = nnkStmtList.newTree(fnDef)
  elif defined(onlyEntry):
    result = nnkStmtList.newTree(entry)
  else:
    result = nnkStmtList.newTree(fnDef, entry)

proc lowerPass(): Pass =
  ## Final pass, pattern-matches the IR and emits the device fn + thin entry
  ## through this module's derivation-first tile programs.
  Pass(name: "lower",
       description: "pattern-match IR and emit device code",
       kind: pkLowering,
       emit: lowerEmit)

# ─────────────────────────────────────────────────────────────────────────
# the ++= macro
# ─────────────────────────────────────────────────────────────────────────
proc buildPasses(attrs: NimNode): seq[Pass] =
  ## Schedule block -> ordered pass list, composed in written order.
  var sched: seq[Pass] = @[]
  if attrs != nil:
    doAssert attrs.kind == nnkStmtList and attrs.len == 1 and
      attrs[0].kind == nnkCall and attrs[0][0].strVal == "schedule",
      "the do block must contain exactly one schedule: section"
    for st in attrs[0][1]:
      case st.kind
      of nnkCall:
        let callee = st[0].strVal
        case callee
        of "tile":
          let loop = st[1].strVal
          let size = (if st.len >= 3 and st[2].kind == nnkIntLit:
                        st[2].intVal.int
                      else:
                        DefaultTile)
          sched.add(tilePass(loop, size))
        of "cache":
          var tensors: seq[string] = @[]
          for arg in st[1 .. ^1]:
            if arg.kind == nnkIdent:
              tensors.add(arg.strVal)
          sched.add(cacheAnnotatePass(tensors))
          sched.add(cachePlanPass())
        of "parallel":
          sched.add(parallelPass(st[1].strVal))
        else:
          error("unknown schedule property " & callee, st)
      of nnkAsgn:
        doAssert st[0].strVal == "threads" and st[1].kind == nnkIntLit,
          "the only schedule assignment is threads = n"
        sched.add(threadsPass(st[1].intVal.int))
      else:
        error("unsupported schedule statement " & $st.kind, st)
  sched.add(lowerPass())
  result = sched

proc genKernel(target, value, attrs: NimNode): NimNode =
  # Emitted function bodies sem eagerly at module scope on this Nim
  # (macro-emitted generic procdef bodies do not defer), so the compile-
  # time target that getTileConfig's backend selection needs must be
  # preset here. The DSL compiles through the Metal eval branch only,
  # the other eval branches being tensorcery's unimplemented placeholders.
  crucibleCompileTarget = ctMetal
  var ir: AlgoIR
  target.parseTarget(ir)
  value.parseValue(ir)
  ir.buildLoops()
  ir.dtype = HostElem.name
  ir.threads = DefaultThreads
  let passes = buildPasses(attrs)
  result = nnkStmtList.newTree(runPipeline(ir, passes))

macro `++=`(target, value: untyped): untyped =
  ## Algorithm statement without a schedule, see the module header.
  genKernel(target, value, nil)

macro `++=`(target, value, attrs: untyped): untyped =
  ## Algorithm statement with a `do: schedule: ...` block.
  genKernel(target, value, attrs)

# ─────────────────────────────────────────────────────────────────────────
# Einsum section 2 (Rocktaschel 2018), all twelve recipes as DSL kernels
# ─────────────────────────────────────────────────────────────────────────
# Schedules are naive where elementwise, 16-tiles plus cache for the GEMM, 32-tiles for the blocked transpose.
# bmm and bilinear take 16-tiles, the contraction runs flat.
# Vectors live in rank-2 (n, 1) views, one read slot addresses axis 0.
# ─────────────────────────────────────────────────────────────────────────

# 2.1 Matrix transpose as a blocked tile copy, a pure reindex via the permutation matcher.
# The schedule-tiled transposed gather and masked store move both block
# sides through tile primitives.
xpose[j, i] ++= a[i, j] | (i, j) do:
  schedule:
    tile(j, 32)
    tile(i, 32)
    threads = 32

# untiled 2.1 for contrast, elementwise tile program with strided column reads
xposeNaive[j, i] ++= a[i, j] | (i, j)

# 2.2 Sum over all indices, the cross-lane all-reduce
totSum[] ++= a[i, j] | (i, j)

# 2.3 Column sum, rows = j, the reduce chunked over the tile cols
colSum[j] ++= a[i, j] | (j, i)

# 2.4 Row sum
rowSum[i] ++= a[i, j] | (i, j)

# 2.5 Matrix-vector product
matvec[i] ++= m[i, k] * v[k] | (i, k)

# 2.6 Matrix-matrix product, the working GEMM schedule (16-tiles + cache)
mm[i, j] ++= m1[i, k] * m2[k, j] | (j, i, k) do:
  schedule:
    tile(i, 16)
    tile(j, 16)
    tile(k, 16)
    cache(m1)
    cache(m2)
    threads = 32

# 2.7 Dot product
dotP[] ++= u[i] * v[i] | (i)

# 2.8 Hadamard product
hadamard[i, j] ++= a[i, j] * b[i, j] | (j, i)

# 2.9 Outer product
outerP[i, j] ++= u[i] * v[j] | (j, i)

# 2.10 Batch matrix multiplication, the batch dim on grid.z, the matrix
# dims tiled at 16 (guarded scalar form, no mma path for three out dims)
bmm[i, j, l] ++= a3[i, j, k] * b3[i, k, l] | (i, j, l, k) do:
  schedule:
    tile(j, 16)
    tile(l, 16)
    threads = 32

# 2.11 Tensor contraction, four out dims unrank a flat per-thread index
contract[i, k, m, o] ++= ca[i, j, k, l, o] * cb[m, j, l, o] |
    (i, k, m, o, j, l) do:
  schedule:
    threads = 32

# 2.12 Bilinear transformation
bilinear[i, j] ++= a[i, k] * b3[j, k, l] * c[i, l] | (i, j, k, l) do:
  schedule:
    tile(j, 16)
    threads = 32

# ─────────────────────────────────────────────────────────────────────────
# reference implementations in laser benchmark style, plain scalar procs over seqs with for loops and manual row-major indexing
# ─────────────────────────────────────────────────────────────────────────
proc refXpose(a: seq[float32], m, n: int): seq[float32] =
  ## Naive scalar transpose (strided column reads).
  result = newSeq[float32](m * n)
  for i in 0 ..< m:
    for j in 0 ..< n:
      result[j * m + i] = a[i * n + j]

proc refTotSum(a: seq[float32], m, n: int): float32 =
  result = 0.0'f32
  for i in 0 ..< m:
    for j in 0 ..< n:
      result += a[i * n + j]

proc refColSum(a: seq[float32], m, n: int): seq[float32] =
  result = newSeq[float32](n)
  for j in 0 ..< n:
    for i in 0 ..< m:
      result[j] += a[i * n + j]

proc refRowSum(a: seq[float32], m, n: int): seq[float32] =
  result = newSeq[float32](m)
  for i in 0 ..< m:
    for j in 0 ..< n:
      result[i] += a[i * n + j]

proc refMatvec(m: seq[float32], v: seq[float32], rows, k: int): seq[float32] =
  result = newSeq[float32](rows)
  for i in 0 ..< rows:
    for kk in 0 ..< k:
      result[i] += m[i * k + kk] * v[kk]

proc refMatmul(a, b: seq[float32], m, k, n: int): seq[float32] =
  result = newSeq[float32](m * n)
  for i in 0 ..< m:
    for j in 0 ..< n:
      for kk in 0 ..< k:
        result[i * n + j] += a[i * k + kk] * b[kk * n + j]

proc refDot(u, v: seq[float32], n: int): float32 =
  result = 0.0'f32
  for i in 0 ..< n:
    result += u[i] * v[i]

proc refHadamard(a, b: seq[float32], m, n: int): seq[float32] =
  result = newSeq[float32](m * n)
  for i in 0 ..< m:
    for j in 0 ..< n:
      result[i * n + j] = a[i * n + j] * b[i * n + j]

proc refOuter(u, v: seq[float32], m, n: int): seq[float32] =
  result = newSeq[float32](m * n)
  for i in 0 ..< m:
    for j in 0 ..< n:
      result[i * n + j] = u[i] * v[j]

proc refBmm(a, b: seq[float32], batch, m, k, n: int): seq[float32] =
  result = newSeq[float32](batch * m * n)
  for bi in 0 ..< batch:
    for i in 0 ..< m:
      for j in 0 ..< n:
        for kk in 0 ..< k:
          result[bi * m * n + i * n + j] +=
            a[bi * m * k + i * k + kk] * b[bi * k * n + kk * n + j]

proc refContract(a, b: seq[float32]): seq[float32] =
  ## C[i,k,m,o] = sum_{j,l} a[i,j,k,l,o] * b[m,j,l,o].
  ## Shapes a (2,3,4,5,8), b (7,3,5,8), out (2,4,7,8).
  const (ni, nj, nk, nl, nm, no) = (2, 3, 4, 5, 7, 8)
  result = newSeq[float32](ni * nk * nm * no)
  for i in 0 ..< ni:
    for k in 0 ..< nk:
      for m in 0 ..< nm:
        for o in 0 ..< no:
          var acc = 0.0'f32
          for j in 0 ..< nj:
            for l in 0 ..< nl:
              acc += a[((((i * nj + j) * nk + k) * nl + l) * no) + o] *
                     b[((m * nj + j) * nl + l) * no + o]
          result[((i * nk + k) * nm + m) * no + o] = acc

proc refBilinear(a, b, c: seq[float32]): seq[float32] =
  ## D[i,j] = sum_{k,l} a[i,k] * b[j,k,l] * c[i,l], shapes a (2,3), b (5,3,7), c (2,7), out (2,5).
  const (ni, nk, nj, nl) = (2, 3, 5, 7)
  result = newSeq[float32](ni * nj)
  for i in 0 ..< ni:
    for j in 0 ..< nj:
      var acc = 0.0'f32
      for k in 0 ..< nk:
        for l in 0 ..< nl:
          acc += a[i * nk + k] * b[(j * nk + k) * nl + l] * c[i * nl + l]
      result[i * nj + j] = acc

when isMainModule:
  import std/[math, random, sequtils, syncio]

  proc main() =
    proc flat(t: tensorcery.Tensor): seq[float32] =
      ## Flat read of a host tensor's elements.
      t.view(float32, bkMetal).toSeq()

    proc maxAbsErr(x, y: seq[float32]): float32 =
      result = 0.0'f32
      for i in 0 ..< x.len:
        result = max(result, abs(x[i] - y[i]))

    proc check(name: string, got, want: seq[float32], tol: float32) =
      let e = maxAbsErr(got, want)
      echo name, ": max abs err ", e,
           (if e <= tol: "  PASS" else: "  FAIL")
      if e > tol:
        for i in 0 ..< min(6, got.len):
          echo "  got[", i, "] ", got[i], "  want ", want[i]
        var worst = 0
        for i in 0 ..< got.len:
          if abs(got[i] - want[i]) > abs(got[worst] - want[worst]):
            worst = i
        echo "  worst flat[", worst, "] got ", got[worst], " want ",
             want[worst]
      doAssert e <= tol, name & " mismatch"

    # host runtime state lives in main's scope so the destructors run
    var rng = initRand(13)

    proc randSeq(n: int): seq[float32] =
      result.newSeq(n)
      for i in 0 ..< n:
        result[i] = rng.gauss()

    let a23H = randSeq(6)
    let a23bH = randSeq(6)
    let m64H = randSeq(64 * 64)
    let mvVH = randSeq(64)
    let m1_64H = randSeq(64 * 64)
    let m2_64H = randSeq(64 * 64)
    let u64H = randSeq(64)
    let v64H = randSeq(64)
    let u3H = randSeq(3)
    let v4H = randSeq(4)
    let a325H = randSeq(3 * 2 * 5)
    let b353H = randSeq(3 * 5 * 3)
    let caH = randSeq(2 * 3 * 4 * 5 * 8)
    let cbH = randSeq(7 * 3 * 5 * 8)
    let bil_aH = randSeq(2 * 3)
    let bil_bH = randSeq(5 * 3 * 7)
    let bil_cH = randSeq(2 * 7)
    let a23 = tensor(a23H, @[2, 3])
    let a23b = tensor(a23bH, @[2, 3])
    let m64 = tensor(m64H, @[64, 64])
    let mvV = tensor(mvVH, @[64, 1])
    let m1_64 = tensor(m1_64H, @[64, 64])
    let m2_64 = tensor(m2_64H, @[64, 64])
    let u64 = tensor(u64H, @[64, 1])
    let v64 = tensor(v64H, @[64, 1])
    let u3 = tensor(u3H, @[3, 1])
    let v4 = tensor(v4H, @[4, 1])
    let a325 = tensor(a325H, @[3, 2, 5])
    let b353 = tensor(b353H, @[3, 5, 3])
    let ca = tensor(caH, @[2, 3, 4, 5, 8])
    let cb = tensor(cbH, @[7, 3, 5, 8])
    let bil_a = tensor(bil_aH, @[2, 3])
    let bil_b = tensor(bil_bH, @[5, 3, 7])
    let bil_c = tensor(bil_cH, @[2, 7])

    var engine = setBackend(bkMetal)

    echo "=== einsum section 2, DSL kernels vs scalar references ==="
    echo "product reductions tolerate ~2e-6 (the C reference contracts the"
    echo "multiply-add into an FMA, Metal rounds mul and add separately,"
    echo "and the tile reductions reorder the summation)"
    echo "xpose runs the blocked tile copy, xposeNaive the elementwise form"
    check "2.1  transpose   ", engine.eval(xpose(a23)).flat,
          refXpose(a23H, 2, 3), 0.0'f32
    check "2.1  xposeNaive ", engine.eval(xposeNaive(a23)).flat,
          refXpose(a23H, 2, 3), 0.0'f32
    check "2.2  sum        ", engine.eval(totSum(a23)).flat,
          @[refTotSum(a23H, 2, 3)], 2e-6'f32
    check "2.3  colsum     ", engine.eval(colSum(a23)).flat,
          refColSum(a23H, 2, 3), 2e-6'f32
    check "2.4  rowsum     ", engine.eval(rowSum(a23)).flat,
          refRowSum(a23H, 2, 3), 2e-6'f32
    check "2.5  matvec     ", engine.eval(matvec(m64, mvV)).flat,
          refMatvec(m64H, mvVH, 64, 64), 1e-5'f32
    check "2.6  matmul     ", engine.eval(mm(m1_64, m2_64)).flat,
          refMatmul(m1_64H, m2_64H, 64, 64, 64), 1e-5'f32
    check "2.7  dot        ", engine.eval(dotP(u64, v64)).flat,
          @[refDot(u64H, v64H, 64)], 2e-6'f32
    check "2.8  hadamard   ", engine.eval(hadamard(a23, a23b)).flat,
          refHadamard(a23H, a23bH, 2, 3), 0.0'f32
    check "2.9  outer      ", engine.eval(outerP(u3, v4)).flat,
          refOuter(u3H, v4H, 3, 4), 0.0'f32
    check "2.10 bmm        ", engine.eval(bmm(a325, b353)).flat,
          refBmm(a325H, b353H, 3, 2, 5, 3), 2e-6'f32
    check "2.11 contraction", engine.eval(contract(ca, cb)).flat,
          refContract(caH, cbH), 2e-6'f32
    check "2.12 bilinear   ", engine.eval(bilinear(bil_a, bil_b, bil_c)).flat,
          refBilinear(bil_aH, bil_bH, bil_cH), 2e-6'f32

    echo "=== einsum section 2 green (12 recipes) ==="

  main()