This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| from functorch import make_fx | |
| from torch.nn.utils import _stateless | |
| import torch_mlir | |
| from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend | |
| class Foo(torch.nn.Module): |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| from functorch import make_fx | |
| from torch.nn.utils import _stateless | |
| import torch_mlir | |
| class Foo(torch.nn.Module): | |
| def __init__(self): |
pashu123
/ forward.mlir
Created
May 5, 2022 10:27
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #loc0 = loc(unknown) | |
| module attributes {torch.debug_module_name = "forward"} { | |
| func private @__torch__.torch.fx.graph_module.forward.__code_getter(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module.forward"> loc(unknown)) -> !torch.str { | |
| %1 = torch.prim.GetAttr %arg0["_code"] : !torch.nn.Module<"__torch__.torch.fx.graph_module.forward"> -> !torch.str loc(#loc0) | |
| return %1 : !torch.str loc(#loc0) | |
| } loc(#loc0) | |
| func private @__torch__.torch.fx.graph_module.forward.forward(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module.forward"> loc(unknown), %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[3,3],f32>} loc(unknown), %arg2: !torch.tensor {torch.type_bound = !torch.vtensor<[3],f32>} loc(unknown), %arg3: !torch.tensor {torch.type_bound = !torch.vtensor<[3,3],f32>} loc(unknown)) -> !torch.tuple<tensor, tensor> { | |
| %float-1.000000e-02 = torch.constant.float -1.000000e-02 loc(#loc1) | |
| %true_0 = torch.constant.bool true loc(#loc2) | |
| %false = torch.constant.bool false loc(#loc3) |
pashu123
/ bert_training.mlir
Created
May 9, 2022 11:02
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| module attributes {torch.debug_module_name = "forward"} { | |
| func private @__torch__.torch.fx.graph_module.forward.__code_getter(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module.forward">) -> !torch.str { | |
| %2 = torch.prim.GetAttr %arg0["_code"] : !torch.nn.Module<"__torch__.torch.fx.graph_module.forward"> -> !torch.str | |
| return %2 : !torch.str | |
| } | |
| func private @__torch__.torch.fx.graph_module.forward.forward(%arg0: !torch.nn.Module<"__torch__.torch.fx.graph_module.forward">, %arg1: !torch.tensor {torch.type_bound = !torch.vtensor<[384],f32>}, %arg2: !torch.tensor {torch.type_bound = !torch.vtensor<[384],f32>}, %arg3: !torch.tensor {torch.type_bound = !torch.vtensor<[512,384],f32>}, %arg4: !torch.tensor {torch.type_bound = !torch.vtensor<[2,384],f32>}, %arg5: !torch.tensor {torch.type_bound = !torch.vtensor<[30522,384],f32>}, %arg6: !torch.tensor {torch.type_bound = !torch.vtensor<[384],f32>}, %arg7: !torch.tensor {torch.type_bound = !torch.vtensor<[384],f32>}, %arg8: !torch.tensor {torch.type_bo |
pashu123
/ torch_sharding.py
Created
May 13, 2022 15:43
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import os | |
| import torch | |
| import torch.distributed as dist | |
| import torch.multiprocessing as mp | |
| from torch.fx import symbolic_trace, replace_pattern | |
| def all_reduce(inp): | |
| dist.all_reduce(inp) |
pashu123
/ resnet50_script.py
Created
May 23, 2022 05:35
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from PIL import Image | |
| import requests | |
| import torch | |
| import torchvision.models as models | |
| from torchvision import transforms | |
| import sys | |
| from shark.shark_runner import SharkInference | |
| from torch.ao.quantization import get_default_qconfig | |
| import torch.quantization.quantize_fx as quantize_fx |
pashu123
/ minilm_jit.py
Created
May 26, 2022 14:18
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| from shark.shark_inference import SharkInference | |
| torch.manual_seed(0) | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") | |
| class MiniLMSequenceClassification(torch.nn.Module): |
pashu123
/ minilm_jit.py
Created
May 26, 2022 14:19
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| from shark.shark_inference import SharkInference | |
| torch.manual_seed(0) | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") | |
| class MiniLMSequenceClassification(torch.nn.Module): |
pashu123
/ minilm_jit.py
Created
May 26, 2022 14:27
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| from shark.shark_inference import SharkInference | |
| torch.manual_seed(0) | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased") | |
| class MiniLMSequenceClassification(torch.nn.Module): |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| %216 = torch.aten.linear %result0, %178, %179 : !torch.vtensor<[1,128,384],f32>, !torch.vtensor<[384,384],f32>, !torch.vtensor<[384],f32> -> !torch.vtensor<[1,128,384],f32> | |
| %217 = torch.aten.linear %result0, %177, %3 : !torch.vtensor<[1,128,384],f32>, !torch.vtensor<[384,384],f32>, !torch.vtensor<[384],f32> -> !torch.vtensor<[1,128,384],f32> | |
| %224 = torch.aten.linear %result0, %175, %176 : !torch.vtensor<[1,128,384],f32>, !torch.vtensor<[384,384],f32>, !torch.vtensor<[384],f32> -> !torch.vtensor<[1,128,384],f32> | |
| %250 = torch.aten.linear %249, %173, %174 : !torch.vtensor<[1,128,384],f32>, !torch.vtensor<[384,384],f32>, !torch.vtensor<[384],f32> -> !torch.vtensor<[1,128,384],f32> | |
| %252 = torch.aten.linear %result0_4, %169, %170 : !torch.vtensor<[1,128,384],f32>, !torch.vtensor<[1536,384],f32>, !torch.vtensor<[1536],f32> -> !torch.vtensor<[1,128,1536],f32> | |
| %254 = torch.aten.linear %253, %167, %168 : !torch.vtensor<[1,128,1536],f32>, !torch.vtensor<[384,1536],f32>, !torch.vtensor<[384],f32> |