load an MLIR file as a model using torch-mlir

Question

I have translated a torch model into MLIR using torch-mlir as in the github example:

def save_module():
    resnet18 = models.resnet18(pretrained=True)
    resnet18.eval()    
    module = torch_mlir.compile(resnet18, torch.ones(1, 3, 224, 224), output_type="torch")
    open("resnet18torch.mlir", "w").write(str(module))

After that, I changed the MLIR and now, I would like to do the opposite and load this MLIR as a module in my python code to continue to compile it like that:

src = open("resnet18torch.mlir", "r").read()
#transform src to module
backend = refbackend.RefBackendLinalgOnTensorsBackend()
compiled = backend.compile(module)
jit_module = backend.load(compiled)
predictions(resnet18.forward, jit_module.forward, img, labels)

I just need a line in the second code to load the MLIR as a module but I couldn't find anything about it on the internet. Does anyone know how to do it?

Answer 1

I finally figured out how to load a torch mlir and how to reuse it with the help of someone on the LLVM discord. Here is the usage with the torch_mlir example (torchscript_resnet18.py):

import torch_mlir

def load_module():
    #load the torch mlir
    src = open("resnet18torch.mlir", "r").read()
    with torch_mlir.ir.Context() as ctx:
        torch_mlir.dialects.torch.register_dialect(ctx)
        with torch_mlir.ir.Location.unknown() as loc:
            module = torch_mlir.ir.Module.parse(src)
    #translate the torch-mlir to the linalg-on-tensors dialect
    torch_mlir.run_pipeline_with_repro_report(module, "builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline)", "Lowering Torch Backend IR -> Linalg-on-Tensors Backend IR")

    #print(module)
    #compile the module
    backend = refbackend.RefBackendLinalgOnTensorsBackend()
    compiled = backend.compile(module)
    jit_module = backend.load(compiled)
    return jit_module

#save_module()
jit_module = load_module()
predictions(jit_module.forward, img, labels)