[Relax][PyTorch] Add support for torch.export.ExportedProgram in Relax PyTorch Frontend

[mshr-h]

Part of #17346
To keep the PR small, this implementation only supports the Core ATen operators used in AlexNet. I'll add more support after the PR is merged.

I ran the AlexNet model from Torchvision and confirmed that the inference outputs from PyTorch and Relax match.
Here's a repro.

test_alexnet.py

import torch
from torch.export import export
from tvm.contrib.download import download_testdata
from torchvision.models import get_model, get_model_weights
from torchvision.io import read_image
import tvm
from tvm import relax
import tvm.testing
from tvm.relax.frontend.torch import from_exported_program

# prepare sample image
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_name = "cat.png"
img_path = download_testdata(img_url, img_name, module="data")
image_tensor = read_image(img_path)

# load model from torchvision
model_name = "alexnet"
torch_model = get_model(model_name, weights="DEFAULT").eval()
weights = get_model_weights(model_name).DEFAULT
transforms = weights.transforms()

batch = transforms(image_tensor).unsqueeze(0)
example_args = (batch,)

# PyTorch
exported_program = export(torch_model, args=example_args)
expected: torch.Tensor = exported_program.module()(*example_args)

# Relax
target = "llvm"
dev = tvm.cpu()
mod = from_exportedprogram(exported_program)
mod = tvm.relax.transform.DecomposeOpsForInference()(mod)
exe = relax.build(mod, target=target)
vm = relax.VirtualMachine(exe, dev)
tvm_args = [tvm.nd.from_dlpack(x.contiguous()) for x in example_args]
tvm_output = vm["main"](*tvm_args)[0]
actual = torch.from_numpy(tvm_output.numpy())

# see if the outputs match
print(f"allclose: {torch.allclose(actual, expected, rtol=1e-5, atol=1e-5, equal_nan=True)}")

[mshr-h]

It seems like the CI is using the old docker image tlcpack/ci-cpu:20240105-165030-51bdaec6, which contains PyTorch 2.0. But the torch.export is introduced in PyTorch 2.1 so we need to upgrade it.
The dockerfile is already updated in #17338 so I guess all we have to do is to upload the new docker image to tlcpack/ci-cpu and update the image name in the jenkins file at ci/jenkins/unity_jenkinsfile.groovy.
cc @yongwww @tqchen

[tqchen]

we need to update https://github.com/apache/tvm/blob/main/ci/jenkins/docker-images.ini

[mshr-h]

Looks like it's still using old docker image.
https://ci.tlcpack.ai/blue/organizations/jenkins/tvm-unity/detail/PR-17396/8/pipeline
cc @tqchen

[tqchen]

Ah, https://github.com/apache/tvm/blob/main/ci/jenkins/unity_jenkinsfile.groovy#L35 needs to be manually updated for this file

[mshr-h]

@tvm-bot rerun

[mshr-h]

Can you take a look at it? Thanks!
@yongwww @tqchen

[tqchen]

overall LGTM, just a minor namimg nit. Thanks @mshr-h !