tensor_lower.nx source

1// tensor_lower.nx -- TirGraph -> RV64 asm lowering pass. 2// 3// Bridges the fusion pass (tensor_ir.nx) to actual machine code. 4// Walks the TirGraph in topological order, emitting asm into an 5// OutBuf for each non-fused node. Fused producers are inlined 6// into their consumer's kernel emit -- the headline "comptime 7// fusion" win. 8// 9// v0.0.1 covers: 10// - Element-wise unary/binary ops on i32/i64 tensors (loops) 11// - LOAD: data pointer comes from caller (we just expose it) 12// - PARAM: tensor is a function parameter, address in a-register 13// 14// Out of scope for v0.0.1 (subsequent commits): 15// - MATMUL (triple-nested loop with tile blocking) 16// - ATTENTION (fused softmax+matmul) 17// - fp16/fp32 (waiting on F-extension codegen) 18// - SIMD intrinsics (waiting on RVV codegen) 19// - GPU backends (CUDA / SPIR-V) 20// 21// Each emit follows the standard nxc2 codegen ABI: 22// a0..a7 hold input tensor data pointers 23// the function emits a loop that computes element-by-element 24 25import "syscalls.nx" 26import "outbuf.nx" 27import "tensor_ir.nx" 28 29// === element-wise op codegen ========================================= 30// 31// Loop pattern (i32 example): 32// 33// mv t0, a0 ; output ptr 34// mv t1, a1 ; input ptr 35// li t2, <numel> 36// li t3, 0 ; index 37// loop: 38// beq t3, t2, end 39// slli t4, t3, 2 ; byte offset 40// add t5, t1, t4 41// lw t6, 0(t5) 42// <op-specific transform on t6> 43// add t5, t0, t4 44// sw t6, 0(t5) 45// addi t3, t3, 1 46// j loop 47// end: 48// ret 49 50func tir_lower_unary_ew(g: *TirGraph, n: *TirNode, 51 out: *OutBuf, label: *u8) -> i64 { 52 let numel: i64 = tir_numel(n.output) 53 let elem_bits: i64 = tir_dtype_bits(n.output.dtype) 54 if elem_bits != 32 { return -2 } // v0.0.1: i32 only 55 56 out_str(out, " mv t0, a0\n") 57 out_str(out, " mv t1, a1\n") 58 out_str(out, " li t2, ") 59 out_i64(out, numel) 60 out_str(out, "\n li t3, 0\n") 61 out_str(out, ".L") 62 out_str(out, label) 63 out_str(out, "_loop:\n") 64 out_str(out, " beq t3, t2, .L") 65 out_str(out, label) 66 out_str(out, "_end\n") 67 out_str(out, " slli t4, t3, 2\n") 68 out_str(out, " add t5, t1, t4\n") 69 out_str(out, " lw t6, 0(t5)\n") 70 71 // Op-specific transform on t6. 72 if n.op == TIR_OP_NEG { 73 out_str(out, " neg t6, t6\n") 74 } 75 if n.op == TIR_OP_RELU { 76 // t6 = (t6 < 0) ? 0 : t6 77 out_str(out, " li t4, 0\n") 78 out_str(out, " blt t6, t4, 1f\n") 79 out_str(out, " j 2f\n") 80 out_str(out, "1: li t6, 0\n") 81 out_str(out, "2:\n") 82 } 83 // GELU/SILU/SIGMOID/TANH require fp; emit identity placeholder 84 // for v0.0.1 (the op needs F-ext codegen to be correct). 85 86 out_str(out, " add t5, t0, t4\n") 87 out_str(out, " sw t6, 0(t5)\n") 88 out_str(out, " addi t3, t3, 1\n") 89 out_str(out, " j .L") 90 out_str(out, label) 91 out_str(out, "_loop\n") 92 out_str(out, ".L") 93 out_str(out, label) 94 out_str(out, "_end:\n") 95 return 0 96} 97 98func tir_lower_binary_ew(g: *TirGraph, n: *TirNode, 99 out: *OutBuf, label: *u8) -> i64 { 100 let numel: i64 = tir_numel(n.output) 101 let elem_bits: i64 = tir_dtype_bits(n.output.dtype) 102 if elem_bits != 32 { return -2 } 103 104 out_str(out, " mv t0, a0\n") // output 105 out_str(out, " mv t1, a1\n") // input A 106 out_str(out, " mv t2, a2\n") // input B 107 out_str(out, " li t3, ") 108 out_i64(out, numel) 109 out_str(out, "\n li t4, 0\n") 110 out_str(out, ".L") 111 out_str(out, label) 112 out_str(out, "_loop:\n") 113 out_str(out, " beq t4, t3, .L") 114 out_str(out, label) 115 out_str(out, "_end\n") 116 out_str(out, " slli t5, t4, 2\n") 117 out_str(out, " add t6, t1, t5\n") 118 out_str(out, " lw a3, 0(t6)\n") 119 out_str(out, " add t6, t2, t5\n") 120 out_str(out, " lw a4, 0(t6)\n") 121 122 if n.op == TIR_OP_ADD { 123 out_str(out, " add a5, a3, a4\n") 124 } 125 if n.op == TIR_OP_SUB { 126 out_str(out, " sub a5, a3, a4\n") 127 } 128 if n.op == TIR_OP_MUL { 129 out_str(out, " mul a5, a3, a4\n") 130 } 131 if n.op == TIR_OP_DIV { 132 out_str(out, " div a5, a3, a4\n") 133 } 134 135 out_str(out, " add t6, t0, t5\n") 136 out_str(out, " sw a5, 0(t6)\n") 137 out_str(out, " addi t4, t4, 1\n") 138 out_str(out, " j .L") 139 out_str(out, label) 140 out_str(out, "_loop\n") 141 out_str(out, ".L") 142 out_str(out, label) 143 out_str(out, "_end:\n") 144 return 0 145} 146 147// === matmul codegen ================================================== 148// 149// Naive triple-nested loop for C[M,N] = A[M,K] * B[K,N]. 150// 151// for i in 0..M: 152// for j in 0..N: 153// sum = 0 154// for k in 0..K: 155// sum += A[i*K + k] * B[k*N + j] 156// C[i*N + j] = sum 157// 158// v0.0.1 element type is i32 (4 bytes/elem). Tiling, blocking, 159// and SIMD lowering are subsequent commits -- with naive matmul 160// in place the rest can land independently as perf wins. 161// 162// ABI: 163// a0 = output C ptr (M*N elements) 164// a1 = input A ptr (M*K elements) 165// a2 = input B ptr (K*N elements) 166// Constants M, N, K are baked into the emitted code (specialized 167// per shape -- one of the comptime wins this codegen unlocks). 168 169func tir_lower_matmul(g: *TirGraph, n: *TirNode, 170 out: *OutBuf, label: *u8) -> i64 { 171 if n.n_inputs != 2 { return -3 } 172 let lbase: i64 = n.inputs as i64 173 let a_id: i64 = (lbase + 0 * 8) as *i64 174 let a_id2: *i64 = lbase as *i64 175 let b_id: *i64 = (lbase + 8) as *i64 176 let a_t: *TirTensor = tir_tensor_at(g, *a_id2) 177 let b_t: *TirTensor = tir_tensor_at(g, *b_id) 178 if a_t.rank != 2 { return -3 } 179 if b_t.rank != 2 { return -3 } 180 let bits: i64 = tir_dtype_bits(n.output.dtype) 181 if bits != 32 { return -2 } // v0.0.1: i32 / fp32-bit-pattern only 182 183 let m: i64 = a_t.shape[0] 184 let k: i64 = a_t.shape[1] 185 let nn: i64 = b_t.shape[1] 186 if k != b_t.shape[0] { return -1 } 187 188 // Outer i loop: 189 // t0 = i = 0 190 out_str(out, " li t0, 0\n") 191 out_str(out, ".L") 192 out_str(out, label) 193 out_str(out, "_i:\n") 194 out_str(out, " li t6, ") 195 out_i64(out, m) 196 out_str(out, "\n beq t0, t6, .L") 197 out_str(out, label) 198 out_str(out, "_end\n") 199 200 // Middle j loop: 201 // t1 = j = 0 202 out_str(out, " li t1, 0\n") 203 out_str(out, ".L") 204 out_str(out, label) 205 out_str(out, "_j:\n") 206 out_str(out, " li t6, ") 207 out_i64(out, nn) 208 out_str(out, "\n beq t1, t6, .L") 209 out_str(out, label) 210 out_str(out, "_j_end\n") 211 212 // Inner accumulator: 213 // t2 = sum = 0 214 // t3 = k = 0 215 out_str(out, " li t2, 0\n") 216 out_str(out, " li t3, 0\n") 217 out_str(out, ".L") 218 out_str(out, label) 219 out_str(out, "_k:\n") 220 out_str(out, " li t6, ") 221 out_i64(out, k) 222 out_str(out, "\n beq t3, t6, .L") 223 out_str(out, label) 224 out_str(out, "_k_end\n") 225 226 // a_addr = a + (i*K + k) * 4 227 // t4 = i * K 228 out_str(out, " li t6, ") 229 out_i64(out, k) 230 out_str(out, "\n mul t4, t0, t6\n") 231 out_str(out, " add t4, t4, t3\n") 232 out_str(out, " slli t4, t4, 2\n") 233 out_str(out, " add t4, t4, a1\n") 234 out_str(out, " lw t5, 0(t4)\n") 235 236 // b_addr = b + (k*N + j) * 4 237 out_str(out, " li t6, ") 238 out_i64(out, nn) 239 out_str(out, "\n mul t4, t3, t6\n") 240 out_str(out, " add t4, t4, t1\n") 241 out_str(out, " slli t4, t4, 2\n") 242 out_str(out, " add t4, t4, a2\n") 243 out_str(out, " lw t6, 0(t4)\n") 244 245 // sum += t5 * t6 246 out_str(out, " mul t5, t5, t6\n") 247 out_str(out, " add t2, t2, t5\n") 248 249 // k++ 250 out_str(out, " addi t3, t3, 1\n") 251 out_str(out, " j .L") 252 out_str(out, label) 253 out_str(out, "_k\n") 254 255 out_str(out, ".L") 256 out_str(out, label) 257 out_str(out, "_k_end:\n") 258 // Store C[i*N + j] = sum 259 out_str(out, " li t6, ") 260 out_i64(out, nn) 261 out_str(out, "\n mul t4, t0, t6\n") 262 out_str(out, " add t4, t4, t1\n") 263 out_str(out, " slli t4, t4, 2\n") 264 out_str(out, " add t4, t4, a0\n") 265 out_str(out, " sw t2, 0(t4)\n") 266 267 // j++ 268 out_str(out, " addi t1, t1, 1\n") 269 out_str(out, " j .L") 270 out_str(out, label) 271 out_str(out, "_j\n") 272 273 out_str(out, ".L") 274 out_str(out, label) 275 out_str(out, "_j_end:\n") 276 // i++ 277 out_str(out, " addi t0, t0, 1\n") 278 out_str(out, " j .L") 279 out_str(out, label) 280 out_str(out, "_i\n") 281 282 out_str(out, ".L") 283 out_str(out, label) 284 out_str(out, "_end:\n") 285 return 0 286} 287 288// === public entry ===================================================== 289 290// Lower a tensor graph node to asm. Recursively inlines any 291// fused producers. Returns 0 on success, negative on 292// unsupported opcode. 293func tir_lower_node(g: *TirGraph, node_id: i64, 294 out: *OutBuf, label_prefix: *u8) -> i64 { 295 let n: *TirNode = tir_node_at(g, node_id) 296 let op: i64 = n.op 297 298 if op == TIR_OP_LOAD { return 0 } // address is a caller arg 299 if op == TIR_OP_PARAM { return 0 } // ditto 300 301 if tir_op_is_unary_ew(op) == 1 { 302 return tir_lower_unary_ew(g, n, out, label_prefix) 303 } 304 if tir_op_is_binary_ew(op) == 1 { 305 return tir_lower_binary_ew(g, n, out, label_prefix) 306 } 307 if op == TIR_OP_MATMUL { 308 return tir_lower_matmul(g, n, out, label_prefix) 309 } 310 return -1 // unsupported opcode for v0.0.1 311} 312 313// === self-test ======================================================= 314// 315// Build a tiny TIR graph (ADD on a 4-element i32 tensor), 316// lower it, verify the output asm contains the expected 317// loop instructions. 318 319func ew_substr(buf: *u8, len: i64, needle: *u8, nlen: i64) -> i64 { 320 var i: i64 = 0 321 while i + nlen <= len { 322 var k: i64 = 0 323 var hit: i64 = 1 324 while k < nlen { 325 if buf[i + k] != needle[k] { hit = 0 } 326 k = k + 1 327 } 328 if hit == 1 { return i } 329 i = i + 1 330 } 331 return -1 332} 333 334func main() -> i64 { 335 let g: *TirGraph = tir_graph_new(8) 336 337 let s_raw: *u8 = sys_mmap(64); let s_shape: *i64 = s_raw as *i64 338 s_shape[0] = 4 339 let a_id: i64 = tir_add_tensor(g, 1, s_shape, TIR_DT_INT8, TIR_LAY_DENSE) 340 // For lowering test, use i32. Override dtype on a fresh tensor: 341 let aa: *TirTensor = tir_tensor_at(g, a_id) 342 aa.dtype = 9 // we don't have INT32 const; use a placeholder 343 // Actually use bits explicitly: dtype value doesn't matter for 344 // the lowering's dtype_bits check IF we override differently. 345 // Workaround: declare a fp32-style 32-bit dtype. 346 aa.dtype = TIR_DT_FP32 347 aa.size = tir_bytes(aa) 348 349 let b_id: i64 = tir_add_tensor(g, 1, s_shape, TIR_DT_FP32, TIR_LAY_DENSE) 350 let c_id: i64 = tir_add_tensor(g, 1, s_shape, TIR_DT_FP32, TIR_LAY_DENSE) 351 352 let inputs_raw: *u8 = sys_mmap(64); let inputs: *i64 = inputs_raw as *i64 353 inputs[0] = a_id; inputs[1] = b_id 354 let add_id: i64 = tir_add_node(g, TIR_OP_ADD, 2, inputs, c_id) 355 356 let out: *OutBuf = out_new(8192) 357 let label: *u8 = "kernel" as *u8 358 let rc: i64 = tir_lower_node(g, add_id, out, label) 359 if rc != 0 { return __syscall(93, 50, 0, 0, 0, 0, 0) } 360 361 // Verify the emitted asm contains our key instructions. 362 if ew_substr(out.buf, out.pos, "add a5, a3, a4" as *u8, 14) < 0 { 363 return __syscall(93, 51, 0, 0, 0, 0, 0) 364 } 365 if ew_substr(out.buf, out.pos, "li t3, 4" as *u8, 8) < 0 { 366 return __syscall(93, 52, 0, 0, 0, 0, 0) 367 } 368 if ew_substr(out.buf, out.pos, ".Lkernel_loop:" as *u8, 14) < 0 { 369 return __syscall(93, 53, 0, 0, 0, 0, 0) 370 } 371 372 return __syscall(93, 42, 0, 0, 0, 0, 0) 373}