[Tripy] Permit eval while tracing, but do not update trace.
#443
Conversation
|
I am not fully certain how the storage mechanism described in the subsequent comment on #409 should work. If the main issue is recompilation caused by the fact that we not updating the trace graph if we evaluate while tracing, perhaps one approach could be to insert the Storage op anyway but restore the original graph after compilation. |
tripy/tripy/frontend/tensor.py
Outdated
| @@ -265,8 +269,15 @@ def __repr__(self) -> str: | |||
|
|
|||
| data_list = self.tolist() | |||
|
|
|||
| assert isinstance(self.trace_tensor.producer, Storage) | |||
| data_shape = self.trace_tensor.producer.shape | |||
| if isinstance(self.trace_tensor.producer, Storage): | |||
There was a problem hiding this comment.
Does the memref have shape information? If so, it would probably be cleaner to use the memref directly rather than tolist().
tripy/tripy/frontend/tensor.py
Outdated
| if not self.trace_tensor.is_compile_tracer: | ||
| Storage.build_internal([], [self.trace_tensor], data) |
There was a problem hiding this comment.
The problem with this is if we evaluate multiple tensors in a compiled graph, we'll get quadratic time complexity. What I suggested in #409 (comment) is to store the evaluated result but only use it for evaluation and not tracing. I'm not sure exactly how that would work, but I expect it would require changes to how we trace during eval().
Could you elaborate on this approach? Sounds promising. |
|
I haven't tried implementing this yet, but if we could record which ops we changed during the trace, we could avoid recompilation during the trace and then change them back afterwards. |
… not update its trace tensor, however
…t to avoid slowing down other evaluations
slyubomirsky
force-pushed
the
no-error-for-eval-while-tracing
branch
from
dbeccd0 to
f89fd44
Compare
| if REVERT_GRAPH_AFTER_COMPILING is not None: | ||
| REVERT_GRAPH_AFTER_COMPILING.append((self.trace_tensor, self.trace_tensor.producer)) |
There was a problem hiding this comment.
This might be a bit of a clumsy approach to keeping around the old producers (required a global var from the compiler). I wasn't sure where else the data structure could reside, since evaluation is handled as a method on tensors. It's a tricky issue because the tensor evaluated could be anywhere in the middle of the graph and the interface to compile is an opaque Callable.
There was a problem hiding this comment.
What about storing it on the trace tensors themselves? original_producer or something like that?
There was a problem hiding this comment.
I guess in that case we would need to do a DFS from the output and swap the producer before stepping further
There was a problem hiding this comment.
Yeah searching the graph might be necessary if there isn't anywhere else we could record the tensors that have been evaluated.
|
Closing per discussion: We've decided that there aren't any cases we could think of where evaluating while tracing is useful, especially since evaluating without using the result (e.g., just printing) already worked. We can revisit if such a case does present itself. |
Addresses #409. Evaluation while tracing still gives a warning but does not alter the graph, so it does not produce any errors.