MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / trace_without_host

Method trace_without_host

imperative/python/megengine/jit/tracing.py:379–487  ·  view source on GitHub ↗
(self, *args, **kwargs)

Source from the content-addressed store, hash-verified

377 return outputs
378
379 def trace_without_host(self, *args, **kwargs):
380 from ..traced_module.pytree import tree_flatten, SUPPORTED_LEAF_CLS
381 from ..module import Module
382 from ..utils.module_utils import get_expand_structure
383 from ..tensor import Tensor
384 from ..optimizer import Optimizer
385
386 assert self.without_host and not self.overall
387 global active_trace
388 symbolic_shape = None
389 outputs = None
390 if self.traced and self.third_party_backend:
391 return self.compile_and_exec(*args, **kwargs)
392 try:
393 active_trace = self
394 self._trace.enter()
395 if self._trace.compiled():
396 arglist = self.flatten_inputs(*args, **kwargs)
397 idx = 0
398 inp_dict = {}
399 for a in arglist:
400 if isinstance(a, RawTensor):
401 inp_dict[self.arg_list[idx]] = a
402 idx += 1
403 self._trace.put_datas(inp_dict)
404 outlist = []
405 for i in self.out_list:
406 if i == -1:
407 if not hasattr(self, "outdef"):
408 outlist.append(None)
409 else:
410 outlist.append(self._trace.get_data(i))
411 keep_vars = []
412 for i in self.keeped_activation:
413 keep_vars.append(self._trace.get_data(i))
414
415 outputs = (
416 self.outdef.unflatten(outlist)
417 if hasattr(self, "outdef")
418 else outlist
419 )
420 if keep_vars:
421 return outputs, keep_vars
422 else:
423 return outputs
424
425 arg_list = self.flatten_inputs(*args, **kwargs)
426 for i, arg in enumerate(arg_list):
427 arg_list[i]._reset(get_marked_input_tensor(self.input_num, arg))
428 self.arg_list.append(self.input_num)
429 self.input_num += 1
430 arg_ids = [id(i) for i in arg_list]
431 del arg_list
432 self.input_need_update_dict = {}
433 origin_reset = Tensor._reset
434
435 def tensor_reset_hook(obj, other):
436 if id(obj) in arg_ids:

Callers 1

__call__Method · 0.95

Calls 13

compile_and_execMethod · 0.95
flatten_inputsMethod · 0.95
setup_envMethod · 0.95
unset_envMethod · 0.95
_process_outputsMethod · 0.95
tree_flattenFunction · 0.85
strFunction · 0.85
set_symbolic_shapeFunction · 0.50
enterMethod · 0.45
appendMethod · 0.45
unflattenMethod · 0.45
_resetMethod · 0.45

Tested by

no test coverage detected