nx_eqsat_congruence_bench.nx source

1// nx_eqsat_congruence_bench.nx -- NON-TOY CONGRUENCE-HEAVY race vs egg 0.11.0. 2// 3// THE RACE (engine-only, SAME expression both sides): 4// egg target = /tmp/race_egg/examples/nx_cong.rs, the congruence-heavy graph: 5// base a = (* x 8), b = (<< x 3) -- mul-pow2 merges a~b 6// N pairs la_i = (- a c_i), lb_i = (- b c_i) -- become CONGRUENT after a~b 7// root = (+ (+ ... (+ 0 (+ la_0 lb_0)) ...) (+ la_{N-1} lb_{N-1})) 8// so after the single base merge, the engine's congruence closure must propagate 9// N parent merges (la_i ~ lb_i) -- the exact reachability work egg's deferred 10// rebuild does and the pre-MEET nx_eqsat could NOT do. This is NOT the toy 11// (* x 8) single-rewrite race; congruence FIRES N times per iteration. 12// 13// OUR side, per timed iteration (the SAME end-to-end optimize pipeline egg ran, 14// no proxy, no membership/provenance machinery egg lacks): 15// nx_eqsat_init (fresh e-graph) 16// enable_meet (egg hashcons + congruence rebuild ON) 17// build a, b, N parent pairs, root (egg's RecExpr build) 18// nx_eqsat_saturate (egg's Runner::run -- rules + rebuild to fixpoint) 19// nx_eqsat_recompute_best (egg's Extractor bottom-up cost) 20// nx_eqsat_extract_best_node(root) (egg's find_best) 21// 22// FAIL LOUD: before timing, one verified iteration asserts congruence actually 23// fired -- find(la_i)==find(lb_i) for ALL i (== egg's internal congruence). A 24// fast wrong answer (congruence NOT propagated) loses, exit nonzero. 25// 26// Output (one line): 27// NX_CONG_ENGINE_ONLY N=<n> ITERS=<it> CONG_MERGES=<n> TOTAL_US=<t> US_PER_OP_NS=<u> 28// (US_PER_OP_NS = total_us*1000/iters = nanoseconds/op; divide by 1000 for us/op.) 29// 30// SOVEREIGN: no .sh; runs on the pinned NishiLang compiler. license_tier: ORIGINAL 31 32import "nx_eqsat.nx" 33const K_MAGIC_2000: i64 = 2000 34const K_MAGIC_2048: i64 = 2048 35const K_MAGIC_4096: i64 = 4096 36const K_MAGIC_8192: i64 = 8192 37 38func _emit_num(v: i64) -> i64 { 39 let b: *u8 = sys_mmap(28); var n: i64 = v; if n < 0 { n = 0 - n } 40 let t2: *u8 = sys_mmap(28); var t: i64 = 0 41 if n == 0 { t2[0] = 48; t = 1 } 42 while n > 0 { t2[t] = 48 + (n % 10); n = n / 10; t = t + 1 } 43 var i: i64 = 0; while i < t { b[i] = t2[t - 1 - i]; i = i + 1 } 44 b[t] = 32; sys_write(1, b, t + 1); return 0 45} 46func _emit_str(s: *u8, len: i64) -> i64 { sys_write(1, s, len); return 0 } 47func _nl() -> i64 { let z: *u8 = sys_mmap(2); z[0] = 10; sys_write(1, z, 1); return 0 } 48 49// Build the congruence-heavy graph in `g` and return the count of parent pairs that 50// became congruent (find(la_i)==find(lb_i)). Equals N on success. The caller hands 51// reusable scratch + a list of (la,lb) ids it allocates; we record them so the 52// caller can re-verify. Returns >=0 merges or a negative sentinel on overflow. 53// la_out/lb_out are caller arrays of length >= n (one timed iter rebuilds them). 54func _one_iter(g: *NxEGraph, 55 nodes: *NxENode, cap_nodes: i64, 56 classes: *NxEClass, cap_cls: i64, 57 hc: *HashMap, par_node: *i64, par_cls: *i64, cap_par: i64, 58 worklist: *i64, cap_work: i64, 59 n: i64, la_out: *i64, lb_out: *i64) -> i64 { 60 if nx_eqsat_init(g, nodes, cap_nodes, classes, cap_cls) != NX_EQSAT_OK { return 0 - 100 } 61 nx_hmap_clear(hc) 62 if nx_eqsat_enable_meet(g, hc, par_node, par_cls, cap_par, worklist, cap_work) != NX_EQSAT_OK { return 0 - 101 } 63 let x: i64 = nx_eqsat_add_var(g, 0) 64 let c8: i64 = nx_eqsat_add_const(g, 8) 65 let c3: i64 = nx_eqsat_add_const(g, 3) 66 let a: i64 = nx_eqsat_add_binary(g, NX_EQ_OP_MUL, x, c8) // base a = (* x 8) 67 let b: i64 = nx_eqsat_add_binary(g, NX_EQ_OP_SHL, x, c3) // base b = (<< x 3) 68 if a < 0 { return 0 - 102 } 69 if b < 0 { return 0 - 103 } 70 // Parents over a/b and the fold all use XOR -- which has NO rewrite rule on 71 // EITHER engine, so no identity/self rule re-fires on the congruence-collapsed 72 // (xor y y) shape and the saturation converges cleanly on both sides. Every 73 // la_i = (xor a c_i) still becomes CONGRUENT to lb_i = (xor b c_i) after a~b. 74 var root: i64 = 0 - 1 75 var i: i64 = 0 76 while i < n { 77 let ci: i64 = nx_eqsat_add_const(g, 100 + i) 78 let la: i64 = nx_eqsat_add_binary(g, NX_EQ_OP_XOR, a, ci) // parent over a 79 let lb: i64 = nx_eqsat_add_binary(g, NX_EQ_OP_XOR, b, ci) // parent over b 80 if la < 0 { return 0 - 104 } 81 if lb < 0 { return 0 - 105 } 82 la_out[i] = la 83 lb_out[i] = lb 84 if root < 0 { root = la } else { root = nx_eqsat_add_binary(g, NX_EQ_OP_XOR, root, la) } 85 if root < 0 { return 0 - 106 } 86 root = nx_eqsat_add_binary(g, NX_EQ_OP_XOR, root, lb) 87 if root < 0 { return 0 - 106 } 88 i = i + 1 89 } 90 // SATURATE: mul_pow2 merges a~b, then the deferred congruence rebuild propagates 91 // la_i ~ lb_i for all i (the N congruence merges egg also performs). 92 let sat: i64 = nx_eqsat_saturate(g, 64) 93 if sat != NX_EQSAT_SATURATED { return 0 - 107 } 94 // EXTRACT (egg's Extractor) -- exercise the same read-the-answer work. 95 if nx_eqsat_recompute_best(g) != NX_EQSAT_OK { return 0 - 108 } 96 let best: i64 = nx_eqsat_extract_best_node(g, root) 97 if best < 0 { return 0 - 109 } 98 // count congruence merges: la_i and lb_i in the same canonical class. 99 var merges: i64 = 0 100 var j: i64 = 0 101 while j < n { 102 if nx_eqsat_find(g, la_out[j]) == nx_eqsat_find(g, lb_out[j]) { merges = merges + 1 } 103 j = j + 1 104 } 105 return merges 106} 107 108func main() -> i64 { 109 let n: i64 = 64 // parent pairs (matches egg default N=64) 110 let iters: i64 = K_MAGIC_2000 // matches egg iters for N=64 111 112 // Generous fixed caps: per iter we add ~ 5 + 4*N nodes/classes. N=64 -> ~261. 113 let cap_nodes: i64 = K_MAGIC_2048 114 let cap_cls: i64 = K_MAGIC_2048 115 let nodes: *NxENode = sys_mmap(cap_nodes * 64) as *NxENode 116 let classes: *NxEClass = sys_mmap(cap_cls * 32) as *NxEClass 117 let g: *NxEGraph = sys_mmap(128) as *NxEGraph 118 // hashcons sized power-of-2 well above nodes/0.7. 119 let hc: *HashMap = nx_hmap_alloc(K_MAGIC_4096) 120 if (hc as i64) == 0 { sys_exit(90); return 90 } 121 let cap_par: i64 = K_MAGIC_8192 122 let par_node: *i64 = sys_mmap(cap_par * 8) as *i64 123 let par_cls: *i64 = sys_mmap(cap_par * 8) as *i64 124 let cap_work: i64 = K_MAGIC_8192 125 let worklist: *i64 = sys_mmap(cap_work * 8) as *i64 126 let la_out: *i64 = sys_mmap(n * 8) as *i64 127 let lb_out: *i64 = sys_mmap(n * 8) as *i64 128 129 // ---- CORRECTNESS GATE (FAIL LOUD before timing): congruence MUST fire N times. 130 let merges: i64 = _one_iter(g, nodes, cap_nodes, classes, cap_cls, 131 hc, par_node, par_cls, cap_par, worklist, cap_work, 132 n, la_out, lb_out) 133 if merges != n { sys_exit(80); return 80 } // every parent pair congruence-merged 134 135 // ---- TIMED LOOP ------------------------------------------------------------- 136 let t0: i64 = sys_clock_now_us() 137 var k: i64 = 0 138 while k < iters { 139 let m: i64 = _one_iter(g, nodes, cap_nodes, classes, cap_cls, 140 hc, par_node, par_cls, cap_par, worklist, cap_work, 141 n, la_out, lb_out) 142 if m != n { sys_exit(81); return 81 } // FAIL LOUD mid-run (no fast wrong answer) 143 k = k + 1 144 } 145 let t1: i64 = sys_clock_now_us() 146 let total_us: i64 = t1 - t0 147 let us_per_op_ns: i64 = (total_us * 1000) / iters 148 149 // ---- REPORT ----------------------------------------------------------------- 150 let s_tag: *u8 = sys_mmap(32) 151 s_tag[0]=78; s_tag[1]=88; s_tag[2]=95; s_tag[3]=67; s_tag[4]=79; s_tag[5]=78; s_tag[6]=71 // "NX_CONG" 152 s_tag[7]=95; s_tag[8]=69; s_tag[9]=78; s_tag[10]=71; s_tag[11]=73; s_tag[12]=78; s_tag[13]=69; s_tag[14]=95; s_tag[15]=79; s_tag[16]=78; s_tag[17]=76; s_tag[18]=89; s_tag[19]=32 // "_ENGINE_ONLY " 153 _emit_str(s_tag, 20) 154 let s_n: *u8 = sys_mmap(4); s_n[0]=78; s_n[1]=61 // "N=" 155 _emit_str(s_n, 2); _emit_num(n) 156 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=" 157 _emit_str(s_it, 6); _emit_num(iters) 158 let s_cm: *u8 = sys_mmap(16); s_cm[0]=67; s_cm[1]=79; s_cm[2]=78; s_cm[3]=71; s_cm[4]=95; s_cm[5]=77; s_cm[6]=69; s_cm[7]=82; s_cm[8]=71; s_cm[9]=69; s_cm[10]=83; s_cm[11]=61 // "CONG_MERGES=" 159 _emit_str(s_cm, 12); _emit_num(merges) 160 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=" 161 _emit_str(s_to, 9); _emit_num(total_us) 162 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=" 163 _emit_str(s_up, 13); _emit_num(us_per_op_ns) 164 _nl() 165 166 sys_exit(0); return 0 167}