gpu_fake

# A GPU program on a machine with no GPU: nine kernels, two sides, one
# handler that is a table in host memory.
#
# The last line of the gallery is not a kernel at all: it is the host asking
# the module what globals it declares (knife TG), which is the one `Gpu`
# operation that names something inside a module rather than a buffer the
# program allocated.
#
# The device side is nine ordinary Dawn functions whose only effect is
# `!Dev`, each the smallest shape that shows one thing the dialect can do
# and each named for the knife that added it (docs/tile-backend-design.md
# 7). Running one computes nothing: `packages/tileir`'s recording handler
# writes down the operations it raises, and the record renders as Tile IR
# text or encodes as the bytecode `tileiras` assembles into a cubin.
#
# The host side is `run_kernel` -- allocate, upload, launch, download, free
# -- whose only effect is `!Gpu`. `std/gpu`'s `with_gpu_real` answers those
# operations from the CUDA driver; `with_gpu_fake` answers them from a table
# in host memory and is pure, so this program and its tests need no device,
# no driver and no `!io`.
#
# The fake device never runs a kernel body. It looks the launch's name up in
# a map and calls the host reference registered under it, so every answer
# below is a reference (`vadd_ref` from `std/gpu`, the rest from
# `packages/tileref`) and not this kernel's transcript: each test
# holds a reference against numbers written down by hand, and it is the
# device that holds the reference against the kernel. It is not a stub
# either. It issues and retires handles, refuses a format it cannot hold, a
# kernel it does not know and a launch of the wrong arity, and rounds what
# it stores to the buffer's format.
#
# Run: dawn run examples/projects/gpu_fake

use std/gpu
use std/gpu.{Gpu, WideRefFn, last_out, with_gpu_fake, reference_kernels, alloc, copy_from_host, copy_to_host, launch, free,
  sync, vadd_ref, buffer_arg, erase }
use std/dtype.{F64}
use std/narrow.{F32, F8E4M3FN}
use std/int/i32.{I32}
use tileref/ref.{loop_count_ref, assert_pass_ref, global_table_ref, alloca_scratch_ref, shape_ops_ref,
  dtype_convert_ref }
use std/io
use std/list
use std/map
use std/narrow
use std/str
use tileir/dev.{Dev, Param, Tile, param, block_id, tile_at, load, store, add, load_cell, store_cell,
  for_arch, hint_latency, d_loop, carry, get, set, d_reduce, broadcast, DYN_DIM,
  i_const, add_i, mul_i, div_i, rem_i, lt_i, le_i, max_i, eq_i, select, idx_const, idx_add,
  d_assert, d_print, d_global, global_ptrs, lanes,
  ptr_offset, load_ptrs, store_ptrs, alloca_ptrs,
  float_to_float, extract, cat, permute_tile }
use tileir/prog.{TileProg, trace_kernel, ByPtr, trace3, cells, In, Out}
use tileir/render.{render}

## The tile the two wide kernels cover, and the tile the others do. Eight
## lanes is what lets a whole answer sit in a test as numbers a reader can
## check; the shapes are otherwise the ones `scripts/tile-golden` pins.
const WIDTH: Int = 128
const N: Int = 8

# ---- the device side ----

## Knife 1, the first milestone: one 128-wide tile per block, two loads, one
## add, one store. `!Dev` and nothing else, so there is no host memory here
## and no way to look at a number. The parameters' cells (`vadd_text` below)
## are the addressing, so the body names no shape and no offset.
fn vadd(a: Param[Float], b: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let ta = load_cell(a)
  let tb = load_cell(b)
  store_cell(out, add(ta, tb))
}

## Knife T15: `vadd` with `optimization_hints` on the first load and on the
## store. A hint is advice about latency, so by definition it cannot change
## what the kernel computes; the second load carries none, which puts both
## states of the flag bit in one kernel.
fn vadd_hinted(a: Param[Float], b: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let blk = tile_at(block_id(0), WIDTH)
  let ta = load(a, blk, [WIDTH], hints: [for_arch("sm_86", [hint_latency(4)])])
  let tb = load(b, blk, [WIDTH])
  store(out, blk, add(ta, tb), hints: [for_arch("sm_86", [hint_latency(7)])])
}

