class TcCompilationUnit(object):
def __init__(self):
self.cu = ATenCompilationUnit()
self.tc_lang = None
self.compilation_cache = {}
def define(self, tc_lang):
self.tc_lang = tc_lang
self.cu.define(tc_lang)
# we could have multiple TC definitions and want to run one of them
def compile(self, name, inputs, **kwargs):
# append the language so that we can use it for creating Autotuner object
# to load options cache
kwargs["tc_lang"] = self.tc_lang
if "type" not in kwargs:
kwargs["type"] = 'forward'
options = get_options_from_kwargs(name, *inputs, **kwargs)
handle = self.cu.compile(name, inputs, options)
return handle
def run(self, handle, name, inputs, **kwargs):
outputs = []
if "outputs" in kwargs and kwargs["outputs"] is not None:
outputs = kwargs["outputs"]
if not isinstance(outputs, list):
outputs = [outputs]
self.cu.run(name, inputs, outputs, handle)
return outputs
def compile_and_run(self, name, inputs, **kwargs):
handle = self.compile(name, inputs, **kwargs)
return self.run(handle, name, inputs, **kwargs)
def __init__(self):
self.cu = ATenCompilationUnit()
self.tc_lang = None
self.compilation_cache = {}