runtime.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  1. import modelverse_kernel.primitives as primitive_functions
  2. class JitCompilationFailedException(Exception):
  3. """A type of exception that is raised when the jit fails to compile a function."""
  4. pass
  5. MUTABLE_FUNCTION_KEY = "mutable"
  6. """A dictionary key for functions that are mutable."""
  7. FUNCTION_BODY_KEY = "body"
  8. """A dictionary key for function bodies."""
  9. KWARGS_PARAMETER_NAME = "kwargs"
  10. """The name of the kwargs parameter in jitted functions."""
  11. CALL_FUNCTION_NAME = "__call_function"
  12. """The name of the '__call_function' function, in the jitted function scope."""
  13. GET_INPUT_FUNCTION_NAME = "__get_input"
  14. """The name of the '__get_input' function, in the jitted function scope."""
  15. JIT_THUNK_CONSTANT_FUNCTION_NAME = "__jit_thunk_constant_function"
  16. """The name of the jit_thunk_constant_function function in the JIT's global context."""
  17. JIT_THUNK_GLOBAL_FUNCTION_NAME = "__jit_thunk_global"
  18. """The name of the jit_thunk_global function in the JIT's global context."""
  19. JIT_REJIT_FUNCTION_NAME = "__jit_rejit"
  20. """The name of the rejit function in the JIT's global context."""
  21. JIT_COMPILE_FUNCTION_BODY_FAST_FUNCTION_NAME = "__jit_compile_function_body_fast"
  22. """The name of the compile_function_body_fast function in the JIT's global context."""
  23. UNREACHABLE_FUNCTION_NAME = "__unreachable"
  24. """The name of the unreachable function in the JIT's global context."""
  25. LOCALS_NODE_NAME = "jit_locals"
  26. """The name of the node that is connected to all JIT locals in a given function call."""
  27. LOCALS_EDGE_NAME = "jit_locals_edge"
  28. """The name of the edge that connects the LOCALS_NODE_NAME node to a user root."""
  29. GLOBAL_NOT_FOUND_MESSAGE_FORMAT = "Not found as global: %s"
  30. """The format of the 'not found as global' message. Takes a single argument."""
  31. BYTECODE_INTERPRETER_ORIGIN_NAME = "bytecode-interpreter"
  32. """The origin name for functions that were produced by the bytecode interpreter."""
  33. BASELINE_JIT_ORIGIN_NAME = "baseline-jit"
  34. """The origin name for functions that were produced by the baseline JIT."""
  35. FAST_JIT_ORIGIN_NAME = "fast-jit"
  36. """The origin name for functions that were produced by the fast JIT."""
  37. def format_stack_frame(function_name, debug_info, origin='unknown'):
  38. """Formats a stack frame, which consists of a function name, debug
  39. information and an origin."""
  40. if function_name is None:
  41. function_name = 'unknown function'
  42. if debug_info is None:
  43. debug_info = '[unknown location] '
  44. return '%sin %s (%s)' % (debug_info, function_name, origin)
  45. def format_trace_message(debug_info, function_name, origin='unknown'):
  46. """Creates a formatted trace message."""
  47. return 'TRACE: %s' % format_stack_frame(function_name, debug_info, origin)
  48. def call_function(function_id, named_arguments, **kwargs):
  49. """Runs the function with the given id, passing it the specified argument dictionary."""
  50. task_root = kwargs['task_root']
  51. kernel = kwargs['mvk']
  52. body_id, is_mutable = yield [
  53. ("RD", [function_id, FUNCTION_BODY_KEY]),
  54. ("RD", [function_id, MUTABLE_FUNCTION_KEY])]
  55. # Try to jit the function here. We might be able to avoid building the stack
  56. # frame.
  57. def handle_jit_failed(_):
  58. """Interprets the function."""
  59. interpreter_args = {'body_id' : body_id, 'named_arguments' : named_arguments}
  60. interpreter_args.update(kwargs)
  61. yield [("TAIL_CALL_KWARGS", [interpret_function_body, interpreter_args])]
  62. if is_mutable is not None:
  63. kernel.jit.mark_no_jit(body_id)
  64. yield [("TAIL_CALL_ARGS", [handle_jit_failed, ()])]
  65. else:
  66. kernel.jit.mark_entry_point(body_id)
  67. yield [("TRY", [])]
  68. yield [("CATCH", [JitCompilationFailedException, handle_jit_failed])]
  69. # Try to compile.
  70. compiled_func, = yield [("CALL_ARGS", [kernel.jit_compile, (task_root, body_id)])]
  71. yield [("END_TRY", [])]
  72. # Add the keyword arguments to the argument dictionary.
  73. named_arguments.update(kwargs)
  74. # Run the function.
  75. yield [("TAIL_CALL_KWARGS", [compiled_func, named_arguments])]
  76. def interpret_function(function_id, named_arguments, **kwargs):
  77. """Makes the interpreter run the function with the given id for the specified
  78. argument dictionary."""
  79. body_id, = yield [("RD", [function_id, FUNCTION_BODY_KEY])]
  80. args = {'body_id' : body_id, named_arguments : named_arguments}
  81. args.update(kwargs)
  82. yield [("TAIL_CALL_KWARGS", [interpret_function_body, args])]
  83. def interpret_function_body(body_id, named_arguments, **kwargs):
  84. """Makes the interpreter run the function body with the given id for the specified
  85. argument dictionary."""
  86. task_root = kwargs['task_root']
  87. kernel = kwargs['mvk']
  88. user_frame, = yield [("RD", [task_root, "frame"])]
  89. inst, = yield [("RD", [user_frame, "IP"])]
  90. kernel.jit.mark_entry_point(body_id)
  91. # Create a new stack frame.
  92. frame_link, new_phase, new_frame, new_evalstack, new_symbols, \
  93. new_returnvalue, intrinsic_return = \
  94. yield [("RDE", [task_root, "frame"]),
  95. ("CNV", ["init"]),
  96. ("CN", []),
  97. ("CN", []),
  98. ("CN", []),
  99. ("CN", []),
  100. ("CN", [])
  101. ]
  102. _, _, _, _, _, _, _, _, _, _ = \
  103. yield [("CD", [task_root, "frame", new_frame]),
  104. ("CD", [new_frame, "evalstack", new_evalstack]),
  105. ("CD", [new_frame, "symbols", new_symbols]),
  106. ("CD", [new_frame, "returnvalue", new_returnvalue]),
  107. ("CD", [new_frame, "caller", inst]),
  108. ("CD", [new_frame, "phase", new_phase]),
  109. ("CD", [new_frame, "IP", body_id]),
  110. ("CD", [new_frame, "prev", user_frame]),
  111. ("CD", [
  112. new_frame,
  113. primitive_functions.EXCEPTION_RETURN_KEY,
  114. intrinsic_return]),
  115. ("DE", [frame_link])
  116. ]
  117. # Put the parameters in the new stack frame's symbol table.
  118. (parameter_vars, parameter_names, _), = yield [
  119. ("CALL_ARGS", [kernel.jit.jit_signature, (body_id,)])]
  120. parameter_dict = dict(zip(parameter_names, parameter_vars))
  121. for (key, value) in named_arguments.items():
  122. param_var = parameter_dict[key]
  123. variable, = yield [("CN", [])]
  124. yield [("CD", [variable, "value", value])]
  125. symbol_edge, = yield [("CE", [new_symbols, variable])]
  126. yield [("CE", [symbol_edge, param_var])]
  127. taskname = kwargs['taskname']
  128. def exception_handler(ex):
  129. # print('Returning from interpreted function. Result: %s' % ex.result)
  130. raise primitive_functions.PrimitiveFinished(ex.result)
  131. # Create an exception handler to catch and translate InterpretedFunctionFinished.
  132. yield [("TRY", [])]
  133. yield [("CATCH", [primitive_functions.InterpretedFunctionFinished, exception_handler])]
  134. while 1:
  135. result, = yield [("CALL_ARGS", [kernel.execute_rule, (taskname,)])]
  136. # An instruction has completed. Forward it.
  137. yield result
  138. class UnreachableCodeException(Exception):
  139. """The type of exception that is thrown when supposedly unreachable code is executed."""
  140. pass
  141. def unreachable():
  142. """Marks unreachable code."""
  143. raise UnreachableCodeException('An unreachable statement was reached.')
  144. def get_input(**parameters):
  145. """Retrieves input."""
  146. mvk = parameters["mvk"]
  147. task_root = parameters["task_root"]
  148. while 1:
  149. yield [("CALL_ARGS", [mvk.input_init, (task_root,)])]
  150. # Finished
  151. if mvk.success:
  152. # Got some input, so we can access it
  153. raise primitive_functions.PrimitiveFinished(mvk.input_value)
  154. else:
  155. # No input, so yield None but don't stop
  156. yield None