Introduce torch.onnx.dynamo_export API

[abock]

This is the first phase of the new ONNX exporter API for exporting from TorchDynamo and FX, and represents the beginning of a new era for exporting ONNX from PyTorch.

The API here is a starting point upon which we will layer more capability and expressiveness in subsequent phases. This first phase introduces the following into torch.onnx:

dynamo_export(
    model: torch.nn.Module,
    /,
    *model_args,
    export_options: Optional[ExportOptions] = None,
    **model_kwargs,
) -> ExportOutput:
    ...

class ExportOptions:
    opset_version: Optional[int] = None
    dynamic_shapes: Optional[bool] = None
    logger: Optional[logging.Logger] = None

class ExportOutputSerializer(Protocol):
    def serialize(
        self,
        export_output: ExportOutput,
        destination: io.BufferedIOBase,
    ) -> None:
        ...

class ExportOutput:
    model_proto: onnx.ModelProto

    def save(
        self,
        destination: Union[str, io.BufferedIOBase],
        *,
        serializer: Optional[ExportOutputSerializer] = None,
    ) -> None:
        ...

In addition to the API in the first commit on this PR, we have a few experiments for exporting Dynamo and FX to ONNX that this PR rationalizes through the new Exporter API and adjusts tests to use the new API.

• A base FXGraphModuleExporter exporter from which all derive:
  • DynamoExportExporter: uses dynamo.export to acquire FX graph
  • DynamoOptimizeExporter: uses dynamo.optimize to acquire FX graph
  • FXSymbolicTraceExporter: uses FX symbolic tracing

The dynamo_export API currently uses DynamoOptimizeExporter.

Next Steps (subsequent PRs):

• Combine DynamoExportExporter and DynamoOptimizeExporter into a single DynamoExporter.
• Make it easy to test FXSymbolicTraceExporter through the same API; eventually FXSymbolicTraceExporter goes away entirely when the Dynamo approach works for large models. We want to keep FXSymbolicTraceExporter around for now for experimenting and internal use.
• Parameterize (on ExportOptions) and consolidate Dynamo exporter tests.
  • This PR intentionally leaves the existing tests unchanged as much as possible except for the necessary plumbing.
• Subsequent API phases:
  • Diagnostics
  • Registry, dispatcher, and Custom Ops
  • Passes
  • Dynamic shapes

Fixes #94774

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/97920

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit f8ed7cc:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[titaiwangms]

Good work

[shubhambhokare1]

LGTM 👍

[abock]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / linux-bionic-py3.8-clang9-slow / filter

Details for Dev Infra team

Raised by workflow job

[abock]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here