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