# One Collatz step per lane, saturating at 1: a lane that has reached 1 stays
# there, which is what makes "every lane is 1" a condition that stays true.
fn collatz_step(v: Tile[I32], s: Tile[I32]) -> (Tile[I32], Tile[I32]) !Dev = {
  let one = i_const(1)
  let zero = i_const(0)
  let two = i_const(2)
  let running = lt_i(one, v)
  let even = eq_i(rem_i(v, two), zero)
  let moved = select(even, div_i(v, two), add_i(mul_i(v, i_const(3)), one))
  (select(running, moved, v), add_i(s, select(running, one, zero)))
}

## Knife T5: `loop` and `break`, the one region shape whose exit the region
## itself decides. The loop carries the values and their step counts and
## stops when the largest value left is 1, so the trip count is DATA: no
## constant here says how many iterations there are, and the number reaches
## the host because the kernel reduces it and stores it.
fn collatz(x: Param[I32], out: Param[I32]) -> Unit !Dev = {
  let shape = [N]
  let blk = tile_at(idx_const(0), N)
  let v = carry(load(x, blk, shape))
  let steps = carry(broadcast(i_const(0), shape))
  d_loop {
    let biggest = d_reduce(v.get(), 1.0, (acc, e) => max_i(acc, e))
    let (nv, ns) = collatz_step(v.get(), steps.get())
    v.set(nv)
    steps.set(ns)
    le_i(biggest, i_const(1))
  }
  store(out, blk, steps.get())
  store(out, tile_at(idx_const(1), N), broadcast(d_reduce(steps.get(), 0.0, (acc, e) => max_i(acc, e)), shape))
}

## Knife T6: an assertion and a print. A launch that fires an `assert` is
## told to the host as a failed synchronisation, and the corpus keeps every
## lane under the limit so this one does not fire; `print_tko` writes to the
## PROCESS's standard output on a real device. Neither is visible from here,
## the fake device never running this body, and both are in the rendered
## text, which is the layer an example can hold.
fn guarded(x: Param[I32], out: Param[I32]) -> Unit !Dev = {
  let blk = tile_at(block_id(0), N)
  let v = load(x, blk, [N])
  d_assert(lt_i(v, i_const(1000)), "gpu_fake: a lane reached the limit")
  d_print("guarded sum=%d\n", [d_reduce(v, 0.0, (acc, e) => add_i(acc, e))])
  store(out, blk, add_i(v, i_const(1)))
}

## The two tables `table` declares, written out here rather than taken from
## `tileref/ref.global_table_a`, which is what the host reference reads: if the
## two ever part company the answers stop matching, so the duplication is a
## second opinion and not a copy.
fn table_a() -> List[Float] = [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]

fn table_b() -> List[Float] = [-0.25, -1.25, -4.25, -9.25, -16.25, -25.25, -36.25, -49.25]

## Knife T7: static globals, the first data in a module that no host
## uploaded. The tables live in the bytecode's Global section and
## `get_global` answers a pointer to each; `@tbl_b` asks for a 256-byte
## alignment and `@tbl_a` takes the dialect's default, so one module carries
## two different values in that field.
fn table(x: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let s = [N]
  d_global("tbl_a", F64, table_a())
  d_global("tbl_b", F64, table_b(), align: 256)
  let blk = tile_at(block_id(0), N)
  let off = lanes(blk, s)
  let tx = load(x, blk, s)
  let pa = ptr_offset(global_ptrs("tbl_a", F64, s), off)
  let pb = ptr_offset(global_ptrs("tbl_b", F64, s), off)
  store(out, blk, add(load_ptrs(pa), tx))
  store(out, idx_add(blk, idx_const(N)), add(load_ptrs(pb), tx))
}

