Some more experimentation, this time using 2019 Lux IR as a base (https://github.com/mratsim/laser/blob/master/laser/lux_compiler/core/lux_types.nim, https://github.com/mratsim/laser/blob/master/laser/lux_compiler/core/README.md, https://github.com/mratsim/laser/blob/master/laser/lux_compiler/core/challenges.md, https://github.com/mratsim/laser/blob/master/laser/lux_compiler/core/low_level_design_considerations.md )
proc mm(m1, m2: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
var j = domain(64)
var k = domain(64)
result[i, j] += m1[i, k] * m2[k, j] do:
tile(i, 16)
tile(j, 16)
tile(k, 16)
cache(m1)
cache(m2)
threads = 32The mma pass recognizes the canonical shape from the IR and dispatches to the staged tensor-core kernel; delete the schedule lines and the same statement runs as a generic tile program (slower, identical result).
proc softmaxMvp(x: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(rows)
var j = domain(cols)
maxRow[i] = x[i, j] do:
reduce(max)
probs[i, j] = exp2((x[i, j] - maxRow[i]) * 1.4426950408889634) do:
tile(j, 16)
sumRow[i] = probs[i, j] do:
tile(j, 16)
result[i, j] = probs[i, j] / sumRow[i] do:
tile(j, 16)The stages feed each other; the fusion pass groups them into one kernel (intermediates in registers, no scratch buffers). Numerically stable (max-shifted).
proc gqa(q, kCache, vCache, btL: LuxTensor): LuxTensor {.luxtile.} =
var t = domain(T)
var kv = domain(Nkv)
var rp = domain(4)
var d = domain(D)
var b = domain(PB)
var jc = domain(8)
let q4 = q.view(T, Nkv, 4, D)
scores[t, kv, b, jc, rp] = where(8*b + jc <= btL[t*5 + 4],
q4[t, kv, rp, d] * kCache[btL[t*5 + b], kv, jc*16 + d],
-1.0e7) do:
reduce(sum)
maxNew[t, kv, rp] = scores[t, kv, b, jc, rp] do:
loopNest(t, kv, rp, jc, b)
guard(8*b <= btL[t*5 + 4])
reduce(max)
probs[t, kv, b, jc, rp] = exp2(scores[t, kv, b, jc, rp] - maxNew[t, kv, rp])
sumPartial[t, kv, b, rp] = probs[t, kv, b, jc, rp] do:
reduce(sum)
sumTotal[t, kv, rp] = sumPartial[t, kv, b, rp] do:
reduce(sum)
outSum[t, kv, rp, d] = probs[t, kv, b, jc, rp] *
vCache[btL[t*5 + b], kv*128 + jc*16 + d] do:
tile(d, 16)
reduce(sum)
result[t, kv, rp, d] = outSum[t, kv, rp, d] / sumTotal[t, kv, rp]
# ─────────────────────────────────────────────────────────────────────────What's inside, each a general IR concept, none attention-specific:
btL[t*5 + b]— a runtime tensor read in an index (TensorRead): the paged block table drives the k/v cache origins. Paged fetch, no gather op.where(cond, then, else)— per-lane select: the banded causal mask.guard(pred)— a chunk-loop condition: flash block-skip, fully-masked blocks never run.lNewreadinglSumoverb— recurrence groups: stages sharing a recurrence domain form one chunk loop, stage order = visibility, accumulators carried in registers.var rp = domain(4)— declaring an iterator the shape can't derive (the GQA head split).
The 128-block Walsh-Hadamard transform (7 butterfly levels, ±1 sign scales, 1/√128 norm)
is expressible today: each level is two statements with affine indices
(s_lo[i, a] = s_in[i*2, a] + s_in[i*2, 2, a]-shaped), the ±1 sign vectors and the norm are
ordinary elementwise stages. What that costs is 15 materialized stage buffers, so the
transform runs multi-pass, memory-bound — correct but not the in-place single-kernel shape
the CUDA kernel uses.
Missing to reach the production form:
- a pure time-loop recurrence: butterfly levels iterate a
lvlthat appears in neither the LHS nor any read — a plain loop that re-executes the stage bodies per level, no combine, no seed. The recurrence-group rule currently assumes chunk-reduce semantics; Lux's ownrecur(T)concept (the Gauss-Seidel challenge's time loop) is the missing generalization. Small and general, not Hadamard-specific. - exact twiddled (bit-reversal/XOR) indexing is out by design — the index grammar is
affine-only,
i ^ 2^kis rejected. For the 128-block transform this is immaterial: the untwiddled FWHT is the twiddled Hadamard up to a row permutation (identical spectrum, identical incoherence effect), which is why our libtorch port doesn't twiddle either. If an exact twiddle ever becomes load-bearing, the fix is the recorded per-lane gather (a permutation-table read), the same extension MoE row-gathers want.
Full code
test_luxtiles.nim
## Einsum section 2 (Rocktaschel 2018), twelve luxtiles recipes checked
## against scalar references on Metal.
##
## Run:
##
## nim c -r -d:release --outdir:build/wip/luxtiles \
## --nimcache:nimcache:wip/luxtiles --hints:off --warnings:off \
## workspace/experimental/tensorcery_luxtiles/test_luxtiles.nim # then run the binary
##
## Recipe notes:
##
## - luxtile statements read like the baseline kfunc recipes, `+=`
## accumulating and `=` defining
## - the domains are declared at each recipe's top, `var i = domain(n)`,
## the declaration order the loop order
## - iterators with derived ranges (the conv2d out dims, shift2x's j) stay implicit
##
## Schedules and layout:
##
## - schedules are naive where elementwise, the GEMM taking 16-tiles
## plus cache and the transpose 32-tiles per side
## - vectors live in rank-2 (n, 1) views, one read slot per axis 0,
## xposeNaive attaching its schedule Fn-level through sched(fn)
import
std/[math, random, sequtils, syncio],
workspace/experimental/tensorcery_luxtiles/luxtiles_compiler
# Declared domains open each recipe as `var i = domain(n)`, the loop
# order the declaration order. The reduce set is the statement's domains
# minus its LHS, never written. Iterator ranges are read-derived where
# no declaration binds them, for example the conv2d out dims plus
# shift2x's j, the derivation machinery their range source.
# 2.1 Matrix transpose as a blocked tile copy, 32-tiles per side
proc xpose(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
result[j, i] += a[i, j] do:
tile(j, 32)
tile(i, 32)
threads = 32
# untiled 2.1 for contrast, the elementwise tile program with strided
# column reads, its schedule attached Fn-level
proc xposeNaive(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
sched(result) do:
threads = 32
result[j, i] += a[i, j]
# 2.2 Sum over all indices, the cross-lane all-reduce
proc totSum(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
result[] += a[i, j]
# 2.3 Column sum, rows = j, the reduce chunked over the tile cols
proc colSum(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
result[j] += a[i, j]
# 2.4 Row sum
proc rowSum(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
result[i] += a[i, j]
# 2.5 Matrix-vector product
proc matvec(m, v: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
var k = domain(64)
result[i] += m[i, k] * v[k]
# 2.6 Matrix-matrix product, the working GEMM schedule (16-tiles + cache)
proc mm(m1, m2: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
var j = domain(64)
var k = domain(64)
result[i, j] += m1[i, k] * m2[k, j] do:
tile(i, 16)
tile(j, 16)
tile(k, 16)
cache(m1)
cache(m2)
threads = 32
# Composition, two statements in one function:
# - the second statement reads the first's target local `t1` under
# the name it was given
# - each statement derives its kernel at compile time
# - each statement launches at runtime
proc chain2(m1, m2: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
var j = domain(64)
result[j, i] += m1[i, j]
proc chain(m1, m2: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
var j = domain(64)
var k = domain(64)
t1[j, i] += m1[i, j]
result[i, j] += t1[i, k] * m2[k, j] do:
tile(i, 16)
tile(j, 16)
tile(k, 16)
cache(t1)
cache(m2)
threads = 32
# 2.7 Dot product
proc dotP(u, v: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(64)
result[] += u[i] * v[i]
# 2.8 Hadamard product
proc hadamard(a, b: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(3)
result[i, j] += a[i, j] * b[i, j]
# 2.9 Outer product
proc outerP(u, v: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(3)
var j = domain(4)
result[i, j] += u[i] * v[j]
# 2.10 Batch matrix multiplication, the batch dim a free out dim on grid.z
proc bmm(a3, b3: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(3)
var j = domain(2)
var l = domain(3)
var k = domain(5)
result[i, j, l] += a3[i, j, k] * b3[i, k, l] do:
tile(j, 16)
tile(l, 16)
tile(k, 16)
threads = 32
# 2.11 Tensor contraction, four out dims unranking a flat per-thread index
proc contract(ca, cb: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var k = domain(4)
var m = domain(7)
var o = domain(8)
var j = domain(3)
var l = domain(5)
result[i, k, m, o] += ca[i, j, k, l, o] * cb[m, j, l, o] do:
threads = 32
# 2.12 Bilinear transformation
proc bilinear(a, b3, c: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
var j = domain(5)
var k = domain(3)
var l = domain(7)
result[i, j] += a[i, k] * b3[j, k, l] * c[i, l] do:
tile(j, 16)
threads = 32
# 2.13 Read-map test, one read index a constant-times-iterator sum
# (2*j + 1), j staying implicit with its range deriving from the read axis
proc shift2x(a: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(2)
result[i, j] += a[i, j*2 + 1] do:
tile(i, 32)
tile(j, 32)
threads = 32
# 2.14 Convolution as implicit GEMM, spatial reads linear in the out dims
# (index forms like `oh*stride + kh - pad`), stride and pad schedule params.
# The derived out dims (oh, ow) and the per-site kernel axes (co, ci)
# stay implicit, the recipe running at two input shapes
proc conv2d(inp, wgt: LuxTensor): LuxTensor {.luxtile.} =
var kh = domain(3)
var kw = domain(3)
result[co, oh, ow] += inp[oh*stride + kh - pad, ow*stride + kw - pad, ci] *
wgt[kh, kw, ci, co] do:
param(stride, 2)
param(pad, 1)
tile(co, 16)
tile(ow, 16)
tile(ci, 16)
threads = 32
# ─────────────────────────────────────────────────────────────────────────
# References, plain scalar procs over seqs with 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 refConv2d(inp, wgt: seq[float32], h, w, c, kh, kw, co, s, p: int): seq[float32] =
## O[co, oh, ow] = sum_{kh, kw, ci} I[oh*s + kh - p, ow*s + kw - p, ci] *
## W[kh, kw, ci, co], zero padding, out range (h + p - 1) div s + 1
## per spatial axis.
let ohn = (h + p - 1) div s + 1
let own = (w + p - 1) div s + 1
result = newSeq[float32](co * ohn * own)
for coI in 0 ..< co:
for oy in 0 ..< ohn:
for ox in 0 ..< own:
var acc = 0.0'f32
for ky in 0 ..< kh:
for kx in 0 ..< kw:
for ci in 0 ..< c:
let iy = oy * s + ky - p
let ix = ox * s + kx - p
if iy >= 0 and iy < h and ix >= 0 and ix < w:
acc += inp[(iy * w + ix) * c + ci] *
wgt[((ky * kw + kx) * c + ci) * co + coI]
result[(coI * ohn + oy) * own + ox] = acc
proc refShift2x(a: seq[float32], m, n: int): seq[float32] =
## o[i, j] = a[i, 2j + 1], out width the derived range (n - 2) div 2 + 1.
let w = (n - 2) div 2 + 1
result = newSeq[float32](m * w)
for i in 0 ..< m:
for j in 0 ..< w:
result[i * w + j] = a[i * n + 2 * j + 1]
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
# ─────────────────────────────────────────────────────────────────────────
# Fixture runner, host fixtures and checks
# ─────────────────────────────────────────────────────────────────────────
when isMainModule:
proc flat(t: LuxTensor): seq[float32] =
t.view(float32).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"
proc main() =
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])
let a29H = randSeq(2 * 9)
let a29 = tensor(a29H, @[2, 9])
let cinpH = randSeq(9 * 9 * 4)
let cwgtH = randSeq(3 * 3 * 4 * 2)
let cinp = tensor(cinpH, @[9, 9, 4])
let cwgt = tensor(cwgtH, @[3, 3, 4, 2])
let cinp2H = randSeq(34 * 34 * 32)
let cwgt2H = randSeq(3 * 3 * 32 * 4)
let cinp2 = tensor(cinp2H, @[34, 34, 32])
let cwgt2 = tensor(cwgt2H, @[3, 3, 32, 4])
var engine = setBackend(bkMetal)
echo "=== einsum section 2, luxtiles stage IR 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); conv2d takes 3e-5"
echo "on the big case, its reduce space being 3-deep (2e-6 held small)"
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.6 chain2 ", engine.eval(chain2(m1_64, m2_64)).flat,
refXpose(m1_64H, 64, 64), 0.0'f32
check "2.6 chain ", engine.eval(chain(m1_64, m2_64)).flat,
refMatmul(refXpose(m1_64H, 64, 64), 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
check "2.13 shift2x ", engine.eval(shift2x(a29)).flat,
refShift2x(a29H, 2, 9), 0.0'f32
check "2.14 conv2d ", engine.eval(conv2d(cinp, cwgt)).flat,
refConv2d(cinpH, cwgtH, 9, 9, 4, 3, 3, 2, 2, 1), 2e-6'f32
check "2.14 conv2d big ", engine.eval(conv2d(cinp2, cwgt2)).flat,
refConv2d(cinp2H, cwgt2H, 34, 34, 32, 3, 3, 4, 2, 1), 3e-5'f32
echo "=== einsum section 2 green (luxtiles stage IR) ==="
main()test_luxtiles_softmax.nim
## Stable softmax and sdpa (online softmax over KV blocks) as luxtiles
## stage IR, checked against scalar references on Metal.
##
## Run:
##
## nim c -r -d:release --outdir:build/wip/luxtiles_softmax \
## --nimcache:nimcache:wip/luxtiles_softmax --hints:off --warnings:off \
## workspace/experimental/tensorcery_luxtiles/test_luxtiles_softmax.nim # then run the binary
##
## Shapes and tolerance:
## - softmax runs 16x32 float32, tolerance 1e-5
## - sdpa runs the 8x32-head geometry, q 8x32, k/v 64x32, eight 8-wide KV blocks, float32, tolerance 1e-5
## - exp lowers through the exp2 unit, stated exp2(x * 1.4426950408889634)
##
## Recipe notes:
##
## - the domains are declared at each recipe's top, `var i = domain(rows)`,
## the declaration order the loop order
## - the reducer is declared in the stage schedule, reduce(max) beside
## the default sum, a reduce(max) stage running the max chain
## - the KV-block iterator b appears only in the maxNew stage's reads
## (scores[i, 8*b + jc]), never on its LHS, Stage.recurrence holding it,
## the lowering running the serial block loop, acc in registers
##
## The max chain is idempotent, so the (b, jc) block/chunk pairs
## re-reading the same columns still yield the exact running max.
##
## Online correction notes:
## - the l and o accumulators sum over the block-chunked j, the online
## correction telescoping to the final-max form with dense scores
## - the correction factors are the expressiveness gap, stating them
## needs the fused single-kernel loop
import
std/[math, random, sequtils, syncio],
workspace/experimental/tensorcery_luxtiles/luxtiles_compiler
const Log2e = 1.4426950408889634'f32
# ─────────────────────────────────────────────────────────────────────────
# Stable softmax, four stages over one recipe
# ─────────────────────────────────────────────────────────────────────────
# maxRow[i] = max_j x[i, j] (max reducer)
# probs[i, j] = exp2((x[i, j] - maxRow[i]) * Log2e)
# (elementwise, maxRow broadcast)
# sumRow[i] = sum_j probs[i, j] (sum reduce, block-chunked)
# result[i, j] = probs[i, j] / sumRow[i]
# (elementwise, sumRow broadcast)
proc softmaxMvp(x: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(rows)
var j = domain(cols)
maxRow[i] = x[i, j] do:
reduce(max)
probs[i, j] = exp2((x[i, j] - maxRow[i]) * 1.4426950408889634) do:
tile(j, 16)
sumRow[i] = probs[i, j] do:
tile(j, 16)
result[i, j] = probs[i, j] / sumRow[i] do:
tile(j, 16)
# ─────────────────────────────────────────────────────────────────────────
# sdpa recipe, online softmax over the 8-wide KV blocks
# ─────────────────────────────────────────────────────────────────────────
# qScaled[i, d] = q[i, d] * (Log2e / sqrt(D)) scale fold
# scores[i, j] = sum_d qScaled[i, d] * k[j, d] the score tile
# maxNew[i] = max(mIn[i], scores[i, 8b + jc]) the running max,
# seeded by the -inf mIn tensor, the block iterator
# b carrying it
# maxOld[i] = maxNew[i] the prior chunk's max,
# read forward through the carried register
# rescale[i] = exp2(maxOld[i] - maxNew[i]) the online rescale factor
# probs[i, j] = exp2(scores[i, j] - maxNew[i]) the Log2e scale sits on qScaled
# sumPartial[i] = sum_j probs[i, j] block-chunked chain
# normNew[i] = normNew[i]*rescale[i] + sumPartial[i] the running normalizer
# seeded from an implicit zero, no lIn tensor
# outNew[i, dd] = outNew[i, dd]*rescale[i] +
# sum_j probs[i, j]*v[j, dd] the o accumulator
# carries itself the same way
# result[i, dd] = outNew[i, dd] / normNew[i]
proc sdpaMvp(q, k, v, mIn: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(qseq)
var j = domain(kv)
var d = domain(dHead)
var b = domain(8)
var jc = domain(8)
var dd = domain(dHead)
qScaled[i, d] = q[i, d] * 0.25503486164919736
scores[i, j] = qScaled[i, d] * k[j, d] do:
tile(j, 16)
tile(d, 16)
cache(qScaled)
cache(k)
threads = 32
maxNew[i] = max(mIn[i], scores[i, 8*b + jc]) do:
tile(jc, 8)
reduce(max)
maxOld[i] = maxNew[i]
rescale[i] = exp2(maxOld[i] - maxNew[i])
probs[i, j] = exp2(scores[i, j] - maxNew[i]) do:
tile(j, 8)
sumPartial[i] = probs[i, j] do:
tile(j, 8)
normNew[i] = normNew[i]*rescale[i] + sumPartial[i]
outNew[i, dd] = outNew[i, dd]*rescale[i] + probs[i, j] * v[j, dd] do:
tile(dd, 16)
tile(j, 16)
cache(probs)
cache(v)
threads = 32
result[i, dd] = outNew[i, dd] / normNew[i]
# ─────────────────────────────────────────────────────────────────────────
# Brittleness detector, a NEW 2-stage recurrence chain whose
# sum-of-squares to normalize pair fuses by the same rule
# with zero new lowering code
# ─────────────────────────────────────────────────────────────────────────
proc rmsnorm(x: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(rows)
var j = domain(cols)
squares[i, j] = x[i, j] * x[i, j] do:
tile(j, 16)
sumSquares[i] = squares[i, j] do:
tile(j, 16)
result[i, j] = x[i, j] / sumSquares[i] do:
tile(j, 16)
# ─────────────────────────────────────────────────────────────────────────
# References, plain scalar procs over seqs with row-major indexing
# ─────────────────────────────────────────────────────────────────────────
proc refSoftmax(x: seq[float32], m, n: int): seq[float32] =
## Row-wise stable softmax, exp over the shifted row.
result = newSeq[float32](m * n)
for i in 0 ..< m:
var mx = -3.402823466e38'f32
for j in 0 ..< n:
mx = max(mx, x[i * n + j])
var l = 0.0'f32
for j in 0 ..< n:
let e = exp(x[i * n + j] - mx)
result[i * n + j] = e
l += e
for j in 0 ..< n:
result[i * n + j] /= l
proc refSdpa(q, k, v: seq[float32], qs, kv, d: int): seq[float32] =
## o[i] = softmax(q_i K^T / sqrt(D)) V over the whole KV axis, the exact
## value that the block-chunked online recurrence telescopes into.
let scale = 1.0'f32 / sqrt(d.float32)
result = newSeq[float32](qs * d)
for i in 0 ..< qs:
var att = newSeq[float32](kv)
var mx = -3.402823466e38'f32
for j in 0 ..< kv:
var acc = 0.0'f32
for dd in 0 ..< d:
acc += q[i * d + dd] * k[j * d + dd]
att[j] = acc * scale
mx = max(mx, att[j])
var l = 0.0'f32
for j in 0 ..< kv:
att[j] = exp(att[j] - mx)
l += att[j]
for dd in 0 ..< d:
var acc = 0.0'f32
for j in 0 ..< kv:
acc += att[j] * v[j * d + dd]
result[i * d + dd] = acc / l
# ─────────────────────────────────────────────────────────────────────────
# Fixture runner, host fixtures and checks
# ─────────────────────────────────────────────────────────────────────────
when isMainModule:
proc flat(t: LuxTensor): seq[float32] =
t.view(float32).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"
proc main() =
var rng = initRand(41)
proc randSeq(n: int): seq[float32] =
result.newSeq(n)
for i in 0 ..< n:
result[i] = rng.gauss()
const (rows, cols, qseq, kv, dHead) = (16, 32, 8, 64, 32)
let xH = randSeq(rows * cols)
let x = tensor(xH, @[rows, cols])
let qH = randSeq(qseq * dHead)
let kH = randSeq(kv * dHead)
let vH = randSeq(kv * dHead)
let q = tensor(qH, @[qseq, dHead])
let k = tensor(kH, @[kv, dHead])
let v = tensor(vH, @[kv, dHead])
var mInH = newSeq[float32](qseq)
for i in 0 ..< qseq:
mInH[i] = -3.402823466e38'f32
let mIn = tensor(mInH, @[qseq, 1])
var engine = setBackend(bkMetal)
echo "=== stable softmax + sdpa, luxtiles stage IR vs scalar references ==="
echo "exp runs through the exp2 unit and the reductions reorder, 1e-5"
check "S.1 softmax ", engine.eval(softmaxMvp(x)).flat,
refSoftmax(xH, rows, cols), 1e-7'f32
check "S.2 sdpa ", engine.eval(sdpaMvp(q, k, v, mIn)).flat,
refSdpa(qH, kH, vH, qseq, kv, dHead), 1e-6'f32
check "S.3 rmsnorm ", engine.eval(rmsnorm(x)).flat,
(block:
var o = newSeq[float32](rows * cols)
for i in 0 ..< rows:
var acc = 0.0'f32
for j in 0 ..< cols: acc += xH[i * cols + j] * xH[i * cols + j]
for j in 0 ..< cols: o[i * cols + j] = xH[i * cols + j] / acc
o), 1e-5'f32
echo "=== softmax + sdpa green (luxtiles stage IR) ==="
main()test_luxtiles_paged.nim
## Paged-KV and atomic features as luxtiles recipes, checked against scalar
## references on Metal.
##
## Run:
##
## nim c -r -d:release --outdir:build/wip/luxtiles_paged \
## --nimcache:nimcache:wip/luxtiles_paged --hints:off --warnings:off \
## workspace/experimental/tensorcery_luxtiles/test_luxtiles_paged.nim # then run the binary
##
## Micro-tests:
##
## - P.1 paged KV gather, fetching the block table per threadgroup through
## a runtime scalar read in an index form (TensorRead), the tile read
## being the ordinary strided read from a [maxPages, pageSize, D] cache
## - P.2 atomic scatter histogram (MoE group counting), combining via
## uint-domain atomic add at the computed index, the count domain
## declared by the atomic directive, counts read back as uint bit patterns
## - P.3 banded causal mask, the per-lane mux select where(j <= i, x, -inf),
## with and without the stage's chunk guard skipping fully-masked column
## chunks (flash block-skip)
##
## Recorded gaps, asserted with clear messages at derivation:
##
## - a scalar bound to a tile-plane dimension needs a per-lane gather
## - a flat paged cache layout (form 8*page[b] + r over [P*8, D]) has no
## derivable within-page range, so the 3D cache layout binds r to an axis
import
std/[math, random, sequtils, syncio],
workspace/experimental/tensorcery_luxtiles/luxtiles_compiler
const Log2e = 1.4426950408889634'f32
const NegMax = -3.402823466e38'f32
# ─────────────────────────────────────────────────────────────────────────
# Recipes, one per micro-test
# ─────────────────────────────────────────────────────────────────────────
# P.1 the paged KV gather, page table indexed at the batch coordinate
proc pageGather(cache, page: LuxTensor): LuxTensor {.luxtile.} =
var b = domain(4)
var r = domain(8)
var d = domain(16)
result[b, r, d] = cache[page[b], r, d]
# P.2 the MoE group-count scatter, count[e] += 1 per token's expert
proc moeCount(topk: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(37)
result[topk[i]] += 1 do:
atomic(result, 4)
# P.3a softmax over the banded-masked scores
proc maskSoftmax(x: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(12)
var j = domain(16)
scores[i, j] = where(j <= i, x[i, j], -3.402823466e38)
maxRow[i] = scores[i, j] do:
reduce(max)
probs[i, j] = exp2((scores[i, j] - maxRow[i]) * 1.4426950408889634)
sumRow[i] = probs[i, j]
result[i, j] = probs[i, j] / sumRow[i]
# P.3b the same, column chunks fully past the mask boundary skipped
proc maskSoftmaxSkip(x: LuxTensor): LuxTensor {.luxtile.} =
var i = domain(12)
var j = domain(16)
scores[i, j] = where(j <= i, x[i, j], -3.402823466e38) do:
guard(j <= i + 7)
maxRow[i] = scores[i, j] do:
reduce(max)
probs[i, j] = exp2((scores[i, j] - maxRow[i]) * 1.4426950408889634)
sumRow[i] = probs[i, j]
result[i, j] = probs[i, j] / sumRow[i]
# P.4 causal paged GQA attention, decode + prefill rows in one batch.
#
# Layout:
#
# - q [Tpad, Hq, D], the GQA head map kv*Nr + rp exposed over
# the zero-copy view `q.view(T, Nkv, 4, D)`, alias reads using plain dims
# - per-token block tables packed flat into btL [Tpad*(PB+1)], row t
# holding PB page ids and the causal limit at column PB
# - k/v caches [maxPages, Nkv, 8*D] with the within-page (j, d) axes folded
# row-major, the read form jc*D + d
#
# Queries arrive pre-scaled by scale*log2e host-side, exp2(sc - mNew)
# being the softmax's exponential. Tolerance 1e-5 covers f32 exp2 over
# ~40 accumulated terms.
#
# Per-stage online softmax collapse, maxNew = max over the summed scores'
# (b, jc) cells, the row scale factor 1 with no prior chunk state,
# sumPartial and sumTotal plain sums over the attended blocks.
# Masked cells carry -1e7 through where() and exp2 underflows them to +0.0.
#
# Stage pipeline:
#
# - scores, the masked d-summed scores
# - maxNew, max over scores' (b, jc) cells reducing over the prior stage
# - probs, exp2(scores - maxNew)
#
# Remaining stages:
#
# - sumPartial reduces probs over jc and sumTotal over b
# - outSum contracts probs with the paged v cache, result dividing
# that sum by sumTotal
#
# maxNew's guard skips chunks fully past the band (flash block-skip).
# loopNest(jc, b) keeps b the chunk loop, the declaration order (b, jc)
# otherwise chunking jc while the guard binds a serial iterator
proc gqa(q, kCache, vCache, btL: LuxTensor): LuxTensor {.luxtile.} =
var t = domain(T)
var kv = domain(Nkv)
var rp = domain(4)
var d = domain(D)
var b = domain(PB)
var jc = domain(8)
let q4 = q.view(T, Nkv, 4, D)
scores[t, kv, b, jc, rp] = where(8*b + jc <= btL[t*5 + 4],
q4[t, kv, rp, d] * kCache[btL[t*5 + b], kv, jc*16 + d],
-1.0e7) do:
reduce(sum)
maxNew[t, kv, rp] = scores[t, kv, b, jc, rp] do:
loopNest(t, kv, rp, jc, b)
guard(8*b <= btL[t*5 + 4])
reduce(max)
probs[t, kv, b, jc, rp] = exp2(scores[t, kv, b, jc, rp] - maxNew[t, kv, rp])
sumPartial[t, kv, b, rp] = probs[t, kv, b, jc, rp] do:
reduce(sum)
sumTotal[t, kv, rp] = sumPartial[t, kv, b, rp] do:
reduce(sum)
outSum[t, kv, rp, d] = probs[t, kv, b, jc, rp] *
vCache[btL[t*5 + b], kv*128 + jc*16 + d] do:
tile(d, 16)
reduce(sum)
result[t, kv, rp, d] = outSum[t, kv, rp, d] / sumTotal[t, kv, rp]
# ─────────────────────────────────────────────────────────────────────────
# References, plain scalar procs over seqs
# ─────────────────────────────────────────────────────────────────────────
proc refMaskSoftmax(x: seq[float32], m, n: int): seq[float32] =
## Row-wise softmax over the banded-masked scores, masked cells contributing nothing.
result = newSeq[float32](m * n)
for i in 0 ..< m:
var mx = NegMax
for j in 0 ..< n:
if j <= i: mx = max(mx, x[i * n + j])
var l = 0.0'f32
for j in 0 ..< n:
var e = 0.0'f32
if j <= i: e = exp(x[i * n + j] - mx)
result[i * n + j] = e
l += e
for j in 0 ..< n:
result[i * n + j] /= l
proc refGqa(qH, kH, vH, btLH: seq[float32], T, Nkv, Nr, PB, D: int): seq[float32] =
## Causal paged GQA attention.
##
## - the packed block table's last column is the causal limit
## - the caches [maxPages, Nkv, 8*D] fold the (j, d) axes row-major
## - the query comes pre-scaled by scale*log2e
result = newSeq[float32](T * Nkv * Nr * D)
for t in 0 ..< T:
let band = btLH[t * (PB + 1) + PB].int
for kv in 0 ..< Nkv:
for rp in 0 ..< Nr:
let h = kv * Nr + rp
var m = -3.402823466e38'f32
for b in 0 ..< PB:
if 8 * b > band: continue
let pg = btLH[t * (PB + 1) + b].int
for jc in 0 ..< 8:
if 8 * b + jc > band: continue
for d in 0 ..< D:
m = max(m, qH[(t * (Nkv * Nr) + h) * D + d] *
kH[((pg * Nkv + kv) * 8 + jc) * D + d])
var l = 0.0'f32
var o = newSeq[float32](D)
for b in 0 ..< PB:
if 8 * b > band: continue
let pg = btLH[t * (PB + 1) + b].int
for jc in 0 ..< 8:
if 8 * b + jc > band: continue
var sc = 0.0'f32
for d in 0 ..< D:
sc += qH[(t * (Nkv * Nr) + h) * D + d] *
kH[((pg * Nkv + kv) * 8 + jc) * D + d]
let pw = exp2(sc - m)
l += pw
for d in 0 ..< D:
o[d] += pw * vH[((pg * Nkv + kv) * 8 + jc) * D + d]
for d in 0 ..< D:
result[((t * Nkv + kv) * Nr + rp) * D + d] = o[d] / l
# ─────────────────────────────────────────────────────────────────────────
# Fixture runner, host fixtures and checks
# ─────────────────────────────────────────────────────────────────────────
when isMainModule:
proc flatF(t: LuxTensor): seq[float32] =
t.view(float32).toSeq()
proc main() =
var rng = initRand(11)
proc randSeq(n: int): seq[float32] =
result.newSeq(n)
for i in 0 ..< n: result[i] = rng.gauss()
var engine = setBackend(bkMetal)
echo "=== paged gather, atomic scatter, banded mask (luxtiles) ==="
block pagedGather:
const (P, R, D, B) = (6, 8, 16, 4)
let cacheH = randSeq(P * R * D)
let pageH: seq[float32] = @[3.0'f32, 0.0, 5.0, 1.0]
let cache = tensor(cacheH, @[P, R, D])
let page = tensor(pageH, @[B, 1])
let got = flatF(engine.eval(pageGather(cache, page)))
for b in 0 ..< B:
let p = pageH[b].int
for r in 0 ..< R:
for d in 0 ..< D:
doAssert abs(got[(b * R + r) * D + d] -
cacheH[(p * R + r) * D + d]) < 1e-6'f32,
"paged gather mismatch at " & $b & "," & $r & "," & $d
echo "P.1 paged gather : PASS"
block atomicHistogram:
const (nTok, nExp) = (37, 4)
var topkH: seq[float32]
var want: array[nExp, int]
for i in 0 ..< nTok:
let e = rng.rand(0 ..< nExp)
topkH.add(e.float32)
inc want[e]
let topk = tensor(topkH, @[nTok, 1])
let o = engine.eval(moeCount(topk))
for e in 0 ..< nExp:
let bits = cast[uint32](o.view(float32)[e])
doAssert bits.uint64 == want[e].uint64, "histogram mismatch at " & $e
echo "P.2 atomic scatter: PASS"
block bandedMask:
const (M, N) = (12, 16)
var xH: seq[float32]
for i in 0 ..< M * N: xH.add(rng.gauss())
let x = tensor(xH, @[M, N])
let want = refMaskSoftmax(xH, M, N)
let oA = engine.eval(maskSoftmax(x))
let oB = engine.eval(maskSoftmaxSkip(x))
for pair in [(oA, "P.3 banded mask "), (oB, "P.3b mask + skip ")]:
let (o, name) = pair
var e1 = 0.0'f32
for i in 0 ..< M * N:
e1 = max(e1, abs(flatF(o)[i] - want[i]))
doAssert e1 <= 1e-6'f32, name & " mismatch " & $e1
echo name, ": max abs err ", e1, " PASS"
block gqaPaged:
const (T, D, Nkv, Nr, PB, Pmax) = (12, 16, 2, 4, 4, 6)
let scale = 0.25'f32
var qH: seq[float32]
for i in 0 ..< T * Nkv * Nr * D:
qH.add(rng.gauss() * scale * Log2e)
var kH: seq[float32]
for i in 0 ..< Pmax * Nkv * 8 * D: kH.add(rng.gauss())
var vH: seq[float32]
for i in 0 ..< Pmax * Nkv * 8 * D: vH.add(rng.gauss())
# per-token block table, packed [T, PB+1] with PB page ids and a causal limit,
# a cached length over ragged decode + prefill rows
var btLH: seq[float32]
for t in 0 ..< T:
let band = 3 + (t * 5) mod 20
for b in 0 ..< PB:
btLH.add(((t + b) mod Pmax).float32)
btLH.add(band.float32)
let q = tensor(qH, @[T, Nkv * Nr, D])
let kCache = tensor(kH, @[Pmax, Nkv, 8 * D])
let vCache = tensor(vH, @[Pmax, Nkv * 8 * D])
let btL = tensor(btLH, @[T * (PB + 1), 1])
let want = refGqa(qH, kH, vH, btLH, T, Nkv, Nr, PB, D)
let got = flatF(engine.eval(gqa(q, kCache, vCache, btL)))
doAssert got.len == want.len, "gqa output size " & $got.len &
" != " & $want.len
var e1 = 0.0'f32
for i in 0 ..< want.len:
e1 = max(e1, abs(got[i] - want[i]))
doAssert e1 <= 1e-5'f32, "P.4 gqa mismatch " & $e1
echo "P.4 gqa paged : max abs err ", e1, " PASS"
echo "=== paged + atomic + mask green (luxtiles stage IR) ==="
main()luxtiles_types.nim
import
std/[macros, tables],
workspace/data_structures/src/small_seqs,
workspace/crucible
# MetalEngine and its private fields are not on crucible's public
# export surface, the {.all.} import names them; the alias keeps the
# crucible `metal` macro unshadowed, the compiler exporting the macro,
# this module only feeding `TensorContext` and `setBackend`.
import workspace/crucible/src/runtime/engines/metal {.all.} as metalEngine
## luxtiles_types, the Lux statement IR with tiles as the value base.
##
## Structure mirrors lux_types.nim (laser lux_compiler) section for section,
## type names matching Lux where semantics match:
##
## Unary / Binary / Ternary kinds -> LuxNodeKind -> ScheduleKind
## -> Iter -> Invariant -> Fn -> Call -> Stage -> StageSchedule / FnSchedule
##
## Value base is tiles, scalars only in the index layer:
##
## - a stage-definition read `a[i, j]` yields a register tile
## shaped by the statement's tile plan
## - scalar expressions remain only in the index layer, read indices,
## domain bounds, schedule parameters
## - tensor op kinds sit at the end of LuxNodeKind, the lowered subset
## being the einsum grammar plus the tilewise maps (Add, Mul, reads)
type
# ###########################################
#
# Internal AST Representation
#
# ###########################################
UnaryOpKind* = enum
Ln
Exp
BinaryOpKind* = enum
# Must return a scalar for scalar expr check
Add
Mul
Sub # Index-layer subtraction, stride/pad forms like `oh*stride - pad`
Eq
Le # Index-layer comparison, the banded causal mask's condition
TernaryOpKind* = enum
# Must return a scalar for scalar expr check
Fma # FusedMultiplyAdd, the fused multiply-add op
Mux # Multiplexer / Selector for example max(x, 0) == mux(x>0, x, 0)
TileOpKind* = enum
## Tile op vocabulary over the statement's tile plan, lowered onto
## the ceramic tile atoms.
##
## Per-lane maps, axis reductions over tiles, plus `opMma`
## the tile contraction primitive.
opMul
opAdd
opAddScaled # fma, the scaled-accumulate tile op
opExp
opExp2
opSub
opDiv
opMax
opMux
opRowMax
opRowSum
opMma
ReduceKind* = enum
## Reducer of a stage's reduction, the parameter of the existing
## reduction lowering beside sum, not a new node.
rkSum
rkMax
LuxNodeKind* = enum
## Computation graph / abstract syntax tree nodes
## representing a Luxtiles computation.
# Order separates expression nodes from statement nodes.
# ################### High-level #########################
# Functions are the high-level concepts of Lux
Func # Everything is a function
# to promote composability
# ################### Expressions #########################
# ----- Kinds with no "children" field, scalar invariants
IntLit # Integer immediate (known at compile-time)
FloatLit # Float immediate (known at compile-time)
IntParam # Integer environment parameter (known at run-time, invariant during function execution)
FloatParam # Float environment parameter (known at run-time, invariant during function execution)
BoolParam # Bool environment parameter (known at run-time, invariant during function execution)
# Affine loop expression
Domain # Iteration domain, one loop index
# Scalar expressions built-ins
BinOpKind # Built-in binary operations
# ----- Kinds carrying the "children" field, expressions and statements
BinOp #
# Tensor/Function spatial indexing and properties
Access # Access a single tile of a Tensor/Func, the value-base read
DimSize # Get the size of one of the Tensor/Func dimensions
TensorRead # Runtime tensor read in index position, a scalar-term read
# (paged-KV block tables, MoE expert gathers)
# Function calls, external routines
ExternCall # Extern function call, like CPUInfo
# ################### Statements #########################
# Statements are generated when the high-level functional AST
# is lowered to an AST that more closely matches Nim's AST.
# General statements, lists of statements
StatementList # A list of statements, the general container
# Scalar statements, assignment and checks
Assign # Assign one value to one location
Check # Runtime check, an assertion over scalars
# Affine loop statements
AffineFor # Affine for loop
AffineIf # Affine if, a guarded stage application
# ################### Tile value layer #########################
# Tensor-typed expressions over the statement's tile plan, at the end
# so the Lux prefix stays byte-comparable.
#
# Index-layer leaves:
# IntLit, IntParam, Domain, BinOp, plus TensorRead for a runtime tensor
# read in index position, a read slot's scalar term of its IndexRef.
TileOp # Tile-typed op, kind in TileOpKind, operands in children
Id* = int
LuxNode* = ref object
## One computation-graph node, a variant over LuxNodeKind.
id*: Id
case kind*: LuxNodeKind
of IntLit:
intVal*: int
of FloatLit:
floatVal*: float
of IntParam, FloatParam:
symParam*: string
of Func:
fn*: Fn
of Domain:
# Iteration domain, its value during execution the loop index.
# Example, "for i in 0 ..< 10" is symbol "i" with start 0,
# stop 10 exclusive and step 1.
iter*: Iter
of BinOpKind:
bopKind*: BinaryOpKind
of TileOp:
tileOp*: TileOpKind
operands*: seq[LuxNode]
# Tile-typed operands, one per map argument or reduction source
else:
# Access and TensorRead carry
# [Func, index...] over their Fn's value
children*: seq[LuxNode]
ScheduleKind* = enum
ScReduce
ScParallel
ScVectorize
ScUnroll
ScStoreLoc
ScComputeLoc
ScOrder
ScPrefetch
ScPipeline # Interleaves N successive loops to avoid pipeline stalls
# and creates extra accumulators for reductions
# To design in Lux, wired here for the tile-program passes:
ScTile # Split a loop into a step-n outer loop and a range-n inner loop
ScCache # Stage an operand through threadgroup memory
ScParam # Bind a compile-time schedule parameter into index forms
ScThreads # Threadgroup thread count
ScAtomic # Atomic scatter stage, the combine a uint-domain atomic
# add at the computed index, range by the directive
ScRange # Fix an iterator's range (range(loop, n)), when no read binds it
ScGuard # Stage condition, a chunk-loop guard affine in the chunk
# iterator + invariants, the flash block-skip
# ###########################################
#
# High-level data structures
#
# ###########################################
Iter* = ref object
## One iterator symbol over a statement's iteration domain.
domId*: Id # Shouldn't be necessary
# But when passing to the backend
# The reference is lost
symbol*: string
start*, stop*, step*: LuxNode
InvariantKind* = enum
ikInt
ikFloat
RangeKind* = enum rkAxis, rkAffine, rkPin, rkDecl
# rkPin marks a range given by a schedule `range(loop, n)` directive
# or a literal `var i = domain(n)` declaration, for iterators no read
# binds directly, one inside a `j - 8*b` paged row form whose co-iter
# bounds the axis per cell, not globally
# rkDecl marks a declared domain's runtime range, `var t = domain(T)`
# holding T a host value at the eval site and the kernel carrying it
# as an int32 scalar argument
Range* = object
## An iterator's range:
## - an axis of a tensor's runtime shape
## - a derived output range when the iterator appears only inside
## a `2*i + 1`-shaped index over a read axis
##
## Derivation maximizes `iter` while keeping the raw index inside the read axis:
##
## range = ((dim - 1) - tmin) div coef + 1
case kind*: RangeKind
of rkAxis:
tensor*: string ## LuxTensor whose shape provides the range.
axis*: int ## Axis of that shape.
of rkAffine:
affineOf*: tuple[tensor: string, axis: int]
## Read axis and tensor the form binds.
coef*: int32 ## Positive coefficient of the iterator in the form.
tmin*: int32 ## Minimum of the form's non-coef terms at the serial iters' lower bounds.
of rkPin:
pin*: int32 ## Range value from a schedule `range` directive or a literal declaration.
of rkDecl:
hostNode*: NimNode ## Host expression of the declared range,
## resolved at the eval site.
## at the eval site (the splice scope).
argName*: string ## The kernel's int32 scalar formal carrying it.
Invariant* = ref object
## Runtime invariant input once a kernel is called.
kind*: InvariantKind
symbol*: string
FnOutputKind* = enum
fnScalar
fnTensor
Fn* = ref object
## A tensor function, one stage per redefinition of the value.
##
## - every stage definition and read evaluates on tiles, shaped
## by the statement's tile plan
## The first stage defines the value on the whole domain
## - a stage whose LHS indices do not mention an iterator recurs
## over that iterator, the accumulator-feedback form like
## `m[i] += max(m[i], row_max(s))` over the k domain
## - input tensors enter as Fn with no stages, reads gathering
## host buffer tiles
symbol*: string
stages*: seq[Stage]
schedule*: FnSchedule
viewOf*: string
## For a view alias, the base tensor's name, "" for a real tensor.
declOrder*: seq[string]
## Declared domains in declaration order, the recipe's loop order.
decls*: Table[string, Range]
## Declared domains' ranges by symbol, feeding every stage's IR,
## literal rkPin values and runtime rkDecl values alike.
outputKind*: FnOutputKind
# ###########################################
#
# Low-level data structures
#
# ###########################################
Call* = ref object
## Indexing of a Fn, `A[i, j]` or `A[i, j] = expression`.
## Ref object, living long enough for the assignation.
fn*: Fn
params*: seq[LuxNode]
Stage* = object
## One definition of a Fn's value.
definition*: LuxNode
params*: seq[LuxNode]
# Domain or specific location where this stage applies
recurrence*: seq[LuxNode]
# A stage repeated over a non-spatial domain, which the LHS
# indices and the definition never mention.
# Example, a Gauss-Seidel smoother over the time domain t:
# ------------------------------------
# for t in 0 ..< timeSteps:
# for i in 1 ..< N-1:
condition*: LuxNode
# Guard of the stage's application, affine in iterators + invariants,
# lowering to tile-level validity over the read axes
# (page-attention-style masks use the field later)
reduce*: ReduceKind
# Reducer of the stage's reduction, sum by default, max chosen
# for running-max stages (declared `reduce(max)` in the schedule).
schedule*: StageSchedule
# Transformations for this stage's lowering, in written order.
# An empty schedule uses the Fn-level defaults instead.
iterOrder*: seq[string]
# Stage iterators in loop order, the declared domains first,
# declaration order, the implicit ones after, first-use order.
ScheduleDirective* = object
## One schedule transformation, a pass factory at lowering time.
##
## - name, the loop / tensor / parameter name for the kind
## - value, the tile size, parameter value or thread count
## - names, the loop order for the loopNest directive (ScOrder)
kind*: ScheduleKind
name*: string
value*: int32
names*: seq[string]
exprNode*: NimNode ## Guard predicate, ScGuard directives only.
StageSchedule* = object
## One stage's schedule transformations, composed in written order.
directives*: seq[ScheduleDirective]
FnSchedule* = object
## Fn-level default transformations for stages with no schedule
## directives written for the Fn.
directives*: seq[ScheduleDirective]
# ###########################################
#
# Tile op vocabulary, reducers, schedule directives
#
# ###########################################
# Appended types beyond the Lux-mirrored span, the tile vocabulary, the reducers,
# the schedule directives, and the lowering records moved in from the compiler and pass modules.
## Tile layer additions over Lux, each born and consumed here:
##
## TileOpKind, born in dsl valueExpr (exp2 / exp / max calls)
## | valueFromLux -> ExprIR tilewise-map kinds
## v emitTileValue -> ceramic per-lane tile atoms -> MSL
##
## ReduceKind, born in dsl parseSchedule's reduce(sum | max)
## | appendStage -> Stage.reduce
## v stageToIR -> AlgoIR.reduceKind -> max chain or row_max epilogue
##
## ScheduleDirective, one per trailing do-block statement, written
## | StageSchedule / FnSchedule carry them in written order
## v buildPasses -> one pass factory per directive -> runPipeline
# ###########################################
#
# Lowering IR, stage parse to kernel program
#
# ###########################################
# One stage's lowering records, born and consumed at macro expansion,
# nothing surviving past the (device fn, thin entry) pair.
#
# stageToIR schedule passes mmaOptPass lowerEmit
# | parse | tile / cache / | hasMma + | plan or
# v v threads, in place v mmaMatch v match
# AlgoIR + reads --> annotated AlgoIR ----> annotated AlgoIR --> device fn
# + thin entry
type
# Read index slot, born in indexFromLux, consumed by the gather
# and range derivations, dying with the stage's lowering.
IdxKind* = enum ikIter, ikConst, ikAffine
## One read index slot:
## - an iterator name
## - an integer constant
## - a constant-times-iterator sum like `3*i + k - 1`
## (`i * j` as an index is rejected, only constant-times-iterator sums work)
# indexFromLux -> IndexRef (one per read slot)
# | iterRange / hasBinding / slotIterCoef read it
# v emitGather turns it into the gd view geometry
ScalarTerm* = ref object
## One runtime tensor read inside an index form, its loaded value
## weighted into the read slot's offset:
##
## ScalarTerm, born in dsl indexExpr (a bracket expr in index position)
## | indexFromLux -> IndexRef.scalars
## v emitGather loads it per tile origin -> base offset
tensor*: string
coef*: int32 ## Weight of the loaded value in the slot's offset.
indices*: seq[IndexRef] ## The read's own index slots, affine forms over iterators.
IndexRef* = object
case kind*: IdxKind
of ikIter: name*: string
of ikConst: value*: int32
of ikAffine:
terms*: seq[tuple[coef: int32, iter: string]]
## One (coefficient, iterator) pair per iterator the form binds.
cst*: int32
## Constant term, may be negative (a padding offset).
scalars*: seq[ScalarTerm]
## Runtime tensor reads in the form for every kind, an empty seq
## meaning the affine-only form.
AffineForm* = object
## Parse-time accumulator for one index expression.
## `parseIndex` normalizes it into an `IndexRef`.
coefs*: seq[tuple[coef: int32, iter: string]]
cst*: int32
scalars*: seq[ScalarTerm]
## Runtime tensor reads in the form, coefficient-scaled by affScale.
# valueFromLux -> ReadIR, one per read in expression order
# | makeTilePlan maps its slots onto the tile plane
# v emitGather gathers it, dead after the kernel body emits
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).
# Expression tree, born in valueFromLux from the stage definition.
# valueFromLux -> ExprIR
# | emitTileValue walks it into per-lane tile atoms
# v the combine epilogue consumes the root tile
ExprKind* = enum xeRead, xeAdd, xeMul, xeSub, xeDiv, xeExp2, xeMax, xeLit, xeMux
ExprIR* = ref object
## Expression tree over reads, `+` and `*` plus the tilewise maps
## sub / div / exp2 / max, float literals. Unary exp2 leaves `rhs` nil.
case kind*: ExprKind
of xeRead: read*: ReadIR
of xeAdd, xeMul, xeSub, xeDiv, xeExp2, xeMax, xeMux:
lhs*, rhs*: ExprIR
cond*: LuxNode
## Mux condition, an index-layer tree over iterators and invariants,
## nil on every other kind.
of xeLit: value*: float32
# Loop nest, born in buildLoops from the out dims and reduce iters.
# buildLoops -> LoopIR list, outer -> inner
# | tilePass splits a loop into outer step / inner range
# v emitTileKernel wraps the for loops around the body
LoopRole* = enum lrOut, lrReduce
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"`.
range*: Range ## For an outer loop, the logical range it iterates.
tile*: int ## 0 = untiled.
innerOf*: string ## "" for outer loops, the parent loop for a tile split.
role*: LoopRole
# Staged-buffer plan, born in cachePlanPass from the schedule tiles.
# cacheAnnotatePass marks tensors -> cachePlanPass sizes buffers
# | smemDecl + the mma staged emitter stage them
# v consumed as threadgroup-memory tiles per atom slab
CachePlan* = object
tensor*: string
buffer*: seq[tuple[iter: string, tile: int]]
## Staged buffer dims, one per tiled iterator appearing in the tensor's reads:
##
## - an untiled out dim contributes none, its coordinate coming from grid.z or a serial loop, never from a buffer dim
# Matched mma shape, born in the mma pass's matchMma,
# from the schedule-tiled dims and the two operand reads' index slots.
# matchMma -> MmaMatch (ra == -1 means no match)
# | isStagedEligible picks staged vs register
# v emitMmaStaged / emitMmaRegister consume it in lowerEmit
MmaMatch* = object
## Matched mma shape, see the mma pass module (`luxtiles_passes_mma`).
##
## - a zero value (`ra == -1`) means no match
ra*, rb*: int = -1
## Read indices of the two operands, `ra` carries the accumulator's rows, `rb` its cols, -1 without a match.
rowDim*, colDim*: string
## Out dims carried by `ra`/`rb`, the accumulator's rows and cols.
freeDims*: seq[string]
## Remaining out dims in IR order. The first two split grid.z between
## them (z div / z mod the ranges), the rest become serial for-loops.
chunkIter*: string
## Reduce iterator tiled along the tile plane's k axis.
serialIters*: seq[string]
## Remaining reduce iters, serial loops around the chunk loop whose
## coordinates are added to the reads' base offsets.
tiles*: tuple[m, n, k: int]
## Schedule tiles over rowDim, colDim and the chunk iter.
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, the engine shim'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`.
hasMma*: bool ## Set by the mma pass on a canonical mma shape, no match keeping the generic tile program.
mmaMatch*: MmaMatch ## The recognized shape, valid only with `hasMma`.
reduceKind*: ReduceKind ## The stage's reducer, max running the max chain and row_max epilogue.
scatter*: bool ## Atomic scatter stage, the combine a uint-domain atomic add.
scatterValue*: int32 ## The count the scatter adds (literal definition).
scatterRange*: int32 ## The output's range, declared by the atomic directive.
recurrence*: seq[string] ## Iters the stage recurs over (absent from its LHS), recorded for the pass trace.
guardNode*: LuxNode
## Chunk-loop condition (Stage.condition), affine in the chunk
## iterator + invariants, nil without a guard directive.
threads*: int
rangePins*: Table[string, int32]
## Literal range values, the schedule `range(loop, n)` directive
## overriding a declaration's literal, iterRange consulting them
## before the read derivation.
declRanges*: Table[string, Range]
## Declared domains' runtime ranges (rkDecl), iterRange consulting
## them after the literal values and before the read derivation.
declArgs*: seq[tuple[argName: string, hostNode: NimNode]]
## Declared ranges' scalar kernel args. The entry passes them after
## the tensor views, the prologue binding each hostNode to its name.
# Pass framework, a schedule is a seq[Pass] run in written order
# Transform and analysis passes rewrite or annotate the IR in place.
# buildPasses (one Pass per schedule directive) -> runPipeline
# | run(ir) transforms / annotates, trace in KernelTrace
# v emit(ir) lowers, the device fn + thin entry node pair
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`.
# Tile plan, the generic tile program's geometry, derived once per
# stage from the IR's out dims, reduce space and schedule tiles.
# makeTilePlan (AlgoIR + schedule) -> TilePlan
# | emitGather, per-read gd view + loadTileBounded geometry
# v emitTileKernel consumes it, the generic tile program
TilePlan* = object
## Per-read view geometry and per-axis roles of one tile program.
rowDim*, colDim*: string ## Out dims mapped to the tile plane, for an all-reduce the row/col chunk iters.
tileR*, tileC*: int ## Tile plane ranges, 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 (added to 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.
# ###########################################
#
# Fusion planning records
#
# ###########################################
# Fusion groups, the fusion pass's records, born in the recipe macro's
# statement collection and consumed by the group lowering:
#
# StageRef list (one per recipe statement)
# | planFusionGroups (luxtiles_passes_fusion)
# v FusionGroup, kind + stage indices
# | fkChain (one recurrence domain) -> one device kernel + entry,
# | the generic multi-stage tile program
# v fkNone -> per-stage lowering (luxtiles_compiler's lowerStage)
type
StageRef* = ref object
## One recipe statement's lowering record, the fusion planner's input.
name*: string ## Target name, the Fn's symbol.
fn*: Fn ## The statement's Fn, `stageIdx` its one lowered stage.
sched*: StageSchedule ## The statement's stage schedule.
stageIdx*: int ## Index of the stage this statement added.
FusionKind* = enum
## Lowering kind of a fusion group.
fkNone ## No fusion, every stage its own kernel launch.
fkChain
## Stages over one recurrence domain, one kernel, a chunk loop per
## visibility pass executing the stages in stage order per chunk.
FusionGroup* = object
## One stage group, the statement indices it covers and the lowering kind.
stages*: seq[int] ## Statement indices, in recipe order.
kind*: FusionKind
description*: string ## Lowering description for the pass trace.
# Engine shim, host LuxTensor, Metal TensorContext
# ─────────────────────────────────────────────────────────────────────────
# Runtime surface the compiler lowers onto, `HostElem` naming the element
# dtype of every kernel.
# - `LuxTensor`, a device buffer behind row-major metadata
# - `TensorContext`, one Metal engine owning the compiled-unit cache
#
# `setBackend` accepts a `BackendKind`, Metal is the only wired backend.
#
# Engine-shim types beyond the Lux-mirrored span, created, consumed,
# dying across the recipe lifecycle:
#
# TensorRef / TensorImpl[T] / Metadata (a SmallSeq shape-stride pair)
# | allocImpl -> the device buffer behind a LuxTensor
# v tensor / makeHostTensor / view / operandArg
#
# LuxTensor, the host handle held by recipe formals and eval results
# | engine.eval launch -> run(backend, kernel, PtrArg, LaunchConfig)
# v result binding, dropped with the statement expansion
#
# TensorContext, one live Metal engine, recipes' launch site
# | setBackend -> engine.eval(...)
# v program end
#
import workspace/crucible/src/runtime/engines/metal {.all.} as metalEngine
# Engine shim, host LuxTensor, Metal TensorContext
# ─────────────────────────────────────────────────────────────────────────
# Runtime surface the compiler lowers onto, `HostElem` naming the element
# dtype of every kernel.
# - `LuxTensor`, a device buffer behind row-major metadata
# - `TensorContext`, one Metal engine owning the compiled-unit cache
#
# `setBackend` accepts a `BackendKind`, Metal is the only wired backend.
proc posix_memalign(memptr: ptr pointer; alignment, size: csize_t): cint
{.importc: "posix_memalign", header: "<stdlib.h>".}
type
TensorRef* = ref object of RootObj
## Type-erased device allocation behind a host LuxTensor.
Metadata = SmallSeq[4, int32]
## Axis ranges or strides of a host tensor, up to rank 4 inline.
TensorImpl*[T] = ref object of TensorRef
buf*: ptr UncheckedArray[T]
shape*: Metadata
strides*: Metadata
LuxTensor* = object
## Host tensor of the module, a device buffer behind row-major
## metadata held by the engine shim.
shape*: Metadata
strides*: Metadata
impl: TensorRef
TensorContext* = object
## One Metal engine context, the recipes' launch site.
metal*: MetalEngine
proc setBackend*(hw: BackendKind): TensorContext =
## Live engine context for the backend `hw`.
case hw
of bkMetal:
TensorContext(metal: bkMetal.init())
else:
raiseAssert "setBackend: backend not implemented"
func numel*(shape: Metadata): int32 =
## Element count over a shape's axis ranges.
result = 1
for n in shape.items:
result *= n
func numel*(t: LuxTensor): int32 = numel(t.shape)
proc toMetadata*(vals: openArray[int32]): Metadata =
## Metadata built from a literal axis list, spilling past the inline capacity.
result = Metadata.init()
for v in vals:
result.add v
proc rowMajorStrides(shape: Metadata): Metadata =
## Row-major strides for any rank.
## Stride 1 on the last axis, each earlier axis stepping by the product of the ranges to its right.
var st: seq[int32]
st.newSeq(shape.len)
var acc: int32 = 1
for ax in countdown(shape.len - 1, 0):
st[ax] = acc
acc *= shape[ax]
st.toMetadata()
proc allocImpl*[T](shape: Metadata): TensorImpl[T] =
## Page-aligned, zero-filled device allocation sized for `shape`.
let sh = shape
let n = numel(sh)
let pageSize = 16384
let bytes = (n.int * sizeof(T) + pageSize - 1) div pageSize * pageSize
var raw: pointer
doAssert posix_memalign(addr raw, pageSize.csize_t, bytes.csize_t) == 0,
"posix_memalign failed"
zeroMem(raw, bytes)
TensorImpl[T](buf: cast[ptr UncheckedArray[T]](raw), shape: sh,
strides: rowMajorStrides(sh))
proc tensor*(dataArr: seq[float32], shape: openArray[int]): LuxTensor =
## Host tensor built from row-major data, uploaded to the device.
var sh32: seq[int32]
for v in shape.items:
sh32.add v.int32
let sh = sh32.toMetadata()
doAssert dataArr.len == numel(sh).int
let impl = allocImpl[float32](sh)
copyMem(impl.buf, unsafeAddr dataArr[0], numel(sh).int * sizeof(float32))
LuxTensor(shape: sh, strides: rowMajorStrides(sh), impl: impl)
template view(impl: TensorImpl): untyped =
## Flat openArray over a device allocation.
impl.buf.toOpenArray(0, numel(impl.shape).int - 1)
template view*(t: LuxTensor, T: typedesc): untyped =
## Typed device view of a host tensor's flat elements.
doAssert t.impl != nil
TensorImpl[T](t.impl).view()
func operandArg*(t: LuxTensor): PtrArg[float32] =
## Marshalling blob for a host tensor operand, its device buffer pointer.
doAssert t.impl != nil
let impl = TensorImpl[float32](t.impl)
PtrArg[float32](buf: impl.buf, len: numel(impl.shape).int, off: 0)
proc makeHostTensor*(shape, strides: Metadata; impl: TensorRef): LuxTensor =
## Host tensor constructor the `eval` expansion's scope-resolved idents
## find here, named apart from the ceramic `LuxTensor`.
LuxTensor(shape: shape, strides: strides, impl: impl)
proc view*(t: LuxTensor, shape: varargs[int]): LuxTensor =
## Zero-copy row-major alias of `t` reinterpreted over `shape`.
##
## Args:
## - t, the base tensor
## - shape, the new axes, the element counts matching t.numel
##
## Returns the alias, its element buffer shared with t, reads using
## plain dims where an affine form like `kv*4 + rp` would be needed.
##
## let q4 = q.view(T, Nkv, 4, D)
var sh: seq[int32]
for d in shape.items:
sh.add(d.int32)
let shM = sh.toMetadata()
doAssert numel(shM) == numel(t),
"view: the new axes' element count " & $numel(shM) &
" does not match the tensor's " & $numel(t)
makeHostTensor(shM, rowMajorStrides(shM), t.impl)
# Element dtype every kernel derives from, the buffer element type
# the engine shim's marshalling (`operandArg`) hands the device.
type HostElem* = typeof(operandArg(default(LuxTensor)).buf[][0])
static:
doAssert HostElem is float32, "substrate elem type should be float32"
# Statement parse record of the luxtile macro's per-statement walk, macro-internal:
#
# luxtile statement walk -> parseStmt -> StmtParts
# | target / value -> indexExpr / valueExpr
# v appendStage, dies after the stage appends
type
StmtParts* = object
## One recipe statement's raw AST parts, macro-internal.
target*: NimNode ## LHS name[d0, d1, ...], the brackets node
accumulate*: bool ## += stage append, = stage definition
value*: NimNode ## the definition expression
schedBlk*: NimNode ## the trailing do block or nil
# Mma emitter geometry record for the staged/register emitters' affine slot:
#
# matchMma + reads -> PlaneGeometry (mma pass) -> loadTile* origin nodes
# v dies with the emitter's statement list
type
PlaneGeometry* = tuple[stride, origin, limit, baseExtra: NimNode,
affine: bool]
## Stride, origin, limit and base offset of one read axis. `baseExtra`
## holds the affine part, added to the view's base offset.
# Fusion-resolution records for the group planner's identity map and classes:
#
# resolveAxes (statement reads) -> PhysAxis table, one per iterator
# | chunkPlan / classifyStage pick the chunk domain and classes
# v GroupFacts, the emitter's lowering input
type
PhysAxis* = tuple[tensor: string, axis: int]
## One physical tensor axis, the identity of a logical dimension.
StageClass* = enum
scFold, scElemTile, scElemRow, scPriorCopy, scElemRowCarry,
scSweepReduce, scChunkReduce, scChunkReduceFused, scMmaFresh, scMmaCarry
GroupFacts* = ref object
## Lowered form's inputs, per-stage classes, tiles and roles.
stages*: seq[int] ## Statement indices, recipe order.
irs*: seq[AlgoIR]
names*: seq[string]
classes*: seq[StageClass]
passes*: seq[int]
rowDim*: string ## Shared row out dim.
rowAxis*: PhysAxis ## Row dim's physical axis.
axes*: Table[string, PhysAxis] ## Iterators resolving to physical axes.
decls*: Table[string, Range] ## Declared domains among the stages' iters.
tileR*: int ## Shared row tile.
cs*: int ## Chunk size, the chunk domain's chunk width.
dAxis*: PhysAxis ## Chunk domain's physical axis.
dIters*: seq[string] ## Iterators resolving to the chunk domain.
extRank*: seq[tuple[tensor: string, rank: int]]
## External inputs of the group kernel, first-appearance order.
readsOf*: seq[seq[int]] ## Per stage, indices of the stages reading its target.
tileC*: seq[int] ## Per stage, the column tile, 0 one-dimensional.
stageBuf*: seq[string] ## Per stage, the value buffer, "" until emitted.
priorOf*: seq[int] ## Per scPriorCopy stage, the snapshot producer.
mmaAIdx*, mmaBIdx*: seq[int] ## Per mma stage, the A/B read slots.
mmaScaleIdx*: seq[int] ## Per scMmaCarry stage, the rescale rv's stage, -1 scalar.
mmaScaleLit*: seq[float32] ## Per scMmaCarry stage, the rescale literal.
carryScaleIdx*: seq[int] ## Per scElemRowCarry stage, the scale rv's stage, -1 scalar.
carryScaleLit*: seq[float32] ## Per scElemRowCarry stage, the scale literal.
fusedSrc*: seq[int] ## Per scElemRowCarry stage, the fused reduce stage, -1 none.
convIter*: seq[string] ## Per mma stage, the contraction iterator.
convTile*: seq[int] ## Per mma stage, the contraction tile.
foldBase*: seq[string] ## Per scFold stage, the base read's tensor.
foldScale*: seq[float32] ## Per scFold stage, the literal scale.
srcExpr*: seq[ExprIR] ## Per reduce stage, the per-chunk source value.
seedOf*: seq[seq[string]] ## Per stage, the seed reads' tensors.
needsChunk*: seq[bool] ## Per stage, recomputed per chunk.
hasCarry*: bool ## Any self-carrying stage, sweep reduces
## then read as running values, 3-arg row ops.
finalIsChunked*: bool ## Final stage recomputes per chunk, storing per chunk.
threads*: int ## Threadgroup width from the stage schedules.
groupCarries*: bool ## Any carry stage in the group.
# Group-kernel emission state of the fusion emitter's running builder:
#
# lowerGroup -> Emitter (body, passBody, counter)
# | emitStageValue / valTile / valRow append fragments
# v emitDeviceFunc + emitEntry consume the body, Emitter dies
type
Emitter* = object
## One group kernel's emission state.
gf*: GroupFacts
body*: NimNode ## Whole body under construction.
passBody*: NimNode ## The current pass's chunk body.
counter*: int ## Fresh temporary names.luxtiles_dsl.nim
## luxtiles_dsl, the authoring layer turning Nim AST into Lux IR.
##
## primitives.nim analog of the lux_dsl API. The `luxtile` macro living
## in luxtiles_compiler walks a recipe proc's statements and drives
## these builders:
##
## - one Fn per target name
## - one Stage per statement
##
## var i = domain(M)
## result[i, j] += m1[i, k] * m2[k, j] do:
## tile(i, 16) # the stage schedule
##
## | form | contract |
## | ----------------------------- | ------------------------------------------------------------------------------------------------- |
## | `var i = domain(n)` | a declared domain, declaration order the loop order, n a literal or a host value at the eval site |
## | `let v = t.view(d0, d1, ...)` | a zero-copy axis rename of a parameter |
## | `A[i, j] = expr` | stage 0 of Fn A, the full-domain definition |
## | `A[i, j] += expr` | the next stage on the same Fn |
## | `do:` trailing block | the stage schedule, `loopNest(a, b, ...)` in it the non-default loop order |
## | `sched(A) do: tile(i, 16)` | the Fn-level default schedule |
##
## The reduce set of a statement is its domains minus its LHS
## iterators and stays derived, never written.
import
std/[macros, tables],
luxtiles_types
var nodeIdCounter {.compileTime.}: int
proc newLux(kind: LuxNodeKind): LuxNode =
## Fresh node of `kind`, children filled by the builders.
inc nodeIdCounter
LuxNode(kind: kind, id: nodeIdCounter)
proc newLuxInt*(v: int): LuxNode =
## Integer immediate node.
var n = newLux(IntLit)
n.intVal = v
n
proc newLuxFloat*(v: float): LuxNode =
## Float immediate node.
var n = newLux(FloatLit)
n.floatVal = v
n
proc newLuxIntParam*(sym: string): LuxNode =
## Schedule parameter node, resolved to a value by the pass pipeline.
var n = newLux(IntParam)
n.symParam = sym
n
proc newFn*(symbol: string): Fn =
## Tensor function with no stages yet, an input tensor until its first
## statement defines a stage.
Fn(symbol: symbol, outputKind: fnTensor)
proc newIter*(symbol: string): Iter =
## Iterator symbol over one statement's domain, bounds deriving
## from the reads at lowering time, start 0 and unit step always.
inc nodeIdCounter
Iter(domId: nodeIdCounter, symbol: symbol,
start: newLuxInt(0), stop: newLuxInt(0), step: newLuxInt(1))
proc domainNode(iter: Iter): LuxNode =
## Index-layer node carrying an Iter.
var n = newLux(Domain)
n.iter = iter
n
func isIterator*(n: LuxNode): bool =
## Whether the node carries an Iter (a Domain).
n.kind == Domain
proc newCall*(fn: Fn, indices: seq[LuxNode]): Call =
## Tensor read at `indices`, one index slot per tensor axis.
Call(fn: fn, params: indices)
proc accessNode*(call: Call): LuxNode =
## Tile-valued read, Access(fn, indices...), the expression leaf.
var n = newLux(Access)
var f = newLux(Func)
f.fn = call.fn
n.children.add(f)
for p in call.params:
n.children.add(p)
n
proc binOpNode*(op: BinaryOpKind, a, b: LuxNode): LuxNode =
## Binary op node, children[0] carrying the op kind.
var n = newLux(BinOp)
var k = newLux(BinOpKind)
k.bopKind = op
n.children = @[k, a, b]
n
proc `+`*(a, b: LuxNode): LuxNode =
## Tile-typed or index-typed sum.
binOpNode(Add, a, b)
proc `*`*(a, b: LuxNode): LuxNode =
## Tile-typed or index-typed product.
binOpNode(Mul, a, b)
proc tileOpNode*(op: TileOpKind, operands: seq[LuxNode]): LuxNode =
## Tile-typed op node, kind and per-lane operands or reduction source.
var n = newLux(TileOp)
n.tileOp = op
n.operands = operands
n
proc `~`*(a: LuxNode): LuxNode =
## Index-layer negation of a form.
var z = newLuxInt(0)
result = binOpNode(Sub, z, a)
proc indexExpr*(n: NimNode, tensors: Table[string, Fn], iters: var Table[string, Iter], params: Table[string, int32]): LuxNode
proc condExpr*(n: NimNode, iters: Table[string, Iter], params: Table[string, int32],
tensors: Table[string, Fn] = initTable[string, Fn]()): LuxNode =
## Index-layer condition node for stage guards and per-lane muxes:
##
## - +- and * forms over iterators, literals and schedule parameters
## - comparisons <, <=, ==, the comparison the tree root or a subtree
## - a plain form valid too, truthy when nonzero
# walk the condition node tree:
proc walk(n: NimNode): LuxNode =
case n.kind
of nnkIntLit:
result = newLuxInt(n.intVal)
of nnkFloatLit:
result = newLuxFloat(n.floatVal)
of nnkIdent, nnkSym:
let s = n.strVal
doAssert s in iters or s in params,
"luxtiles: condition symbol '" & s & "' is neither an iterator " &
"nor a bound schedule parameter"
if s in iters:
var d = newLux(Domain)
d.iter = iters[s]
result = d
else:
result = newLuxIntParam(s)
of nnkPar:
doAssert n.len == 1, "unsupported parenthesized condition"
result = walk(n[0])
of nnkBracketExpr:
# a runtime tensor read in condition position, a scalar predicate term
# (the banded mask's per-token limit)
let sym = n[0].strVal
doAssert sym in tensors,
"luxtiles: condition scalar read of '" & sym & "' is neither a " &
"recipe parameter nor an earlier Fn target"
var idxs: seq[LuxNode] = @[]
var itersCopy = iters
for i in 1 ..< n.len:
idxs.add(indexExpr(n[i], tensors, itersCopy, params))
result = newLux(TensorRead)
result.children.add(newLux(Func))
result.children[0].fn = tensors[sym]
for ix in idxs:
result.children.add(ix)
of nnkInfix:
let op = n[0].strVal
case op
of "<=", "<", "==":
let k = (case op
of "<=": Le
of "<": Sub # a < b lowers to the truthy subtraction, b - a > 0
else: Eq)
if op == "<":
result = binOpNode(Sub, walk(n[2]), walk(n[1]))
else:
result = binOpNode(k, walk(n[1]), walk(n[2]))
of "+", "-":
let l = walk(n[1])
let r = walk(n[2])
result = (if op == "+": l + r else: l + (~r))
of "*":
result = walk(n[1]) * walk(n[2])
else:
error("unsupported condition operator " & op, n)
else:
error("unsupported condition node " & $n.kind, n)
walk(n)
# ─────────────────────────────────────────────────────────────────────────
# Nim AST -> LuxNode builders
# ─────────────────────────────────────────────────────────────────────────
# `iters` binds iterator symbols, created on first use and shared
# across the proc's statements.
# `params` carries the schedule parameter values.
# Everything else in index position is a constant or an error.
proc indexExpr*(n: NimNode, tensors: Table[string, Fn], iters: var Table[string, Iter], params: Table[string, int32]): LuxNode =
## Returns one read or LHS index as an index-layer LuxNode:
##
## - an iterator symbol (Domain)
## - an integer immediate or schedule parameter
## - a +- and * form over those, plus a runtime tensor read in index
## position (TensorRead), the scalar term of a read slot's offset
# walk the node tree:
proc walk(n: NimNode, iters: var Table[string, Iter], params: Table[string, int32]): LuxNode =
case n.kind
of nnkIntLit:
newLuxInt(n.intVal)
of nnkIdent, nnkSym:
let s = n.strVal
if s in iters:
domainNode(iters[s])
elif s in params:
newLuxIntParam(s)
else:
# An unknown symbol binds a new iterator symbol, ranges
# deriving from the reads at lowering time and stray symbols
# failing there with a clear range error
iters[s] = newIter(s)
domainNode(iters[s])
of nnkPar:
doAssert n.len == 1, "unsupported parenthesized index"
walk(n[0], iters, params)
of nnkPrefix:
doAssert n[0].strVal == "-", "unsupported index operator " & n[0].strVal
~(walk(n[1], iters, params))
of nnkInfix:
let op = n[0].strVal
case op
of "+", "-":
let l = walk(n[1], iters, params)
let r = walk(n[2], iters, params)
(if op == "+": l + r else: l + (~r))
of "*":
walk(n[1], iters, params) * walk(n[2], iters, params)
else:
error("unsupported index operator " & op, n)
newLuxInt(0)
else:
case n.kind
of nnkBracketExpr:
let sym = n[0].strVal
doAssert sym in tensors,
"luxtiles: scalar-term read of '" & sym & "' is neither a recipe " &
"parameter nor an earlier Fn target"
var idxs: seq[LuxNode] = @[]
if n.len == 2 and n[1].kind == nnkPar:
for i in 0 ..< n[1].len:
idxs.add(walk(n[1][i], iters, params))
else:
for i in 1 ..< n.len:
idxs.add(walk(n[i], iters, params))
var f = newLux(Func)
f.fn = tensors[sym]
var tr = newLux(TensorRead)
tr.children.add(f)
for ix in idxs:
tr.children.add(ix)
tr
else:
error("unsupported index expression node " & $n.kind, n)
newLuxInt(0)
walk(n, iters, params)
proc valueExpr*(n: NimNode, tensors: var Table[string, Fn], iters: var Table[string, Iter], params: Table[string, int32]): LuxNode =
## Returns one stage definition as a value-layer LuxNode tree:
##
## - tensor reads, Access nodes over Fn
## - tile sums and products (BinOp)
## - float or int literals
# walk the node tree:
proc walk(n: NimNode, tensors: var Table[string, Fn], iters: var Table[string, Iter], params: Table[string, int32]): LuxNode =
case n.kind
of nnkInfix:
let op = n[0].strVal
case op
of "+", "*", "-", "/":
let l = walk(n[1], tensors, iters, params)
let r = walk(n[2], tensors, iters, params)
case op
of "+": result = l + r
of "*": result = l * r
of "-": result = tileOpNode(opSub, @[l, r])
else: result = tileOpNode(opDiv, @[l, r])
else:
error("unsupported value operator " & op, n)
result = newLuxFloat(0.0)
of nnkBracketExpr:
let sym = n[0].strVal
doAssert sym in tensors,
"luxtiles: read of '" & sym & "' is neither a recipe parameter nor " &
"an earlier Fn target"
var indices: seq[LuxNode] = @[]
if n.len == 2 and n[1].kind == nnkPar:
for i in 0 ..< n[1].len:
indices.add(indexExpr(n[1][i], tensors, iters, params))
else:
for i in 1 ..< n.len:
indices.add(indexExpr(n[i], tensors, iters, params))
doAssert indices.len in 1 .. 5,
"reads carry 1 to 5 index slots"
result = accessNode(newCall(tensors[sym], indices))
of nnkCall:
let callee = n[0].strVal
doAssert callee in ["exp2", "exp", "max", "where"],
"unsupported value call " & callee &
" (the tile vocabulary is exp2, exp, max, where)"
if callee == "where":
doAssert n.len == 4,
"where takes a condition and two tile values"
let thn = walk(n[2], tensors, iters, params)
let els = walk(n[3], tensors, iters, params)
result = tileOpNode(opMux, @[
condExpr(n[1], iters, params, tensors), thn, els])
else:
let op = (case callee
of "exp2": opExp2
of "exp": opExp
else: opMax)
var operands: seq[LuxNode] = @[]
for i in 1 ..< n.len:
operands.add(walk(n[i], tensors, iters, params))
result = tileOpNode(op, operands)
of nnkPrefix:
# unary minus on a literal, the tile vocabulary's negative fill
# (an -inf mask value)
doAssert n[0].strVal == "-" and n[1].kind in {nnkFloatLit, nnkFloat32Lit,
nnkFloat64Lit, nnkIntLit},
"only -literal is supported in value position"
if n[1].kind == nnkIntLit:
result = newLuxInt(-n[1].intVal)
else:
result = newLuxFloat(-n[1].floatVal)
of nnkPar:
doAssert n.len == 1, "unsupported parenthesized expression"
result = walk(n[0], tensors, iters, params)
of nnkFloatLit, nnkFloat32Lit, nnkFloat64Lit:
result = newLuxFloat(n.floatVal)
of nnkIntLit:
result = newLuxInt(n.intVal)
else:
error("unsupported expression node " & $n.kind, n)
result = newLuxFloat(0.0)
walk(n, tensors, iters, params)
proc `range`*(loop: string, n: int): ScheduleDirective =
## Literal range lock for one loop, `range(jc, 8)` in a schedule block.
## A declaration `var i = domain(n)` carries the same literal recipe-wide
## and this directive overrides the declaration for one stage.
ScheduleDirective(kind: ScRange, name: loop, value: int32(n))
proc domainRange*(symbol: string, arg: NimNode): Range =
## A declared domain's range from its written argument.
## - an integer literal locks the range (rkPin)
## - any other expression is a host value resolved at the eval site
## (rkDecl), the kernels needing it carrying it as an int32 scalar
## argument named symbol & "R"
##
## Returns the Range.
if arg.kind == nnkIntLit:
Range(kind: rkPin, pin: int32(arg.intVal))
else:
Range(kind: rkDecl, hostNode: arg.copyNimTree(),
argName: symbol & "R")
proc parseSchedule*(blk: NimNode): StageSchedule =
## A `do:` block into a StageSchedule, one directive per statement.
##
## - tile(loop, n), cache(t), range(loop, n), loopNest(a, b, ...)
## - param(name, v), threads = n
doAssert blk.kind == nnkStmtList, "the schedule must be a do block"
for st in blk:
case st.kind
of nnkCall:
let callee = st[0].strVal
case callee
of "tile":
doAssert st.len in 2 .. 3 and st[1].kind == nnkIdent,
"tile must be written tile(loop) or tile(loop, n)"
let size = (if st.len == 3: st[2].intVal.int else: 0)
result.directives.add(ScheduleDirective(
kind: ScTile, name: st[1].strVal, value: int32(size)))
of "cache":
doAssert st.len >= 2, "cache takes tensor names"
for arg in st[1 .. ^1]:
doAssert arg.kind == nnkIdent, "cache takes tensor names"
result.directives.add(ScheduleDirective(
kind: ScCache, name: arg.strVal))
of "range":
doAssert st.len == 3 and st[1].kind == nnkIdent and
st[2].kind == nnkIntLit,
"range must be written range(loop, n)"
result.directives.add(`range`(st[1].strVal, st[2].intVal.int))
of "loopNest":
doAssert st.len >= 2,
"loopNest takes the loop names in the wanted order"
var names: seq[string] = @[]
for arg in st[1 .. ^1]:
doAssert arg.kind == nnkIdent,
"loopNest takes loop names"
names.add(arg.strVal)
result.directives.add(ScheduleDirective(
kind: ScOrder, names: names))
of "reduce":
doAssert st.len == 2 and st[1].kind in {nnkIdent, nnkSym} and
st[1].strVal in ["sum", "max"],
"reduce must be written reduce(sum) or reduce(max)"
result.directives.add(ScheduleDirective(
kind: ScReduce, name: st[1].strVal))
of "atomic":
doAssert st.len == 3 and st[1].kind == nnkIdent and st[2].kind == nnkIntLit,
"atomic must be written atomic(target, range)"
result.directives.add(ScheduleDirective(
kind: ScAtomic, name: st[1].strVal, value: int32(st[2].intVal)))
of "guard":
doAssert st.len == 2,
"guard takes one predicate expression"
result.directives.add(ScheduleDirective(
kind: ScGuard, exprNode: st[1]))
of "param":
doAssert st.len == 3 and st[1].kind == nnkIdent and
st[2].kind == nnkIntLit,
"param must be written param(name, integer)"
result.directives.add(ScheduleDirective(
kind: ScParam, name: st[1].strVal, value: int32(st[2].intVal)))
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"
result.directives.add(ScheduleDirective(
kind: ScThreads, value: int32(st[1].intVal)))
else:
error("unsupported schedule statement " & $st.kind, st)
proc stageSize*(sched: StageSchedule, name: string): int =
## Tile size bound for `name` by the schedule, 0 when untiled.
for d in sched.directives:
if d.kind == ScTile and d.name == name and d.value != 0:
return d.value.int
0
proc hasCache*(sched: StageSchedule, tensor: string): bool =
## Whether the schedule stages `tensor` through threadgroup memory.
for d in sched.directives:
if d.kind == ScCache and d.name == tensor:
return true
false
proc schedParams*(sched: StageSchedule): Table[string, int32] =
## Parameter bindings of the schedule, keyed by name.
for d in sched.directives:
if d.kind == ScParam:
result[d.name] = d.value
proc stageThreads*(sched: StageSchedule): int =
## Thread count of the schedule, 0 when unset.
for d in sched.directives:
if d.kind == ScThreads:
return d.value.int
0
proc appendStage*(fn: Fn, lhsIndices: seq[LuxNode], definition: LuxNode, iters: seq[Iter], condition: LuxNode, schedule: StageSchedule) =
## Appends one stage to `fn`.
##
## - definition applies on the LHS domain
## - iterators absent from the LHS indices become the recurrence space
var st: Stage
st.definition = definition
for ix in lhsIndices:
st.params.add(ix)
for it in iters:
var onLhs = false
for ix in lhsIndices:
if ix.isIterator and ix.iter.symbol == it.symbol: onLhs = true
if not onLhs:
st.recurrence.add(domainNode(it))
st.condition = condition
st.iterOrder = @[]
for it in iters:
st.iterOrder.add(it.symbol)
for d in schedule.directives:
if d.kind == ScReduce:
st.reduce = (if d.name == "max": rkMax else: rkSum)
st.schedule = schedule
fn.stages.add(st)
proc collectIterSymbols*(n: LuxNode): seq[string] =
## Returns iterator symbols in appearance order from a depth-first
## walk of the definition tree, reduce iters with LHS symbols first.
var acc: seq[string] = @[]
proc walk(n: LuxNode, acc: var seq[string]) =
case n.kind
of Domain:
if n.iter.symbol notin acc:
acc.add(n.iter.symbol)
of BinOp, Access, TensorRead:
for c in n.children:
walk(c, acc)
of TileOp:
for c in n.operands:
walk(c, acc)
else:
discard
walk(n, acc)
accluxtiles_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.
## luxtiles_compiler, the Lux statement IR lowered to Metal device functions.
##
## A recipe proc's body is a sequence of stage statements over Iter
## symbols:
##
## - one Fn per target name, one stage per definition statement,
## schedules coming from trailing do-blocks
## - a pass pipeline over the stage's AlgoIR lowers every stage into
## one Metal device fn + thin entry, mma an optimization pass,
## fusion grouping a recipe's stages into one launch
##
## Compile-time dataflow, stage to MSL
##
## var i = domain(M)
## A[i, j] += m1[i, k] * m2[k, j] do:
## tile(i, 16) # the stage schedule
##
## luxtile statement walk -> stageToIR, the AlgoIR
##
## -> schedule passes -> makeTilePlan, per-read view geometry
## -> emitTileKernel, the derived tile program
## -> emitDeviceFunc + emitEntry, evalCall, MSL
##
## Runtime lifecycle, stage to result
##
## recipe proc statement
## | expansion, fnDef + entry, result binding
##
## v allocImpl + ingest, once per kernel
## | run, a LaunchConfig grid over blk
##
## v host result tensor rebound to the target local
import
std/[macros, macrocache, strutils, tables, typetraits],
workspace/experimental/tensorcery_luxtiles/luxtiles_types,
workspace/experimental/tensorcery_luxtiles/luxtiles_dsl,
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/ptr_arithmetic,
workspace/ceramic/src/tensor_datatypes,
workspace/ceramic/src/layout_algebra,
workspace/ceramic/src/tile_algebra,
workspace/ceramic/src/tile_algebra/tile_io_bounded,
workspace/crucible,
workspace/crucible/src/codegen/gpu_compiler
# ── Launch expansion, alloc → ingest-once → run ──────────────────────────
# One statement expansion turns the lowering's launch spec (DescSpecs) plus
# the host call arguments into one output allocation, one wrapper ingest
# and one run, the result tensor bound to the statement's target ident.
#
# buildGlobalFn wraps the device fn node in a `{.global.}` proc
#
# device fn node
# | wrapper, TensorView params marshalled to raw pointers plus
# | shape and stride scalars
#
# We need `{.smem.}` declared inside a kernel, so with `inlineBody` the body
# splices through a fresh template re-typed against the wrapper scope.
const DescSpecs = CacheTable"tc.desc"
## Per-kernel launch spec fields, written by the lowering pass's thin
## entry and read by the statement expansion, `smemDecls` optional.
proc descFields(fnName: string): Table[string, NimNode] =
doAssert DescSpecs.hasKey(fnName), "desc spec for " & fnName
result = initTable[string, NimNode]()
for eq in DescSpecs[fnName]:
doAssert eq.kind == nnkExprEqExpr
result[$eq[0]] = eq[1]
proc substituteIdents(n: NimNode, bindings: Table[string, NimNode]): NimNode =
if n.kind in {nnkSym, nnkIdent} and n.strVal in bindings:
return bindings[n.strVal]
if n.kind in {nnkEmpty, nnkIdent, nnkSym} or n.len == 0:
return n
var changed = false
var kids: seq[NimNode] = @[]
for i in 0 ..< n.len:
let k = substituteIdents(n[i], bindings)
if k != n[i]:
changed = true
kids.add(k)
if not changed:
return n
result = newNimNode(n.kind, n)
for k in kids:
result.add(k)
proc tupleConstrToPar(n: NimNode): NimNode =
## Type-position `(int, int)` tuple nodes rewritten to the paren form,
## value tuples degrading to `void`.
if n.kind == nnkTupleConstr:
result = nnkPar.newTree()
for ch in n:
result.add(tupleConstrToPar(ch))
elif n.kind in {nnkBracketExpr, nnkPtrTy, nnkPar} and n.len > 0:
result = newNimNode(n.kind)
for ch in n:
result.add(tupleConstrToPar(ch))
else:
result = n
proc buildGlobalFn(name: string, origSym: NimNode, impl: NimNode,
binds: NimNode, inlineBody: bool,
prelude = nnkStmtList.newTree(),
staticArgs: seq[NimNode] = @[]): NimNode =
## `{.global.}` wrapper proc around a device fn. Formals in order:
##
## o buf | tensor bufs... | sc (one packed int32 array) | scratch | geometry
##
## Packed sc layout, 2*rank ints per view:
##
## shapes 0..rank-1 then strides 0..rank-1, per formal in order,
## then one int per scalar formal
##
## The expansion side (`evalCall`) packs the same layout, walking
## the same formals.
doAssert impl.kind in {nnkProcDef, nnkTemplateDef},
"buildGlobalFn: expected a proc or template kernel body"
var bindings = initTable[string, NimNode]()
for arg in binds:
bindings[$arg[0]] = arg[1]
var formal = nnkFormalParams.newTree(newEmptyNode())
var marshals = nnkStmtList.newTree()
# the packed scalar array's formal, appended after the buffer pointers
var scOffset = 0
proc scAt(k: int): NimNode =
nnkBracketExpr.newTree(ident"sc", newLit(k))
var scalarFormals: seq[tuple[name: string, off: int]] = @[]
var scratchFormals: seq[NimNode] = @[]
# pass 1 walks the buffer formals and the sc layout, one view's
# 2*rank ints, then one int per scalar formal (geometry stays formal)
for i in 1 ..< impl[3].len:
let idefs = impl[3][i]
let ptype = substituteIdents(idefs[idefs.len - 2], bindings)
for j in 0 ..< idefs.len - 2:
let pname = idefs[j]
let pnameStr = pname.strVal
if ptype.kind == nnkBracketExpr and $ptype[0] == "TensorView":
bindings[pnameStr] = newIdentNode(pnameStr)
let elT = ptype[1]
let rank = ptype[2].len
let bufName = newIdentNode(pnameStr & "Buf")
formal.add(nnkIdentDefs.newTree(bufName,
nnkPtrTy.newTree(nnkBracketExpr.newTree(ident"UncheckedArray", elT)),
newEmptyNode()))
var shapeArgs = nnkPar.newTree()
var strideArgs = nnkPar.newTree()
for ax in 0 ..< rank:
shapeArgs.add(nnkCall.newTree(ident"int", scAt(scOffset + ax)))
for ax in 0 ..< rank:
strideArgs.add(nnkCall.newTree(ident"int", scAt(scOffset + rank + ax)))
marshals.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
newIdentNode(pnameStr),
# Annotated type. The GPU code generator does not infer a let
# local's type through `make_view`, so the marshalled view needs
# the annotation where spliced kernel bodies touch `.data`.
tupleConstrToPar(substituteIdents(ptype, bindings)),
nnkCall.newTree(ident"make_view", bufName,
nnkCall.newTree(ident"make_layout", shapeArgs, strideArgs)))))
scOffset += 2 * rank
elif ptype.kind == nnkPtrTy:
# threadgroup scratch pointer, a real kernel formal, appended after sc
scratchFormals.add(nnkIdentDefs.newTree(pname, ptype, newEmptyNode()))
elif $ptype == "int32" and pnameStr in ["tid", "nthr", "bx", "by", "bz"]:
formal.add(nnkIdentDefs.newTree(pname, ptype, newEmptyNode()))
else:
# a scalar launch value, packed into sc
scalarFormals.add((pnameStr, scOffset))
bindings[pnameStr] = scAt(scOffset)
inc scOffset
formal.add(nnkIdentDefs.newTree(ident"sc",
nnkPtrTy.newTree(nnkBracketExpr.newTree(ident"UncheckedArray",
ident"int32")),
newEmptyNode()))
for sf in scratchFormals:
formal.add(sf)
if inlineBody:
let tplSym = genSym(nskTemplate, name & "Inline")
var tplFormal = nnkFormalParams.newTree(newEmptyNode())
var tplCall = nnkCall.newTree(tplSym)
for i in 1 ..< impl[3].len:
let idefs = impl[3][i]
let ptype = tupleConstrToPar(
substituteIdents(idefs[idefs.len - 2], bindings))
var idNames = nnkIdentDefs.newTree()
for j in 0 ..< idefs.len - 2:
let pname = idefs[j]
idNames.add(newIdentNode(pname.strVal))
tplCall.add(bindings[pname.strVal])
idNames.add(ident"auto")
idNames.add(newEmptyNode())
tplFormal.add(idNames)
let tplGenParams = (if staticArgs.len > 0 and impl[2].kind == nnkGenericParams:
substituteIdents(impl[2], bindings)
else:
newEmptyNode())
let tplCallHead: NimNode = (if staticArgs.len > 0:
block:
var h = nnkBracketExpr.newTree(tplSym)
for sa in staticArgs:
h.add(sa)
h
else:
tplSym)
marshals.add(nnkTemplateDef.newTree(
tplSym, tplGenParams, newEmptyNode(),
tplFormal, newEmptyNode(), newEmptyNode(),
substituteIdents(impl[6], bindings)))
if prelude.len > 0:
# Hoisted declarations (e.g. smem arrays) must live in the kernel
# function itself, not in the spliced template, and the spliced
# template captures them by name at its instantiation site.
# A `{.smem.}` pragma on a var inside an instantiated template
# expands to a call node the codegen's address-space resolver rejects.
for st in prelude:
marshals.add(st)
var tplCallNode = nnkCall.newTree(tplCallHead)
for i in 1 ..< tplCall.len:
tplCallNode.add(tplCall[i])
marshals.add(tplCallNode)
else:
var callArgs = @[origSym]
for i in 1 ..< impl[3].len:
let idefs = impl[3][i]
for j in 0 ..< idefs.len - 2:
callArgs.add(bindings[idefs[j].strVal])
marshals.add(nnkCall.newTree(callArgs))
result = nnkProcDef.newTree(
newIdentNode(name), newEmptyNode(), newEmptyNode(),
formal,
nnkPragma.newTree(ident"global"),
newEmptyNode(),
marshals)
proc flattenExprs(n: NimNode, outExprs: var seq[NimNode]) =
if n.kind == nnkPrefix and n[0].kind in {nnkIdent, nnkSym} and
n[0].strVal == "@":
for e in n[1]:
outExprs.add(e)
elif n.kind == nnkCall and n[0].strVal == "@":
for i in 1 ..< n.len:
outExprs.add(n[i])
elif n.kind == nnkCall and n[0].kind == nnkBracketExpr and
n[0][0].strVal == "@":
for i in 1 ..< n.len:
outExprs.add(n[i])
elif n.kind in {nnkBracket, nnkTupleConstr}:
for e in n:
outExprs.add(e)
elif n.kind == nnkPar and n.len > 1:
for e in n:
flattenExprs(e, outExprs)
elif n.kind == nnkPar and n.len == 1:
flattenExprs(n[0], outExprs)
else:
outExprs.add(n)
proc viewArgsTuple(arg: NimNode, rank: int): NimNode =
## Marshalling tuple for one host tensor argument, the raw buffer plus
## shape/stride scalars per axis, rank-matched to the device formal.
var parts: seq[NimNode] = @[nnkCall.newTree(ident"operandArg", arg)]
for ax in 0 ..< rank:
# Plain BracketExpr nodes, a quote-built `arg.shape[ax]` resolving into
# symbol-choice nodes whose name the expression flattener cannot read.
for field in ["shape", "strides"]:
parts.add(nnkBracketExpr.newTree(
nnkDotExpr.newTree(arg, ident(field)), newLit(ax)))
nnkPar.newTree(parts)
proc evalExpansion(call, fnSym, engineExpr, cgIdent: NimNode, kernelName = ""): NimNode =
## One launch expansion, output allocated from the runtime shape, wrapper
## ingested once, run binding the output blob to the target ident.
var sym = fnSym
if sym.kind in {nnkOpenSymChoice, nnkClosedSymChoice}:
sym = sym[0]
let impl = sym.getImpl()
doAssert impl.kind == nnkTemplateDef,
"eval: expected a template kernel body (the DSL derives templates only)"
let fields = descFields((if kernelName.len > 0: kernelName else: $sym))
let outParam = fields["output"].strVal
# Host call args bind onto the device fn params positionally.
# The output param and static params take no host arg.
var selfReads: seq[string] = @[]
if "selfReads" in fields:
for n in fields["selfReads"]:
selfReads.add(n.strVal)
var bindings = initTable[string, NimNode]()
var blobs: seq[string] = @[]
var blobRanks = initTable[string, int]()
var scalarNames: seq[string] = @[]
# the output formal's rank drives the output marshalling tuple
var outRank = 2
for pi in 1 ..< impl[3].len:
let idefs = impl[3][pi]
let pt = idefs[idefs.len - 2]
for j in 0 ..< idefs.len - 2:
if idefs[j].strVal == outParam and pt.kind == nnkBracketExpr and
$pt[0] == "TensorView":
outRank = pt[2].len
let outShape = genSym(nskLet, "outShape")
let outTensor = genSym(nskLet, "outTensor")
let outArgSym = genSym(nskVar, "outArg")
var pos = 1
for pi in 1 ..< impl[3].len:
let idefs = impl[3][pi]
let ptype = idefs[idefs.len - 2]
for j in 0 ..< idefs.len - 2:
let name = idefs[j].strVal
if name == outParam or ptype.kind == nnkStaticTy:
continue
let isView = ptype.kind == nnkBracketExpr and $ptype[0] == "TensorView"
if name in selfReads:
# A self-reading (recurrence) stage reads the freshly allocated,
# zero-filled output buffer as its implicit zero seed, no host arg.
bindings[name] = outTensor
if isView: blobs.add(name)
continue
doAssert pos < call.len,
"eval: no host arg for device param " & name
bindings[name] = call[pos]
inc pos
if isView:
blobs.add(name)
else:
scalarNames.add(name)
doAssert pos == call.len, "eval: unconsumed host args"
let fnName = $fnSym
let kname = (if "kernel" in fields: fields["kernel"].strVal else: fnName)
let fnNameLit = quote do: `fnName`
let knameLit = quote do: `kname`
let outShapeRaw = ident"toMetadata"
.newCall(substituteIdents(fields["outShape"], bindings))
var argCode = nnkStmtList.newTree()
# Packed sc layout, matching the wrapper's:
# - the output's 2*rank ints first, `run`'s own argument the output buffer
# - then per formal, 2*rank ints for a view, one int for a scalar
var scVals: seq[NimNode] = @[]
# the output's ints first, matching the wrapper's sc layout
for ax in 0 ..< outRank:
scVals.add(nnkBracketExpr.newTree(
nnkDotExpr.newTree(outTensor, ident"shape"), newLit(ax)))
for ax in 0 ..< outRank:
scVals.add(nnkBracketExpr.newTree(
nnkDotExpr.newTree(outTensor, ident"strides"), newLit(ax)))
for pi in 1 ..< impl[3].len:
let idefs = impl[3][pi]
let ptype = idefs[idefs.len - 2]
for j in 0 ..< idefs.len - 2:
let name = idefs[j].strVal
if name == outParam or ptype.kind == nnkStaticTy:
continue
if ptype.kind == nnkBracketExpr and $ptype[0] == "TensorView":
let argExpr = substituteIdents(bindings[name], bindings)
let rank = ptype[2].len
for ax in 0 ..< rank:
scVals.add(nnkBracketExpr.newTree(
nnkDotExpr.newTree(argExpr, ident"shape"), newLit(ax)))
for ax in 0 ..< rank:
scVals.add(nnkBracketExpr.newTree(
nnkDotExpr.newTree(argExpr, ident"strides"), newLit(ax)))
elif name in ["tid", "nthr", "bx", "by", "bz"]:
continue
else:
scVals.add(substituteIdents(bindings[name], bindings))
doAssert scVals.len > 0, "eval: a kernel with no packed scalars is impossible"
let scArr = genSym(nskLet, "scArr")
var scLit = nnkBracket.newTree()
for v in scVals:
scLit.add(v)
argCode.add(nnkLetSection.newTree(nnkIdentDefs.newTree(scArr,
newEmptyNode(), scLit)))
var tupleNode = nnkPar.newTree()
for name in blobs:
let argExpr = substituteIdents(bindings[name], bindings)
let aSym = genSym(nskLet, "operand")
argCode.add quote do:
let `aSym` = `argExpr`
tupleNode.add(nnkCall.newTree(ident"operandArg", aSym))
let scArgSym = genSym(nskLet, "scArg")
let scLenLit = newLit(scVals.len)
argCode.add quote do:
let `scArgSym` = PtrArg[int32](
buf: cast[ptr UncheckedArray[int32]](addr `scArr`[0]),
len: `scLenLit`, off: 0)
tupleNode.add(scArgSym)
# DSL kernels are template bodies with statics baked in as literals.
# The wrapper, with the body spliced in, is built here at macro time,
# since a template symbol cannot flow through a macro argument as a value.
let prelude = (if "smemDecls" in fields:
nnkStmtList.newTree(fields["smemDecls"])
else:
nnkStmtList.newTree())
var staticArgs: seq[NimNode] = @[]
if "statics" in fields:
for eq in fields["statics"]:
staticArgs.add(substituteIdents(eq[1], bindings))
let genFnCall = buildGlobalFn(kname, sym, impl, nnkBracket.newTree(), true,
prelude, staticArgs)
let grid = substituteIdents(fields["grid"], bindings)
let blk = substituteIdents(fields["blk"], bindings)
result = quote do:
let `outShape` = `outShapeRaw`
if `engineExpr`.resident != `fnNameLit`:
let unit = `cgIdent`:
`genFnCall`
when defined(crucibleDumpMsl):
echo "=== MSL ===\n", unit, "=== END MSL ==="
`engineExpr`.ingest(unit)
`engineExpr`.resident = `fnNameLit`
let outImpl = allocImpl[float32](`outShape`)
let `outTensor` = makeHostTensor(outImpl.shape, outImpl.strides,
outImpl)
var `outArgSym` = PtrArg[float32](buf: outImpl.buf,
len: numel(`outShape`).int, off: 0)
`argCode`
when defined(launchTrace):
echo "launch: ", `knameLit`
run(`engineExpr`, `knameLit`, `outArgSym`, `tupleNode`,
LaunchConfig(grid: Dim3(x: `grid`[0], y: `grid`[1], z: `grid`[2]),
blk: Dim3(x: `blk`[0], y: `blk`[1], z: `blk`[2])))
`outTensor`
macro evalCall(e: TensorContext, fn: typed, call: untyped,
kernelName: static string = ""): LuxTensor =
## Backend-dispatched launch of the kernel `fn` over `call`'s host arguments,
## `kernelName` overriding the spec lookup key on re-splices.
result = evalExpansion(call, fn, nnkDotExpr.newTree(e, ident"metal"),
ident"metal", kernelName)
func tdim*(t: LuxTensor, 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
# ── Compile-time caches and kernel constants ─────────────────────────────
# Derivation runs once per kernel name, the caches holding the pass trace,
# the schedule provenance and the derived (device fn, thin entry) pair.
#
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`.
KernelCache = CacheTable"tc.kernel"
## Derived (device fn, thin entry) node pairs, keyed by kernel name.
## The second statement with the same name reuses, not re-derives.
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)
# ─────────────────────────────────────────────────────────────────────────
# Nest IR, one record per algorithm statement
# ─────────────────────────────────────────────────────────────────────────
# `AlgoIR` (types in luxtiles_types) carries the parse of one statement,
# each pass reading and annotating it in place until the lowering consumes it.
# - parse builds the reads, the expression tree and the loop nest
# - schedule passes annotate tiles, cache plans and thread count
# - mmaOptPass annotates `hasMma` + `mmaMatch` or leaves them zero
#
# parseTarget / parseValue / parseIndex, buildLoops -> AlgoIR
# | schedule passes (tilePass / cacheAnnotatePass / cachePlanPass / threadsPass)
#
# v mmaOptPass, hasMma + mmaMatch
# | lowerEmit, TilePlan or MmaMatch consumed
#
# ─────────────────────────────────────────────────────────────────────────
# ─────────────────────────────────────────────────────────────────────────
# Pass framework, a schedule is a seq[Pass] run in written order
# ─────────────────────────────────────────────────────────────────────────
# Transform and analysis passes rewrite or annotate the IR in place.
# The single lowering pass emits the tile program and its form description,
# `runPipeline` recording the IR after every pass in `KernelTrace`.
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 & "["
case lp.range.kind
of rkAxis:
if lp.range.tensor.len > 0:
s &= "ax" & $lp.range.axis & "(" & lp.range.tensor & ")"
of rkAffine:
s &= "aff" & $lp.range.coef & "(" & lp.range.affineOf.tensor &
".ax" & $lp.range.affineOf.axis & ") tmin=" & $lp.range.tmin
of rkPin:
s &= "pin=" & $lp.range.pin
of rkDecl:
s &= "decl(" & lp.range.argName & ")"
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
of ikAffine:
var f = ""
for t in ix.terms:
f &= (if f.len == 0: "" else: "+") & $t.coef & "*" & t.iter
f & (if ix.cst != 0: (if ix.cst < 0: "-" else: "+") & $abs(ix.cst)
else: ""))
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 &
(if ir.recurrence.len > 0: " | recurrence: " & ir.recurrence.join(",")
else: "") &
(if ir.reduceKind == rkMax: " | reduce=max" else: "")
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"
if not KernelTrace.hasKey(ir.name):
# a fusion-planned stage is analyzed twice (plan + emit), the first trace standing
KernelTrace[ir.name] = newLit(trace)
# ─────────────────────────────────────────────────────────────────────────
# Schedule micro-passes over the schedule AST
# ─────────────────────────────────────────────────────────────────────────
# One pass per schedule property, tile / cache / cachePlan / threads, each
# a small in-place transform or analysis over the IR.
proc tilePass(loop: string, size: int): Pass =
## Splits `loop` into an outer step-`size` loop and an inner range-`size` loop, read indices stay logical (lowering rewrites them).
Pass(name: "tile(" & loop & "," & $size & ")",
description: "split loop '" & loop & "' into outer step " & $size &
" and inner range " & $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",
range: ir.loops[i].range,
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:
if lp.tile > 0:
plan.buffer.add((ix.name, lp.tile))
elif lp.role == lrReduce:
doAssert false,
"cache: reduce iterator '" & ix.name &
"' staging '" & t & "' must be tiled"
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 affAdd(a, b: AffineForm): AffineForm =
## Sums two linear forms, merging same-iterator coefficients.
result = a
result.cst += b.cst
for t in b.scalars:
result.scalars.add(t)
for t in b.coefs:
var merged = false
for i in 0 ..< result.coefs.len:
if result.coefs[i].iter == t.iter:
result.coefs[i].coef += t.coef
merged = true
if not merged:
result.coefs.add(t)
proc affNeg(a: AffineForm): AffineForm =
## Negates a linear form, scalar-term weights included.
result.cst = -a.cst
for t in a.scalars:
result.scalars.add(ScalarTerm(tensor: t.tensor, coef: -t.coef,
indices: t.indices))
for t in a.coefs:
result.coefs.add((-t.coef, t.iter))
proc affScale(a: AffineForm, k: int32): AffineForm =
## Scales a linear form by an integer constant, scalar terms included.
result.cst = a.cst * k
for t in a.scalars:
result.scalars.add(ScalarTerm(tensor: t.tensor, coef: t.coef * k,
indices: t.indices))
for t in a.coefs:
result.coefs.add((t.coef * k, t.iter))
proc isConst(a: AffineForm): bool =
a.coefs.len == 0 and a.scalars.len == 0
proc iterRange*(ir: AlgoIR, iter: string, skip = ""): Range =
## An iterator's range, a schedule `range(loop, n)` value first, else:
##
## - the (tensor, axis) of its first binding read slot
## - or the derived range of its first affine binding, none when unbound
## - reads of the stage's own target (`skip`) ignored, ranges derive
## from external reads
if iter in ir.rangePins:
return Range(kind: rkPin, pin: ir.rangePins[iter])
if ir.declRanges.hasKey(iter):
return ir.declRanges[iter]
for rd in ir.reads:
if rd.tensor == skip: continue
for k, ix in rd.indices:
if ix.kind == ikIter and ix.name == iter:
return Range(kind: rkAxis, tensor: rd.tensor, axis: k)
for st in ix.scalars:
for j, six in st.indices:
if six.kind == ikIter and six.name == iter:
return Range(kind: rkAxis, tensor: st.tensor, axis: j)
for rd in ir.reads:
if rd.tensor == skip: continue
for k, ix in rd.indices:
# a form with a runtime scalar term (a TensorRead) cannot bound
# the iterator, its range coming from another binding read
if ix.scalars.len > 0: continue
if ix.kind == ikAffine:
for t in ix.terms:
if t.iter == iter:
doAssert t.coef > 0,
"iterator '" & iter & "' binds with a non-positive coefficient; " &
"no derived range exists for it"
for t2 in ix.terms:
doAssert t2.coef > 0,
"iterator '" & iter & "' shares an index form with a " &
"negatively weighted iterator; no derived range exists"
return Range(kind: rkAffine, affineOf: (rd.tensor, k),
coef: t.coef, tmin: ix.cst)
raiseAssert("no read binds iterator '" & iter & "'; its range 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, range: ir.iterRange(nm),
role: (if nm in ir.reduceIters: lrReduce else: lrOut)))
# ─────────────────────────────────────────────────────────────────────────
# ─────────────────────────────────────────────────────────────────────────
# Stage parse, the LuxNode IR into AlgoIR
# ─────────────────────────────────────────────────────────────────────────
# One stage's definition tree and index forms are LuxNode, walked into
# the AlgoIR the pass pipeline consumes.
#
# Stage (definition, params, recurrence) -> AlgoIR
# | outDims from the LHS domain nodes
# | iters, LHS symbols first then the definition's read order
# | valueFromLux / indexFromLux, reads + ExprIR + IndexRef
# | buildLoops, the loop nest
#
# Index-layer forms, accepted and rejected:
#
# | form | verdict |
# | -------------------- | ------------------------------------ |
# | 3*i + k - 1 | accepted, an iterator form |
# | oh*stride + kh - pad | accepted, stride and pad as params |
# | i * j | rejected, a product of two iterators |
#
# TensorRead extension point:
# the read-index representation is an index-layer LuxNode tree, not
# a closed sum:
#
# - a runtime tensor value in index position extends indexFromLux
# with one branch, threading its value through the emitters
# - paged-KV block tables and MoE expert gathers need exactly that branch
proc indexFromLux(n: LuxNode, params: Table[string, int32]): IndexRef =
## Returns one read index from an index-layer LuxNode tree, staying
## in the simplest kind:
##
## - an iterator
## - a constant
## - a constant times iterator sum
var f: AffineForm
proc walk(n: LuxNode): AffineForm =
case n.kind
of IntLit:
AffineForm(cst: int32(n.intVal))
of IntParam:
doAssert n.symParam in params,
"index parameter '" & n.symParam & "' is not bound by the schedule"
AffineForm(cst: params[n.symParam])
of Domain:
AffineForm(coefs: @[(1'i32, n.iter.symbol)])
of TensorRead:
# A runtime tensor read in index position, its value weighting
# the slot's offset, the slot's stride scaling the loaded value
var st = ScalarTerm(tensor: n.children[0].fn.symbol, coef: 1)
for i in 1 ..< n.children.len:
st.indices.add(indexFromLux(n.children[i], params))
AffineForm(scalars: @[st])
of BinOp:
case n.children[0].bopKind
of Add:
affAdd(walk(n.children[1]), walk(n.children[2]))
of Sub:
affAdd(walk(n.children[1]), affNeg(walk(n.children[2])))
of Mul:
let l = walk(n.children[1])
let r = walk(n.children[2])
if l.isConst and not r.isConst:
affScale(r, l.cst)
elif r.isConst and not l.isConst:
affScale(l, r.cst)
elif l.isConst and r.isConst:
AffineForm(cst: l.cst * r.cst)
else:
raiseAssert("index may not multiply two iterators: " & $n.kind)
else:
raiseAssert("unsupported index-layer op")
else:
raiseAssert("unsupported index node kind " & $n.kind)
f = walk(n)
var coefs: seq[tuple[coef: int32, iter: string]]
for t in f.coefs:
if t.coef != 0:
coefs.add(t)
f.coefs = coefs
var res: IndexRef
if f.coefs.len == 0:
res = IndexRef(kind: ikConst, value: f.cst)
elif f.coefs.len == 1 and f.coefs[0].coef == 1 and f.cst == 0 and
f.scalars.len == 0:
res = IndexRef(kind: ikIter, name: f.coefs[0].iter)
else:
res = IndexRef(kind: ikAffine, terms: f.coefs, cst: f.cst)
res.scalars = f.scalars
res
proc foldParams*(n: LuxNode, params: Table[string, int32]): LuxNode =
## Condition params folded to literals so emitters emit plain arithmetic.
if n.kind == IntParam:
doAssert n.symParam in params
return newLuxInt(params[n.symParam].int)
result = LuxNode(kind: n.kind, id: n.id)
case n.kind
of BinOp:
result.children = @[n.children[0], foldParams(n.children[1], params),
foldParams(n.children[2], params)]
of TileOp:
for c in n.operands: result.operands.add(foldParams(c, params))
else:
result = n
proc valueFromLux(n: LuxNode, ir: var AlgoIR, params: Table[string, int32]): ExprIR =
## Returns the expression tree over reads from one definition tree,
## holding `+` and `*` over tensor reads (Access over Fn) plus
## float or int literals.
case n.kind
of Access:
var rd = ReadIR(tensor: n.children[0].fn.symbol)
for i in 1 ..< n.children.len:
rd.indices.add(indexFromLux(n.children[i], params))
doAssert rd.indices.len in 1 .. 5,
"reads carry 1 to 5 index slots"
ir.reads.add(rd)
ExprIR(kind: xeRead, read: rd)
of BinOp:
doAssert n.children[0].bopKind in {Add, Mul},
"only + and * are supported in this phase's lowering"
var e = ExprIR(kind: (if n.children[0].bopKind == Add: xeAdd
else: xeMul))
e.lhs = valueFromLux(n.children[1], ir, params)
e.rhs = valueFromLux(n.children[2], ir, params)
e
of TileOp:
case n.tileOp
of opMux:
doAssert n.operands.len == 3,
"mux takes a condition and two tile values"
var e = ExprIR(kind: xeMux, cond: foldParams(n.operands[0], params))
e.lhs = valueFromLux(n.operands[1], ir, params)
e.rhs = valueFromLux(n.operands[2], ir, params)
e
of opSub, opDiv, opMax:
doAssert n.operands.len == 2,
"sub / div / max take two tile operands"
var e = ExprIR(kind: (case n.tileOp
of opSub: xeSub
of opDiv: xeDiv
else: xeMax))
e.lhs = valueFromLux(n.operands[0], ir, params)
e.rhs = valueFromLux(n.operands[1], ir, params)
e
of opExp2, opExp:
doAssert n.operands.len == 1,
"exp2 / exp take one tile operand"
var e = ExprIR(kind: xeExp2)
e.lhs = valueFromLux(n.operands[0], ir, params)
if n.tileOp == opExp:
# exp lowers through the exp2 unit, x * Log2e, tensorcery's op form
var m = ExprIR(kind: xeMul)
m.lhs = e.lhs
m.rhs = ExprIR(kind: xeLit, value: 1.4426950408889634'f32)
e.lhs = m
e
else:
raiseAssert("unsupported tile op in the value layer: " & $n.tileOp)
of FloatLit:
ExprIR(kind: xeLit, value: float32(n.floatVal))
of IntLit:
ExprIR(kind: xeLit, value: float32(n.intVal))
else:
raiseAssert("unsupported definition node kind " & $n.kind)
ExprIR(kind: xeLit)
proc stageToIR(fn: Fn, stageIdx: int, sched: StageSchedule, kernelName: string): AlgoIR =
## Returns one stage's AlgoIR, out dims taken from the LHS domain,
## the iteration space LHS symbols first then the definition's
## read order, the reductions being whatever remains afterwards.
let st = fn.stages[stageIdx]
result.name = kernelName & "_" & fn.symbol &
(if stageIdx > 0: "_s" & $stageIdx else: "")
result.target = fn.symbol
# the atomic scatter form, the LHS index being a computed TensorRead,
# the schedule declaring the output range, the combine a uint-domain
# atomic add of the literal value
var atomicDir: ScheduleDirective
for d in sched.directives:
if d.kind == ScAtomic: atomicDir = d
if atomicDir.kind == ScAtomic:
result.scatter = true
result.scatterRange = atomicDir.value
doAssert fn.stages[stageIdx].definition.kind == IntLit,
"the scatter stage's value must be an integer literal (the count added)"
result.scatterValue = int32(fn.stages[stageIdx].definition.intVal)
doAssert st.params.len == 1 and st.params[0].kind == TensorRead,
"the atomic stage's LHS is one computed index, count[idxExpr]"
var lhsRd = ReadIR(tensor: st.params[0].children[0].fn.symbol)
for i in 1 ..< st.params[0].children.len:
lhsRd.indices.add(indexFromLux(st.params[0].children[i], schedParams(sched)))
result.reads.add(lhsRd)
doAssert st.iterOrder.len >= 1,
"the atomic stage's scatter index must bind a declared domain " &
"or an iterator"
result.iters = st.iterOrder
for d in st.recurrence:
result.recurrence.add(d.iter.symbol)
result.reduceKind = st.reduce
result.dtype = HostElem.name
result.buildLoops()
return
for p in st.params:
doAssert p.isIterator, "LHS indices must be iterator symbols"
result.outDims.add(p.iter.symbol)
doAssert result.outDims.len <= 5, "at most five output dimensions"
var used = result.outDims
for s in collectIterSymbols(st.definition):
if s notin used:
used.add(s)
for s in st.iterOrder:
doAssert s in used,
"the stage's iter order misses iterator '" & s & "'"
for s in used:
doAssert s in st.iterOrder,
"the stage's iter order carries iterator '" & s &
"' beyond the statement's iterators"
result.iters = st.iterOrder
for d in st.recurrence:
result.recurrence.add(d.iter.symbol)
result.reduceKind = st.reduce
if st.condition != nil:
result.guardNode = foldParams(st.condition, schedParams(sched))
# Declared domains, two settlement paths:
# - a literal argument joins rangePins first, a later schedule `range`
# directive overriding it per stage
# - a runtime argument lands on declRanges plus the scalar args
for s in result.iters:
if fn.declOrder.contains(s):
let dr = fn.decls[s]
case dr.kind
of rkPin:
result.rangePins[s] = dr.pin
of rkDecl:
result.declRanges[s] = Range(kind: rkDecl, hostNode: dr.hostNode,
argName: kernelName & "_" & s & "R")
result.declArgs.add((result.declRanges[s].argName, dr.hostNode))
else:
discard
for d in sched.directives:
if d.kind == ScRange:
result.rangePins[d.name] = d.value
let params = schedParams(sched)
result.expr = valueFromLux(st.definition, result, params)
result.buildLoops()
result.dtype = HostElem.name
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
# Node construction helpers for the lowering
# ─────────────────────────────────────────────────────────────────────────
# Node builders for the emitters. We need late name binding, so the builders
# emit plain idents, not quote-do gensyms, for every name a sibling fragment or the wrapper references.
# Late name binding means the name resolves when the fragments splice together,
# not at node build time. Intra-fragment temporaries may gensym.
#
# Views and operators wrap under callable names (`mkGdView`, `offsetBuf`),
# method-call syntax and imported operator templates not resolving inside
# macro-emitted procedure bodies, named templates and procs resolving there.
proc id(s: string): NimNode = newIdentNode(s)
func rowStride*[T, Sh, St](v: TensorView[T, Sh, St]): int32 {.inline.} =
## Row stride of a marshalled device view.
int32(v.layout.stride[0])
func colStride*[T, Sh, St](v: TensorView[T, Sh, St]): int32 {.inline.} =
## Column stride of a marshalled device view.
int32(v.layout.stride[1])
proc mkGdView*[T](buf: ptr UncheckedArray[T], sb, sd, sr, sc: int32): GlView[T] =
## Global-buffer view of dynamic (-1) ranges with the given strides,
## tile_algebra's `gd` in proc form.
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] =
## Buffer pointer advanced by `off` elements.
##
## - wrapping arithmetic, a negative offset subtracting like pointer
## `+%` does (an affine plane slot's serial term entering the base)
cast[ptr UncheckedArray[T]](cast[uint64](
cast[int64](p) + cast[int64](off) * int64(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 rangeExpr(e: Range, device: bool): NimNode =
## Loop range, `tdim(t, k)` in host clause position,
## `int32(t.layout.shape[k])` in kernel-body position.
##
## - affine ranges emit their derivation formula
## - a declared domain's runtime value emits the scalar arg's name,
## the prologue let binding it on the host, the kernel formal
## carrying it on the device
if e.kind == rkDecl:
return id(e.argName)
if e.kind == rkPin:
return callI32(newLit(e.pin.int))
if e.kind == rkAffine:
let num = infixF("-",
infixF("-", rangeExpr(Range(kind: rkAxis, tensor: e.affineOf.tensor,
axis: e.affineOf.axis), device), i32lit(1)),
i32lit(e.tmin.int))
return infixF("+", infixF("div", num, i32lit(e.coef.int)), i32lit(1))
var dimN: NimNode
if device:
dimN = callI32(layoutField(e.tensor, "shape", e.axis))
else:
dimN = nnkCall.newTree(id"tdim", id(e.tensor), newLit(e.axis))
return dimN
let num = infixF("-", infixF("-", dimN, i32lit(1)), i32lit(e.tmin.int))
infixF("+", infixF("div", num, i32lit(e.coef.int)), i32lit(1))
# ─────────────────────────────────────────────────────────────────────────
# Tile-program derivation from the IR + schedule
# ─────────────────────────────────────────────────────────────────────────
#
# AlgoIR + schedule
# | TilePlan (makeTilePlan)
#
# v per-read gather, gd view + loadTileBounded (emitGather)
# | tile combine, mul / convert / addScaled (emitTileValue)
# v accumulator addScaled or storeTileMasked (emitTileKernel)
#
# Tile plan, the last two out dims mapping onto the register-tile plane, earlier
# out dims onto the gd view's batch and depth slots, the reduce space into
# one chunked iter plus serial iters, serial coordinates added to the reads'
# base offsets, zero out dims yielding the all-reduce variants.
#
# Per read axis, VARYING keeps the axis stride with the tile origin carrying
# the grid or chunk coordinate and the load limit at the axis range,
# BROADCAST has stride 0 and origin 0, loadTileBounded's limits absolute
# over the view.
#
# A value with no operand (just a number) is a reported gap, tiles cannot
# be filled from a number alone.
#
# Affine plane slot example, index `2*j + 1` over plane iter j
#
# View coordinate stays the plane iterator, shifted by v0 = ceil(-T/coef)
# (T the slot's serial and constant terms), keeping the origin
# coord*tileExt - v0 and the limit ((dim - 1 - T) div coef + 1) - v0 inside
# the read axis, (coef*v0 + T) * stride added to the base offset.
proc hasBinding(rd: ReadIR, name: string): bool =
## Whether any read slot's index form binds `name`.
##
## Affine forms and scalar-term indices included.
for ix in rd.indices:
if ix.kind == ikIter and ix.name == name:
return true
if ix.kind == ikAffine:
for t in ix.terms:
if t.iter == name:
return true
for st in ix.scalars:
for six in st.indices:
if six.kind == ikIter and six.name == name:
return true
if six.kind == ikAffine:
for t in six.terms:
if t.iter == name:
return true
false
proc slotIterCoef(rd: ReadIR, slot: int, name: string): int32 =
## Coefficient of iterator `name` in read slot `slot`'s index form, 0
## when the slot does not bind it.
let ix = rd.indices[slot]
case ix.kind
of ikIter: (if ix.name == name: 1'i32 else: 0'i32)
of ikAffine:
for t in ix.terms:
if t.iter == name:
return t.coef
0
else: 0
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
of ikAffine:
var f = ""
for t in ix.terms:
f &= (if f.len == 0: "" else: "+") & $t.coef & "*" & t.iter
f & (if ix.cst != 0: (if ix.cst < 0: "-" else: "+") & $abs(ix.cst)
else: ""))
result &= ")"
proc affineSlot(rd: ReadIR, name: string): int =
## Slot whose affine index form binds `name`, -1 when the slot binding it is
## pure or absent.
for k, ix in rd.indices:
if ix.kind == ikAffine:
for t in ix.terms:
if t.iter == name:
return k
-1
proc slotTermNode(rd: ReadIR, slot: int, planeIter: string, resolve: Table[string, NimNode], skipIters: seq[string] = @[]): NimNode =
## An affine slot's non-plane terms as one absolute offset node.
## Returns nil when pure.
##
## - the slot's serial iterators' coordinates and constant summed
## - `skipIters` drops terms the caller accounts elsewhere, prefix dims
## sitting on the gd view's batch and depth slot strides
## - the base must not repeat them while the origin, limit and v0 shift
## still see them, prefix dims sharing the slot included
let ix = rd.indices[slot]
doAssert ix.kind == ikAffine
var acc: NimNode = nil
proc add(c: int, node: NimNode) =
var term = node
if c != 1:
term = infixF("*", i32lit(c), term)
acc = (if acc == nil: term else: infixF("+", acc, term))
for t in ix.terms:
if t.iter == planeIter or t.iter in skipIters:
continue
doAssert t.iter in resolve,
"index form term '" & t.iter & "' of '" & rd.tensor & "' is neither on " &
"the tile plane nor a resolved serial iterator"
add(t.coef.int, resolve[t.iter])
if ix.cst != 0:
add(ix.cst.int, i32lit(1))
acc
proc makeTilePlan(ir: AlgoIR): TilePlan =
## TilePlan derived from the IR's out dims and reduce iters, the out dims
## mapped onto the gd view's plane and batch/depth slots.
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]` with the reduce domain declared")
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 > 3:
raiseAssert(
"reported gap: output rank above 5 has no grid axis for the " &
"prefix dims, the gd view's batch and depth slots cover two and " &
"the gather base carries the third. Reproduce with a six-dim " &
"recipe")
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.hasBinding(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
# L>=2 stages chunk the last bound reduce iter too, the tile plane
# staying (row, col), the accumulator sweeping across chunk iterations
# and the stage's chunk guard applying to it
if result.chunkIter.len == 0:
for it in countdown(ir.reduceIters.len - 1, 0):
var bound = false
for rd in ir.reads:
if rd.hasBinding(ir.reduceIters[it]): bound = true
if bound:
result.chunkIter = ir.reduceIters[it]
break
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`, weighted
## by each slot's index-form coefficient, 0 when unbound (a broadcast axis).
var acc: NimNode = nil
for k in 0 ..< rd.indices.len:
let c = rd.slotIterCoef(k, name)
if c == 0: continue
var term = callI32(layoutField(rd.tensor, "stride", k))
if c != 1:
term = infixF("*", i32lit(c.int), term)
acc = (if acc == nil: term else: infixF("+", acc, term))
(if acc == nil: i32lit(0) else: acc)
proc readBase(rd: ReadIR, resolve: Table[string, NimNode], skip: seq[int] = @[]): NimNode =
## int32 offset added to the gather's buffer pointer, the resolved serial
## iterators and constants weighted in, `skip` slots contributing nothing.
var acc: NimNode = nil
proc addTerm(c: int, node: NimNode, k: int) =
var term = infixF("*", node, callI32(layoutField(rd.tensor, "stride", k)))
if c != 1:
term = infixF("*", i32lit(c), term)
acc = (if acc == nil: term else: infixF("+", acc, term))
for k in 0 ..< rd.indices.len:
if k in skip: continue
let ix = rd.indices[k]
case ix.kind
of ikConst:
addTerm(1, i32lit(ix.value.int), k)
of ikIter:
if ix.name in resolve:
addTerm(1, resolve[ix.name], k)
of ikAffine:
for t in ix.terms:
if t.iter in resolve:
addTerm(t.coef.int, resolve[t.iter], k)
if ix.cst != 0:
addTerm(ix.cst.int, i32lit(1), k)
(if acc == nil: i32lit(0) else: acc)
proc intLitVal(n: NimNode): int =
(if n.kind == nnkIntLit: n.intVal.int else: int32.low.int)
proc affinePlaneGeometry(rd: ReadIR, slot: int, planeIter: string, coord: NimNode, tileExt: int, resolve: Table[string, NimNode]): tuple[stride, v0, origin, limit: NimNode] =
## Stride, view coordinate, origin and limit of one affine plane slot,
## the raw index coef*planeIter + T staying valid inside the read axis.
##
## The view coordinate shifts down by v0 = ceil(-T/coef) so the read
## axis' first valid cell sits at the tile's first lane:
##
## - negative coordinates zero-fill through the bounded loader
## - the load limit is the axis bound in the shifted coordinate
let coef = rd.slotIterCoef(slot, planeIter).int
doAssert coef > 0,
"a non-positive plane coefficient has no monotone lane mapping"
result.stride = rd.slotStride(planeIter)
let tN = rd.slotTermNode(slot, planeIter, resolve)
var v0: NimNode
if tN == nil:
v0 = i32lit(0)
elif tN.kind == nnkIntLit and tN.intVal >= 0:
v0 = i32lit(0)
elif tN.kind == nnkIntLit:
v0 = i32lit((-tN.intVal + coef - 1) div coef)
else:
v0 = nnkIfExpr.newTree(
nnkElifBranch.newTree(infixF("<", tN, i32lit(0)),
infixF("div", infixF("+", infixF("*", i32lit(-1), tN),
i32lit(coef - 1)), i32lit(coef))),
nnkElseExpr.newTree(i32lit(0)))
result.v0 = v0
var origin = infixF("*", coord, i32lit(tileExt))
if not (v0.kind == nnkIntLit and v0.intVal == 0):
origin = infixF("-", origin, v0)
result.origin = origin
let dimN = callI32(layoutField(rd.tensor, "shape", slot))
var num = infixF("-", dimN, i32lit(1))
if tN != nil:
num = infixF("-", num, tN)
var limit = infixF("+", infixF("div", num, i32lit(coef)), i32lit(1))
if not (v0.kind == nnkIntLit and v0.intVal == 0):
limit = infixF("-", limit, v0)
result.limit = limit
proc planeBaseTerm(rd: ReadIR, slot: int, planeIter: string, v0: NimNode, resolve: Table[string, NimNode],
skipIters: seq[string] = @[]): NimNode =
## A plane slot's base-offset contribution, (coef*v0 + T) * stride.
## T covers the slot's serial and constant terms, nil when nothing contributes.
## `skipIters` drops terms the view's batch and depth slot strides already carry.
let ix = rd.indices[slot]
doAssert ix.kind == ikAffine
var coefV0: NimNode = nil
if not (v0.kind == nnkIntLit and v0.intVal == 0):
let coef = rd.slotIterCoef(slot, planeIter).int
coefV0 = infixF("*", i32lit(coef), v0)
let tN = rd.slotTermNode(slot, planeIter, resolve, skipIters)
var acc: NimNode = nil
for cand in [coefV0, tN]:
if cand == nil: continue
if cand.kind == nnkIntLit and cand.intVal == 0: continue
acc = (if acc == nil: cand else: infixF("+", acc, cand))
if acc == nil: return nil
var term = callI32(layoutField(rd.tensor, "stride", slot))
term = infixF("*", acc, term)
term
proc gdViewNode(viewName, tensorName: string, bStr, dStr, rStr, cStr, base: NimNode): NimNode =
## Emits `let <viewName> = mkGdView(offsetBuf(<tensor>.data, base), ...)`, 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 with derived
## strides and `loadTileBounded` at the tile-grid-coordinate origin.
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
var rVaries, cVaries: bool
var rSlot = affineSlot(rd, plan.rowDim)
let cPlaneName = (if plan.colDim.len > 0: plan.colDim else: plan.chunkIter)
var cSlot = (if cPlaneName.len > 0: affineSlot(rd, cPlaneName) else: -1)
if plan.flatRect or plan.allReduce:
doAssert rSlot < 0 and cSlot < 0,
"reported gap: affine index forms in the all-reduce tile programs " &
"have no derived origin; reproduce with a semi-affine single-read " &
"all-reduce such as `s[] ++= a[2*i + 1]` with i declared"
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 range
rStr = id"atomN"
cStr = i32lit(1)
rOrig = id"rchunk"
cOrig = i32lit(0)
limR = id"flatRows"
limC = id"atomN"
else:
rVaries = rd.hasBinding(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))
if plan.colDim.len > 0:
cVaries = rd.hasBinding(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.hasBinding(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 range as the load limit.
# Broadcast axes take the tile range, loadTileBounded's limits being
# absolute over the view.
limR = (if rVaries:
callI32(rangeExpr(ir.iterRange(plan.rowDim), device = true))
else:
id"tileRc")
if plan.colDim.len > 0 and cVaries:
limC = callI32(rangeExpr(ir.iterRange(plan.colDim), device = true))
elif plan.chunkIter.len > 0 and cVaries:
limC = callI32(rangeExpr(ir.iterRange(plan.chunkIter), device = true))
else:
limC = id"tileCc"
# An affine plane slot keeps its plane iterator as the view coordinate,
# shifted by v0, the stride scaled by the form's coefficient.
# (coef*v0 + T) is added to the base offset, origin and limit absolute.
# The plane geometry's term resolution sees the prefix dims' coordinates
# (bCoord/dCoord). Their base contribution sits on the gd view's batch
# and depth slot strides, so planeBaseTerm drops those terms.
var geoRes = initTable[string, NimNode]()
for nm, v in resolve:
geoRes[nm] = v
if plan.prefixDims.len >= 1: geoRes[plan.prefixDims[0]] = id"bCoord"
if plan.prefixDims.len >= 2: geoRes[plan.prefixDims[1]] = id"dCoord"
if plan.chunkIter.len > 0: geoRes[plan.chunkIter] = id"chunk"
# A plane slot may mix prefix dims with serial terms, the prefix dims
# sitting on the view's batch and depth slot strides, the serial terms
# on the base offset. The v0 shift (computed over the full term sum)
# stays consistent per threadgroup, origin and limit being per-threadgroup.
var rPlane, cPlane: tuple[stride, v0, origin, limit: NimNode]
if rSlot >= 0:
rPlane = affinePlaneGeometry(rd, rSlot, plan.rowDim, id"by",
plan.tileR, geoRes)
if cSlot >= 0:
let cCoord = (if plan.colDim.len > 0: id"bx" else: id"chunk")
cPlane = affinePlaneGeometry(rd, cSlot, cPlaneName, cCoord,
plan.tileC, geoRes)
let loadAbs = rSlot >= 0 or cSlot >= 0
var skipSlots: seq[int] = @[]
if rSlot >= 0: skipSlots.add(rSlot)
if cSlot >= 0: skipSlots.add(cSlot)
var base = rd.readBase(resolve, skipSlots)
block:
var extras: seq[NimNode] = @[]
if rSlot >= 0: extras.add(planeBaseTerm(rd, rSlot, plan.rowDim, rPlane.v0,
geoRes, plan.prefixDims))
if cSlot >= 0: extras.add(planeBaseTerm(rd, cSlot, cPlaneName, cPlane.v0,
geoRes, plan.prefixDims))
for e in extras:
if e != nil and not (e.kind == nnkIntLit and e.intVal == 0):
base = (if base.kind == nnkIntLit and base.intVal == 0: e
else: infixF("+", base, e))
var stmts = nnkStmtList.newTree()
# Scalar terms load per tile origin, their values weighting the base offset.
# Every bound iterator resolves to a threadgroup coordinate or to a chunk/serial loop variable, re-run per tile.
block:
var sres = initTable[string, NimNode]()
for nm, v in resolve:
sres[nm] = v
if plan.prefixDims.len >= 1: sres[plan.prefixDims[0]] = id"bCoord"
if plan.prefixDims.len >= 2: sres[plan.prefixDims[1]] = id"dCoord"
if plan.chunkIter.len > 0: sres[plan.chunkIter] = id"chunk"
var idx1 = 0
for k in 0 ..< rd.indices.len:
# plane slots contribute through the base as well, their scalar
# terms never entering the plane geometry's origin and limit
for st in rd.indices[k].scalars:
var acc: NimNode = nil
proc addI(c: int, node: NimNode) =
var term = node
if c != 1:
term = infixF("*", i32lit(c), term)
acc = (if acc == nil: term else: infixF("+", acc, term))
for j, six in st.indices:
case six.kind
of ikConst:
if six.value != 0: addI(1, i32lit(six.value.int))
of ikIter, ikAffine:
var terms: seq[tuple[coef: int32, name: string]]
var cst2 = 0'i32
if six.kind == ikIter:
terms = @[(1'i32, six.name)]
else:
terms = @[]
for t2 in six.terms:
terms.add((t2.coef, t2.iter))
cst2 = six.cst
for t2 in terms:
doAssert t2.name in sres,
"scalar term of '" & rd.tensor & "' binds '" & t2.name &
"', which is no serial loop, chunk loop, prefix dim or " &
"tile-plane dimension of the stage"
# each term weighted by the scalar tensor's slot stride,
# the flat element index like any read
addI(t2.coef.int, nnkInfix.newTree(id"*", sres[t2.name],
callI32(layoutField(st.tensor, "stride", j))))
if cst2 != 0:
addI(cst2.int, i32lit(1))
let idxN = (if acc == nil: i32lit(0) else: acc)
let scName = "sc" & $rdIdx & "_" & $k & "_" & $idx1
inc idx1
var contrib = infixF("*", id(scName), callI32(layoutField(t, "stride", k)))
if st.coef != 1:
contrib = infixF("*", i32lit(st.coef.int), contrib)
# the loaded value converts to an index, the host element being
# a float while offsets are element counts
stmts.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id(scName),
newEmptyNode(),
callI32(nnkBracketExpr.newTree(dot(id(st.tensor), "data"), idxN)))))
base = (if base.kind == nnkIntLit and base.intVal == 0: contrib
else: infixF("+", base, contrib))
stmts.add(gdViewNode(gdName, t, bStr, dStr, rStr, cStr, base))
stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tileName),
nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
newEmptyNode())))
if loadAbs:
# absolute origins, per axis kind
# - bound to a linear index form (like 2*oh + kh), the origin comes
# from that form's geometry
# - a plain iterator axis, the tile-unit origin converts to absolute units
# - a broadcast axis, the origin is zero (the read never moves)
# The view coordinate j' = j - v0 can start below zero when the slot's
# serial term is negative (a paged row form j - 8*b over the chunk b):
# clamp the load window at zero, the out-of-window cells zero-filling,
# so the plane geometry's absolute bounds stay intact.
if rVaries:
rOrig = (if rSlot >= 0: rPlane.origin
else: infixF("*", id"by", id"tileRc"))
if rSlot >= 0:
rOrig = rPlane.origin
limR = rPlane.limit
else:
rOrig = i32lit(0)
limR = id"tileRc"
if cVaries:
let cCoord = (if plan.colDim.len > 0: id"bx" else: id"chunk")
cOrig = (if cSlot >= 0: cPlane.origin
else: infixF("*", cCoord, id"tileCc"))
if cSlot >= 0:
cOrig = cPlane.origin
limC = cPlane.limit
else:
cOrig = i32lit(0)
limC = id"tileCc"
stmts.add(nnkCall.newTree(nnkDotExpr.newTree(id(tileName),
ident"loadTileBoundedAbs"), 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))
else:
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 tileFrag(t: NimNode, lane: NimNode): NimNode =
## Lane cell of a tile, `t.frags[n][m].frag[v]`.
nnkBracketExpr.newTree(
nnkDotExpr.newTree(
nnkBracketExpr.newTree(
nnkBracketExpr.newTree(nnkDotExpr.newTree(t, ident"frags"), ident"n"),
ident"m"),
ident"frag"),
lane)
proc emitLaneMap2*(stmts: NimNode, dst, src1, src2, rows, cols: NimNode, f: proc(x, y: NimNode): NimNode) =
## Per-lane dst[i] = f(src1[i], src2[i]) over the (n, m, v) fragment loops:
##
## - emitter-owned x/y gensyms, never the expansion caller's scope idents
## (an untyped map2 lambda binds them there)
## - rows / cols the tile's lane-range expressions (tileRc, tileCc)
let x = genSym(nskLet, "laneX")
let y = genSym(nskLet, "laneY")
var inner = nnkStmtList.newTree(
nnkLetSection.newTree(nnkIdentDefs.newTree(x, newEmptyNode(),
tileFrag(src1, ident"v"))),
nnkLetSection.newTree(nnkIdentDefs.newTree(y, newEmptyNode(),
tileFrag(src2, ident"v"))),
nnkAsgn.newTree(tileFrag(dst, ident"v"), f(x, y)))
var mid = nnkStmtList.newTree(inner)
mid = nnkForStmt.newTree(ident"v",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32", ident"atomVpt")), mid)
var outer = nnkForStmt.newTree(ident"m",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32",
infixF("div", cols, ident"atomN"))), mid)
outer = nnkForStmt.newTree(ident"n",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32",
infixF("div", rows, ident"atomM"))), outer)
stmts.add(outer)
proc laneRowExpr(laneRow: NimNode): NimNode =
## A lane cell's row coordinate, the tile row the (n, lane) position reads.
infixF("+", infixF("*", id"by", id"tileRc"),
infixF("+", laneRow, infixF("*", ident"n", ident"atomM")))
proc laneColExpr(laneCol: NimNode, colBase: NimNode): NimNode =
## A lane cell's column coordinate, the chunk or grid column base plus
## the lane's column and the (m, v) fragment offsets.
var col = infixF("+", laneCol, infixF("*", ident"m", ident"atomN"))
col = infixF("+", col, ident"v")
infixF("+", colBase, col)
proc condScalarLoads(n: LuxNode, resolve: Table[string, NimNode], stmts: NimNode): Table[string, NimNode] =
## Loads a condition's TensorRead scalar terms per tile origin, keyed
## by node id.
##
## `condToNim` and `guardToNim` reference them by that key.
##
## - the scalar's flat element index is stride-weighted like any read
## - its index iterators resolve through `resolve` (prefix coordinates, chunk loop variable, serial iters)
var accT = initTable[string, NimNode]()
proc addI(c: int, node: NimNode): NimNode =
var term = node
if c != 1:
term = infixF("*", i32lit(c), term)
term
proc walk(x: LuxNode) =
case x.kind
of TensorRead:
let key = "sc" & $x.id
if key in accT: return
var sacc: NimNode = nil
var tensor = x.children[0].fn.symbol
for i in 1 ..< x.children.len:
# each slot's index terms add their stride-weighted coordinate
# contribution to the flat element index, slot i-1 being axis i-1
let ix = indexFromLux(x.children[i], initTable[string, int32]())
let ax = i - 1
var terms: seq[tuple[coef: int32, name: string]]
var cst = 0'i32
case ix.kind
of ikConst:
if ix.value != 0:
cst = ix.value.int32
of ikIter:
terms = @[(1'i32, ix.name)]
of ikAffine:
terms = @[]
for t in ix.terms:
terms.add((t.coef, t.iter))
cst = ix.cst
for t in terms:
doAssert t.name in resolve,
"condition scalar of '" & tensor & "' binds '" & t.name &
"', which is no prefix dim, chunk loop or serial iterator"
var term = infixF("*", resolve[t.name],
callI32(layoutField(tensor, "stride", ax)))
if t.coef != 1:
term = infixF("*", i32lit(t.coef.int), term)
sacc = (if sacc == nil: term else: infixF("+", sacc, term))
if cst != 0:
sacc = (if sacc == nil: i32lit(cst.int)
else: infixF("+", sacc, i32lit(cst.int)))
let idxN = (if sacc == nil: i32lit(0) else: sacc)
let load = callI32(nnkBracketExpr.newTree(dot(id(tensor), "data"), idxN))
stmts.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id(key),
newEmptyNode(), load)))
accT[key] = id(key)
of BinOp:
walk(x.children[1])
walk(x.children[2])
else: discard
walk(n)
result = accT
proc guardToNim*(n: LuxNode, resolve: Table[string, NimNode], scalars = initTable[string, NimNode]()): NimNode =
## Returns the Nim expression of a stage-guard condition, parameters
## already folded:
##
## - iterators resolved through the caller's coordinate table, the chunk
## loop variable and the grid coordinates
## - TensorRead scalar terms through the loaded lets `condScalarLoads` built
## - comparisons as infix ops, `a < b` already normalized to the truthy
## `b - a` form by the DSL
# walk the condition tree:
proc walk(x: LuxNode): NimNode =
case x.kind
of IntLit:
i32lit(x.intVal)
of Domain:
doAssert x.iter.symbol in resolve,
"the stage guard binds '" & x.iter.symbol &
"', which is not the chunk iterator or a grid coordinate"
resolve[x.iter.symbol]
of TensorRead:
doAssert ("sc" & $x.id) in scalars,
"the stage guard's scalar terms must be loaded before the guard"
scalars["sc" & $x.id]
of BinOp:
let l = walk(x.children[1])
let r = walk(x.children[2])
case x.children[0].bopKind
of Add: infixF("+", l, r)
of Sub: infixF("-", l, r)
of Mul: infixF("*", l, r)
of Le: infixF("<=", l, r)
of Eq: infixF("==", l, r)
else:
raiseAssert("unsupported stage guard node kind " & $x.kind)
walk(n)
proc condToNim*(n: LuxNode, laneCoords: Table[string, char], colBase: NimNode, laneRow: NimNode, laneCol: NimNode, resolve = initTable[string, NimNode](), scalars = initTable[string, NimNode]()): NimNode =
## Returns the mux condition as a lane-resolved Nim expression:
##
## - row-dim iters resolve to the row coordinate (the n loop), col-dim
## and chunk iters to the column coordinate (the m, v loops)
## - other iterators resolve through `resolve` (prefix coordinates, chunk loop variable, serial iters) and TensorRead scalar terms
## through the lets `condScalarLoads` built, plane-constant per cell
## - comparisons as infix ops, `a < b` already normalized to the truthy
## `b - a` form by the DSL
# walk the condition tree:
proc walk(x: LuxNode): NimNode =
case x.kind
of IntLit:
i32lit(x.intVal)
of Domain:
case laneCoords.getOrDefault(x.iter.symbol, '?')
of 'r': laneRowExpr(laneRow)
of 'c': laneColExpr(laneCol, colBase)
else:
doAssert x.iter.symbol in resolve,
"the mux condition binds '" & x.iter.symbol &
"', which is neither the tile plane's row, its column, nor a " &
"resolved per-cell coordinate"
resolve[x.iter.symbol]
of TensorRead:
doAssert ("sc" & $x.id) in scalars,
"the mux condition's scalar terms must be loaded before the select"
scalars["sc" & $x.id]
of BinOp:
let l = walk(x.children[1])
let r = walk(x.children[2])
case x.children[0].bopKind
of Add: infixF("+", l, r)
of Sub: infixF("-", l, r)
of Mul: infixF("*", l, r)
of Le: infixF("<=", l, r)
of Eq: infixF("==", l, r)
else:
raiseAssert("unsupported condition node kind " & $x.kind)
walk(n)
proc emitLaneMux*(stmts: NimNode, dst, thenT, thenV, elseT, elseV: NimNode, cond: LuxNode, laneCoords: Table[string, char], colBase: NimNode,
resolve = initTable[string, NimNode](), scalars = initTable[string, NimNode]()) =
## Per-lane select over the (n, m, v) fragment loops:
##
## The select's contract sets every dst cell, reading every src cell unchanged.
## The condition's iterators resolve to lane coordinates inside the (n, m, v) loops,
## a constant branch marked by the zero tile node, its value node the per-cell constant.
var usesRow = false
var usesCol = false
proc scan(x: LuxNode) =
case x.kind
of Domain:
case laneCoords.getOrDefault(x.iter.symbol, '?')
of 'r': usesRow = true
of 'c': usesCol = true
else: discard
of BinOp:
scan(x.children[1]); scan(x.children[2])
else: discard
scan(cond)
var pre = nnkStmtList.newTree()
var laneRow: NimNode
var laneCol: NimNode
if usesRow or usesCol:
pre.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id"muxLane", newEmptyNode(),
nnkCall.newTree(id"int", id"tid"))))
pre.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id"muxCell", newEmptyNode(),
nnkCall.newTree(id"toIntVal", nnkCall.newTree(
nnkCall.newTree(nnkDotExpr.newTree(id"atom", ident"getLayoutA")),
nnkPar.newTree(id"muxLane", newLit(0)))))))
if usesRow: laneRow = infixF("mod", id"muxCell", id"atomM")
if usesCol: laneCol = infixF("div", id"muxCell", id"atomM")
# the lane coordinate lets only when the condition binds a plane iter,
# an unbound lane expr being nil
var inner = nnkStmtList.newTree(pre)
if usesRow:
inner.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id"muxRow",
newEmptyNode(), laneRowExpr(laneRow))))
if usesCol:
inner.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id"muxCol",
newEmptyNode(), laneColExpr(laneCol, colBase))))
inner.add(
nnkLetSection.newTree(nnkIdentDefs.newTree(id"muxCnd", newEmptyNode(),
condToNim(cond, laneCoords, colBase, laneRow, laneCol, resolve,
scalars))),
nnkLetSection.newTree(nnkIdentDefs.newTree(id"thenCell", newEmptyNode(),
(if thenT.kind == nnkIntLit and thenT.intVal == 0: thenV
else: tileFrag(thenT, ident"v")))),
nnkLetSection.newTree(nnkIdentDefs.newTree(id"elseCell", newEmptyNode(),
(if elseT.kind == nnkIntLit and elseT.intVal == 0: elseV
else: tileFrag(elseT, ident"v")))),
nnkAsgn.newTree(tileFrag(dst, ident"v"),
nnkIfExpr.newTree(
nnkElifBranch.newTree(id"muxCnd", id"thenCell"),
nnkElseExpr.newTree(id"elseCell"))))
var mid = nnkStmtList.newTree(inner)
mid = nnkForStmt.newTree(ident"v",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32", ident"atomVpt")), mid)
var outer = nnkForStmt.newTree(ident"m",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32",
infixF("div", id"tileCc", ident"atomN"))), mid)
outer = nnkForStmt.newTree(ident"n",
nnkInfix.newTree(id"..<", i32lit(0),
nnkCall.newTree(id"int32",
infixF("div", id"tileRc", ident"atomM"))), outer)
stmts.add(outer)
proc emitTileValue(e: ExprIR, tileOf: Table[string, string], stmts: NimNode, counter: var int,
laneCoords: Table[string, char] = initTable[string, char](),
colBase: NimNode = nil,
resolve = initTable[string, NimNode](),
condScalars = initTable[string, NimNode]()):
tuple[tile: string, scalar: NimNode] =
## Statement value expression as one tile:
##
## - operand tiles combine through mul and addScaled
## - float literals carry the scalar multiplier
## - the mux's condition reads its iterators through `laneCoords`,
## the column base being `colBase`
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, laneCoords, colBase, resolve, condScalars)
let (rt, rsc) = emitTileValue(e.rhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
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` with i, j declared")
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, laneCoords, colBase, resolve, condScalars)
let (rt, rsc) = emitTileValue(e.rhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
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` with i, j declared")
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]` with i, j declared"
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)
of xeSub, xeDiv, xeMax:
let (lt0, lsc) = emitTileValue(e.lhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
let (rt0, rsc) = emitTileValue(e.rhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
doAssert lt0.len > 0 and rt0.len > 0,
"reported gap: a scalar operand of a tilewise map has no broadcast " &
"fill; reproduce with `o[i, j] ++= a[i, j] - 1.0 | (i, j)`"
let lt = applyScalar(lt0, lsc, stmts, counter)
let rt = applyScalar(rt0, rsc, stmts, counter)
let tmp = "wt" & $counter
inc counter
stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tmp),
nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
newEmptyNode())))
let f: proc(x, y: NimNode): NimNode = (case e.kind
of xeSub: (proc (x, y: NimNode): NimNode = infixF("-", x, y))
of xeDiv: (proc (x, y: NimNode): NimNode = infixF("/", x, y))
else: (proc (x, y: NimNode): NimNode =
nnkCall.newTree(id"max", x, y)))
emitLaneMap2(stmts, id(tmp), id(lt), id(rt), id"tileRc", id"tileCc", f)
(tile: tmp, scalar: nil)
of xeMux:
let (tt, tsc) = emitTileValue(e.lhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
let (et, esc) = emitTileValue(e.rhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
# a literal branch carries its value per cell, the zero tile node
# marking the constant branch
var tn: NimNode = i32lit(0)
var tv: NimNode = nil
if tt.len > 0:
tn = id(applyScalar(tt, tsc, stmts, counter))
else:
tv = tsc
var en: NimNode = i32lit(0)
var ev: NimNode = nil
if et.len > 0:
en = id(applyScalar(et, esc, stmts, counter))
else:
ev = esc
let tmp = "xt" & $counter
inc counter
stmts.add(nnkVarSection.newTree(nnkIdentDefs.newTree(id(tmp),
nnkCall.newTree(id"rt_l", id"T", id"tileRc", id"tileCc"),
newEmptyNode())))
emitLaneMux(stmts, id(tmp), tn, tv, en, ev, e.cond, laneCoords, colBase,
resolve, condScalars)
(tile: tmp, scalar: nil)
of xeExp2:
let (lt0, lsc) = emitTileValue(e.lhs, tileOf, stmts, counter, laneCoords, colBase, resolve, condScalars)
doAssert lt0.len > 0,
"reported gap: exp2 of a literal has no broadcast fill"
let lt = applyScalar(lt0, lsc, stmts, counter)
let tmp = "wt" & $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(id"exp2", id(tmp), id(lt)))
(tile: tmp, scalar: nil)
# Threadgroup staging primitive, a declaration the entry hoists and the device fn receives as a threadgroup-pointer parameter.
proc smemDecl(name: string, rows, cols: int): NimNode =
## `var name: array[rows * cols, HostElem] {.smem.}` declaration, hoisted
## into the entry via the launch spec, passed to the device fn.
nnkVarSection.newTree(nnkIdentDefs.newTree(
nnkPragmaExpr.newTree(id(name), nnkPragma.newTree(id"smem")),
nnkBracketExpr.newTree(id"array", newLit(rows * cols), id"HostElem"),
newEmptyNode()))
# ─────────────────────────────────────────────────────────────────────────
# mma pass boundary
# ─────────────────────────────────────────────────────────────────────────
# The mma optimization pass lives in `luxtiles_passes_mma`, reading this module's
# IR and helper surface. The import sits below every helper the pass module reads,
# the sibling compiling against the declarations above it.
#
# schedule passes -> mmaOptPass (luxtiles_passes_mma) -> lowerEmit
#
# A pass-matched statement lowers to the staged (threadgroup memory)
# or register mma form, everything else to the generic tile program.
# `-d:noMmaPass` drops the pass from every pipeline.
export AlgoIR, MmaMatch, ReadIR, IndexRef, IdxKind, Range, RangeKind,
ExprKind, ExprIR, LoopRole, LoopIR, CachePlan, Pass, PassKind
export outerTile, iterRange, rangeExpr, infixF, i32lit, dot, callI32,
layoutField, gdViewNode, slotStride, readBase, affineSlot, slotTermNode,
affinePlaneGeometry, planeBaseTerm, smemDecl, mkGdView, offsetBuf,
hasBinding, slotIterCoef, tileFrag, emitLaneMap2, emitLaneMux,
laneRowExpr, laneColExpr, condToNim, guardToNim, foldParams
import workspace/experimental/tensorcery_luxtiles/luxtiles_passes_mma
# ─────────────────────────────────────────────────────────────────────────
# Kernel body, the derived tile program
# ─────────────────────────────────────────────────────────────────────────
proc condScalarMap(e: ExprIR, resolve: Table[string, NimNode], stmts: NimNode): Table[string, NimNode] =
## Loads every mux condition's TensorRead scalar terms of an expression
## tree into `stmts`, keyed by node id, for `condToNim` to reference.
var acc = initTable[string, NimNode]()
proc walk(e: ExprIR) =
if e == nil: return
if e.kind == xeMux:
let loaded = condScalarLoads(e.cond, resolve, stmts)
for k, v in loaded:
acc[k] = v
return
if e.kind in {xeAdd, xeMul, xeSub, xeDiv, xeExp2, xeMax}:
walk(e.lhs)
walk(e.rhs)
walk(e)
result = acc
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.
##
## - the emitter names the tile ranges, grid coordinates, loop variables,
## prefix-dim coordinates, plane ranges, output stride and lane cell
# A two-out-dim max reduce lowers like any other reduce, the acc's
# columns being the col dim and the reduce space serial or chunked
# past it. Only a chunk iter owning the col dim's columns has no form
var body = nnkStmtList.newTree()
# These consts are built directly, the later fragments referencing them
# by name (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, bz split over the batch and depth ranges
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"bRange", newEmptyNode(),
rangeExpr(ir.iterRange(plan.prefixDims[1]), device = true))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"bCoord", newEmptyNode(), infixF("div", id"bz", id"bRange"))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"dCoord", newEmptyNode(), infixF("mod", id"bz", id"bRange"))))
of 3:
# gridZ = e0*e1*e2, prefix0 major, the third prefix's coordinate
# advancing the gathers' base offsets (the gd view has two slots)
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"dbRange", newEmptyNode(),
infixF("*",
rangeExpr(ir.iterRange(plan.prefixDims[1]), device = true),
rangeExpr(ir.iterRange(plan.prefixDims[2]), device = true)))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"bCoord", newEmptyNode(), infixF("div", id"bz", id"dbRange"))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"dRange", newEmptyNode(),
rangeExpr(ir.iterRange(plan.prefixDims[2]), device = true))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"dCoord", newEmptyNode(), infixF("mod",
infixF("div", id"bz", id"dRange"),
rangeExpr(ir.iterRange(plan.prefixDims[1]), device = true)))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"zCoord", newEmptyNode(), infixF("mod", id"bz", id"dRange"))))
else: discard
let L = ir.outDims.len
if L >= 1:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"rowRange", newEmptyNode(),
rangeExpr(ir.iterRange(plan.rowDim), device = true))))
if L >= 2:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"colRange", newEmptyNode(),
rangeExpr(ir.iterRange(plan.colDim), device = true))))
if plan.flatRect:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"flatRows", newEmptyNode(),
callI32(infixF("div",
infixF("+",
rangeExpr(ir.iterRange(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())))
if ir.reduceKind == rkMax:
# the max chain seeds the accumulator with the most-negative finite
# fp32, any finite tile value replacing it on the first combine
body.add(nnkCall.newTree(id"map", id(accName), id(accName),
nnkPrefix.newTree(id"-", newLit(3.402823466e38'f32))))
else:
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)
# the third prefix dim advances the gathers' base offsets, resolving
# like a serial iterator's coordinate
if plan.prefixDims.len >= 3: resolve[plan.prefixDims[2]] = id"zCoord"
# with a col dim the chunk iter's cells are nowhere in the tile, its non-plane
# slot terms entering the gather base per chunk step.
# Without a col dim the chunk iter IS the tile's column axis and the cPlane covers it.
if plan.chunkIter.len > 0 and plan.colDim.len > 0:
resolve[plan.chunkIter] = id"chunk"
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
# mux lane coordinates, the row and column dims of the tile plane
var laneCoords = initTable[string, char]()
if plan.rowDim.len > 0: laneCoords[plan.rowDim] = 'r'
if plan.colDim.len > 0: laneCoords[plan.colDim] = 'c'
elif plan.chunkIter.len > 0: laneCoords[plan.chunkIter] = 'c'
let muxColBase = (if plan.colDim.len > 0: infixF("*", id"bx", id"tileCc")
elif plan.chunkIter.len > 0: infixF("*", id"chunk", id"tileCc")
else: i32lit(0))
# the condition context resolves the prefix dims' coordinates too, the value side's gather base terms (dStr) already carrying them
var condRes = resolve
for pi, pd in plan.prefixDims:
condRes[pd] = (case pi
of 0: id"bCoord"
of 1: id"dCoord"
else: id"zCoord")
let condScalars = condScalarMap(ir.expr, condRes, inner)
let (valueTile, valueScalar) = emitTileValue(ir.expr, tileOf, inner,
counter, laneCoords, muxColBase, condRes, condScalars)
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:
if ir.reduceKind == rkMax:
emitLaneMap2(inner, id(accName), id(accName), id(value),
id"tileRc", id"tileCc",
(proc (x, y: NimNode): NimNode = nnkCall.newTree(id"max", x, y)))
else:
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)
var oBase: NimNode = i32lit(0)
if L >= 3: oBat = callI32(layoutField("o", "stride", 0))
if L >= 4: oDep = callI32(layoutField("o", "stride", 1))
if L >= 5: oBase = infixF("*", id"zCoord",
callI32(layoutField("o", "stride", 2)))
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, oBase))
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"rowRange",
infixF("*", id"by", id"tileRc"))
else: id"rowRange"),
(if L >= 2: infixF("-", id"colRange",
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",
rangeExpr(ir.iterRange(plan.rowDim), device = true), stmts)
if plan.colDim.len > 0:
stmts = wrapFor("cchunk",
rangeExpr(ir.iterRange(plan.colDim), device = true), stmts)
elif plan.chunkIter.len > 0:
var chunkBody = stmts
if ir.guardNode != nil:
# the stage's condition guards whole chunk iterations, the flash
# block-skip, the chunk body running only when the predicate holds
var gres = initTable[string, NimNode]()
for it, v in resolve: gres[it] = v
gres[plan.chunkIter] = id"chunk"
if plan.rowDim.len > 0: gres[plan.rowDim] = id"by"
if plan.colDim.len > 0: gres[plan.colDim] = id"bx"
for pi in 0 ..< plan.prefixDims.len:
gres[plan.prefixDims[pi]] = (if pi == 0: id"bCoord"
elif pi == 1: id"dCoord"
else: id"zCoord")
# the guard's scalar terms load inside the chunk loop before the predicate, chunk-dependent scalars re-reading per iteration
let gsc = condScalarLoads(ir.guardNode, gres, stmts)
var gscNode = nnkStmtList.newTree()
for k, _ in gsc:
gscNode.add(stmts[stmts.len - 1])
stmts.del(stmts.len - 1)
chunkBody = nnkStmtList.newTree(gscNode,
nnkIfStmt.newTree(nnkElifBranch.newTree(
guardToNim(ir.guardNode, gres, gsc), stmts)))
# the chunk iter's cells are the tile's columns when the plane has no
# col dim (each chunk covering tileCc chunk values). With a col dim
# the chunk's cells are nowhere in the tile and the loop steps one
# chunk value per iteration, the reads' chunk term advancing per step
let chunkBound = (if plan.colDim.len == 0:
infixF("div",
infixF("+",
rangeExpr(ir.iterRange(plan.chunkIter), device = true),
infixF("-", id"tileCc", newLit(1))),
id"tileCc")
else:
rangeExpr(ir.iterRange(plan.chunkIter), device = true))
stmts = wrapFor("chunk", chunkBound, chunkBody)
for it in countdown(plan.serialIters.len - 1, 0):
stmts = wrapFor(plan.serialIters[it],
rangeExpr(ir.iterRange(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"
if ir.reduceKind == rkMax:
body.add(quote do:
var `rvI`: rv(HostElem, tileRc, tileCc)
`rvI`.neg_infty()
`rvI`.row_max(`accI`, `rvI`))
else:
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"))
if ir.reduceKind == rkMax:
body.add(quote do:
if tid == 0:
var total = -3.402823466e38'f32
for i in 0'i32 ..< int32(`combineBound`):
total = max(total, part[i])
o.data[0] = total)
else:
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"
if ir.reduceKind == rkMax:
body.add(quote do:
var `rvI`: rv(HostElem, tileRc, tileCc)
`rvI`.neg_infty()
`rvI`.row_max(`accI`, `rvI`))
else:
body.add(quote do:
var `rvI`: rv(HostElem, tileRc, tileCc)
`rvI`.row_sum(`accI`))
body.add(quote do:
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 < rowRange:
o.data[row * os1] = `rvI`.data[n * atomVpt])
elif L >= 2 and hasReduce:
var oBat, oDep: NimNode = i32lit(0)
var oBase: NimNode = i32lit(0)
if L >= 3: oBat = callI32(layoutField("o", "stride", 0))
if L >= 4: oDep = callI32(layoutField("o", "stride", 1))
if L >= 5: oBase = infixF("*", id"zCoord",
callI32(layoutField("o", "stride", 2)))
body.add(gdViewNode("gdO", "o", oBat, oDep,
callI32(layoutField("o", "stride", L - 2)),
callI32(layoutField("o", "stride", L - 1)), oBase))
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"rowRange", infixF("*", id"by", id"tileRc")),
infixF("-", id"colRange", infixF("*", id"bx", id"tileCc"))))
result.body = body
# ─────────────────────────────────────────────────────────────────────────
# Kernel emission, device fn + thin entry and the compile-time AST echo
# ─────────────────────────────────────────────────────────────────────────
# `lowerEmit` lowers the generic tile program to the (device fn, thin entry)
# pair and registers the launch spec the statement expansion reads.
#
# lowerEmit, TilePlan -> emitDeviceFunc (proc <name>Fn[T]) -> compile-time echo
#
# v emitEntry, geometry lets + call + DescSpecs -> compile-time 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 scalarTensors(ir: AlgoIR): seq[string] =
## Tensors read in index position (scalar terms), first-appearance order.
for rd in ir.reads:
for ix in rd.indices:
for st in ix.scalars:
if st.tensor notin result:
result.add(st.tensor)
proc condScalarTensors(ir: AlgoIR): seq[string] =
## Tensors read by mux conditions and the stage guard, so their views
## reach the kernel's buffer arguments like `scalarTensors`' reads.
var acc: seq[string] = @[]
proc walkLux(x: LuxNode) =
if x == nil: return
case x.kind
of TensorRead:
let t = x.children[0].fn.symbol
if t notin acc:
acc.add(t)
of BinOp:
walkLux(x.children[1])
walkLux(x.children[2])
else: discard
proc walkExpr(e: ExprIR) =
if e == nil: return
if e.kind == xeMux:
walkLux(e.cond)
return
if e.kind in {xeAdd, xeMul, xeSub, xeDiv, xeExp2, xeMax}:
walkExpr(e.lhs)
walkExpr(e.rhs)
walkLux(ir.guardNode)
walkExpr(ir.expr)
result = acc
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)
for t in ir.scalarTensors() & ir.condScalarTensors():
if t notin result:
result.add(t)
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.
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 d in ir.declArgs:
formals.add(nnkIdentDefs.newTree(id(d.argName), ident"int32",
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 or `{.device.}` pragma, the 7-child procdef shape passing
# semantic analysis where the 8-child parser shape does not, the pragma not
# instantiating here, crucible's pull path attaching device semantics.
let genParams = nnkGenericParams.newTree(
nnkIdentDefs.newTree(ident"T", newEmptyNode(), newEmptyNode()))
when defined(pocTrace):
var fs: seq[string] = @[]
for f in formals[1 .. ^1]:
for j in 0 ..< f.len - 2: fs.add($f[j])
echo "DEVFN formals ", ir.name, ": ", fs
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, registering the launch spec `evalCall` reads.
# Spec key, the kernel name, the template symbol suffixed `Entry` to stay
# distinct from the kernel wrapper embedding the statement.
let entryName = ir.name & "Entry"
let spec = nnkBracket.newTree(
nnkExprEqExpr.newTree(ident"output", ident"o"),
nnkExprEqExpr.newTree(ident"kernel", newLit(ir.name)),
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()))
for d in ir.declArgs:
formals.add(nnkIdentDefs.newTree(id(d.argName), ident"int32",
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 d in ir.declArgs:
call.add(id(d.argName))
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(entryName), newEmptyNode(), newEmptyNode(),
formals, newEmptyNode(), newEmptyNode(),
lowered)
echo "####################[ derived kernel: " & ir.name & " ]####################"
echo result.repr
echo "####################[ end derived kernel: " & ir.name & " ]####################"
proc tileProgramName(ir: AlgoIR): string =
## Lowering description of the generic tile program, from the IR's
## out dims, reduce iterators and read shapes.
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)"
proc lowerEmit(ir: var AlgoIR, pass: Pass): NimNode =
## Lowering core, one tile program per statement, generic plan or matched
## staged or register mma form, never an error path either way.
let mt = ir.mmaMatch
var body: NimNode
var smemDecls = nnkStmtList.newTree()
var scratch: seq[string] = @[]
let L = ir.outDims.len
var gridX, gridY, gridZ: NimNode
var staged = false
var plan: TilePlan
if ir.scatter:
# Atomic scatter program, one thread per source element:
# idx = bx*threads + tid, guard idx < range,
# cell = int32(<indexTensor>.data[idx]),
# atomic_fetch_add_explicit on the bit-cast output cell
pass.description = "atomic scatter program, one thread per source element"
body = nnkStmtList.newTree()
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(id"iIdx", newEmptyNode(),
infixF("+", infixF("*", id"bx", i32lit(ir.threads)), id"tid"))))
let ext = rangeExpr(ir.iterRange(ir.iters[0]), device = true)
body.add(nnkIfStmt.newTree(nnkElifBranch.newTree(
infixF("<", id"iIdx", ext), nnkStmtList.newTree(
nnkLetSection.newTree(nnkIdentDefs.newTree(id"cell", newEmptyNode(),
callI32(nnkBracketExpr.newTree(
dot(id(ir.reads[0].tensor), "data"), id"iIdx")))),
nnkPragma.newTree(nnkExprColonExpr.newTree(id"emit", nnkBracket.newTree(
newLit("atomic_fetch_add_explicit(reinterpret_cast<volatile device atomic_uint*>(&o.data["),
id"cell", newLit("]), "), newLit($ir.scatterValue & "U"),
newLit(", memory_order_relaxed);"))))))))
gridX = nnkCall.newTree(ident"cdiv",
rangeExpr(ir.iterRange(ir.iters[0]), device = false),
i32lit(ir.threads))
gridY = newLit(1)
gridZ = newLit(1)
elif not ir.hasMma:
pass.description = tileProgramName(ir)
plan = ir.makeTilePlan()
if ir.reduceKind == rkMax and plan.chunkIter.len > 0 and
plan.chunkIter == plan.colDim:
raiseAssert(
"reported gap: a max reduce chunking its own col dim has no " &
"generic lowering, the reduce space and the col dim both owning " &
"the tile columns; reproduce with `o[i, j] ++= max(a[i, j], 0.0)`")
(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",
rangeExpr(ir.iterRange(plan.colDim, ir.target), device = false),
i32lit(plan.tileC))
else:
gridX = newLit(1)
gridY = nnkCall.newTree(ident"cdiv",
rangeExpr(ir.iterRange(plan.rowDim, ir.target), device = false),
i32lit(plan.tileR))
case plan.prefixDims.len
of 1:
gridZ = rangeExpr(ir.iterRange(plan.prefixDims[0], ir.target), device = false)
of 2:
gridZ = infixF("*",
rangeExpr(ir.iterRange(plan.prefixDims[0], ir.target), device = false),
rangeExpr(ir.iterRange(plan.prefixDims[1], ir.target), device = false))
of 3:
gridZ = infixF("*",
rangeExpr(ir.iterRange(plan.prefixDims[0], ir.target), device = false),
infixF("*",
rangeExpr(ir.iterRange(plan.prefixDims[1], ir.target), device = false),
rangeExpr(ir.iterRange(plan.prefixDims[2], ir.target), device = false)))
else:
gridZ = newLit(1)
else:
staged = ir.isStagedEligible(mt)
pass.description = mmaFormName(ir, mt, staged)
if staged:
(body, smemDecls, scratch) = ir.emitMmaStaged(mt)
else:
body = ir.emitMmaRegister(mt)
gridX = nnkCall.newTree(ident"cdiv",
rangeExpr(ir.iterRange(mt.colDim, ir.target), device = false),
i32lit(mt.tiles[1]))
gridY = nnkCall.newTree(ident"cdiv",
rangeExpr(ir.iterRange(mt.rowDim, ir.target), device = false),
i32lit(mt.tiles[0]))
case mt.freeDims.len
of 0:
gridZ = newLit(1)
of 1:
gridZ = rangeExpr(ir.iterRange(mt.freeDims[0], ir.target), device = false)
else:
gridZ = infixF("*",
rangeExpr(ir.iterRange(mt.freeDims[0], ir.target), device = false),
rangeExpr(ir.iterRange(mt.freeDims[1], ir.target), device = false))
let grid = nnkPar.newTree(gridX, gridY, gridZ)
var outDimsHost: seq[NimNode] = @[]
if ir.scatter:
outDimsHost = @[callI32(i32lit(ir.scatterRange.int)), callI32(i32lit(1))]
elif L == 0:
outDimsHost = @[callI32(i32lit(1)), callI32(i32lit(1))]
elif L == 1:
# a one-out-dim target stores (n, 1), the range on axis 0, so a later
# stage's one-slot read addresses it, the flat order unchanged
outDimsHost = @[rangeExpr(ir.iterRange(ir.outDims[0], ir.target), device = false),
callI32(i32lit(1))]
else:
for d in 0 ..< L:
outDimsHost.add(rangeExpr(ir.iterRange(ir.outDims[d], ir.target),
device = false))
let outShape = nnkPrefix.newTree(ident"@", nnkBracket.newTree(outDimsHost))
# schedule provenance, which schedule values produced which tile types
var provenance: string
if ir.hasMma:
let (bm, bn, bk) = mt.tiles
provenance = "tiles y=" & $bm & " x=" & $bn & " k=" & $bk &
(if mt.freeDims.len > 0:
", free/batch dims " & mt.freeDims.join(",") & " on grid.z, "
else:
", ") &
(if staged:
"cache: smem staged tiles, acc rt_l(HostElem, " & $bm & ", " &
$bn & "), gemm per-atom-slab"
else:
"register-only: ta rt_l(HostElem, " & $bm & ", " & $bk & "), tb rt_r(" &
$bk & ", " & $bn & "), acc rt_l(HostElem, " & $bm & ", " & $bn & ")")
else:
if L == 0:
provenance = (if plan.flatRect:
"all-reduce, one threadgroup, atom-rectangle unrank + partial combine"
else:
"all-reduce, one threadgroup, chunked tiles + partial combine")
else:
provenance = "tile program " & $plan.tileR & "x" & $plan.tileC &
", derived gather/combine, atom from getTileConfig(HostElem, HostElem)" &
(if ir.reduceKind == rkMax: ", max reducer" else: "") &
(if ir.recurrence.len > 0: ", recurrence over " & ir.recurrence.join(",")
else: "")
if not SchedInfo.hasKey(ir.name):
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, the staged or register mma program when mmaOptPass matched,
## the generic tile program otherwise.
Pass(name: "lower",
description: "emit the tile program as a device fn + thin entry",
kind: pkLowering,
emit: lowerEmit)
# ─────────────────────────────────────────────────────────────────────────
# ─────────────────────────────────────────────────────────────────────────
# Host launch, the stage statement expansion
# ─────────────────────────────────────────────────────────────────────────
# One stage per launch, or one fused group per launch:
#
# - a Fn's first stage defines its buffer
# - a later stage reads the previous buffer, the previous stage's
# target local rebound by name, allocating a fresh output
# - cross-stage ordering needs no barrier, separate kernel launches
# ordering the stages
#
# stage -> stageToIR -> buildPasses -> runPipeline -> expansion
# recipe statements -> planFusionGroups -> lowerFusedGroup -> expansion
#
# buildPasses composes the stage's schedule into a pass list in written order,
# mma optimization (luxtiles_passes_mma) before the lowering, fusion
# planning (luxtiles_passes_fusion) above the per-stage pipelines.
# Generic tile program = the semantics, a pass on top an optimization.
proc buildPasses(sched: StageSchedule, withLower = true): seq[Pass] =
## Stage schedule -> ordered pass list, composed in written order.
## Lowering pass present only when `withLower` is true, the fusion
## planner's analysis running the same pipeline without it.
result = @[]
for d in sched.directives:
case d.kind
of ScTile:
result.add(tilePass(d.name,
(if d.value == 0: DefaultTile else: d.value.int)))
of ScCache:
result.add(cacheAnnotatePass(@[d.name]))
result.add(cachePlanPass())
of ScThreads:
result.add(threadsPass(d.value.int))
of ScParam:
discard # resolved into the index forms at parse time
of ScRange:
discard # resolved into the IR's rangePins at IR build (stageToIR)
of ScReduce:
discard # resolved into the stage's reducer at IR build (appendStage)
of ScAtomic:
discard # resolved into the IR's scatter fields at stage parse
of ScGuard:
discard # resolved into the IR's guardNode at stage parse
of ScOrder:
discard # resolved into the stage's iter order at parse (loopNest)
else:
raiseAssert("schedule kind " & $d.kind & " has no pass wired yet")
when not defined(noMmaPass):
result.add(mmaOptPass())
if withLower:
result.add(lowerPass())
proc analyzeStage*(fn: Fn, stageIdx: int, sched: StageSchedule, kernelName: string): AlgoIR =
## One stage's AlgoIR, schedule + mma passes applied and no lowering,
## the structural-facts source of the fusion planner.
crucibleCompileTarget = ctMetal
result = stageToIR(fn, stageIdx, sched, kernelName)
let threads = stageThreads(sched)
result.threads = (if threads > 0: threads else: DefaultThreads)
discard runPipeline(result, buildPasses(sched, withLower = false))
# ─────────────────────────────────────────────────────────────────────────
# Fusion pass boundary
# ─────────────────────────────────────────────────────────────────────────
# The fusion pass lives in `luxtiles_passes_fusion`, reading this module's
# analysis and emission surface, the import sitting below every helper
# the pass module reads, the sibling compiling against the declarations above it.
#
# analyzeStage (per-stage pipeline, no lowering) -> planFusionGroups
# | fused group -> emitDeviceFunc + emitEntry, one launch per group
# v ffNone group's stages -> lowerStage, one launch per stage
#
# A fusable group lowers into one kernel, its intermediates staying
# inside it as rt tiles / rv row vectors. `-d:noFuse` keeps every stage
# its own launch.
export DefaultTile, DefaultThreads, DescSpecs, KernelTrace, SchedInfo,
buildPasses, emitDeviceFunc, emitEntry, analyzeStage, stageThreads,
StmtParts, Emitter, GroupFacts, PhysAxis, StageClass, PlaneGeometry
import workspace/experimental/tensorcery_luxtiles/luxtiles_passes_fusion
var respliceCount {.compileTime.} = 0
## Counter over re-splices, suffixing the re-spliced fn and entry names.
proc retitled(n: NimNode, mapping: seq[tuple[old: string, fresh: NimNode]]): NimNode =
## Copy of `n` with ident and symbol nodes named in `mapping` replaced,
## for re-splicing a derived kernel's fn and entry under fresh names.
result = n.copyNimTree()
proc walk(node: NimNode) =
for i in 0 ..< node.len:
let ch = node[i]
if ch.kind in {nnkIdent, nnkSym}:
for mp in mapping:
if ch.strVal == mp.old:
node[i] = mp.fresh
else:
walk(ch)
walk(result)
proc lowerStageParts(fn: Fn, stageIdx: int, sched: StageSchedule, kernelName: string,
externals: seq[string] = @[], views: seq[string] = @[]):
tuple[fnDef, entry, launch: NimNode, target, entryName, fnName: string] =
## Runs one stage through the pass pipeline, returning the device fn,
## the thin entry and the launch call node, one kernel per stage.
##
## Derivation happens once per kernel name, re-splicing any stored
## nodes on each repeat.
##
## The eval boundary assembles the program's pieces, nothing
## launching at this boundary.
crucibleCompileTarget = ctMetal
var ir = stageToIR(fn, stageIdx, sched, kernelName)
let threads = stageThreads(sched)
ir.threads = (if threads > 0: threads else: DefaultThreads)
let passes = buildPasses(sched)
let entryBase = ir.name & "Entry"
var fnDef: NimNode
var entry: NimNode
var entryName: NimNode
if KernelCache.hasKey(ir.name):
# Same kernel named twice:
# reuse the stored nodes with a fresh suffix, so the two statements
# do not redefine each other's symbols.
inc respliceCount
let suffix = $respliceCount
let mapping = @[(ir.name & "Fn", id(ir.name & "Fn" & suffix)),
(entryBase, id(entryBase & suffix))]
fnDef = retitled(KernelCache[ir.name][0], mapping)
entry = retitled(KernelCache[ir.name][1], mapping)
entryName = id(entryBase & suffix)
else:
let derived = runPipeline(ir, passes)
KernelCache[ir.name] = nnkBracket.newTree(derived[0], derived[1])
fnDef = derived[0]
entry = derived[1]
entryName = id(entryBase)
var selfReads: seq[string] = @[]
for t in ir.tensorParams():
if t == ir.target:
selfReads.add(t)
var call = nnkCall.newTree(entryName)
for t in ir.tensorParams():
if t notin selfReads:
call.add(id((if t in views: kernelName & "_" & t
elif t in externals: t
else: kernelName & "_" & t)))
for d in ir.declArgs:
call.add(id(d.argName))
if selfReads.len > 0:
var selfNodes: seq[NimNode] = @[]
for t in selfReads:
selfNodes.add(newLit(t))
DescSpecs[ir.name].add(nnkExprEqExpr.newTree(ident"selfReads",
nnkBracket.newTree(selfNodes)))
result.launch = nnkCall.newTree(id"evalCall", id"engine", entryName, call,
newLit(ir.name))
result.fnDef = fnDef
result.entry = entry
result.target = ir.target
result.entryName = (if entryName.kind in {nnkIdent, nnkSym}:
entryName.strVal else: "")
# the device fn's ident mirrors the entry's name, suffix included
result.fnName = entryName.strVal[0 ..< entryName.strVal.len - 5] & "Fn"
# ─────────────────────────────────────────────────────────────────────────
# The luxtile recipe header
# ─────────────────────────────────────────────────────────────────────────
# The recipe proc's body is a sequence of stage statements over Iter
# symbols. The macro runs before the body is analyzed:
#
# - it walks the raw statements
# - it builds the Fn IR
# - it lowers stage by stage
#
# var i = domain(n) declares a domain, declaration order the loop order
# A[i, j] = expr stage 0 of Fn A, definition
# A[i, j] += expr stage n of Fn A, append
# a trailing do-block stage schedule, a StageSchedule
# loopNest(a, b, ...) the non-default loop order, a schedule directive
# sched(A) with do-block Fn-level default schedule, a FnSchedule
#
# Statement AST shapes:
#
# A[i,j] += expr with do-block sched
# -> nnkInfix(+=, target, infix, do-block)
# A[i,j] = expr
# -> nnkAsgn(target, expr)
# sched(A) with do-block sched
# -> nnkCall(sched, A, block)
proc parseStmt(st: NimNode): StmtParts =
## One body statement's target, accumulate flag, value and schedule block.
var lhs, rhs: NimNode
case st.kind
of nnkInfix:
doAssert st[0].strVal == "+=",
"luxtiles: unsupported infix statement " & st[0].strVal
result.accumulate = true
lhs = st[1]
rhs = st[2]
if st.len >= 4:
result.schedBlk = st[3]
of nnkAsgn:
lhs = st[0]
rhs = st[1]
if lhs.kind == nnkInfix and lhs[0].strVal == "+=":
result.accumulate = true
lhs = lhs[1]
else:
error("luxtiles: unsupported statement " & $st.kind, st)
doAssert lhs.kind == nnkBracketExpr and lhs.len >= 1,
"the target must be written name[d0, d1, ...] (or name[] for all-reduce)"
result.target = lhs
doAssert not (rhs.kind == nnkInfix and rhs[0].strVal == "|"),
"the `| (...)` iteration-domain tail is superseded: declare the " &
"domains at the recipe top with `var i = domain(n)` and order them " &
"with the loopNest schedule directive"
# A trailing do-block can sit on the value node as its last child,
# a call or the += infix value, stripped here into the schedule block
if result.schedBlk == nil and rhs.kind in {nnkCall, nnkInfix,
nnkBracketExpr} and rhs[^1].kind == nnkStmtList and
(rhs.kind != nnkInfix or rhs.len >= 4):
result.schedBlk = rhs[^1]
var stripped = newNimNode(rhs.kind, rhs)
for i in 0 ..< rhs.len - 1:
stripped.add(rhs[i])
rhs = (if rhs.kind == nnkCall and stripped.len == 1: stripped[0]
else: stripped)
result.value = rhs
# ─────────────────────────────────────────────────────────────────────────
# The eval boundary, the only launch site
# ─────────────────────────────────────────────────────────────────────────
# The recipe records its program, device fn + thin entry derivations
# and launch call nodes in group order.
# engine.eval splices it at the call site, substituting the engine
# and the recipe formals.
# The last launch's value becomes the expansion's tensor.
#
# recipe proc (records) --> ProgramCache --> engine.eval(recipe(args))
# -> program spliced, launches in order, intermediates as locals
#
# A recipe evaluated at several sites re-splices its kernels under
# fresh name suffixes, the launch specs staying keyed by kernel name.
const
ProgramCache = CacheTable"tc.program"
## Per-recipe recorded program, written by the luxtile macro
## and spliced by `engine.eval`.
var recipeFormals {.compileTime.}: Table[string, seq[string]]
## Recipe formal names, the eval substitution's binding keys.
var recipeKernels {.compileTime.}: Table[string, seq[string]]
## Entry names a recipe's program defines, the re-splice targets.
var evalInvocation {.compileTime.} = 0
## Count over eval sites of one compilation, suffixing repeat splices.
macro eval*(e: TensorContext, call: untyped): LuxTensor =
## Launch site of a recorded luxtiles recipe.
##
## - the program splices here, its launches running in order
## - nothing launches at record time
doAssert call.kind == nnkCall and call.len >= 2 and
call[0].kind in {nnkIdent, nnkSym}, "eval: recipe call expected"
let recipe = call[0].strVal
doAssert ProgramCache.hasKey(recipe),
"eval: '" & recipe & "' is not a luxtiles recipe"
let formals = recipeFormals[recipe]
doAssert call.len - 1 == formals.len,
"eval: argument count mismatch for '" & recipe & "'"
inc evalInvocation
var prog = ProgramCache[recipe].copyNimTree()
if evalInvocation > 1:
var mapping: seq[tuple[old: string, fresh: NimNode]] = @[]
for kn in recipeKernels[recipe]:
mapping.add((kn, id(kn & $evalInvocation)))
prog = retitled(prog, mapping)
# the launch specs re-titled with the program, the grid and out-shape
# expressions carrying the declared ranges' arg names, the spec copies
# landing under the suffixed keys the launches then look up
for kn in recipeKernels[recipe]:
if DescSpecs.hasKey(kn):
DescSpecs[kn & $evalInvocation] = retitled(DescSpecs[kn], mapping)
proc renameKernelKey(n: NimNode, suffix: string, keys: seq[string]) =
for i in 0 ..< n.len:
let ch = n[i]
if ch.kind == nnkCall and ch.len >= 5 and
ch[0].kind == nnkIdent and ch[0].strVal == "evalCall" and
ch[4].kind == nnkStrLit and ch[4].strVal in keys:
ch[4] = newLit(ch[4].strVal & suffix)
else:
renameKernelKey(ch, suffix, keys)
renameKernelKey(prog, $evalInvocation, recipeKernels[recipe])
var bindings = initTable[string, NimNode]()
bindings["engine"] = e
for k, f in formals:
bindings[f] = call[k + 1]
# the launch nodes only, device fn and entry defs holding no evalCall
# a typed macro parameter cannot be captured by a closure, so it
# flows in as an argument
proc walk(n: NimNode, engineNode: NimNode, subs: Table[string, NimNode]) =
for i in 0 ..< n.len:
let ch = n[i]
if ch.kind == nnkCall and ch.len >= 4 and ch[0].kind == nnkIdent and
ch[0].strVal == "evalCall":
ch[1] = engineNode
for a in 1 ..< ch[3].len:
let arg = ch[3][a]
if arg.kind in {nnkIdent, nnkSym} and $arg in subs:
ch[3][a] = subs[$arg]
else:
walk(ch, engineNode, subs)
walk(prog, e, bindings)
# the prologue's view lets reference the recipe's formals by name,
# the substitution binding them to the call site's arguments
prog[0] = substituteIdents(prog[0], bindings)
result = prog
macro luxtile*(p: untyped): untyped =
## Recipe header for the luxtiles DSL, statements recording into Fn stages.
##
## - the pass pipeline lowers the recorded program once
## - the program caches for the engine eval, the only launch site
## - nothing launches at recipe-record time
doAssert p.kind == nnkProcDef, "luxtile: recipe proc expected"
let params = p[3]
doAssert p[0].kind in {nnkIdent, nnkPostfix}, "luxtile: recipe name expected"
let procName = (if p[0].kind == nnkIdent: p[0].strVal else: p[0][1].strVal)
# the read surface, every formal an input tensor
var tensors: Table[string, Fn]
var formals: seq[string] = @[]
for i in 1 ..< params.len:
for j in 0 ..< params[i].len - 2:
let nm = params[i][j].strVal
doAssert nm != "engine",
"luxtile: the engine flows from engine.eval, recipes take tensors only"
tensors[nm] = newFn(nm)
formals.add(nm)
var iterTable: Table[string, Iter]
var declOrder: seq[string] = @[]
var recipeDecls: Table[string, Range] = initTable[string, Range]()
var prologue = nnkStmtList.newTree()
# statement collection, the fusion planner's input, the lowering
# (per stage or per fused group) running after the parse
var staged: seq[StageRef] = @[]
var views: seq[string] = @[]
for st in p[^1]:
if st.kind == nnkVarSection:
# Recipe-declared domains, `var i = domain(n)`, declaration
# order the loop order, the reduce set never written per statement
for def in st:
doAssert def.kind == nnkIdentDefs and def.len == 3 and
def[1].kind == nnkEmpty and def[2].kind == nnkCall and
def[2].len == 2 and def[2][0].strVal == "domain",
"declared domains are written `var i = domain(n)`, n an integer " &
"literal or a host value at the eval site"
let sym = def[0].strVal
doAssert sym notin iterTable,
"the domain '" & sym & "' is declared twice"
recipeDecls[sym] = domainRange(sym, def[2][1])
declOrder.add(sym)
iterTable[sym] = newIter(sym)
continue
if st.kind == nnkLetSection:
# Zero-copy views, `let q4 = q.view(...)`, a program-level alias
# of a recipe parameter, the launches passing the alias by name
for def in st:
doAssert def.kind == nnkIdentDefs and def.len == 3 and
def[1].kind == nnkEmpty and def[2].kind == nnkCall and
def[2].len >= 2 and def[2][0].kind == nnkDotExpr and
$def[2][0][1] == "view" and def[2][0][0].strVal in tensors and
tensors[def[2][0][0].strVal].stages.len == 0,
"views are written `let v = t.view(d0, d1, ...)` over a recipe " &
"parameter"
let alias = def[0].strVal
doAssert alias notin tensors,
"the view name '" & alias & "' is already a recipe name"
tensors[alias] = newFn(alias)
tensors[alias].viewOf = def[2][0][0].strVal
views.add(alias)
prologue.add(st)
continue
if st.kind == nnkCall and st.len in 2 .. 3 and
st[0].kind in {nnkIdent, nnkSym} and st[0].strVal == "sched":
# sched(fn) with do-block, the Fn-level default schedule
let fnName = st[1].strVal
if fnName notin tensors:
tensors[fnName] = newFn(fnName)
let blk = (if st.len >= 3: st[2] else: nnkStmtList.newTree())
tensors[fnName].schedule =
FnSchedule(directives: parseSchedule(blk).directives)
continue
let parts = parseStmt(st)
let name = parts.target[0].strVal
doAssert name notin formals,
"luxtiles: '" & name & "' is a recipe parameter, not a target"
var fn: Fn
if name in tensors:
fn = tensors[name]
else:
fn = newFn(name)
tensors[name] = fn
var sched: StageSchedule
if parts.schedBlk != nil:
sched = parseSchedule(parts.schedBlk)
else:
sched = StageSchedule(directives: fn.schedule.directives)
let params = schedParams(sched)
var lhsNodes: seq[LuxNode] = @[]
for i in 1 ..< parts.target.len:
lhsNodes.add(indexExpr(parts.target[i], tensors, iterTable, params))
let def = valueExpr(parts.value, tensors, iterTable, params)
# the statement's iterators, the declared domains first by declaration
# order and the implicit ones after by first-use order, from the LHS
# symbols through the definition to the LHS index forms' own symbols
var used: seq[string] = @[]
for ix in lhsNodes:
if ix.kind == Domain and ix.iter.symbol notin used:
used.add(ix.iter.symbol)
for s in collectIterSymbols(def):
if s notin used:
used.add(s)
for ix in lhsNodes:
for s in collectIterSymbols(ix):
if s notin used:
used.add(s)
var itersOrder: seq[Iter] = @[]
for s in declOrder:
if s in used:
itersOrder.add(iterTable[s])
for s in used:
if s notin declOrder:
itersOrder.add(iterTable[s])
# loopNest, the non-default loop order, the listed names first,
# the remaining iterators keeping their relative order
for d in sched.directives:
if d.kind != ScOrder: continue
var ordered: seq[Iter] = @[]
var listed: seq[string] = @[]
for nm in d.names:
doAssert nm in used,
"loopNest names '" & nm & "', which the statement never uses"
doAssert nm notin listed, "loopNest names '" & nm & "' twice"
listed.add(nm)
for it in itersOrder:
if it.symbol == nm:
ordered.add(it)
break
for it in itersOrder:
if it.symbol notin d.names:
ordered.add(it)
itersOrder = ordered
var guardCond: LuxNode = nil
for d in sched.directives:
if d.kind == ScGuard:
guardCond = condExpr(d.exprNode, iterTable, params, tensors)
fn.appendStage(lhsNodes, def, itersOrder, guardCond, sched)
staged.add(StageRef(name: name, fn: fn, sched: sched,
stageIdx: fn.stages.len - 1))
# the recipe's declared domains and views reach every Fn, the stage
# IR builds reading them
for _, fn in tensors:
fn.decls = recipeDecls
fn.declOrder = declOrder
# program assembly, every group's (device fn, thin entry) derivation
# running here, the launch call nodes recorded for the eval boundary
var program = nnkStmtList.newTree(
nnkStmtList.newTree())
var kernels: seq[string] = @[]
var launches: seq[tuple[launch: NimNode, target: string]] = @[]
for grp in planFusionGroups(staged, procName):
if grp.kind == fkNone:
for si in grp.stages:
let sr = staged[si]
let (fnDef, entry, launch, target, entryName, fnName) =
lowerStageParts(sr.fn, sr.stageIdx, sr.sched, procName, formals,
views)
program.add(fnDef)
program.add(entry)
launches.add((launch, target))
kernels.add(entryName)
kernels.add(fnName)
kernels.add(fnName[0 ..< fnName.len - 2])
else:
let (fnDef, entry, launch, _) =
lowerGroup(staged, grp.stages[0], grp.stages[^1] + 1, procName, formals, views)
program.add(fnDef)
program.add(entry)
launches.add((launch, staged[grp.stages[^1]].name))
kernels.add(procName & "_fusedFn")
kernels.add(procName & "_fusedEntry")
kernels.add(procName & "_fused")
doAssert launches[^1].target == "result",
"the recipe's final statement must target result"
# the program prologue, the views as written, then the declared ranges'
# host lets, one per declared domain carrying a runtime value, resolved
# at the eval site
for st in prologue:
# the view aliases' lets carry the recipe-prefixed name, the launches
# referencing them by that name at the eval site
for i in 0 ..< st.len:
st[i][0] = id(procName & "_" & st[i][0].strVal)
program[0].add(st)
for s in declOrder:
if recipeDecls[s].kind == rkDecl:
program[0].add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id(procName & "_" & s & "R"), newEmptyNode(),
nnkCall.newTree(ident"int32", recipeDecls[s].hostNode))))
for k in 0 ..< launches.len - 1:
# intermediate targets carry the recipe-prefixed name, the readers
# referencing them by that name at the eval site
program.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id((if launches[k].target == "result": launches[k].target
else: procName & "_" & launches[k].target)),
newEmptyNode(), launches[k].launch)))
program.add(launches[^1].launch)
for s in declOrder:
if recipeDecls[s].kind == rkDecl:
kernels.add(procName & "_" & s & "R")
ProgramCache[procName] = program
recipeFormals[procName] = formals
recipeKernels[procName] = kernels
result = p.copyNimTree()
# the recipe proc is a record, the engine flowing from engine.eval
# and the body staying empty of launches
result[3][0] = newEmptyNode()
result[^1] = nnkStmtList.newTree(
nnkCall.newTree(id"doAssert", id"false",
newLit("recipe '" & procName &
"' is a program record; launch it with engine.eval")))
# ─────────────────────────────────────────────────────────────────────────
# Authoring surface, what a recipe proc's scope must see
# ─────────────────────────────────────────────────────────────────────────
# The statement expansions splice into the recipe's scope, resolving plain
# idents there, so the module re-exports the engine shim, the tile
# vocabulary and the layout helpers the expansions reference, mirroring
# the in-module scope the baseline recipes compiled against.
# Recipe target names are proc-prefixed, so no kernel name collides.
export setBackend, tensor, view, allocImpl, toMetadata, makeHostTensor, numel
export operandArg, tdim, cdiv, mkGdView, rowStride, colStride
export TensorContext, HostElem, LuxTensor
export luxtile, evalCall, eval, metal
export BackendKind
export crucible
export small_seqs
export int_tuples, layouts, layout_constructors, layout_indexing
export ptr_arithmetic, tensor_datatypes, layout_algebra
export tile_algebra, tile_io_boundedluxtiles_passes_mma.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.
## luxtiles_passes_mma, the mma optimization pass over the luxtiles stage IR.
##
## Matches a canonical mma shape from the stage IR's schedule-tiled dims
## and the two operand reads' index slots.
##
## The matched form lowers into a staged (threadgroup memory) or register
## kernel body:
##
## - the generic tile program stays the semantics
## - the pass is an optimization, `-d:noMmaPass` compiling it out
##
## AlgoIR + schedule tiles -> matchMma -> MmaMatch
## -> staged (schedule-cached, pure-indexed) | register otherwise
## -> device fn + thin entry (luxtiles_compiler's lowerEmit)
import
std/[macros, tables],
workspace/experimental/tensorcery_luxtiles/luxtiles_types,
workspace/experimental/tensorcery_luxtiles/luxtiles_compiler,
workspace/ceramic/src/tile_algebra,
workspace/ceramic/src/tile_algebra/tile_io_bounded,
workspace/crucible
# Node construction helper emitting plain idents for late name binding, sibling fragments resolving each other's names at splice time.
proc id(s: string): NimNode = newIdentNode(s)
# ─────────────────────────────────────────────────────────────────────────
# mma free dims, the broadcast/batch out dims
# ─────────────────────────────────────────────────────────────────────────
# Example, bmm result[i, j, l] ++= a3[i, j, k] * b3[i, k, l]
#
# i is a free dim. Both reads bind it, the schedule neither tiles it
# nor sums over it, and the kernel runs one i per grid.z slice.
#
# Of the free dims, the first two split grid.z between them
# (i = z div range1, j = z mod range1). From the third on, each free
# dim becomes a normal for-loop whose loop variable is added to every
# read's index and the output's offset.
proc outDimPos(ir: AlgoIR, d: string): int =
## Axis of the output tensor carrying out dim `d`, -1 when absent.
for i, od in ir.outDims:
if od == d: return i
-1
proc freeRange(ir: AlgoIR, d: string): Range =
## Range of a free dim, the first read binding it purely, the derived
## range of its affine binding or the output tensor's own axis.
for rd in ir.reads:
for k, ix in rd.indices:
if ix.kind == ikIter and ix.name == d:
return Range(kind: rkAxis, tensor: rd.tensor, axis: k)
for rd in ir.reads:
for ix in rd.indices:
if ix.kind == ikAffine:
for t in ix.terms:
if t.iter == d:
return ir.iterRange(d)
Range(kind: rkAxis, tensor: ir.target, axis: ir.outDimPos(d))
proc freeCoordPlan(ir: AlgoIR, freeDims: seq[string], body: NimNode):
tuple[resolve: Table[string, NimNode], serial: seq[string]] =
## Free dims' coordinate lets emitted into `body`, the resolve table
## (iterator → coordinate, added to the reads' base offsets) returned.
if freeDims.len == 1:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"bCoord", newEmptyNode(), id"bz")))
result.resolve[freeDims[0]] = id"bCoord"
elif freeDims.len >= 2:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"bRange", newEmptyNode(),
rangeExpr(ir.freeRange(freeDims[1]), device = true))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"bCoord", newEmptyNode(), infixF("div", id"bz", id"bRange"))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"dCoord", newEmptyNode(), infixF("mod", id"bz", id"bRange"))))
result.resolve[freeDims[0]] = id"bCoord"
result.resolve[freeDims[1]] = id"dCoord"
if freeDims.len > 2:
for d in freeDims[2 ..^ 1]:
result.resolve[d] = id(d)
result.serial.add(d)
proc wrapSerialFree(ir: AlgoIR, serial: seq[string], s: NimNode): NimNode =
## Wraps `s` in one serial for-loop per deeper free dim, innermost first.
result = s
for i in countdown(serial.len - 1, 0):
result = nnkForStmt.newTree(id(serial[i]),
nnkInfix.newTree(id"..<", newLit(0'i32),
callI32(rangeExpr(ir.freeRange(serial[i]), device = true))),
nnkStmtList.newTree(result))
proc outFreeBase(ir: AlgoIR, freeDims: seq[string], resolve: Table[string, NimNode]): NimNode =
## Output tensor's base offset for the deeper free dims, `coord * stride`
## per bound dim, zero when the free dims sit on the view's batch/depth slots.
result = i32lit(0)
if freeDims.len <= 2: return
for d in freeDims[2 ..^ 1]:
let term = infixF("*", resolve[d],
callI32(layoutField(ir.target, "stride", ir.outDimPos(d))))
result = (if result.kind == nnkIntLit and result.intVal == 0:
term
else:
infixF("+", result, term))
proc outViewStrides(ir: AlgoIR, freeDims: seq[string]): tuple[b, d: NimNode] =
## Output gd view's batch/depth strides, the first free dims' strides
## when present, dummy `o.rowStride`/0 otherwise.
result.b = (if freeDims.len >= 1:
callI32(layoutField("o", "stride", ir.outDimPos(freeDims[0])))
else:
nnkDotExpr.newTree(id"o", ident"rowStride"))
result.d = (if freeDims.len >= 2:
callI32(layoutField("o", "stride", ir.outDimPos(freeDims[1])))
else:
i32lit(0))
proc outPlaneStride(ir: AlgoIR, d: string): NimNode =
## Output gd view's row or col stride, the out dim's own axis stride.
callI32(layoutField("o", "stride", ir.outDimPos(d)))
# ─────────────────────────────────────────────────────────────────────────
# mmaOptPass matcher, schedule-tiled dims and read index slots to MmaMatch
# ─────────────────────────────────────────────────────────────────────────
# `matchMma` reads the IR's schedule-tiled dims and the two operand reads'
# index slots, recognizing the canonical mma shape. A zero match
# (`ra == -1`) keeps the generic tile program in place.
#
# Match contract, mm as the example, `result[i, j] ++= m1[i, k] * m2[k, j]`
# tiled 16x16 over i, j, k chunked at 16
#
# - the accumulator's rows and cols are the two schedule-tiled out dims,
# tiles 16 or 32, tiles outside [16, 32] matching nothing
# - `ra` binds the row dim exactly once, `rb` the col dim exactly once, two
# different reads, a read carrying both matching nothing
# - the remaining out dims are free, the first two splitting grid.z, deeper
# ones running as serial for-loops
#
# Exactly one reduce iterator is schedule-tiled (the chunk iter), the reduce
# iters binding at most once per read, at least once across the operands.
#
# An index form like `2*oh + kh - pad` may bind one row or col dim plus
# serial reduce iters, free dims and constants, never two of the row/col dims.
proc bindCount(rd: ReadIR, nm: string): int =
## How many of a read's slot index forms bind iterator `nm`, affine forms
## counted per term.
result = 0
for ix in rd.indices:
if ix.kind == ikIter and ix.name == nm:
inc result
elif ix.kind == ikAffine:
for t in ix.terms:
if t.iter == nm:
inc result
proc boundSlot(rd: ReadIR, nm: string): int =
## Iterator `nm`'s one binding slot, -1 when unbound or multiply bound.
var slot = -1
for k, ix in rd.indices:
if ix.kind == ikIter and ix.name == nm:
if slot >= 0: return -1
slot = k
elif ix.kind == ikAffine:
for t in ix.terms:
if t.iter == nm:
if slot >= 0: return -1
slot = k
slot
proc matchMma(ir: AlgoIR): MmaMatch =
## Mma shape recognized from the reads' index slots, one per axis position
## within a read (see `IndexRef`), zero keeping the generic tile program.
result.ra = -1
result.rb = -1
if ir.scatter or ir.expr == nil or ir.expr.kind != xeMul: return
if ir.expr.lhs.kind != xeRead or ir.expr.rhs.kind != xeRead: return
if ir.reads.len != 2: return
# A scalar-term read (a runtime value in an index form) has no slot
# geometry for the mma operands, the generic tile program lowering it
for rd in ir.reads:
for ix in rd.indices:
if ix.scalars.len > 0: return
var plane: seq[string]
for d in ir.outDims:
if ir.outerTile(d) in [16, 32]:
plane.add(d)
if plane.len != 2: return
let rowDim = plane[0]
let colDim = plane[1]
var ra = -1
var rb = -1
for r in 0 ..< 2:
if bindCount(ir.reads[r], rowDim) == 1:
ra = r
if bindCount(ir.reads[r], colDim) == 1:
rb = r
if ra < 0 or rb < 0 or ra == rb: return
var chunkIter = ""
var serialIters: seq[string] = @[]
for it in ir.reduceIters:
if ir.outerTile(it) in [16, 32]:
if chunkIter.len > 0: return
chunkIter = it
else:
serialIters.add(it)
if chunkIter.len == 0: return
for it in ir.reduceIters:
if bindCount(ir.reads[ra], it) > 1 or bindCount(ir.reads[rb], it) > 1:
return
if bindCount(ir.reads[ra], it) == 0 and bindCount(ir.reads[rb], it) == 0:
return
# a slot binding two plane iterators has no single view axis
for r in 0 ..< 2:
let rs = ir.reads[r].boundSlot(rowDim)
let cs = ir.reads[r].boundSlot(colDim)
let ks = ir.reads[r].boundSlot(chunkIter)
if rs >= 0 and cs >= 0 and rs == cs: return
if rs >= 0 and ks >= 0 and rs == ks: return
if cs >= 0 and ks >= 0 and cs == ks: return
var freeDims: seq[string] = @[]
for d in ir.outDims:
if d != rowDim and d != colDim:
freeDims.add(d)
let (tm, tn, tk) = (ir.outerTile(rowDim), ir.outerTile(colDim),
ir.outerTile(chunkIter))
if tm notin [16, 32] or tn notin [16, 32] or tk notin [16, 32]: return
result.ra = ra
result.rb = rb
result.rowDim = rowDim
result.colDim = colDim
result.freeDims = freeDims
result.chunkIter = chunkIter
result.serialIters = serialIters
result.tiles = (tm, tn, tk)
# ─────────────────────────────────────────────────────────────────────────
# MmaMatch → kernel body, staged or register. Staged copies operand tiles
# through threadgroup memory per reduce chunk and accumulates with the gemm
# slab expansion, register loads bounded tiles per chunk and calls mma_AB.
#
# MmaMatch, staged or register form
# | staged, emitCopyFromGlobal + emitGemm + acc
#
# v register, bounded loads + mma_AB per chunk
# | emitDeviceFunc + emitEntry, both forms
#
# Register form example, `result[i, j] ++= a[i, k] * b[k, j]` with tiles
# 16x16, k in chunks of 16
#
# Each thread's tiles of `a` and `b` multiply into its result tile, edges
# reading past the tensor contributing zero.
#
# Reads indexed through a linear form (`oh*s + kh - p` in conv2d) load absolute,
# bounds-checked against the read tensor's own range.
# A linear form carrying a free dim gets one whole-body validity guard.
proc emitCopyFromGlobal(outBody: var NimNode, dst: string, dims: tuple[r, c: int], ten: string, base, rStride, cStride, rE, cE, rowLim, colLim: NimNode) =
## Strided flat staging loop, the threadgroup filling `dst` row-major, one
## (r, c) tile of the read's global plane, out-of-range lanes zero-fill.
let dstN = id(dst)
let tenN = id(ten)
let rL = newLit(dims.r)
let cL = newLit(dims.c)
let i = genSym(nskVar, "stageI")
let row = infixF("+", rE,
nnkInfix.newTree(id"div", i, callI32(newLit(dims.c))))
let col = infixF("+", cE,
nnkInfix.newTree(id"mod", i, callI32(newLit(dims.c))))
var off = infixF("+", infixF("*", row, rStride), infixF("*", col, cStride))
if not (base.kind == nnkIntLit and base.intVal == 0):
off = infixF("+", base, off)
outBody.add quote do:
var `i` = int32(thread_index_in_threadgroup)
while `i` < int32(`rL` * `cL`):
if `row` < `rowLim` and `col` < `colLim`:
`dstN`[`i`] = `tenN`.data[`off`]
else:
`dstN`[`i`] = zeroTileValue(HostElem)
`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.
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, frag: string,
dims: tuple[r, c: int], bStr, dStr, rStr, cStr, base, bCoord, dCoord,
rE, cE, validM, validN: NimNode) =
## Masked tile store of the accumulator into the global output, ragged
## tiles writing only the in-range cells.
let gd = genSym(nskLet, "gdStoreView")
let rT = genSym(nskLet, "tileR")
let cT = genSym(nskLet, "tileC")
let frL = newLit(dims.r)
let fcL = newLit(dims.c)
let fragN = id(frag)
let buf = (if base.kind == nnkIntLit and base.intVal == 0:
nnkDotExpr.newTree(id"o", ident"data")
else:
nnkCall.newTree(id"offsetBuf",
nnkDotExpr.newTree(id"o", ident"data"), base))
outBody.add quote do:
let `gd` = mkGdView(`buf`, `bStr`, `dStr`, `rStr`, `cStr`)
let `rT` = `rE` div int32(`frL`)
let `cT` = `cE` div int32(`fcL`)
storeTileMasked(`gd`, `fragN`, (`bCoord`, `dCoord`, `rT`, `cT`),
`validM`, `validN`)
proc emitMmaStaged*(ir: AlgoIR, mt: MmaMatch):
tuple[body, smemDecls: NimNode, scratch: seq[string]] =
## Staged mma kernel body for the matched shape `mt`, one threadgroup tile
## per cached operand staged from global memory per reduce chunk.
let (bm, bn, bk) = mt.tiles
let rdA = ir.reads[mt.ra]
let rdB = ir.reads[mt.rb]
let tenA = rdA.tensor
let tenB = rdB.tensor
let bufA = "sm_" & tenA
let bufB = "sm_" & tenB
let frag = "acc"
let reduceVar = "it"
let r0 = ir.reduceIters[0]
let yOrig = infixF("*", id"by", i32lit(bm))
let xOrig = infixF("*", id"bx", i32lit(bn))
let kOrig = infixF("*", id(reduceVar), i32lit(bk))
var smem = nnkStmtList.newTree(smemDecl(bufA, bm, bk),
smemDecl(bufB, bk, bn))
var body = nnkStmtList.newTree()
let (resolve, serialFree) = ir.freeCoordPlan(mt.freeDims, body)
for spec in [("rowRange", rangeExpr(ir.iterRange(mt.rowDim), device = true)),
("colRange", rangeExpr(ir.iterRange(mt.colDim), device = true)),
("kRange", rangeExpr(ir.iterRange(r0), device = true))]:
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id(spec[0]), newEmptyNode(), spec[1])))
body.add(nnkVarSection.newTree(nnkIdentDefs.newTree(
id(frag), nnkCall.newTree(id"rt_l", id"T", i32lit(bm), i32lit(bn)),
newEmptyNode())))
body.add(nnkCall.newTree(dot(id(frag), "zero")))
var chunkBody = nnkStmtList.newTree()
emitCopyFromGlobal(chunkBody, bufA, (bm, bk), tenA,
rdA.readBase(resolve), rdA.slotStride(mt.rowDim), rdA.slotStride(r0),
yOrig, kOrig, id"rowRange", id"kRange")
emitCopyFromGlobal(chunkBody, bufB, (bk, bn), tenB,
rdB.readBase(resolve), rdB.slotStride(r0), rdB.slotStride(mt.colDim),
kOrig, xOrig, id"kRange", id"colRange")
chunkBody.add(nnkCall.newTree(id"syncthreads"))
emitGemm(chunkBody, bufA, bufB, frag, (bm, bk), (bk, bn), (bm, bn))
chunkBody.add(nnkCall.newTree(id"syncthreads"))
body.add(nnkForStmt.newTree(id(reduceVar),
nnkInfix.newTree(id"..<", newLit(0'i32),
callI32(infixF("div",
infixF("-",
infixF("+",
rangeExpr(ir.iterRange(r0), device = true), i32lit(bk)),
newLit(1)),
i32lit(bk)))),
chunkBody))
let (bStrO, dStrO) = ir.outViewStrides(mt.freeDims)
var storeStmts = nnkStmtList.newTree()
emitCopyToGlobal(storeStmts, frag, (bm, bn), bStrO, dStrO,
ir.outPlaneStride(mt.rowDim), ir.outPlaneStride(mt.colDim),
ir.outFreeBase(mt.freeDims, resolve),
(if mt.freeDims.len >= 1: id"bCoord" else: i32lit(0)),
(if mt.freeDims.len >= 2: id"dCoord" else: i32lit(0)),
yOrig, xOrig,
infixF("-", id"rowRange", yOrig), infixF("-", id"colRange", xOrig))
body.add(storeStmts)
result.body = ir.wrapSerialFree(serialFree, body)
result.smemDecls = smem
result.scratch = @[bufA, bufB]
proc slotGuard(rd: ReadIR, slot: int, resolve: Table[string, NimNode]): NimNode =
## Whole-tile validity predicate of an affine slot at the base offset,
## `0 <= coef*it + ... + cst < dim` with no tile-plane iterator.
let ix = rd.indices[slot]
doAssert ix.kind == ikAffine
var idx: NimNode = nil
for t in ix.terms:
doAssert t.iter in resolve,
"index form term '" & t.iter & "' of '" & rd.tensor & "' is neither " &
"on the tile plane nor a resolved serial or free iterator"
var term = resolve[t.iter]
if t.coef != 1:
term = infixF("*", i32lit(t.coef.int), term)
idx = (if idx == nil: term else: infixF("+", idx, term))
if ix.cst != 0 or idx == nil:
idx = infixF("+", (if idx == nil: i32lit(0) else: idx),
i32lit(ix.cst.int))
let dimN = callI32(layoutField(rd.tensor, "shape", slot))
infixF("and", infixF(">=", idx, i32lit(0)), infixF("<", idx, dimN))
proc emitMmaRegister*(ir: AlgoIR, mt: MmaMatch): NimNode =
## Register-only mma kernel body for the matched shape `mt`, every thread
## multiplying from its own registers, no shared staging memory.
let (bm, bn, bk) = mt.tiles
let rdA = ir.reads[mt.ra]
let rdB = ir.reads[mt.rb]
let tenA = rdA.tensor
let tenB = rdB.tensor
let gdA = "gd_" & tenA
let gdB = "gd_" & tenB
let ta = "ta_" & tenA
let tb = "tb_" & tenB
let acc = "acc"
let reduceVar = "it"
let chunk = mt.chunkIter
let yOrig = infixF("*", id"by", i32lit(bm))
let xOrig = infixF("*", id"bx", i32lit(bn))
var stmts = nnkStmtList.newTree()
var (resolve, serialFree) = ir.freeCoordPlan(mt.freeDims, stmts)
for it in mt.serialIters:
resolve[it] = id(it)
for spec in [("rowRange", rangeExpr(ir.iterRange(mt.rowDim), device = true)),
("colRange", rangeExpr(ir.iterRange(mt.colDim), device = true)),
("kRange", rangeExpr(ir.iterRange(chunk), device = true))]:
stmts.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id(spec[0]), newEmptyNode(), spec[1])))
# One read's plane-axis geometry, a pure slot keeping the tile-unit origin
# and the plane iterator range, an affine slot keeping the v0-shifted
# plane iterator with absolute limits, (coef*v0 + T) added to the base.
proc planeAxis(rd: ReadIR, planeIter: string, coord: NimNode, tileExt: int, limitIdent: string): PlaneGeometry =
let slot = rd.boundSlot(planeIter)
doAssert slot >= 0,
"tile plane iterator '" & planeIter & "' has no single slot in '" &
rd.tensor & "'"
result.affine = rd.indices[slot].kind == ikAffine
result.stride = rd.slotStride(planeIter)
if result.affine:
let geometry = affinePlaneGeometry(rd, slot, planeIter, coord, tileExt,
resolve)
result.stride = geometry.stride
result.origin = geometry.origin
result.limit = geometry.limit
result.baseExtra = planeBaseTerm(rd, slot, planeIter, geometry.v0,
resolve)
else:
result.origin = coord
result.limit = id(limitIdent)
var axRow = planeAxis(rdA, mt.rowDim, id"by", bm, "rowRange")
var axChunkA = planeAxis(rdA, chunk, id(reduceVar), bk, "kRange")
var axColB = planeAxis(rdB, mt.colDim, id"bx", bn, "colRange")
var axChunkB = planeAxis(rdB, chunk, id(reduceVar), bk, "kRange")
# a read loading absolute converts its pure plane axes' tile-unit origins
if axRow.affine or axChunkA.affine:
if not axRow.affine:
axRow.origin = infixF("*", id"by", i32lit(bm))
if not axChunkA.affine:
axChunkA.origin = infixF("*", id(reduceVar), i32lit(bk))
if axColB.affine or axChunkB.affine:
if not axColB.affine:
axColB.origin = infixF("*", id"bx", i32lit(bn))
if not axChunkB.affine:
axChunkB.origin = infixF("*", id(reduceVar), i32lit(bk))
# whole-tile guards, affine slots added to the base offsets
var guards: seq[NimNode] = @[]
block:
let rowSlot = rdA.boundSlot(mt.rowDim)
let chunkSlotA = rdA.boundSlot(chunk)
for slot, ix in rdA.indices:
if ix.kind == ikAffine and slot != rowSlot and slot != chunkSlotA:
guards.add(slotGuard(rdA, slot, resolve))
block:
let colSlot = rdB.boundSlot(mt.colDim)
let chunkSlotB = rdB.boundSlot(chunk)
for slot, ix in rdB.indices:
if ix.kind == ikAffine and slot != colSlot and slot != chunkSlotB:
guards.add(slotGuard(rdB, slot, resolve))
proc loadNameProc(useAbs: bool): NimNode =
(if useAbs: ident"loadTileBoundedAbs" else: ident"loadTileBounded")
let loadAName = loadNameProc(axRow.affine or axChunkA.affine)
let loadBName = loadNameProc(axColB.affine or axChunkB.affine)
var accDecl = 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()))
let zeroCall = nnkCall.newTree(dot(id(acc), "zero"))
proc baseWith(rd: ReadIR, planeIters: seq[string], axes: seq[PlaneGeometry]): NimNode =
## A read's base offset.
## `readBase` runs with the plane slots skipped and the affine plane
## slots' (coef*v0 + T) parts added back.
var slots: seq[int] = @[]
for pi in planeIters:
let sl = rd.boundSlot(pi)
if sl >= 0: slots.add(sl)
var base = rd.readBase(resolve, slots)
for ax in axes:
if ax.baseExtra != nil and
not (ax.baseExtra.kind == nnkIntLit and ax.baseExtra.intVal == 0):
base = (if base.kind == nnkIntLit and base.intVal == 0: ax.baseExtra
else: infixF("+", base, ax.baseExtra))
base
var gdAlet = gdViewNode(gdA, tenA,
nnkDotExpr.newTree(id(tenA), ident"rowStride"), i32lit(0),
axRow.stride, axChunkA.stride,
baseWith(rdA, @[mt.rowDim, chunk], @[axRow, axChunkA]))
var gdBlet = gdViewNode(gdB, tenB,
nnkDotExpr.newTree(id(tenB), ident"rowStride"), i32lit(0),
axColB.stride, axChunkB.stride,
baseWith(rdB, @[mt.colDim, chunk], @[axColB, axChunkB]))
var chunkBody = nnkStmtList.newTree(
nnkCall.newTree(nnkDotExpr.newTree(id(ta), loadAName), id(gdA),
nnkPar.newTree(newLit(0), newLit(0), axRow.origin, axChunkA.origin),
axRow.limit, axChunkA.limit),
nnkCall.newTree(nnkDotExpr.newTree(id(tb), loadBName), id(gdB),
nnkPar.newTree(newLit(0), newLit(0), axColB.origin, axChunkB.origin),
axColB.limit, axChunkB.limit),
nnkCall.newTree(dot(id(acc), "mma_AB"), id(ta), id(tb)))
var chunkLoop = nnkForStmt.newTree(id(reduceVar),
nnkInfix.newTree(id"..<", newLit(0'i32),
callI32(infixF("div",
infixF("-",
infixF("+",
rangeExpr(ir.iterRange(chunk), device = true), i32lit(bk)),
newLit(1)),
i32lit(bk)))),
chunkBody)
let (bStrO, dStrO) = ir.outViewStrides(mt.freeDims)
var storelets = nnkStmtList.newTree(
gdViewNode("gdO", "o", bStrO, dStrO,
ir.outPlaneStride(mt.rowDim), ir.outPlaneStride(mt.colDim),
ir.outFreeBase(mt.freeDims, resolve)),
nnkCall.newTree(dot(id"gdO", "storeTileMasked"), id(acc),
nnkPar.newTree(
(if mt.freeDims.len >= 1: id"bCoord" else: i32lit(0)),
(if mt.freeDims.len >= 2: id"dCoord" else: i32lit(0)),
id"by", id"bx"),
infixF("-", id"rowRange", yOrig),
infixF("-", id"colRange", xOrig)))
if mt.serialIters.len == 0 and guards.len == 0:
# no serial reduce iters and no guards, the statement order matches
# the plain (non-affine) emission exactly
stmts.add(gdAlet)
stmts.add(gdBlet)
stmts.add(accDecl)
stmts.add(zeroCall)
stmts.add(chunkLoop)
stmts.add(storelets)
else:
stmts.add(accDecl)
stmts.add(zeroCall)
var inner = nnkStmtList.newTree(gdAlet, gdBlet, chunkLoop)
for g in guards:
inner = nnkIfStmt.newTree(nnkElifBranch.newTree(g, inner))
stmts.add(ir.wrapSerialFree(mt.serialIters, inner))
stmts.add(storelets)
result = ir.wrapSerialFree(serialFree, stmts)
# ─────────────────────────────────────────────────────────────────────────
# Pass wiring, eligibility, form description and the optimization pass
# ─────────────────────────────────────────────────────────────────────────
proc isStagedEligible*(ir: AlgoIR, mt: MmaMatch): bool =
## Staged form for a matched mma shape with both schedule-cached operands
## and one pure-iterator reduce, anything wider loading registers only.
if mt.ra < 0: return false
if mt.serialIters.len > 0: return false
for rd in ir.reads:
if rd.tensor notin ir.cached: return false
for ix in rd.indices:
if ix.kind == ikAffine: return false
true
proc mmaFormName*(ir: AlgoIR, mt: MmaMatch, staged: bool): string =
## Lowering description of the matched mma tile program.
var s = (if staged: "mma kernel, smem-staged operands + gemm accumulation"
else: "mma kernel, register tiles (loadTileBounded + mma_AB)")
if mt.freeDims.len > 0:
s &= ", free/batch dims on grid.z"
s
proc mmaOptPass*(): Pass =
## Optimization pass recognizing the canonical mma shape into `ir.mmaMatch`,
## an uncovered program keeping the generic tile program, correct and slower.
Pass(name: "mmaOpt",
description: "canonical mma shape recognition, annotate or fall through",
kind: pkTransform,
run: proc(ir: var AlgoIR) =
ir.mmaMatch = ir.matchMma()
ir.hasMma = ir.mmaMatch.ra >= 0)luxtiles_passes_fusion.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/LICENSE-2.0).
## at your option. This file may not be copied, modified, or distributed except according to those terms.
## luxtiles_passes_fusion, the fusion pass over the luxtiles stage IR.
##
## Stages over one recurrence domain group into ONE kernel,
## the body being a chunk loop per visibility pass, the stages
## executing in stage order.
##
## Intermediates live in rt tiles / rv row vectors.
##
## recipe statements --> planFusionGroups --> FusionGroup list
## --> fkChain --> buildGroupFacts -> lowerGroup
## one device fn + thin entry, per pass one chunk loop
## over the chunk domain, carries in rv/tile, self and sibling
## references resolving to registers, one kernel launch per group
## --> fkNone --> one kernel per stage (lowerStageParts) per group
##
## Visibility over the chunk domain (Lux recurrence semantics):
## the read-resolution contract inside a group.
## | the stage reads | resolves to |
## | ------------------------------------------------- | ------------------------------------ |
## | its own target (self-reference) | the prior chunk's value |
## | a sibling stage's output, same pass | that stage's same-chunk value |
## | a target declared LATER in stage order | its prior value, a Gauss-Seidel |
## | | snapshot, `mOld[i] = mNew[i]` |
## | a reduce over the chunk domain itself (softmax m) | the running accumulator, |
## | | same-pass consumers read it per |
## | | chunk only in a carry group, else |
## | | the value completes at the sweep end |
## | a reduce over the in-chunk slice (mNew over jc) | a fresh per-chunk value |
## | a self-carrying stage without an explicit seed | zero (lNew, oNew), the recipe |
## | | states the rescale explicitly |
##
## Eligibility facts, checked on the analyzed stages (analyzeStage),
## any failure falling back to the per-stage lowering:
##
## - the stages chain by reads, one shared row dim
## - one shared chunk-domain axis, a carry stage and per-stage shape support
## - no statement outside the group reads an intermediate (noEscapes)
##
## Any failure falls back to the per-stage lowering.
##
## Groups containing mma stages need the mma pass, `-d:noMmaPass`
## refusing them and `-d:noFuse` compiling the pass out.
##
## A group stage presenting a canonical two-operand contraction
## lowers through mma_AB inside the chunk body.
##
## The in-group matcher reads the group's own tile plan, group
## tiles including sub-16 chunk widths.
##
## The standalone matchMma of luxtiles_passes_mma stays untouched
## for singleton stages.
##
## Ragged tails:
##
## - loads zero-fill through loadTileBounded, tile-unit origins
## with absolute limits
## - per-chunk tile values masked after computation, 0 for sum
## consumers and -inf for max consumers
## - stores run through storeTileMasked
import
std/[macros, macrocache, strutils, tables],
workspace/experimental/tensorcery_luxtiles/luxtiles_types,
workspace/experimental/tensorcery_luxtiles/luxtiles_compiler,
workspace/crucible/src/codegen/gpu_compiler
# Node construction helper emitting plain idents for late name binding,
# sibling fragments resolving each other's names at splice time.
proc id(s: string): NimNode = newIdentNode(s)
# ─────────────────────────────────────────────────────────────────────────
# Node builders, the fragments' shared shape
# ─────────────────────────────────────────────────────────────────────────
proc declTile(name: string, r, c: int): NimNode =
## Declares `var name: rt_l(T, r, c)`, the mma A-fragment shape.
nnkVarSection.newTree(nnkIdentDefs.newTree(id(name),
nnkCall.newTree(id"rt_l", id"T", newLit(r), newLit(c)), newEmptyNode()))
proc declRv(name: string, r, c: int): NimNode =
## `var name: rv(T, r, c)` declaration, c the source tile's column range.
nnkVarSection.newTree(nnkIdentDefs.newTree(id(name),
nnkCall.newTree(id"rv", id"T", newLit(r), newLit(c)), newEmptyNode()))
proc declRtR(name: string, k, n: int): NimNode =
## `var name: rt_r(T, k, n)` declaration, the mma B-fragment shape.
nnkVarSection.newTree(nnkIdentDefs.newTree(id(name),
nnkCall.newTree(id"rt_r", id"T", newLit(k), newLit(n)), newEmptyNode()))
proc callOp(recv, op: string, args: varargs[NimNode]): NimNode =
var call = nnkCall.newTree(nnkDotExpr.newTree(id(recv), id(op)))
for a in args: call.add(a)
call
proc negInfNode(): NimNode =
## Most-negative finite fp32, the max reduction's neutral value.
nnkPrefix.newTree(id"-", newLit(3.402823466e38'f32))
proc readTensor(e: ExprIR): string =
## Tensor of a read expression, "" otherwise.
(if e.kind == xeRead: e.read.tensor else: "")
proc boundSlotIn(rd: ReadIR, nm: string): int =
## One slot binding `nm`, -1 when unbound or multiply bound.
var slot = -1
for k, ix in rd.indices:
if ix.kind == ikIter and ix.name == nm:
if slot >= 0: return -1
slot = k
elif ix.kind == ikAffine:
for t in ix.terms:
if t.iter == nm:
if slot >= 0: return -1
slot = k
slot
proc pureIters(rd: ReadIR, names: seq[string]): bool =
## Read slots all plain iterators, bound to `names` in order.
if rd.indices.len != names.len: return false
for k, ix in rd.indices:
if ix.kind != ikIter or ix.name != names[k]: return false
true
proc symbolsWalk(n: LuxNode, acc: var seq[string]) =
## Symbols of one stage definition node, appended in appearance order.
case n.kind
of Access:
let s = n.children[0].fn.symbol
if s notin acc: acc.add(s)
for c in 1 ..< n.children.len:
symbolsWalk(n.children[c], acc)
of BinOp:
for c in n.children:
symbolsWalk(c, acc)
of TileOp:
for c in n.operands:
symbolsWalk(c, acc)
else:
discard
proc readSymbols(fn: Fn, stageIdx: int): seq[string] =
## Tensor symbols a stage's definition reads, in appearance order.
symbolsWalk(fn.stages[stageIdx].definition, result)
proc groupThreads(stmts: seq[StageRef], fromI, untilI: int): int =
## First thread count bound by the group's stage schedules, 0 when unset.
for k in fromI ..< untilI:
let t = stageThreads(stmts[k].sched)
if t > 0: return t
0
proc noEscapes(stmts: seq[StageRef], fromI, untilI: int): bool =
## No statement outside the index range reads one of the group's
## intermediate targets, the final target's reads allowed.
var inner: seq[string] = @[]
for k in fromI ..< untilI - 1:
inner.add(stmts[k].name)
if inner.len == 0: return true
for s in 0 ..< stmts.len:
if s >= fromI and s < untilI: continue
for sym in readSymbols(stmts[s].fn, stmts[s].stageIdx):
if sym in inner: return false
true
proc bindsIter(rd: ReadIR, nm: string): bool =
## Whether any read slot's index form binds `nm`.
for ix in rd.indices:
if ix.kind == ikIter and ix.name == nm: return true
if ix.kind == ikAffine:
for t in ix.terms:
if t.iter == nm: return true
false
proc isScaledRead(e: ExprIR, base: var string, scale: var float32): bool =
## `lit * read` or `read * lit` over one two-slot pure read.
if e.kind != xeMul: return false
var rd: ExprIR
if e.lhs.kind == xeLit and e.rhs.kind == xeRead:
scale = e.lhs.value; rd = e.rhs
elif e.lhs.kind == xeRead and e.rhs.kind == xeLit:
scale = e.rhs.value; rd = e.lhs
else:
return false
if rd.read.indices.len != 2: return false
for ix in rd.read.indices:
if ix.kind != ikIter: return false
base = rd.read.tensor
true
# ─────────────────────────────────────────────────────────────────────────
# Group facts, per-stage classes and the chunk-domain plan
# ─────────────────────────────────────────────────────────────────────────
# One physical axis (tensor, read slot) identifies a logical
# dimension across the recipe's statements, a stage's internal
# reads inheriting the producer's out-dim axis, external reads
# binding their own tensor axis.
#
# Chunk domain D, the physical axis most stages reduce over.
# A stage's D-binding iterators split into the chunk-loop
# iterator (the untiled one) and the in-chunk iterator
# (the tiled one).
#
# Stage classes, the emitter's per-stage lowering rules.
#
# scFold scalar-scaled external read, folded into the mma
# operand load consuming it
# scElemTile elementwise (rowDim, col), rt tile per chunk
# scElemRow elementwise (rowDim), rv per chunk
# scPriorCopy copy of a LATER carry stage's value, the snapshot
# taken before that stage updates
# scElemRowCarry self-referencing (rowDim) elementwise, rv carry
# scSweepReduce reduce over the chunk-loop iterator, accumulator
# across chunks, rv at the sweep end
# scChunkReduce reduce over the in-chunk slice, rv per chunk,
# combine with the seed rv when seeded
# scChunkReduceFused scChunkReduce fused into its sole consumer's
# self-carry combine (row-reduce form)
# scMmaFresh two-operand contraction, fresh rt tile per chunk
# scMmaCarry self-carrying contraction accumulate, rescale + mma
proc targetIndex(gf: GroupFacts, nm: string): int =
## Stage index of target `nm`, -1 when not a group target.
for i, n in gf.names:
if n == nm: return i
-1
proc isExternal(gf: GroupFacts, nm: string): bool =
## Whether `nm` lies outside the group, being a recipe input
## or an earlier program target.
gf.targetIndex(nm) < 0
proc stageTileR(gf: GroupFacts, i: int): int =
## Effective row tile of stage `i`, the schedule tile or DefaultTile.
let t = gf.irs[i].outerTile(gf.rowDim)
(if t == 0: DefaultTile else: t)
proc stageDIters(gf: GroupFacts, i: int): seq[string] =
## Stage `i`'s iterators resolving to the chunk domain.
for od in gf.irs[i].outDims:
if od in gf.dIters: result.add(od)
for ri in gf.irs[i].reduceIters:
if ri in gf.dIters and ri notin result: result.add(ri)
proc isCarry(gf: GroupFacts, i: int): bool =
## Whether stage `i` carries its value across chunks, a prior
## value existing for a forward read.
##
## - a plain chunk reduce is a progressive accumulator, its row
## vector persisting across chunks, zero-seeded without a seed read
## - a fused one exists only inside its carry's update, carrying
## no prior value
gf.classes[i] == scElemRowCarry or gf.classes[i] == scMmaCarry or
gf.classes[i] == scSweepReduce or gf.classes[i] == scChunkReduce
proc maxConsumer(gf: GroupFacts, i: int): bool =
## Whether any reader of stage `i` is a max reduction, its padding
## masking to -inf.
for r in gf.readsOf[i]:
if gf.classes[r] in {scChunkReduce, scSweepReduce} and
gf.irs[r].reduceKind == rkMax:
return true
false
proc resolveAxes(stmts: seq[StageRef], irs: seq[AlgoIR]): Table[string, PhysAxis] =
## Logical iterator -> physical axis mapping. External reads bind
## their own tensor axis, internal reads inherit the producer's axis.
var external: seq[string] = @[]
for si in 0 ..< stmts.len:
for rd in irs[si].reads:
if rd.tensor notin external: external.add(rd.tensor)
var targets: seq[string] = @[]
for si in 0 ..< stmts.len:
if stmts[si].name notin targets: targets.add(stmts[si].name)
for t in targets:
let k = external.find(t)
if k >= 0: external.delete(k)
var pending: seq[tuple[iter: string, producer: string, slot: int]] = @[]
for si in 0 ..< stmts.len:
let ir = irs[si]
for rd in ir.reads:
let isExternal = rd.tensor in external
for slot, ix in rd.indices:
case ix.kind
of ikIter:
if isExternal:
if ix.name notin result: result[ix.name] = (rd.tensor, slot)
else:
pending.add((ix.name, rd.tensor, slot))
of ikAffine:
for t in ix.terms:
if isExternal:
if t.iter notin result: result[t.iter] = (rd.tensor, slot)
else:
pending.add((t.iter, rd.tensor, slot))
else: discard
for pd in pending:
if pd.iter in result: continue
for si in 0 ..< stmts.len:
if stmts[si].name != pd.producer: continue
if pd.slot >= irs[si].outDims.len: break
let od = irs[si].outDims[pd.slot]
if od in result: result[pd.iter] = result[od]
break
proc classifyStage(gf: GroupFacts, i: int, msg: var string): bool =
## Stage `i`'s class from its analyzed shape, false with `msg` when the facts do
## not support it. Folds, prior copies and fusions resolve in `resolveDeferred`,
## the counterpart stages' classes not being known at this point.
let ir = gf.irs[i]
let me = gf.names[i]
var selfRead = false
for rd in ir.reads:
if rd.tensor == me: selfRead = true
if ir.reduceIters.len == 0:
# Elementwise stages, no reducer over any iterator
if selfRead:
if ir.outDims.len == 1:
gf.classes[i] = scElemRowCarry
return true
msg = "self-carrying stage '" & me & "' must be one-dimensional"
return false
case ir.outDims.len
of 1:
for rd in ir.reads:
if gf.isExternal(rd.tensor):
msg = "one-dimensional stage '" & me &
"' reading external tensors is unsupported in a group"
return false
gf.classes[i] = scElemRow
return true
of 2:
let col = ir.outDims[1]
let finalCol = (if gf.irs[^1].outDims.len == 2: gf.irs[^1].outDims[1]
else: "")
if col in gf.dIters:
gf.classes[i] = scElemTile
gf.tileC[i] = gf.cs
return true
if i == gf.names.len - 1 and col == finalCol:
for rd in ir.reads:
if gf.isExternal(rd.tensor):
msg = "the final stage '" & me &
"' reading external tensors needs a chunk-domain column"
return false
gf.classes[i] = scElemTile
gf.tileC[i] = (if ir.outerTile(col) == 0: DefaultTile
else: ir.outerTile(col))
return true
# A non-chunk, non-final column defers to the fold resolution,
# a scalar fold feeding an mma A operand is the only supported form
gf.classes[i] = scElemTile
return true
else:
msg = "stage '" & me & "' has more than two out dims"
return false
# Reduce stages, the chunk-domain and non-chunk reducer iterators
var dR: seq[string] = @[]
for ri in ir.reduceIters:
if ri in gf.dIters: dR.add(ri)
var nonD: seq[string] = @[]
for ri in ir.reduceIters:
if ri notin gf.dIters: nonD.add(ri)
# Mma fresh, a two-operand contraction over a tiled non-chunk iterator
if ir.expr.kind == xeMul and ir.reads.len == 2 and nonD.len == 1:
var aRead, bRead = -1
for r, rd in ir.reads:
if rd.bindsIter(ir.outDims[0]): aRead = r
elif rd.bindsIter(ir.outDims[1]): bRead = r
if aRead >= 0 and bRead >= 0 and aRead != bRead:
if ir.outDims[1] notin gf.dIters:
msg = "mma stage '" & me & "' has an unsupported column dim"
return false
if ir.outerTile(nonD[0]) == 0:
msg = "mma stage '" & me & "' must tile its contraction iterator"
return false
gf.classes[i] = scMmaFresh
gf.mmaAIdx[i] = aRead
gf.mmaBIdx[i] = bRead
gf.convIter[i] = nonD[0]
gf.convTile[i] = ir.outerTile(nonD[0])
gf.tileC[i] = gf.cs
return true
# Mma carry, self-carrying rescale plus contraction accumulate over D
if selfRead and ir.expr.kind == xeAdd and dR.len == 1 and nonD.len == 0:
var selfSide, otherSide: ExprIR
for cand in [ir.expr.lhs, ir.expr.rhs]:
if cand.kind != xeMul: continue
let selfHere = (cand.lhs.kind == xeRead and
cand.lhs.read.tensor == me) or
(cand.rhs.kind == xeRead and cand.rhs.read.tensor == me)
if selfHere and selfSide == nil:
selfSide = cand
elif not selfHere and cand.lhs.kind == xeRead and
cand.rhs.kind == xeRead and otherSide == nil:
otherSide = cand
if selfSide != nil and otherSide != nil:
var scaleIdx = -1
var scaleLit = 1.0'f32
let se = (if selfSide.lhs.kind == xeRead and
selfSide.lhs.read.tensor == me: selfSide.rhs
elif selfSide.rhs.kind == xeRead and
selfSide.rhs.read.tensor == me: selfSide.lhs
else: ExprIR(kind: xeLit, value: 1.0'f32))
if se.kind == xeRead:
scaleIdx = gf.targetIndex(se.read.tensor)
if scaleIdx < 0:
msg = "mma stage '" & me & "' rescales from an external tensor"
return false
elif se.kind == xeLit:
scaleLit = se.value
gf.classes[i] = scMmaCarry
gf.mmaScaleIdx[i] = scaleIdx
gf.mmaScaleLit[i] = scaleLit
gf.convIter[i] = dR[0]
gf.convTile[i] = 0
gf.tileC[i] = (if ir.outerTile(ir.outDims[1]) == 0: DefaultTile
else: ir.outerTile(ir.outDims[1]))
return true
if ir.outDims.len != 1:
msg = "reduce stage '" & me & "' must be one-dimensional"
return false
# A contraction-shaped source outside the mma forms is refused,
# the elementwise hadamard lowering would compute the wrong math
if ir.expr.kind == xeMul and ir.reads.len == 2:
msg = "contraction stage '" & me & "' has an unsupported mma shape"
return false
if nonD.len > 0:
msg = "reduce stage '" & me & "' reduces over a non-chunk iterator"
return false
var inChunk = ""
var chunkIt = ""
for di in dR:
if gf.irs[i].outerTile(di) > 0: inChunk = di
else: chunkIt = di
# Seed split, a max/add over an external row-bound read seeds
# the accumulator, the reduce source being the remaining operand,
# shared by the sweep and chunk forms
var srcE = ir.expr
if ir.expr.kind in {xeMax, xeAdd} and ir.reads.len == 2:
for r, rd in ir.reads:
if rd.tensor != me and gf.isExternal(rd.tensor) and
rd.pureIters(@[ir.outDims[0]]):
gf.seedOf[i].add(rd.tensor)
let lhsIsSeed = ir.expr.lhs.kind == xeRead and
ir.expr.lhs.read.tensor == rd.tensor
srcE = (if lhsIsSeed: ir.expr.rhs else: ir.expr.lhs)
# Sweep reduce, the reduction covering the chunk-loop iterator,
# the accumulator completing at the sweep end, both the full-axis
# form (chunk iter + in-chunk slice) and the bare form
if chunkIt.len > 0:
gf.classes[i] = scSweepReduce
gf.srcExpr[i] = srcE
return true
# Chunk reduce, the in-chunk slice only, fresh per chunk, seeded
# or fused into a carry later
if inChunk.len > 0:
gf.classes[i] = scChunkReduce
gf.srcExpr[i] = srcE
return true
msg = "reduce stage '" & me & "' has an unsupported reduction shape"
false
proc resolveDeferred(gf: GroupFacts, msg: var string): bool =
## Fold, prior-copy and chunk-reduce-fusion resolution over the assigned classes,
## the self-carry shapes and the mma-carry operand reads.
# Fold, one scaled external read consumed only by an mma A operand
for i in 0 ..< gf.names.len:
let ir = gf.irs[i]
var foldBase = ""
var foldScale: float32
if ir.reduceIters.len == 0 and ir.outDims.len == 2 and
ir.reads.len == 1 and ir.expr.isScaledRead(foldBase, foldScale) and
gf.isExternal(foldBase):
if gf.readsOf[i].len == 1 and gf.classes[gf.readsOf[i][0]] == scMmaFresh:
let mi = gf.readsOf[i][0]
if gf.irs[mi].reads[gf.mmaAIdx[mi]].tensor == gf.names[i]:
gf.classes[i] = scFold
gf.foldBase[i] = foldBase
gf.foldScale[i] = foldScale
continue
msg = "the fold '" & gf.names[i] & "' is not the mma stage's A operand"
return false
msg = "the scalar fold '" & gf.names[i] &
"' must feed exactly one mma stage's A operand"
return false
# A deferred 2-D elementwise stage that did not fold has no
# supported lowering, its column being neither the chunk domain
# nor the final stage's column
for i in 0 ..< gf.names.len:
if gf.classes[i] != scElemTile: continue
let ir = gf.irs[i]
if ir.outDims.len != 2: continue
let col = ir.outDims[1]
let finalCol = (if gf.irs[^1].outDims.len == 2: gf.irs[^1].outDims[1]
else: "")
if col notin gf.dIters and not (i == gf.names.len - 1 and col == finalCol):
msg = "elementwise stage '" & gf.names[i] &
"' has an unsupported column dim; only scalar folds feeding " &
"an mma A operand are supported there"
return false
# Prior copy, a bare copy of a carry stage's value snapshotting
# the prior chunk's value (the forward read) at any stage position,
# and a chunk reduce whose only reader is a self-carrying stage
# fusing into it
for i in 0 ..< gf.names.len:
let ir = gf.irs[i]
if ir.reduceIters.len == 0 and ir.reads.len == 1 and
ir.outDims.len == 1 and ir.expr.kind == xeRead:
let pi = gf.targetIndex(ir.reads[0].tensor)
if pi >= 0:
if not gf.isCarry(pi):
msg = "the forward read in '" & gf.names[i] & "' targets '" &
ir.reads[0].tensor & "', which carries no prior value"
return false
gf.classes[i] = scPriorCopy
gf.priorOf[i] = pi
continue
if gf.classes[i] == scChunkReduce and gf.seedOf[i].len == 0:
if gf.readsOf[i].len == 1:
let ci = gf.readsOf[i][0]
if gf.classes[ci] == scElemRowCarry:
gf.classes[i] = scChunkReduceFused
gf.fusedSrc[ci] = i
continue
# Elementwise self-carry shape, the scale operand and fused source
for i in 0 ..< gf.names.len:
if gf.classes[i] != scElemRowCarry: continue
let ir = gf.irs[i]
let me = gf.names[i]
var root = ir.expr
var scaleIdx = -1
var scaleLit = 1.0'f32
var fused = -1
if root.kind == xeAdd:
for cand in [root.lhs, root.rhs]:
if cand.kind == xeMul:
let selfHere = (cand.lhs.kind == xeRead and
cand.lhs.read.tensor == me) or
(cand.rhs.kind == xeRead and cand.rhs.read.tensor == me)
if selfHere:
let se = (if cand.lhs.kind == xeRead and
cand.lhs.read.tensor == me: cand.rhs
else: cand.lhs)
if se.kind == xeRead:
scaleIdx = gf.targetIndex(se.read.tensor)
if scaleIdx < 0:
msg = "self-carry '" & me & "' rescales from an external tensor"
return false
elif se.kind == xeLit:
scaleLit = se.value
else:
msg = "self-carry '" & me & "' has an unsupported scale operand"
return false
elif cand.kind == xeRead:
let fi2 = gf.targetIndex(cand.read.tensor)
if fi2 >= 0 and gf.classes[fi2] == scChunkReduceFused:
fused = fi2
else:
msg = "self-carry '" & me & "' has an unsupported addend"
return false
gf.carryScaleIdx[i] = scaleIdx
gf.carryScaleLit[i] = scaleLit
if fused >= 0:
gf.fusedSrc[i] = fused
else:
msg = "self-carry '" & me &
"' combines only with a fused chunk reduction"
return false
# Mma-carry operand reads, the A/B read slots
for i in 0 ..< gf.names.len:
if gf.classes[i] != scMmaCarry: continue
let ir = gf.irs[i]
var aIdx, bIdx = -1
for r, rd in ir.reads:
let pi = gf.targetIndex(rd.tensor)
if pi >= 0 and gf.classes[pi] in {scElemTile, scMmaFresh}:
aIdx = r
elif pi < 0:
bIdx = r
if aIdx < 0 or bIdx < 0:
msg = "mma stage '" & gf.names[i] &
"' needs one in-group tile operand and one external operand"
return false
gf.mmaAIdx[i] = aIdx
gf.mmaBIdx[i] = bIdx
true
proc chunkPlan(gf: GroupFacts, msg: var string): bool =
## Chunk size, per-stage column tiles and the chunk-domain index forms.
var cs = 0
for i in 0 ..< gf.irs.len:
let ir = gf.irs[i]
for di in gf.stageDIters(i):
let t = ir.outerTile(di)
if t > 0 and (cs == 0 or t < cs): cs = t
if cs == 0: cs = DefaultTile
gf.cs = cs
for i in 0 ..< gf.irs.len:
let ir = gf.irs[i]
if gf.stageDIters(i).len > 2:
msg = "stage '" & gf.names[i] &
"' binds more than two chunk-domain iterators"
return false
if gf.stageDIters(i).len == 2:
var tiled = 0
for d in gf.stageDIters(i):
if ir.outerTile(d) > 0: inc tiled
if tiled != 1:
msg = "stage '" & gf.names[i] &
"' must tile exactly one of its chunk-domain iterators"
return false
# Column tiles for the stages classification did not place
if gf.tileC[i] == 0 and ir.outDims.len == 2:
let col = ir.outDims[1]
let finalCol = (if gf.irs[^1].outDims.len == 2: gf.irs[^1].outDims[1]
else: "")
if col in gf.dIters:
# Chunk tile wins over the stage's schedule hint, the chunk
# loop stepping in cs units so the column tiles at cs
gf.tileC[i] = gf.cs
elif col == finalCol:
gf.tileC[i] = (if ir.outerTile(col) == 0: DefaultTile
else: ir.outerTile(col))
# Affine chunk-domain forms must be cs*chunk + in-chunk exactly
for i in 0 ..< gf.irs.len:
for rd in gf.irs[i].reads:
if gf.isExternal(rd.tensor): continue
for ix in rd.indices:
if ix.kind != ikAffine: continue
var dTerms: seq[tuple[coef: int32, iter: string]] = @[]
for t in ix.terms:
if t.iter in gf.dIters: dTerms.add(t)
if dTerms.len != 2: continue
var csCoef: int32
for t in dTerms:
# Chunk-loop iterator's coefficient is cs, the in-chunk
# iterator being the tiled one with coefficient 1
if gf.irs[i].outerTile(t.iter) == 0: csCoef = t.coef
if csCoef.int != gf.cs:
msg = "the chunk-domain index form of '" & gf.names[i] &
"' must be cs*chunk + in-chunk with cs = " & $gf.cs
return false
true
proc stageNeedsChunk(gf: GroupFacts, i: int, seen: var seq[int]): bool =
## Whether stage `i`'s value is recomputed per chunk.
##
## - reductions and mma forms inherently
## - elementwise stages through their operand closure, the chunk
## domain gathers, per-chunk temps and running reduces
if i in seen: return false
seen.add(i)
if gf.classes[i] in {scSweepReduce, scChunkReduce, scChunkReduceFused,
scMmaFresh, scMmaCarry, scPriorCopy}:
return true
if gf.classes[i] notin {scElemTile, scElemRow}: return false
for rd in gf.irs[i].reads:
let pi = gf.targetIndex(rd.tensor)
if pi < 0:
for di in gf.dIters:
if rd.bindsIter(di): return true
elif gf.classes[pi] in {scElemTile, scMmaFresh, scChunkReduce,
scChunkReduceFused, scPriorCopy}:
return true
elif gf.classes[pi] in {scElemRow, scElemRowCarry} and
gf.stageNeedsChunk(pi, seen):
return true
false
proc finalStageChunked(gf: GroupFacts, i: int, seen: var seq[int]): bool =
## Whether the FINAL stage recomputes per chunk, its closure
## holding a per-chunk temp or a chunk-domain gather.
##
## Carries and prior reads resolve to their sweep-end values,
## triggering no recompute.
if i in seen: return false
seen.add(i)
if gf.classes[i] in {scMmaFresh}: return true
if gf.classes[i] notin {scElemTile, scElemRow}: return false
for rd in gf.irs[i].reads:
let pi = gf.targetIndex(rd.tensor)
if pi < 0:
for di in gf.dIters:
if rd.bindsIter(di): return true
elif gf.classes[pi] in {scElemTile, scMmaFresh}:
return true
elif gf.classes[pi] in {scElemRow, scElemRowCarry} and
gf.finalStageChunked(pi, seen):
return true
# Carries, sweeps, chunk reduces and prior copies hold per-chunk
# valid or sweep-end values, they do not force the recompute
else:
discard
false
proc buildGroupFacts*(stmts: seq[StageRef], fromI, untilI: int, kernelName: string): GroupFacts =
## Facts of the statement range as one fusion group, nil when any fact fails,
## the fallback being the per-stage lowering.
var msg = ""
template fail(): untyped =
return nil
when defined(noFuse):
fail()
let n = untilI - fromI
if n < 2: return nil
# A stage reading a scalar-term index (a runtime value in an index form)
# has no per-chunk base-offset facts in the group kernel yet, so per-stage
# lowering handles it
for si in fromI ..< untilI:
let air = analyzeStage(stmts[si].fn, stmts[si].stageIdx,
stmts[si].sched, kernelName)
# atomic scatter stages and scalar-term reads lower per stage
if air.scatter: return nil
for rd in air.reads:
for ix in rd.indices:
if ix.scalars.len > 0: return nil
let gf = GroupFacts(stages: @[])
for i in fromI ..< untilI:
gf.stages.add(i)
for k in 0 ..< n:
let si = gf.stages[k]
let ir = analyzeStage(stmts[si].fn, stmts[si].stageIdx,
stmts[si].sched, kernelName)
gf.irs.add(ir)
gf.names.add(stmts[si].name)
gf.classes.add(scElemTile)
gf.passes.add(0)
gf.readsOf.add(@[])
gf.tileC.add(0)
gf.stageBuf.add("")
gf.priorOf.add(-1)
gf.mmaAIdx.add(-1)
gf.mmaBIdx.add(-1)
gf.mmaScaleIdx.add(-1)
gf.mmaScaleLit.add(1.0'f32)
gf.carryScaleIdx.add(-1)
gf.carryScaleLit.add(1.0'f32)
gf.fusedSrc.add(-1)
gf.convIter.add("")
gf.convTile.add(0)
gf.foldBase.add("")
gf.foldScale.add(1.0'f32)
gf.srcExpr.add(nil)
gf.seedOf.add(@[])
gf.needsChunk.add(false)
# One shared row dim, at most two out dims per stage
gf.rowDim = gf.irs[0].outDims[0]
for i in 0 ..< n:
if gf.irs[i].outDims.len < 1 or gf.irs[i].outDims.len > 2 or
gf.irs[i].outDims[0] != gf.rowDim:
msg = "stage '" & gf.names[i] & "' has no shared row dim"
fail()
# Readers of each target, in-group only
for i in 0 ..< n:
for rd in gf.irs[i].reads:
let pi = gf.targetIndex(rd.tensor)
if pi >= 0 and pi != i: gf.readsOf[pi].add(i)
## Physical axes and the chunk domain, over the group's slice only
var gstmts: seq[StageRef]
for k in 0 ..< n: gstmts.add(stmts[gf.stages[k]])
gf.axes = resolveAxes(gstmts, gf.irs)
for k in 0 ..< n:
for sym, dr in gstmts[k].fn.decls:
if sym in gf.irs[k].iters:
if dr.kind == rkDecl:
gf.decls[sym] = Range(kind: rkDecl, hostNode: dr.hostNode,
argName: kernelName & "_" & sym & "R")
else:
gf.decls[sym] = dr
# Chunk domain, the axis a recurrence read binds through an affine
# (cs*chunk + in-chunk) form declaring the chunking explicitly,
# otherwise the reduce axes' stage votes decide
var bestAxis: PhysAxis
var bestCount = 0
var affineAxes: seq[PhysAxis] = @[]
for i in 0 ..< n:
for rd in gf.irs[i].reads:
for ix in rd.indices:
if ix.kind != ikAffine: continue
for t in ix.terms:
if t.iter in gf.axes:
let a = gf.axes[t.iter]
if a notin affineAxes: affineAxes.add(a)
if affineAxes.len == 1:
bestAxis = affineAxes[0]
bestCount = 1
else:
var axisCount = initTable[PhysAxis, int]()
for i in 0 ..< n:
for ri in gf.irs[i].reduceIters:
if ri in gf.axes:
let a = gf.axes[ri]
if a notin axisCount: axisCount[a] = 0
inc axisCount[a]
if axisCount[a] > bestCount:
bestCount = axisCount[a]
bestAxis = a
if bestCount == 0:
msg = "no chunk domain: no recurrence read and no reduce axis"
fail()
gf.dAxis = bestAxis
for it, a in gf.axes:
if a == bestAxis: gf.dIters.add(it)
if gf.rowDim notin gf.axes:
msg = "the row dim '" & gf.rowDim & "' resolves to no physical axis"
fail()
gf.rowAxis = gf.axes[gf.rowDim]
# Classification, then the deferred folds / prior copies / fusions
for i in 0 ..< n:
if not gf.classifyStage(i, msg): fail()
if not gf.resolveDeferred(msg): fail()
if not gf.chunkPlan(msg): fail()
# No statement outside the range may read a group intermediate,
# the fused kernel materializing only its final target
if not stmts.noEscapes(fromI, untilI):
msg = "an outside statement reads a group intermediate"
fail()
# A carry stage must exist, else the group adds nothing
gf.groupCarries = false
for i in 0 ..< n:
if gf.isCarry(i) or gf.classes[i] == scSweepReduce:
gf.groupCarries = true
if not gf.groupCarries:
msg = "no carry stage: the group adds nothing over per-stage launches"
fail()
# Mma stages need the mma pass
when defined(noMmaPass):
for c in gf.classes:
if c in {scMmaFresh, scMmaCarry}:
msg = "the group has mma stages under -d:noMmaPass"
fail()
# Final stage's store geometry must resolve
let fi = n - 1
for od in gf.irs[fi].outDims:
if od != gf.rowDim and od notin gf.dIters and od notin gf.axes:
msg = "the final stage's out dim '" & od & "' resolves to no axis"
fail()
# Shared effective row tile
let tr = gf.stageTileR(0)
for i in 1 ..< n:
if gf.stageTileR(i) != tr:
msg = "stages disagree on the row tile"
fail()
gf.tileR = tr
# A carry stage makes the group a recurrence, its sweep reduces
# updating per chunk, the running value being what the recipe reads
for c in gf.classes:
if c in {scElemRowCarry, scMmaCarry}: gf.hasCarry = true
# Visibility passes, per-chunk producers readable same-pass, sweep
# producers forcing the consumer one pass later. In a recurrence
# group the sweep reduces are progressive 3-arg row ops, readable
# as running values
proc perChunk(c: StageClass): bool =
# A fused chunk reduce exists only inside its carry's update,
# a plain chunk reduce being a progressive accumulator whose
# statement value completes at the sweep end
c in {scElemTile, scElemRow, scPriorCopy, scElemRowCarry,
scChunkReduceFused, scMmaFresh} or
(gf.hasCarry and c == scSweepReduce)
for i in 0 ..< n:
for rd in gf.irs[i].reads:
let pi = gf.targetIndex(rd.tensor)
if pi < 0 or pi >= i: continue
if perChunk(gf.classes[pi]):
gf.passes[i] = max(gf.passes[i], gf.passes[pi])
else:
gf.passes[i] = max(gf.passes[i], gf.passes[pi] + 1)
for i in 0 ..< n:
var seen: seq[int] = @[]
gf.needsChunk[i] = gf.stageNeedsChunk(i, seen)
var seenF: seq[int] = @[]
gf.finalIsChunked = gf.finalStageChunked(fi, seenF)
# External inputs, first-appearance order, widest read rank
for i in 0 ..< n:
for rd in gf.irs[i].reads:
if gf.isExternal(rd.tensor):
var found = false
for e in gf.extRank:
if e.tensor == rd.tensor: found = true
if not found:
gf.extRank.add((rd.tensor, max(2, rd.indices.len)))
else:
for k in 0 ..< gf.extRank.len:
if gf.extRank[k].tensor == rd.tensor:
gf.extRank[k].rank =
max(gf.extRank[k].rank, max(2, rd.indices.len))
gf.threads = stmts.groupThreads(fromI, untilI)
result = gf
proc planFusionGroups*(stmts: seq[StageRef], kernelName: string): seq[FusionGroup] =
## Recipe statements partitioned into fusion groups, singleton groups lowering
## one kernel per stage, all of them under -d:noFuse.
when defined(noFuse):
for i in 0 ..< stmts.len:
result.add(FusionGroup(stages: @[i], kind: fkNone,
description: "no fusion, -d:noFuse"))
return
var i = 0
while i < stmts.len:
var g = FusionGroup(stages: @[i], kind: fkNone, description: "")
var advance = 1
# Longest run of consecutive stages whose facts hold
var until = stmts.len
while until > i + 1:
let gf = stmts.buildGroupFacts(i, until, kernelName)
if gf != nil:
var desc = "group fused: "
for k in 0 ..< gf.names.len:
desc &= (if k > 0: "," else: "") & gf.names[k]
desc &= " -> one launch, chunk domain " & gf.dAxis.tensor & "[" &
$gf.dAxis.axis & "] x" & $gf.cs
g = FusionGroup(stages: gf.stages, kind: fkChain, description: desc)
advance = until - i
break
dec until
result.add(g)
inc i, advance
# ─────────────────────────────────────────────────────────────────────────
# Group kernel body, the generic multi-stage tile program
# ─────────────────────────────────────────────────────────────────────────
#
# consts (atom, tiles) -> ranges -> declares (carries) -> seed inits
# -> per pass, one chunk loop over the stages in stage order,
# sweep ends after the loop
# -> the final stage's store
#
# Buffers, tl_<target> tiles, rv_<target> row vectors,
# gt<k>/tm<k> fresh per-chunk temporaries.
# A re-emitted value, a later pass re-reading a per-chunk temp,
# gets its own fresh buffer.
#
# In-place tile updates (row ops, scalar folds, exp2) stay safe
# while every reader of that producer precedes the updating stage,
# otherwise the value copies through a fresh buffer first.
proc valueBuffer(gf: GroupFacts, i: int, fresh: string = ""): string =
## Stage `i`'s value buffer, the fresh name overriding when re-emitted.
(if fresh.len > 0: fresh
elif gf.stageBuf[i].len > 0: gf.stageBuf[i]
elif gf.tileC[i] > 0: "tl_" & gf.names[i]
else: "rv_" & gf.names[i])
proc fresh(self: var Emitter, prefix: string): string =
result = prefix & $self.counter
inc self.counter
proc emitStageValue(self: var Emitter, i: int, pass: int, freshOverride = ""): tuple[tile, rv: string]
## Forward declaration of the stage emission, `valTile` re-emitting
## per-chunk temporaries for later passes through it.
proc canMutate(gf: GroupFacts, producer, by: int): bool =
## In-place update of `producer`'s buffer by stage `by` is safe
## when every reader of the producer precedes `by`.
for r in gf.readsOf[producer]:
if r > by: return false
true
proc emitPadFill(self: var Emitter, tile: string, padNegInf: bool) =
## Out-of-range cells of a per-chunk tile value masked, 0 or -inf, covering rows
## and columns, the chunk-domain column width is tileCc.
let pad = (if padNegInf: negInfNode() else: newLit(0.0'f32))
let tileN = id(tile)
self.passBody.add(quote do:
for n in 0'i32 ..< int32(tileRc div atomM):
for m in 0'i32 ..< int32(tileCc div atomN):
for v in 0'i32 ..< int32(atomVpt):
if by * tileRc + laneRow + n * atomM >= rowRange or
cb * tileCc + laneCol + m * atomN + v >= colRange:
`tileN`.frags[n][m].frag[v] = `pad`)
proc emitSeedGather(self: var Emitter, rv: string, seedT: string) =
## Per-row seed load into a row vector, guarded by the row range.
let stride = callI32(layoutField(seedT, "stride", 0))
let rvN = id(rv)
let seedN = id(seedT)
self.passBody.add(quote do:
for n in 0'i32 ..< int32(tileRc div atomM):
let row = by * tileRc + laneRow + n * atomM
if row < rowRange:
let seedVal = `seedN`.data[row * `stride`]
for v in 0'i32 ..< int32(atomVpt):
`rvN`.data[n * atomVpt + v] = seedVal)
proc emitGatherTile(self: var Emitter, rd: ReadIR, colTile: int, padNegInf: bool): string =
## One external read's per-chunk gather, tile-unit origins (by, cb)
## and absolute limits, zero-filled tails, -inf padding when a max
## reduce consumes the tile.
let name = self.fresh("gt")
let gd = "gd" & $self.counter
var base: NimNode = i32lit(0)
for slot, ix in rd.indices:
if ix.kind == ikConst and ix.value != 0:
base = infixF("+", base,
infixF("*", i32lit(ix.value.int),
callI32(layoutField(rd.tensor, "stride", slot))))
let bStr = (if rd.indices.len >= 4:
callI32(layoutField(rd.tensor, "stride", 0)) else: i32lit(0))
let dStr = (if rd.indices.len >= 5:
callI32(layoutField(rd.tensor, "stride", 1)) else: i32lit(0))
let rStr = callI32(layoutField(rd.tensor, "stride", 0))
let cStr = callI32(layoutField(rd.tensor, "stride", 1))
self.passBody.add(gdViewNode(gd, rd.tensor, bStr, dStr, rStr, cStr, base))
self.passBody.add(declTile(name, self.gf.tileR, colTile))
self.passBody.add(callOp(name, "loadTileBounded", id(gd),
nnkPar.newTree(newLit(0), newLit(0), id"by", id"cb"),
id"rowRange", id"colRange"))
if padNegInf:
self.emitPadFill(name, true)
name
proc copyTile(self: var Emitter, src: string, r, c: int): string =
## Fresh copy of a tile, the tile copy atom.
let dst = self.fresh("cp")
self.passBody.add(declTile(dst, r, c))
self.passBody.add(callOp(dst, "copy", id(src)))
dst
proc valTile(self: var Emitter, i: int, e: ExprIR, pass: int,
scalar: var NimNode, srcPadMax: bool,
freshOverride = ""): tuple[tile, rv: string] =
## Tile-path value of an expression, tile and rv operands
## combined through tile ops, row ops and scalar multiplies.
##
## Tile operands from earlier passes re-emit their computation
## into the current body.
case e.kind
of xeLit:
scalar = (if scalar == nil: newLit(e.value)
else: infixF("*", scalar, newLit(e.value)))
("", "")
of xeRead:
let rd = e.read
let pi = self.gf.targetIndex(rd.tensor)
if pi < 0:
doAssert rd.indices.len == 2 and
rd.indices[0].kind == ikIter and
rd.indices[0].name == self.gf.rowDim,
"group gathers bind (row, column) iterators"
let tile = self.emitGatherTile(rd, self.gf.cs, srcPadMax)
(tile, "")
elif rd.tensor == self.gf.names[i]:
(self.gf.valueBuffer(i, freshOverride), "")
elif self.gf.classes[pi] in {scElemTile, scMmaFresh, scMmaCarry}:
if pass > self.gf.passes[pi] and
self.gf.classes[pi] in {scElemTile, scMmaFresh}:
let rn = self.fresh("re")
let (t, _) = self.emitStageValue(pi, pass, rn)
return (t, "")
# Fresh override names the re-emitted stage's own buffer only,
# a producer read as an operand keeping its stored buffer
(self.gf.valueBuffer(pi), "")
else:
("", self.gf.valueBuffer(pi))
of xeMux:
# Per-lane select, the condition's row and column iterators resolving to the lane cell's coordinates (banded causal masks).
# A literal branch carries its value per cell, the zero tile node marking the constant branch.
var thenScalar: NimNode = nil
let (tt0, _) = self.valTile(i, e.lhs, pass, thenScalar, srcPadMax,
freshOverride)
var elseScalar: NimNode = nil
let (et0, _) = self.valTile(i, e.rhs, pass, elseScalar, srcPadMax,
freshOverride)
var tn: NimNode = i32lit(0)
var tv: NimNode = nil
if tt0.len > 0: tn = id(tt0)
else: tv = thenScalar
var en: NimNode = i32lit(0)
var ev: NimNode = nil
if et0.len > 0: en = id(et0)
else: ev = elseScalar
let tmp = self.fresh("xm")
self.passBody.add(declTile(tmp, self.gf.tileR, self.gf.tileC[i]))
var laneCoords = initTable[string, char]()
laneCoords[self.gf.rowDim] = 'r'
for dIt in self.gf.dIters:
laneCoords[dIt] = 'c'
let colBase = infixF("*", id"cb", id"tileCc")
emitLaneMux(self.passBody, id(tmp), tn, tv, en, ev, e.cond,
laneCoords, colBase)
return (tmp, "")
of xeMul, xeAdd, xeSub, xeDiv, xeMax:
let (lt, lrv) = self.valTile(i, e.lhs, pass, scalar, srcPadMax,
freshOverride)
let (rt, rrv) = self.valTile(i, e.rhs, pass, scalar, srcPadMax,
freshOverride)
if e.kind == xeMul:
if lt.len == 0 and rt.len == 0:
doAssert lrv.len == 0 and rrv.len == 0,
"a row vector cannot multiply a scalar in a tile expression"
return ("", "")
if lt.len == 0 and lrv.len == 0: return (rt, rrv)
if rt.len == 0 and rrv.len == 0: return (lt, lrv)
if rrv.len > 0:
let pi = self.gf.targetIndex(readTensor(e.rhs))
let base = (if pi >= 0 and not self.gf.canMutate(pi, i):
self.copyTile(lt, self.gf.tileR, self.gf.tileC[i]) else: lt)
self.passBody.add(callOp(base, "mul_row", id(base), id(rrv)))
return (base, "")
if lrv.len > 0:
self.passBody.add(callOp(rt, "mul_row", id(rt), id(lrv)))
return (rt, "")
let tmp = self.fresh("tm")
self.passBody.add(declTile(tmp, self.gf.tileR, self.gf.tileC[i]))
self.passBody.add(callOp(tmp, "mul", id(lt), id(rt)))
return (tmp, "")
if lt.len == 0 and rt.len == 0:
# Scalar op scalar, the literal factors staying in `scalar`
return ("", "")
var ltile = lt
var rtile = rt
if ltile.len == 0:
# Scalar-lhs summands commute into the tile path, rv-lhs has no
# row op at all
doAssert rrv.len == 0 and e.kind == xeAdd,
"a scalar or row-vector lhs needs a tilewise add to commute"
(ltile, rtile) = (rtile, ltile)
doAssert rtile.len > 0 or rrv.len > 0,
"a literal summand of a tilewise map has no broadcast form"
if rrv.len > 0:
doAssert e.kind in {xeSub, xeDiv},
"max with a row vector has no row op"
let op = (if e.kind == xeSub: "sub_row" else: "div_row")
let pi = self.gf.targetIndex(readTensor(e.lhs))
let base = (if pi >= 0 and not self.gf.canMutate(pi, i):
self.copyTile(ltile, self.gf.tileR, self.gf.tileC[i]) else: ltile)
self.passBody.add(callOp(base, op, id(base), id(rrv)))
return (base, "")
doAssert lrv.len == 0, "a row vector cannot be the mapped operand"
let tmp = self.fresh("tm")
self.passBody.add(declTile(tmp, self.gf.tileR, self.gf.tileC[i]))
case e.kind
of xeSub, xeDiv, xeMax:
let f: proc(x, y: NimNode): NimNode = (case e.kind
of xeSub: (proc (x, y: NimNode): NimNode = infixF("-", x, y))
of xeDiv: (proc (x, y: NimNode): NimNode = infixF("/", x, y))
else: (proc (x, y: NimNode): NimNode =
nnkCall.newTree(id"max", x, y)))
emitLaneMap2(self.passBody, id(tmp), id(ltile), id(rtile),
id"tileRc",
(if self.gf.tileC[i] == self.gf.cs: id"tileCc" else: id"tileOc"), f)
else:
self.passBody.add(callOp(tmp, "convert", id(ltile)))
self.passBody.add(callOp(tmp, "addScaled", id(rtile), newLit(1.0'f32)))
(tmp, "")
of xeExp2:
let (lt, lrv) = self.valTile(i, e.lhs, pass, scalar, srcPadMax,
freshOverride)
if lrv.len > 0:
self.passBody.add(callOp(lrv, "exp2", id(lrv)))
return ("", lrv)
# A pending scalar factor is the operand's scale, exp2 seeing
# the scaled operand, the factor applying here, not above the unary
var lt2 = lt
if scalar != nil and lt2.len > 0:
let ts = self.fresh("sc")
self.passBody.add(declTile(ts, self.gf.tileR, self.gf.tileC[i]))
self.passBody.add(callOp(ts, "mul", id(lt2), scalar))
scalar = nil
lt2 = ts
let pi = self.gf.targetIndex(readTensor(e.lhs))
let base = (if pi >= 0 and not self.gf.canMutate(pi, i):
self.copyTile(lt2, self.gf.tileR, self.gf.tileC[i]) else: lt2)
self.passBody.add(callOp(base, "exp2", id(base)))
(base, "")
proc valRow(self: var Emitter, i: int, e: ExprIR, pass: int): string =
## Row-vector path of an elementwise stage's value, rv ops only.
case e.kind
of xeRead:
let pi = self.gf.targetIndex(e.read.tensor)
doAssert pi >= 0, "a row-vector read binds an in-group target"
self.gf.valueBuffer(pi)
of xeLit:
raiseAssert("a literal has no row-vector broadcast form")
of xeMul, xeSub:
if e.lhs.kind == xeLit:
let tmp = self.valRow(i, e.rhs, pass)
self.passBody.add(callOp(tmp, "mul", id(tmp), newLit(e.lhs.value)))
return tmp
if e.rhs.kind == xeLit:
let tmp = self.valRow(i, e.lhs, pass)
self.passBody.add(callOp(tmp, "mul", id(tmp), newLit(e.rhs.value)))
return tmp
let la = self.valRow(i, e.lhs, pass)
let rb = self.valRow(i, e.rhs, pass)
let tmp = self.fresh("tr")
self.passBody.add(declRv(tmp, self.gf.tileR, self.gf.cs))
self.passBody.add(callOp(tmp,
(if e.kind == xeSub: "sub" else: "mul"), id(la), id(rb)))
tmp
of xeExp2:
let la = self.valRow(i, e.lhs, pass)
let tmp = self.fresh("tr")
self.passBody.add(declRv(tmp, self.gf.tileR, self.gf.cs))
self.passBody.add(callOp(tmp, "exp2", id(la)))
tmp
else:
raiseAssert("unsupported row-vector op " & $e.kind)
proc groupRange(gf: GroupFacts, iter: string, axis: PhysAxis): Range =
## Group range for `iter`, the declared domain's runtime value
## when the iterator is declared and the physical axis read otherwise.
if gf.decls.hasKey(iter):
return gf.decls[iter]
Range(kind: rkAxis, tensor: axis.tensor, axis: axis.axis)
proc emitMmaFresh(self: var Emitter, i: int): string =
## Fresh contraction tile, zeroed with the contraction kk loop
## loading the A and B operand tiles (the fold applied to A), one
## mma per slab, the padding masked for the consumers.
##
## The A tile rt_l rows the row dim, cols the contraction.
## The B tile rt_r rows the column dim, cols the contraction,
## the K axis sitting in the gd column slot.
let gf = self.gf
let ir = gf.irs[i]
let att = valueBuffer(gf, i)
let aT = self.fresh("ma")
let bT = self.fresh("mb")
let convTile = gf.convTile[i]
## Contraction range comes from the resolved physical axis,
## the read tensor being a folded internal (qs) in general.
let ca = gf.axes[gf.convIter[i]]
let convExt = rangeExpr(groupRange(gf, gf.convIter[i], ca), device = true)
let rdA0 = ir.reads[gf.mmaAIdx[i]]
let rdB = ir.reads[gf.mmaBIdx[i]]
var aTensor = rdA0.tensor
var aGeom = rdA0
var foldScale: NimNode = nil
let fi2 = gf.targetIndex(rdA0.tensor)
if fi2 >= 0 and gf.classes[fi2] == scFold:
aGeom = gf.irs[fi2].reads[0]
aTensor = gf.foldBase[fi2]
foldScale = newLit(gf.foldScale[fi2])
let aRSlot = aGeom.boundSlotIn(gf.rowDim)
let aKSlot = aGeom.boundSlotIn(gf.convIter[i])
let bNSlot = rdB.boundSlotIn(ir.outDims[1])
let bKSlot = rdB.boundSlotIn(gf.convIter[i])
doAssert aRSlot >= 0 and aKSlot >= 0 and aRSlot != aKSlot and
bNSlot >= 0 and bKSlot >= 0 and bNSlot != bKSlot,
"the mma operands bind the plane and contraction iterators"
let convBound = callI32(infixF("div",
infixF("+", convExt, infixF("-", newLit(convTile), newLit(1'i32))),
newLit(convTile)))
self.passBody.add(declTile(aT, gf.tileR, convTile))
self.passBody.add(declRtR(bT, convTile, gf.cs))
self.passBody.add(callOp(att, "zero"))
var kkBody = nnkStmtList.newTree()
kkBody.add(gdViewNode("gd" & aT, aTensor, i32lit(0), i32lit(0),
callI32(layoutField(aTensor, "stride", aRSlot)),
callI32(layoutField(aTensor, "stride", aKSlot)), i32lit(0)))
kkBody.add(callOp(aT, "loadTileBounded", id("gd" & aT),
nnkPar.newTree(newLit(0), newLit(0), id"by", id"kk"),
id"rowRange", convExt))
if foldScale != nil:
kkBody.add(callOp(aT, "mul", id(aT), foldScale))
kkBody.add(gdViewNode("gd" & bT, rdB.tensor, i32lit(0), i32lit(0),
callI32(layoutField(rdB.tensor, "stride", bNSlot)),
callI32(layoutField(rdB.tensor, "stride", bKSlot)), i32lit(0)))
kkBody.add(callOp(bT, "loadTileBounded", id("gd" & bT),
nnkPar.newTree(newLit(0), newLit(0), id"cb", id"kk"),
id"colRange", convExt))
kkBody.add(callOp(att, "mma_AB", id(aT), id(bT)))
self.passBody.add(nnkForStmt.newTree(id"kk",
nnkInfix.newTree(id"..<", newLit(0'i32), convBound),
nnkStmtList.newTree(kkBody)))
if gf.maxConsumer(i):
self.emitPadFill(att, true)
att
proc emitMmaCarry(self: var Emitter, i: int) =
## Self-carrying contraction accumulate, rescale by row vector
## or scalar then one mma over the chunk slice.
##
## The B operand's rows are the final column (N), its cols
## being the chunk-domain slice (K).
let gf = self.gf
let ir = gf.irs[i]
let acc = valueBuffer(gf, i)
if gf.mmaScaleIdx[i] >= 0:
self.passBody.add(callOp(acc, "mul_row", id(acc),
id(valueBuffer(gf, gf.mmaScaleIdx[i]))))
elif gf.mmaScaleLit[i] != 1.0'f32:
self.passBody.add(callOp(acc, "mul", id(acc), newLit(gf.mmaScaleLit[i])))
let rdB = ir.reads[gf.mmaBIdx[i]]
let bT = self.fresh("mb")
let kSlot = rdB.boundSlotIn(gf.convIter[i])
let nSlot = rdB.boundSlotIn(ir.outDims[1])
doAssert kSlot >= 0 and nSlot >= 0 and kSlot != nSlot,
"the mma B operand binds the contraction and the column dim"
let kStride = callI32(layoutField(rdB.tensor, "stride", kSlot))
let nStride = callI32(layoutField(rdB.tensor, "stride", nSlot))
let nRange = callI32(layoutField(rdB.tensor, "shape", nSlot))
let kLimit = nnkCall.newTree(id"min",
infixF("+", infixF("*", id"cb", newLit(gf.cs)), newLit(gf.cs)),
id"colRange")
self.passBody.add(declRtR(bT, gf.cs, gf.tileC[i]))
self.passBody.add(gdViewNode("gd" & bT, rdB.tensor, i32lit(0), i32lit(0),
nStride, kStride, i32lit(0)))
self.passBody.add(callOp(bT, "loadTileBounded", id("gd" & bT),
nnkPar.newTree(newLit(0), newLit(0), id"bx", id"cb"),
nRange, kLimit))
self.passBody.add(callOp(acc, "mma_AB",
id(valueBuffer(gf, gf.targetIndex(
ir.reads[gf.mmaAIdx[i]].tensor))), id(bT)))
proc emitStore(self: var Emitter, i: int) =
# Group's final value stored to the output, per chunk when the final
# stage recomputes per chunk, once otherwise.
let gf = self.gf
let ir = gf.irs[i]
let oBat, oDep: NimNode = i32lit(0)
let oRow = callI32(layoutField("o", "stride", 0))
let oCol = (if ir.outDims.len == 2: callI32(layoutField("o", "stride", 1))
else: i32lit(0))
self.passBody.add(gdViewNode("gdO", "o", oBat, oDep, oRow, oCol, i32lit(0)))
if ir.outDims.len == 2 and gf.finalIsChunked:
self.passBody.add(callOp("gdO", "storeTileMasked",
id(valueBuffer(gf, i)),
nnkPar.newTree(newLit(0), newLit(0), id"by", id"cb"),
infixF("-", id"rowRange", infixF("*", id"by", id"tileRc")),
infixF("-", id"colRange", infixF("*", id"cb", id"tileCc"))))
elif ir.outDims.len == 2:
self.passBody.add(callOp("gdO", "storeTileMasked",
id(valueBuffer(gf, i)),
nnkPar.newTree(newLit(0), newLit(0), id"by", id"bx"),
infixF("-", id"rowRange", infixF("*", id"by", id"tileRc")),
infixF("-", id"colRange2",
infixF("*", id"bx", newLit(gf.tileC[i])))))
else:
let rv = valueBuffer(gf, i)
let os1 = callI32(layoutField("o", "stride", 1))
let rvN = id(rv)
self.passBody.add(quote do:
for n in 0'i32 ..< int32(tileRc div atomM):
let row = by * tileRc + laneRow + n * atomM
if row < rowRange:
o.data[row * `os1`] = `rvN`.data[n * atomVpt])
proc emitStageValue(self: var Emitter, i: int, pass: int, freshOverride = ""): tuple[tile, rv: string] =
## One stage's value computed into the current pass body, its buffer names
## returned (tile for tile classes, rv for row classes).
let gf = self.gf
let ir = gf.irs[i]
let me = gf.names[i]
if gf.classes[i] == scFold:
return ("", "")
if freshOverride.len == 0 and gf.stageBuf[i].len == 0:
gf.stageBuf[i] = valueBuffer(gf, i)
case gf.classes[i]
of scElemTile:
var scalar: NimNode = nil
let (t, _) = self.valTile(i, ir.expr, pass, scalar, false, freshOverride)
doAssert t.len > 0, "elementwise stage '" & me & "' has no tile value"
var buf = t
if scalar != nil:
buf = self.fresh("ts")
self.passBody.add(declTile(buf, gf.tileR, gf.tileC[i]))
self.passBody.add(callOp(buf, "mul", id(t), scalar))
if freshOverride.len == 0:
gf.stageBuf[i] = buf
if gf.readsOf[i].len > 0 and gf.tileC[i] == gf.cs:
self.emitPadFill(buf, gf.maxConsumer(i))
result = (buf, "")
of scElemRow:
let rvName = self.valRow(i, ir.expr, pass)
if rvName != gf.stageBuf[i]:
self.passBody.add(callOp(gf.stageBuf[i], "copy", id(rvName)))
result = ("", gf.stageBuf[i])
of scPriorCopy:
self.passBody.add(callOp(gf.stageBuf[i], "copy",
id(valueBuffer(gf, gf.priorOf[i]))))
result = ("", gf.stageBuf[i])
of scElemRowCarry:
let dst = gf.stageBuf[i]
if gf.carryScaleIdx[i] >= 0:
self.passBody.add(callOp(dst, "mul", id(dst),
id(valueBuffer(gf, gf.carryScaleIdx[i]))))
elif gf.carryScaleLit[i] != 1.0'f32:
self.passBody.add(callOp(dst, "mul", id(dst),
newLit(gf.carryScaleLit[i])))
if gf.fusedSrc[i] >= 0:
let fi2 = gf.fusedSrc[i]
var scalar: NimNode = nil
let (srcTile, _) = self.valTile(fi2, gf.srcExpr[fi2], pass, scalar,
false)
if gf.irs[fi2].reduceKind == rkMax:
self.passBody.add(callOp(dst, "row_max", id(srcTile), id(dst)))
else:
self.passBody.add(callOp(dst, "row_sum", id(srcTile), id(dst)))
result = ("", dst)
of scSweepReduce:
var scalar: NimNode = nil
let (srcTile, _) = self.valTile(i, gf.srcExpr[i], pass, scalar,
gf.irs[i].reduceKind == rkMax)
if gf.hasCarry:
# Running value is the recipe's read, the 3-arg form folding in the prior accumulator.
let dst = valueBuffer(gf, i)
if ir.reduceKind == rkMax:
self.passBody.add(callOp(dst, "row_max", id(srcTile), id(dst)))
else:
self.passBody.add(callOp(dst, "row_sum", id(srcTile), id(dst)))
result = ("", dst)
else:
let acc = "acc_" & me
if ir.reduceKind == rkMax:
emitLaneMap2(self.passBody, id(acc), id(acc), id(srcTile),
id"tileRc", id"tileCc",
(proc (x, y: NimNode): NimNode = nnkCall.newTree(id"max", x, y)))
else:
self.passBody.add(callOp(acc, "addScaled", id(srcTile),
newLit(1.0'f32)))
result = ("", valueBuffer(gf, i))
of scChunkReduce:
let dst = gf.stageBuf[i]
var scalar: NimNode = nil
let (srcTile, _) = self.valTile(i, gf.srcExpr[i], pass, scalar,
gf.irs[i].reduceKind == rkMax)
if ir.reduceKind == rkMax:
self.passBody.add(callOp(dst, "row_max", id(srcTile), id(dst)))
else:
self.passBody.add(callOp(dst, "row_sum", id(srcTile), id(dst)))
result = ("", dst)
of scChunkReduceFused:
result = ("", "")
of scMmaFresh:
let att = self.emitMmaFresh(i)
gf.stageBuf[i] = att
result = (att, "")
of scMmaCarry:
self.emitMmaCarry(i)
result = (gf.stageBuf[i], "")
of scFold:
result = ("", "")
if i == gf.names.len - 1 and freshOverride.len == 0:
self.emitStore(i)
proc emitSweepEnd(self: var Emitter, i: int) =
## Sweep-end reduction of a sweep stage's accumulator into its row vector,
## after the pass's chunk loop.
let gf = self.gf
let rv = valueBuffer(gf, i)
let acc = "acc_" & gf.names[i]
if gf.irs[i].reduceKind == rkMax:
self.body.add(callOp(rv, "neg_infty"))
self.body.add(callOp(rv, "row_max", id(acc), id(rv)))
else:
self.body.add(callOp(rv, "row_sum", id(acc)))
proc lowerGroup*(stmts: seq[StageRef], fromI, untilI: int, kernelName: string,
externals = @[""], views = @[""]): tuple[fnDef, entry, launch: NimNode, threads: int] =
## One fusion group's device fn + thin entry + launch node, the chunk loop
## per pass, stage order inside, carries in registers.
let gf = stmts.buildGroupFacts(fromI, untilI, kernelName)
if gf == nil:
raiseAssert("lowerGroup: the fusion facts vanished between plan and lower")
# crucibleCompileTarget = ctMetal
let n = gf.names.len
let fi = n - 1
let finalIr = gf.irs[fi]
var body = nnkStmtList.newTree(
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"atomK", newEmptyNode(),
nnkCall.newTree(dot(id"atom", "getK"))),
nnkConstDef.newTree(id"atomVpt", newEmptyNode(),
nnkCall.newTree(dot(id"atom", "getVpt")))),
nnkConstSection.newTree(
nnkConstDef.newTree(id"tileRc", newEmptyNode(), newLit(gf.tileR)),
nnkConstDef.newTree(id"tileCc", newEmptyNode(), newLit(gf.cs)),
nnkConstDef.newTree(id"tileOc", newEmptyNode(), newLit(gf.tileC[fi]))),
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")))))
var tileAssert = infixF("and",
infixF("==", infixF("mod", id"tileRc", id"atomM"), newLit(0)),
infixF("==", infixF("mod", id"tileCc", id"atomN"), newLit(0)))
tileAssert = infixF("and", tileAssert,
infixF("==", infixF("mod", id"tileOc", id"atomN"), newLit(0)))
for i in 0 ..< n:
if gf.classes[i] == scMmaFresh:
tileAssert = infixF("and", tileAssert,
infixF("==", infixF("mod", newLit(gf.convTile[i]), id"atomK"),
newLit(0)))
body.add(nnkStaticStmt.newTree(nnkStmtList.newTree(
nnkCall.newTree(id"doAssert", tileAssert,
newLit("chunk, column and contraction tiles must be multiples of the atom's M, N and K")))))
# Ranges, row, chunk-domain and final column, a declared domain's
# runtime value taking its scalar arg's name
let rowRangeR = groupRange(gf, gf.rowDim, gf.rowAxis)
var chunkRangeR = Range(kind: rkAxis, tensor: gf.dAxis.tensor,
axis: gf.dAxis.axis)
for di in gf.dIters:
if gf.decls.hasKey(di):
chunkRangeR = gf.decls[di]
break
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"rowRange", newEmptyNode(), rangeExpr(rowRangeR, device = true))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"colRange", newEmptyNode(), rangeExpr(chunkRangeR, device = true))))
if finalIr.outDims.len == 2:
let fca = gf.axes[finalIr.outDims[1]]
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"colRange2", newEmptyNode(),
rangeExpr(groupRange(gf, finalIr.outDims[1], fca), device = true))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"lane", newEmptyNode(), nnkCall.newTree(id"int32",
ident"thread_index_in_threadgroup"))))
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"))))
body.add(nnkLetSection.newTree(nnkIdentDefs.newTree(
id"laneCol", newEmptyNode(), infixF("div", id"cell", id"atomM"))))
# Declares the carries and sweep accumulators, plus seed inits
var em = Emitter(gf: gf, body: body, passBody: nil)
for i in 0 ..< n:
case gf.classes[i]
of scMmaCarry:
em.body.add(declTile(valueBuffer(gf, i), gf.tileR, gf.tileC[i]))
em.body.add(callOp(valueBuffer(gf, i), "zero"))
of scMmaFresh:
em.body.add(declTile(valueBuffer(gf, i), gf.tileR, gf.tileC[i]))
of scElemRowCarry:
em.body.add(declRv(valueBuffer(gf, i), gf.tileR, gf.cs))
em.body.add(callOp(valueBuffer(gf, i), "zero"))
of scElemRow, scPriorCopy:
em.body.add(declRv(valueBuffer(gf, i), gf.tileR, gf.cs))
of scSweepReduce:
if gf.hasCarry:
em.body.add(declRv(valueBuffer(gf, i), gf.tileR, gf.cs))
if gf.irs[i].reduceKind == rkMax:
em.body.add(callOp(valueBuffer(gf, i), "neg_infty"))
else:
em.body.add(callOp(valueBuffer(gf, i), "zero"))
# Seed rows are chunk-independent, gathering in the decls
em.passBody = em.body
for seedT in gf.seedOf[i]:
em.emitSeedGather(valueBuffer(gf, i), seedT)
em.passBody = nil
else:
em.body.add(declTile("acc_" & gf.names[i], gf.tileR, gf.cs))
if gf.irs[i].reduceKind == rkMax:
em.body.add(callOp("acc_" & gf.names[i], "map",
id("acc_" & gf.names[i]), negInfNode()))
else:
em.body.add(callOp("acc_" & gf.names[i], "zero"))
em.body.add(declRv(valueBuffer(gf, i), gf.tileR, gf.cs))
of scChunkReduce:
em.body.add(declRv(valueBuffer(gf, i), gf.tileR, gf.cs))
if gf.irs[i].reduceKind == rkMax:
em.body.add(callOp(valueBuffer(gf, i), "neg_infty"))
else:
em.body.add(callOp(valueBuffer(gf, i), "zero"))
# Seed rows are chunk-independent, gathering in the decls
em.passBody = em.body
for seedT in gf.seedOf[i]:
em.emitSeedGather(valueBuffer(gf, i), seedT)
em.passBody = nil
else:
discard
# Passes, ascending, each its own chunk loop when any of its stages
# recomputes per chunk
var maxPass = 0
for p in gf.passes:
maxPass = max(maxPass, p)
for p in 0 .. maxPass:
var passStages: seq[int] = @[]
for i in 0 ..< n:
if gf.passes[i] == p and gf.classes[i] != scFold:
passStages.add(i)
if passStages.len == 0: continue
var loopPass = false
for i in passStages:
if gf.needsChunk[i]: loopPass = true
em.passBody = nnkStmtList.newTree()
# A prior copy must snapshot the carried value before the same
# pass's update writes it, so they emit ahead of stage order
for i in passStages:
if gf.classes[i] == scPriorCopy:
discard em.emitStageValue(i, p)
for i in passStages:
if gf.classes[i] == scPriorCopy: continue
discard em.emitStageValue(i, p)
var passChunk = em.passBody
if loopPass:
passChunk = nnkForStmt.newTree(id"cb",
nnkInfix.newTree(id"..<", newLit(0'i32),
callI32(infixF("div",
infixF("+", id"colRange", newLit(gf.cs - 1)), newLit(gf.cs)))),
nnkStmtList.newTree(passChunk))
body.add(passChunk)
for i in passStages:
if gf.classes[i] == scSweepReduce and not gf.hasCarry:
em.emitSweepEnd(i)
## Fused kernel's records and launch
var desc = "fused group: " & gf.names.join(",") &
" -> one launch, chunk domain " & gf.dAxis.tensor & "[" &
$gf.dAxis.axis & "] x" & $gf.cs
let fusedName = kernelName & "_fused"
var fir = AlgoIR(name: fusedName, target: "o",
threads: (if gf.threads > 0: gf.threads else: DefaultThreads))
fir.outDims = finalIr.outDims
for e in gf.extRank:
var idx: seq[IndexRef] = @[]
for k in 0 ..< e.rank:
idx.add(IndexRef(kind: ikConst, value: 0))
fir.reads.add(ReadIR(tensor: e.tensor, indices: idx))
var gridX: NimNode
let gridY = nnkCall.newTree(id"cdiv",
rangeExpr(rowRangeR, device = false), newLit(gf.tileR))
if gf.finalIsChunked or finalIr.outDims.len < 2:
gridX = newLit(1)
else:
let fca = gf.axes[finalIr.outDims[1]]
gridX = nnkCall.newTree(id"cdiv",
rangeExpr(groupRange(gf, finalIr.outDims[1], fca), device = false),
newLit(gf.tileC[fi]))
var shape: seq[NimNode] = @[]
shape.add(rangeExpr(rowRangeR, device = false))
if finalIr.outDims.len == 2:
let fca = gf.axes[finalIr.outDims[1]]
shape.add(rangeExpr(groupRange(gf, finalIr.outDims[1], fca),
device = false))
else:
shape.add(newLit(1))
let grid = nnkPar.newTree(gridX, gridY, newLit(1))
let outShape = nnkPrefix.newTree(ident"@", nnkBracket.newTree(shape))
KernelTrace[fusedName] = newLit(fusedName & ": " & desc & "\n")
SchedInfo[fusedName] = newLit(desc)
for ir in gf.irs:
for d in ir.declArgs:
var known = false
for e in fir.declArgs:
if e.argName == d.argName: known = true
if not known:
fir.declArgs.add(d)
let fnDef = emitDeviceFunc(fir, body, @[])
let entry = emitEntry(fir, grid, outShape, nnkStmtList.newTree(), @[])
let entryName = id(fusedName & "Entry")
var call = nnkCall.newTree(entryName)
for e in gf.extRank:
call.add(id((if e.tensor in views: kernelName & "_" & e.tensor
elif e.tensor in externals: e.tensor
else: kernelName & "_" & e.tensor)))
for d in fir.declArgs:
call.add(id(d.argName))
let launch = nnkCall.newTree(id"evalCall", id"engine", entryName, call,
newLit(fusedName))
result = (fnDef, entry, launch, fir.threads)
## Forward declaration resolved, the emitter stage emission
## and its value walker re-emitting per-chunk temporaries
## inside the later visibility passes.