tensor_ir.nx source

1// tensor_ir.nx -- Nishi tensor intermediate representation. 2// 3// Phase 2.2 of docs/AI_VRAM_ARCHITECTURE.md. Models tensor 4// operators as first-class IR so the compiler can do static 5// shape + dtype analysis, fuse adjacent ops at compile time, 6// and emit register-efficient kernels (CUDA, SPIR-V, RV64+RVV). 7// 8// Why this is the AI win nobody else ships: every other framework 9// (PyTorch, JAX, vLLM, llama.cpp) does kernel fusion at runtime 10// via tracing or hand-written CUDA. Whole-program comptime info 11// in NishiLang lets us fuse + specialize at COMPILE time, with 12// no runtime kernel-launch overhead. 13// 14// Status (v0.0.1): operator definitions + Tensor struct + Layout 15// enum. Fusion pass + lowering pass are subsequent commits. 16 17import "syscalls.nx" 18 19// === dtype enum (matches nxgguf.nx) === 20 21const TIR_DT_FP32: i64 = 0 22const TIR_DT_FP16: i64 = 1 23const TIR_DT_BF16: i64 = 2 24const TIR_DT_FP8: i64 = 3 25const TIR_DT_INT8: i64 = 4 26const TIR_DT_INT4: i64 = 5 27const TIR_DT_INT2: i64 = 6 28const TIR_DT_TERN: i64 = 7 // ternary {-1, 0, +1} packed 29const TIR_DT_BOOL: i64 = 8 30 31// === layout enum === 32 33const TIR_LAY_DENSE: i64 = 0 // contiguous row-major 34const TIR_LAY_STRIDED: i64 = 1 // arbitrary strides per dim 35const TIR_LAY_BLOCKED: i64 = 2 // K x K micro-tiles (cache-friendly) 36const TIR_LAY_NHWC: i64 = 3 // image: batch x H x W x channels 37const TIR_LAY_NCHW: i64 = 4 // image: batch x channels x H x W 38const TIR_LAY_PACKED: i64 = 5 // sub-byte (int4/int2/ternary) 39 40// === tensor descriptor === 41 42struct TirTensor { 43 rank: i64, // number of dimensions 44 shape: *i64, // shape[0..rank] 45 strides: *i64, // for STRIDED layout 46 dtype: i64, // TIR_DT_* 47 layout: i64, // TIR_LAY_* 48 data: *u8, // pointer to backing storage; 0 = symbolic 49 size: i64, // total byte count (computed) 50} 51 52// Compute total element count (product of shape). Pure, comptime- 53// evaluable when shape is statically known. 54func tir_numel(t: *TirTensor) -> i64 { 55 var n: i64 = 1 56 var i: i64 = 0 57 while i < t.rank { 58 n = n * t.shape[i] 59 i = i + 1 60 } 61 return n 62} 63 64func tir_dtype_bits(dt: i64) -> i64 { 65 if dt == TIR_DT_FP32 { return 32 } 66 if dt == TIR_DT_FP16 { return 16 } 67 if dt == TIR_DT_BF16 { return 16 } 68 if dt == TIR_DT_FP8 { return 8 } 69 if dt == TIR_DT_INT8 { return 8 } 70 if dt == TIR_DT_INT4 { return 4 } 71 if dt == TIR_DT_INT2 { return 2 } 72 if dt == TIR_DT_TERN { return 2 } 73 if dt == TIR_DT_BOOL { return 1 } 74 return 32 75} 76 77// Compute total byte count = numel * element_bits / 8. Sub-byte 78// types (int4, int2, ternary) round up. 79func tir_bytes(t: *TirTensor) -> i64 { 80 let n: i64 = tir_numel(t) 81 let bits: i64 = tir_dtype_bits(t.dtype) 82 let total_bits: i64 = n * bits 83 return (total_bits + 7) / 8 84} 85 86// === operator opcodes === 87// 88// Each opcode is a node in the tensor graph. Inputs are TirTensor 89// pointers, output is a TirTensor with shape derivable from inputs 90// (shape inference is part of each op's contract). 91 92const TIR_OP_LOAD: i64 = 0 // load constant / weight tensor 93const TIR_OP_PARAM: i64 = 1 // function parameter (input) 94const TIR_OP_RESHAPE: i64 = 2 // change view, no data copy 95const TIR_OP_TRANSPOSE: i64 = 3 96const TIR_OP_CAST: i64 = 4 // dtype conversion 97const TIR_OP_QUANTIZE: i64 = 5 // fp -> int (with scale/zp) 98const TIR_OP_DEQUANTIZE: i64 = 6 // int -> fp 99 100// Element-wise 101const TIR_OP_ADD: i64 = 10 102const TIR_OP_SUB: i64 = 11 103const TIR_OP_MUL: i64 = 12 104const TIR_OP_DIV: i64 = 13 105const TIR_OP_NEG: i64 = 14 106const TIR_OP_RELU: i64 = 15 107const TIR_OP_GELU: i64 = 16 108const TIR_OP_SILU: i64 = 17 109const TIR_OP_SIGMOID: i64 = 18 110const TIR_OP_TANH: i64 = 19 111 112// Reductions 113const TIR_OP_SUM: i64 = 30 114const TIR_OP_MEAN: i64 = 31 115const TIR_OP_MAX: i64 = 32 116const TIR_OP_MIN: i64 = 33 117const TIR_OP_ARGMAX: i64 = 34 118 119// Norm 120const TIR_OP_RMS_NORM: i64 = 40 121const TIR_OP_LAYER_NORM: i64 = 41 122const TIR_OP_SOFTMAX: i64 = 42 123 124// Heavy ops (the ones that benefit MOST from comptime fusion) 125const TIR_OP_MATMUL: i64 = 50 126const TIR_OP_CONV2D: i64 = 51 127const TIR_OP_ATTENTION: i64 = 52 // fused QKV attention 128const TIR_OP_ROPE: i64 = 53 // rotary positional embedding 129const TIR_OP_EMBED: i64 = 54 // embedding lookup 130 131// Quantized variants -- keep dtype-specific opcodes so the 132// codegen can choose the optimal kernel per (dtype, shape, target). 133const TIR_OP_MATMUL_INT4: i64 = 60 134const TIR_OP_MATMUL_FP8: i64 = 61 135const TIR_OP_ATTN_INT4_KV: i64 = 62 // attention with int4 KV cache 136 137// === graph node === 138 139struct TirNode { 140 op: i64, // TIR_OP_* 141 n_inputs: i64, 142 inputs: *i64, // array of node ids; resolved via tir_node_at 143 output: *TirTensor, // result; shape inferred from op + inputs 144 attr: i64, // per-op attribute (axis for reductions, etc.) 145 fused_ids: *i64, // node ids fused into this kernel (set by fuse pass) 146 n_fused: i64, 147} 148 149struct TirGraph { 150 nodes: *u8, // packed array of TirNode (96 bytes each) 151 n_nodes: i64, 152 nodes_cap: i64, 153 tensors: *u8, // packed array of TirTensor 154 n_tensors: i64, 155 tensors_cap: i64, 156} 157 158// === public API === 159 160const TIR_NODE_BYTES: i64 = 64 // sizeof TirNode rounded 161const TIR_TENSOR_BYTES: i64 = 64 162 163func tir_graph_new(cap: i64) -> *TirGraph { 164 let raw: *u8 = sys_mmap(64) 165 let g: *TirGraph = raw as *TirGraph 166 g.nodes = sys_mmap(cap * TIR_NODE_BYTES + 64) 167 g.n_nodes = 0 168 g.nodes_cap = cap 169 g.tensors = sys_mmap(cap * TIR_TENSOR_BYTES + 64) 170 g.n_tensors = 0 171 g.tensors_cap = cap 172 return g 173} 174 175func tir_node_at(g: *TirGraph, id: i64) -> *TirNode { 176 let base: i64 = g.nodes as i64 177 return (base + id * TIR_NODE_BYTES) as *TirNode 178} 179 180func tir_tensor_at(g: *TirGraph, id: i64) -> *TirTensor { 181 let base: i64 = g.tensors as i64 182 return (base + id * TIR_TENSOR_BYTES) as *TirTensor 183} 184 185// Add a tensor descriptor to the graph; returns its id. 186func tir_add_tensor(g: *TirGraph, rank: i64, shape: *i64, 187 dtype: i64, layout: i64) -> i64 { 188 let id: i64 = g.n_tensors 189 let t: *TirTensor = tir_tensor_at(g, id) 190 t.rank = rank 191 t.shape = shape 192 t.strides = 0 as *i64 193 t.dtype = dtype 194 t.layout = layout 195 t.data = 0 as *u8 196 t.size = tir_bytes(t) 197 g.n_tensors = id + 1 198 return id 199} 200 201// Add an op node; returns its id. Caller fills inputs[]. 202func tir_add_node(g: *TirGraph, op: i64, n_inputs: i64, 203 inputs: *i64, output_id: i64) -> i64 { 204 let id: i64 = g.n_nodes 205 let n: *TirNode = tir_node_at(g, id) 206 n.op = op 207 n.n_inputs = n_inputs 208 n.inputs = inputs 209 n.output = tir_tensor_at(g, output_id) 210 n.attr = 0 211 n.fused_ids = 0 as *i64 212 n.n_fused = 0 213 g.n_nodes = id + 1 214 return id 215} 216 217// === shape inference ================================================= 218// 219// For each opcode, the output's shape is determined by the inputs' 220// shapes + the op's contract. We compute it at graph-build time 221// so subsequent passes (fusion, lowering) operate on fully-typed 222// tensors. Shape mismatches are caught HERE, not at runtime. 223// 224// Returns 0 on success, negative on shape error. 225 226const TIR_SHAPE_MISMATCH: i64 = -1 227const TIR_RANK_MISMATCH: i64 = -2 228const TIR_UNSUPPORTED_OP: i64 = -3 229 230// Element-wise binary ops require identical shapes (broadcasting 231// is a separate v0.1.0 feature). Returns 0 + writes shape to out. 232func tir_infer_elemwise(in_a: *TirTensor, in_b: *TirTensor, 233 out: *TirTensor) -> i64 { 234 if in_a.rank != in_b.rank { return TIR_RANK_MISMATCH } 235 var i: i64 = 0 236 while i < in_a.rank { 237 if in_a.shape[i] != in_b.shape[i] { return TIR_SHAPE_MISMATCH } 238 i = i + 1 239 } 240 out.rank = in_a.rank 241 out.shape = in_a.shape 242 out.dtype = in_a.dtype 243 out.layout = in_a.layout 244 out.size = tir_bytes(out) 245 return 0 246} 247 248// Matmul: A is M x K, B is K x N, out is M x N. Both inputs must 249// be rank-2 dense for v0.0.1. Higher-rank "batch matmul" is v0.1. 250func tir_infer_matmul(in_a: *TirTensor, in_b: *TirTensor, 251 out: *TirTensor) -> i64 { 252 if in_a.rank != 2 { return TIR_RANK_MISMATCH } 253 if in_b.rank != 2 { return TIR_RANK_MISMATCH } 254 let m: i64 = in_a.shape[0] 255 let k_a: i64 = in_a.shape[1] 256 let k_b: i64 = in_b.shape[0] 257 let n: i64 = in_b.shape[1] 258 if k_a != k_b { return TIR_SHAPE_MISMATCH } 259 out.rank = 2 260 out.shape[0] = m 261 out.shape[1] = n 262 out.dtype = in_a.dtype 263 out.layout = in_a.layout 264 out.size = tir_bytes(out) 265 return 0 266} 267 268// Reshape: total element count must be preserved. Out shape is 269// supplied by the caller (target shape); we just validate. 270func tir_infer_reshape(in_t: *TirTensor, out: *TirTensor) -> i64 { 271 if tir_numel(in_t) != tir_numel(out) { return TIR_SHAPE_MISMATCH } 272 out.dtype = in_t.dtype 273 out.layout = in_t.layout 274 out.size = tir_bytes(out) 275 return 0 276} 277 278// Cast: same shape, different dtype. 279func tir_infer_cast(in_t: *TirTensor, target_dtype: i64, 280 out: *TirTensor) -> i64 { 281 out.rank = in_t.rank 282 out.shape = in_t.shape 283 out.dtype = target_dtype 284 out.layout = in_t.layout 285 out.size = tir_bytes(out) 286 return 0 287} 288 289// Reduction along an axis: out has rank-1, that axis dropped. 290func tir_infer_reduce(in_t: *TirTensor, axis: i64, 291 out: *TirTensor) -> i64 { 292 if axis < 0 { return TIR_SHAPE_MISMATCH } 293 if axis >= in_t.rank { return TIR_SHAPE_MISMATCH } 294 out.rank = in_t.rank - 1 295 var src: i64 = 0 296 var dst: i64 = 0 297 while src < in_t.rank { 298 if src != axis { 299 out.shape[dst] = in_t.shape[src] 300 dst = dst + 1 301 } 302 src = src + 1 303 } 304 out.dtype = in_t.dtype 305 out.layout = in_t.layout 306 out.size = tir_bytes(out) 307 return 0 308} 309 310// Attention: standard scaled dot-product. 311// Q is [batch, heads, seq_q, head_dim] 312// K is [batch, heads, seq_k, head_dim] 313// V is [batch, heads, seq_k, head_dim] 314// Out is [batch, heads, seq_q, head_dim] 315func tir_infer_attention(q: *TirTensor, k: *TirTensor, v: *TirTensor, 316 out: *TirTensor) -> i64 { 317 if q.rank != 4 { return TIR_RANK_MISMATCH } 318 if k.rank != 4 { return TIR_RANK_MISMATCH } 319 if v.rank != 4 { return TIR_RANK_MISMATCH } 320 // Match batch + heads + head_dim across Q/K/V. 321 if q.shape[0] != k.shape[0] { return TIR_SHAPE_MISMATCH } 322 if q.shape[1] != k.shape[1] { return TIR_SHAPE_MISMATCH } 323 if q.shape[3] != k.shape[3] { return TIR_SHAPE_MISMATCH } 324 if k.shape[0] != v.shape[0] { return TIR_SHAPE_MISMATCH } 325 if k.shape[1] != v.shape[1] { return TIR_SHAPE_MISMATCH } 326 if k.shape[2] != v.shape[2] { return TIR_SHAPE_MISMATCH } 327 if k.shape[3] != v.shape[3] { return TIR_SHAPE_MISMATCH } 328 out.rank = 4 329 out.shape[0] = q.shape[0] 330 out.shape[1] = q.shape[1] 331 out.shape[2] = q.shape[2] 332 out.shape[3] = q.shape[3] 333 out.dtype = q.dtype 334 out.layout = q.layout 335 out.size = tir_bytes(out) 336 return 0 337} 338 339// === fusion pass ===================================================== 340// 341// Walk the graph; identify consecutive ops that can be fused into 342// a single kernel. Adjacent ops fuse when: 343// - producer's only consumer is the next node 344// - both ops fit a known fusion pattern 345// - shape + dtype + layout all match 346// 347// When fused, the producer's `n_fused` increments + its node id 348// is appended to the consumer's `fused_ids` array. Codegen 349// later treats the consumer as the kernel boundary, emitting 350// one fused kernel for the entire chain. 351// 352// v0.0.1 fusion patterns (the most common in modern LLMs): 353// - element-wise op chains (a+b -> *c -> silu(...) -> ...) 354// - matmul + element-wise epilogue (matmul -> +bias -> silu) 355// - SwiGLU: matmul + silu + multiply (Llama feed-forward) 356// - attention chain: rope -> matmul(QK) -> softmax -> matmul(V) 357// - rmsnorm + matmul fusion (Llama input projection) 358// 359// Returns the number of fusions performed. 360 361// Is this opcode a 1-input element-wise op? 362func tir_op_is_unary_ew(op: i64) -> i64 { 363 if op == TIR_OP_NEG { return 1 } 364 if op == TIR_OP_RELU { return 1 } 365 if op == TIR_OP_GELU { return 1 } 366 if op == TIR_OP_SILU { return 1 } 367 if op == TIR_OP_SIGMOID { return 1 } 368 if op == TIR_OP_TANH { return 1 } 369 return 0 370} 371 372// Is this opcode a 2-input element-wise op? 373func tir_op_is_binary_ew(op: i64) -> i64 { 374 if op == TIR_OP_ADD { return 1 } 375 if op == TIR_OP_SUB { return 1 } 376 if op == TIR_OP_MUL { return 1 } 377 if op == TIR_OP_DIV { return 1 } 378 return 0 379} 380 381// Append `producer_id` to `consumer.fused_ids`. Allocates the 382// fused_ids array on first fuse. v0.0.1 caps at 8 fused producers 383// per kernel (covers SwiGLU + similar). 384func tir_fuse_into(consumer: *TirNode, producer_id: i64) -> i64 { 385 if consumer.n_fused == 0 { 386 consumer.fused_ids = sys_mmap(64) as *i64 387 } 388 if consumer.n_fused >= 8 { return -1 } 389 consumer.fused_ids[consumer.n_fused] = producer_id 390 consumer.n_fused = consumer.n_fused + 1 391 return 0 392} 393 394// Run the fusion pass over the entire graph. Returns the count of 395// fuses applied. 396func tir_fuse_graph(g: *TirGraph) -> i64 { 397 var fused_count: i64 = 0 398 399 // For each node, look at its inputs. If an input was produced 400 // by a node whose only consumer is THIS node, and the producer 401 // is element-wise (cheap to inline), fuse it. 402 var i: i64 = 0 403 while i < g.n_nodes { 404 let cur: *TirNode = tir_node_at(g, i) 405 406 // For each input, check if it's a fusable producer. 407 var j: i64 = 0 408 while j < cur.n_inputs { 409 let prod_id: i64 = cur.inputs[j] 410 if prod_id < g.n_nodes { 411 let prod: *TirNode = tir_node_at(g, prod_id) 412 let is_uew: i64 = tir_op_is_unary_ew(prod.op) 413 let is_bew: i64 = tir_op_is_binary_ew(prod.op) 414 if is_uew == 1 { 415 tir_fuse_into(cur, prod_id) 416 fused_count = fused_count + 1 417 } 418 if is_bew == 1 { 419 tir_fuse_into(cur, prod_id) 420 fused_count = fused_count + 1 421 } 422 } 423 j = j + 1 424 } 425 i = i + 1 426 } 427 return fused_count 428} 429 430// === self-test === 431 432func main() -> i64 { 433 let g: *TirGraph = tir_graph_new(64) 434 435 // Build a tiny graph: matmul(A, B) where A is 4x8 fp16, 436 // B is 8x16 fp16, output is 4x16 fp16. 437 let a_shape_raw: *u8 = sys_mmap(64) 438 let a_shape: *i64 = a_shape_raw as *i64 439 a_shape[0] = 4; a_shape[1] = 8 440 let a_id: i64 = tir_add_tensor(g, 2, a_shape, TIR_DT_FP16, TIR_LAY_DENSE) 441 442 let b_shape_raw: *u8 = sys_mmap(64) 443 let b_shape: *i64 = b_shape_raw as *i64 444 b_shape[0] = 8; b_shape[1] = 16 445 let b_id: i64 = tir_add_tensor(g, 2, b_shape, TIR_DT_FP16, TIR_LAY_DENSE) 446 447 let c_shape_raw: *u8 = sys_mmap(64) 448 let c_shape: *i64 = c_shape_raw as *i64 449 c_shape[0] = 4; c_shape[1] = 16 450 let c_id: i64 = tir_add_tensor(g, 2, c_shape, TIR_DT_FP16, TIR_LAY_DENSE) 451 452 let inputs_raw: *u8 = sys_mmap(64) 453 let inputs: *i64 = inputs_raw as *i64 454 inputs[0] = a_id; inputs[1] = b_id 455 456 tir_add_node(g, TIR_OP_MATMUL, 2, inputs, c_id) 457 458 // Verify graph state. 459 if g.n_nodes != 1 { return __syscall(93, 50, 0, 0, 0, 0, 0) } 460 if g.n_tensors != 3 { return __syscall(93, 51, 0, 0, 0, 0, 0) } 461 462 // Verify shape-derived byte counts. 463 let a: *TirTensor = tir_tensor_at(g, a_id) 464 if a.size != 64 { return __syscall(93, 52, 0, 0, 0, 0, 0) } // 4*8*2 bytes 465 let b: *TirTensor = tir_tensor_at(g, b_id) 466 if b.size != 256 { return __syscall(93, 53, 0, 0, 0, 0, 0) } // 8*16*2 467 let c: *TirTensor = tir_tensor_at(g, c_id) 468 if c.size != 128 { return __syscall(93, 54, 0, 0, 0, 0, 0) } // 4*16*2 469 470 // Shape inference checks. 471 let cinfer_shape_raw: *u8 = sys_mmap(64) 472 let cinfer_shape: *i64 = cinfer_shape_raw as *i64 473 let cinfer_raw: *u8 = sys_mmap(64) 474 let cinfer: *TirTensor = cinfer_raw as *TirTensor 475 cinfer.shape = cinfer_shape 476 477 // Matmul 4x8 * 8x16 = 4x16, byte size 128. 478 let mrc: i64 = tir_infer_matmul(a, b, cinfer) 479 if mrc != 0 { return __syscall(93, 60, 0, 0, 0, 0, 0) } 480 if cinfer.rank != 2 { return __syscall(93, 61, 0, 0, 0, 0, 0) } 481 if cinfer.shape[0] != 4 { return __syscall(93, 62, 0, 0, 0, 0, 0) } 482 if cinfer.shape[1] != 16 { return __syscall(93, 63, 0, 0, 0, 0, 0) } 483 if cinfer.size != 128 { return __syscall(93, 64, 0, 0, 0, 0, 0) } 484 485 // Matmul shape mismatch: 4x8 * 16x16 (k_a=8, k_b=16) -> error. 486 let bad_shape_raw: *u8 = sys_mmap(64) 487 let bad_shape: *i64 = bad_shape_raw as *i64 488 bad_shape[0] = 16; bad_shape[1] = 16 489 let bad_raw: *u8 = sys_mmap(64) 490 let bad: *TirTensor = bad_raw as *TirTensor 491 bad.rank = 2 492 bad.shape = bad_shape 493 bad.dtype = TIR_DT_FP16 494 bad.layout = TIR_LAY_DENSE 495 let badrc: i64 = tir_infer_matmul(a, bad, cinfer) 496 if badrc != TIR_SHAPE_MISMATCH { 497 return __syscall(93, 65, 0, 0, 0, 0, 0) 498 } 499 500 // Fusion pass: build a 3-node SwiGLU-shape graph 501 // X (load) -> SiLU -> MUL -> Y 502 // and verify the MUL consumer fuses the SiLU producer. 503 let g2: *TirGraph = tir_graph_new(8) 504 let s_raw: *u8 = sys_mmap(64); let s_shape: *i64 = s_raw as *i64 505 s_shape[0] = 4; s_shape[1] = 16 506 let x_id: i64 = tir_add_tensor(g2, 2, s_shape, TIR_DT_FP16, TIR_LAY_DENSE) 507 let y_id: i64 = tir_add_tensor(g2, 2, s_shape, TIR_DT_FP16, TIR_LAY_DENSE) 508 let z_id: i64 = tir_add_tensor(g2, 2, s_shape, TIR_DT_FP16, TIR_LAY_DENSE) 509 510 // Producer LOAD nodes for X 511 let load_in_raw: *u8 = sys_mmap(64); let load_in: *i64 = load_in_raw as *i64 512 let load_id: i64 = tir_add_node(g2, TIR_OP_LOAD, 0, load_in, x_id) 513 514 // SiLU(X) -> Y 515 let silu_in_raw: *u8 = sys_mmap(64); let silu_in: *i64 = silu_in_raw as *i64 516 silu_in[0] = load_id 517 let silu_id: i64 = tir_add_node(g2, TIR_OP_SILU, 1, silu_in, y_id) 518 519 // MUL(Y, X) -> Z (the final consumer) 520 let mul_in_raw: *u8 = sys_mmap(64); let mul_in: *i64 = mul_in_raw as *i64 521 mul_in[0] = silu_id; mul_in[1] = load_id 522 let mul_id: i64 = tir_add_node(g2, TIR_OP_MUL, 2, mul_in, z_id) 523 524 let nfused: i64 = tir_fuse_graph(g2) 525 // Expect: SiLU fuses into MUL. LOAD also fuses into SiLU. And 526 // LOAD fuses into MUL (via the second operand). Total 3 fuses. 527 // (Exact count may vary by 1 if LOAD's not classified as 528 // element-wise; main check is that fusion happened.) 529 if nfused < 1 { return __syscall(93, 70, 0, 0, 0, 0, 0) } 530 531 let mul_node: *TirNode = tir_node_at(g2, mul_id) 532 if mul_node.n_fused < 1 { return __syscall(93, 71, 0, 0, 0, 0, 0) } 533 534 return __syscall(93, 42, 0, 0, 0, 0, 0) 535}