## Knives T3 and T9: the narrow element formats. Segment 0 reads a buffer OF
## f8E4M3FN and widens it, segment 1 takes an f64 and goes there and back;
## the storage form and the conversion answer the same numbers, both landing
## on the same eight-bit grid. The subword formats T9 added (`i4`,
## `f4E2M1FN`) have no buffer format at all, so they are not here.
fn narrowed(x: Param[F8E4M3FN], y: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let s = [N]
  let blk = tile_at(block_id(0), N)
  store(out, blk, float_to_float(load(x, blk, s), F64))
  store(out, idx_add(blk, idx_const(N)), float_to_float(float_to_float(load(y, blk, s), F8E4M3FN), F64))
}

## Knife T4: the attribute domains, of which a directed rounding is the one
## whose effect a reader can see. The same two binary32 operands are added
## twice, towards each infinity, so the two segments BRACKET the exact sum:
## a lane on the f32 grid gets one number twice, and a lane off it gets the
## two neighbours it falls between.
fn interval(x: Param[Float], y: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let s = [N]
  let blk = tile_at(block_id(0), N)
  let a = float_to_float(load(x, blk, s), F32)
  let b = float_to_float(load(y, blk, s), F32)
  store(out, blk, float_to_float(add(a, b, rounding: narrow.Down), F64))
  store(out, idx_add(blk, idx_const(N)), float_to_float(add(a, b, rounding: narrow.Up), F64))
}

## Knife T10: `alloca`, memory that belongs to the block and to nothing else.
## Every other pointer here comes from a parameter the host uploaded or a
## global the module declared; this is the third source and the only one no
## host can see, so the corpus has to make the memory ANSWER rather than
## show it. The kernel writes `x + x` into the scratch, reads it back in the
## same launch and adds the input, so the answer is `3 * x` and a scratch
## that read back zeros would answer `x`. The count is not a power of two
## and the alignment is: the dialect's rule, and why they cannot be swapped.
fn scratch(x: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let s = [N]
  let blk = tile_at(block_id(0), N)
  let p = ptr_offset(alloca_ptrs(F64, 12, 8, s), lanes(idx_const(0), s))
  let tx = load(x, blk, s)
  store_ptrs(p, add(tx, tx))
  store(out, blk, add(load_ptrs(p), tx))
}

## Knife T2: the three shape operations, on one block of 128 f64 read twice.
## As a 16 by 8 tile, `extract` takes two quarters and `cat` puts them back
## the other way round; as a 4 by 4 by 8 one, `permute` transposes the first
## two axes. Both source dimensions have two slices and the rank-3 tile's
## first two dimensions are equal, so every wrong spelling here is still
## well typed: a wrong ANSWER rather than a refusal.
fn shapes(x: Param[Float], out: Param[Float]) -> Unit !Dev = {
  let b = block_id(0)
  let base = tile_at(b, 128)
  let obase = tile_at(b, 192)
  let t2 = load(x, base, [16, 8])
  let lower = extract(t2, [8, 4], [idx_const(1), idx_const(0)])
  let upper = extract(t2, [8, 4], [idx_const(0), idx_const(0)])
  store(out, obase, cat(lower, upper, 0))
  let t3 = load(x, base, [4, 4, 8])
  store(out, idx_add(obase, idx_const(64)), permute_tile(t3, [1, 0, 2]))
}

# ---- the records ----
#
# `trace3` installs the recording handler, calls the body once with one
# marker a parameter and answers the record; `render` turns it into the text
# a reader (and a golden) can hold. `vadd`'s parameters are cut into cells of
# WIDTH, as many as the grid has; the others are `trace_kernel`'s, whose
# list is the entry signature and whose every parameter is a `Shared`, the
# pointer path's: each writes two regions a block, a pointer of its own or
# through a hint, which no cell describes. A `param` in such a body that
# disagrees with the list is refused here rather than by the assembler.

fn vadd_prog() -> TileProg = {
  let g = cells([DYN_DIM], [WIDTH])
  let (p, _e) = trace3("vadd", In(F64, g), In(F64, g), Out(F64, g), vadd)
  p
}

fn vadd_text() -> String = vadd_prog().render()

fn vadd_hinted_text() -> String =
  trace_kernel("vadd_hinted", [ByPtr("f64"), ByPtr("f64"), ByPtr("f64")],
    () => vadd_hinted(param(F64, 0), param(F64, 1), param(F64, 2))).render()

fn collatz_text() -> String =
  trace_kernel("collatz", [ByPtr("i32"), ByPtr("i32")], () => collatz(param(I32, 0), param(I32, 1)))
  .render()

fn guarded_text() -> String =
  trace_kernel("guarded", [ByPtr("i32"), ByPtr("i32")], () => guarded(param(I32, 0), param(I32, 1)))
  .render()

fn table_text() -> String =
  trace_kernel("table", [ByPtr("f64"), ByPtr("f64")], () => table(param(F64, 0), param(F64, 1)))
  .render()

fn narrowed_text() -> String =
  trace_kernel("narrowed", [ByPtr("f8E4M3FN"), ByPtr("f64"), ByPtr("f64")],
    () => narrowed(param(F8E4M3FN, 0), param(F64, 1), param(F64, 2))).render()

fn interval_text() -> String =
  trace_kernel("interval", [ByPtr("f64"), ByPtr("f64"), ByPtr("f64")],
    () => interval(param(F64, 0), param(F64, 1), param(F64, 2))).render()

fn scratch_text() -> String =
  trace_kernel("scratch", [ByPtr("f64"), ByPtr("f64")], () => scratch(param(F64, 0), param(F64, 1)))
  .render()

fn shapes_text() -> String =
  trace_kernel("shapes", [ByPtr("f64"), ByPtr("f64")], () => shapes(param(F64, 0), param(F64, 1)))
  .render()

# ---- the host side ----

## A buffer the host is about to make: its element format and what goes into
## it. The LAST one a launch is given is the one a reference writes back to,
## so every list below ends with the output buffer's initial contents.
alias Buf = (String, List[Float])

## Three buffers, one launch, one answer back, with `DeviceBuffer[Float]` carrying
## the format in its type. Which device this is, is the handler's business;
## `run` names only `!Gpu`.
fn run(xs: List[Float], ys: List[Float]) -> Result[List[Float], ForeignError] !Gpu = {
  let n = len(xs)
  let a = alloc(F64, n)?
  let b = alloc(F64, n)?
  let out = alloc(F64, n)?
  copy_from_host(a, xs)?
  copy_from_host(b, ys)?
  launch("vadd", blocks(n), [erase(buffer_arg(a)), erase(buffer_arg(b)), erase(buffer_arg(out))])?
  sync()?
  let got = copy_to_host(out)?
  free(a)?
  free(b)?
  free(out)?
  Ok(got)
}

## One block per tile. The fake device does not consult the grid -- a
## reference computes the whole vector at once -- so this is the number the
## real device would need, kept honest by the kernel's width.
fn blocks(n: Int) -> Int = (n + WIDTH - 1) / WIDTH

# Allocate and fill one buffer per pair, in order, over the raw operations
# rather than over `DeviceBuffer[D]`, because one launch here takes buffers of two
# formats and a list cannot hold two types.
fn alloc_upload(bufs: List[Buf]) -> Result[List[Int], ForeignError] !Gpu = {
  var hs: List[Int] = []
  for (dt, data) in bufs {
    let h = gpu.gpu_alloc(dt, len(data))?
    gpu.gpu_copy_from_host(h, gpu.encode_floats(dt, data))?
    hs = hs ++ [h]
  }
  Ok(hs)
}

## One kernel over `bufs` on a `grid` by 1 by 1 grid: allocate, upload,
## launch, wait, read the LAST buffer back, release. A refusal is the answer,
## and the handles issued before it are left to the handler, which here is
## pure and takes them with it.
fn run_kernel(name: String, grid: Int, bufs: List[Buf]) -> Result[List[Float], ForeignError] !Gpu = {
  let hs = alloc_upload(bufs)?
  launch(name, grid, list.map(hs, h => gpu.Buffer(h)))?
  sync()?
  let (last_dt, _last) = bufs[len(bufs) - 1]
  let got = gpu.decode_floats(last_dt, gpu.gpu_copy_to_host(hs[len(hs) - 1])?)
  for h in hs {
    let _ = gpu.gpu_free(h)
  }
  Ok(got)
}

