forked from M-Labs/artiq
1
0
Fork 0

embedding: specialize inherited functions.

Fixes #414.
This commit is contained in:
whitequark 2016-05-16 15:59:09 +00:00
parent 640022122b
commit 355af3e569
4 changed files with 172 additions and 93 deletions

View File

@ -22,6 +22,26 @@ from .transforms.asttyped_rewriter import LocalExtractor
def coredevice_print(x): print(x)
class SpecializedFunction:
def __init__(self, instance_type, host_function):
self.instance_type = instance_type
self.host_function = host_function
def __eq__(self, other):
if isinstance(other, tuple):
return (self.instance_type == other[0] or
self.host_function == other[1])
else:
return (self.instance_type == other.instance_type or
self.host_function == other.host_function)
def __ne__(self, other):
return not self == other
def __hash__(self):
return hash((self.instance_type, self.host_function))
class EmbeddingMap:
def __init__(self):
self.object_current_key = 0
@ -31,17 +51,14 @@ class EmbeddingMap:
self.function_map = {}
# Types
def store_type(self, typ, instance_type, constructor_type):
self.type_map[typ] = (instance_type, constructor_type)
def store_type(self, host_type, instance_type, constructor_type):
self.type_map[host_type] = (instance_type, constructor_type)
def retrieve_type(self, typ):
return self.type_map[typ]
def retrieve_type(self, host_type):
return self.type_map[host_type]
def has_type(self, typ):
return typ in self.type_map
def iter_types(self):
return self.type_map.values()
def has_type(self, host_type):
return host_type in self.type_map
# Functions
def store_function(self, function, ir_function_name):
@ -50,6 +67,9 @@ class EmbeddingMap:
def retrieve_function(self, function):
return self.function_map[function]
def specialize_function(self, instance_type, host_function):
return SpecializedFunction(instance_type, host_function)
# Objects
def store_object(self, obj_ref):
obj_id = id(obj_ref)
@ -65,12 +85,22 @@ class EmbeddingMap:
return self.object_forward_map[obj_key]
def iter_objects(self):
return self.object_forward_map.keys()
for obj_id in self.object_forward_map.keys():
obj_ref = self.object_forward_map[obj_id]
if isinstance(obj_ref, (pytypes.FunctionType, pytypes.MethodType,
pytypes.BuiltinFunctionType, SpecializedFunction)):
continue
elif isinstance(obj_ref, type):
_, obj_typ = self.type_map[obj_ref]
else:
obj_typ, _ = self.type_map[type(obj_ref)]
yield obj_id, obj_ref, obj_typ
def has_rpc(self):
return any(filter(lambda x: inspect.isfunction(x) or inspect.ismethod(x),
self.object_forward_map.values()))
class ASTSynthesizer:
def __init__(self, embedding_map, value_map, quote_function=None, expanded_from=None):
self.source = ""
@ -128,7 +158,8 @@ class ASTSynthesizer:
begin_loc=begin_loc, end_loc=end_loc,
loc=begin_loc.join(end_loc))
elif inspect.isfunction(value) or inspect.ismethod(value) or \
isinstance(value, pytypes.BuiltinFunctionType):
isinstance(value, pytypes.BuiltinFunctionType) or \
isinstance(value, SpecializedFunction):
if inspect.ismethod(value):
quoted_self = self.quote(value.__self__)
function_type = self.quote_function(value.__func__, self.expanded_from)
@ -139,7 +170,7 @@ class ASTSynthesizer:
loc = quoted_self.loc.join(name_loc)
return asttyped.QuoteT(value=value, type=method_type,
self_loc=quoted_self.loc, loc=loc)
else:
else: # function
function_type = self.quote_function(value, self.expanded_from)
quote_loc = self._add('`')
@ -417,7 +448,7 @@ class StitchingInferencer(Inferencer):
# def f(self): pass
# we want f to be defined on the class, not on the instance.
attributes = object_type.constructor.attributes
attr_value = attr_value.__func__
attr_value = SpecializedFunction(object_type, attr_value.__func__)
else:
attributes = object_type.attributes
@ -582,26 +613,6 @@ class Stitcher:
break
old_typedtree_hash = typedtree_hash
# For every host class we embed, fill in the function slots
# with their corresponding closures.
for instance_type, constructor_type in self.embedding_map.iter_types():
# Do we have any direct reference to a constructor?
if len(self.value_map[constructor_type]) > 0:
# Yes, use it.
constructor, _constructor_loc = self.value_map[constructor_type][0]
else:
# No, extract one from a reference to an instance.
instance, _instance_loc = self.value_map[instance_type][0]
constructor = type(instance)
for attr in constructor_type.attributes:
if types.is_function(constructor_type.attributes[attr]):
synthesizer = self._synthesizer()
ast = synthesizer.assign_attribute(constructor, attr,
getattr(constructor, attr))
synthesizer.finalize()
self._inject(ast)
# After we have found all functions, synthesize a module to hold them.
source_buffer = source.Buffer("", "<synthesized>")
self.typedtree = asttyped.ModuleT(
@ -619,11 +630,16 @@ class Stitcher:
quote_function=self._quote_function)
def _quote_embedded_function(self, function, flags):
if not hasattr(function, "artiq_embedded"):
raise ValueError("{} is not an embedded function".format(repr(function)))
if isinstance(function, SpecializedFunction):
host_function = function.host_function
else:
host_function = function
if not hasattr(host_function, "artiq_embedded"):
raise ValueError("{} is not an embedded function".format(repr(host_function)))
# Extract function source.
embedded_function = function.artiq_embedded.function
embedded_function = host_function.artiq_embedded.function
source_code = inspect.getsource(embedded_function)
filename = embedded_function.__code__.co_filename
module_name = embedded_function.__globals__['__name__']
@ -652,7 +668,13 @@ class Stitcher:
function_node = parser.file_input().body[0]
# Mangle the name, since we put everything into a single module.
function_node.name = "{}.{}".format(module_name, function.__qualname__)
full_function_name = "{}.{}".format(module_name, host_function.__qualname__)
if isinstance(function, SpecializedFunction):
instance_type = function.instance_type
function_node.name = "_Z{}{}I{}{}Ezz".format(len(full_function_name), full_function_name,
len(instance_type.name), instance_type.name)
else:
function_node.name = "_Z{}{}zz".format(len(full_function_name), full_function_name)
# Record the function in the function map so that LLVM IR generator
# can handle quoting it.
@ -808,64 +830,75 @@ class Stitcher:
return function_type
def _quote_rpc(self, function, loc):
if isinstance(function, SpecializedFunction):
host_function = function.host_function
else:
host_function = function
ret_type = builtins.TNone()
if isinstance(function, pytypes.BuiltinFunctionType):
if isinstance(host_function, pytypes.BuiltinFunctionType):
pass
elif isinstance(function, pytypes.FunctionType) or isinstance(function, pytypes.MethodType):
if isinstance(function, pytypes.FunctionType):
signature = inspect.signature(function)
elif (isinstance(host_function, pytypes.FunctionType) or \
isinstance(host_function, pytypes.MethodType)):
if isinstance(host_function, pytypes.FunctionType):
signature = inspect.signature(host_function)
else:
# inspect bug?
signature = inspect.signature(function.__func__)
signature = inspect.signature(host_function.__func__)
if signature.return_annotation is not inspect.Signature.empty:
ret_type = self._extract_annot(function, signature.return_annotation,
ret_type = self._extract_annot(host_function, signature.return_annotation,
"return type", loc, is_syscall=False)
else:
assert False
function_type = types.TRPC(ret_type, service=self.embedding_map.store_object(function))
function_type = types.TRPC(ret_type,
service=self.embedding_map.store_object(host_function))
self.functions[function] = function_type
return function_type
def _quote_function(self, function, loc):
if isinstance(function, SpecializedFunction):
host_function = function.host_function
else:
host_function = function
if function in self.functions:
pass
elif not hasattr(function, "artiq_embedded"):
elif not hasattr(host_function, "artiq_embedded"):
self._quote_rpc(function, loc)
elif function.artiq_embedded.function is not None:
if function.__name__ == "<lambda>":
elif host_function.artiq_embedded.function is not None:
if host_function.__name__ == "<lambda>":
note = diagnostic.Diagnostic("note",
"lambda created here", {},
self._function_loc(function.artiq_embedded.function))
self._function_loc(host_function.artiq_embedded.function))
diag = diagnostic.Diagnostic("fatal",
"lambdas cannot be used as kernel functions", {},
loc,
notes=[note])
self.engine.process(diag)
core_name = function.artiq_embedded.core_name
core_name = host_function.artiq_embedded.core_name
if core_name is not None and self.dmgr.get(core_name) != self.core:
note = diagnostic.Diagnostic("note",
"called from this function", {},
loc)
diag = diagnostic.Diagnostic("fatal",
"this function runs on a different core device '{name}'",
{"name": function.artiq_embedded.core_name},
self._function_loc(function.artiq_embedded.function),
{"name": host_function.artiq_embedded.core_name},
self._function_loc(host_function.artiq_embedded.function),
notes=[note])
self.engine.process(diag)
self._quote_embedded_function(function,
flags=function.artiq_embedded.flags)
elif function.artiq_embedded.syscall is not None:
flags=host_function.artiq_embedded.flags)
elif host_function.artiq_embedded.syscall is not None:
# Insert a storage-less global whose type instructs the compiler
# to perform a system call instead of a regular call.
self._quote_syscall(function, loc)
elif function.artiq_embedded.forbidden is not None:
elif host_function.artiq_embedded.forbidden is not None:
diag = diagnostic.Diagnostic("fatal",
"this function cannot be called as an RPC", {},
self._function_loc(function),
self._function_loc(host_function),
notes=self._call_site_note(loc, is_syscall=True))
self.engine.process(diag)
else:

