A reimplementation of STOKE in haskell.
STOKE is a superoptimizer that finds optimal instruction sequences to perform a particular task by using MCMC methods to explore the large search space.
What we do is to "MCMC-sample the space of possible programs, with a scoring function based on program equivalence and performance". Broken down, the steps are:
- Start with the original program
p - Perturb it slightly to get a program
q(add an instruction, remove some instructions, change an instruction) - Assign a score to
qby sendingpandqrandom inputs, and checking how often they answer the same. - If they answered the same for all random inputs, ask an SMT solver
nicely if
pandqare equivalent. If she's in a good mood and the universe is kind, the answer is yes. - Now, score
qbased on the factors of:- Are
pandqsemantically equivalent? - On how many inputs
pandqhad the same answer? - is
qfaster thanp?
- Are
- Now, either pick
qto be the newpor stay atpdepending onq's score. - Repeat ~10,000 times.
- Order best
qs visited.
The results are quite impressive: Even with code as naive as what I've written, we are able to regain peephole optimisations that compilers perform essentially "for free".
*** original: (nparams: 0 | [IPush 2,IPush 3,IAdd])***
[IPush 5] | score: 2.5 // constant folding
[IPush 2,IPush 3,IAdd] | score: 2.125
[IPush 2,IPush 3,ISwap,IAdd] | score: 2.0625
*** original: (nparams: 1 | [IPush 2,IMul])***
[IDup,IAdd] | score: 2.25 // strength reduction: 2 * x -> x + x
[IDup,ISwap,IAdd] | score: 2.125
[IDup,ISwap,ISwap,IAdd] | score: 2.0625
[IDup,ISwap,ISwap,ISwap,IAdd] | score: 2.03125
*** original: (nparams: 1 | progInsts = [IDup,IAnd])***
[] | score: 3.0 // constant folding: x & x == x
[IDup,IAnd] | score: 2.25
[IDup,ISwap,IAnd] | score: 2.125