## The `add(.., rounding: Down)` / `rounding: Up` pair on the host. The sum of two binary32
## values is EXACT in binary64, so this rounds the exact sum in the named
## direction and `std/narrow`'s rounder is the whole of the arithmetic. It
## is the one reference `std/gpu` does not already carry.
fn interval_ref(_dtypes: List[String], ins: List[List[Float]]) -> List[Float] = {
  let x = ins[0]
  let y = ins[1]
  list.map(range(0, len(ins[2])), k => {
    let i = k % N
    let exact = narrow.round_f32(x[i]) + narrow.round_f32(y[i])
    narrow.round_f32(exact, mode: if k < N { narrow.Down } else { narrow.Up })
  })
}

## How many rounds `loop_count_ref` may take before it gives up. The corpus
## needs sixteen, so this never stops it; it is there so that a reference can
## never be the thing that hangs.
const REF_FUEL: Int = 4096

## The kernels the fake device knows: the three `std/gpu` registers for
## every caller, plus one per kernel above. An entry is the number of
## buffers the launch takes and the reference that answers it; every
## reference but `interval_ref` is `std/gpu`'s, with the lengths these
## kernels were written for fixed here.
fn devices() -> Map[String, (Int, WideRefFn)] =
  list.fold([("vadd_hinted", (3, last_out(vadd_ref))),
    ("collatz", (2, last_out((d, i) => loop_count_ref(N, REF_FUEL, 2, d, i)))),
    ("guarded", (2, last_out((d, i) => assert_pass_ref(N, d, i)))),
    ("table", (2, last_out((d, i) => global_table_ref(N, d, i)))),
    ("narrowed", (3, last_out((d, i) => dtype_convert_ref(N, "f8E4M3FN", d, i)))),
    ("interval", (3, last_out(interval_ref))),
    ("scratch", (2, last_out((d, i) => alloca_scratch_ref(N, d, i)))),
    ("shapes", (2, last_out((d, i) => shape_ops_ref(1, d, i))))],
    reference_kernels(), (m, e) => {
      let (name, entry) = e
      map.insert(m, name, entry)
    })

## `run_kernel` on the fake device: pure, so a test may call it and so may
## `main`.
fn on_fake(name: String, grid: Int, bufs: List[Buf]) -> Result[List[Float], ForeignError] =
  with_gpu_fake(devices(), () => run_kernel(name, grid, bufs))

# ---- the corpora, and what each kernel answers ----

fn zeros(n: Int) -> List[Float] = list.map(range(0, n), _i => 0.0)

fn left() -> List[Float] = [1.0, 2.0, 3.0, 4.0]

fn right() -> List[Float] = [10.0, 20.0, 30.0, 40.0]

fn vadd_hinted_answer() -> Result[List[Float], ForeignError] =
  on_fake("vadd_hinted", 1, [("f64", left()), ("f64", right()), ("f64", zeros(4))])

## The chains of 1 to 8 are 0, 1, 7, 2, 5, 8, 16 and 3 steps long, which is a
## fact about the numbers and not about this program.
fn collatz_answer() -> Result[List[Float], ForeignError] =
  on_fake("collatz", 1, [("i32", [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]),
    ("i32", zeros(2 * N))])

fn guarded_answer() -> Result[List[Float], ForeignError] =
  on_fake("guarded", 1, [("i32", [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0]),
    ("i32", zeros(N))])

fn table_answer() -> Result[List[Float], ForeignError] =
  on_fake("table", 1, [("f64", list.map(range(0, N), _i => 10.0)), ("f64", zeros(2 * N))])

## Values `f8E4M3FN` does not hold exactly, next to values it does: the grid
## is 1 + k/8 scaled by a power of two, so 1.1 lands on 1.125 and 0.3 on
## 0.3125, while 240 and 448 (the format's largest finite value) are exact.
fn narrow_corpus() -> List[Float] = [1.1, 0.3, 2.0, 6.0, -1.1, 240.0, 448.0, 0.0]

fn narrowed_answer() -> Result[List[Float], ForeignError] =
  on_fake("narrowed", 1, [("f8E4M3FN", narrow_corpus()), ("f64", narrow_corpus()),
    ("f64", zeros(2 * N))])

## One binary32 ulp at 1.0. The first three lanes below add less than this
## to 1.0, so their sums fall strictly between two f32 values and the two
## directions have to differ there; the rest are on the grid.
const F32_ULP_AT_ONE: Float = 1.1920928955078125e-7

fn interval_answer() -> Result[List[Float], ForeignError] =
  on_fake("interval", 1, [("f64", list.map(range(0, N), _i => 1.0)),
    ("f64", [9.31322574615478515625e-10, 2.98023223876953125e-8, 5.9604644775390625e-8,
      1.1920928955078125e-7, 1.0, 0.5, 0.25, 0.0]),
    ("f64", zeros(2 * N))])

fn scratch_answer() -> Result[List[Float], ForeignError] =
  on_fake("scratch", 1, [("f64", [1.0, 2.0, 3.0, 4.0, 0.5, -1.0, 8.0, 0.0]),
    ("f64", zeros(N))])

## The input is its own index, so a permutation's answer names where each
## lane came from.
fn shapes_answer() -> Result[List[Float], ForeignError] =
  on_fake("shapes", 1, [("f64", list.map(range(0, 128), i => to_float(i))),
    ("f64", zeros(192))])

# ---- the gallery ----

fn head(xs: List[Float], n: Int) -> List[Float] =
  list.map(range(0, if n < len(xs) { n } else { len(xs) }), i => xs[i])

fn shown(r: Result[List[Float], ForeignError]) -> String = match r {
  Ok(xs) -> "${len(xs)} lanes, first four ${head(xs, 4)}"
  Err(e) -> "refused: ${e.kind}"
}

## Knife TG: the two globals `table` declares, as the HOST sees them. This
## is the one operation of `Gpu` that names something inside a module rather
## than a buffer the program allocated: `module_global` answers a buffer
## over the global's own storage, so the host can read the table the module
## was born with, write it, and have the next launch read the write. On a
## real device it is the CUDA driver's `cuModuleGetGlobal`; here it is
## `with_gpu_fake` with `globals:`, which is what the module declares.
fn module_globals() -> Map[String, List[(String, String, List[Float])]] =
  map.from([("table", [("tbl_a", "f64", table_a()), ("tbl_b", "f64", table_b())])])

## What the host can name in `table`'s module: the first table read straight
## out of it, and the refusal for a name it does not declare.
fn symbols() -> Result[(List[Float], String), ForeignError] !Gpu = {
  let a = gpu.module_global(F64, "table", "tbl_a")?
  let held = copy_to_host(a)?
  let missing = match gpu.module_global(F64, "table", "no_such_table") {
    Ok(_) -> "found"
    Err(e) -> e.kind
  }
  Ok((held, missing))
}

fn symbols_answer() -> Result[(List[Float], String), ForeignError] =
  with_gpu_fake(devices(), symbols, globals: module_globals())

fn symbols_line() -> String = match symbols_answer() {
  Ok(p) -> {
    let (held, missing) = p
    "@tbl_a holds ${head(held, 4)}, an undeclared name says ${missing}"
  }
  Err(e) -> "refused: ${e.kind}"
}

## One line per kernel, in knife order: what it is for, and what the fake
## device answered.
fn gallery() -> List[(String, String)] = [
  ("T2  shapes    extract, cat, permute", shapes_answer().shown()),
  ("T3  narrowed  f8E4M3FN storage and conversion", narrowed_answer().shown()),
  ("T4  interval  addf towards each infinity", interval_answer().shown()),
  ("T5  collatz   a loop the data stops", collatz_answer().shown()),
  ("T6  guarded   assert and print_tko", guarded_answer().shown()),
  ("T7  table     two globals nobody uploaded", table_answer().shown()),
  ("T10 scratch   alloca, written and read back", scratch_answer().shown()),
  ("T15 hinted    vadd with optimization_hints", vadd_hinted_answer().shown()),
  ("TG  symbols   a module global the host looks up", symbols_line())]

