jit.py 43 KB

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