jit.py 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565
  1. import modelverse_kernel.primitives as primitive_functions
  2. import modelverse_jit.tree_ir as tree_ir
  3. KWARGS_PARAMETER_NAME = "remainder"
  4. """The name of the kwargs parameter in jitted functions."""
  5. class JitCompilationFailedException(Exception):
  6. """A type of exception that is raised when the jit fails to compile a function."""
  7. pass
  8. class ModelverseJit(object):
  9. """A high-level interface to the modelverse JIT compiler."""
  10. def __init__(self, max_instructions=None):
  11. self.todo_entry_points = set()
  12. self.no_jit_entry_points = set()
  13. self.jitted_entry_points = {}
  14. self.jit_globals = {
  15. 'PrimitiveFinished' : primitive_functions.PrimitiveFinished
  16. }
  17. self.jit_count = 0
  18. self.max_instructions = 30 if max_instructions is None else max_instructions
  19. def mark_entry_point(self, body_id):
  20. """Marks the node with the given identifier as a function entry point."""
  21. if body_id not in self.no_jit_entry_points and body_id not in self.jitted_entry_points:
  22. self.todo_entry_points.add(body_id)
  23. def is_entry_point(self, body_id):
  24. """Tells if the node with the given identifier is a function entry point."""
  25. return body_id in self.todo_entry_points or \
  26. body_id in self.no_jit_entry_points or \
  27. body_id in self.jitted_entry_points
  28. def is_jittable_entry_point(self, body_id):
  29. """Tells if the node with the given identifier is a function entry point that
  30. has not been marked as non-jittable."""
  31. return body_id in self.todo_entry_points or \
  32. body_id in self.jitted_entry_points
  33. def mark_no_jit(self, body_id):
  34. """Informs the JIT that the node with the given identifier is a function entry
  35. point that must never be jitted."""
  36. self.no_jit_entry_points.add(body_id)
  37. if body_id in self.todo_entry_points:
  38. self.todo_entry_points.remove(body_id)
  39. def register_compiled(self, body_id, compiled):
  40. """Registers a compiled entry point with the JIT."""
  41. self.jitted_entry_points[body_id] = compiled
  42. if body_id in self.todo_entry_points:
  43. self.todo_entry_points.remove(body_id)
  44. def try_jit(self, body_id, parameter_list):
  45. """Tries to jit the function defined by the given entry point id and parameter list."""
  46. if body_id in self.jitted_entry_points:
  47. # We have already compiled this function.
  48. raise primitive_functions.PrimitiveFinished(self.jitted_entry_points[body_id])
  49. elif body_id in self.no_jit_entry_points:
  50. # We're not allowed to jit this function or have tried and failed before.
  51. raise JitCompilationFailedException(
  52. 'Cannot jit function at %d because it is marked non-jittable.' % body_id)
  53. try:
  54. gen = AnalysisState(self.max_instructions).analyze(body_id)
  55. inp = None
  56. while True:
  57. inp = yield gen.send(inp)
  58. except primitive_functions.PrimitiveFinished as ex:
  59. constructed_body = ex.result
  60. except JitCompilationFailedException as ex:
  61. self.mark_no_jit(body_id)
  62. raise JitCompilationFailedException(
  63. '%s (function at %d)' % (ex.message, body_id))
  64. # Wrap the IR in a function definition, give it a unique name.
  65. constructed_function = tree_ir.DefineFunctionInstruction(
  66. 'jit_func%d' % self.jit_count,
  67. parameter_list + ['**' + KWARGS_PARAMETER_NAME],
  68. constructed_body.simplify())
  69. self.jit_count += 1
  70. # Convert the function definition to Python code, and compile it.
  71. exec(str(constructed_function), self.jit_globals)
  72. # Extract the compiled function from the JIT global state.
  73. compiled_function = self.jit_globals[constructed_function.name]
  74. print(constructed_function)
  75. # Save the compiled function so we can reuse it later.
  76. self.jitted_entry_points[body_id] = compiled_function
  77. raise primitive_functions.PrimitiveFinished(compiled_function)
  78. class AnalysisState(object):
  79. """The state of a bytecode analysis call graph."""
  80. def __init__(self, max_instructions=None):
  81. self.analyzed_instructions = set()
  82. self.max_instructions = max_instructions
  83. def get_local_name(self, local_id):
  84. """Gets the name for a local with the given id."""
  85. return 'local%d' % local_id
  86. def retrieve_user_root(self):
  87. """Creates an instruction that stores the user_root variable
  88. in a local."""
  89. return tree_ir.StoreLocalInstruction(
  90. 'user_root',
  91. tree_ir.LoadIndexInstruction(
  92. tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME),
  93. tree_ir.LiteralInstruction('user_root')))
  94. def analyze(self, instruction_id):
  95. """Tries to build an intermediate representation from the instruction with the
  96. given id."""
  97. # Check the analyzed_instructions set for instruction_id to avoid
  98. # infinite loops.
  99. if instruction_id in self.analyzed_instructions:
  100. raise JitCompilationFailedException('Cannot jit non-tree instruction graph.')
  101. elif (self.max_instructions is not None and
  102. len(self.analyzed_instructions) > self.max_instructions):
  103. raise JitCompilationFailedException('Maximal number of instructions exceeded.')
  104. self.analyzed_instructions.add(instruction_id)
  105. instruction_val, = yield [("RV", [instruction_id])]
  106. instruction_val = instruction_val["value"]
  107. if instruction_val in self.instruction_analyzers:
  108. gen = self.instruction_analyzers[instruction_val](self, instruction_id)
  109. try:
  110. inp = None
  111. while True:
  112. inp = yield gen.send(inp)
  113. except StopIteration:
  114. raise Exception(
  115. "Instruction analyzer (for '%s') finished without returning a value!" %
  116. (instruction_val))
  117. except primitive_functions.PrimitiveFinished as outer_e:
  118. # Check if the instruction has a 'next' instruction.
  119. next_instr, = yield [("RD", [instruction_id, "next"])]
  120. if next_instr is None:
  121. raise outer_e
  122. else:
  123. gen = self.analyze(next_instr)
  124. try:
  125. inp = None
  126. while True:
  127. inp = yield gen.send(inp)
  128. except primitive_functions.PrimitiveFinished as inner_e:
  129. raise primitive_functions.PrimitiveFinished(
  130. tree_ir.CompoundInstruction(
  131. outer_e.result,
  132. inner_e.result))
  133. else:
  134. raise JitCompilationFailedException(
  135. "Unknown instruction type: '%s'" % (instruction_val))
  136. def analyze_all(self, instruction_ids):
  137. """Tries to compile a list of IR trees from the given list of instruction ids."""
  138. results = []
  139. for inst in instruction_ids:
  140. gen = self.analyze(inst)
  141. try:
  142. inp = None
  143. while True:
  144. inp = yield gen.send(inp)
  145. except primitive_functions.PrimitiveFinished as ex:
  146. results.append(ex.result)
  147. raise primitive_functions.PrimitiveFinished(results)
  148. def analyze_return(self, instruction_id):
  149. """Tries to analyze the given 'return' instruction."""
  150. retval_id, = yield [("RD", [instruction_id, 'value'])]
  151. if retval_id is None:
  152. raise primitive_functions.PrimitiveFinished(
  153. tree_ir.ReturnInstruction(
  154. tree_ir.EmptyInstruction()))
  155. else:
  156. gen = self.analyze(retval_id)
  157. try:
  158. inp = None
  159. while True:
  160. inp = yield gen.send(inp)
  161. except primitive_functions.PrimitiveFinished as ex:
  162. raise primitive_functions.PrimitiveFinished(
  163. tree_ir.ReturnInstruction(ex.result))
  164. def analyze_if(self, instruction_id):
  165. """Tries to analyze the given 'if' instruction."""
  166. cond, true, false = yield [
  167. ("RD", [instruction_id, "cond"]),
  168. ("RD", [instruction_id, "then"]),
  169. ("RD", [instruction_id, "else"])]
  170. gen = self.analyze_all(
  171. [cond, true]
  172. if false is None
  173. else [cond, true, false])
  174. try:
  175. inp = None
  176. while True:
  177. inp = yield gen.send(inp)
  178. except primitive_functions.PrimitiveFinished as ex:
  179. if false is None:
  180. cond_r, true_r = ex.result
  181. false_r = tree_ir.EmptyInstruction()
  182. else:
  183. cond_r, true_r, false_r = ex.result
  184. raise primitive_functions.PrimitiveFinished(
  185. tree_ir.SelectInstruction(
  186. tree_ir.ReadValueInstruction(cond_r),
  187. true_r,
  188. false_r))
  189. def analyze_while(self, instruction_id):
  190. """Tries to analyze the given 'while' instruction."""
  191. cond, body = yield [
  192. ("RD", [instruction_id, "cond"]),
  193. ("RD", [instruction_id, "body"])]
  194. gen = self.analyze_all([cond, body])
  195. try:
  196. inp = None
  197. while True:
  198. inp = yield gen.send(inp)
  199. except primitive_functions.PrimitiveFinished as ex:
  200. cond_r, body_r = ex.result
  201. raise primitive_functions.PrimitiveFinished(
  202. tree_ir.LoopInstruction(
  203. tree_ir.CompoundInstruction(
  204. tree_ir.SelectInstruction(
  205. tree_ir.ReadValueInstruction(cond_r),
  206. tree_ir.EmptyInstruction(),
  207. tree_ir.BreakInstruction()),
  208. body_r)))
  209. def analyze_constant(self, instruction_id):
  210. """Tries to analyze the given 'constant' (literal) instruction."""
  211. node_id, = yield [("RD", [instruction_id, "node"])]
  212. raise primitive_functions.PrimitiveFinished(
  213. tree_ir.LiteralInstruction(node_id))
  214. def analyze_output(self, instruction_id):
  215. """Tries to analyze the given 'output' instruction."""
  216. # The plan is to basically generate this tree:
  217. #
  218. # value = <some tree>
  219. # last_output, last_output_link, new_last_output = \
  220. # yield [("RD", [user_root, "last_output"]),
  221. # ("RDE", [user_root, "last_output"]),
  222. # ("CN", []),
  223. # ]
  224. # _, _, _, _ = \
  225. # yield [("CD", [last_output, "value", value]),
  226. # ("CD", [last_output, "next", new_last_output]),
  227. # ("CD", [user_root, "last_output", new_last_output]),
  228. # ("DE", [last_output_link])
  229. # ]
  230. # yield None
  231. value_id, = yield [("RD", [instruction_id, "value"])]
  232. gen = self.analyze(value_id)
  233. try:
  234. inp = None
  235. while True:
  236. inp = yield gen.send(inp)
  237. except primitive_functions.PrimitiveFinished as ex:
  238. value_local = tree_ir.StoreLocalInstruction('value', ex.result)
  239. store_user_root = self.retrieve_user_root()
  240. last_output = tree_ir.StoreLocalInstruction(
  241. 'last_output',
  242. tree_ir.ReadDictionaryValueInstruction(
  243. store_user_root.create_load(),
  244. tree_ir.LiteralInstruction('last_output')))
  245. last_output_link = tree_ir.StoreLocalInstruction(
  246. 'last_output_link',
  247. tree_ir.ReadDictionaryEdgeInstruction(
  248. store_user_root.create_load(),
  249. tree_ir.LiteralInstruction('last_output')))
  250. new_last_output = tree_ir.StoreLocalInstruction(
  251. 'new_last_output',
  252. tree_ir.CreateNodeInstruction())
  253. result = tree_ir.create_block(
  254. value_local,
  255. store_user_root,
  256. last_output,
  257. last_output_link,
  258. new_last_output,
  259. tree_ir.CreateDictionaryEdgeInstruction(
  260. last_output.create_load(),
  261. tree_ir.LiteralInstruction('value'),
  262. value_local.create_load()),
  263. tree_ir.CreateDictionaryEdgeInstruction(
  264. last_output.create_load(),
  265. tree_ir.LiteralInstruction('next'),
  266. new_last_output.create_load()),
  267. tree_ir.CreateDictionaryEdgeInstruction(
  268. store_user_root.create_load(),
  269. tree_ir.LiteralInstruction('last_output'),
  270. new_last_output.create_load()),
  271. tree_ir.DeleteEdgeInstruction(last_output_link.create_load()),
  272. tree_ir.NopInstruction())
  273. raise primitive_functions.PrimitiveFinished(result)
  274. def analyze_input(self, _):
  275. """Tries to analyze the given 'input' instruction."""
  276. # The plan is to generate this tree:
  277. #
  278. # value = None
  279. # while True:
  280. # if value is None:
  281. # yield None # nop
  282. # else:
  283. # break
  284. #
  285. # _input = yield [("RD", [user_root, "input"])]
  286. # value = yield [("RD", [_input, "value"])]
  287. #
  288. # _next = yield [("RD", [_input, "next"])]
  289. # yield [("CD", [user_root, "input", _next])]
  290. # yield [("DN", [_input])]
  291. user_root = self.retrieve_user_root()
  292. _input = tree_ir.StoreLocalInstruction(
  293. '_input',
  294. tree_ir.ReadDictionaryValueInstruction(
  295. user_root.create_load(),
  296. tree_ir.LiteralInstruction('input')))
  297. value = tree_ir.StoreLocalInstruction(
  298. 'value',
  299. tree_ir.ReadDictionaryValueInstruction(
  300. _input.create_load(),
  301. tree_ir.LiteralInstruction('value')))
  302. raise primitive_functions.PrimitiveFinished(
  303. tree_ir.CompoundInstruction(
  304. tree_ir.create_block(
  305. user_root,
  306. value.create_store(tree_ir.LiteralInstruction(None)),
  307. tree_ir.LoopInstruction(
  308. tree_ir.create_block(
  309. tree_ir.SelectInstruction(
  310. tree_ir.BinaryInstruction(
  311. value.create_load(),
  312. 'is',
  313. tree_ir.LiteralInstruction(None)),
  314. tree_ir.NopInstruction(),
  315. tree_ir.BreakInstruction()),
  316. _input,
  317. value)),
  318. tree_ir.CreateDictionaryEdgeInstruction(
  319. user_root.create_load(),
  320. tree_ir.LiteralInstruction('input'),
  321. tree_ir.ReadDictionaryValueInstruction(
  322. _input.create_load(),
  323. tree_ir.LiteralInstruction('next'))),
  324. tree_ir.DeleteNodeInstruction(_input.create_load())),
  325. value.create_load()))
  326. def analyze_resolve(self, instruction_id):
  327. """Tries to analyze the given 'resolve' instruction."""
  328. var_id, = yield [("RD", [instruction_id, "var"])]
  329. var_name, = yield [("RV", [var_id])]
  330. # To resolve a variable, we'll do something along the
  331. # lines of:
  332. #
  333. # if 'local_var' in locals():
  334. # tmp = local_var
  335. # else:
  336. # _globals, = yield [("RD", [user_root, "globals"])]
  337. # global_var, = yield [("RD", [_globals, var_name])]
  338. #
  339. # if global_var is None:
  340. # raise Exception("Runtime error: global '%s' not found" % (var_name))
  341. #
  342. # tmp = global_var
  343. user_root = self.retrieve_user_root()
  344. global_var = tree_ir.StoreLocalInstruction(
  345. 'global_var',
  346. tree_ir.ReadDictionaryValueInstruction(
  347. tree_ir.ReadDictionaryValueInstruction(
  348. user_root.create_load(),
  349. tree_ir.LiteralInstruction('globals')),
  350. tree_ir.LiteralInstruction(var_name)))
  351. err_block = tree_ir.SelectInstruction(
  352. tree_ir.BinaryInstruction(
  353. global_var.create_load(),
  354. 'is',
  355. tree_ir.LiteralInstruction(None)),
  356. tree_ir.RaiseInstruction(
  357. tree_ir.CallInstruction(
  358. tree_ir.LoadLocalInstruction('Exception'),
  359. [tree_ir.LiteralInstruction(
  360. "Runtime error: global '%s' not found" % var_name)
  361. ])),
  362. tree_ir.EmptyInstruction())
  363. name = self.get_local_name(var_id)
  364. raise primitive_functions.PrimitiveFinished(
  365. tree_ir.SelectInstruction(
  366. tree_ir.LocalExistsInstruction(name),
  367. tree_ir.LoadLocalInstruction(name),
  368. tree_ir.CompoundInstruction(
  369. tree_ir.create_block(
  370. user_root,
  371. global_var,
  372. err_block),
  373. global_var.create_load())))
  374. def analyze_declare(self, instruction_id):
  375. """Tries to analyze the given 'declare' function."""
  376. var_id, = yield [("RD", [instruction_id, "var"])]
  377. name = self.get_local_name(var_id)
  378. # The following logic declares a local:
  379. #
  380. # if 'local_name' not in locals():
  381. # local_name, = yield [("CN", [])]
  382. raise primitive_functions.PrimitiveFinished(
  383. tree_ir.SelectInstruction(
  384. tree_ir.LocalExistsInstruction(name),
  385. tree_ir.EmptyInstruction(),
  386. tree_ir.StoreLocalInstruction(
  387. name,
  388. tree_ir.CreateNodeInstruction())))
  389. def analyze_global(self, instruction_id):
  390. """Tries to analyze the given 'global' (declaration) instruction."""
  391. var_id, = yield [("RD", [instruction_id, "var"])]
  392. var_name, = yield [("RV", [var_id])]
  393. # To resolve a variable, we'll do something along the
  394. # lines of:
  395. #
  396. # _globals, = yield [("RD", [user_root, "globals"])]
  397. # global_var = yield [("RD", [_globals, var_name])]
  398. #
  399. # if global_var is None:
  400. # global_var, = yield [("CN", [])]
  401. # yield [("CD", [_globals, var_name, global_var])]
  402. #
  403. # tmp = global_var
  404. user_root = self.retrieve_user_root()
  405. _globals = tree_ir.StoreLocalInstruction(
  406. '_globals',
  407. tree_ir.ReadDictionaryValueInstruction(
  408. user_root.create_load(),
  409. tree_ir.LiteralInstruction('globals')))
  410. global_var = tree_ir.StoreLocalInstruction(
  411. 'global_var',
  412. tree_ir.ReadDictionaryValueInstruction(
  413. _globals.create_load(),
  414. tree_ir.LiteralInstruction(var_name)))
  415. raise primitive_functions.PrimitiveFinished(
  416. tree_ir.CompoundInstruction(
  417. tree_ir.create_block(
  418. user_root,
  419. _globals,
  420. global_var,
  421. tree_ir.SelectInstruction(
  422. tree_ir.BinaryInstruction(
  423. global_var.create_load(),
  424. 'is',
  425. tree_ir.LiteralInstruction(None)),
  426. tree_ir.create_block(
  427. global_var.create_store(
  428. tree_ir.CreateNodeInstruction()),
  429. tree_ir.CreateDictionaryEdgeInstruction(
  430. _globals.create_load(),
  431. tree_ir.LiteralInstruction(var_name),
  432. global_var.create_load())),
  433. tree_ir.EmptyInstruction())),
  434. global_var.create_load()))
  435. def analyze_assign(self, instruction_id):
  436. """Tries to analyze the given 'assign' instruction."""
  437. var_id, value_id = yield [("RD", [instruction_id, "var"]),
  438. ("RD", [instruction_id, "value"])]
  439. try:
  440. gen = self.analyze_all([var_id, value_id])
  441. inp = None
  442. while True:
  443. inp = yield gen.send(inp)
  444. except primitive_functions.PrimitiveFinished as ex:
  445. var_r, value_r = ex.result
  446. # Assignments work like this:
  447. #
  448. # value_link = yield [("RDE", [variable, "value"])]
  449. # _, _ = yield [("CD", [variable, "value", value]),
  450. # ("DE", [value_link])]
  451. variable = tree_ir.StoreLocalInstruction('variable', var_r)
  452. value = tree_ir.StoreLocalInstruction('value', value_r)
  453. value_link = tree_ir.StoreLocalInstruction(
  454. 'value_link',
  455. tree_ir.ReadDictionaryEdgeInstruction(
  456. variable.create_load(),
  457. tree_ir.LiteralInstruction('value')))
  458. raise primitive_functions.PrimitiveFinished(
  459. tree_ir.create_block(
  460. variable,
  461. value,
  462. value_link,
  463. tree_ir.CreateDictionaryEdgeInstruction(
  464. variable.create_load(),
  465. tree_ir.LiteralInstruction('value'),
  466. value.create_load()),
  467. tree_ir.DeleteEdgeInstruction(
  468. value_link.create_load())))
  469. def analyze_access(self, instruction_id):
  470. """Tries to analyze the given 'access' instruction."""
  471. var_id, = yield [("RD", [instruction_id, "var"])]
  472. try:
  473. gen = self.analyze(var_id)
  474. inp = None
  475. while True:
  476. inp = yield gen.send(inp)
  477. except primitive_functions.PrimitiveFinished as ex:
  478. var_r = ex.result
  479. # Accessing a variable is pretty easy. It really just boils
  480. # down to reading the value corresponding to the 'value' key
  481. # of the variable.
  482. #
  483. # value, = yield [("RD", [returnvalue, "value"])]
  484. raise primitive_functions.PrimitiveFinished(
  485. tree_ir.ReadDictionaryValueInstruction(
  486. var_r,
  487. tree_ir.LiteralInstruction('value')))
  488. instruction_analyzers = {
  489. 'if' : analyze_if,
  490. 'while' : analyze_while,
  491. 'return' : analyze_return,
  492. 'constant' : analyze_constant,
  493. 'resolve' : analyze_resolve,
  494. 'declare' : analyze_declare,
  495. 'global' : analyze_global,
  496. 'assign' : analyze_assign,
  497. 'access' : analyze_access,
  498. 'output' : analyze_output,
  499. 'input' : analyze_input
  500. }