fn vadd_line() -> String = match with_gpu_fake(reference_kernels(), () => run(left(), right())) {
  Ok(zs) -> "${zs}"
  Err(e) -> "refused: ${e.kind}"
}

pub fn main() -> Unit !io = {
  io.print(vadd_text())
  io.println("vadd on the fake device: ${vadd_line()}")
  for row in gallery() {
    let (label, line) = row
    io.println("${label}: ${line}")
  }
}

# ---- the records are a function of the bodies ----

test "the record is a function of the body, not of the run" {
  assert vadd_prog() == vadd_prog()
}

test "the kernel renders as one Tile IR entry" {
  let text = vadd_text()
  assert str.contains(text, "cuda_tile.module @m")
  assert str.contains(text, "entry @vadd")
  assert str.contains(text, "tile<128xf64>")
}

test "T15: the hints are on the first load and the store, and nowhere else" {
  let text = vadd_hinted_text()
  assert str.contains(text, "load_ptr_tko weak %12 token=%0 optimization_hints=<sm_86 = {latency = 4}>")
  assert str.contains(text, "store_ptr_tko weak %23, %20 token=%19 optimization_hints=<sm_86 = {latency = 7}>")
  assert str.contains(text, "load_ptr_tko weak %17 token=%14 : tile<128xptr<f64>>")
}

test "T5: the loop carries two tiles and leaves through a break" {
  let text = collatz_text()
  assert str.contains(text, "loop iter_values(%17 = %11, %18 = %13, %19 = %12)")
  assert str.contains(text, "break %17, %18, %19 : tile<8xi32>, tile<8xi32>, token")
}

test "T6: the assertion and the print spell the dialect's own assembly" {
  let text = guarded_text()
  assert str.contains(text, "assert %16, \"gpu_fake: a lane reached the limit\" : tile<8xi1>")
  assert str.contains(text, "print_tko \"guarded sum=%d\\n\", %17 token=%14 : tile<i32> -> token")
}

test "T7: the globals print above the entry, the alignment only on the second" {
  let text = table_text()
  assert str.contains(text, "global @tbl_a <f64: [1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]> : tile<8xf64>")
  assert str.contains(text, "global @tbl_b alignment = 256 <f64: [-0.25,")
  assert str.contains(text, "get_global @tbl_a : tile<ptr<f64>>")
}

test "T3: the conversion names both formats and the rounding it carries" {
  let text = narrowed_text()
  assert str.contains(text, "ftof %13 rounding<nearest_even> : tile<8xf8E4M3FN> -> tile<8xf64>")
  assert str.contains(text, "ftof %25 rounding<nearest_even> : tile<8xf64> -> tile<8xf8E4M3FN>")
}

test "T4: one pair of operands, two rounding directions" {
  let text = interval_text()
  assert str.contains(text, "addf %15, %21 rounding<negative_inf> : tile<8xf32>")
  assert str.contains(text, "addf %15, %21 rounding<positive_inf> : tile<8xf32>")
}

test "T10: the allocation carries its element count and its alignment" {
  assert str.contains(scratch_text(), "alloca num_elem = 12, alignment = 8 : tile<ptr<f64>>")
}

test "T2: the three shape operations spell their shapes" {
  let text = shapes_text()
  assert str.contains(text, "extract %23[%25, %26] : tile<16x8xf64> -> tile<8x4xf64>")
  assert str.contains(text, "cat %27, %30 dim = 0 : tile<8x4xf64>, tile<8x4xf64> -> tile<16x4xf64>")
  assert str.contains(text, "permute %69 [1, 0, 2] : tile<4x4x8xf64> -> tile<4x4x8xf64>")
}

# ---- what the fake device answers ----

test "the same program on the fake device adds the vectors" {
  assert with_gpu_fake(reference_kernels(), () => run(left(), right())) ==
    Ok([11.0, 22.0, 33.0, 44.0])
}

test "T15: a hint does not change the answer" {
  assert vadd_hinted_answer() == Ok([11.0, 22.0, 33.0, 44.0])
}

