jit.py 29 KB

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