tensor_ir.nx
buildroot/runtime/tensor_ir.nx
about
tensor_ir.nx -- Nishi tensor intermediate representation.
Phase 2.2 of docs/AI_VRAM_ARCHITECTURE.md. Models tensor
operators as first-class IR so the compiler can do static
shape + dtype analysis, fuse adjacent ops at compile time,
and emit register-efficient kernels (CUDA, SPIR-V, RV64+RVV).
Why this is the AI win nobody else ships: every other framework
(PyTorch, JAX, vLLM, llama.cpp) does kernel fusion at runtime
via tracing or hand-written CUDA. Whole-program comptime info
in NishiLang lets us fuse + specialize at COMPILE time, with
no runtime kernel-launch overhead.
Status (v0.0.1): operator definitions + Tensor struct + Layout
enum. Fusion pass + lowering pass are subsequent commits.
dependencies 1 imports · 1 importers
imports: syscalls.nx
imported by: tensor_lower.nx
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| 42 | struct TirTensor { |
| 139 | struct TirNode { |
| 149 | struct TirGraph { |
consts
| 21 | const TIR_DT_FP32: i64 = 0 |
| 22 | const TIR_DT_FP16: i64 = 1 |
| 23 | const TIR_DT_BF16: i64 = 2 |
| 24 | const TIR_DT_FP8: i64 = 3 |
| 25 | const TIR_DT_INT8: i64 = 4 |
| 26 | const TIR_DT_INT4: i64 = 5 |
| 27 | const TIR_DT_INT2: i64 = 6 |
| 28 | const TIR_DT_TERN: i64 = 7 // ternary {-1, 0, +1} packed |
| 29 | const TIR_DT_BOOL: i64 = 8 |
| 33 | const TIR_LAY_DENSE: i64 = 0 // contiguous row-major |
| 34 | const TIR_LAY_STRIDED: i64 = 1 // arbitrary strides per dim |
| 35 | const TIR_LAY_BLOCKED: i64 = 2 // K x K micro-tiles (cache-friendly) |
| 36 | const TIR_LAY_NHWC: i64 = 3 // image: batch x H x W x channels |
| 37 | const TIR_LAY_NCHW: i64 = 4 // image: batch x channels x H x W |
| 38 | const TIR_LAY_PACKED: i64 = 5 // sub-byte (int4/int2/ternary) |
| 92 | const TIR_OP_LOAD: i64 = 0 // load constant / weight tensor |
| 93 | const TIR_OP_PARAM: i64 = 1 // function parameter (input) |
| 94 | const TIR_OP_RESHAPE: i64 = 2 // change view, no data copy |
| 95 | const TIR_OP_TRANSPOSE: i64 = 3 |
| 96 | const TIR_OP_CAST: i64 = 4 // dtype conversion |
| 97 | const TIR_OP_QUANTIZE: i64 = 5 // fp -> int (with scale/zp) |
| 98 | const TIR_OP_DEQUANTIZE: i64 = 6 // int -> fp |
| 101 | const TIR_OP_ADD: i64 = 10 |
| 102 | const TIR_OP_SUB: i64 = 11 |
| 103 | const TIR_OP_MUL: i64 = 12 |
| 104 | const TIR_OP_DIV: i64 = 13 |
| 105 | const TIR_OP_NEG: i64 = 14 |
| 106 | const TIR_OP_RELU: i64 = 15 |
| 107 | const TIR_OP_GELU: i64 = 16 |
| 108 | const TIR_OP_SILU: i64 = 17 |
| 109 | const TIR_OP_SIGMOID: i64 = 18 |
| 110 | const TIR_OP_TANH: i64 = 19 |
| 113 | const TIR_OP_SUM: i64 = 30 |
| 114 | const TIR_OP_MEAN: i64 = 31 |
| 115 | const TIR_OP_MAX: i64 = 32 |
| 116 | const TIR_OP_MIN: i64 = 33 |
| 117 | const TIR_OP_ARGMAX: i64 = 34 |
| 120 | const TIR_OP_RMS_NORM: i64 = 40 |
| 121 | const TIR_OP_LAYER_NORM: i64 = 41 |
| 122 | const TIR_OP_SOFTMAX: i64 = 42 |
| 125 | const TIR_OP_MATMUL: i64 = 50 |
| 126 | const TIR_OP_CONV2D: i64 = 51 |
| 127 | const TIR_OP_ATTENTION: i64 = 52 // fused QKV attention |
| 128 | const TIR_OP_ROPE: i64 = 53 // rotary positional embedding |
| 129 | const TIR_OP_EMBED: i64 = 54 // embedding lookup |
| 133 | const TIR_OP_MATMUL_INT4: i64 = 60 |
| 134 | const TIR_OP_MATMUL_FP8: i64 = 61 |
| 135 | const TIR_OP_ATTN_INT4_KV: i64 = 62 // attention with int4 KV cache |
| 160 | const TIR_NODE_BYTES: i64 = 64 // sizeof TirNode rounded |
| 161 | const TIR_TENSOR_BYTES: i64 = 64 |
| 226 | const TIR_SHAPE_MISMATCH: i64 = -1 |
| 227 | const TIR_RANK_MISMATCH: i64 = -2 |
| 228 | const TIR_UNSUPPORTED_OP: i64 = -3 |
functions
| 54 | func tir_numel(t: *TirTensor) -> i64 { |
| 64 | func tir_dtype_bits(dt: i64) -> i64 { |
| 79 | func tir_bytes(t: *TirTensor) -> i64 { |
| 163 | func tir_graph_new(cap: i64) -> *TirGraph { |
| 175 | func tir_node_at(g: *TirGraph, id: i64) -> *TirNode { |
| 180 | func tir_tensor_at(g: *TirGraph, id: i64) -> *TirTensor { |
| 186 | func tir_add_tensor(g: *TirGraph, rank: i64, shape: *i64, |
| 202 | func tir_add_node(g: *TirGraph, op: i64, n_inputs: i64, |
| 232 | func tir_infer_elemwise(in_a: *TirTensor, in_b: *TirTensor,
calls 1: tir_bytes |
| 250 | func tir_infer_matmul(in_a: *TirTensor, in_b: *TirTensor, |
| 270 | func tir_infer_reshape(in_t: *TirTensor, out: *TirTensor) -> i64 { |
| 279 | func tir_infer_cast(in_t: *TirTensor, target_dtype: i64,
calls 1: tir_bytes |
| 290 | func tir_infer_reduce(in_t: *TirTensor, axis: i64,
calls 1: tir_bytes |
| 315 | func tir_infer_attention(q: *TirTensor, k: *TirTensor, v: *TirTensor,
calls 1: tir_bytes |
| 362 | func tir_op_is_unary_ew(op: i64) -> i64 { |
| 373 | func tir_op_is_binary_ew(op: i64) -> i64 { |
| 384 | func tir_fuse_into(consumer: *TirNode, producer_id: i64) -> i64 {
called by 1: tir_fuse_graph |
| 396 | func tir_fuse_graph(g: *TirGraph) -> i64 { |
| 432 | func main() -> i64 { |