Using TorchScript Converters in the FX Frontend

[narendasan]

TS Converters for FX

TL;DR

We want to be able to bootstrap the FX frontend with all the work we have done for Torch Script writing aten -> TRT converters.

Goal(s)

Right now TorchScript supports many more operations natively that FX.

Usecases

Proposed APIs / UX

The user would not know the difference as long as they use the proxy tensor tracer

Example Workflow

TBD

Limitations

TBD

Internal Implementation

Design

We need to replicate the ConversionCtx from the TorchScript frontend in FX, being able to package a handle for the network, populate the args, package the node, and then read the resultant layers out of the context map.

Extensions Required to Core API implementations

We need to add a bunch of wrappers to the Python API for the FX frontend to call

Data Structures

TBD

Details specific for TorchScript Support

N/A

Details specific for FX support

N/A

Implementation Phases

Prototype

• Write a FX converter to TS converter glue layer to package inputs to the converter
• Write a ConversionCtx -> FX extractor

MVP (1.5.0 / 1.6.0)

• One or two converters implemented in TS but called from FX

Extension Phase 1 (<TARGET RELEASE VERSION>)

• Stubs for the rest of the conversion lib

[narendasan]

IMO this feature should be deprioritized for the torch2trt -> FX bootstrap work @apbose is working on. That is a much lighter lift for similar outcomes