Support AOTInductor compilation and packaging with Torch-TensorRT modules

[peri044]

TL;DR

We want to use AOTInductor to compile and package Torch-TensorRT compiled modules to produce shared object files which can be used in Pythonless deployment. This could be especially beneficial for Jetson platforms.
Reference: https://pytorch.org/docs/main/torch.compiler_aot_inductor.html

Goal(s)

Ensure we support AOTInductor to export Torch-TRT compiled programs using torch._inductor.aoti_compile_and_package.

Usecases

• Embedded deployment
• Cases where people want stable and actively maintained C++ deployment solution (We have Torchscript currently but it is deprecated).

Proposed APIs/UX

We could probably get this feature for free (atleast from a frontend perspective). Here's the workflow I have in my mind.

import os
import torch

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(10, 16)
        self.relu = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(16, 1)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x

with torch.no_grad():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = Model().to(device=device)
    example_inputs=(torch.randn(8, 10, device=device),)
    batch_dim = torch.export.Dim("batch", min=1, max=1024)
    # [Optional] Specify the first dimension of the input x as dynamic.
    exported = torch.export.export(model, example_inputs, dynamic_shapes={"x": {0: batch_dim}})
    # Compile with Torch-TRT and export
    trt_gm = torch_tensorrt.dynamo.compile(model, example_inputs, enabled_precisions={torch.float16})
    trt_ep = torch.export.export(trt_gm, example_inputs, dynamic_shapes={"x": {0: batch_dim}})
    
    output_path = torch._inductor.aoti_compile_and_package(
        trt_ep,
        example_inputs,
        # [Optional] Specify the generated shared library path. If not specified,
        # the generated artifact is stored in your system temp directory.
        package_path=os.path.join(os.getcwd(), "model.pt2"),
    )

We have recently added the re-export feature #3262 which uses torch.export to create first class ExportedPrograms from Torch-TensorRT compilation. As long as these exported programs are valid and torchbind objects are supported by AOTI, I would think the above workflow would work.

Limitations

• Cross compilation might not be supported. (cross compile for windows)

Considerations

• In the case of custom ops which use TRT plugins, do the plugin library (.so) get linked to this ?
• In the case of fallback, does AOTInductor recompile the Pytorch segments with inductor ?

Data Structures

None at this time

Details specific for TorchScript Support

Torchscript support won't be affected.

Details specific for FX support

All this will be applicable to dynamo frontend. ir=fx isn't active so it doesn't get affected.

Implementation Phases

Prototype -

MVP (<TARGET RELEASE VERSION>)

Extension Phase 1 (<TARGET RELEASE VERSION>)

Extension Phase 2 (<TARGET RELEASE VERSION>)

[peri044]

@angelayi Could you take a look at let us know if the above workflow is a potential solution ?