nx_superopt_strength.nx source

1// nx_superopt_strength.nx -- a GENERATOR: it AUTHORS strength-reduction lowerings 2// for `x * C` by SEARCH, not by me hand-writing each rule. Given a constant C, it 3// enumerates candidate implementations -- a single shift (x<<a), or a shift-op pair 4// ((x<<a) + (x<<b) / (x<<a) - (x<<b)), or the fallback multiply -- VERIFIES each is 5// 1:1 equal to x*C over many inputs, scores each by the measured cost model 6// (mul > shift+add, per nx_measure_cost), and RETURNS the cheapest verified one. 7// 8// THIS is a building capability: the team discovers `x*7 = (x<<3)-x`, `x*9 = 9// (x<<3)+x`, `x*1024 = x<<10` ITSELF -- I wrote the searcher; the search wrote the 10// answers. It is the Generator leg the codegen loop consumes (each authored 11// lowering is a verified, cheaper codegen rewrite). Refs: Warren, Hacker's Delight 12// 2e ch.8 (strength reduction); the e-graph superopt idea (egg, Willsey POPL 2021). 13 14import "nx_syscalls.nx" 15const SO_MAGIC_19088743: i64 = 19088743 16const SO_MAGIC_6364136223846793005: i64 = 6364136223846793005 17const SO_MAGIC_1442695040888963407: i64 = 1442695040888963407 18 19const SO_PLAN_MUL: i64 = 0 20const SO_PLAN_SHL: i64 = 1 21const SO_PLAN_SHIFTOP: i64 = 2 22const SO_ADD: i64 = 0 23const SO_SUB: i64 = 1 24 25// cost model (measured ordering: mul costs several adds; a shift ~ an add). 26const SO_MUL_COST: i64 = 4 27const SO_SHL_COST: i64 = 1 28const SO_OP_COST: i64 = 1 29const SO_SHIFT_MAX: i64 = 62 30 31// evaluate a candidate plan on input x. 32func so_eval(kind: i64, a: i64, op: i64, b: i64, C: i64, x: i64) -> i64 { 33 if kind == SO_PLAN_MUL { return x * C } 34 if kind == SO_PLAN_SHL { return x << a } 35 let ta: i64 = x << a 36 let tb: i64 = x << b 37 if op == SO_ADD { return ta + tb } 38 return ta - tb 39} 40 41// cost of a plan (lower = better). 42func so_cost(kind: i64, a: i64, b: i64) -> i64 { 43 if kind == SO_PLAN_MUL { return SO_MUL_COST } 44 if kind == SO_PLAN_SHL { if a == 0 { return 0 } return SO_SHL_COST } 45 var c: i64 = SO_OP_COST 46 if a > 0 { c = c + SO_SHL_COST } 47 if b > 0 { c = c + SO_SHL_COST } 48 return c 49} 50 51// VERIFY a candidate computes x*C for all test inputs (1:1). 0 = no, 1 = yes. 52func so_verify(kind: i64, a: i64, op: i64, b: i64, C: i64) -> i64 { 53 if so_eval(kind, a, op, b, C, 0) != 0 { return 0 } 54 if so_eval(kind, a, op, b, C, 1) != C { return 0 } 55 var s: i64 = SO_MAGIC_19088743 56 var t: i64 = 0 57 while t < 32 { 58 s = s * SO_MAGIC_6364136223846793005 + SO_MAGIC_1442695040888963407 59 let x: i64 = s 60 if so_eval(kind, a, op, b, C, x) != x * C { return 0 } 61 t = t + 1 62 } 63 return 1 64} 65 66// THE SEARCH: author the cheapest verified lowering of x*C. Writes the winning 67// plan into out[0..3] = {kind, a, op, b}; returns its cost. 68func so_best(C: i64, out: *i64) -> i64 { 69 var best_cost: i64 = SO_MUL_COST 70 out[0] = SO_PLAN_MUL; out[1] = 0; out[2] = 0; out[3] = 0 // multiply always works 71 // single shift 72 var a: i64 = 0 73 while a < SO_SHIFT_MAX { 74 if so_verify(SO_PLAN_SHL, a, 0, 0, C) == 1 { 75 let c: i64 = so_cost(SO_PLAN_SHL, a, 0) 76 if c < best_cost { best_cost = c; out[0] = SO_PLAN_SHL; out[1] = a; out[2] = 0; out[3] = 0 } 77 } 78 a = a + 1 79 } 80 // (x<<a) op (x<<b) 81 a = 0 82 while a < SO_SHIFT_MAX { 83 var b: i64 = 0 84 while b < SO_SHIFT_MAX { 85 var op: i64 = 0 86 while op < 2 { 87 if so_verify(SO_PLAN_SHIFTOP, a, op, b, C) == 1 { 88 let c: i64 = so_cost(SO_PLAN_SHIFTOP, a, b) 89 if c < best_cost { best_cost = c; out[0] = SO_PLAN_SHIFTOP; out[1] = a; out[2] = op; out[3] = b } 90 } 91 op = op + 1 92 } 93 b = b + 1 94 } 95 a = a + 1 96 } 97 return best_cost 98}