nx_superopt_opt.nx source

1// nx_superopt_opt.nx -- a COST-OPTIMAL (exhaustive-shortest) superoptimizer over the 2// full integer/bit op set {add,sub,mul,shl,sar,xor,and,or}. Given a function as 3// input/output examples it returns the SHORTEST straight-line program that matches, 4// searching length 1 then 2 then 3 -- so the result is provably minimal within the 5// bound (the superoptimizer's defining property: optimality, which heuristic 6// compilers sacrifice). Verified over FULL-RANGE 64-bit inputs (sign bit + negatives), 7// not toy ranges. This is the engine for an honest 1:1 board vs gcc -O2. 8// 9// Refs: Massalin, "Superoptimizer," AFIPS/ASPLOS 1987; Warren, Hacker's Delight 2e 10// (the branchless bit-tricks); GNU superopt. Slots: 0=x, 1=0, 2=1, then computed. 11 12import "nx_syscalls.nx" 13 14const SO_ADD: i64 = 0 15const SO_SUB: i64 = 1 16const SO_MUL: i64 = 2 17const SO_SHL: i64 = 3 18const SO_SAR: i64 = 4 19const SO_XOR: i64 = 5 20const SO_AND: i64 = 6 21const SO_OR: i64 = 7 22// fused x86 lea = base + index*scale, scale in {2,4,8} (one instruction; base==index gives 23// *3/*5/*9). Used by the multiplier-chain generator + the emitter to match gcc's lea. 24const SO_LEA2: i64 = 9 25const SO_LEA4: i64 = 10 26const SO_LEA8: i64 = 11 27 28func so_is_shift(o: i64) -> i64 { if o == SO_SHL { return 1 } if o == SO_SAR { return 1 } return 0 } 29 30// shift immediates worth trying (small + the sign positions). 31func so_imm(k: i64) -> i64 { 32 if k == 0 { return 1 } 33 if k == 1 { return 2 } 34 if k == 2 { return 3 } 35 if k == 3 { return 4 } 36 if k == 4 { return 8 } 37 if k == 5 { return 16 } 38 if k == 6 { return 31 } 39 if k == 7 { return 32 } 40 return 63 41} 42const SO_NIMM: i64 = 9 43 44// s = caller scratch slot buffer (size >= L+4), reused across millions of calls 45// (allocating here would mmap per eval and is fatally slow). 46func so_eval(op: *i64, a: *i64, b: *i64, L: i64, x: i64, s: *i64) -> i64 { 47 s[0] = x; s[1] = 0; s[2] = 1 48 var t: i64 = 0 49 while t < L { 50 let o: i64 = op[t] 51 let va: i64 = s[a[t]] 52 var r: i64 = 0 53 if o == SO_ADD { r = va + s[b[t]] } 54 if o == SO_SUB { r = va - s[b[t]] } 55 if o == SO_MUL { r = va * s[b[t]] } 56 if o == SO_SHL { r = va << b[t] } 57 if o == SO_SAR { r = va >> b[t] } 58 if o == SO_XOR { r = va ^ s[b[t]] } 59 if o == SO_AND { r = va & s[b[t]] } 60 if o == SO_OR { r = va | s[b[t]] } 61 s[t + 3] = r 62 t = t + 1 63 } 64 return s[L + 2] 65} 66 67func so_matches(op: *i64, a: *i64, b: *i64, L: i64, ex_x: *i64, ex_y: *i64, nex: i64, s: *i64) -> i64 { 68 var k: i64 = 0 69 while k < nex { if so_eval(op, a, b, L, ex_x[k], s) != ex_y[k] { return 0 } k = k + 1 } 70 return 1 71} 72 73// recursive exhaustive enumeration of program positions t..L-1; nslot = slots 74// available at this position (grows by 1 per instruction). Returns 1 on a match. 75func so_enum(op: *i64, a: *i64, b: *i64, t: i64, L: i64, nslot: i64, 76 ex_x: *i64, ex_y: *i64, nex: i64, s: *i64) -> i64 { 77 if t == L { 78 if so_matches(op, a, b, L, ex_x, ex_y, nex, s) == 1 { return 1 } 79 return 0 80 } 81 var o: i64 = 0 82 while o < 8 { 83 op[t] = o 84 var av: i64 = 0 85 while av < nslot { 86 a[t] = av 87 if so_is_shift(o) == 1 { 88 var ki: i64 = 0 89 while ki < SO_NIMM { 90 b[t] = so_imm(ki) 91 if so_enum(op, a, b, t + 1, L, nslot + 1, ex_x, ex_y, nex, s) == 1 { return 1 } 92 ki = ki + 1 93 } 94 } else { 95 var bv: i64 = 0 96 while bv < nslot { 97 b[t] = bv 98 if so_enum(op, a, b, t + 1, L, nslot + 1, ex_x, ex_y, nex, s) == 1 { return 1 } 99 bv = bv + 1 100 } 101 } 102 av = av + 1 103 } 104 o = o + 1 105 } 106 return 0 107} 108 109// find the SHORTEST program (length 1..maxL) matching the examples. Returns L (op 110// count), or 0 if none. Writes the program into op/a/b. Allocates ONE scratch. 111func so_find(ex_x: *i64, ex_y: *i64, nex: i64, op: *i64, a: *i64, b: *i64, maxL: i64) -> i64 { 112 let s: *i64 = sys_mmap(8 * (maxL + 6)) 113 var L: i64 = 1 114 while L <= maxL { 115 if so_enum(op, a, b, 0, L, 3, ex_x, ex_y, nex, s) == 1 { return L } 116 L = L + 1 117 } 118 return 0 119}