[MHLO] bert-tiny and resnet18 example from torchscript to mhlo

Co-authored-by: Bairen Yi <[email protected]>
Co-authored-by: Jiawei Wu <[email protected]>
Co-authored-by: Tianyou Guo <[email protected]>
Co-authored-by: Xu Yan <[email protected]>
Co-authored-by: Ziheng Jiang <[email protected]>

Author: Vremold

examples/torchscript_mhlo_backend_resnet.py

@@ -0,0 +1,13 @@
+import torch
+import torchvision.models as models
+import torch_mlir
+
+model = models.resnet18(pretrained=True)
+model.eval()
+data = torch.randn(2,3,200,200)
+out_mhlo_mlir_path = "./resnet18_mhlo.mlir"
+
+module = torch_mlir.compile(model, data, output_type=torch_mlir.OutputType.MHLO, use_tracing=False)
+with open(out_mhlo_mlir_path, "w", encoding="utf-8") as outf:
+    outf.write(str(module))
+print(f"MHLO IR of resent18 successfully written into {out_mhlo_mlir_path}")

examples/torchscript_mhlo_backend_tinybert.py

@@ -0,0 +1,23 @@
+import torch
+import torch_mlir
+
+from transformers import BertForMaskedLM
+
+# Wrap the bert model to avoid multiple returns problem
+class BertTinyWrapper(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.bert = BertForMaskedLM.from_pretrained("prajjwal1/bert-tiny", return_dict=False)
+    
+    def forward(self, data):
+        return self.bert(data)[0]
+
+model = BertTinyWrapper()
+model.eval()
+data = torch.randint(30522, (2, 128))
+out_mhlo_mlir_path = "./bert_tiny_mhlo.mlir"
+
+module = torch_mlir.compile(model, data, output_type=torch_mlir.OutputType.MHLO, use_tracing=True)
+with open(out_mhlo_mlir_path, "w", encoding="utf-8") as outf:
+    outf.write(str(module))
+print(f"MHLO IR of tiny bert successfully written into {out_mhlo_mlir_path}")