test "T5: the step counts are the Collatz chains, and the round count is the longest" {
  assert collatz_answer() ==
    Ok([0.0, 1.0, 7.0, 2.0, 5.0, 8.0, 16.0, 3.0,
      16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0, 16.0])
}

test "T6: an assertion that holds leaves an ordinary launch" {
  assert guarded_answer() == Ok([11.0, 21.0, 31.0, 41.0, 51.0, 61.0, 71.0, 81.0])
}

test "TG: the host reads a module global by name, and is refused a name the module has not got" {
  # `@tbl_a` is the table the module was BORN with: nothing uploaded it, and
  # what comes back is what `table_a` declares. The second half is the whole
  # of what a symbol lookup can refuse.
  assert symbols_answer() == Ok((table_a(), "gpu.no_module_symbol"))
}

test "T7: the answer is the two tables, plus ten" {
  assert table_answer() ==
    Ok([11.0, 11.5, 12.0, 12.5, 13.0, 13.5, 14.0, 14.5,
      9.75, 8.75, 5.75, 0.75, -6.25, -15.25, -26.25, -39.25])
}

test "T3: both segments land on the eight-bit grid, and on the same points" {
  let rounded = [1.125, 0.3125, 2.0, 6.0, -1.125, 240.0, 448.0, 0.0]
  assert narrowed_answer() == Ok(rounded ++ rounded)
}

test "T4: the two directions bracket the exact sum, and differ only off the grid" {
  let down = [1.0, 1.0, 1.0, 1.0 + F32_ULP_AT_ONE, 2.0, 1.5, 1.25, 1.0]
  let up = [1.0 + F32_ULP_AT_ONE, 1.0 + F32_ULP_AT_ONE, 1.0 + F32_ULP_AT_ONE,
    1.0 + F32_ULP_AT_ONE, 2.0, 1.5, 1.25, 1.0]
  assert interval_answer() == Ok(down ++ up)
  for i in range(0, N) {
    assert down[i] <= up[i]
  }
}

test "T10: the scratch answers what was put in it, so the answer is three times the input" {
  assert scratch_answer() == Ok([3.0, 6.0, 9.0, 12.0, 1.5, -3.0, 24.0, 0.0])
}

## The probes are hand-computed from what the operations MEAN, the input
## being its own index: `cat(extract[1, 0], extract[0, 0])` makes row `i` of
## the first 64 lanes row `(i + 8) % 16` of the source cut to four columns,
## and `permute [1, 0, 2]` makes lane `[a, c, e]` of the next 128 lane
## `[c, a, e]`. The sum is the bulk check: a permutation moves numbers
## without inventing them, so 0 through 127 still sum to 8128.
test "T2: the shapes are permuted, not invented" {
  match shapes_answer() {
    Err(_) -> { assert false }
    Ok(got) -> {
      assert len(got) == 192
      assert got[0] == 64.0
      assert got[4] == 72.0
      assert got[63] == 59.0
      assert got[72] == 32.0
      assert list.fold(range(64, 192), 0.0, (acc, i) => acc + got[i]) == 8128.0
    }
  }
}

# ---- what the device refuses ----

test "a kernel the device does not have is refused, and nothing is computed" {
  match with_gpu_fake(map.empty(), () => run(left(), right())) {
    Ok(_) -> { assert false }
    Err(e) -> { assert e.kind == "gpu.no_kernel" }
  }
}

test "a launch with the wrong number of buffers is refused before anything runs" {
  match on_fake("collatz", 1, [("i32", zeros(N)), ("i32", zeros(N)), ("i32", zeros(2 * N))]) {
    Ok(_) -> { assert false }
    Err(e) -> { assert e.kind == "gpu.bad_arity" }
  }
}

test "a format the device cannot hold is refused at the allocation" {
  match on_fake("scratch", 1, [("i4", zeros(N)), ("f64", zeros(N))]) {
    Ok(_) -> { assert false }
    Err(e) -> { assert e.kind == "gpu.unsupported_dtype" }
  }
}

test "one block covers a tile, and a vector longer than a tile needs more" {
  assert blocks(1) == 1
  assert blocks(128) == 1
  assert blocks(129) == 2
}