nx_eqsat_constfold_bench.nx source

1// nx_eqsat_constfold_bench.nx -- CONST-FOLD-HEAVY race vs egg 0.11.0. 2// 3// THE RACE (engine-only, SAME const-heavy expression both sides): 4// egg target = nxc2/bench/nx_constfold_egg_bench.rs (copy into 5// /tmp/race_egg/examples/), the const-fold ANALYSIS graph build(n): 6// acc = 1; for i in 0..N: leaf=(i%7)+2; acc = acc {+,*,-} leaf (cycle by i%3) 7// so EVERY subtree is a known constant -- the const-fold analysis must evaluate + 8// collapse all N nodes, folding the whole chain to a single Num. egg does this via 9// its ConstantFold Analysis (make/merge/modify); WE do it via make-on-add + a CONST 10// merge through the union chokepoint citing NX_EQSAT_RULE_CONSTFOLD. 11// 12// OUR side, per timed iteration (the SAME end-to-end optimize pipeline egg ran): 13// nx_eqsat_init (fresh e-graph) 14// enable_constfold (egg's ConstantFold analysis ON) 15// build the all-const chain (egg's RecExpr build; fold fires at add-time) 16// nx_eqsat_saturate (egg's Runner::run -- analysis to fixpoint) 17// nx_eqsat_recompute_best (egg's Extractor bottom-up cost) 18// nx_eqsat_extract_best_node(root) (egg's find_best -> a single Num) 19// 20// FAIL LOUD: before timing, one verified iteration asserts the WHOLE chain folded 21// -- root's canonical best_node is a CONST (cost 0). A fast wrong answer (chain not 22// folded) loses, exit nonzero. 23// 24// Output (one line): 25// NX_CONSTFOLD_ENGINE_ONLY N=<n> ITERS=<it> FOLD=<v> TOTAL_US=<t> US_PER_OP_NS=<u> 26// (US_PER_OP_NS = total_us*1000/iters = nanoseconds/op; /1000 for us/op.) 27// 28// SOVEREIGN: no .sh; runs on the pinned NishiLang compiler. license_tier: ORIGINAL 29 30import "nx_eqsat.nx" 31const K_MAGIC_2000: i64 = 2000 32const K_MAGIC_4096: i64 = 4096 33 34func _emit_num(v: i64) -> i64 { 35 let b: *u8 = sys_mmap(28); var n: i64 = v; if n < 0 { n = 0 - n; sys_write(1, "-" as *u8, 1) } 36 let t2: *u8 = sys_mmap(28); var t: i64 = 0 37 if n == 0 { t2[0] = 48; t = 1 } 38 while n > 0 { t2[t] = 48 + (n % 10); n = n / 10; t = t + 1 } 39 var i: i64 = 0; while i < t { b[i] = t2[t - 1 - i]; i = i + 1 } 40 b[t] = 32; sys_write(1, b, t + 1); return 0 41} 42func _emit_str(s: *u8, len: i64) -> i64 { sys_write(1, s, len); return 0 } 43func _nl() -> i64 { let z: *u8 = sys_mmap(2); z[0] = 10; sys_write(1, z, 1); return 0 } 44 45// Build the all-const chain in `g`, run the analysis to fixpoint, extract. Writes 46// the folded constant value to fold_out[0] and returns 1 on success (the whole 47// chain collapsed to ONE CONST), or 0 on failure. Mirrors egg's build(n) + 48// Runner::run(&[]) (NO rewrite rules -- const-fold is an ANALYSIS, not a rewrite): 49// egg's analysis runs during build + rebuild; OUR analysis runs eagerly at add-time 50// (make-on-add), so the equivalent of egg's empty-rules run is just the (already- 51// done) eager fold + an extract. We do NOT run the 7 strength-reduction inline 52// rules -- that would be EXTRA work egg's const-fold bench never does (unfair). The 53// folded value is read off the root's canonical best_node, which the CONST merge 54// makes the cheapest node (CONST cost 0). cycle +,*,- over leaves like egg's build. 55func _one_iter(g: *NxEGraph, nodes: *NxENode, cap_nodes: i64, 56 classes: *NxEClass, cap_cls: i64, n: i64, fold_out: *i64) -> i64 { 57 if nx_eqsat_init(g, nodes, cap_nodes, classes, cap_cls) != NX_EQSAT_OK { return 0 } 58 if nx_eqsat_enable_constfold(g) != NX_EQSAT_OK { return 0 } 59 var acc: i64 = nx_eqsat_add_const(g, 1) 60 if acc < 0 { return 0 } 61 var i: i64 = 0 62 while i < n { 63 let leaf: i64 = nx_eqsat_add_const(g, (i % 7) + 2) 64 if leaf < 0 { return 0 } 65 let sel: i64 = i % 3 66 if sel == 0 { acc = nx_eqsat_add_binary(g, NX_EQ_OP_ADD, acc, leaf) } 67 if sel == 1 { acc = nx_eqsat_add_binary(g, NX_EQ_OP_MUL, acc, leaf) } 68 if sel == 2 { acc = nx_eqsat_add_binary(g, NX_EQ_OP_SUB, acc, leaf) } 69 if acc < 0 { return 0 } 70 i = i + 1 71 } 72 // EXTRACT (egg's Extractor::find_best) -- the const-fold merges already happened 73 // eagerly at add-time, so recompute_best + extract reads the fully-folded root. 74 if nx_eqsat_recompute_best(g) != NX_EQSAT_OK { return 0 } 75 let best: i64 = nx_eqsat_extract_best_node(g, acc) 76 if best < 0 { return 0 } 77 // the whole chain must have folded to a single CONST (egg's fully-folded root). 78 if g.nodes[best].op != NX_EQ_OP_CONST { return 0 } 79 fold_out[0] = g.nodes[best].payload 80 return 1 81} 82 83func main() -> i64 { 84 let n: i64 = 64 // chain length (matches egg default N=64) 85 let iters: i64 = K_MAGIC_2000 // matches egg iters for N=64 86 87 // Per iter we add ~ 1 + 2*N nodes (op + its folded CONST per step) + interned 88 // CONST leaves. N=64 -> well under 1024; size caps generously. 89 let cap_nodes: i64 = K_MAGIC_4096 90 let cap_cls: i64 = K_MAGIC_4096 91 let nodes: *NxENode = sys_mmap(cap_nodes * 64) as *NxENode 92 let classes: *NxEClass = sys_mmap(cap_cls * 32) as *NxEClass 93 let g: *NxEGraph = sys_mmap(128) as *NxEGraph 94 95 // ---- CORRECTNESS GATE (FAIL LOUD before timing): the chain MUST fully fold. 96 let fbuf: *i64 = sys_mmap(8) as *i64 97 if _one_iter(g, nodes, cap_nodes, classes, cap_cls, n, fbuf) != 1 { sys_exit(80); return 80 } 98 let fold: i64 = fbuf[0] 99 100 // ---- TIMED LOOP ------------------------------------------------------------- 101 let t0: i64 = sys_clock_now_us() 102 var k: i64 = 0 103 while k < iters { 104 if _one_iter(g, nodes, cap_nodes, classes, cap_cls, n, fbuf) != 1 { sys_exit(81); return 81 } // FAIL LOUD mid-run (no fast wrong answer) 105 k = k + 1 106 } 107 let t1: i64 = sys_clock_now_us() 108 let total_us: i64 = t1 - t0 109 let us_per_op_ns: i64 = (total_us * 1000) / iters 110 111 // ---- REPORT ----------------------------------------------------------------- 112 let tag: *u8 = sys_mmap(40) 113 // "NX_CONSTFOLD_ENGINE_ONLY " 114 tag[0]=78; tag[1]=88; tag[2]=95; tag[3]=67; tag[4]=79; tag[5]=78; tag[6]=83; tag[7]=84 115 tag[8]=70; tag[9]=79; tag[10]=76; tag[11]=68; tag[12]=95; tag[13]=69; tag[14]=78; tag[15]=71 116 tag[16]=73; tag[17]=78; tag[18]=69; tag[19]=95; tag[20]=79; tag[21]=78; tag[22]=76; tag[23]=89; tag[24]=32 117 _emit_str(tag, 25) 118 let s_n: *u8 = sys_mmap(4); s_n[0]=78; s_n[1]=61 // "N=" 119 _emit_str(s_n, 2); _emit_num(n) 120 let s_it: *u8 = sys_mmap(8); s_it[0]=73; s_it[1]=84; s_it[2]=69; s_it[3]=82; s_it[4]=83; s_it[5]=61 // "ITERS=" 121 _emit_str(s_it, 6); _emit_num(iters) 122 let s_fd: *u8 = sys_mmap(8); s_fd[0]=70; s_fd[1]=79; s_fd[2]=76; s_fd[3]=68; s_fd[4]=61 // "FOLD=" 123 _emit_str(s_fd, 5); _emit_num(fold) 124 let s_to: *u8 = sys_mmap(16); s_to[0]=84; s_to[1]=79; s_to[2]=84; s_to[3]=65; s_to[4]=76; s_to[5]=95; s_to[6]=85; s_to[7]=83; s_to[8]=61 // "TOTAL_US=" 125 _emit_str(s_to, 9); _emit_num(total_us) 126 let s_up: *u8 = sys_mmap(24); s_up[0]=85; s_up[1]=83; s_up[2]=95; s_up[3]=80; s_up[4]=69; s_up[5]=82; s_up[6]=95; s_up[7]=79; s_up[8]=80; s_up[9]=95; s_up[10]=78; s_up[11]=83; s_up[12]=61 // "US_PER_OP_NS=" 127 _emit_str(s_up, 13); _emit_num(us_per_op_ns) 128 _nl() 129 130 sys_exit(0); return 0 131}