|
@@ -1,274 +0,0 @@
|
|
|
-import math
|
|
|
-import keyword
|
|
|
-from collections import defaultdict
|
|
|
-import modelverse_kernel.primitives as primitive_functions
|
|
|
-import modelverse_jit.runtime as jit_runtime
|
|
|
-
|
|
|
-# Import JitCompilationFailedException because it used to be defined
|
|
|
-# in this module.
|
|
|
-JitCompilationFailedException = jit_runtime.JitCompilationFailedException
|
|
|
-
|
|
|
-class ModelverseJit(object):
|
|
|
- """A high-level interface to the modelverse JIT compiler."""
|
|
|
- def __init__(self):
|
|
|
- self.todo_entry_points = set()
|
|
|
- self.no_jit_entry_points = set()
|
|
|
- self.jitted_parameters = {}
|
|
|
- self.jit_globals = {
|
|
|
- 'PrimitiveFinished' : primitive_functions.PrimitiveFinished,
|
|
|
- }
|
|
|
- # jitted_entry_points maps body ids to values in jit_globals.
|
|
|
- self.jitted_entry_points = {}
|
|
|
- # global_functions maps global value names to body ids.
|
|
|
- self.global_functions = {}
|
|
|
- # global_functions_inv maps body ids to global value names.
|
|
|
- self.global_functions_inv = {}
|
|
|
- # jitted_function_aliases maps body ids to known aliases.
|
|
|
- self.jitted_function_aliases = defaultdict(set)
|
|
|
- self.jit_count = 0
|
|
|
- self.compilation_dependencies = {}
|
|
|
- self.jit_enabled = True
|
|
|
-
|
|
|
- def set_jit_enabled(self, is_enabled=True):
|
|
|
- """Enables or disables the JIT."""
|
|
|
- self.jit_enabled = is_enabled
|
|
|
-
|
|
|
- def allow_direct_calls(self, is_allowed=True):
|
|
|
- """Allows or disallows direct calls from jitted to jitted code."""
|
|
|
- self.direct_calls_allowed = is_allowed
|
|
|
-
|
|
|
- def use_input_function(self, is_enabled=True):
|
|
|
- """Configures the JIT to compile 'input' instructions as function calls."""
|
|
|
- self.input_function_enabled = is_enabled
|
|
|
-
|
|
|
- def enable_tracing(self, is_enabled=True):
|
|
|
- """Enables or disables tracing for jitted code."""
|
|
|
- self.tracing_enabled = is_enabled
|
|
|
-
|
|
|
- def enable_source_maps(self, is_enabled=True):
|
|
|
- """Enables or disables the creation of source maps for jitted code. Source maps
|
|
|
- convert lines in the generated code to debug information.
|
|
|
-
|
|
|
- Source maps are enabled by default."""
|
|
|
- self.source_maps_enabled = is_enabled
|
|
|
-
|
|
|
- def enable_nop_insertion(self, is_enabled=True):
|
|
|
- """Enables or disables nop insertion for jitted code. If enabled, the JIT will
|
|
|
- insert nops at loop back-edges. Inserting nops sacrifices performance to
|
|
|
- keep the jitted code from blocking the thread of execution and consuming
|
|
|
- all resources; nops give the Modelverse server an opportunity to interrupt
|
|
|
- the currently running code."""
|
|
|
- self.nop_insertion_enabled = is_enabled
|
|
|
-
|
|
|
- def set_function_body_compiler(self, compile_function_body):
|
|
|
- """Sets the function that the JIT uses to compile function bodies."""
|
|
|
- self.compile_function_body = compile_function_body
|
|
|
-
|
|
|
- def mark_entry_point(self, body_id):
|
|
|
- """Marks the node with the given identifier as a function entry point."""
|
|
|
- if body_id not in self.no_jit_entry_points and body_id not in self.jitted_entry_points:
|
|
|
- self.todo_entry_points.add(body_id)
|
|
|
-
|
|
|
- def is_entry_point(self, body_id):
|
|
|
- """Tells if the node with the given identifier is a function entry point."""
|
|
|
- return body_id in self.todo_entry_points or \
|
|
|
- body_id in self.no_jit_entry_points or \
|
|
|
- body_id in self.jitted_entry_points
|
|
|
-
|
|
|
- def is_jittable_entry_point(self, body_id):
|
|
|
- """Tells if the node with the given identifier is a function entry point that
|
|
|
- has not been marked as non-jittable. This only returns `True` if the JIT
|
|
|
- is enabled and the function entry point has been marked jittable, or if
|
|
|
- the function has already been compiled."""
|
|
|
- return ((self.jit_enabled and body_id in self.todo_entry_points) or
|
|
|
- self.has_compiled(body_id))
|
|
|
-
|
|
|
- def has_compiled(self, body_id):
|
|
|
- """Tests if the function belonging to the given body node has been compiled yet."""
|
|
|
- return body_id in self.jitted_entry_points
|
|
|
-
|
|
|
- def get_compiled_name(self, body_id):
|
|
|
- """Gets the name of the compiled version of the given body node in the JIT
|
|
|
- global state."""
|
|
|
- if body_id in self.jitted_entry_points:
|
|
|
- return self.jitted_entry_points[body_id]
|
|
|
- else:
|
|
|
- return None
|
|
|
-
|
|
|
- def mark_no_jit(self, body_id):
|
|
|
- """Informs the JIT that the node with the given identifier is a function entry
|
|
|
- point that must never be jitted."""
|
|
|
- self.no_jit_entry_points.add(body_id)
|
|
|
- if body_id in self.todo_entry_points:
|
|
|
- self.todo_entry_points.remove(body_id)
|
|
|
-
|
|
|
- def generate_name(self, infix, suggested_name=None):
|
|
|
- """Generates a new name or picks the suggested name if it is still
|
|
|
- available."""
|
|
|
- if suggested_name is not None \
|
|
|
- and suggested_name not in self.jit_globals \
|
|
|
- and not keyword.iskeyword(suggested_name):
|
|
|
- self.jit_count += 1
|
|
|
- return suggested_name
|
|
|
- else:
|
|
|
- function_name = 'jit_%s%d' % (infix, self.jit_count)
|
|
|
- self.jit_count += 1
|
|
|
- return function_name
|
|
|
-
|
|
|
- def generate_function_name(self, body_id, suggested_name=None):
|
|
|
- """Generates a new function name or picks the suggested name if it is still
|
|
|
- available."""
|
|
|
- if suggested_name is None:
|
|
|
- suggested_name = self.get_global_name(body_id)
|
|
|
-
|
|
|
- return self.generate_name('func', suggested_name)
|
|
|
-
|
|
|
- def register_global(self, body_id, global_name):
|
|
|
- """Associates the given body id with the given global name."""
|
|
|
- self.global_functions[global_name] = body_id
|
|
|
- self.global_functions_inv[body_id] = global_name
|
|
|
-
|
|
|
- def get_global_name(self, body_id):
|
|
|
- """Gets the name of the global function with the given body id.
|
|
|
- Returns None if no known global exists with the given id."""
|
|
|
- if body_id in self.global_functions_inv:
|
|
|
- return self.global_functions_inv[body_id]
|
|
|
- else:
|
|
|
- return None
|
|
|
-
|
|
|
- def get_global_body_id(self, global_name):
|
|
|
- """Gets the body id of the global function with the given name.
|
|
|
- Returns None if no known global exists with the given name."""
|
|
|
- if global_name in self.global_functions:
|
|
|
- return self.global_functions[global_name]
|
|
|
- else:
|
|
|
- return None
|
|
|
-
|
|
|
- def register_compiled(self, body_id, compiled_function, function_name=None):
|
|
|
- """Registers a compiled entry point with the JIT."""
|
|
|
- # Get the function's name.
|
|
|
- actual_function_name = self.generate_function_name(body_id,
|
|
|
- function_name)
|
|
|
- # Map the body id to the given parameter list.
|
|
|
- self.jitted_entry_points[body_id] = actual_function_name
|
|
|
- self.jit_globals[actual_function_name] = compiled_function
|
|
|
- if function_name is not None:
|
|
|
- self.register_global(body_id, function_name)
|
|
|
-
|
|
|
- if body_id in self.todo_entry_points:
|
|
|
- self.todo_entry_points.remove(body_id)
|
|
|
-
|
|
|
- def import_value(self, value, suggested_name=None):
|
|
|
- """Imports the given value into the JIT's global scope, with the given suggested name.
|
|
|
- The actual name of the value (within the JIT's global scope) is returned."""
|
|
|
- actual_name = self.generate_name('import', suggested_name)
|
|
|
- self.jit_globals[actual_name] = value
|
|
|
- return actual_name
|
|
|
-
|
|
|
- def __lookup_compiled_body_impl(self, body_id):
|
|
|
- """Looks up a compiled function by body id. Returns a matching function,
|
|
|
- or None if no function was found."""
|
|
|
- if body_id is not None and body_id in self.jitted_entry_points:
|
|
|
- return self.jit_globals[self.jitted_entry_points[body_id]]
|
|
|
- else:
|
|
|
- return None
|
|
|
-
|
|
|
- def __lookup_external_body_impl(self, global_name, body_id):
|
|
|
- """Looks up an external function by global name. Returns a matching function,
|
|
|
- or None if no function was found."""
|
|
|
- if global_name is not None and self.compiled_function_lookup is not None:
|
|
|
- result = self.compiled_function_lookup(global_name)
|
|
|
- if result is not None and body_id is not None:
|
|
|
- self.register_compiled(body_id, result, global_name)
|
|
|
-
|
|
|
- return result
|
|
|
- else:
|
|
|
- return None
|
|
|
-
|
|
|
- def lookup_compiled_body(self, body_id):
|
|
|
- """Looks up a compiled function by body id. Returns a matching function,
|
|
|
- or None if no function was found."""
|
|
|
- result = self.__lookup_compiled_body_impl(body_id)
|
|
|
- if result is not None:
|
|
|
- return result
|
|
|
- else:
|
|
|
- global_name = self.get_global_name(body_id)
|
|
|
- return self.__lookup_external_body_impl(global_name, body_id)
|
|
|
-
|
|
|
- def lookup_compiled_function(self, global_name):
|
|
|
- """Looks up a compiled function by global name. Returns a matching function,
|
|
|
- or None if no function was found."""
|
|
|
- body_id = self.get_global_body_id(global_name)
|
|
|
- result = self.__lookup_compiled_body_impl(body_id)
|
|
|
- if result is not None:
|
|
|
- return result
|
|
|
- else:
|
|
|
- return self.__lookup_external_body_impl(global_name, body_id)
|
|
|
-
|
|
|
- def jit_signature(self, body_id):
|
|
|
- """Acquires the signature for the given body id node, which consists of the
|
|
|
- parameter variables, parameter name and a flag that tells if the given function
|
|
|
- is mutable."""
|
|
|
- if body_id not in self.jitted_parameters:
|
|
|
- signature_id, = yield [("RRD", [body_id, jit_runtime.FUNCTION_BODY_KEY])]
|
|
|
- signature_id = signature_id[0]
|
|
|
- param_set_id, is_mutable = yield [
|
|
|
- ("RD", [signature_id, "params"]),
|
|
|
- ("RD", [signature_id, jit_runtime.MUTABLE_FUNCTION_KEY])]
|
|
|
- if param_set_id is None:
|
|
|
- self.jitted_parameters[body_id] = ([], [], is_mutable)
|
|
|
- else:
|
|
|
- param_name_ids, = yield [("RDK", [param_set_id])]
|
|
|
- param_names = yield [("RV", [n]) for n in param_name_ids]
|
|
|
- #NOTE Patch up strange links...
|
|
|
- param_names = [i for i in param_names if i is not None]
|
|
|
- param_vars = yield [("RD", [param_set_id, k]) for k in param_names]
|
|
|
-
|
|
|
- #NOTE that variables might not be in the correct order, as we just read them out!
|
|
|
- lst = sorted([(name, var) for name, var in zip(param_names, param_vars)])
|
|
|
- param_vars = [i[1] for i in lst]
|
|
|
- param_names = [i[0] for i in lst]
|
|
|
- self.jitted_parameters[body_id] = (param_vars, param_names, is_mutable)
|
|
|
-
|
|
|
- raise primitive_functions.PrimitiveFinished(self.jitted_parameters[body_id])
|
|
|
-
|
|
|
- def check_jittable(self, body_id, suggested_name=None):
|
|
|
- """Checks if the function with the given body id is obviously non-jittable. If it's
|
|
|
- non-jittable, then a `JitCompilationFailedException` exception is thrown."""
|
|
|
- if body_id is None:
|
|
|
- raise ValueError('body_id cannot be None: ' + suggested_name)
|
|
|
- elif body_id in self.no_jit_entry_points:
|
|
|
- # We're not allowed to jit this function or have tried and failed before.
|
|
|
- raise JitCompilationFailedException(
|
|
|
- 'Cannot jit function %s at %d because it is marked non-jittable.' % (
|
|
|
- '' if suggested_name is None else "'" + suggested_name + "'",
|
|
|
- body_id))
|
|
|
- elif not self.jit_enabled:
|
|
|
- # We're not allowed to jit anything.
|
|
|
- raise JitCompilationFailedException(
|
|
|
- 'Cannot jit function %s at %d because the JIT has been disabled.' % (
|
|
|
- '' if suggested_name is None else "'" + suggested_name + "'",
|
|
|
- body_id))
|
|
|
-
|
|
|
- def jit_define_function(self, function_name, function_def):
|
|
|
- """Converts the given tree-IR function definition to Python code, defines it,
|
|
|
- and extracts the resulting function."""
|
|
|
- # The comment below makes pylint shut up about our (hopefully benign) use of exec here.
|
|
|
- # pylint: disable=I0011,W0122
|
|
|
- if self.jit_code_log_function is not None:
|
|
|
- self.jit_code_log_function(function_def)
|
|
|
-
|
|
|
- # Convert the function definition to Python code, and compile it.
|
|
|
- code_generator = tree_ir.PythonGenerator()
|
|
|
- function_def.generate_python_def(code_generator)
|
|
|
- source_map_name = self.get_source_map_name(function_name)
|
|
|
- if source_map_name is not None:
|
|
|
- self.jit_globals[source_map_name] = code_generator.source_map_builder.source_map
|
|
|
- exec(str(code_generator), self.jit_globals)
|
|
|
-
|
|
|
- # Extract the compiled function from the JIT global state.
|
|
|
- return self.jit_globals[function_name]
|
|
|
-
|
|
|
- def jit_delete_function(self, function_name):
|
|
|
- """Deletes the function with the given function name."""
|
|
|
- del self.jit_globals[function_name]
|