jit.py 41 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968
  1. import modelverse_kernel.primitives as primitive_functions
  2. import modelverse_jit.tree_ir as tree_ir
  3. import modelverse_jit.runtime as jit_runtime
  4. KWARGS_PARAMETER_NAME = "kwargs"
  5. """The name of the kwargs parameter in jitted functions."""
  6. INTERPRET_FUNCTION_NAME = "__interpret_function"
  7. """The name of the '__interpret_function' function, in the jitted function scope."""
  8. def get_parameter_names(compiled_function):
  9. """Gets the given compiled function's parameter names."""
  10. if hasattr(compiled_function, '__code__'):
  11. return compiled_function.__code__.co_varnames[
  12. :compiled_function.__code__.co_argcount]
  13. elif hasattr(compiled_function, '__init__'):
  14. return get_parameter_names(compiled_function.__init__)[1:]
  15. else:
  16. raise ValueError("'compiled_function' must be a function or a type.")
  17. def apply_intrinsic(intrinsic_function, named_args):
  18. """Applies the given intrinsic to the given sequence of named arguments."""
  19. param_names = get_parameter_names(intrinsic_function)
  20. if tuple(param_names) == tuple([n for n, _ in named_args]):
  21. # Perfect match. Yay!
  22. return intrinsic_function(**dict(named_args))
  23. else:
  24. # We'll have to store the arguments into locals to preserve
  25. # the order of evaluation.
  26. stored_args = [(name, tree_ir.StoreLocalInstruction(None, arg)) for name, arg in named_args]
  27. arg_value_dict = dict([(name, arg.create_load()) for name, arg in stored_args])
  28. store_instructions = [instruction for _, instruction in stored_args]
  29. return tree_ir.CompoundInstruction(
  30. tree_ir.create_block(*store_instructions),
  31. intrinsic_function(**arg_value_dict))
  32. class JitCompilationFailedException(Exception):
  33. """A type of exception that is raised when the jit fails to compile a function."""
  34. pass
  35. class ModelverseJit(object):
  36. """A high-level interface to the modelverse JIT compiler."""
  37. def __init__(self, max_instructions=50, compiled_function_lookup=None):
  38. self.todo_entry_points = set()
  39. self.no_jit_entry_points = set()
  40. self.jitted_entry_points = {}
  41. self.jitted_parameters = {}
  42. self.jit_globals = {
  43. 'PrimitiveFinished' : primitive_functions.PrimitiveFinished,
  44. INTERPRET_FUNCTION_NAME : jit_runtime.interpret_function
  45. }
  46. self.jit_count = 0
  47. self.max_instructions = max_instructions
  48. self.compiled_function_lookup = compiled_function_lookup
  49. # jit_intrinsics is a function name -> intrinsic map.
  50. self.jit_intrinsics = {}
  51. self.compilation_dependencies = {}
  52. self.jit_enabled = True
  53. self.direct_calls_allowed = True
  54. self.tracing_enabled = False
  55. def set_jit_enabled(self, is_enabled=True):
  56. """Enables or disables the JIT."""
  57. self.jit_enabled = is_enabled
  58. def allow_direct_calls(self, is_allowed=True):
  59. """Allows or disallows direct calls from jitted to jitted code."""
  60. self.direct_calls_allowed = is_allowed
  61. def enable_tracing(self, is_enabled=True):
  62. """Enables or disables tracing for jitted code."""
  63. self.tracing_enabled = is_enabled
  64. def mark_entry_point(self, body_id):
  65. """Marks the node with the given identifier as a function entry point."""
  66. if body_id not in self.no_jit_entry_points and body_id not in self.jitted_entry_points:
  67. self.todo_entry_points.add(body_id)
  68. def is_entry_point(self, body_id):
  69. """Tells if the node with the given identifier is a function entry point."""
  70. return body_id in self.todo_entry_points or \
  71. body_id in self.no_jit_entry_points or \
  72. body_id in self.jitted_entry_points
  73. def is_jittable_entry_point(self, body_id):
  74. """Tells if the node with the given identifier is a function entry point that
  75. has not been marked as non-jittable. This only returns `True` if the JIT
  76. is enabled and the function entry point has been marked jittable, or if
  77. the function has already been compiled."""
  78. return ((self.jit_enabled and body_id in self.todo_entry_points) or
  79. self.has_compiled(body_id))
  80. def has_compiled(self, body_id):
  81. """Tests if the function belonging to the given body node has been compiled yet."""
  82. return body_id in self.jitted_entry_points
  83. def get_compiled_name(self, body_id):
  84. """Gets the name of the compiled version of the given body node in the JIT
  85. global state."""
  86. return self.jitted_entry_points[body_id]
  87. def mark_no_jit(self, body_id):
  88. """Informs the JIT that the node with the given identifier is a function entry
  89. point that must never be jitted."""
  90. self.no_jit_entry_points.add(body_id)
  91. if body_id in self.todo_entry_points:
  92. self.todo_entry_points.remove(body_id)
  93. def generate_function_name(self, suggested_name=None):
  94. """Generates a new function name or picks the suggested name if it is still
  95. available."""
  96. if suggested_name is not None and suggested_name not in self.jit_globals:
  97. self.jit_count += 1
  98. return suggested_name
  99. else:
  100. function_name = 'jit_func%d' % self.jit_count
  101. self.jit_count += 1
  102. return function_name
  103. def register_compiled(self, body_id, compiled_function, function_name=None):
  104. """Registers a compiled entry point with the JIT."""
  105. # Get the function's name.
  106. function_name = self.generate_function_name(function_name)
  107. # Map the body id to the given parameter list.
  108. self.jitted_entry_points[body_id] = function_name
  109. self.jit_globals[function_name] = compiled_function
  110. if body_id in self.todo_entry_points:
  111. self.todo_entry_points.remove(body_id)
  112. def lookup_compiled_function(self, name):
  113. """Looks up a compiled function by name. Returns a matching function,
  114. or None if no function was found."""
  115. if name is None:
  116. return None
  117. elif name in self.jit_globals:
  118. return self.jit_globals[name]
  119. elif self.compiled_function_lookup is not None:
  120. return self.compiled_function_lookup(name)
  121. else:
  122. return None
  123. def get_intrinsic(self, name):
  124. """Tries to find an intrinsic version of the function with the
  125. given name."""
  126. if name in self.jit_intrinsics:
  127. return self.jit_intrinsics[name]
  128. else:
  129. return None
  130. def register_intrinsic(self, name, intrinsic_function):
  131. """Registers the given intrisic with the JIT. This will make the JIT replace calls to
  132. the function with the given entry point by an application of the specified function."""
  133. self.jit_intrinsics[name] = intrinsic_function
  134. def register_binary_intrinsic(self, name, operator):
  135. """Registers an intrinsic with the JIT that represents the given binary operation."""
  136. self.register_intrinsic(name, lambda a, b: tree_ir.CreateNodeWithValueInstruction(
  137. tree_ir.BinaryInstruction(
  138. tree_ir.ReadValueInstruction(a),
  139. operator,
  140. tree_ir.ReadValueInstruction(b))))
  141. def register_unary_intrinsic(self, name, operator):
  142. """Registers an intrinsic with the JIT that represents the given unary operation."""
  143. self.register_intrinsic(name, lambda a: tree_ir.CreateNodeWithValueInstruction(
  144. tree_ir.UnaryInstruction(
  145. operator,
  146. tree_ir.ReadValueInstruction(a))))
  147. def jit_parameters(self, body_id):
  148. """Acquires the parameter list for the given body id node."""
  149. if body_id not in self.jitted_parameters:
  150. signature_id, = yield [("RRD", [body_id, "body"])]
  151. signature_id = signature_id[0]
  152. param_set_id, = yield [("RD", [signature_id, "params"])]
  153. if param_set_id is None:
  154. self.jitted_parameters[body_id] = ([], [])
  155. else:
  156. param_name_ids, = yield [("RDK", [param_set_id])]
  157. param_names = yield [("RV", [n]) for n in param_name_ids]
  158. param_vars = yield [("RD", [param_set_id, k]) for k in param_names]
  159. self.jitted_parameters[body_id] = (param_vars, param_names)
  160. raise primitive_functions.PrimitiveFinished(self.jitted_parameters[body_id])
  161. def jit_compile(self, user_root, body_id, suggested_name=None):
  162. """Tries to jit the function defined by the given entry point id and parameter list."""
  163. # The comment below makes pylint shut up about our (hopefully benign) use of exec here.
  164. # pylint: disable=I0011,W0122
  165. if body_id in self.jitted_entry_points:
  166. # We have already compiled this function.
  167. raise primitive_functions.PrimitiveFinished(
  168. self.jit_globals[self.jitted_entry_points[body_id]])
  169. elif body_id in self.no_jit_entry_points:
  170. # We're not allowed to jit this function or have tried and failed before.
  171. raise JitCompilationFailedException(
  172. 'Cannot jit function %s at %d because it is marked non-jittable.' % (
  173. '' if suggested_name is None else "'" + suggested_name + "'",
  174. body_id))
  175. # Generate a name for the function we're about to analyze, and pretend that
  176. # it already exists. (we need to do this for recursive functions)
  177. function_name = self.generate_function_name(suggested_name)
  178. self.jitted_entry_points[body_id] = function_name
  179. self.jit_globals[function_name] = None
  180. try:
  181. gen = self.jit_parameters(body_id)
  182. inp = None
  183. while True:
  184. inp = yield gen.send(inp)
  185. except primitive_functions.PrimitiveFinished as ex:
  186. parameter_ids, parameter_list = ex.result
  187. param_dict = dict(zip(parameter_ids, parameter_list))
  188. body_param_dict = dict(zip(parameter_ids, [p + "_ptr" for p in parameter_list]))
  189. dependencies = set([body_id])
  190. self.compilation_dependencies[body_id] = dependencies
  191. try:
  192. gen = AnalysisState(
  193. self, body_id, user_root, body_param_dict,
  194. self.max_instructions).analyze(body_id)
  195. inp = None
  196. while True:
  197. inp = yield gen.send(inp)
  198. except primitive_functions.PrimitiveFinished as ex:
  199. del self.compilation_dependencies[body_id]
  200. constructed_body = ex.result
  201. except JitCompilationFailedException as ex:
  202. del self.compilation_dependencies[body_id]
  203. for dep in dependencies:
  204. self.mark_no_jit(dep)
  205. if dep in self.jitted_entry_points:
  206. del self.jitted_entry_points[dep]
  207. raise JitCompilationFailedException(
  208. "%s (function '%s' at %d)" % (ex.message, function_name, body_id))
  209. # Write a prologue and prepend it to the generated function body.
  210. prologue_statements = []
  211. for (key, val) in param_dict.items():
  212. arg_ptr = tree_ir.StoreLocalInstruction(
  213. body_param_dict[key],
  214. tree_ir.CreateNodeInstruction())
  215. prologue_statements.append(arg_ptr)
  216. prologue_statements.append(
  217. tree_ir.CreateDictionaryEdgeInstruction(
  218. arg_ptr.create_load(),
  219. tree_ir.LiteralInstruction('value'),
  220. tree_ir.LoadLocalInstruction(val)))
  221. constructed_body = tree_ir.create_block(
  222. *(prologue_statements + [constructed_body]))
  223. # Wrap the IR in a function definition, give it a unique name.
  224. constructed_function = tree_ir.DefineFunctionInstruction(
  225. function_name,
  226. parameter_list + ['**' + KWARGS_PARAMETER_NAME],
  227. constructed_body.simplify())
  228. # Convert the function definition to Python code, and compile it.
  229. exec(str(constructed_function), self.jit_globals)
  230. # Extract the compiled function from the JIT global state.
  231. compiled_function = self.jit_globals[function_name]
  232. print(constructed_function)
  233. raise primitive_functions.PrimitiveFinished(compiled_function)
  234. class AnalysisState(object):
  235. """The state of a bytecode analysis call graph."""
  236. def __init__(self, jit, body_id, user_root, local_mapping, max_instructions=None):
  237. self.analyzed_instructions = set()
  238. self.function_vars = set()
  239. self.local_vars = set()
  240. self.body_id = body_id
  241. self.max_instructions = max_instructions
  242. self.user_root = user_root
  243. self.jit = jit
  244. self.local_mapping = local_mapping
  245. def get_local_name(self, local_id):
  246. """Gets the name for a local with the given id."""
  247. if local_id not in self.local_mapping:
  248. self.local_mapping[local_id] = 'local%d' % local_id
  249. return self.local_mapping[local_id]
  250. def register_local_var(self, local_id):
  251. """Registers the given variable node id as a local."""
  252. if local_id in self.function_vars:
  253. raise JitCompilationFailedException(
  254. "Local is used as target of function call.")
  255. self.local_vars.add(local_id)
  256. def register_function_var(self, local_id):
  257. """Registers the given variable node id as a function."""
  258. if local_id in self.local_vars:
  259. raise JitCompilationFailedException(
  260. "Local is used as target of function call.")
  261. self.function_vars.add(local_id)
  262. def retrieve_user_root(self):
  263. """Creates an instruction that stores the user_root variable
  264. in a local."""
  265. return tree_ir.StoreLocalInstruction(
  266. 'user_root',
  267. tree_ir.LoadIndexInstruction(
  268. tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME),
  269. tree_ir.LiteralInstruction('user_root')))
  270. def load_kernel(self):
  271. """Creates an instruction that loads the Modelverse kernel."""
  272. return tree_ir.LoadIndexInstruction(
  273. tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME),
  274. tree_ir.LiteralInstruction('mvk'))
  275. def analyze(self, instruction_id):
  276. """Tries to build an intermediate representation from the instruction with the
  277. given id."""
  278. # Check the analyzed_instructions set for instruction_id to avoid
  279. # infinite loops.
  280. if instruction_id in self.analyzed_instructions:
  281. raise JitCompilationFailedException('Cannot jit non-tree instruction graph.')
  282. elif (self.max_instructions is not None and
  283. len(self.analyzed_instructions) > self.max_instructions):
  284. raise JitCompilationFailedException('Maximum number of instructions exceeded.')
  285. self.analyzed_instructions.add(instruction_id)
  286. instruction_val, = yield [("RV", [instruction_id])]
  287. instruction_val = instruction_val["value"]
  288. if instruction_val in self.instruction_analyzers:
  289. # If tracing is enabled, then this would be an appropriate time to
  290. # retrieve the debug information.
  291. if self.jit.tracing_enabled:
  292. debug_info, = yield [("RD", [instruction_id, "__debug"])]
  293. if debug_info is not None:
  294. debug_info, = yield [("RV", [debug_info])]
  295. else:
  296. debug_info = None
  297. try:
  298. gen = self.instruction_analyzers[instruction_val](self, instruction_id)
  299. inp = None
  300. while True:
  301. inp = yield gen.send(inp)
  302. except StopIteration:
  303. raise Exception(
  304. "Instruction analyzer (for '%s') finished without returning a value!" %
  305. (instruction_val))
  306. except primitive_functions.PrimitiveFinished as outer_e:
  307. # Check if the instruction has a 'next' instruction.
  308. next_instr, = yield [("RD", [instruction_id, "next"])]
  309. outer_result = tree_ir.with_debug_info_trace(outer_e.result, debug_info)
  310. if next_instr is None:
  311. raise primitive_functions.PrimitiveFinished(outer_result)
  312. else:
  313. gen = self.analyze(next_instr)
  314. try:
  315. inp = None
  316. while True:
  317. inp = yield gen.send(inp)
  318. except primitive_functions.PrimitiveFinished as inner_e:
  319. raise primitive_functions.PrimitiveFinished(
  320. tree_ir.CompoundInstruction(
  321. outer_result,
  322. inner_e.result))
  323. else:
  324. raise JitCompilationFailedException(
  325. "Unknown instruction type: '%s'" % (instruction_val))
  326. def analyze_all(self, instruction_ids):
  327. """Tries to compile a list of IR trees from the given list of instruction ids."""
  328. results = []
  329. for inst in instruction_ids:
  330. gen = self.analyze(inst)
  331. try:
  332. inp = None
  333. while True:
  334. inp = yield gen.send(inp)
  335. except primitive_functions.PrimitiveFinished as ex:
  336. results.append(ex.result)
  337. raise primitive_functions.PrimitiveFinished(results)
  338. def analyze_return(self, instruction_id):
  339. """Tries to analyze the given 'return' instruction."""
  340. retval_id, = yield [("RD", [instruction_id, 'value'])]
  341. if retval_id is None:
  342. raise primitive_functions.PrimitiveFinished(
  343. tree_ir.ReturnInstruction(
  344. tree_ir.EmptyInstruction()))
  345. else:
  346. gen = self.analyze(retval_id)
  347. try:
  348. inp = None
  349. while True:
  350. inp = yield gen.send(inp)
  351. except primitive_functions.PrimitiveFinished as ex:
  352. raise primitive_functions.PrimitiveFinished(
  353. tree_ir.ReturnInstruction(ex.result))
  354. def analyze_if(self, instruction_id):
  355. """Tries to analyze the given 'if' instruction."""
  356. cond, true, false = yield [
  357. ("RD", [instruction_id, "cond"]),
  358. ("RD", [instruction_id, "then"]),
  359. ("RD", [instruction_id, "else"])]
  360. try:
  361. gen = self.analyze_all(
  362. [cond, true]
  363. if false is None
  364. else [cond, true, false])
  365. inp = None
  366. while True:
  367. inp = yield gen.send(inp)
  368. except primitive_functions.PrimitiveFinished as ex:
  369. if false is None:
  370. cond_r, true_r = ex.result
  371. false_r = tree_ir.EmptyInstruction()
  372. else:
  373. cond_r, true_r, false_r = ex.result
  374. raise primitive_functions.PrimitiveFinished(
  375. tree_ir.SelectInstruction(
  376. tree_ir.ReadValueInstruction(cond_r),
  377. true_r,
  378. false_r))
  379. def analyze_while(self, instruction_id):
  380. """Tries to analyze the given 'while' instruction."""
  381. cond, body = yield [
  382. ("RD", [instruction_id, "cond"]),
  383. ("RD", [instruction_id, "body"])]
  384. try:
  385. gen = self.analyze_all([cond, body])
  386. inp = None
  387. while True:
  388. inp = yield gen.send(inp)
  389. except primitive_functions.PrimitiveFinished as ex:
  390. cond_r, body_r = ex.result
  391. raise primitive_functions.PrimitiveFinished(
  392. tree_ir.LoopInstruction(
  393. tree_ir.CompoundInstruction(
  394. tree_ir.SelectInstruction(
  395. tree_ir.ReadValueInstruction(cond_r),
  396. tree_ir.EmptyInstruction(),
  397. tree_ir.BreakInstruction()),
  398. body_r)))
  399. def analyze_constant(self, instruction_id):
  400. """Tries to analyze the given 'constant' (literal) instruction."""
  401. node_id, = yield [("RD", [instruction_id, "node"])]
  402. raise primitive_functions.PrimitiveFinished(
  403. tree_ir.LiteralInstruction(node_id))
  404. def analyze_output(self, instruction_id):
  405. """Tries to analyze the given 'output' instruction."""
  406. # The plan is to basically generate this tree:
  407. #
  408. # value = <some tree>
  409. # last_output, last_output_link, new_last_output = \
  410. # yield [("RD", [user_root, "last_output"]),
  411. # ("RDE", [user_root, "last_output"]),
  412. # ("CN", []),
  413. # ]
  414. # _, _, _, _ = \
  415. # yield [("CD", [last_output, "value", value]),
  416. # ("CD", [last_output, "next", new_last_output]),
  417. # ("CD", [user_root, "last_output", new_last_output]),
  418. # ("DE", [last_output_link])
  419. # ]
  420. # yield None
  421. value_id, = yield [("RD", [instruction_id, "value"])]
  422. try:
  423. gen = self.analyze(value_id)
  424. inp = None
  425. while True:
  426. inp = yield gen.send(inp)
  427. except primitive_functions.PrimitiveFinished as ex:
  428. value_local = tree_ir.StoreLocalInstruction('value', ex.result)
  429. store_user_root = self.retrieve_user_root()
  430. last_output = tree_ir.StoreLocalInstruction(
  431. 'last_output',
  432. tree_ir.ReadDictionaryValueInstruction(
  433. store_user_root.create_load(),
  434. tree_ir.LiteralInstruction('last_output')))
  435. last_output_link = tree_ir.StoreLocalInstruction(
  436. 'last_output_link',
  437. tree_ir.ReadDictionaryEdgeInstruction(
  438. store_user_root.create_load(),
  439. tree_ir.LiteralInstruction('last_output')))
  440. new_last_output = tree_ir.StoreLocalInstruction(
  441. 'new_last_output',
  442. tree_ir.CreateNodeInstruction())
  443. result = tree_ir.create_block(
  444. value_local,
  445. store_user_root,
  446. last_output,
  447. last_output_link,
  448. new_last_output,
  449. tree_ir.CreateDictionaryEdgeInstruction(
  450. last_output.create_load(),
  451. tree_ir.LiteralInstruction('value'),
  452. value_local.create_load()),
  453. tree_ir.CreateDictionaryEdgeInstruction(
  454. last_output.create_load(),
  455. tree_ir.LiteralInstruction('next'),
  456. new_last_output.create_load()),
  457. tree_ir.CreateDictionaryEdgeInstruction(
  458. store_user_root.create_load(),
  459. tree_ir.LiteralInstruction('last_output'),
  460. new_last_output.create_load()),
  461. tree_ir.DeleteEdgeInstruction(last_output_link.create_load()),
  462. tree_ir.NopInstruction())
  463. raise primitive_functions.PrimitiveFinished(result)
  464. def analyze_input(self, _):
  465. """Tries to analyze the given 'input' instruction."""
  466. # Possible alternative to the explicit syntax tree:
  467. #
  468. # raise primitive_functions.PrimitiveFinished(
  469. # tree_ir.JitCallInstruction(
  470. # tree_ir.LoadGlobalInstruction('__get_input'),
  471. # [],
  472. # tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME)))
  473. # The plan is to generate this tree:
  474. #
  475. # value = None
  476. # while True:
  477. # _input = yield [("RD", [user_root, "input"])]
  478. # value = yield [("RD", [_input, "value"])]
  479. #
  480. # if value is None:
  481. # kwargs['mvk'].success = False # to avoid blocking
  482. # yield None # nop/interrupt
  483. # else:
  484. # break
  485. #
  486. # _next = yield [("RD", [_input, "next"])]
  487. # yield [("CD", [user_root, "input", _next])]
  488. # yield [("DN", [_input])]
  489. user_root = self.retrieve_user_root()
  490. _input = tree_ir.StoreLocalInstruction(
  491. None,
  492. tree_ir.ReadDictionaryValueInstruction(
  493. user_root.create_load(),
  494. tree_ir.LiteralInstruction('input')))
  495. value = tree_ir.StoreLocalInstruction(
  496. None,
  497. tree_ir.ReadDictionaryValueInstruction(
  498. _input.create_load(),
  499. tree_ir.LiteralInstruction('value')))
  500. raise primitive_functions.PrimitiveFinished(
  501. tree_ir.CompoundInstruction(
  502. tree_ir.create_block(
  503. user_root,
  504. value.create_store(tree_ir.LiteralInstruction(None)),
  505. tree_ir.LoopInstruction(
  506. tree_ir.create_block(
  507. _input,
  508. value,
  509. tree_ir.SelectInstruction(
  510. tree_ir.BinaryInstruction(
  511. value.create_load(),
  512. 'is',
  513. tree_ir.LiteralInstruction(None)),
  514. tree_ir.create_block(
  515. tree_ir.StoreMemberInstruction(
  516. self.load_kernel(),
  517. 'success',
  518. tree_ir.LiteralInstruction(False)),
  519. tree_ir.NopInstruction()),
  520. tree_ir.BreakInstruction()))),
  521. tree_ir.CreateDictionaryEdgeInstruction(
  522. user_root.create_load(),
  523. tree_ir.LiteralInstruction('input'),
  524. tree_ir.ReadDictionaryValueInstruction(
  525. _input.create_load(),
  526. tree_ir.LiteralInstruction('next'))),
  527. tree_ir.DeleteNodeInstruction(_input.create_load())),
  528. value.create_load()))
  529. def analyze_resolve(self, instruction_id):
  530. """Tries to analyze the given 'resolve' instruction."""
  531. var_id, = yield [("RD", [instruction_id, "var"])]
  532. var_name, = yield [("RV", [var_id])]
  533. # To resolve a variable, we'll do something along the
  534. # lines of:
  535. #
  536. # if 'local_var' in locals():
  537. # tmp = local_var
  538. # else:
  539. # _globals, = yield [("RD", [user_root, "globals"])]
  540. # global_var, = yield [("RD", [_globals, var_name])]
  541. #
  542. # if global_var is None:
  543. # raise Exception("Runtime error: global '%s' not found" % (var_name))
  544. #
  545. # tmp = global_var
  546. name = self.get_local_name(var_id)
  547. if var_name is None:
  548. raise primitive_functions.PrimitiveFinished(
  549. tree_ir.LoadLocalInstruction(name))
  550. user_root = self.retrieve_user_root()
  551. global_var = tree_ir.StoreLocalInstruction(
  552. 'global_var',
  553. tree_ir.ReadDictionaryValueInstruction(
  554. tree_ir.ReadDictionaryValueInstruction(
  555. user_root.create_load(),
  556. tree_ir.LiteralInstruction('globals')),
  557. tree_ir.LiteralInstruction(var_name)))
  558. err_block = tree_ir.SelectInstruction(
  559. tree_ir.BinaryInstruction(
  560. global_var.create_load(),
  561. 'is',
  562. tree_ir.LiteralInstruction(None)),
  563. tree_ir.RaiseInstruction(
  564. tree_ir.CallInstruction(
  565. tree_ir.LoadGlobalInstruction('Exception'),
  566. [tree_ir.LiteralInstruction(
  567. "Runtime error: global '%s' not found" % var_name)
  568. ])),
  569. tree_ir.EmptyInstruction())
  570. raise primitive_functions.PrimitiveFinished(
  571. tree_ir.SelectInstruction(
  572. tree_ir.LocalExistsInstruction(name),
  573. tree_ir.LoadLocalInstruction(name),
  574. tree_ir.CompoundInstruction(
  575. tree_ir.create_block(
  576. user_root,
  577. global_var,
  578. err_block),
  579. global_var.create_load())))
  580. def analyze_declare(self, instruction_id):
  581. """Tries to analyze the given 'declare' function."""
  582. var_id, = yield [("RD", [instruction_id, "var"])]
  583. self.register_local_var(var_id)
  584. name = self.get_local_name(var_id)
  585. # The following logic declares a local:
  586. #
  587. # if 'local_name' not in locals():
  588. # local_name, = yield [("CN", [])]
  589. raise primitive_functions.PrimitiveFinished(
  590. tree_ir.SelectInstruction(
  591. tree_ir.LocalExistsInstruction(name),
  592. tree_ir.EmptyInstruction(),
  593. tree_ir.StoreLocalInstruction(
  594. name,
  595. tree_ir.CreateNodeInstruction())))
  596. def analyze_global(self, instruction_id):
  597. """Tries to analyze the given 'global' (declaration) instruction."""
  598. var_id, = yield [("RD", [instruction_id, "var"])]
  599. var_name, = yield [("RV", [var_id])]
  600. # To resolve a variable, we'll do something along the
  601. # lines of:
  602. #
  603. # _globals, = yield [("RD", [user_root, "globals"])]
  604. # global_var = yield [("RD", [_globals, var_name])]
  605. #
  606. # if global_var is None:
  607. # global_var, = yield [("CN", [])]
  608. # yield [("CD", [_globals, var_name, global_var])]
  609. #
  610. # tmp = global_var
  611. user_root = self.retrieve_user_root()
  612. _globals = tree_ir.StoreLocalInstruction(
  613. '_globals',
  614. tree_ir.ReadDictionaryValueInstruction(
  615. user_root.create_load(),
  616. tree_ir.LiteralInstruction('globals')))
  617. global_var = tree_ir.StoreLocalInstruction(
  618. 'global_var',
  619. tree_ir.ReadDictionaryValueInstruction(
  620. _globals.create_load(),
  621. tree_ir.LiteralInstruction(var_name)))
  622. raise primitive_functions.PrimitiveFinished(
  623. tree_ir.CompoundInstruction(
  624. tree_ir.create_block(
  625. user_root,
  626. _globals,
  627. global_var,
  628. tree_ir.SelectInstruction(
  629. tree_ir.BinaryInstruction(
  630. global_var.create_load(),
  631. 'is',
  632. tree_ir.LiteralInstruction(None)),
  633. tree_ir.create_block(
  634. global_var.create_store(
  635. tree_ir.CreateNodeInstruction()),
  636. tree_ir.CreateDictionaryEdgeInstruction(
  637. _globals.create_load(),
  638. tree_ir.LiteralInstruction(var_name),
  639. global_var.create_load())),
  640. tree_ir.EmptyInstruction())),
  641. global_var.create_load()))
  642. def analyze_assign(self, instruction_id):
  643. """Tries to analyze the given 'assign' instruction."""
  644. var_id, value_id = yield [("RD", [instruction_id, "var"]),
  645. ("RD", [instruction_id, "value"])]
  646. try:
  647. gen = self.analyze_all([var_id, value_id])
  648. inp = None
  649. while True:
  650. inp = yield gen.send(inp)
  651. except primitive_functions.PrimitiveFinished as ex:
  652. var_r, value_r = ex.result
  653. # Assignments work like this:
  654. #
  655. # value_link = yield [("RDE", [variable, "value"])]
  656. # _, _ = yield [("CD", [variable, "value", value]),
  657. # ("DE", [value_link])]
  658. variable = tree_ir.StoreLocalInstruction(None, var_r)
  659. value = tree_ir.StoreLocalInstruction(None, value_r)
  660. value_link = tree_ir.StoreLocalInstruction(
  661. 'value_link',
  662. tree_ir.ReadDictionaryEdgeInstruction(
  663. variable.create_load(),
  664. tree_ir.LiteralInstruction('value')))
  665. raise primitive_functions.PrimitiveFinished(
  666. tree_ir.create_block(
  667. variable,
  668. value,
  669. value_link,
  670. tree_ir.CreateDictionaryEdgeInstruction(
  671. variable.create_load(),
  672. tree_ir.LiteralInstruction('value'),
  673. value.create_load()),
  674. tree_ir.DeleteEdgeInstruction(
  675. value_link.create_load())))
  676. def analyze_access(self, instruction_id):
  677. """Tries to analyze the given 'access' instruction."""
  678. var_id, = yield [("RD", [instruction_id, "var"])]
  679. try:
  680. gen = self.analyze(var_id)
  681. inp = None
  682. while True:
  683. inp = yield gen.send(inp)
  684. except primitive_functions.PrimitiveFinished as ex:
  685. var_r = ex.result
  686. # Accessing a variable is pretty easy. It really just boils
  687. # down to reading the value corresponding to the 'value' key
  688. # of the variable.
  689. #
  690. # value, = yield [("RD", [returnvalue, "value"])]
  691. raise primitive_functions.PrimitiveFinished(
  692. tree_ir.ReadDictionaryValueInstruction(
  693. var_r,
  694. tree_ir.LiteralInstruction('value')))
  695. def analyze_direct_call(self, callee_id, callee_name, first_parameter_id):
  696. """Tries to analyze a direct 'call' instruction."""
  697. self.register_function_var(callee_id)
  698. body_id, = yield [("RD", [callee_id, "body"])]
  699. # Make this function dependent on the callee.
  700. if body_id in self.jit.compilation_dependencies:
  701. self.jit.compilation_dependencies[body_id].add(self.body_id)
  702. # Figure out if the function might be an intrinsic.
  703. intrinsic = self.jit.get_intrinsic(callee_name)
  704. if intrinsic is None:
  705. compiled_func = self.jit.lookup_compiled_function(callee_name)
  706. if compiled_func is None:
  707. # Compile the callee.
  708. try:
  709. gen = self.jit.jit_compile(self.user_root, body_id, callee_name)
  710. inp = None
  711. while True:
  712. inp = yield gen.send(inp)
  713. except primitive_functions.PrimitiveFinished as ex:
  714. pass
  715. else:
  716. self.jit.register_compiled(body_id, compiled_func, callee_name)
  717. # Get the callee's name.
  718. compiled_func_name = self.jit.get_compiled_name(body_id)
  719. # This handles the corner case where a constant node is called, like
  720. # 'call(constant(9), ...)'. In this case, `callee_name` is `None`
  721. # because 'constant(9)' doesn't give us a name. However, we can look up
  722. # the name of the function at a specific node. If that turns out to be
  723. # an intrinsic, then we still want to pick the intrinsic over a call.
  724. intrinsic = self.jit.get_intrinsic(compiled_func_name)
  725. # Analyze the argument dictionary.
  726. try:
  727. gen = self.analyze_arguments(first_parameter_id)
  728. inp = None
  729. while True:
  730. inp = yield gen.send(inp)
  731. except primitive_functions.PrimitiveFinished as ex:
  732. named_args = ex.result
  733. if intrinsic is not None:
  734. raise primitive_functions.PrimitiveFinished(
  735. apply_intrinsic(intrinsic, named_args))
  736. else:
  737. raise primitive_functions.PrimitiveFinished(
  738. tree_ir.JitCallInstruction(
  739. tree_ir.LoadGlobalInstruction(compiled_func_name),
  740. named_args,
  741. tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME)))
  742. def analyze_arguments(self, first_argument_id):
  743. """Analyzes the parameter-to-argument mapping started by the specified first argument
  744. node."""
  745. next_param = first_argument_id
  746. named_args = []
  747. while next_param is not None:
  748. param_name_id, = yield [("RD", [next_param, "name"])]
  749. param_name, = yield [("RV", [param_name_id])]
  750. param_val_id, = yield [("RD", [next_param, "value"])]
  751. try:
  752. gen = self.analyze(param_val_id)
  753. inp = None
  754. while True:
  755. inp = yield gen.send(inp)
  756. except primitive_functions.PrimitiveFinished as ex:
  757. named_args.append((param_name, ex.result))
  758. next_param, = yield [("RD", [next_param, "next_param"])]
  759. raise primitive_functions.PrimitiveFinished(named_args)
  760. def analyze_indirect_call(self, func_id, first_arg_id):
  761. """Analyzes a call to an unknown function."""
  762. # First off, let's analyze the callee and the argument list.
  763. try:
  764. gen = self.analyze(func_id)
  765. inp = None
  766. while True:
  767. inp = yield gen.send(inp)
  768. except primitive_functions.PrimitiveFinished as ex:
  769. func_val = ex.result
  770. try:
  771. gen = self.analyze_arguments(first_arg_id)
  772. inp = None
  773. while True:
  774. inp = yield gen.send(inp)
  775. except primitive_functions.PrimitiveFinished as ex:
  776. named_args = ex.result
  777. # Call the __interpret_function function to run the interpreter, like so:
  778. #
  779. # __interpret_function(function_id, { first_param_name : first_param_val, ... }, **kwargs)
  780. #
  781. dict_literal = tree_ir.DictionaryLiteralInstruction(
  782. [(tree_ir.LiteralInstruction(key), val) for key, val in named_args])
  783. raise primitive_functions.PrimitiveFinished(
  784. tree_ir.JitCallInstruction(
  785. tree_ir.LoadGlobalInstruction(INTERPRET_FUNCTION_NAME),
  786. [('function_id', func_val), ('named_arguments', dict_literal)],
  787. tree_ir.LoadLocalInstruction(KWARGS_PARAMETER_NAME)))
  788. def try_analyze_direct_call(self, func_id, first_param_id):
  789. """Tries to analyze the given 'call' instruction as a direct call."""
  790. if not self.jit.direct_calls_allowed:
  791. raise JitCompilationFailedException('Direct calls are not allowed by the JIT.')
  792. # Figure out what the 'func' instruction's type is.
  793. func_instruction_op, = yield [("RV", [func_id])]
  794. if func_instruction_op['value'] == 'access':
  795. # 'access(resolve(var))' instructions are translated to direct calls.
  796. access_value_id, = yield [("RD", [func_id, "var"])]
  797. access_value_op, = yield [("RV", [access_value_id])]
  798. if access_value_op['value'] == 'resolve':
  799. resolved_var_id, = yield [("RD", [access_value_id, "var"])]
  800. resolved_var_name, = yield [("RV", [resolved_var_id])]
  801. # Try to look up the name as a global.
  802. _globals, = yield [("RD", [self.user_root, "globals"])]
  803. global_var, = yield [("RD", [_globals, resolved_var_name])]
  804. global_val, = yield [("RD", [global_var, "value"])]
  805. if global_val is not None:
  806. gen = self.analyze_direct_call(
  807. global_val, resolved_var_name, first_param_id)
  808. inp = None
  809. while True:
  810. inp = yield gen.send(inp)
  811. # PrimitiveFinished exception will bubble up from here.
  812. elif func_instruction_op['value'] == 'constant':
  813. # 'const(func_id)' instructions are also translated to direct calls.
  814. function_val_id, = yield [("RD", [func_id, "node"])]
  815. gen = self.analyze_direct_call(
  816. function_val_id, None, first_param_id)
  817. inp = None
  818. while True:
  819. inp = yield gen.send(inp)
  820. raise JitCompilationFailedException(
  821. "Cannot JIT function calls that target an unknown value as direct calls.")
  822. def analyze_call(self, instruction_id):
  823. """Tries to analyze the given 'call' instruction."""
  824. func_id, first_param_id, = yield [("RD", [instruction_id, "func"]),
  825. ("RD", [instruction_id, "params"])]
  826. try:
  827. # Try to analyze the call as a direct call.
  828. gen = self.try_analyze_direct_call(func_id, first_param_id)
  829. inp = None
  830. while 1:
  831. inp = yield gen.send(inp)
  832. # PrimitiveFinished exception will bubble up from here.
  833. except JitCompilationFailedException:
  834. # Looks like we'll have to compile it as an indirect call.
  835. gen = self.analyze_indirect_call(func_id, first_param_id)
  836. inp = None
  837. while True:
  838. inp = yield gen.send(inp)
  839. # PrimitiveFinished exception will bubble up from here.
  840. instruction_analyzers = {
  841. 'if' : analyze_if,
  842. 'while' : analyze_while,
  843. 'return' : analyze_return,
  844. 'constant' : analyze_constant,
  845. 'resolve' : analyze_resolve,
  846. 'declare' : analyze_declare,
  847. 'global' : analyze_global,
  848. 'assign' : analyze_assign,
  849. 'access' : analyze_access,
  850. 'output' : analyze_output,
  851. 'input' : analyze_input,
  852. 'call' : analyze_call
  853. }