Dumping various passes of IRs in tvm

Hi, What are the different IRs that get generated when a Python model is compiled on a CPU? I understand that torch script gets converted to Relay IR and I know how to dump it too.

Could you please let me know what IRs gets generated after Relay IR, is it tensor expression IR? After tensor expression IR, is LLVM IR generated?

How do I dump the IRs and its various passes that get generated for the following Python program? Do we have currently have such an option? Basically I want to understand the various optimizations that the tvm compiler performs on various IRs that it generates.

import tvm import torch from tvm import relay from tvm.contrib import graph_runtime

class MultilayerLinearResidual(torch.nn.Module): def init(self, input_size, hidden_size, n_layers=1): super().init() self.layers = torch.nn.ModuleList() self.layers.append(torch.nn.Linear(input_size, hidden_size)) for __ in range(n_layers-1): self.layers.append(torch.nn.Linear(hidden_size, hidden_size))

def forward(self, x):
    for i, l in enumerate(self.layers):
        y = torch.relu(l(x))
        if i != 0:
            y += x
        x = y
    return x

def get_example(example_name): if example_name == ‘MultilayerLinearResidual’: Din = 64 H = 128 layers = 3 B = 1 model = MultilayerLinearResidual(Din, H, n_layers=layers) x = torch.randn(B, Din) else: raise ValueError(‘Example {} is not defined’.format(example_name)) return model, x def main(): example_name = ‘MultilayerLinearResidual’ model, x = get_example(example_name) shape_list = [(‘input’, x.shape)] print(’\n\nTRACED GRAPH:’) traced_model = torch.jit.trace(model, x).eval() print(traced_model.graph) traced_mod, traced_params = relay.frontend.from_pytorch(traced_model, shape_list) target = ‘llvm’ target_host = ‘llvm’ ctx = tvm.cpu(0) with relay.build_config(opt_level=3): graph, lib, params = relay.build(traced_mod, target=target, target_host=target_host, params=traced_params)

m = graph_runtime.create(graph, lib, ctx)
input_nd = x.numpy().astype('float32')
m.set_input('input', tvm.nd.array(input_nd))
m.set_input(**params)
m.run()
tvm_output = m.get_output(0)
y = traced_model(x)
print("Relay graph", graph)
print("traced model is ", y)