View File

@ -891,20 +891,27 @@ class Inferencer(algorithm.Visitor):
typ_optargs = typ.optargs
typ_ret = typ.ret
else:
typ = types.get_method_function(typ)
if types.is_var(typ):
typ_self = types.get_method_self(typ)
typ_func = types.get_method_function(typ)
if types.is_var(typ_func):
return # not enough info yet
elif types.is_rpc(typ):
self._unify(node.type, typ.ret,
elif types.is_rpc(typ_func):
self._unify(node.type, typ_func.ret,
node.loc, None)
return
elif typ.arity() == 0:
elif typ_func.arity() == 0:
return # error elsewhere
typ_arity = typ.arity() - 1
typ_args = OrderedDict(list(typ.args.items())[1:])
typ_optargs = typ.optargs
typ_ret = typ.ret
method_args = list(typ_func.args.items())
self_arg_name, self_arg_type = method_args[0]
self._unify(self_arg_type, typ_self,
node.loc, None)
typ_arity = typ_func.arity() - 1
typ_args = OrderedDict(method_args[1:])
typ_optargs = typ_func.optargs
typ_ret = typ_func.ret
passed_args = dict()

View File

@ -378,11 +378,19 @@ class LLVMIRGenerator:
llfunty = self.llty_of_type(typ, bare=True)
llfun = ll.Function(self.llmodule, llfunty, name)
llretty = self.llty_of_type(typ.ret, for_return=True)
llretty = self.llty_of_type(typ.find().ret, for_return=True)
if self.needs_sret(llretty):
llfun.args[0].add_attribute('sret')
return llfun
def get_function_with_undef_env(self, typ, name):
llfun = self.get_function(typ, name)
llclosure = ll.Constant(self.llty_of_type(typ), [
ll.Constant(llptr, ll.Undefined),
llfun
])
return llclosure
def map(self, value):
if isinstance(value, (ir.Argument, ir.Instruction, ir.BasicBlock)):
return self.llmap[value]
@ -408,19 +416,10 @@ class LLVMIRGenerator:
def emit_attribute_writeback(self):
llobjects = defaultdict(lambda: [])
for obj_id in self.embedding_map.iter_objects():
obj_ref = self.embedding_map.retrieve_object(obj_id)
if isinstance(obj_ref, (pytypes.FunctionType, pytypes.MethodType,
pytypes.BuiltinFunctionType)):
continue
elif isinstance(obj_ref, type):
_, typ = self.embedding_map.retrieve_type(obj_ref)
else:
typ, _ = self.embedding_map.retrieve_type(type(obj_ref))
for obj_id, obj_ref, obj_typ in self.embedding_map.iter_objects():
llobject = self.llmodule.get_global("O.{}".format(obj_id))
if llobject is not None:
llobjects[typ].append(llobject.bitcast(llptr))
llobjects[obj_typ].append(llobject.bitcast(llptr))
llrpcattrty = self.llcontext.get_identified_type("A")
llrpcattrty.elements = [lli32, llptr, llptr]
@ -695,8 +694,8 @@ class LLVMIRGenerator:
llglobal = self.llmodule.get_global(name)
else:
llglobal = ll.GlobalVariable(self.llmodule, llty, name)
llglobal.linkage = "private"
if llvalue is not None:
llglobal.linkage = "private"
llglobal.initializer = llvalue
return llglobal
@ -705,7 +704,7 @@ class LLVMIRGenerator:
llty = self.llty_of_type(typ).pointee
return self.get_or_define_global("C.{}".format(typ.name), llty)
def get_global_closure(self, typ, attr):
def get_global_closure_ptr(self, typ, attr):
closure_type = typ.attributes[attr]
assert types.is_constructor(typ)
assert types.is_function(closure_type) or types.is_rpc(closure_type)
@ -713,7 +712,13 @@ class LLVMIRGenerator:
return None
llty = self.llty_of_type(typ.attributes[attr])
llclosureptr = self.get_or_define_global("F.{}.{}".format(typ.name, attr), llty)
return self.get_or_define_global("F.{}.{}".format(typ.name, attr), llty)
def get_global_closure(self, typ, attr):
llclosureptr = self.get_global_closure_ptr(typ, attr)
if llclosureptr is None:
return None
# LLVM's GlobalOpt pass only considers for SROA the globals that
# are used only by GEPs, so we have to do this stupid hack.
llenvptr = self.llbuilder.gep(llclosureptr, [self.llindex(0), self.llindex(0)])
@ -721,12 +726,12 @@ class LLVMIRGenerator:
return [llenvptr, llfunptr]
def load_closure(self, typ, attr):
llclosureptrs = self.get_global_closure(typ, attr)
if llclosureptrs is None:
llclosureparts = self.get_global_closure(typ, attr)
if llclosureparts is None:
return ll.Constant(llunit, [])
# See above.
llenvptr, llfunptr = llclosureptrs
llenvptr, llfunptr = llclosureparts
llenv = self.llbuilder.load(llenvptr)
llfun = self.llbuilder.load(llfunptr)
llclosure = ll.Constant(ll.LiteralStructType([llenv.type, llfun.type]), ll.Undefined)
@ -735,10 +740,10 @@ class LLVMIRGenerator:
return llclosure
def store_closure(self, llclosure, typ, attr):
llclosureptrs = self.get_global_closure(typ, attr)
assert llclosureptrs is not None
llclosureparts = self.get_global_closure(typ, attr)
assert llclosureparts is not None
llenvptr, llfunptr = llclosureptrs
llenvptr, llfunptr = llclosureparts
llenv = self.llbuilder.extract_value(llclosure, 0)
llfun = self.llbuilder.extract_value(llclosure, 1)
self.llbuilder.store(llenv, llenvptr)
@ -1343,6 +1348,12 @@ class LLVMIRGenerator:
llty = self.llty_of_type(typ)
if types.is_constructor(typ) or types.is_instance(typ):
if types.is_instance(typ):
# Make sure the class functions are quoted, as this has the side effect of
# initializing the global closures.
self._quote(type(value), typ.constructor,
lambda: path() + ['__class__'])
llglobal = None
llfields = []
for attr in typ.attributes:
@ -1359,8 +1370,18 @@ class LLVMIRGenerator:
self.llobject_map[value_id] = llglobal
else:
llfields.append(self._quote(getattr(value, attr), typ.attributes[attr],
lambda: path() + [attr]))
attrvalue = getattr(value, attr)
is_class_function = (types.is_constructor(typ) and
types.is_function(typ.attributes[attr]) and
not types.is_c_function(typ.attributes[attr]))
if is_class_function:
attrvalue = self.embedding_map.specialize_function(typ.instance, attrvalue)
llattrvalue = self._quote(attrvalue, typ.attributes[attr],
lambda: path() + [attr])
llfields.append(llattrvalue)
if is_class_function:
llclosureptr = self.get_global_closure_ptr(typ, attr)
llclosureptr.initializer = llattrvalue
llglobal.initializer = ll.Constant(llty.pointee, llfields)
llglobal.linkage = "private"
@ -1400,12 +1421,8 @@ class LLVMIRGenerator:
# RPC and C functions have no runtime representation.
return ll.Constant(llty, ll.Undefined)
elif types.is_function(typ):
llfun = self.get_function(typ.find(), self.embedding_map.retrieve_function(value))
llclosure = ll.Constant(self.llty_of_type(typ), [
ll.Constant(llptr, ll.Undefined),
llfun
])
return llclosure
return self.get_function_with_undef_env(typ.find(),
self.embedding_map.retrieve_function(value))
elif types.is_method(typ):
llclosure = self._quote(value.__func__, types.get_method_function(typ),
lambda: path() + ['__func__'])

View File

@ -0,0 +1,22 @@
# RUN: %python -m artiq.compiler.testbench.embedding %s
from artiq.language.core import *
from artiq.language.types import *
class a:
@kernel
def f(self):
print(self.x)
return None
class b(a):
x = 1
class c(a):
x = 2
bi = b()
ci = c()
@kernel
def entrypoint():
bi.f()
ci.f()