michael@0: # vim: set ts=8 sts=4 et sw=4 tw=99:
michael@0: # This Source Code Form is subject to the terms of the Mozilla Public
michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/.
michael@0: 
michael@0: import os, re
michael@0: import tempfile
michael@0: import subprocess
michael@0: import sys, math
michael@0: import datetime
michael@0: import random
michael@0: 
michael@0: def realpath(k):
michael@0:     return os.path.realpath(os.path.normpath(k))
michael@0: 
michael@0: class UCTNode:
michael@0:     def __init__(self, loop):
michael@0:         self.children = None
michael@0:         self.loop = loop
michael@0:         self.visits = 1
michael@0:         self.score = 0
michael@0: 
michael@0:     def addChild(self, child):
michael@0:         if self.children == None:
michael@0:             self.children = []
michael@0:         self.children.append(child)
michael@0: 
michael@0:     def computeUCB(self, coeff):
michael@0:         return (self.score / self.visits) + math.sqrt(coeff / self.visits)
michael@0: 
michael@0: class UCT:
michael@0:     def __init__(self, benchmark, bestTime, enableLoops, loops, fd, playouts):
michael@0:         self.bm = benchmark
michael@0:         self.fd = fd
michael@0:         self.numPlayouts = playouts
michael@0:         self.maxNodes = self.numPlayouts * 20
michael@0:         self.loops = loops
michael@0:         self.enableLoops = enableLoops
michael@0:         self.maturityThreshold = 20
michael@0:         self.originalBest = bestTime
michael@0:         self.bestTime = bestTime
michael@0:         self.bias = 20
michael@0:         self.combos = []
michael@0:         self.zobrist = { }
michael@0:         random.seed()
michael@0: 
michael@0:     def expandNode(self, node, pending):
michael@0:         for loop in pending:
michael@0:             node.addChild(UCTNode(loop))
michael@0:             self.numNodes += 1
michael@0:             if self.numNodes >= self.maxNodes:
michael@0:                 return False
michael@0:         return True
michael@0: 
michael@0:     def findBestChild(self, node):
michael@0:         coeff = self.bias * math.log(node.visits)
michael@0:         bestChild = None
michael@0:         bestUCB = -float('Infinity')
michael@0: 
michael@0:         for child in node.children:
michael@0:             ucb = child.computeUCB(coeff)
michael@0:             if ucb >= bestUCB:
michael@0:                 bestUCB = ucb
michael@0:                 bestChild = child
michael@0: 
michael@0:         return child
michael@0: 
michael@0:     def playout(self, history):
michael@0:         queue = []
michael@0:         for i in range(0, len(self.loops)):
michael@0:             queue.append(random.randint(0, 1))
michael@0:         for node in history:
michael@0:             queue[node.loop] = not self.enableLoops
michael@0:         zash = 0
michael@0:         for i in range(0, len(queue)):
michael@0:             if queue[i]:
michael@0:                 zash |= (1 << i)
michael@0:         if zash in self.zobrist:
michael@0:             return self.zobrist[zash]
michael@0: 
michael@0:         self.bm.generateBanList(self.loops, queue)
michael@0:         result = self.bm.treeSearchRun(self.fd, ['-m', '-j'], 3)
michael@0:         self.zobrist[zash] = result
michael@0:         return result
michael@0: 
michael@0:     def step(self, loopList):
michael@0:         node = self.root
michael@0:         pending = loopList[:]
michael@0:         history = [node]
michael@0: 
michael@0:         while True:
michael@0:             # If this is a leaf node...
michael@0:             if node.children == None:
michael@0:                 # And the leaf node is mature...
michael@0:                 if node.visits >= self.maturityThreshold:
michael@0:                     # If the node can be expanded, keep spinning.
michael@0:                     if self.expandNode(node, pending) and node.children != None:
michael@0:                         continue
michael@0: 
michael@0:                 # Otherwise, this is a leaf node. Run a playout.
michael@0:                 score = self.playout(history)
michael@0:                 break
michael@0: 
michael@0:             # Find the best child.
michael@0:             node = self.findBestChild(node)
michael@0:             history.append(node)
michael@0:             pending.remove(node.loop)
michael@0: 
michael@0:         # Normalize the score.
michael@0:         origScore = score
michael@0:         score = (self.originalBest - score) / self.originalBest
michael@0: 
michael@0:         for node in history:
michael@0:             node.visits += 1
michael@0:             node.score += score
michael@0: 
michael@0:         if int(origScore) < int(self.bestTime):
michael@0:             print('New best score: {0:f}ms'.format(origScore))
michael@0:             self.combos = [history]
michael@0:             self.bestTime = origScore
michael@0:         elif int(origScore) == int(self.bestTime):
michael@0:             self.combos.append(history)
michael@0: 
michael@0:     def run(self):
michael@0:         loopList = [i for i in range(0, len(self.loops))]
michael@0:         self.numNodes = 1
michael@0:         self.root = UCTNode(-1)
michael@0:         self.expandNode(self.root, loopList)
michael@0: 
michael@0:         for i in range(0, self.numPlayouts):
michael@0:             self.step(loopList)
michael@0: 
michael@0:         # Build the expected combination vector.
michael@0:         combos = [ ]
michael@0:         for combo in self.combos:
michael@0:             vec = [ ]
michael@0:             for i in range(0, len(self.loops)):
michael@0:                 vec.append(int(self.enableLoops))
michael@0:             for node in combo:
michael@0:                 vec[node.loop] = int(not self.enableLoops)
michael@0:             combos.append(vec)
michael@0: 
michael@0:         return [self.bestTime, combos]
michael@0: 
michael@0: class Benchmark:
michael@0:     def __init__(self, JS, fname):
michael@0:         self.fname = fname
michael@0:         self.JS = JS
michael@0:         self.stats = { }
michael@0:         self.runList = [ ]
michael@0: 
michael@0:     def run(self, fd, eargs):
michael@0:         args = [self.JS]
michael@0:         args.extend(eargs)
michael@0:         args.append(fd.name)
michael@0:         return subprocess.check_output(args).decode()
michael@0: 
michael@0:     #    self.stats[name] = { }
michael@0:     #    self.runList.append(name)
michael@0:     #    for line in output.split('\n'):
michael@0:     #        m = re.search('line (\d+): (\d+)', line)
michael@0:     #        if m:
michael@0:     #            self.stats[name][int(m.group(1))] = int(m.group(2))
michael@0:     #        else:
michael@0:     #            m = re.search('total: (\d+)', line)
michael@0:     #            if m:
michael@0:     #                self.stats[name]['total'] = m.group(1)
michael@0: 
michael@0:     def winnerForLine(self, line):
michael@0:         best = self.runList[0]
michael@0:         bestTime = self.stats[best][line]
michael@0:         for run in self.runList[1:]:
michael@0:             x = self.stats[run][line]
michael@0:             if x < bestTime:
michael@0:                 best = run
michael@0:                 bestTime = x
michael@0:         return best
michael@0: 
michael@0:     def chart(self):
michael@0:         sys.stdout.write('{0:7s}'.format(''))
michael@0:         sys.stdout.write('{0:15s}'.format('line'))
michael@0:         for run in self.runList:
michael@0:             sys.stdout.write('{0:15s}'.format(run))
michael@0:         sys.stdout.write('{0:15s}\n'.format('best'))
michael@0:         for c in self.counters:
michael@0:             sys.stdout.write('{0:10d}'.format(c))
michael@0:             for run in self.runList:
michael@0:                 sys.stdout.write('{0:15d}'.format(self.stats[run][c]))
michael@0:             sys.stdout.write('{0:12s}'.format(''))
michael@0:             sys.stdout.write('{0:15s}'.format(self.winnerForLine(c)))
michael@0:             sys.stdout.write('\n')
michael@0: 
michael@0:     def preprocess(self, lines, onBegin, onEnd):
michael@0:         stack = []
michael@0:         counters = []
michael@0:         rd = open(self.fname, 'rt')
michael@0:         for line in rd:
michael@0:             if re.search('\/\* BEGIN LOOP \*\/', line):
michael@0:                 stack.append([len(lines), len(counters)])
michael@0:                 counters.append([len(lines), 0])
michael@0:                 onBegin(lines, len(lines))
michael@0:             elif re.search('\/\* END LOOP \*\/', line):
michael@0:                 old = stack.pop()
michael@0:                 onEnd(lines, old[0], len(lines))
michael@0:                 counters[old[1]][1] = len(lines)
michael@0:             else:
michael@0:                 lines.append(line)
michael@0:         return [lines, counters]
michael@0: 
michael@0:     def treeSearchRun(self, fd, args, count = 5):
michael@0:         total = 0
michael@0:         for i in range(0, count):
michael@0:             output = self.run(fd, args)
michael@0:             total += int(output)
michael@0:         return total / count
michael@0: 
michael@0:     def generateBanList(self, counters, queue):
michael@0:         if os.path.exists('/tmp/permabans'):
michael@0:             os.unlink('/tmp/permabans')
michael@0:         fd = open('/tmp/permabans', 'wt')
michael@0:         for i in range(0, len(counters)):
michael@0:             for j in range(counters[i][0], counters[i][1] + 1):
michael@0:                 fd.write('{0:d} {1:d}\n'.format(j, int(queue[i])))
michael@0:         fd.close()
michael@0: 
michael@0:     def internalExhaustiveSearch(self, params):
michael@0:         counters = params['counters']
michael@0: 
michael@0:         # iterative algorithm to explore every combination
michael@0:         ncombos = 2 ** len(counters)
michael@0:         queue = []
michael@0:         for c in counters:
michael@0:             queue.append(0)
michael@0: 
michael@0:         fd = params['fd']
michael@0:         bestTime = float('Infinity')
michael@0:         bestCombos = []
michael@0: 
michael@0:         i = 0
michael@0:         while i < ncombos:
michael@0:             temp = i
michael@0:             for j in range(0, len(counters)):
michael@0:                 queue[j] = temp & 1
michael@0:                 temp = temp >> 1
michael@0:             self.generateBanList(counters, queue)
michael@0: 
michael@0:             t = self.treeSearchRun(fd, ['-m', '-j'])
michael@0:             if (t < bestTime):
michael@0:                 bestTime = t
michael@0:                 bestCombos = [queue[:]]
michael@0:                 print('New best time: {0:f}ms'.format(t))
michael@0:             elif int(t) == int(bestTime):
michael@0:                 bestCombos.append(queue[:])
michael@0: 
michael@0:             i = i + 1
michael@0: 
michael@0:         return [bestTime, bestCombos]
michael@0: 
michael@0:     def internalTreeSearch(self, params):
michael@0:         fd = params['fd']
michael@0:         methodTime = params['methodTime']
michael@0:         tracerTime = params['tracerTime']
michael@0:         combinedTime = params['combinedTime']
michael@0:         counters = params['counters']
michael@0: 
michael@0:         # Build the initial loop data.
michael@0:         # If the method JIT already wins, disable tracing by default.
michael@0:         # Otherwise, enable tracing by default.
michael@0:         if methodTime < combinedTime:
michael@0:             enableLoops = True
michael@0:         else:
michael@0:             enableLoops = False
michael@0: 
michael@0:         enableLoops = False
michael@0: 
michael@0:         uct = UCT(self, combinedTime, enableLoops, counters[:], fd, 50000)
michael@0:         return uct.run()
michael@0: 
michael@0:     def treeSearch(self):
michael@0:         fd, counters = self.ppForTreeSearch()
michael@0: 
michael@0:         os.system("cat " + fd.name + " > /tmp/k.js")
michael@0: 
michael@0:         if os.path.exists('/tmp/permabans'):
michael@0:             os.unlink('/tmp/permabans')
michael@0:         methodTime = self.treeSearchRun(fd, ['-m'])
michael@0:         tracerTime = self.treeSearchRun(fd, ['-j'])
michael@0:         combinedTime = self.treeSearchRun(fd, ['-m', '-j'])
michael@0: 
michael@0:         #Get a rough estimate of how long this benchmark will take to fully compute.
michael@0:         upperBound = max(methodTime, tracerTime, combinedTime)
michael@0:         upperBound *= 2 ** len(counters)
michael@0:         upperBound *= 5    # Number of runs
michael@0:         treeSearch = False
michael@0:         if (upperBound < 1000):
michael@0:             print('Estimating {0:d}ms to test, so picking exhaustive '.format(int(upperBound)) +
michael@0:                   'search.')
michael@0:         else:
michael@0:             upperBound = int(upperBound / 1000)
michael@0:             delta = datetime.timedelta(seconds = upperBound)
michael@0:             if upperBound < 180:
michael@0:                 print('Estimating {0:d}s to test, so picking exhaustive '.format(int(upperBound)))
michael@0:             else:
michael@0:                 print('Estimating {0:s} to test, so picking tree search '.format(str(delta)))
michael@0:                 treeSearch = True
michael@0: 
michael@0:         best = min(methodTime, tracerTime, combinedTime)
michael@0: 
michael@0:         params = {
michael@0:                     'fd': fd,
michael@0:                     'counters': counters,
michael@0:                     'methodTime': methodTime,
michael@0:                     'tracerTime': tracerTime,
michael@0:                     'combinedTime': combinedTime
michael@0:                  }
michael@0: 
michael@0:         print('Method JIT:  {0:d}ms'.format(int(methodTime)))
michael@0:         print('Tracing JIT: {0:d}ms'.format(int(tracerTime)))
michael@0:         print('Combined:    {0:d}ms'.format(int(combinedTime)))
michael@0: 
michael@0:         if 1 and treeSearch:
michael@0:             results = self.internalTreeSearch(params)
michael@0:         else:
michael@0:             results = self.internalExhaustiveSearch(params)
michael@0: 
michael@0:         bestTime = results[0]
michael@0:         bestCombos = results[1]
michael@0:         print('Search found winning time {0:d}ms!'.format(int(bestTime)))
michael@0:         print('Combos at this time: {0:d}'.format(len(bestCombos)))
michael@0: 
michael@0:         #Find loops that traced every single time
michael@0:         for i in range(0, len(counters)):
michael@0:             start = counters[i][0]
michael@0:             end = counters[i][1]
michael@0:             n = len(bestCombos)
michael@0:             for j in bestCombos:
michael@0:                 n -= j[i]
michael@0:             print('\tloop @ {0:d}-{1:d} traced {2:d}% of the time'.format(
michael@0:                     start, end, int(n / len(bestCombos) * 100)))
michael@0: 
michael@0:     def ppForTreeSearch(self):
michael@0:         def onBegin(lines, lineno):
michael@0:             lines.append('GLOBAL_THINGY = 1;\n')
michael@0:         def onEnd(lines, old, lineno):
michael@0:             lines.append('GLOBAL_THINGY = 1;\n')
michael@0: 
michael@0:         lines = ['var JINT_START_TIME = Date.now();\n',
michael@0:                  'var GLOBAL_THINGY = 0;\n']
michael@0: 
michael@0:         lines, counters = self.preprocess(lines, onBegin, onEnd)
michael@0:         fd = tempfile.NamedTemporaryFile('wt')
michael@0:         for line in lines:
michael@0:             fd.write(line)
michael@0:         fd.write('print(Date.now() - JINT_START_TIME);\n')
michael@0:         fd.flush()
michael@0:         return [fd, counters]
michael@0: 
michael@0:     def preprocessForLoopCounting(self):
michael@0:         def onBegin(lines, lineno):
michael@0:             lines.append('JINT_TRACKER.line_' + str(lineno) + '_start = Date.now();\n')
michael@0: 
michael@0:         def onEnd(lines, old, lineno):
michael@0:             lines.append('JINT_TRACKER.line_' + str(old) + '_end = Date.now();\n')
michael@0:             lines.append('JINT_TRACKER.line_' + str(old) + '_total += ' + \
michael@0:                          'JINT_TRACKER.line_' + str(old) + '_end - ' + \
michael@0:                          'JINT_TRACKER.line_' + str(old) + '_start;\n')
michael@0: 
michael@0:         lines, counters = self.preprocess(onBegin, onEnd)
michael@0:         fd = tempfile.NamedTemporaryFile('wt')
michael@0:         fd.write('var JINT_TRACKER = { };\n')
michael@0:         for c in counters:
michael@0:             fd.write('JINT_TRACKER.line_' + str(c) + '_start = 0;\n')
michael@0:             fd.write('JINT_TRACKER.line_' + str(c) + '_end = 0;\n')
michael@0:             fd.write('JINT_TRACKER.line_' + str(c) + '_total = 0;\n')
michael@0:         fd.write('JINT_TRACKER.begin = Date.now();\n')
michael@0:         for line in lines:
michael@0:             fd.write(line)
michael@0:         fd.write('JINT_TRACKER.total = Date.now() - JINT_TRACKER.begin;\n')
michael@0:         for c in self.counters:
michael@0:             fd.write('print("line ' + str(c) + ': " + JINT_TRACKER.line_' + str(c) +
michael@0:                            '_total);')
michael@0:         fd.write('print("total: " + JINT_TRACKER.total);')
michael@0:         fd.flush()
michael@0:         return fd
michael@0: 
michael@0: if __name__ == '__main__':
michael@0:     script_path = os.path.abspath(__file__)
michael@0:     script_dir = os.path.dirname(script_path)
michael@0:     test_dir = os.path.join(script_dir, 'tests')
michael@0:     lib_dir = os.path.join(script_dir, 'lib')
michael@0: 
michael@0:     # The [TESTS] optional arguments are paths of test files relative
michael@0:     # to the jit-test/tests directory.
michael@0: 
michael@0:     from optparse import OptionParser
michael@0:     op = OptionParser(usage='%prog [options] JS_SHELL test')
michael@0:     (OPTIONS, args) = op.parse_args()
michael@0:     if len(args) < 2:
michael@0:         op.error('missing JS_SHELL and test argument')
michael@0:     # We need to make sure we are using backslashes on Windows.
michael@0:     JS = realpath(args[0])
michael@0:     test = realpath(args[1])
michael@0: 
michael@0:     bm = Benchmark(JS, test)
michael@0:     bm.treeSearch()
michael@0:     # bm.preprocess()
michael@0:     # bm.run('mjit', ['-m'])
michael@0:     # bm.run('tjit', ['-j'])
michael@0:     # bm.run('m+tjit', ['-m', '-j'])
michael@0:     # bm.chart()
michael@0: