Skip to content

Scoped tracing #3906

@ezyang

Description

@ezyang

General strategy:

  • We're going to introduce a new, special attribute, "impl_graph", available for all nodes, which contains a subgraph that could be inlined inplace of this node and preserve the semantics. The intent is that when we start tracing in a "scope", we'll create a single node, and then start tracing into the impl_graph subgraph for internal tracing that occurs. We call a node with "impl_graph" set an "aggregate node"
  • How do we decide what graph we're tracing into? We're going to switch forward tracing to use TLS. Now the thread local state variable tracks what impl_graph you should be tracing into at any given point in time. In lieu of doing this properly, a global variable should suffice: https://github.com/pytorch/pytorch/pull/3016/files is a good place to start, as it does tracing in this way.
  • How do we export a graph with these aggregate nodes? We first test if we know how to directly export the node. If so, we export and discard "impl_graph". Otherwise, we inline "impl_graph" and rerun the export on it.

cc @houseroad @spandantiwari @lara-hdr @BowenBao @neginraoof

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNot as big of a feature, but technically not a bug. Should be easy to fixmodule: onnxRelated to torch.onnxonnx-needs-infoneeds information from the author / reporter before ONNX team can take actiontriagedThis issue has been looked at a team member, and triaged and prioritized into an appropriate module

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions