def start_define_computation(self, computation_decl):
self.exop_codegen.append("class {}(HetrLocals, ConvLocals):",
computation_decl.computation_op.name)
with indenting(self.exop_codegen):
self.exop_codegen.append("def __init__(self, **kwargs):")
with indenting(self.exop_codegen):
if is_tracing_enabled():
self.exop_codegen.append("""
self.__profiler_start__ = list()
self.__profiler_stop__ = list()
""")
self.exop_codegen.append('super({}, self).__init__(**kwargs)',
computation_decl.computation_op.name)
for exop in computation_decl.exop_block:
output_decl = exop.output_decls[0] if len(
exop.output_decls) > 0 else None
# TODO better way to deal with multiple values
self.exop_codegen.exop = exop
self.exop_codegen.allocate_op(exop.op, output_decl,
*exop.input_decls)
self.exop_codegen.endl()
self.exop_codegen.indent(1)
self.exop_codegen.append("def __call__(self):")
self.exop_codegen.indent(1)
self.codegen_define_length = self.exop_codegen.code_length
def generate_op_pre(self, op):
# exop = self.exop
# self.append("\n# {} pre", exop.name)
# for input_decl in exop.input_decls:
# input_decl_name = 'a_'+input_decl.tensor.tensor_name
# self.append("# arg {}", input_decl_name)
# for output_decl in exop.output_decls:
# output_decl_name = 'a_'+output_decl.tensor.tensor_name
# self.append("# output_decl {}", val_name)
if is_tracing_enabled():
self.append("self.__profiler_start__.append(monotonic())")
def generate_op_post(self, op):
# exop = self.exop
# self.append("print('{{}}'.format('{}'))", op.name)
# for input_decl in exop.input.decls:
# input_decl_name = 'a_'+input_decl.tensor.tensor_name
# self.append("print(' arg {} = {{}}'.format({}))", input_decl_name, arg_name)
# for val in exop.output_decls:
# output_decl_name = 'a_'+val.tensor.tensor_name
# self.append("# output_decl {}", output_decl_name)
# self.append("print(' output_decl {} = {{}}'.format({}))", \
# output_decl_name, output_decl_name)
if is_tracing_enabled():
self.append("self.__profiler_stop__.append(monotonic())")
def __call__(self, *args, **kwargs):
"""
Executes the computation passing args in to the function.
"""
args = self.unpack_args_or_feed_dict(args, kwargs)
# TODO Should this be automatic?
self.transformer.initialize()
# Get the parameters to the device
self.transformer.host_to_device(self, self.computation_op.parameters,
args)
self.executor()
# Generate a timeline for the computation
if is_tracing_enabled():
self.generate_profile(self.executor.__profiler_start__,
self.executor.__profiler_stop__)
# TODO Should copy this out of the device to a destination when it is not scalar
def value(op):
"""
Returns the computed value of op, or None if it has no value.
:param op:
:return: Return value for op.
"""
if op.is_tensor_op and 'skip_returns' not in op.metadata:
return self.transformer.device_to_host(self, op)
else:
return None
if isinstance(self.computation_op.returns, Op):
return value(self.computation_op.returns)
elif isinstance(self.computation_op.returns,
(collections.Sequence, OrderedSet)):
return tuple(value(op) for op in self.computation_op.returns)
elif isinstance(self.computation_op.returns, collections.Set):
result = dict()
for op in self.computation_op.returns:
result[op] = value(op)
return result
else:
return None
