dig into torch-mlir

dig_torch_mlir.py

import torch
from torch import nn

# import torch_mlir
from torch_mlir.passmanager import PassManager
from torch_mlir.dialects.torch.importer.jit_ir import ClassAnnotator, ModuleBuilder


torch.manual_seed(42)

module = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
    nn.LogSoftmax(1),
)
module = torch.jit.script(module)
arg = torch.ones(32, 784)
# module = torch.jit.trace(module, torch.ones(32, 784))
# module = torch.jit.trace_module(module, dict(forward=arg))

### module = torch_mlir.compile(module, torch.ones(32, 784), output_type="raw")
class_annotator = ClassAnnotator()
class_annotator.exportNone(module._c._type())
class_annotator.exportPath(module._c._type(), ["forward"])
annotations = [None, (arg.shape, arg.dtype, True)]  # self
class_annotator.annotateArgs(module._c._type(), ["forward"], annotations)
mb = ModuleBuilder()
mb.import_module(module._c, class_annotator)
module = mb.module

# module = torch_mlir.compile(module, torch.ones(32, 784), output_type="torch")
with module.context:
    # option_string = "{backend-legal-ops=aten.warn}"
    # pm = PassManager.parse(f"builtin.module(torchscript-module-to-torch-backend-pipeline{option_string})")
    pm = PassManager.parse("builtin.module(torchscript-module-to-torch-backend-pipeline)")
    pm.run(module.operation)

### module = torch_mlir.compile(module, torch.ones(32, 784), output_type="linalg_on_tensors")
with module.context:
    pm = PassManager.parse("builtin.module(torch-backend-to-linalg-on-tensors-backend-pipeline)")
    pm.run(module.operation)

# from torch_mlir_e2e_test.linalg_on_tensors_backends.refbackend import LOWERING_PIPELINE
LOWERING_PIPELINE = "builtin.module(" + ",".join([
    "func.func(refback-generalize-tensor-pad)",
    "func.func(linalg-fuse-elementwise-ops)",
    # Bufferize.
    "func.func(scf-bufferize)",
    "func.func(tm-tensor-bufferize)",
    "func.func(empty-tensor-to-alloc-tensor)",
    "func.func(linalg-bufferize)",
    "func-bufferize",
    "arith-bufferize",
    "refback-mlprogram-bufferize",
    "func.func(tensor-bufferize)",
    "func.func(finalizing-bufferize)",
    "func.func(buffer-deallocation)",
]) + ")"
with module.context:
    pm = PassManager.parse(LOWERING_PIPELINE)
    pm.run(module.operation)
# print(module)

striped = module.operation.get_asm(large_elements_limit=10, enable_debug_info=False)
print(striped)