Examples

Programs the compiler already accepts

These .kina files live in tests/compiler/examples/. They are compile tests, not a standard library. Where a golden exists, the example page also shows the textual MLIR the OCaml compiler emits. storylm0_forward.kina is a tiny causal-attention forward graph.storylm0_nll.kina composes host log_softmax and take_along_axis; it is not training. The Python StoryLM lab under experiments/storylm/ is a separate Hugging Face evaluation harness.

ad_basic.kina

Restricted reverse-mode AD example

let f = fn(x: tensor<F32, [4, 4], @host>) -> tensor<F32, [4, 4], @host> {
  let W = tensor<F32, [4, 4], @host> { 1.0 } in
  let b = tensor<F32, [4, 4], @host> { 1.0 } in
  let mm = W @ x in
  let linear = mm + b in
  let act = relu(linear) in
  act
} in
let x_input = tensor<F32, [4, 4], @host> { 1.0 } in
handle f(x_input) with reverseMode

11 lines · MLIR golden

ad_gather.kina

Restricted reverse-mode 2D gather VJP; I64 ids have no cotangent

let f = fn(
  emb: tensor<F32, [8, 4], @host>,
  ids: tensor<I64, [1, 4], @host>
) -> tensor<F32, [1, 4, 4], @host> {
  let y = emb.gather(ids) in
  y
…

18 lines · MLIR golden

ad_log_softmax.kina

Restricted reverse-mode log_softmax VJP

let f = fn(x: tensor<F32, [2, 3], @host>) -> tensor<F32, [2, 3], @host> {
  let y = x.log_softmax(axis=1) in
  y
} in
let x = tensor<F32, [2, 3], @host> in
handle f(x) with reverseMode

7 lines · MLIR golden

ad_matmul4.kina

Restricted reverse-mode rank-4 matmul VJP

let f = fn(
  q: tensor<F32, [1, 2, 4, 2], @host>,
  kt: tensor<F32, [1, 2, 2, 4], @host>
) -> tensor<F32, [1, 2, 4, 4], @host> {
  let s = q @ kt in
  s
…

18 lines · MLIR golden

ad_mlp.kina

Restricted reverse-mode AD example

let mlp_forward = fn(x: tensor<F32, [2, 2], @host>, w1: tensor<F32, [2, 2], @host>, w2: tensor<F32, [2, 2], @host>) !{Diff} -> tensor<F32, [2, 2], @host> {
  let active_w1 = perform Diff.active(w1) in
  let active_w2 = perform Diff.active(w2) in
  let hidden = x @ active_w1 in
  let hidden_sq = hidden * hidden in
  let out = hidden_sq @ active_w2 in
…

20 lines

ad_neg.kina

Restricted reverse-mode neg VJP; @main returns (primal, g)

let f = fn(x: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  let y = -x in
  y
} in
let x = tensor<F32, [2, 2], @host> in
handle f(x) with reverseMode

7 lines · MLIR golden

ad_neg_grad.kina

reverseMode wrapper fn returning (primal, grad) from a .so

let f = fn(x: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  let y = -x in
  y
} in
let ad = fn(x: tensor<F32, [2, 2], @host>)
  -> (tensor<F32, [2, 2], @host>, tensor<F32, [2, 2], @host>) {
  handle f(x) with reverseMode
} in
let x = tensor<F32, [2, 2], @host> in
ad(x)

11 lines · MLIR golden

ad_nll.kina

Restricted reverse-mode NLL VJP; @main returns (loss, g)

let nll = fn(logits: tensor<F32, [1, 4, 8], @host>) -> tensor<F32, [1, 4], @host> {
  let targets = tensor<I64, [1, 4], @host> in
  let log_probs = logits.log_softmax(axis=2) in
  let picked = log_probs.take_along_axis(targets, axis=2) in
  let loss = -picked in
  loss
} in
let logits = tensor<F32, [1, 4, 8], @host> in
handle nll(logits) with reverseMode

10 lines · MLIR golden

ad_reduce_sum.kina

Restricted reverse-mode keepdims reduce_sum VJP

let f = fn(x: tensor<F32, [2, 3], @host>) -> tensor<F32, [2, 1], @host> {
  let y = x.reduce_keep(op=sum, axis=1) in
  y
} in
let step = fn(x: tensor<F32, [2, 3], @host>) -> (
  tensor<F32, [2, 1], @host>,
…

14 lines · MLIR golden

ad_reshape.kina

Restricted reverse-mode reshape VJP

let f = fn(x: tensor<F32, [1, 4, 4], @host>) -> tensor<F32, [1, 4, 2, 2], @host> {
  let y = x.reshape([1, 4, 2, 2]) in
  y
} in
let grad = fn(x: tensor<F32, [1, 4, 4], @host>) -> tensor<F32, [1, 4, 4], @host> {
  let (y, g) = handle f(x) with reverseMode in
  g
} in
let x = tensor<F32, [1, 4, 4], @host> in
grad(x)

11 lines · MLIR golden

ad_sgd_step.kina

Name reverseMode grads with let (y, g) and apply SGD

let f = fn(w: tensor<F32, [4, 8], @host>) -> tensor<F32, [4, 8], @host> {
  let y = -w in
  y
} in
let step = fn(
  w: tensor<F32, [4, 8], @host>,
…

17 lines · MLIR golden

ad_stdlib_test.kina

Restricted reverse-mode AD example

let x = tensor<F32, [1, 32, 64], @host> { 1.0 } in
let gamma = tensor<F32, [1, 1, 64], @host> { 1.0 } in

// Test RMSNorm Gradients
let out1 = handle rmsnorm(x, gamma) with reverseMode in

…

19 lines

adamw_step.kina

Functional AdamW as existing arithmetic; not a primitive

let adamw_step = fn(
  w: tensor<F32, [4, 8], @host>,
  g: tensor<F32, [4, 8], @host>,
  m: tensor<F32, [4, 8], @host>,
  v: tensor<F32, [4, 8], @host>,
  lr: tensor<F32, [1, 1], @host>,
…

63 lines · MLIR golden

add_2d.kina

Compiler example

(tensor<F32, [2, 3]> + tensor<F32, [2, 3]>)

2 lines · MLIR golden

argmax_axis1.kina

Host reduce argmax; I64 result; lowest-index tie-break

tensor<F32, [2, 3]>.reduce(op=argmax, axis=1)

2 lines · MLIR golden

argmax_keep_axis1.kina

Host reduce argmax; I64 result; lowest-index tie-break

tensor<F32, [2, 3]>.reduce_keep(op=argmax, axis=1)

2 lines · MLIR golden

bitwise_and_i32.kina

Compiler example

tensor<I32, [2, 3]> & tensor<I32, [2, 3]>

2 lines · MLIR golden

broadcast_row.kina

Compiler example

tensor<F32, [1, 3]>.broadcast([2, 3])

2 lines · MLIR golden

cast_bool_i32.kina

Dtype cast

tensor<Bool, [2]>.cast(I32)

2 lines · MLIR golden

cast_f16_bf16.kina

Dtype cast

tensor<F16, [2, 2]>.cast(BF16)

2 lines · MLIR golden

cast_i32.kina

Dtype cast

tensor<F32, [2, 4]>.cast(I32)

2 lines · MLIR golden

cast_i8_i32.kina

Dtype cast

tensor<I8, [3]>.cast(I32)

2 lines · MLIR golden

concat_axis0.kina

Compiler example

tensor<F32, [2, 3]>.concat(axis=0, tensor<F32, [3, 3]>)

2 lines · MLIR golden

div_2d.kina

Compiler example

(tensor<F32, [2, 3]> / tensor<F32, [2, 3]>)

2 lines · MLIR golden

effects_test.kina

Early control-flow / effect example

let run = fn(val : tensor<F32, [], @host>) !{State} -> tensor<F32, [], @host> {
  let current = perform State.get() in
  let next = current + val in
  next
} in
handle run(tensor<F32, [], @host>) with {
  return(x) => x,
  State.get(), k => k(tensor<F32, [], @host>)
}

10 lines

eq_f32.kina

Compiler example

(tensor<F32, [2, 2]> == tensor<F32, [2, 2]>)

2 lines · MLIR golden

frem_f32.kina

Compiler example

tensor<F32, [2, 3]> frem tensor<F32, [2, 3]>

2 lines · MLIR golden

frem_f8e4m3.kina

Compiler example

tensor<F8E4M3, [2, 3]> frem tensor<F8E4M3, [2, 3]>

2 lines · MLIR golden

gather_2d.kina

Compiler example

tensor<F32, [10, 20], @host>.gather(tensor<I64, [4, 5], @host>)

2 lines

if_else.kina

Early control-flow / effect example

let relu_custom = fn(x: tensor<F32, [4, 4], @host>, scalar_cond: tensor<Bool, [], @host>) -> tensor<F32, [4, 4], @host> {
  let zero = tensor<F32, [4, 4], @host> in
  let final = if (scalar_cond) { x } else { zero } in
  final
} in
let x_input = tensor<F32, [4, 4], @host> in
let cond = tensor<Bool, [], @host> in
relu_custom(x_input, cond)

9 lines · MLIR golden

invalid_add_shape_mismatch.kina

Rejected by the type checker

(tensor<F32, [2, 3]> + tensor<F32, [2, 4]>)

2 lines

invalid_argmax_bool.kina

Rejected by the type checker

tensor<Bool, [2, 3]>.reduce(op=argmax, axis=1)

2 lines

invalid_argmax_device.kina

Rejected by the type checker

tensor<F32, [2, 3], @device(0)>.reduce(op=argmax, axis=1)

2 lines

invalid_dim_zero.kina

Rejected by the type checker

tensor<F32, [0, 4]>

2 lines

invalid_log_softmax_axis.kina

Rejected by the type checker

tensor<F32, [2, 3], @host>.log_softmax(axis=2)

2 lines

invalid_log_softmax_dtype.kina

Rejected by the type checker

tensor<I32, [2, 3], @host>.log_softmax(axis=1)

2 lines

invalid_mul_shape_mismatch.kina

Rejected by the type checker

(tensor<F32, [2, 3]> * tensor<F32, [2, 4]>)

2 lines

invalid_reduce_mean_ad.kina

Rejected by the type checker

let f = fn(x: tensor<F32, [2, 3], @host>) -> tensor<F32, [2, 1], @host> {
  let y = x.reduce_keep(op=mean, axis=1) in
  y
} in
let x = tensor<F32, [2, 3], @host> in
handle f(x) with reverseMode

7 lines

invalid_requires_mismatch.kina

Rejected by the type checker

requires [B = B_prime] compute {
  tensor<F32, [B, M, K], @host> @ tensor<F32, [B_prime, K_prime, N], @host>
}

4 lines

invalid_reshape_volume_mismatch.kina

Rejected by the type checker

tensor<F32, [2, 3]>.reshape([4, 2])

2 lines

invalid_reverse_mode_return.kina

Rejected by the type checker

let f = fn(x: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  let y = -x in
  y
} in
let ad = fn(x: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  handle f(x) with reverseMode
} in
let x = tensor<F32, [2, 2], @host> in
ad(x)

10 lines

invalid_sgd_shape.kina

Rejected by the type checker

let w = tensor<F32, [4, 8], @host> in
let g = tensor<F32, [4, 8], @host> in
let lr = tensor<F32, [1, 1], @host> in
w - lr * g

5 lines

invalid_take_along_axis_axis.kina

Rejected by the type checker

tensor<F32, [1, 4, 8], @host>.take_along_axis(tensor<I64, [1, 4], @host>, axis=3)

2 lines

invalid_take_along_axis_device.kina

Rejected by the type checker

tensor<F32, [2, 3], @device(0)>.take_along_axis(tensor<I64, [2], @device(0)>, axis=1)

2 lines

invalid_take_along_axis_indices.kina

Rejected by the type checker

tensor<F32, [1, 4, 8], @host>.take_along_axis(tensor<I32, [1, 4], @host>, axis=2)

2 lines

invalid_take_along_axis_rank.kina

Rejected by the type checker

tensor<F32, [4], @host>.take_along_axis(tensor<I64, [4], @host>, axis=0)

2 lines

invalid_take_along_axis_shape.kina

Rejected by the type checker

tensor<F32, [1, 4, 8], @host>.take_along_axis(tensor<I64, [1, 3], @host>, axis=2)

2 lines

invalid_tuple_let_arity.kina

Rejected by the type checker

let f = fn(w: tensor<F32, [4, 8], @host>) -> tensor<F32, [4, 8], @host> {
  let y = -w in
  y
} in
let w = tensor<F32, [4, 8], @host> in
let (y, g, extra) = handle f(w) with reverseMode in
y

8 lines

invalid_tuple_let_duplicate.kina

Rejected by the type checker

let f = fn(w: tensor<F32, [4, 8], @host>) -> tensor<F32, [4, 8], @host> {
  let y = -w in
  y
} in
let w = tensor<F32, [4, 8], @host> in
let (g, g) = handle f(w) with reverseMode in
g

8 lines

invalid_tuple_let_not_tuple.kina

Rejected by the type checker

let w = tensor<F32, [4, 8], @host> in
let (y, g) = w in
y

4 lines

llm_generate.kina

Loop-carried matrix sketch, not token generation

let generate = fn(steps: tensor<I32, [], @host>, start_token: tensor<F32, [4, 4], @host>, init_kv: tensor<F32, [16, 4], @host>, start_pos: tensor<I32, [], @host>, step_size: tensor<I32, [], @host>) -> (tensor<F32, [4, 4], @host>, tensor<F32, [16, 4], @host>, tensor<I32, [], @host>) {
  let final = while (pos < steps) [tok = start_token, kv = init_kv, pos = start_pos] {
    let w = tensor<F32, [4, 4], @host> in
    let next_tok = tok @ w in
    let next_kv = kv + kv in
    let next_pos = pos + step_size in
…

17 lines

log_softmax_axis0.kina

Stable log-softmax on a compile-time axis

tensor<F32, [2, 3], @host>.log_softmax(axis=0)

2 lines · MLIR golden

log_softmax_axis1.kina

Stable log-softmax on a compile-time axis

tensor<F32, [2, 3], @host>.log_softmax(axis=1)

2 lines · MLIR golden

log_softmax_rank3.kina

Stable log-softmax on a compile-time axis

tensor<F32, [1, 4, 8], @host>.log_softmax(axis=2)

2 lines · MLIR golden

matmul_2d.kina

Matrix multiply

(tensor<F32, [2, 3]> @ tensor<F32, [3, 4]>)

2 lines · MLIR golden

matmul_batch.kina

Matrix multiply

(tensor<F32, [1, 2, 3]> @ tensor<F32, [1, 3, 4]>)

2 lines · MLIR golden

max_f8e4m3.kina

Compiler example

max(tensor<F8E4M3, [2, 2]>, tensor<F8E4M3, [2, 2]>)

2 lines · MLIR golden

max_i32.kina

Compiler example

max(tensor<I32, [2, 3]>, tensor<I32, [2, 3]>)

2 lines · MLIR golden

min_f32.kina

Compiler example

min(tensor<F32, [1, 4]>, tensor<F32, [1, 4]>)

2 lines · MLIR golden

mlp_forward.kina

MLP forward compile example

requires [B = B_prime, In = In_prime, Hidden = Hidden_out] compute {
    let x = tensor<F32, [B, In], @host> in
    let w1 = tensor<F32, [In_prime, Hidden], @host> in
    let w2 = tensor<F32, [Hidden_out, Out], @host> in
    
    // First linear layer
…

17 lines

mlp_forward_gpu.kina

MLP forward compile example

let mlp_gpu = fn(x: tensor<F32, [128, 256], @device(0)>, w1: tensor<F32, [256, 512], @device(0)>, w2: tensor<F32, [512, 10], @device(0)>) -> tensor<F32, [128, 10], @device(0)> {
    let h1 = x @ w1 in
    let h1_relu = relu(h1) in
    let out = h1_relu @ w2 in
    out
} in
let x = tensor<F32, [128, 256], @device(0)> { 1.0 } in
let w1 = tensor<F32, [256, 512], @device(0)> { 1.0 } in
let w2 = tensor<F32, [512, 10], @device(0)> { 1.0 } in
mlp_gpu(x, w1, w2)

11 lines

mlp_forward_static.kina

MLP forward compile example

let mlp = fn(x: tensor<F32, [128, 256], @host>, w1: tensor<F32, [256, 512], @host>, w2: tensor<F32, [512, 10], @host>) -> tensor<F32, [128, 10], @host> {
    let h1 = x @ w1 in
    let h1_relu = relu(h1) in
    let out = h1_relu @ w2 in
    out
} in
let x = tensor<F32, [128, 256], @host> { 1.0 } in
let w1 = tensor<F32, [256, 512], @host> { 1.0 } in
let w2 = tensor<F32, [512, 10], @host> { 1.0 } in
mlp(x, w1, w2)

11 lines

mul_2d.kina

Compiler example

(tensor<F32, [2, 3]> * tensor<F32, [2, 3]>)

2 lines · MLIR golden

pow_f32.kina

Compiler example

(tensor<F32, [2, 2]> ** tensor<F32, [2, 2]>)

2 lines · MLIR golden

pow_f8e4m3.kina

Compiler example

(tensor<F8E4M3, [2, 2]> ** tensor<F8E4M3, [2, 2]>)

2 lines · MLIR golden

reduce_keep_mean_axis1.kina

Reduction

tensor<F32, [2, 3]>.reduce_keep(op=mean, axis=1)

2 lines · MLIR golden

reduce_max_axis1.kina

Reduction

tensor<F32, [2, 3]>.reduce(op=max, axis=1)

2 lines · MLIR golden

reduce_max_f8_axis0.kina

Reduction

tensor<F8E4M3, [2, 3]>.reduce(op=max, axis=0)

2 lines · MLIR golden

reduce_mean_f8_axis0.kina

Reduction

tensor<F8E4M3, [2, 3]>.reduce(op=mean, axis=0)

2 lines · MLIR golden

reduce_min_axis0.kina

Reduction

tensor<I32, [2, 3]>.reduce(op=min, axis=0)

2 lines · MLIR golden

reduce_sum_axis1.kina

Reduction

tensor<F32, [2, 3]>.reduce(op=sum, axis=1)

2 lines · MLIR golden

rem_i32.kina

Compiler example

(tensor<I32, [2, 4]> % tensor<I32, [2, 4]>)

2 lines · MLIR golden

repeat_interleave_axis0.kina

Compiler example

tensor<F32, [2, 3], @host>.repeat_interleave(axis=0, reps=2)

2 lines

reshape_2d_to_3d.kina

Compiler example

tensor<F32, [2, 3]>.reshape([1, 2, 3])

2 lines · MLIR golden

scalar_bool.kina

Compiler example

tensor<Bool, []>

2 lines · MLIR golden

sgd_step.kina

Functional SGD as broadcast(lr)*g subtracted from w

let sgd_step = fn(
  w: tensor<F32, [4, 8], @host>,
  g: tensor<F32, [4, 8], @host>,
  lr: tensor<F32, [1, 1], @host>
) -> tensor<F32, [4, 8], @host> {
  let lr_b = lr.broadcast([4, 8]) in
…

15 lines · MLIR golden

slice_axis0.kina

Compiler example

tensor<F32, [4, 5]>.slice(axis=0, start=0, stop=2)

2 lines · MLIR golden

softmax_axis1.kina

Compiler example

tensor<F32, [2, 3], @host>.softmax(axis=1)

2 lines

state_effect.kina

Early control-flow / effect example

let f = fn(x : tensor<F32, []>) !{State} -> tensor<F32, []> {
    let current = perform State.get() in
    let new_val = current + x in
    let dummy = perform State.set(new_val) in
    new_val
} in
…

14 lines

static_2d.kina

Compiler example

tensor<F32, [2, 2]>

2 lines · MLIR golden

storylm0_ad_step.kina

Decoder graph + embedding SGD; not a trained model

let storylm0_fwd = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
…

72 lines · MLIR golden

storylm0_adamw_step.kina

NLL ∘ decoder + multi-weight AdamW; not a trained model

let adamw_e = fn(
  w: tensor<F32, [8, 4], @host>,
  g: tensor<F32, [8, 4], @host>,
  m: tensor<F32, [8, 4], @host>,
  v: tensor<F32, [8, 4], @host>,
  lr: tensor<F32, [1, 1], @host>,
…

315 lines · MLIR golden

storylm0_forward.kina

Tiny causal-attention forward graph (not a trained model)

let storylm0_forward = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
…

56 lines · MLIR golden

storylm0_generate.kina

Host greedy loop over static-S decode; not while-grown cache

let decode_self = fn(
  token_ids: tensor<I64, [1, 1], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
  wk: tensor<F32, [4, 4], @host>,
…

129 lines · MLIR golden

storylm0_greedy.kina

Host greedy argmax on [B, T, V] logits; not sampling

let greedy = fn(logits: tensor<F32, [1, 4, 8], @host>) -> tensor<I64, [1, 4], @host> {
  let tokens = logits.reduce(op=argmax, axis=2) in
  tokens
} in
let logits = tensor<F32, [1, 4, 8], @host> in
greedy(logits)

7 lines · MLIR golden

storylm0_kv_step.kina

One-step KV concat + attention; not a generation loop

let kv_step = fn(
  q_new: tensor<F32, [1, 2, 1, 2], @host>,
  k_past: tensor<F32, [1, 2, 2, 3], @host>,
  k_new: tensor<F32, [1, 2, 2, 1], @host>,
  v_past: tensor<F32, [1, 2, 3, 2], @host>,
  v_new: tensor<F32, [1, 2, 1, 2], @host>,
…

29 lines · MLIR golden

storylm0_mean_nll_step.kina

Mean NLL ∘ decoder + multi-weight SGD; not reduce_mean

let storylm0_mean_nll = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  targets: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
…

113 lines · MLIR golden

storylm0_nll.kina

Host NLL as -log_softmax.take_along_axis; not training

let nll = fn(
  logits: tensor<F32, [1, 4, 8], @host>,
  targets: tensor<I64, [1, 4], @host>
) -> tensor<F32, [1, 4], @host> {
  let log_probs = logits.log_softmax(axis=2) in
  let picked = log_probs.take_along_axis(targets, axis=2) in
…

12 lines · MLIR golden

storylm0_nll_step.kina

NLL ∘ decoder + multi-weight SGD; not a trained model

let storylm0_nll = fn(
  token_ids: tensor<I64, [1, 4], @host>,
  targets: tensor<I64, [1, 4], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
…

102 lines · MLIR golden

storylm0_prefill.kina

Prefill cache handoff plus one decode step; not T_max scratch

let prefill = fn(
  token_ids: tensor<I64, [1, 3], @host>,
  token_embedding: tensor<F32, [8, 4], @host>,
  causal_mask: tensor<F32, [1, 1, 3, 3], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  wq: tensor<F32, [4, 4], @host>,
…

95 lines · MLIR golden

storylm0_rollout.kina

Host 32-token rollout from a patterned fixture; not a trained model

let greedy_row = fn(
  logits: tensor<F32, [1, 1, 8], @host>
) -> tensor<I64, [1, 1], @host> {
  let tokens = logits.reduce(op=argmax, axis=2) in
  tokens
} in
…

67 lines · MLIR golden

storylm0_sample.kina

Host Gumbel-max sample; seed is the Uniform tensor

let sample = fn(
  logits: tensor<F32, [1, 4, 8], @host>,
  u: tensor<F32, [1, 4, 8], @host>,
  temperature: tensor<F32, [1, 1], @host>
) -> tensor<I64, [1, 4], @host> {
  let log_u = log(u) in
…

20 lines · MLIR golden

storylm0_tmax.kina

Pad live KV to T_max with zeros, then slice the prefix

let pad_s3 = fn(
  k_live: tensor<F32, [1, 2, 2, 3], @host>,
  v_live: tensor<F32, [1, 2, 3, 2], @host>,
  zero: tensor<F32, [1, 1, 1, 1], @host>
) -> (
  tensor<F32, [1, 2, 2, 5], @host>,
…

44 lines · MLIR golden

storylm0_tp_heads.kina

2-way head-parallel attention; not @mesh or @sharded

let tp_heads = fn(
  q: tensor<F32, [1, 2, 4, 2], @host>,
  k_t: tensor<F32, [1, 2, 2, 4], @host>,
  v: tensor<F32, [1, 2, 4, 2], @host>,
  attention_scale: tensor<F32, [1, 1, 1, 1], @host>,
  causal_mask: tensor<F32, [1, 1, 4, 4], @host>
…

65 lines · MLIR golden

storylm0_tp_proj.kina

2-way column-parallel QKV; not @mesh or @sharded

let proj_q0 = fn(
  x: tensor<F32, [1, 4, 4], @host>,
  wq: tensor<F32, [4, 4], @host>
) -> tensor<F32, [1, 1, 4, 2], @host> {
  let wq0 = wq.slice(axis=1, start=0, stop=2) in
  let wq0_b = wq0.broadcast([1, 4, 2]) in
…

85 lines · MLIR golden

storylm0_tp_vocab.kina

2-way column-parallel lm_head; not @mesh or @sharded

let tp_logits = fn(
  residual: tensor<F32, [1, 4, 4], @host>,
  lm_head: tensor<F32, [4, 8], @host>
) -> tensor<F32, [1, 4, 8], @host> {
  let lm0 = lm_head.slice(axis=1, start=0, stop=4) in
  let lm1 = lm_head.slice(axis=1, start=4, stop=8) in
…

58 lines · MLIR golden

sub_2d.kina

Compiler example

(tensor<F32, [2, 3]> - tensor<F32, [2, 3]>)

2 lines · MLIR golden

symbolic_3d.kina

Compiler example

tensor<BF16, [B, S, D]>

2 lines · MLIR golden

symbolic_requires.kina

Compiler example

requires [B = B_prime, K = K_prime] compute {
  tensor<F32, [B, M, K], @host> @ tensor<F32, [B_prime, K_prime, N], @host>
}

4 lines

take_along_axis_interior.kina

Select along a compile-time axis (not 2D gather)

tensor<F32, [1, 8, 4], @host>.take_along_axis(tensor<I64, [1, 4], @host>, axis=1)

2 lines · MLIR golden

take_along_axis_last.kina

Select along a compile-time axis (not 2D gather)

tensor<F32, [1, 4, 8], @host>.take_along_axis(tensor<I64, [1, 4], @host>, axis=2)

2 lines · MLIR golden

transformer_ad.kina

Transformer block / restricted AD example

let fwd = fn(
  x: tensor<F32, [1, 32, 64], @host>, 
  gamma1: tensor<F32, [1, 1, 64], @host>,
  wq: tensor<F32, [1, 64, 64], @host>,
  wk: tensor<F32, [1, 64, 64], @host>,
  wv: tensor<F32, [1, 64, 64], @host>,
…

74 lines · MLIR golden

transformer_block.kina

Transformer block / restricted AD example

requires [B = B_p, Seq = Seq_p, D = D_p, D_ffn = D_ffn_p] compute {
    let run_transformer = fn(
        x: tensor<F32, [1, 32, 64], @host>,
        gamma1: tensor<F32, [1, 1, 64], @host>,
        wq: tensor<F32, [1, 64, 64], @host>,
        wk: tensor<F32, [1, 64, 64], @host>,
…

35 lines

transpose_2d.kina

Compiler example

tensor<F32, [3, 4]>.transpose([1, 0])

2 lines · MLIR golden

unary_bitwise_not_i32.kina

Unary op

!(tensor<I32, [2, 3]>)

2 lines · MLIR golden

unary_exp_f8e4m3.kina

Unary op

exp(tensor<F8E4M3, [2, 3]>)

2 lines · MLIR golden

unary_gelu_f32.kina

Unary op

gelu(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_log_f32.kina

Unary op

log(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_neg_bool.kina

Unary op

-tensor<Bool, []>

2 lines · MLIR golden

unary_neg_i32.kina

Unary op

-tensor<I32, [2, 2]>

2 lines · MLIR golden

unary_relu.kina

Unary op

relu(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_relu_f8e4m3.kina

Unary op

relu(tensor<F8E4M3, [2, 2]>)

2 lines · MLIR golden

unary_rsqrt_f32.kina

Unary op

rsqrt(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_sigmoid_f32.kina

Unary op

sigmoid(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_sigmoid_f8e4m3.kina

Unary op

sigmoid(tensor<F8E4M3, [1, 2]>)

2 lines · MLIR golden

unary_silu_f32.kina

Unary op

silu(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_sqrt_f32.kina

Unary op

sqrt(tensor<F32, [2, 3]>)

2 lines · MLIR golden

unary_tanh_f32.kina

Unary op

tanh(tensor<F32, [2, 3]>)

2 lines · MLIR golden

while_checkpoint.kina

Early control-flow / effect example

let loop_checkpoint_test = fn(steps: tensor<I32, [], @host>, init_val: tensor<F32, [4, 4], @host>, i_init: tensor<I32, [], @host>, step: tensor<I32, [], @host>) -> (tensor<F32, [4, 4], @host>, tensor<I32, [], @host>) {
  let final_val = @checkpoint(k=4) while (i < steps) [val = init_val, i = i_init] {
    let next_val = val + val in
    let next_i = i + step in
    (next_val, next_i)
  } in
…

14 lines · MLIR golden

while_loop.kina

Early control-flow / effect example

let loop_test = fn(steps: tensor<I32, [], @host>, init_val: tensor<F32, [4, 4], @host>, i_init: tensor<I32, [], @host>, step: tensor<I32, [], @host>) -> (tensor<F32, [4, 4], @host>, tensor<I32, [], @host>) {
  let final_val = while (i < steps) [val = init_val, i = i_init] {
    let next_val = val + val in
    let next_i = i + step in
    (next_val, next_i)
  } in
…

14 lines · MLIR golden