Esempio n. 1
0
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)
Esempio n. 2
0
 def __init__(self):
     self.cu = ATenCompilationUnit()
     self.tc_lang = None
     self.compilation_cache = {}