jit.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581
  1. import modelverse_kernel.primitives as primitive_functions
  2. import modelverse_jit.bytecode_ir as bytecode_ir
  3. import modelverse_jit.bytecode_parser as bytecode_parser
  4. import modelverse_jit.bytecode_to_tree as bytecode_to_tree
  5. import modelverse_jit.tree_ir as tree_ir
  6. import modelverse_jit.runtime as jit_runtime
  7. import keyword
  8. # Import JitCompilationFailedException because it used to be defined
  9. # in this module.
  10. JitCompilationFailedException = jit_runtime.JitCompilationFailedException
  11. def map_and_simplify_generator(function, instruction):
  12. """Applies the given mapping function to every instruction in the tree
  13. that has the given instruction as root, and simplifies it on-the-fly.
  14. This is at least as powerful as first mapping and then simplifying, as
  15. maps and simplifications are interspersed.
  16. This function assumes that function creates a generator that returns by
  17. raising a primitive_functions.PrimitiveFinished."""
  18. # First handle the children by mapping on them and then simplifying them.
  19. new_children = []
  20. for inst in instruction.get_children():
  21. new_inst, = yield [("CALL_ARGS", [map_and_simplify_generator, (function, inst)])]
  22. new_children.append(new_inst)
  23. # Then apply the function to the top-level node.
  24. transformed, = yield [("CALL_ARGS", [function, (instruction.create(new_children),)])]
  25. # Finally, simplify the transformed top-level node.
  26. raise primitive_functions.PrimitiveFinished(transformed.simplify_node())
  27. def expand_constant_read(instruction):
  28. """Tries to replace a read of a constant node by a literal."""
  29. if isinstance(instruction, tree_ir.ReadValueInstruction) and \
  30. isinstance(instruction.node_id, tree_ir.LiteralInstruction):
  31. val, = yield [("RV", [instruction.node_id.literal])]
  32. raise primitive_functions.PrimitiveFinished(tree_ir.LiteralInstruction(val))
  33. else:
  34. raise primitive_functions.PrimitiveFinished(instruction)
  35. def optimize_tree_ir(instruction):
  36. """Optimizes an IR tree."""
  37. return map_and_simplify_generator(expand_constant_read, instruction)
  38. def create_function(
  39. function_name, parameter_list, param_dict,
  40. body_param_dict, function_body):
  41. """Creates a function from the given function name, parameter list,
  42. variable-to-parameter name map, variable-to-local name map and
  43. function body."""
  44. # Write a prologue and prepend it to the generated function body.
  45. prologue_statements = []
  46. # Create a LOCALS_NODE_NAME node, and connect it to the user root.
  47. prologue_statements.append(
  48. tree_ir.create_new_local_node(
  49. jit_runtime.LOCALS_NODE_NAME,
  50. tree_ir.LoadIndexInstruction(
  51. tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
  52. tree_ir.LiteralInstruction('user_root')),
  53. jit_runtime.LOCALS_EDGE_NAME))
  54. for (key, val) in param_dict.items():
  55. arg_ptr = tree_ir.create_new_local_node(
  56. body_param_dict[key],
  57. tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
  58. prologue_statements.append(arg_ptr)
  59. prologue_statements.append(
  60. tree_ir.CreateDictionaryEdgeInstruction(
  61. tree_ir.LoadLocalInstruction(body_param_dict[key]),
  62. tree_ir.LiteralInstruction('value'),
  63. tree_ir.LoadLocalInstruction(val)))
  64. constructed_body = tree_ir.create_block(
  65. *(prologue_statements + [function_body]))
  66. # Optimize the function's body.
  67. constructed_body, = yield [("CALL_ARGS", [optimize_tree_ir, (constructed_body,)])]
  68. # Shield temporaries from the GC.
  69. constructed_body = tree_ir.protect_temporaries_from_gc(
  70. constructed_body, tree_ir.LoadLocalInstruction(jit_runtime.LOCALS_NODE_NAME))
  71. # Wrap the IR in a function definition, give it a unique name.
  72. constructed_function = tree_ir.DefineFunctionInstruction(
  73. function_name,
  74. parameter_list + ['**' + jit_runtime.KWARGS_PARAMETER_NAME],
  75. constructed_body)
  76. raise primitive_functions.PrimitiveFinished(constructed_function)
  77. def print_value(val):
  78. """A thin wrapper around 'print'."""
  79. print(val)
  80. class ModelverseJit(object):
  81. """A high-level interface to the modelverse JIT compiler."""
  82. def __init__(self, max_instructions=None, compiled_function_lookup=None):
  83. self.todo_entry_points = set()
  84. self.no_jit_entry_points = set()
  85. self.jitted_parameters = {}
  86. self.jit_globals = {
  87. 'PrimitiveFinished' : primitive_functions.PrimitiveFinished,
  88. jit_runtime.CALL_FUNCTION_NAME : jit_runtime.call_function,
  89. jit_runtime.GET_INPUT_FUNCTION_NAME : jit_runtime.get_input
  90. }
  91. # jitted_entry_points maps body ids to values in jit_globals.
  92. self.jitted_entry_points = {}
  93. # global_functions maps global value names to body ids.
  94. self.global_functions = {}
  95. # global_functions_inv maps body ids to global value names.
  96. self.global_functions_inv = {}
  97. # bytecode_graphs maps body ids to their parsed bytecode graphs.
  98. self.bytecode_graphs = {}
  99. self.jit_count = 0
  100. self.max_instructions = max_instructions
  101. self.compiled_function_lookup = compiled_function_lookup
  102. # jit_intrinsics is a function name -> intrinsic map.
  103. self.jit_intrinsics = {}
  104. self.compilation_dependencies = {}
  105. self.jit_enabled = True
  106. self.direct_calls_allowed = True
  107. self.tracing_enabled = False
  108. self.input_function_enabled = False
  109. self.nop_insertion_enabled = True
  110. self.jit_success_log_function = None
  111. self.jit_code_log_function = None
  112. def set_jit_enabled(self, is_enabled=True):
  113. """Enables or disables the JIT."""
  114. self.jit_enabled = is_enabled
  115. def allow_direct_calls(self, is_allowed=True):
  116. """Allows or disallows direct calls from jitted to jitted code."""
  117. self.direct_calls_allowed = is_allowed
  118. def use_input_function(self, is_enabled=True):
  119. """Configures the JIT to compile 'input' instructions as function calls."""
  120. self.input_function_enabled = is_enabled
  121. def enable_tracing(self, is_enabled=True):
  122. """Enables or disables tracing for jitted code."""
  123. self.tracing_enabled = is_enabled
  124. def enable_nop_insertion(self, is_enabled=True):
  125. """Enables or disables nop insertion for jitted code. The JIT will insert nops at loop
  126. back-edges. Inserting nops sacrifices performance to keep the jitted code from
  127. blocking the thread of execution by consuming all resources; nops give the
  128. Modelverse server an opportunity to interrupt the currently running code."""
  129. self.nop_insertion_enabled = is_enabled
  130. def set_jit_success_log(self, log_function=print_value):
  131. """Configures this JIT instance with a function that prints output to a log.
  132. Success and failure messages for specific functions are then sent to said log."""
  133. self.jit_success_log_function = log_function
  134. def set_jit_code_log(self, log_function=print_value):
  135. """Configures this JIT instance with a function that prints output to a log.
  136. Function definitions of jitted functions are then sent to said log."""
  137. self.jit_code_log_function = log_function
  138. def mark_entry_point(self, body_id):
  139. """Marks the node with the given identifier as a function entry point."""
  140. if body_id not in self.no_jit_entry_points and body_id not in self.jitted_entry_points:
  141. self.todo_entry_points.add(body_id)
  142. def is_entry_point(self, body_id):
  143. """Tells if the node with the given identifier is a function entry point."""
  144. return body_id in self.todo_entry_points or \
  145. body_id in self.no_jit_entry_points or \
  146. body_id in self.jitted_entry_points
  147. def is_jittable_entry_point(self, body_id):
  148. """Tells if the node with the given identifier is a function entry point that
  149. has not been marked as non-jittable. This only returns `True` if the JIT
  150. is enabled and the function entry point has been marked jittable, or if
  151. the function has already been compiled."""
  152. return ((self.jit_enabled and body_id in self.todo_entry_points) or
  153. self.has_compiled(body_id))
  154. def has_compiled(self, body_id):
  155. """Tests if the function belonging to the given body node has been compiled yet."""
  156. return body_id in self.jitted_entry_points
  157. def get_compiled_name(self, body_id):
  158. """Gets the name of the compiled version of the given body node in the JIT
  159. global state."""
  160. return self.jitted_entry_points[body_id]
  161. def mark_no_jit(self, body_id):
  162. """Informs the JIT that the node with the given identifier is a function entry
  163. point that must never be jitted."""
  164. self.no_jit_entry_points.add(body_id)
  165. if body_id in self.todo_entry_points:
  166. self.todo_entry_points.remove(body_id)
  167. def generate_name(self, infix, suggested_name=None):
  168. """Generates a new name or picks the suggested name if it is still
  169. available."""
  170. if suggested_name is not None \
  171. and suggested_name not in self.jit_globals \
  172. and not keyword.iskeyword(suggested_name):
  173. self.jit_count += 1
  174. return suggested_name
  175. else:
  176. function_name = 'jit_%s%d' % (infix, self.jit_count)
  177. self.jit_count += 1
  178. return function_name
  179. def generate_function_name(self, body_id, suggested_name=None):
  180. """Generates a new function name or picks the suggested name if it is still
  181. available."""
  182. if suggested_name is None:
  183. suggested_name = self.get_global_name(body_id)
  184. return self.generate_name('func', suggested_name)
  185. def register_global(self, body_id, global_name):
  186. """Associates the given body id with the given global name."""
  187. self.global_functions[global_name] = body_id
  188. self.global_functions_inv[body_id] = global_name
  189. def get_global_name(self, body_id):
  190. """Gets the name of the global function with the given body id.
  191. Returns None if no known global exists with the given id."""
  192. if body_id in self.global_functions_inv:
  193. return self.global_functions_inv[body_id]
  194. else:
  195. return None
  196. def get_global_body_id(self, global_name):
  197. """Gets the body id of the global function with the given name.
  198. Returns None if no known global exists with the given name."""
  199. if global_name in self.global_functions:
  200. return self.global_functions[global_name]
  201. else:
  202. return None
  203. def register_compiled(self, body_id, compiled_function, function_name=None):
  204. """Registers a compiled entry point with the JIT."""
  205. # Get the function's name.
  206. actual_function_name = self.generate_function_name(body_id, function_name)
  207. # Map the body id to the given parameter list.
  208. self.jitted_entry_points[body_id] = actual_function_name
  209. self.jit_globals[actual_function_name] = compiled_function
  210. if function_name is not None:
  211. self.register_global(body_id, function_name)
  212. if body_id in self.todo_entry_points:
  213. self.todo_entry_points.remove(body_id)
  214. def import_value(self, value, suggested_name=None):
  215. """Imports the given value into the JIT's global scope, with the given suggested name.
  216. The actual name of the value (within the JIT's global scope) is returned."""
  217. actual_name = self.generate_name('import', suggested_name)
  218. self.jit_globals[actual_name] = value
  219. return actual_name
  220. def __lookup_compiled_body_impl(self, body_id):
  221. """Looks up a compiled function by body id. Returns a matching function,
  222. or None if no function was found."""
  223. if body_id is not None and body_id in self.jitted_entry_points:
  224. return self.jit_globals[self.jitted_entry_points[body_id]]
  225. else:
  226. return None
  227. def __lookup_external_body_impl(self, global_name, body_id):
  228. """Looks up an external function by global name. Returns a matching function,
  229. or None if no function was found."""
  230. if self.compiled_function_lookup is not None:
  231. result = self.compiled_function_lookup(global_name)
  232. if result is not None and body_id is not None:
  233. self.register_compiled(body_id, result, global_name)
  234. return result
  235. else:
  236. return None
  237. def lookup_compiled_body(self, body_id):
  238. """Looks up a compiled function by body id. Returns a matching function,
  239. or None if no function was found."""
  240. result = self.__lookup_compiled_body_impl(body_id)
  241. if result is not None:
  242. return result
  243. else:
  244. global_name = self.get_global_name(body_id)
  245. return self.__lookup_external_body_impl(global_name, body_id)
  246. def lookup_compiled_function(self, global_name):
  247. """Looks up a compiled function by global name. Returns a matching function,
  248. or None if no function was found."""
  249. body_id = self.get_global_body_id(global_name)
  250. result = self.__lookup_compiled_body_impl(body_id)
  251. if result is not None:
  252. return result
  253. else:
  254. return self.__lookup_external_body_impl(global_name, body_id)
  255. def get_intrinsic(self, name):
  256. """Tries to find an intrinsic version of the function with the
  257. given name."""
  258. if name in self.jit_intrinsics:
  259. return self.jit_intrinsics[name]
  260. else:
  261. return None
  262. def register_intrinsic(self, name, intrinsic_function):
  263. """Registers the given intrisic with the JIT. This will make the JIT replace calls to
  264. the function with the given entry point by an application of the specified function."""
  265. self.jit_intrinsics[name] = intrinsic_function
  266. def register_binary_intrinsic(self, name, operator):
  267. """Registers an intrinsic with the JIT that represents the given binary operation."""
  268. self.register_intrinsic(name, lambda a, b: tree_ir.CreateNodeWithValueInstruction(
  269. tree_ir.BinaryInstruction(
  270. tree_ir.ReadValueInstruction(a),
  271. operator,
  272. tree_ir.ReadValueInstruction(b))))
  273. def register_unary_intrinsic(self, name, operator):
  274. """Registers an intrinsic with the JIT that represents the given unary operation."""
  275. self.register_intrinsic(name, lambda a: tree_ir.CreateNodeWithValueInstruction(
  276. tree_ir.UnaryInstruction(
  277. operator,
  278. tree_ir.ReadValueInstruction(a))))
  279. def register_cast_intrinsic(self, name, target_type):
  280. """Registers an intrinsic with the JIT that represents a unary conversion operator."""
  281. self.register_intrinsic(name, lambda a: tree_ir.CreateNodeWithValueInstruction(
  282. tree_ir.CallInstruction(
  283. tree_ir.LoadGlobalInstruction(target_type.__name__),
  284. [tree_ir.ReadValueInstruction(a)])))
  285. def jit_signature(self, body_id):
  286. """Acquires the signature for the given body id node, which consists of the
  287. parameter variables, parameter name and a flag that tells if the given function
  288. is mutable."""
  289. if body_id not in self.jitted_parameters:
  290. signature_id, = yield [("RRD", [body_id, jit_runtime.FUNCTION_BODY_KEY])]
  291. signature_id = signature_id[0]
  292. param_set_id, is_mutable = yield [
  293. ("RD", [signature_id, "params"]),
  294. ("RD", [signature_id, jit_runtime.MUTABLE_FUNCTION_KEY])]
  295. if param_set_id is None:
  296. self.jitted_parameters[body_id] = ([], [], is_mutable)
  297. else:
  298. param_name_ids, = yield [("RDK", [param_set_id])]
  299. param_names = yield [("RV", [n]) for n in param_name_ids]
  300. param_vars = yield [("RD", [param_set_id, k]) for k in param_names]
  301. self.jitted_parameters[body_id] = (param_vars, param_names, is_mutable)
  302. raise primitive_functions.PrimitiveFinished(self.jitted_parameters[body_id])
  303. def jit_parse_bytecode(self, body_id):
  304. """Parses the given function body as a bytecode graph."""
  305. if body_id in self.bytecode_graphs:
  306. raise primitive_functions.PrimitiveFinished(self.bytecode_graphs[body_id])
  307. parser = bytecode_parser.BytecodeParser()
  308. result, = yield [("CALL_ARGS", [parser.parse_instruction, (body_id,)])]
  309. self.bytecode_graphs[body_id] = result
  310. raise primitive_functions.PrimitiveFinished(result)
  311. def check_jittable(self, body_id, suggested_name=None):
  312. """Checks if the function with the given body id is obviously non-jittable. If it's
  313. non-jittable, then a `JitCompilationFailedException` exception is thrown."""
  314. if body_id is None:
  315. raise ValueError('body_id cannot be None')
  316. elif body_id in self.jitted_entry_points:
  317. # We have already compiled this function.
  318. raise primitive_functions.PrimitiveFinished(
  319. self.jit_globals[self.jitted_entry_points[body_id]])
  320. elif body_id in self.no_jit_entry_points:
  321. # We're not allowed to jit this function or have tried and failed before.
  322. raise JitCompilationFailedException(
  323. 'Cannot jit function %s at %d because it is marked non-jittable.' % (
  324. '' if suggested_name is None else "'" + suggested_name + "'",
  325. body_id))
  326. elif not self.jit_enabled:
  327. # We're not allowed to jit anything.
  328. raise JitCompilationFailedException(
  329. 'Cannot jit function %s at %d because the JIT has been disabled.' % (
  330. '' if suggested_name is None else "'" + suggested_name + "'",
  331. body_id))
  332. def jit_recompile(self, user_root, body_id, function_name):
  333. """Replaces the function with the given name by compiling the bytecode at the given
  334. body id."""
  335. self.check_jittable(body_id, function_name)
  336. # Generate a name for the function we're about to analyze, and pretend that
  337. # it already exists. (we need to do this for recursive functions)
  338. self.jitted_entry_points[body_id] = function_name
  339. self.jit_globals[function_name] = None
  340. (parameter_ids, parameter_list, is_mutable), = yield [
  341. ("CALL_ARGS", [self.jit_signature, (body_id,)])]
  342. param_dict = dict(zip(parameter_ids, parameter_list))
  343. body_param_dict = dict(zip(parameter_ids, [p + "_ptr" for p in parameter_list]))
  344. dependencies = set([body_id])
  345. self.compilation_dependencies[body_id] = dependencies
  346. def handle_jit_exception(exception):
  347. # If analysis fails, then a JitCompilationFailedException will be thrown.
  348. del self.compilation_dependencies[body_id]
  349. for dep in dependencies:
  350. self.mark_no_jit(dep)
  351. if dep in self.jitted_entry_points:
  352. del self.jitted_entry_points[dep]
  353. failure_message = "%s (function '%s' at %d)" % (
  354. exception.message, function_name, body_id)
  355. if self.jit_success_log_function is not None:
  356. self.jit_success_log_function('JIT compilation failed: %s' % failure_message)
  357. raise JitCompilationFailedException(failure_message)
  358. # Try to analyze the function's body.
  359. yield [("TRY", [])]
  360. yield [("CATCH", [JitCompilationFailedException, handle_jit_exception])]
  361. if is_mutable:
  362. # We can't just JIT mutable functions. That'd be dangerous.
  363. raise JitCompilationFailedException(
  364. "Function was marked '%s'." % jit_runtime.MUTABLE_FUNCTION_KEY)
  365. body_bytecode, = yield [("CALL_ARGS", [self.jit_parse_bytecode, (body_id,)])]
  366. state = bytecode_to_tree.AnalysisState(
  367. self, body_id, user_root, body_param_dict,
  368. self.max_instructions)
  369. constructed_body, = yield [("CALL_ARGS", [state.analyze, (body_bytecode,)])]
  370. yield [("END_TRY", [])]
  371. del self.compilation_dependencies[body_id]
  372. # Wrap the tree IR in a function definition.
  373. constructed_function, = yield [
  374. ("CALL_ARGS",
  375. [create_function,
  376. (function_name, parameter_list, param_dict, body_param_dict, constructed_body)])]
  377. # Convert the function definition to Python code, and compile it.
  378. compiled_function = self.jit_define_function(function_name, constructed_function)
  379. if self.jit_success_log_function is not None:
  380. self.jit_success_log_function(
  381. "JIT compilation successful: (function '%s' at %d)" % (function_name, body_id))
  382. raise primitive_functions.PrimitiveFinished(compiled_function)
  383. def jit_define_function(self, function_name, function_def):
  384. """Converts the given tree-IR function definition to Python code, defines it,
  385. and extracts the resulting function."""
  386. # The comment below makes pylint shut up about our (hopefully benign) use of exec here.
  387. # pylint: disable=I0011,W0122
  388. # Convert the function definition to Python code, and compile it.
  389. exec(str(function_def), self.jit_globals)
  390. if self.jit_code_log_function is not None:
  391. self.jit_code_log_function(function_def)
  392. # Extract the compiled function from the JIT global state.
  393. return self.jit_globals[function_name]
  394. def jit_delete_function(self, function_name):
  395. """Deletes the function with the given function name."""
  396. del self.jit_globals[function_name]
  397. def jit_compile(self, user_root, body_id, suggested_name=None):
  398. """Tries to jit the function defined by the given entry point id and parameter list."""
  399. # Generate a name for the function we're about to analyze, and pretend that
  400. # it already exists. (we need to do this for recursive functions)
  401. function_name = self.generate_function_name(body_id, suggested_name)
  402. yield [("TAIL_CALL_ARGS", [self.jit_recompile, (user_root, body_id, function_name)])]
  403. def jit_thunk(self, get_function_body, global_name=None):
  404. """Creates a thunk from the given IR tree that computes the function's body id.
  405. This thunk is a function that will invoke the function whose body id is retrieved.
  406. The thunk's name in the JIT's global context is returned."""
  407. # The general idea is to first create a function that looks a bit like this:
  408. #
  409. # def jit_get_function_body(**kwargs):
  410. # raise primitive_functions.PrimitiveFinished(<get_function_body>)
  411. #
  412. get_function_body_name = self.generate_name('get_function_body')
  413. get_function_body_func_def, = yield [
  414. ("CALL_ARGS",
  415. [create_function,
  416. (get_function_body_name, [], {}, {}, tree_ir.ReturnInstruction(get_function_body))])]
  417. get_function_body_func = self.jit_define_function(
  418. get_function_body_name, get_function_body_func_def)
  419. # Next, we want to create a thunk that invokes said function, and then replaces itself.
  420. def __jit_thunk(**kwargs):
  421. # Compute the body id, and delete the function that computes the body id; we won't
  422. # be needing it anymore after this call.
  423. body_id = yield [("CALL_KWARGS", [get_function_body_func, kwargs])]
  424. self.jit_delete_function(get_function_body_name)
  425. # Try to associate the global name with the body id, if that's at all possible.
  426. if global_name is not None:
  427. self.register_global(body_id, global_name)
  428. compiled_function = self.lookup_compiled_body(body_id)
  429. if compiled_function is not None:
  430. # Replace this thunk by the compiled function.
  431. self.jit_globals[get_function_body_name] = compiled_function
  432. else:
  433. def __handle_jit_exception(_):
  434. # Replace this thunk by a different thunk: one that calls the interpreter
  435. # directly, without checking if the function is jittable.
  436. (_, parameter_names, _), = yield [
  437. ("CALL_ARGS", [self.jit_signature, (body_id,)])]
  438. def __interpreter_thunk(**new_kwargs):
  439. named_arg_dict = {name : new_kwargs[name] for name in parameter_names}
  440. return jit_runtime.interpret_function_body(
  441. body_id, named_arg_dict, **new_kwargs)
  442. self.jit_globals[get_function_body_name] = __interpreter_thunk
  443. yield [("TRY", [])]
  444. yield [("CATCH", [JitCompilationFailedException, __handle_jit_exception])]
  445. compiled_function, = yield [
  446. ("CALL_ARGS",
  447. [self.jit_recompile, (kwargs['user_root'], body_id, "jit_thunk")])]
  448. yield [("END_TRY", [])]
  449. # Call the compiled function.
  450. yield [("TAIL_CALL_KWARGS", [compiled_function, kwargs])]
  451. thunk_name = self.generate_name('thunk', global_name)
  452. self.jit_globals[thunk_name] = __jit_thunk
  453. raise primitive_functions.PrimitiveFinished(thunk_name)
  454. def jit_thunk_constant(self, body_id):
  455. """Creates a thunk from given body id.
  456. This thunk is a function that will invoke the function whose body id is given.
  457. The thunk's name in the JIT's global context is returned."""
  458. # We might have compiled the function with the given body id already. In that case,
  459. # we need not bother with constructing the thunk; we can return the compiled function
  460. # right away.
  461. if self.lookup_compiled_body(body_id) is not None:
  462. raise primitive_functions.PrimitiveFinished(self.get_compiled_name(body_id))
  463. # Looks like we'll just have to build that thunk after all.
  464. yield [
  465. ("TAIL_CALL_ARGS",
  466. [self.jit_thunk, (tree_ir.LiteralInstruction(body_id),)])]
  467. def jit_thunk_global(self, global_name):
  468. """Creates a thunk from given global name.
  469. This thunk is a function that will invoke the function whose body id is given.
  470. The thunk's name in the JIT's global context is returned."""
  471. # We might have compiled the function with the given name already. In that case,
  472. # we need not bother with constructing the thunk; we can return the compiled function
  473. # right away.
  474. body_id = self.get_global_body_id(global_name)
  475. if body_id is not None and self.lookup_compiled_body(body_id) is not None:
  476. raise primitive_functions.PrimitiveFinished(self.get_compiled_name(body_id))
  477. # Looks like we'll just have to build that thunk after all.
  478. # We want to look up the global function like so
  479. #
  480. # _globals, = yield [("RD", [kwargs['user_root'], "globals"])]
  481. # global_var, = yield [("RD", [_globals, global_name])]
  482. # function_id, = yield [("RD", [global_var, "value"])]
  483. # body_id, = yield [("RD", [function_id, jit_runtime.FUNCTION_BODY_KEY])]
  484. #
  485. yield [
  486. ("TAIL_CALL_ARGS",
  487. [self.jit_thunk,
  488. (tree_ir.ReadDictionaryValueInstruction(
  489. tree_ir.ReadDictionaryValueInstruction(
  490. tree_ir.ReadDictionaryValueInstruction(
  491. tree_ir.ReadDictionaryValueInstruction(
  492. tree_ir.LoadIndexInstruction(
  493. tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
  494. tree_ir.LiteralInstruction('user_root')),
  495. 'globals'),
  496. global_name),
  497. 'value'),
  498. tree_ir.LiteralInstruction(jit_runtime.FUNCTION_BODY_KEY)),
  499. global_name)])]