github/SHARK-Studio
1
1
Fork 0
mirror of https://github.com/nod-ai/SHARK-Studio.git synced 2026-01-09 13:57:54 -05:00
Code Issues Packages Projects Releases Wiki Activity

Divide iree_utils and do module imports on function calls.

Browse Source
This commit is contained in:
Prashant Kumar
2022-06-22 11:13:17 +05:30
parent 08eda2ce35
commit e8aa105b2a
147 changed files with 4705 additions and 2662 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.github/workflows/nightly.yml vendored
View File

@@ -53,7 +53,7 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install flake8 pytest yapf toml
python -m pip install flake8 pytest toml
if [ -f requirements.txt ]; then pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/nightly/cpu -f https://github.com/llvm/torch-mlir/releases -f https://github.com/nod-ai/SHARK-Runtime/releases; fi
- name: Lint with flake8
run: |
@@ -61,7 +61,6 @@ jobs:
flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude shark.venv,lit.cfg.py
# exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude shark.venv,lit.cfg.py
yapf -i --style .style.yapf shark/*.py
- name: Build and validate the package
run: |

15
.github/workflows/python-format.yml vendored Normal file
View File

@@ -0,0 +1,15 @@
name: black-formatter
on: [pull_request]
jobs:
build:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v1
- name: Set up Python 3.10
uses: actions/setup-python@v1
with:
python-version: 3.10.5
- name: Install Black
run: pip install black
- name: Run formatter check on the entire project.
run: black --line-length 80 --check .

3
.github/workflows/test-models.yml vendored
View File

@@ -37,7 +37,7 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install flake8 pytest yapf toml
python -m pip install flake8 pytest toml
- name: Lint with flake8
run: |
@@ -45,7 +45,6 @@ jobs:
flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude lit.cfg.py
# exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics --exclude lit.cfg.py
yapf -i --style .style.yapf shark/*.py
- name: Validate Models
run: |

8
benchmarks/hf_model_benchmark.py
View File

@@ -6,16 +6,16 @@ parser.add_argument(
"--model_name",
type=str,
required=True,
help=
"Specifies name of HF model to benchmark. (For exmaple \"microsoft/MiniLM-L12-H384-uncased\""
help='Specifies name of HF model to benchmark. (For exmaple "microsoft/MiniLM-L12-H384-uncased"',
)
load_args, unknown = parser.parse_known_args()
if __name__ == "__main__":
model_name = load_args.model_name
test_input = torch.randint(2, (1, 128))
shark_module = SharkHFBenchmarkRunner(model_name, (test_input,),
jit_trace=True)
shark_module = SharkHFBenchmarkRunner(
model_name, (test_input,), jit_trace=True
)
shark_module.benchmark_c()
shark_module.benchmark_python((test_input,))
shark_module.benchmark_torch(test_input)

100
benchmarks/hf_transformer.py
View File

@@ -1,8 +1,12 @@
import torch
from shark.shark_runner import SharkBenchmarkRunner
from shark.shark_benchmark_runner import SharkBenchmarkRunner
from shark.parser import shark_args
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from onnxruntime.transformers.benchmark import run_pytorch, run_tensorflow, run_onnxruntime
from onnxruntime.transformers.benchmark import (
run_pytorch,
run_tensorflow,
run_onnxruntime,
)
from onnxruntime.transformers.huggingface_models import MODELS
from onnxruntime.transformers.benchmark_helper import ConfigModifier, Precision
import os
@@ -10,7 +14,6 @@ import psutil
class OnnxFusionOptions(object):
def __init__(self):
self.disable_gelu = False
self.disable_layer_norm = False
@@ -25,17 +28,13 @@ class OnnxFusionOptions(object):
class HuggingFaceLanguage(torch.nn.Module):
def __init__(self, hf_model_name):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
hf_model_name, # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -62,8 +61,16 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
)
self.model_name = model_name
model = HuggingFaceLanguage(model_name)
SharkBenchmarkRunner.__init__(self, model, input, dynamic, self.device,
jit_trace, from_aot, frontend)
SharkBenchmarkRunner.__init__(
self,
model,
input,
dynamic,
self.device,
jit_trace,
from_aot,
frontend,
)
def benchmark_torch(self, inputs):
use_gpu = self.device == "gpu"
@@ -74,10 +81,20 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
sequence_lengths = [inputs.shape[-1]]
cache_dir = os.path.join(".", "cache_models")
verbose = False
result = run_pytorch(use_gpu, [self.model_name], None, config_modifier,
Precision.FLOAT32, num_threads, batch_sizes,
sequence_lengths, shark_args.num_iterations, False,
cache_dir, verbose)
result = run_pytorch(
use_gpu,
[self.model_name],
None,
config_modifier,
Precision.FLOAT32,
num_threads,
batch_sizes,
sequence_lengths,
shark_args.num_iterations,
False,
cache_dir,
verbose,
)
print(
f"ONNX Pytorch-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)
@@ -92,10 +109,19 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
sequence_lengths = [inputs.shape[-1]]
cache_dir = os.path.join(".", "cache_models")
verbose = False
result = run_tensorflow(use_gpu, [self.model_name], None,
config_modifier, Precision.FLOAT32, num_threads,
batch_sizes, sequence_lengths,
shark_args.num_iterations, cache_dir, verbose)
result = run_tensorflow(
use_gpu,
[self.model_name],
None,
config_modifier,
Precision.FLOAT32,
num_threads,
batch_sizes,
sequence_lengths,
shark_args.num_iterations,
cache_dir,
verbose,
)
print(
f"ONNX TF-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)
@@ -105,7 +131,8 @@ class SharkHFBenchmarkRunner(SharkBenchmarkRunner):
print(
f"{self.model_name} is currently not supported in ORT's HF. Check \
https://github.com/microsoft/onnxruntime/blob/master/onnxruntime/python/tools/transformers/huggingface_models.py \
for currently supported models. Exiting benchmark ONNX.")
for currently supported models. Exiting benchmark ONNX."
)
return
use_gpu = self.device == "gpu"
num_threads = psutil.cpu_count(logical=False)
@@ -121,17 +148,34 @@ for currently supported models. Exiting benchmark ONNX.")
use_raw_attention_mask = True
model_fusion_statistics = {}
overwrite = False
model_source = "pt" #Either "pt" or "tf"
model_source = "pt" # Either "pt" or "tf"
provider = None
config_modifier = ConfigModifier(None)
onnx_args = OnnxFusionOptions()
result = run_onnxruntime(
use_gpu, provider, [self.model_name], None, config_modifier,
Precision.FLOAT32, num_threads, batch_sizes, sequence_lengths,
shark_args.num_iterations, input_counts, optimize_onnx,
validate_onnx, cache_dir, onnx_dir, verbose, overwrite,
disable_ort_io_binding, use_raw_attention_mask,
model_fusion_statistics, model_source, onnx_args)
use_gpu,
provider,
[self.model_name],
None,
config_modifier,
Precision.FLOAT32,
num_threads,
batch_sizes,
sequence_lengths,
shark_args.num_iterations,
input_counts,
optimize_onnx,
validate_onnx,
cache_dir,
onnx_dir,
verbose,
overwrite,
disable_ort_io_binding,
use_raw_attention_mask,
model_fusion_statistics,
model_source,
onnx_args,
)
print(
f"ONNX ORT-benchmark:{result[0]['QPS']} iter/second, Total Iterations:{shark_args.num_iterations}"
)

143
benchmarks/tests/test_benchmark.py
View File

@@ -1,19 +1,23 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
import torch
import tensorflow as tf
import numpy as np
import torchvision.models as models
from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel
from transformers import (
AutoModelForSequenceClassification,
BertTokenizer,
TFBertModel,
)
import importlib
import pytest
import unittest
torch.manual_seed(0)
gpus = tf.config.experimental.list_physical_devices('GPU')
gpus = tf.config.experimental.list_physical_devices("GPU")
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
tf.config.experimental.set_memory_growth(gpu, True)
##################### Tensorflow Hugging Face LM Models ###################################
MAX_SEQUENCE_LENGTH = 512
@@ -23,12 +27,11 @@ BATCH_SIZE = 1
tf_bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
class TFHuggingFaceLanguage(tf.Module):
def __init__(self, hf_model_name):
super(TFHuggingFaceLanguage, self).__init__()
# Create a BERT trainer with the created network.
@@ -36,7 +39,8 @@ class TFHuggingFaceLanguage(tf.Module):
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=tf_bert_input)
def forward(self, input_ids, attention_mask, token_type_ids):
@@ -47,15 +51,21 @@ def get_TFhf_model(name):
model = TFHuggingFaceLanguage(name)
tokenizer = BertTokenizer.from_pretrained(name)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
encoded_input = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
test_input = (encoded_input["input_ids"], encoded_input["attention_mask"],
encoded_input["token_type_ids"])
tf.convert_to_tensor(encoded_input[key]), 0
)
test_input = (
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
)
actual_out = model.forward(*test_input)
return model, test_input, actual_out
@@ -64,17 +74,13 @@ def get_TFhf_model(name):
class HuggingFaceLanguage(torch.nn.Module):
def __init__(self, hf_model_name):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
hf_model_name, # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -96,7 +102,6 @@ def get_hf_model(name):
class VisionModule(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
@@ -117,46 +122,56 @@ def get_vision_model(torch_model):
############################# Benchmark Tests ####################################
pytest_benchmark_param = pytest.mark.parametrize(
('dynamic', 'device'),
("dynamic", "device"),
[
pytest.param(False, 'cpu'),
pytest.param(False, "cpu"),
# TODO: Language models are failing for dynamic case..
pytest.param(True, 'cpu', marks=pytest.mark.skip),
pytest.param(False,
'gpu',
marks=pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")),
pytest.param(True,
'gpu',
marks=pytest.mark.skip),
pytest.param(True, "cpu", marks=pytest.mark.skip),
pytest.param(
False,
'vulkan',
"gpu",
marks=pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
),
),
pytest.param(True, "gpu", marks=pytest.mark.skip),
pytest.param(
False,
"vulkan",
marks=pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
),
),
pytest.param(
True,
'vulkan',
"vulkan",
marks=pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)),
])
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
),
),
],
)
@pytest.mark.skipif(importlib.util.find_spec("iree.tools") is None,
reason="Cannot find tools to import TF")
@pytest.mark.skipif(
importlib.util.find_spec("iree.tools") is None,
reason="Cannot find tools to import TF",
)
@pytest_benchmark_param
def test_bench_minilm_torch(dynamic, device):
model, test_input, act_out = get_hf_model(
"microsoft/MiniLM-L12-H384-uncased")
shark_module = SharkInference(model, (test_input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True)
"microsoft/MiniLM-L12-H384-uncased"
)
shark_module = SharkInference(
model,
(test_input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True,
)
try:
# If becnhmarking succesful, assert success/True.
shark_module.compile()
@@ -167,17 +182,21 @@ def test_bench_minilm_torch(dynamic, device):
assert False
@pytest.mark.skipif(importlib.util.find_spec("iree.tools") is None,
reason="Cannot find tools to import TF")
@pytest.mark.skipif(
importlib.util.find_spec("iree.tools") is None,
reason="Cannot find tools to import TF",
)
@pytest_benchmark_param
def test_bench_distilbert(dynamic, device):
model, test_input, act_out = get_TFhf_model("distilbert-base-uncased")
shark_module = SharkInference(model,
test_input,
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True)
shark_module = SharkInference(
model,
test_input,
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True,
)
try:
# If becnhmarking succesful, assert success/True.
shark_module.set_frontend("tensorflow")
@@ -193,12 +212,14 @@ def test_bench_distilbert(dynamic, device):
@pytest_benchmark_param
def test_bench_xlm_roberta(dynamic, device):
model, test_input, act_out = get_TFhf_model("xlm-roberta-base")
shark_module = SharkInference(model,
test_input,
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True)
shark_module = SharkInference(
model,
test_input,
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=True,
)
try:
# If becnhmarking succesful, assert success/True.
shark_module.set_frontend("tensorflow")

26
benchmarks/tests/test_hf_benchmark.py
View File

@@ -9,25 +9,31 @@ torch.manual_seed(0)
# Test running benchmark module without failing.
pytest_benchmark_param = pytest.mark.parametrize(
('dynamic', 'device'),
("dynamic", "device"),
[
pytest.param(False, 'cpu'),
pytest.param(False, "cpu"),
# TODO: Language models are failing for dynamic case..
pytest.param(True, 'cpu', marks=pytest.mark.skip),
])
pytest.param(True, "cpu", marks=pytest.mark.skip),
],
)
@pytest.mark.skipif(importlib.util.find_spec("onnxruntime") is None,
reason="Cannot find ONNXRUNTIME.")
@pytest.mark.skipif(
importlib.util.find_spec("onnxruntime") is None,
reason="Cannot find ONNXRUNTIME.",
)
@pytest_benchmark_param
def test_HFbench_minilm_torch(dynamic, device):
model_name = "bert-base-uncased"
test_input = torch.randint(2, (1, 128))
try:
shark_module = SharkHFBenchmarkRunner(model_name, (test_input,),
jit_trace=True,
dynamic=dynamic,
device=device)
shark_module = SharkHFBenchmarkRunner(
model_name,
(test_input,),
jit_trace=True,
dynamic=dynamic,
device=device,
)
shark_module.benchmark_c()
shark_module.benchmark_python((test_input,))
shark_module.benchmark_torch(test_input)

91
generate_sharktank.py
View File

@@ -18,13 +18,15 @@ import argparse
import iree.compiler.tflite as ireec_tflite
from shark.iree_utils import IREE_TARGET_MAP
class SharkTank:
def __init__(self,
torch_model_list: str = None,
tf_model_list: str = None,
tflite_model_list: str = None,
upload: bool = False):
class SharkTank:
def __init__(
self,
torch_model_list: str = None,
tf_model_list: str = None,
tflite_model_list: str = None,
upload: bool = False,
):
self.torch_model_list = torch_model_list
self.tf_model_list = tf_model_list
self.tflite_model_list = tflite_model_list
@@ -43,30 +45,53 @@ class SharkTank:
# compile and run tfhub tflite
if self.tflite_model_list is not None:
print("Setting up for tflite TMP_DIR")
self.tflite_workdir = os.path.join(os.path.dirname(__file__), "./gen_shark_tank/tflite")
self.tflite_workdir = os.path.join(
os.path.dirname(__file__), "./gen_shark_tank/tflite"
)
print(f"tflite TMP_shark_tank_DIR = {self.tflite_workdir}")
os.makedirs(self.tflite_workdir, exist_ok=True)
with open(self.tflite_model_list) as csvfile:
tflite_reader = csv.reader(csvfile, delimiter=',')
tflite_reader = csv.reader(csvfile, delimiter=",")
for row in tflite_reader:
tflite_model_name = row[0]
tflite_model_link = row[1]
print("tflite_model_name", tflite_model_name)
print("tflite_model_link", tflite_model_link)
tflite_model_name_dir = os.path.join(self.tflite_workdir, str(tflite_model_name))
tflite_model_name_dir = os.path.join(
self.tflite_workdir, str(tflite_model_name)
)
os.makedirs(tflite_model_name_dir, exist_ok=True)
tflite_saving_file = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_tflite.tflite'])
tflite_tosa_file = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_tosa.mlir'])
self.binary = '/'.join([tflite_model_name_dir, str(tflite_model_name)+'_module.bytecode'])
print("Setting up local address for tflite model file: ", tflite_saving_file)
tflite_saving_file = "/".join(
[
tflite_model_name_dir,
str(tflite_model_name) + "_tflite.tflite",
]
)
tflite_tosa_file = "/".join(
[
tflite_model_name_dir,
str(tflite_model_name) + "_tosa.mlir",
]
)
self.binary = "/".join(
[
tflite_model_name_dir,
str(tflite_model_name) + "_module.bytecode",
]
)
print(
"Setting up local address for tflite model file: ",
tflite_saving_file,
)
if os.path.exists(tflite_saving_file):
print(tflite_saving_file, "exists")
else:
print("Download tflite model")
urllib.request.urlretrieve(str(tflite_model_link),
tflite_saving_file)
urllib.request.urlretrieve(
str(tflite_model_link), tflite_saving_file
)
if os.path.exists(tflite_tosa_file):
print("Exists", tflite_tosa_file)
@@ -76,19 +101,35 @@ class SharkTank:
tflite_saving_file,
input_type="tosa",
save_temp_iree_input=tflite_tosa_file,
target_backends=[IREE_TARGET_MAP['cpu']],
import_only=False)
target_backends=[IREE_TARGET_MAP["cpu"]],
import_only=False,
)
if self.upload == True:
print("upload tmp tank to gcp")
os.system('gsutil cp -r ./gen_shark_tank gs://shark_tank/')
if __name__ == '__main__':
os.system("gsutil cp -r ./gen_shark_tank gs://shark_tank/")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--torch_model_list", type=str, default="./tank/torch/torch_model_list.csv")
parser.add_argument("--tf_model_list", type=str, default="./tank/tf/tf_model_list.csv")
parser.add_argument("--tflite_model_list", type=str, default="./tank/tflite/tflite_model_list.csv")
parser.add_argument(
"--torch_model_list",
type=str,
default="./tank/torch/torch_model_list.csv",
)
parser.add_argument(
"--tf_model_list", type=str, default="./tank/tf/tf_model_list.csv"
)
parser.add_argument(
"--tflite_model_list",
type=str,
default="./tank/tflite/tflite_model_list.csv",
)
parser.add_argument("--upload", type=bool, default=False)
args = parser.parse_args()
SharkTank(torch_model_list=args.torch_model_list, tf_model_list=args.tf_model_list,
tflite_model_list=args.tflite_model_list, upload=args.upload)
SharkTank(
torch_model_list=args.torch_model_list,
tf_model_list=args.tf_model_list,
tflite_model_list=args.tflite_model_list,
upload=args.upload,
)

2
setup.py
View File

@@ -26,7 +26,7 @@ setup(
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
],
packages=find_packages(exclude=('examples')),
packages=find_packages(exclude=("examples")),
python_requires=">=3.7",
install_requires=[
"numpy",

29
shark/backward_makefx.py
View File

@@ -22,7 +22,6 @@ import tempfile
class MakeFxModule:
def __init__(self, model, inputs, labels=None, custom_inference_fn=None):
self.model = model
self.inputs = inputs
@@ -52,20 +51,28 @@ class MakeFxModule:
return fx_g
def generate_graph(self):
fx_g = make_fx(self.custom_inference_fn,
decomposition_table=get_decompositions([
torch.ops.aten.embedding_dense_backward,
torch.ops.aten.native_layer_norm_backward,
torch.ops.aten.slice_backward,
torch.ops.aten.select_backward
]))(dict(self.model.named_parameters()),
dict(self.model.named_buffers()), self.inputs)
fx_g = make_fx(
self.custom_inference_fn,
decomposition_table=get_decompositions(
[
torch.ops.aten.embedding_dense_backward,
torch.ops.aten.native_layer_norm_backward,
torch.ops.aten.slice_backward,
torch.ops.aten.select_backward,
]
),
)(
dict(self.model.named_parameters()),
dict(self.model.named_buffers()),
self.inputs,
)
fx_g.graph.set_codegen(torch.fx.graph.CodeGen())
fx_g.recompile()
fx_g = self.change_fx_graph_return_to_tuple(fx_g)
ts_g = torch.jit.script(fx_g)
temp = tempfile.NamedTemporaryFile(suffix='_shark_ts',
prefix='temp_ts_')
temp = tempfile.NamedTemporaryFile(
suffix="_shark_ts", prefix="temp_ts_"
)
ts_g.save(temp.name)
new_ts = torch.jit.load(temp.name)
self.training_graph = new_ts

16
shark/examples/shark_eager/dynamo_demo.py
View File

@@ -8,7 +8,9 @@ try:
from torchdynamo.optimizations.backends import create_backend
from torchdynamo.optimizations.subgraph import SubGraph
except ModuleNotFoundError:
print("Please install TorchDynamo using pip install git+https://github.com/pytorch/torchdynamo")
print(
"Please install TorchDynamo using pip install git+https://github.com/pytorch/torchdynamo"
)
exit()
NUM_ITERS = 10
@@ -24,7 +26,9 @@ def __torch_mlir(fx_graph, *args, **kwargs):
for node in fx_g.graph.nodes:
if node.op == "output":
assert len(node.args) == 1, "Output node must have a single argument"
assert (
len(node.args) == 1
), "Output node must have a single argument"
node_arg = node.args[0]
if isinstance(node_arg, tuple) and len(node_arg) == 1:
node.args = (node_arg[0],)
@@ -41,8 +45,12 @@ def __torch_mlir(fx_graph, *args, **kwargs):
if len(args) == 1 and isinstance(args[0], list):
args = args[0]
linalg_module = compile(ts_graph, args, output_type=OutputType.LINALG_ON_TENSORS)
callable, _ = get_iree_compiled_module(linalg_module, "cuda", func_name="forward")
linalg_module = compile(
ts_graph, args, output_type=OutputType.LINALG_ON_TENSORS
)
callable, _ = get_iree_compiled_module(
linalg_module, "cuda", func_name="forward"
)
def forward(*inputs):
return callable(*inputs)

34
shark/examples/shark_inference/CLIPModel_tf.py
View File

@@ -9,23 +9,24 @@ from shark.shark_inference import SharkInference
clip_vit_inputs = [
tf.TensorSpec(shape=[2, 7], dtype=tf.int32),
tf.TensorSpec(shape=[2, 7], dtype=tf.int32),
tf.TensorSpec(shape=[1, 3, 224, 224], dtype=tf.float32)
tf.TensorSpec(shape=[1, 3, 224, 224], dtype=tf.float32),
]
class CLIPModule(tf.Module):
def __init__(self):
super(CLIPModule, self).__init__()
self.m = TFCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
self.m.predict = lambda x, y, z: self.m(
input_ids=x, attention_mask=y, pixel_values=z)
input_ids=x, attention_mask=y, pixel_values=z
)
@tf.function(input_signature=clip_vit_inputs)
def forward(self, input_ids, attention_mask, pixel_values):
return self.m.predict(input_ids, attention_mask,
pixel_values).logits_per_image
return self.m.predict(
input_ids, attention_mask, pixel_values
).logits_per_image
if __name__ == "__main__":
@@ -35,17 +36,26 @@ if __name__ == "__main__":
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=["a photo of a cat", "a photo of a dog"],
images=image,
return_tensors="tf",
padding=True)
inputs = processor(
text=["a photo of a cat", "a photo of a dog"],
images=image,
return_tensors="tf",
padding=True,
)
shark_module = SharkInference(
CLIPModule(),
(inputs["input_ids"], inputs["attention_mask"], inputs["pixel_values"]))
(inputs["input_ids"], inputs["attention_mask"], inputs["pixel_values"]),
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
print(
shark_module.forward((inputs["input_ids"], inputs["attention_mask"],
inputs["pixel_values"])))
shark_module.forward(
(
inputs["input_ids"],
inputs["attention_mask"],
inputs["pixel_values"],
)
)
)

6
shark/examples/shark_inference/gpt2_tf.py
View File

@@ -13,7 +13,6 @@ gpt2_inputs = [
class GPT2Module(tf.Module):
def __init__(self):
super(GPT2Module, self).__init__()
self.m = TFGPT2Model.from_pretrained("distilgpt2")
@@ -30,9 +29,10 @@ if __name__ == "__main__":
tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
text = "I love the distilled version of models."
inputs = tokenizer(text, return_tensors='tf')
inputs = tokenizer(text, return_tensors="tf")
shark_module = SharkInference(
GPT2Module(), (inputs["input_ids"], inputs["attention_mask"]))
GPT2Module(), (inputs["input_ids"], inputs["attention_mask"])
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
print(shark_module.forward((inputs["input_ids"], inputs["attention_mask"])))

19
shark/examples/shark_inference/minilm_benchmark.py
View File

@@ -7,17 +7,13 @@ tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
class MiniLMSequenceClassification(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -27,9 +23,12 @@ class MiniLMSequenceClassification(torch.nn.Module):
test_input = torch.randint(2, (1, 128))
shark_module = SharkInference(MiniLMSequenceClassification(), (test_input,),
jit_trace=True,
benchmark_mode=True)
shark_module = SharkInference(
MiniLMSequenceClassification(),
(test_input,),
jit_trace=True,
benchmark_mode=True,
)
shark_module.compile()
shark_module.forward((test_input,))

37
shark/examples/shark_inference/minilm_benchmark_tf.py
View File

@@ -9,21 +9,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
# Create a BERT trainer with the created network.
self.m = TFBertModel.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", from_pt=True)
"microsoft/MiniLM-L12-H384-uncased", from_pt=True
)
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=bert_input)
def forward(self, input_ids, attention_mask, token_type_ids):
@@ -33,22 +34,26 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# Prepping Data
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased")
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
encoded_input = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
tf.convert_to_tensor(encoded_input[key]), 0
)
test_input = (encoded_input["input_ids"], encoded_input["attention_mask"],
encoded_input["token_type_ids"])
shark_module = SharkInference(
BertModule(),
test_input,
benchmark_mode=True)
test_input = (
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
)
shark_module = SharkInference(BertModule(), test_input, benchmark_mode=True)
shark_module.set_frontend("tensorflow")
shark_module.compile()
shark_module.benchmark_all(test_input)

18
shark/examples/shark_inference/minilm_jit.py
View File

@@ -7,17 +7,13 @@ tokenizer = AutoTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
class MiniLMSequenceClassification(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -27,9 +23,11 @@ class MiniLMSequenceClassification(torch.nn.Module):
test_input = torch.randint(2, (1, 128)).to(torch.int32)
shark_module = SharkInference(MiniLMSequenceClassification(),
(test_input, test_input, test_input),
jit_trace=True)
shark_module = SharkInference(
MiniLMSequenceClassification(),
(test_input, test_input, test_input),
jit_trace=True,
)
shark_module.compile()
result = shark_module.forward((test_input, test_input, test_input))

42
shark/examples/shark_inference/minilm_tf.py
View File

@@ -9,21 +9,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
# Create a BERT trainer with the created network.
self.m = TFBertModel.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", from_pt=True)
"microsoft/MiniLM-L12-H384-uncased", from_pt=True
)
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=bert_input)
def forward(self, input_ids, attention_mask, token_type_ids):
@@ -33,24 +34,37 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# Prepping Data
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased")
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
encoded_input = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
tf.convert_to_tensor(encoded_input[key]), 0
)
shark_module = SharkInference(
BertModule(),
(encoded_input["input_ids"], encoded_input["attention_mask"],
encoded_input["token_type_ids"]))
(
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
),
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
print(
shark_module.forward(
(encoded_input["input_ids"], encoded_input["attention_mask"],
encoded_input["token_type_ids"])))
(
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
)
)
)

26
shark/examples/shark_inference/resnet50_script.py
View File

@@ -10,19 +10,22 @@ from shark.shark_inference import SharkInference
################################## Preprocessing inputs and model ############
def load_and_preprocess_image(url: str):
headers = {
"User-Agent":
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36"
"User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/50.0.2661.102 Safari/537.36"
}
img = Image.open(requests.get(url, headers=headers,
stream=True).raw).convert("RGB")
img = Image.open(
requests.get(url, headers=headers, stream=True).raw
).convert("RGB")
# preprocessing pipeline
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
preprocess = transforms.Compose(
[
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
),
]
)
img_preprocessed = preprocess(img)
return torch.unsqueeze(img_preprocessed, 0)
@@ -44,7 +47,6 @@ def top3_possibilities(res):
class Resnet50Module(torch.nn.Module):
def __init__(self):
super().__init__()
self.resnet = models.resnet50(pretrained=True)

13
shark/examples/shark_inference/t5_tf.py
View File

@@ -11,12 +11,12 @@ t5_inputs = [
tf.TensorSpec(shape=[1, 10], dtype=tf.int32),
]
class T5Module(tf.Module):
class T5Module(tf.Module):
def __init__(self):
super(T5Module, self).__init__()
self.m = TFT5Model.from_pretrained("t5-small")
self.m.predict = lambda x,y: self.m(input_ids=x, decoder_input_ids=y)
self.m.predict = lambda x, y: self.m(input_ids=x, decoder_input_ids=y)
@tf.function(input_signature=t5_inputs)
def forward(self, input_ids, decoder_input_ids):
@@ -27,12 +27,9 @@ if __name__ == "__main__":
# Prepping Data
tokenizer = T5Tokenizer.from_pretrained("t5-small")
text = "I love the distilled version of models."
inputs = tokenizer(
text, return_tensors="tf"
).input_ids
inputs = tokenizer(text, return_tensors="tf").input_ids
shark_module = SharkInference(
T5Module(), (inputs, inputs))
shark_module = SharkInference(T5Module(), (inputs, inputs))
shark_module.set_frontend("tensorflow")
shark_module.compile()
print(shark_module.forward((inputs,inputs)))
print(shark_module.forward((inputs, inputs)))

1
shark/examples/shark_inference/torch_vision_models_script.py
View File

@@ -4,7 +4,6 @@ from shark.shark_inference import SharkInference
class VisionModule(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model

1
shark/examples/shark_inference/unet_script.py
View File

@@ -4,7 +4,6 @@ from shark_runner import SharkInference
# Currently not supported aten.transpose_conv2d missing.
class UnetModule(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = torch.hub.load(

15
shark/examples/shark_training/bert_training.py
View File

@@ -5,17 +5,13 @@ from shark.shark_runner import SharkTrainer
class MiniLMSequenceClassification(torch.nn.Module):
def __init__(self):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -37,8 +33,9 @@ inp = (torch.randint(2, (1, 128)),)
def forward(params, buffers, args):
params_and_buffers = {**params, **buffers}
_stateless.functional_call(mod, params_and_buffers, args,
{}).sum().backward()
_stateless.functional_call(
mod, params_and_buffers, args, {}
).sum().backward()
optim = torch.optim.SGD(get_sorted_params(params), lr=0.01)
# optim.load_state_dict(optim_state)
optim.step()

23
shark/examples/shark_training/bert_training_load_tf.py
View File

@@ -11,7 +11,8 @@ parser.add_argument(
"--download_mlir_path",
type=str,
default="bert_tf_training.mlir",
help="Specifies path to target mlir file that will be loaded.")
help="Specifies path to target mlir file that will be loaded.",
)
load_args, unknown = parser.parse_known_args()
tf.random.set_seed(0)
@@ -25,20 +26,30 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
file_link = "https://storage.googleapis.com/shark_tank/users/stanley/bert_tf_training.mlir"
response = request.urlretrieve(file_link, load_args.download_mlir_path)
sample_input_tensors = [tf.convert_to_tensor(val, dtype=tf.int32) for val in predict_sample_input]
sample_input_tensors = [
tf.convert_to_tensor(val, dtype=tf.int32)
for val in predict_sample_input
]
num_iter = 10
if not os.path.isfile(load_args.download_mlir_path):
raise ValueError(f"Tried looking for target mlir in {load_args.download_mlir_path}, but cannot be found.")
raise ValueError(
f"Tried looking for target mlir in {load_args.download_mlir_path}, but cannot be found."
)
with open(load_args.download_mlir_path, "rb") as input_file:
bert_mlir = input_file.read()
shark_module = SharkTrainer(
bert_mlir,
(sample_input_tensors,
tf.convert_to_tensor(np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32)))
(
sample_input_tensors,
tf.convert_to_tensor(
np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32
),
),
)
shark_module.set_frontend("mhlo")
shark_module.compile()
start = time.time()

36
shark/examples/shark_training/bert_training_tf.py
View File

@@ -25,17 +25,17 @@ bert_input = [
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=2,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -46,10 +46,12 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels
])
@tf.function(
input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
]
)
def forward(self, inputs, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -67,14 +69,22 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
sample_input_tensors = [
tf.convert_to_tensor(val, dtype=tf.int32)
for val in predict_sample_input
]
sample_input_tensors = [tf.convert_to_tensor(val, dtype=tf.int32) for val in predict_sample_input]
num_iter = 10
shark_module = SharkTrainer(
BertModule(),
(sample_input_tensors,
tf.convert_to_tensor(np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32)))
(
sample_input_tensors,
tf.convert_to_tensor(
np.random.randint(5, size=(BATCH_SIZE)), dtype=tf.int32
),
),
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
start = time.time()

6
shark/examples/shark_training/neural_net_training.py
View File

@@ -4,7 +4,6 @@ from shark.shark_trainer import SharkTrainer
class Foo(torch.nn.Module):
def __init__(self):
super(Foo, self).__init__()
self.l1 = torch.nn.Linear(10, 16)
@@ -28,8 +27,9 @@ def get_sorted_params(named_params):
def forward(params, buffers, args):
params_and_buffers = {**params, **buffers}
_stateless.functional_call(mod, params_and_buffers, args,
{}).sum().backward()
_stateless.functional_call(
mod, params_and_buffers, args, {}
).sum().backward()
optim = torch.optim.SGD(get_sorted_params(params), lr=0.01)
optim.step()
return params, buffers

29
shark/iree_eager_backend.py
View File

@@ -28,9 +28,14 @@ from torch_mlir.eager_mode.torch_mlir_eager_backend import (
TorchMLIREagerBackend,
TensorMetaData,
)
from torch_mlir_e2e_test.eager_backends.refbackend import NUMPY_TO_TORCH_DTYPE_DICT
from torch_mlir_e2e_test.eager_backends.refbackend import (
NUMPY_TO_TORCH_DTYPE_DICT,
)
from shark.iree_utils import get_iree_compiled_module, IREE_DEVICE_MAP
from shark.iree_utils.compile_utils import (
get_iree_compiled_module,
IREE_DEVICE_MAP,
)
class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend):
@@ -46,15 +51,16 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend):
self.iree_device_str = IREE_DEVICE_MAP[device]
self.config = ireert.Config(self.iree_device_str)
def get_torch_metadata(self, tensor: DeviceArray,
kwargs: Dict[str, Any]) -> TensorMetaData:
def get_torch_metadata(
self, tensor: DeviceArray, kwargs: Dict[str, Any]
) -> TensorMetaData:
return TensorMetaData(
size=tensor.shape,
dtype=NUMPY_TO_TORCH_DTYPE_DICT[tensor.dtype.type],
device=torch.device(self.torch_device_str),
requires_grad=tensor.dtype.type
in {np.float, np.float32, np.float64} and
kwargs.get("requires_grad", False),
in {np.float, np.float32, np.float64}
and kwargs.get("requires_grad", False),
)
def compile(self, imported_module: Module):
@@ -64,9 +70,9 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend):
"torch-function-to-torch-backend-pipeline,torch-backend-to-linalg-on-tensors-backend-pipeline",
"EagerMode",
)
callable, _ = get_iree_compiled_module(imported_module,
self.iree_device_str,
func_name=fn_name)
callable, _ = get_iree_compiled_module(
imported_module, self.iree_device_str, func_name=fn_name
)
return callable
def copy_into(self, dst, src):
@@ -76,6 +82,7 @@ class EagerModeIREELinalgOnTensorsBackend(TorchMLIREagerBackend):
def transfer_from_device_to_torch(self, e):
return torch.from_numpy(e.to_host())
def transfer_from_torch_to_device(self,
tensor: torch.Tensor) -> DeviceArray:
def transfer_from_torch_to_device(
self, tensor: torch.Tensor
) -> DeviceArray:
return iree.runtime.asdevicearray(self.config.device, tensor.numpy())

381
shark/iree_utils.py
View File

@@ -1,381 +0,0 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import iree.runtime as ireert
import iree.runtime.scripts.iree_benchmark_module as benchmark_module
import iree.compiler as ireec
from shark.torch_mlir_utils import get_module_name_for_asm_dump
from shark.cuda_utils import get_cuda_sm_cc
from shark.model_annotation import *
import subprocess
import numpy as np
import os
import re
import sys
IREE_DEVICE_MAP = {
"cpu": "local-task",
"gpu": "cuda",
"cuda": "cuda",
"vulkan": "vulkan",
"metal": "vulkan",
"rocm": "rocm"
}
IREE_TARGET_MAP = {
"cpu": "dylib",
"gpu": "cuda",
"cuda": "cuda",
"vulkan": "vulkan",
"metal": "vulkan",
"rocm": "rocm"
}
UNIT_TO_SECOND_MAP = {"ms": 0.001, "s": 1}
def check_device_drivers(device):
"""Checks necessary drivers present for gpu and vulkan devices"""
if (device in ["gpu", "cuda"]):
try:
subprocess.check_output('nvidia-smi')
except Exception:
return True
elif (device in ["metal", "vulkan"]):
try:
subprocess.check_output('vulkaninfo')
except Exception:
return True
elif (device == "cpu"):
return False
# Unknown device.
else:
return True
return False
def get_iree_cpu_args():
find_triple_cmd = "uname -s -m"
os_name, proc_name = subprocess.run(
find_triple_cmd, shell=True, stdout=subprocess.PIPE,
check=True).stdout.decode('utf-8').split()
if os_name == "Darwin":
find_kernel_version_cmd = "uname -r"
kernel_version = subprocess.run(find_kernel_version_cmd,
shell=True,
stdout=subprocess.PIPE,
check=True).stdout.decode('utf-8')
target_triple = f"{proc_name}-apple-darwin{kernel_version}"
elif os_name == "Linux":
target_triple = f"{proc_name}-linux-gnu"
else:
error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)"
raise Exception(error_message)
print(f"Target triple found:{target_triple}")
return [f"-iree-llvm-target-triple={target_triple}"]
def get_iree_gpu_args():
ireert.flags.FUNCTION_INPUT_VALIDATION = False
ireert.flags.parse_flags("--cuda_allow_inline_execution")
sm_arch = get_cuda_sm_cc()
if sm_arch in ['sm_70', 'sm_72', 'sm_75', 'sm_80', 'sm_84', 'sm_86']:
return [
"--iree-hal-cuda-disable-loop-nounroll-wa",
f"--iree-hal-cuda-llvm-target-arch={sm_arch}"
]
else:
return ["--iree-hal-cuda-disable-loop-nounroll-wa"]
def get_vulkan_triple_flag():
vulkan_device_cmd = "vulkaninfo | grep deviceName | awk \'END{{print $NF}}\'"
vulkan_device = run_cmd(vulkan_device_cmd).strip()
if vulkan_device == "M1":
print("Found Apple Device. Using m1-moltenvk-macos")
return "-iree-vulkan-target-triple=m1-moltenvk-macos"
elif vulkan_device == "A100-SXM4-40GB":
print("Found Nvidia Device. Using ampere-rtx3080-linux")
return "-iree-vulkan-target-triple=ampere-rtx3080-linux"
else:
print(
"Optimized kernel for your target device is not added yet. Contact SHARK Admin on discord[https://discord.com/invite/RUqY2h2s9u] or pull up an issue."
)
return None
def get_iree_vulkan_args():
#vulkan_flag = ["--iree-flow-demote-i64-to-i32"]
vulkan_flag = []
vulkan_triple_flag = get_vulkan_triple_flag()
if vulkan_triple_flag is not None:
vulkan_flag.append(vulkan_triple_flag)
return vulkan_flag
def get_iree_device_args(device):
if device == "cpu":
return get_iree_cpu_args()
if device in ["gpu", "cuda"]:
return get_iree_gpu_args()
if device in ["metal", "vulkan"]:
return get_iree_vulkan_args()
return []
def get_iree_frontend_args(frontend):
if frontend in ["torch", "pytorch", "linalg"]:
return ["--iree-llvm-target-cpu-features=host"]
elif frontend in ["tensorflow", "tf", "mhlo"]:
return [
"--iree-llvm-target-cpu-features=host",
"--iree-mhlo-demote-i64-to-i32=false",
"--iree-flow-demote-i64-to-i32"
]
else:
# Frontend not found.
return []
def get_iree_common_args():
return [
"--iree-stream-resource-index-bits=64",
"--iree-vm-target-index-bits=64"
]
def compile_module_to_flatbuffer(module, device, frontend, func_name,
model_config_path):
# Setup Compile arguments wrt to frontends.
input_type = ""
args = get_iree_frontend_args(frontend)
args += get_iree_device_args(device)
args += get_iree_common_args()
if frontend in ["tensorflow", "tf"]:
input_type = "mhlo"
elif frontend in ["mhlo", "tosa"]:
input_type = frontend
elif frontend in ["tflite", "tflite-tosa"]:
input_type = "tosa"
# Annotate the input module with the configs
if model_config_path != None:
# Currently tuned model only works on tf frontend
if frontend in ["tensorflow", "tf"]:
input_module = module.decode('utf-8')
elif frontend in ["pytorch", "torch"]:
input_module = module.operation.get_asm()
with create_context() as ctx:
module = model_annotation(ctx,
input_contents=input_module,
config_path=model_config_path)
module = str(module)
# Compile according to the input type, else just try compiling.
if input_type not in ["mhlo", "tosa"]:
module = str(module)
if input_type != "":
# Currently for MHLO/TOSA.
flatbuffer_blob = ireec.compile_str(
module,
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args,
input_type=input_type)
else:
# Currently for Torch.
flatbuffer_blob = ireec.compile_str(
str(module),
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args)
return flatbuffer_blob
def get_iree_module(flatbuffer_blob, device, func_name):
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
config = ireert.Config(IREE_DEVICE_MAP[device])
ctx = ireert.SystemContext(config=config)
ctx.add_vm_module(vm_module)
ModuleCompiled = ctx.modules.module[func_name]
return ModuleCompiled, config
def get_iree_compiled_module(module,
device: str,
frontend: str = "torch",
func_name: str = "forward",
model_config_path: str = None):
"""Given a module returns the compiled .vmfb and configs"""
flatbuffer_blob = compile_module_to_flatbuffer(module, device, frontend,
func_name, model_config_path)
return get_iree_module(flatbuffer_blob, device, func_name)
def export_iree_module_to_vmfb(module,
device: str,
directory: str,
frontend: str = "torch",
func_name: str = "forward",
model_config_path: str = None):
flatbuffer_blob = compile_module_to_flatbuffer(module, device, frontend,
func_name, model_config_path)
module_name = f"{frontend}_{func_name}_{device}"
filename = os.path.join(directory, module_name + ".vmfb")
print(f"Saved vmfb in {filename}.")
with open(filename, 'wb') as f:
f.write(flatbuffer_blob)
return filename
def export_module_to_mlir_file(module, frontend, directory: str):
mlir_str = module
if frontend in ["tensorflow", "tf", "mhlo", "tflite"]:
mlir_str = module.decode('utf-8')
elif frontend in ["pytorch", "torch"]:
mlir_str = module.operation.get_asm()
filename = os.path.join(directory, "model.mlir")
with open(filename, 'w') as f:
f.write(mlir_str)
print(f"Saved mlir in {filename}.")
return filename
def get_results(compiled_vm, input, config, frontend="torch"):
"""Runs a .vmfb file given inputs and config and returns output."""
device_inputs = input
if frontend in ["torch", "pytorch"]:
device_inputs = [ireert.asdevicearray(config.device, a) for a in input]
if frontend in ["tensorflow", "tf", "tflite", "tflite-tosa"]:
device_inputs = []
for a in input:
if (isinstance(a, list)):
device_inputs.append([
ireert.asdevicearray(config.device, val, dtype=val.dtype)
for val in a
])
else:
device_inputs.append(ireert.asdevicearray(config.device, a))
result = compiled_vm(*device_inputs)
result_tensors = []
if (isinstance(result, tuple)):
for val in result:
result_tensors.append(np.copy(np.asarray(val, val.dtype)))
return result_tensors
elif (isinstance(result, dict)):
data = list(result.items())
res = np.array(data, dtype=object)
return np.copy(res)
else:
return np.copy(np.asarray(result, dtype=result.dtype))
######### Benchmark Related Tools ###########
def tensor_to_type_str(input_tensors: tuple, frontend: str):
"""
Input: A tuple of input tensors i.e tuple(torch.tensor)
Output: list of string that represent mlir types (i.e 1x24xf64)
# TODO: Support more than floats, and ints
"""
list_of_type = []
for input_tensor in input_tensors:
if (isinstance(input_tensor, tuple)):
for val in input_tensor:
type_string = "x".join([str(dim) for dim in val.shape])
if frontend in ["torch", "pytorch"]:
dtype_string = str(val.dtype).replace("torch.", "")
elif frontend in ["tensorflow", "tf"]:
dtype = val.dtype
dtype_string = re.findall('\'[^"]*\'',
str(dtype))[0].replace("\'", "")
regex_split = re.compile("([a-zA-Z]+)([0-9]+)")
match = regex_split.match(dtype_string)
mlir_type_string = str(match.group(1)[0]) + str(match.group(2))
type_string += f"x{mlir_type_string}"
list_of_type.append(type_string)
else:
type_string = "x".join([str(dim) for dim in input_tensor.shape])
if frontend in ["torch", "pytorch"]:
dtype_string = str(input_tensor.dtype).replace("torch.", "")
elif frontend in ["tensorflow", "tf"]:
dtype = input_tensor.dtype
dtype_string = re.findall('\'[^"]*\'',
str(dtype))[0].replace("\'", "")
regex_split = re.compile("([a-zA-Z]+)([0-9]+)")
match = regex_split.match(dtype_string)
mlir_type_string = str(match.group(1)[0]) + str(match.group(2))
type_string += f"x{mlir_type_string}"
list_of_type.append(type_string)
return list_of_type
def build_benchmark_args(input_file: str,
device: str,
input_tensors: tuple,
frontend: str,
training=False):
"""
Inputs: input_file leading to vmfb, input_tensor to function, target device, and whether it is training or not.
Outputs: string that execute benchmark-module on target model.
"""
path = benchmark_module.__path__[0]
benchmarker_path = os.path.join(path, "..", "..", "iree-benchmark-module")
benchmark_cl = [benchmarker_path, f"--module_file={input_file}"]
fn_name = "forward"
if training == True:
# TODO: Replace name of train with actual train fn name.
fn_name = "train"
benchmark_cl.append(f"--entry_function={fn_name}")
benchmark_cl.append(f"--device={IREE_DEVICE_MAP[device]}")
mlir_input_types = tensor_to_type_str(input_tensors, frontend)
for mlir_input in mlir_input_types:
benchmark_cl.append(f"--function_input={mlir_input}")
time_extractor = "| awk \'END{{print $2 $3}}\'"
benchmark_cl.append(time_extractor)
return benchmark_cl
def run_cmd(cmd):
"""
Inputs: cli command string.
"""
try:
result = subprocess.run(cmd,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
check=True)
result_str = result.stdout.decode()
return result_str
except Exception:
sys.exit("Exiting program due to error running:", cmd)
def run_benchmark_module(benchmark_cl):
"""
Run benchmark command, extract result and return iteration/seconds.
Input: benchmark command.
"""
benchmark_path = benchmark_cl[0]
assert os.path.exists(
benchmark_path
), "Cannot find benchmark_module, Please contact SHARK maintainer on discord."
bench_result = run_cmd(' '.join(benchmark_cl))
regex_split = re.compile("([0-9]+[.]*[0-9]*)([a-zA-Z]+)")
match = regex_split.match(bench_result)
time = float(match.group(1))
unit = match.group(2)
return 1.0 / (time * UNIT_TO_SECOND_MAP[unit])

0
shark/iree_utils/__init__.py Normal file
View File

77
shark/iree_utils/_common.py Normal file
View File

@@ -0,0 +1,77 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
## Common utilities to be shared by iree utilities.
import os
import sys
import subprocess
def run_cmd(cmd):
"""
Inputs: cli command string.
"""
try:
result = subprocess.run(
cmd,
shell=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
check=True,
)
result_str = result.stdout.decode()
return result_str
except Exception:
sys.exit("Exiting program due to error running:", cmd)
IREE_DEVICE_MAP = {
"cpu": "local-task",
"gpu": "cuda",
"cuda": "cuda",
"vulkan": "vulkan",
"metal": "vulkan",
"rocm": "rocm",
}
IREE_TARGET_MAP = {
"cpu": "dylib",
"gpu": "cuda",
"cuda": "cuda",
"vulkan": "vulkan",
"metal": "vulkan",
"rocm": "rocm",
}
# Finds whether the required drivers are installed for the given device.
def check_device_drivers(device):
"""Checks necessary drivers present for gpu and vulkan devices"""
if device in ["gpu", "cuda"]:
try:
subprocess.check_output("nvidia-smi")
except Exception:
return True
elif device in ["metal", "vulkan"]:
try:
subprocess.check_output("vulkaninfo")
except Exception:
return True
elif device == "cpu":
return False
# Unknown device.
else:
return True
return False

94
shark/iree_utils/benchmark_utils.py Normal file
View File

@@ -0,0 +1,94 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import iree.runtime.scripts.iree_benchmark_module as benchmark_module
from shark.iree_utils._common import run_cmd, IREE_DEVICE_MAP
import numpy as np
import os
import re
UNIT_TO_SECOND_MAP = {"ms": 0.001, "s": 1}
def tensor_to_type_str(input_tensors: tuple, frontend: str):
"""
Input: A tuple of input tensors i.e tuple(torch.tensor)
Output: list of string that represent mlir types (i.e 1x24xf64)
# TODO: Support more than floats, and ints
"""
list_of_type = []
for input_tensor in input_tensors:
type_string = "x".join([str(dim) for dim in input_tensor.shape])
if frontend in ["torch", "pytorch"]:
dtype_string = str(input_tensor.dtype).replace("torch.", "")
elif frontend in ["tensorflow", "tf"]:
dtype = input_tensor.dtype
dtype_string = re.findall("'[^\"]*'", str(dtype))[0].replace(
"'", ""
)
regex_split = re.compile("([a-zA-Z]+)([0-9]+)")
match = regex_split.match(dtype_string)
mlir_type_string = str(match.group(1)[0]) + str(match.group(2))
type_string += f"x{mlir_type_string}"
list_of_type.append(type_string)
return list_of_type
def build_benchmark_args(
input_file: str,
device: str,
input_tensors: tuple,
frontend: str,
training=False,
):
"""
Inputs: input_file leading to vmfb, input_tensor to function, target device,
and whether it is training or not.
Outputs: string that execute benchmark-module on target model.
"""
path = benchmark_module.__path__[0]
benchmarker_path = os.path.join(path, "..", "..", "iree-benchmark-module")
benchmark_cl = [benchmarker_path, f"--module_file={input_file}"]
# TODO: The function named can be passed as one of the args.
fn_name = "forward"
if training == True:
# TODO: Replace name of train with actual train fn name.
fn_name = "train"
benchmark_cl.append(f"--entry_function={fn_name}")
benchmark_cl.append(f"--device={IREE_DEVICE_MAP[device]}")
mlir_input_types = tensor_to_type_str(input_tensors, frontend)
for mlir_input in mlir_input_types:
benchmark_cl.append(f"--function_input={mlir_input}")
time_extractor = "| awk 'END{{print $2 $3}}'"
benchmark_cl.append(time_extractor)
return benchmark_cl
def run_benchmark_module(benchmark_cl):
"""
Run benchmark command, extract result and return iteration/seconds.
# TODO: Add an example of the benchmark command.
Input: benchmark command.
"""
benchmark_path = benchmark_cl[0]
assert os.path.exists(
benchmark_path
), "Cannot find benchmark_module, Please contact SHARK maintainer on discord."
bench_result = run_cmd(" ".join(benchmark_cl))
regex_split = re.compile("([0-9]+[.]*[0-9]*)([a-zA-Z]+)")
match = regex_split.match(bench_result)
time = float(match.group(1))
unit = match.group(2)
return 1.0 / (time * UNIT_TO_SECOND_MAP[unit])

188
shark/iree_utils/compile_utils.py Normal file
View File

@@ -0,0 +1,188 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import iree.runtime as ireert
import iree.compiler as ireec
from shark.iree_utils._common import IREE_DEVICE_MAP, IREE_TARGET_MAP
from shark.model_annotation import *
import numpy as np
import os
# Get the iree-compile arguments given device.
def get_iree_device_args(device):
if device == "cpu":
from shark.iree_utils.cpu_utils import get_iree_cpu_args
return get_iree_cpu_args()
if device in ["gpu", "cuda"]:
from shark.iree_utils.gpu_utils import get_iree_gpu_args
return get_iree_gpu_args()
if device in ["metal", "vulkan"]:
from shark.iree_utils.vulkan_utils import get_iree_vulkan_args
return get_iree_vulkan_args()
return []
# Get the iree-compiler arguments given frontend.
def get_iree_frontend_args(frontend):
if frontend in ["torch", "pytorch", "linalg"]:
return ["--iree-llvm-target-cpu-features=host"]
elif frontend in ["tensorflow", "tf", "mhlo"]:
return [
"--iree-llvm-target-cpu-features=host",
"--iree-mhlo-demote-i64-to-i32=false",
"--iree-flow-demote-i64-to-i32",
]
else:
# Frontend not found.
return []
# Common args to be used given any frontend or device.
def get_iree_common_args():
return [
"--iree-stream-resource-index-bits=64",
"--iree-vm-target-index-bits=64",
]
def compile_module_to_flatbuffer(
module, device, frontend, func_name, model_config_path
):
# Setup Compile arguments wrt to frontends.
input_type = ""
args = get_iree_frontend_args(frontend)
args += get_iree_device_args(device)
args += get_iree_common_args()
if frontend in ["tensorflow", "tf"]:
input_type = "mhlo"
elif frontend in ["mhlo", "tosa"]:
input_type = frontend
elif frontend in ["tflite", "tflite-tosa"]:
input_type = "tosa"
# TODO: make it simpler.
# Compile according to the input type, else just try compiling.
if input_type not in ["mhlo", "tosa"]:
module = str(module)
if input_type != "":
# Currently for MHLO/TOSA.
flatbuffer_blob = ireec.compile_str(
module,
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args,
input_type=input_type,
)
else:
# Currently for Torch.
flatbuffer_blob = ireec.compile_str(
str(module),
target_backends=[IREE_TARGET_MAP[device]],
extra_args=args,
)
return flatbuffer_blob
def get_iree_module(flatbuffer_blob, device, func_name):
# Returns the compiled module and the configs.
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
config = ireert.Config(IREE_DEVICE_MAP[device])
ctx = ireert.SystemContext(config=config)
ctx.add_vm_module(vm_module)
ModuleCompiled = ctx.modules.module[func_name]
return ModuleCompiled, config
def get_iree_compiled_module(
module,
device: str,
frontend: str = "torch",
func_name: str = "forward",
model_config_path: str = None,
):
"""Given a module returns the compiled .vmfb and configs"""
flatbuffer_blob = compile_module_to_flatbuffer(
module, device, frontend, func_name, model_config_path
)
return get_iree_module(flatbuffer_blob, device, func_name)
def export_iree_module_to_vmfb(
module,
device: str,
directory: str,
frontend: str = "torch",
func_name: str = "forward",
model_config_path: str = None,
):
# Compiles the module given specs and saves it as .vmfb file.
flatbuffer_blob = compile_module_to_flatbuffer(
module, device, frontend, func_name, model_config_path
)
module_name = f"{frontend}_{func_name}_{device}"
filename = os.path.join(directory, module_name + ".vmfb")
print(f"Saved vmfb in {filename}.")
with open(filename, "wb") as f:
f.write(flatbuffer_blob)
return filename
def export_module_to_mlir_file(module, frontend, directory: str):
# TODO: write proper documentation.
mlir_str = module
if frontend in ["tensorflow", "tf", "mhlo", "tflite"]:
mlir_str = module.decode("utf-8")
elif frontend in ["pytorch", "torch"]:
mlir_str = module.operation.get_asm()
filename = os.path.join(directory, "model.mlir")
with open(filename, "w") as f:
f.write(mlir_str)
print(f"Saved mlir in {filename}.")
return filename
def get_results(compiled_vm, input, config, frontend="torch"):
"""Runs a .vmfb file given inputs and config and returns output."""
device_inputs = input
if frontend in ["torch", "pytorch"]:
device_inputs = [ireert.asdevicearray(config.device, a) for a in input]
if frontend in ["tensorflow", "tf", "tflite", "tflite-tosa"]:
device_inputs = []
for a in input:
if isinstance(a, list):
device_inputs.append(
[
ireert.asdevicearray(
config.device, val, dtype=val.dtype
)
for val in a
]
)
else:
device_inputs.append(ireert.asdevicearray(config.device, a))
result = compiled_vm(*device_inputs)
result_tensors = []
if isinstance(result, tuple):
for val in result:
result_tensors.append(np.copy(np.asarray(val, val.dtype)))
return result_tensors
elif isinstance(result, dict):
data = list(result.items())
res = np.array(data, dtype=object)
return np.copy(res)
else:
return np.copy(np.asarray(result, dtype=result.dtype))

44
shark/iree_utils/cpu_utils.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# All the iree_cpu related functionalities go here.
import subprocess
# Get the default cpu args.
def get_iree_cpu_args():
find_triple_cmd = "uname -s -m"
os_name, proc_name = (
subprocess.run(
find_triple_cmd, shell=True, stdout=subprocess.PIPE, check=True
)
.stdout.decode("utf-8")
.split()
)
if os_name == "Darwin":
find_kernel_version_cmd = "uname -r"
kernel_version = subprocess.run(
find_kernel_version_cmd,
shell=True,
stdout=subprocess.PIPE,
check=True,
).stdout.decode("utf-8")
target_triple = f"{proc_name}-apple-darwin{kernel_version}"
elif os_name == "Linux":
target_triple = f"{proc_name}-linux-gnu"
else:
error_message = f"OS Type f{os_name} not supported and triple can't be determined, open issue to dSHARK team please :)"
raise Exception(error_message)
print(f"Target triple found:{target_triple}")
return [f"-iree-llvm-target-triple={target_triple}"]

66
shark/cuda_utils.py → shark/iree_utils/gpu_utils.py
View File

@@ -12,9 +12,27 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# All the iree_gpu related functionalities go here.
import iree.runtime as ireert
import ctypes
#Some constants taken from cuda.h
# Get the default gpu args given the architecture.
def get_iree_gpu_args():
ireert.flags.FUNCTION_INPUT_VALIDATION = False
ireert.flags.parse_flags("--cuda_allow_inline_execution")
# TODO: Give the user_interface to pass the sm_arch.
sm_arch = get_cuda_sm_cc()
if sm_arch in ["sm_70", "sm_72", "sm_75", "sm_80", "sm_84", "sm_86"]:
return [
"--iree-hal-cuda-disable-loop-nounroll-wa",
f"--iree-hal-cuda-llvm-target-arch={sm_arch}",
]
else:
return ["--iree-hal-cuda-disable-loop-nounroll-wa"]
# Some constants taken from cuda.h
CUDA_SUCCESS = 0
CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT = 16
CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR = 39
@@ -23,7 +41,7 @@ CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE = 36
def get_cuda_sm_cc():
libnames = ('libcuda.so', 'libcuda.dylib', 'cuda.dll')
libnames = ("libcuda.so", "libcuda.dylib", "cuda.dll")
for libname in libnames:
try:
cuda = ctypes.CDLL(libname)
@@ -32,10 +50,10 @@ def get_cuda_sm_cc():
else:
break
else:
raise OSError("could not load any of: " + ' '.join(libnames))
raise OSError("could not load any of: " + " ".join(libnames))
nGpus = ctypes.c_int()
name = b' ' * 100
name = b" " * 100
cc_major = ctypes.c_int()
cc_minor = ctypes.c_int()
@@ -47,31 +65,43 @@ def get_cuda_sm_cc():
result = cuda.cuInit(0)
if result != CUDA_SUCCESS:
cuda.cuGetErrorString(result, ctypes.byref(error_str))
print("cuInit failed with error code %d: %s" %
(result, error_str.value.decode()))
print(
"cuInit failed with error code %d: %s"
% (result, error_str.value.decode())
)
return 1
result = cuda.cuDeviceGetCount(ctypes.byref(nGpus))
if result != CUDA_SUCCESS:
cuda.cuGetErrorString(result, ctypes.byref(error_str))
print("cuDeviceGetCount failed with error code %d: %s" %
(result, error_str.value.decode()))
print(
"cuDeviceGetCount failed with error code %d: %s"
% (result, error_str.value.decode())
)
return 1
print("Found %d device(s)." % nGpus.value)
for i in range(nGpus.value):
result = cuda.cuDeviceGet(ctypes.byref(device), i)
if result != CUDA_SUCCESS:
cuda.cuGetErrorString(result, ctypes.byref(error_str))
print("cuDeviceGet failed with error code %d: %s" %
(result, error_str.value.decode()))
print(
"cuDeviceGet failed with error code %d: %s"
% (result, error_str.value.decode())
)
return 1
print("Device: %d" % i)
if cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name),
device) == CUDA_SUCCESS:
print(" Name: %s" % (name.split(b'\0', 1)[0].decode()))
if cuda.cuDeviceComputeCapability(ctypes.byref(cc_major),
ctypes.byref(cc_minor),
device) == CUDA_SUCCESS:
print(" Compute Capability: %d.%d" %
(cc_major.value, cc_minor.value))
if (
cuda.cuDeviceGetName(ctypes.c_char_p(name), len(name), device)
== CUDA_SUCCESS
):
print(" Name: %s" % (name.split(b"\0", 1)[0].decode()))
if (
cuda.cuDeviceComputeCapability(
ctypes.byref(cc_major), ctypes.byref(cc_minor), device
)
== CUDA_SUCCESS
):
print(
" Compute Capability: %d.%d" % (cc_major.value, cc_minor.value)
)
sm = f"sm_{cc_major.value}{cc_minor.value}"
return sm

44
shark/iree_utils/vulkan_utils.py Normal file
View File

@@ -0,0 +1,44 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# All the iree_vulkan related functionalities go here.
from shark.iree_utils._common import run_cmd
def get_vulkan_triple_flag():
vulkan_device_cmd = "vulkaninfo | grep deviceName | awk 'END{{print $NF}}'"
vulkan_device = run_cmd(vulkan_device_cmd).strip()
if vulkan_device == "M1":
print("Found Apple Device. Using m1-moltenvk-macos")
return "-iree-vulkan-target-triple=m1-moltenvk-macos"
elif vulkan_device == "A100-SXM4-40GB":
print("Found Nvidia Device. Using ampere-rtx3080-linux")
return "-iree-vulkan-target-triple=ampere-rtx3080-linux"
else:
print(
"""Optimized kernel for your target device is not added yet.
Contact SHARK Admin on discord[https://discord.com/invite/RUqY2h2s9u]
or pull up an issue."""
)
return None
def get_iree_vulkan_args():
# vulkan_flag = ["--iree-flow-demote-i64-to-i32"]
vulkan_flag = []
vulkan_triple_flag = get_vulkan_triple_flag()
if vulkan_triple_flag is not None:
vulkan_flag.append(vulkan_triple_flag)
return vulkan_flag

61
shark/model_annotation.py
View File

@@ -21,7 +21,8 @@ from iree.compiler import ir
from iree.compiler.transforms import ireec as ireec_trans
MATMUL_OP_NAMES = set(
["linalg.matmul", "linalg.batch_matmul", "mhlo.dot", "mhlo.dot_general"])
["linalg.matmul", "linalg.batch_matmul", "mhlo.dot", "mhlo.dot_general"]
)
idx = 0
@@ -47,7 +48,8 @@ def model_annotation(ctx: ir.Context, *, input_contents: str, config_path: str):
# - Disables verification (already done above)
# - Writes as binary, avoiding costly unicode conversions
sys.stdout.buffer.write(
module.operation.get_asm(assume_verified=True, binary=True))
module.operation.get_asm(assume_verified=True, binary=True)
)
return module
@@ -61,14 +63,21 @@ def walk_children(op: ir.Operation, configs: List[Dict]):
child_op = child_op.operation
if child_op.name in MATMUL_OP_NAMES:
global idx
tile_sizes, pipeline, workgroup_size, \
split_k, pipeline_depth = parse_config(configs[idx])
(
tile_sizes,
pipeline,
workgroup_size,
split_k,
pipeline_depth,
) = parse_config(configs[idx])
add_compilation_info(child_op,
tile_sizes=tile_sizes,
pipeline=pipeline,
workgroup_size=workgroup_size,
pipeline_depth=pipeline_depth)
add_compilation_info(
child_op,
tile_sizes=tile_sizes,
pipeline=pipeline,
workgroup_size=workgroup_size,
pipeline_depth=pipeline_depth,
)
if split_k:
add_split_k(child_op, split_k)
@@ -80,8 +89,11 @@ def walk_children(op: ir.Operation, configs: List[Dict]):
def parse_config(config: Dict):
if config["pipeline"] == "GPU" or config["pipeline"] == "GPU_TENSORCORE":
pipeline = "LLVMGPUMatmulSimt" if config[
"pipeline"] == "GPU" else "LLVMGPUMatmulTensorCore"
pipeline = (
"LLVMGPUMatmulSimt"
if config["pipeline"] == "GPU"
else "LLVMGPUMatmulTensorCore"
)
tile_sizes = [config["work_group_tile_sizes"]]
workgroup_size = config["work_group_sizes"]
try:
@@ -95,8 +107,9 @@ def parse_config(config: Dict):
else:
pipeline = config["pipeline"]
tile_sizes = [
config["work_group_tile_sizes"], config["l1_tile_sizes"],
config["vector_tile_sizes"]
config["work_group_tile_sizes"],
config["l1_tile_sizes"],
config["vector_tile_sizes"],
]
workgroup_size = []
split_k = None
@@ -104,9 +117,13 @@ def parse_config(config: Dict):
return tile_sizes, pipeline, workgroup_size, split_k, pipeline_depth
def add_compilation_info(op: ir.Operation, tile_sizes: List[List[int]],
pipeline: str, workgroup_size: List[int],
pipeline_depth: int):
def add_compilation_info(
op: ir.Operation,
tile_sizes: List[List[int]],
pipeline: str,
workgroup_size: List[int],
pipeline_depth: int,
):
# We don't have a Python binding for CompilationInfo, so we just parse
# its string form.
if pipeline_depth:
@@ -114,13 +131,15 @@ def add_compilation_info(op: ir.Operation, tile_sizes: List[List[int]],
f"#iree_codegen.compilation_info<"
f"lowering_config = <tile_sizes = {repr(tile_sizes)}>, "
f"translation_info = <{pipeline} pipeline_depth = {pipeline_depth}>, "
f"workgroup_size = {repr(workgroup_size)}>")
f"workgroup_size = {repr(workgroup_size)}>"
)
else:
attr = ir.Attribute.parse(
f"#iree_codegen.compilation_info<"
f"lowering_config = <tile_sizes = {repr(tile_sizes)}>, "
f"translation_info = <{pipeline}>, "
f"workgroup_size = {repr(workgroup_size)}>")
f"workgroup_size = {repr(workgroup_size)}>"
)
op.attributes["compilation_info"] = attr
@@ -138,6 +157,6 @@ def create_context() -> ir.Context:
if __name__ == "__main__":
with create_context() as ctx:
model_annotation(ctx,
input_contents=sys.argv[1],
config_path=sys.argv[2])
model_annotation(
ctx, input_contents=sys.argv[1], config_path=sys.argv[2]
)

46
shark/parser.py
View File

@@ -21,7 +21,8 @@ def dir_path(path):
return path
else:
raise argparse.ArgumentTypeError(
f"readable_dir:{path} is not a valid path")
f"readable_dir:{path} is not a valid path"
)
def dir_file(path):
@@ -29,43 +30,52 @@ def dir_file(path):
return path
else:
raise argparse.ArgumentTypeError(
f"readable_file:{path} is not a valid file")
f"readable_file:{path} is not a valid file"
)
parser = argparse.ArgumentParser(description='SHARK runner.')
parser = argparse.ArgumentParser(description="SHARK runner.")
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Device on which shark_runner runs. options are cpu, gpu, and vulkan")
help="Device on which shark_runner runs. options are cpu, gpu, and vulkan",
)
parser.add_argument(
"--repro_dir",
help=
"Directory to which module files will be saved for reproduction or debugging.",
help="Directory to which module files will be saved for reproduction or debugging.",
type=dir_path,
default="/tmp/")
parser.add_argument("--save_mlir",
default=False,
action="store_true",
help="Saves input MLIR module to /tmp/ directory.")
parser.add_argument("--save_vmfb",
default=False,
action="store_true",
help="Saves iree .vmfb module to /tmp/ directory.")
default="/tmp/",
)
parser.add_argument(
"--save_mlir",
default=False,
action="store_true",
help="Saves input MLIR module to /tmp/ directory.",
)
parser.add_argument(
"--save_vmfb",
default=False,
action="store_true",
help="Saves iree .vmfb module to /tmp/ directory.",
)
parser.add_argument(
"--model_config_path",
help="Directory to where the tuned model config file is located.",
default=None)
default=None,
)
parser.add_argument(
"--num_warmup_iterations",
type=int,
default=2,
help="Run the model for the specified number of warmup iterations.")
help="Run the model for the specified number of warmup iterations.",
)
parser.add_argument(
"--num_iterations",
type=int,
default=1,
help="Run the model for the specified number of iterations.")
help="Run the model for the specified number of iterations.",
)
shark_args, unknown = parser.parse_known_args()

166
shark/shark_benchmark_runner.py Normal file
View File

@@ -0,0 +1,166 @@
# Copyright 2020 The Nod Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from shark.shark_runner import SharkRunner
from shark.iree_utils.compile_utils import export_iree_module_to_vmfb
from shark.iree_utils.benchmark_utils import (
build_benchmark_args,
run_benchmark_module,
)
from shark.parser import shark_args
from datetime import datetime
import time
import csv
import os
class SharkBenchmarkRunner(SharkRunner):
# SharkRunner derived class with Benchmarking capabilities.
def __init__(
self,
model,
input: tuple,
dynamic: bool = False,
device: str = None,
jit_trace: bool = False,
from_aot: bool = False,
frontend: str = "torch",
):
SharkRunner.__init__(
self, model, input, dynamic, device, jit_trace, from_aot, frontend
)
if self.vmfb_file == None:
self.vmfb_file = export_iree_module_to_vmfb(
self.model, device, shark_args.repro_dir, frontend
)
self.benchmark_cl = build_benchmark_args(
self.vmfb_file, device, input, frontend, from_aot
)
def benchmark_frontend(self, inputs):
if self.frontend in ["pytorch", "torch"]:
return self.benchmark_torch(inputs)
elif self.frontend in ["tensorflow", "tf"]:
return self.benchmark_tf(inputs)
def benchmark_torch(self, inputs):
inputs = self.input if self.from_aot else inputs
inputs = inputs[0]
for i in range(shark_args.num_warmup_iterations):
self.frontend_model.forward(inputs)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.frontend_model.forward(inputs)
if i == shark_args.num_iterations - 1:
end = time.time()
break
print(
f"Torch benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [
f"{shark_args.num_iterations/(end-begin)}",
f"{((end-begin)/shark_args.num_iterations)*1000}",
]
def benchmark_tf(self, inputs):
for i in range(shark_args.num_warmup_iterations):
self.frontend_model.forward(*inputs)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.frontend_model.forward(*inputs)
if i == shark_args.num_iterations - 1:
end = time.time()
break
print(
f"TF benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [
f"{shark_args.num_iterations/(end-begin)}",
f"{((end-begin)/shark_args.num_iterations)*1000}",
]
def benchmark_c(self):
result = run_benchmark_module(self.benchmark_cl)
print(f"Shark-{self.frontend} C-benchmark:{result} iter/second")
return [f"{result}", f"{1000/result}"]
def benchmark_python(self, inputs):
inputs = self.input if self.from_aot else inputs
input_list = [x for x in inputs]
for i in range(shark_args.num_warmup_iterations):
self.forward(input_list, self.frontend)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.forward(input_list, self.frontend)
if i == shark_args.num_iterations - 1:
end = time.time()
print(
f"Shark-{self.frontend} Python-benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [
f"{shark_args.num_iterations/(end-begin)}",
f"{((end-begin)/shark_args.num_iterations)*1000}",
]
def benchmark_all(self, inputs):
self.benchmark_frontend(inputs)
self.benchmark_python(inputs)
self.benchmark_c()
def benchmark_all_csv(self, inputs, modelname, dynamic, device_str):
field_names = [
"platform",
"model",
"dynamic",
"device",
"iter/sec",
"ms/iter",
"datetime",
]
platforms = ["frontend", "shark_python", "shark_iree_c"]
if not os.path.exists("bench_results.csv"):
with open("bench_results.csv", mode="w", newline="") as f:
writer = csv.writer(f)
writer.writerow(field_names)
with open("bench_results.csv", mode="a", newline="") as f:
writer = csv.DictWriter(f, fieldnames=field_names)
bench_result = {}
bench_result["model"] = modelname
if dynamic == True:
bench_result["dynamic"] = "True"
else:
bench_result["dynamic"] = "False"
bench_result["device"] = device_str
for p in platforms:
if p == "frontend":
bench_result["platform"] = "frontend"
bench_result["iter/sec"] = self.benchmark_frontend(inputs)[
0
]
bench_result["ms/iter"] = self.benchmark_frontend(inputs)[1]
elif p == "shark_python":
bench_result["platform"] = "shark_python"
bench_result["iter/sec"] = self.benchmark_python(inputs)[0]
bench_result["ms/iter"] = self.benchmark_python(inputs)[1]
else:
bench_result["platform"] = "shark_iree_c"
bench_result["iter/sec"] = self.benchmark_c()[0]
bench_result["ms/iter"] = self.benchmark_c()[1]
bench_result["datetime"] = str(datetime.now())
writer.writerow(bench_result)

58
shark/shark_downloader.py
View File

@@ -22,19 +22,20 @@ input_type_to_np_dtype = {
"float64": np.float64,
"bool": np.bool_,
"int32": np.int32,
"int64": np.int64
"int64": np.int64,
}
class SharkDownloader:
def __init__(self,
model_name: str,
tank_url: str = "https://storage.googleapis.com/shark_tank",
local_tank_dir: str = "./../gen_shark_tank/tflite",
model_type: str = "tflite-tosa",
input_json: str = "input.json",
input_type: str = "int32"):
def __init__(
self,
model_name: str,
tank_url: str = "https://storage.googleapis.com/shark_tank",
local_tank_dir: str = "./../gen_shark_tank/tflite",
model_type: str = "tflite-tosa",
input_json: str = "input.json",
input_type: str = "int32",
):
self.model_name = model_name
self.local_tank_dir = local_tank_dir
self.tank_url = tank_url
@@ -65,37 +66,52 @@ class SharkDownloader:
print("load json inputs")
if self.model_type in ["tflite-tosa"]:
args = []
with open(self.input_json, 'r') as f:
with open(self.input_json, "r") as f:
args = json.load(f)
self.inputs = [np.asarray(arg, dtype=self.input_type) for arg in args]
self.inputs = [
np.asarray(arg, dtype=self.input_type) for arg in args
]
else:
print("No json input required for current model. You could call setup_inputs(you_inputs).")
print(
"No json input required for current model. You could call setup_inputs(you_inputs)."
)
return self.inputs
def load_mlir_model(self):
if self.model_type in ["tflite-tosa"]:
workdir = os.path.join(os.path.dirname(__file__), self.local_tank_dir)
workdir = os.path.join(
os.path.dirname(__file__), self.local_tank_dir
)
os.makedirs(workdir, exist_ok=True)
print(f"TMP_MODEL_DIR = {workdir}")
# use model name get dir.
model_name_dir = os.path.join(workdir, str(self.model_name))
if not os.path.exists(model_name_dir):
print("Model has not been download."
"shark_downloader will automatically download by tank_url if provided."
" You can also manually to download the model from shark_tank by yourself.")
print(
"Model has not been download."
"shark_downloader will automatically download by tank_url if provided."
" You can also manually to download the model from shark_tank by yourself."
)
os.makedirs(model_name_dir, exist_ok=True)
print(f"TMP_MODELNAME_DIR = {model_name_dir}")
mlir_url = self.tank_url + "/tflite/" + str(self.model_name) + "/"+str(self.model_name) + '_tosa.mlir'
self.mlir_file = '/'.join(
[model_name_dir, str(self.model_name) + '_tosa.mlir'])
mlir_url = (
self.tank_url
+ "/tflite/"
+ str(self.model_name)
+ "/"
+ str(self.model_name)
+ "_tosa.mlir"
)
self.mlir_file = "/".join(
[model_name_dir, str(self.model_name) + "_tosa.mlir"]
)
if os.path.exists(self.mlir_file):
print("Model has been downloaded before.", self.mlir_file)
else:
print("Download mlir model", mlir_url)
urllib.request.urlretrieve(mlir_url,
self.mlir_file)
urllib.request.urlretrieve(mlir_url, self.mlir_file)
print("Get tosa.mlir model return")
with open(self.mlir_file) as f:

111
shark/shark_importer.py
View File

@@ -15,19 +15,20 @@ from shark.iree_utils import IREE_TARGET_MAP
class SharkImporter:
def __init__(self,
model_name: str=None,
model_path: str=None,
model_type: str = "tflite",
model_source_hub: str = "tfhub",
device: str = None,
dynamic: bool = False,
jit_trace: bool = False,
benchmark_mode: bool = False,
input_details=None,
output_details=None,
tank_url: str = None):
def __init__(
self,
model_name: str = None,
model_path: str = None,
model_type: str = "tflite",
model_source_hub: str = "tfhub",
device: str = None,
dynamic: bool = False,
jit_trace: bool = False,
benchmark_mode: bool = False,
input_details=None,
output_details=None,
tank_url: str = None,
):
self.model_name = model_name
self.model_path = model_path
self.model_type = model_type
@@ -45,7 +46,9 @@ class SharkImporter:
# create tmp model file directory
if self.model_path is None and self.model_name is None:
print("Error. No model_path, No model name,Please input either one.")
print(
"Error. No model_path, No model name,Please input either one."
)
return
if self.model_source_hub == "tfhub":
@@ -56,43 +59,62 @@ class SharkImporter:
print("Error, load tflite model fail")
return
if (self.input_details == None) or \
(self.output_details == None):
print("Setting up tflite interpreter to get model input details")
if (self.input_details == None) or (
self.output_details == None
):
print(
"Setting up tflite interpreter to get model input details"
)
self.tflite_interpreter = tf.lite.Interpreter(
model_path=self.tflite_saving_file)
model_path=self.tflite_saving_file
)
self.tflite_interpreter.allocate_tensors()
# default input initialization
self.input_details, self.output_details = self.get_model_details(
)
(
self.input_details,
self.output_details,
) = self.get_model_details()
inputs = self.generate_inputs(
self.input_details) # device_inputs
self.input_details
) # device_inputs
self.setup_inputs(inputs)
def load_tflite_model(self):
print("Setting up for TMP_DIR")
tflite_workdir = os.path.join(os.path.dirname(__file__), "./../gen_shark_tank/tflite")
tflite_workdir = os.path.join(
os.path.dirname(__file__), "./../gen_shark_tank/tflite"
)
os.makedirs(tflite_workdir, exist_ok=True)
print(f"TMP_TFLITE_DIR = {tflite_workdir}")
# use model name get dir.
tflite_model_name_dir = os.path.join(tflite_workdir, str(self.model_name))
tflite_model_name_dir = os.path.join(
tflite_workdir, str(self.model_name)
)
# TODO Download model from google bucket to tflite_model_name_dir by tank_url
os.makedirs(tflite_model_name_dir, exist_ok=True)
print(f"TMP_TFLITE_MODELNAME_DIR = {tflite_model_name_dir}")
self.tflite_saving_file = '/'.join(
[tflite_model_name_dir, str(self.model_name) + '_tflite.tflite'])
self.tflite_tosa_file = '/'.join(
[tflite_model_name_dir, str(self.model_name) + '_tosa.mlir'])
self.tflite_saving_file = "/".join(
[tflite_model_name_dir, str(self.model_name) + "_tflite.tflite"]
)
self.tflite_tosa_file = "/".join(
[tflite_model_name_dir, str(self.model_name) + "_tosa.mlir"]
)
if os.path.exists(self.tflite_saving_file):
print("Local address for tflite model file Exists: ", self.tflite_saving_file)
print(
"Local address for tflite model file Exists: ",
self.tflite_saving_file,
)
else:
print("No local tflite file, Download tflite model")
if self.model_path is None:
# get model file from tflite_model_list.csv or download from gs://bucket
print("No model_path, get from tflite_model_list.csv")
tflite_model_list_path = os.path.join(os.path.dirname(__file__), "../tank/tflite/tflite_model_list.csv")
tflite_model_list_path = os.path.join(
os.path.dirname(__file__),
"../tank/tflite/tflite_model_list.csv",
)
tflite_model_list = csv.reader(open(tflite_model_list_path))
for row in tflite_model_list:
if str(row[0]) == self.model_name:
@@ -100,12 +122,13 @@ class SharkImporter:
if self.model_path is None:
print("Error, No model path find in tflite_model_list.csv")
return False
urllib.request.urlretrieve(self.model_path,
self.tflite_saving_file)
urllib.request.urlretrieve(self.model_path, self.tflite_saving_file)
if os.path.exists(self.tflite_tosa_file):
print("Exists", self.tflite_tosa_file)
else:
print("No tflite tosa.mlir, please use python generate_sharktank.py to download tosa model")
print(
"No tflite tosa.mlir, please use python generate_sharktank.py to download tosa model"
)
return True
def generate_inputs(self, input_details):
@@ -139,21 +162,25 @@ class SharkImporter:
tosa_model = []
with open(self.tflite_tosa_file) as f:
tosa_model = f.read()
self.shark_module = SharkInference(tosa_model,
self.inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=self.jit_trace)
self.shark_module = SharkInference(
tosa_model,
self.inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=self.jit_trace,
)
self.shark_module.set_frontend("tflite-tosa")
self.shark_module.compile()
else:
# compile and run tfhub tflite
print("Inference tfhub tflite model")
self.shark_module = SharkInference(self.tflite_saving_file,
self.inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=self.jit_trace)
self.shark_module = SharkInference(
self.tflite_saving_file,
self.inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=self.jit_trace,
)
self.shark_module.set_frontend("tflite")
self.shark_module.compile()
elif self.model_source_hub == "huggingface":

79
shark/shark_inference.py
View File

@@ -9,10 +9,10 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend
from shark.parser import shark_args
from shark.shark_runner import SharkRunner, SharkBenchmarkRunner
from shark.shark_runner import SharkRunner
import sys
import os
# Prints to stderr.
@@ -23,13 +23,15 @@ def print_err(*a):
class SharkInference:
"""Inference API targeting pytorch, tensorflow, linalg, mhlo and tosa frontend."""
def __init__(self,
model,
input: tuple,
device: str = None,
dynamic: bool = False,
jit_trace: bool = False,
benchmark_mode: bool = False):
def __init__(
self,
model,
input: tuple,
device: str = None,
dynamic: bool = False,
jit_trace: bool = False,
benchmark_mode: bool = False,
):
self.model = model
self.input = input
self.dynamic = dynamic
@@ -49,27 +51,48 @@ class SharkInference:
# Sets the frontend i.e `pytorch` or `tensorflow`.
def set_frontend(self, frontend: str):
if frontend not in [
"pytorch", "torch", "tensorflow", "tf", "mhlo", "linalg",
"tosa", "tflite", "tflite-tosa"
"pytorch",
"torch",
"tensorflow",
"tf",
"mhlo",
"linalg",
"tosa",
"tflite",
"tflite-tosa",
]:
print_err("frontend not supported.")
else:
self.frontend = frontend
def compile(self):
# Inference do not use AOT.
# Inference do not use AOT. # TODO: Remove the from_aot arg as it's not
# needed.
from_aot = False
if (self.benchmark_mode == True):
self.shark_runner = SharkBenchmarkRunner(self.model, self.input,
self.dynamic, self.device,
self.jit_trace, from_aot,
self.frontend)
if self.benchmark_mode == True:
# Only import shark_benchmark runner when needed.
from shark.shark_benchmark_runner import SharkBenchmarkRunner
self.shark_runner = SharkBenchmarkRunner(
self.model,
self.input,
self.dynamic,
self.device,
self.jit_trace,
from_aot,
self.frontend,
)
else:
self.shark_runner = SharkRunner(self.model, self.input,
self.dynamic, self.device,
self.jit_trace, from_aot,
self.frontend,
self.model_config_path)
self.shark_runner = SharkRunner(
self.model,
self.input,
self.dynamic,
self.device,
self.jit_trace,
from_aot,
self.frontend,
self.model_config_path,
)
# inputs are considered to be np.array.
def forward(self, inputs):
@@ -81,6 +104,9 @@ class SharkInference:
input_list = [x.numpy() for x in inputs]
return self.shark_runner.forward(input_list, self.frontend)
def import_mlir(self, model_name="model", dir=os.getcwd()):
self.shark_runner.import_mlir(model_name, dir)
# Saves the .vmfb module.
def save_module(self, dir=None):
if dir is None:
@@ -89,9 +115,10 @@ class SharkInference:
######### Benchmark Related Functions #########
def benchmark_mode(func):
def inner(self, *args, **kwargs):
assert self.benchmark_mode, "SharkRunner needs to be in benchmark mode to run benchmark methods."
assert (
self.benchmark_mode
), "SharkRunner needs to be in benchmark mode to run benchmark methods."
return func(self, *args, **kwargs)
return inner
@@ -114,4 +141,6 @@ class SharkInference:
@benchmark_mode
def benchmark_all_csv(self, inputs, modelname, dynamic, device_str):
self.shark_runner.benchmark_all_csv(inputs, modelname, dynamic, device_str)
self.shark_runner.benchmark_all_csv(
inputs, modelname, dynamic, device_str
)

208
shark/shark_runner.py
View File

@@ -20,13 +20,14 @@ from torch_mlir_e2e_test.eager_backends.refbackend import EagerModeRefBackend
from shark.iree_eager_backend import EagerModeIREELinalgOnTensorsBackend
from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend
from shark.iree_utils import get_results, get_iree_compiled_module, export_iree_module_to_vmfb, export_module_to_mlir_file, build_benchmark_args, run_benchmark_module
import os
from shark.parser import shark_args
from tqdm import tqdm
from datetime import datetime
import time
import csv
from shark.iree_utils.compile_utils import (
get_iree_compiled_module,
export_iree_module_to_vmfb,
export_module_to_mlir_file,
get_results,
)
import os
class SharkRunner:
@@ -57,24 +58,27 @@ class SharkRunner:
elif self.frontend in ["pytorch", "torch"]:
# get torch-mlir dialect
# self.model = torch.Module
# Lowers in linalg dialect.
# TODO assert
self.model = get_torch_mlir_module(self.model, input, dynamic,
jit_trace, from_aot)
# TODO tosa dialect from torch_module.
self.model = get_torch_mlir_module(
self.model, input, dynamic, jit_trace, from_aot
)
elif self.frontend in ["tensorflow", "tf"]:
# get mhlo dialect
# self.model = tf.Module
# TODO assert
self.model = tfc.compile_module(self.model,
exported_names=[func_name],
import_only=True)
self.model = tfc.compile_module(
self.model, exported_names=[func_name], import_only=True
)
elif self.frontend in ["tflite"]:
print("Setting up for IREE compiler tflite")
# get tosa dialect
# self.model = model.tflite
# TODO assert
self.model = ireec_tflite.compile_file(self.model,
input_type="tosa",
import_only=True)
self.model = ireec_tflite.compile_file(
self.model, input_type="tosa", import_only=True
)
func_name = "main"
# TODO: We can capture the .vmfb module here and later use it for saving
@@ -82,29 +86,44 @@ class SharkRunner:
(
self.iree_compilation_module,
self.iree_config,
) = get_iree_compiled_module(self.model,
self.device,
self.frontend,
func_name=func_name,
model_config_path=model_config_path)
) = get_iree_compiled_module(
self.model,
self.device,
self.frontend,
func_name=func_name,
model_config_path=model_config_path,
)
# Debugging Options:
if shark_args.save_mlir:
export_module_to_mlir_file(self.model, self.frontend,
shark_args.repro_dir)
export_module_to_mlir_file(
self.model, self.frontend, shark_args.repro_dir
)
if shark_args.save_vmfb:
self.vmfb_file = self.save_module(shark_args.repro_dir)
# All the timings and benchmarking can be done here.
def forward(self, input, frontend):
return get_results(self.iree_compilation_module, input,
self.iree_config, frontend)
return get_results(
self.iree_compilation_module, input, self.iree_config, frontend
)
# Saves the .mlir file, can be in tosa, linalg or mhlo dialect.
# torch-mlir can export tosa or linalg dialects.
# tensorflow models get exported to mhlo dialect.
def import_mlir(self, model_name, dir):
filename = os.path.join(dir, f"{model_name}_{self.frontend}.mlir")
with open(filename, "w") as f:
f.write(self.model)
print(f"Saved mlir in {filename}.")
return filename
# TODO: Instead of passing directory and having names decided by the module
# , user may want to save the module with manual names.
def save_module(self, dir=os.getcwd()):
return export_iree_module_to_vmfb(self.model, self.device, dir,
self.frontend)
return export_iree_module_to_vmfb(
self.model, self.device, dir, self.frontend
)
# TODO: Load a module and directly use it, we will need to set the frontend
# in this case.
@@ -112,150 +131,17 @@ class SharkRunner:
pass
# TODO: Document shark_eager mode.
class SharkEagerMode:
def __init__(self, device="cpu"):
if device == "refbackend":
torch_mlir_tensor.backend = EagerModeRefBackend()
else:
torch_mlir_tensor.backend = EagerModeIREELinalgOnTensorsBackend(
device)
device
)
self.guard = enable_torch_dispatch_mode(TorchMLIRTensor)
self.guard.__enter__()
def __del__(self):
self.guard.__exit__(None, None, None)
class SharkBenchmarkRunner(SharkRunner):
# SharkRunner derived class with Benchmarking capabilities.
def __init__(
self,
model,
input: tuple,
dynamic: bool = False,
device: str = None,
jit_trace: bool = False,
from_aot: bool = False,
frontend: str = "torch",
):
SharkRunner.__init__(self, model, input, dynamic, device, jit_trace,
from_aot, frontend)
if (self.vmfb_file == None):
self.vmfb_file = export_iree_module_to_vmfb(self.model, device,
shark_args.repro_dir,
frontend)
self.benchmark_cl = build_benchmark_args(self.vmfb_file, device, input,
frontend, from_aot)
def benchmark_frontend(self, inputs):
if self.frontend in ["pytorch", "torch"]:
return self.benchmark_torch(inputs)
elif self.frontend in ["tensorflow", "tf"]:
return self.benchmark_tf(inputs)
def benchmark_torch(self, inputs):
inputs = self.input if self.from_aot else inputs
inputs = inputs[0]
for i in range(shark_args.num_warmup_iterations):
self.frontend_model.forward(inputs)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.frontend_model.forward(inputs)
if i == shark_args.num_iterations - 1:
end = time.time()
break
print(
f"Torch benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"]
def benchmark_tf(self, inputs):
for i in range(shark_args.num_warmup_iterations):
self.frontend_model.forward(*inputs)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.frontend_model.forward(*inputs)
if i == shark_args.num_iterations - 1:
end = time.time()
break
print(
f"TF benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"]
def benchmark_c(self):
result = run_benchmark_module(self.benchmark_cl)
print(f"Shark-{self.frontend} C-benchmark:{result} iter/second")
return [f"{result}", f"{1000/result}"]
def benchmark_python(self, inputs):
inputs = self.input if self.from_aot else inputs
input_list = [x for x in inputs]
for i in range(shark_args.num_warmup_iterations):
self.forward(input_list, self.frontend)
begin = time.time()
for i in range(shark_args.num_iterations):
out = self.forward(input_list, self.frontend)
if i == shark_args.num_iterations - 1:
end = time.time()
print(
f"Shark-{self.frontend} Python-benchmark:{shark_args.num_iterations/(end-begin)} iter/second, Total Iterations:{shark_args.num_iterations}"
)
return [f"{shark_args.num_iterations/(end-begin)}", f"{((end-begin)/shark_args.num_iterations)*1000}"]
def benchmark_all(self, inputs):
self.benchmark_frontend(inputs)
self.benchmark_python(inputs)
self.benchmark_c()
def benchmark_all_csv(self, inputs, modelname, dynamic, device_str):
field_names = [
'platform',
'model',
'dynamic',
'device',
'iter/sec',
'ms/iter',
'datetime'
]
platforms = [
'frontend',
'shark_python',
'shark_iree_c'
]
if not os.path.exists('bench_results.csv'):
with open('bench_results.csv', mode='w', newline='') as f:
writer = csv.writer(f)
writer.writerow(field_names)
with open('bench_results.csv', mode='a', newline='') as f:
writer = csv.DictWriter(f, fieldnames=field_names)
bench_result = {}
bench_result['model'] = modelname
if dynamic == True:
bench_result['dynamic'] = "True"
else:
bench_result['dynamic'] = "False"
bench_result['device'] = device_str
for p in platforms:
if p == 'frontend':
bench_result['platform'] = "frontend"
bench_result['iter/sec'] = self.benchmark_frontend(inputs)[0]
bench_result['ms/iter'] = self.benchmark_frontend(inputs)[1]
elif p == 'shark_python':
bench_result['platform'] = "shark_python"
bench_result['iter/sec'] = self.benchmark_python(inputs)[0]
bench_result['ms/iter'] = self.benchmark_python(inputs)[1]
else:
bench_result['platform'] = "shark_iree_c"
bench_result['iter/sec'] = self.benchmark_c()[0]
bench_result['ms/iter'] = self.benchmark_c()[1]
bench_result['datetime'] = str(datetime.now())
writer.writerow(bench_result)

53
shark/shark_trainer.py
View File

@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from shark.torch_mlir_utils import get_torch_mlir_module, run_on_refbackend
from shark.iree_utils import get_results, get_iree_compiled_module, export_iree_module_to_vmfb
from shark.parser import shark_args
from shark.shark_runner import SharkRunner
from shark.backward_makefx import MakeFxModule
@@ -56,7 +54,13 @@ class SharkTrainer:
# Sets the frontend i.e `pytorch` or `tensorflow`.
def set_frontend(self, frontend: str):
if frontend not in [
"pytorch", "torch", "tensorflow", "tf", "mhlo", "linalg", "tosa"
"pytorch",
"torch",
"tensorflow",
"tf",
"mhlo",
"linalg",
"tosa",
]:
print_err("frontend not supported.")
else:
@@ -65,22 +69,32 @@ class SharkTrainer:
# Training function is needed in the case of torch_fn.
def compile(self, training_fn=None):
if self.frontend in ["torch", "pytorch"]:
aot_module = MakeFxModule(self.model,
tuple(self.input),
custom_inference_fn=training_fn)
aot_module = MakeFxModule(
self.model, tuple(self.input), custom_inference_fn=training_fn
)
aot_module.generate_graph()
# Returns the backward graph.
training_graph = aot_module.training_graph
weights = self.get_torch_params()
self.shark_runner = SharkRunner(training_graph,
weights + self.input, self.dynamic,
self.device, self.jit_trace,
self.from_aot, self.frontend)
self.shark_runner = SharkRunner(
training_graph,
weights + self.input,
self.dynamic,
self.device,
self.jit_trace,
self.from_aot,
self.frontend,
)
elif self.frontend in ["tensorflow", "tf", "mhlo"]:
self.shark_runner = SharkRunner(self.model, self.input,
self.dynamic, self.device,
self.jit_trace, self.from_aot,
self.frontend)
self.shark_runner = SharkRunner(
self.model,
self.input,
self.dynamic,
self.device,
self.jit_trace,
self.from_aot,
self.frontend,
)
else:
print_err("Unknown frontend")
return
@@ -98,8 +112,9 @@ class SharkTrainer:
params = [x.numpy() for x in params]
print(f"Training started for {num_iters} iterations:")
for i in tqdm(range(num_iters)):
params = self.shark_runner.forward(params + self.input,
self.frontend)
params = self.shark_runner.forward(
params + self.input, self.frontend
)
return params
@@ -109,15 +124,15 @@ class SharkTrainer:
def _train_tf(self, num_iters):
input_list = []
for x in self.input:
if (isinstance(x, list)):
if isinstance(x, list):
nested_list = []
for val in x:
if (isinstance(val, np.ndarray)):
if isinstance(val, np.ndarray):
nested_list.append(val)
else:
nested_list.append(val.numpy())
input_list.append(nested_list)
elif (isinstance(x, np.ndarray)):
elif isinstance(x, np.ndarray):
input_list.append(x)
else:
input_list.append(x.numpy())

46
shark/tests/test_shark_importer.py
View File

@@ -13,38 +13,48 @@ def generate_inputs(input_details):
args = []
args.append(
np.random.randint(low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"]))
np.random.randint(
low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"],
)
)
args.append(
np.ones(shape=input_details[1]["shape"],
dtype=input_details[1]["dtype"]))
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(shape=input_details[2]["shape"],
dtype=input_details[2]["dtype"]))
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
# A specific case can be run by commenting different cases. Runs all the test
# across cpu, gpu and vulkan according to available drivers.
pytest_param = pytest.mark.parametrize(
('dynamic', 'device'),
("dynamic", "device"),
[
pytest.param(False, 'cpu'),
pytest.param(False, "cpu"),
# TODO: Language models are failing for dynamic case..
pytest.param(True, 'cpu', marks=pytest.mark.skip),
])
pytest.param(True, "cpu", marks=pytest.mark.skip),
],
)
@pytest_param
def test_albert(dynamic, device):
my_shark_importer = SharkImporter(model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device=device,
dynamic=dynamic,
jit_trace=True)
my_shark_importer = SharkImporter(
model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device=device,
dynamic=dynamic,
jit_trace=True,
)
input_details, output_details = my_shark_importer.get_model_details()
inputs = generate_inputs(input_details) # device_inputs
my_shark_importer.compile(inputs)

39
shark/torch_mlir_utils.py
View File

@@ -21,8 +21,10 @@ from torch_mlir.dialects.torch.importer.jit_ir import (
ModuleBuilder,
)
from torch_mlir_e2e_test.torchscript.serialization import (
extract_serializable_annotations, apply_serializable_annotations,
SerializableTest)
extract_serializable_annotations,
apply_serializable_annotations,
SerializableTest,
)
from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend
@@ -38,7 +40,8 @@ def get_module_name_for_asm_dump(module):
if not "torch.debug_module_name" in module.operation.attributes:
return "UnnammedModule"
return StringAttr(
module.operation.attributes["torch.debug_module_name"]).value
module.operation.attributes["torch.debug_module_name"]
).value
def get_input_annotations(inputs: tuple, dynamic: bool) -> list:
@@ -65,8 +68,9 @@ def run_on_refbackend(torch_module, inputs):
return jit_module.forward(np_inputs[0])
def shark_jit_trace(module, input: tuple, dynamic: bool,
tracing_required: bool):
def shark_jit_trace(
module, input: tuple, dynamic: bool, tracing_required: bool
):
"""TODO: Include necessary documentation."""
if not tracing_required:
@@ -76,21 +80,26 @@ def shark_jit_trace(module, input: tuple, dynamic: bool,
actual_script = traced_module._actual_script_module
export(actual_script.forward)
annotate_args_decorator = annotate_args(
get_input_annotations(input, dynamic))
get_input_annotations(input, dynamic)
)
annotate_args_decorator(actual_script.forward)
module = torch.jit.script(actual_script)
# TODO: remove saved annotations.pickle
torchscript_module_bytes = module.save_to_buffer({
"annotations.pkl":
pickle.dumps(extract_serializable_annotations(module))
})
serializable_test = SerializableTest(unique_name="",
program=torchscript_module_bytes,
trace=None)
torchscript_module_bytes = module.save_to_buffer(
{
"annotations.pkl": pickle.dumps(
extract_serializable_annotations(module)
)
}
)
serializable_test = SerializableTest(
unique_name="", program=torchscript_module_bytes, trace=None
)
_extra_files = {"annotations.pkl": ""}
module = torch.jit.load(io.BytesIO(serializable_test.program),
_extra_files=_extra_files)
module = torch.jit.load(
io.BytesIO(serializable_test.program), _extra_files=_extra_files
)
# Load the pickled annotations.
annotations = pickle.loads(_extra_files["annotations.pkl"])
apply_serializable_annotations(module, annotations)

40
tank/conftest.py
View File

@@ -1,18 +1,26 @@
def pytest_addoption(parser):
# Attaches SHARK command-line arguments to the pytest machinery.
parser.addoption("--save_mlir",
action="store_true",
default="False",
help="Pass option to save input MLIR")
parser.addoption("--save_vmfb",
action="store_true",
default="False",
help="Pass option to save IREE output .vmfb")
parser.addoption("--benchmark",
action="store_true",
default="False",
help="Pass option to benchmark and write results.csv")
parser.addoption("--save_temps",
action="store_true",
default="False",
help="Saves IREE reproduction artifacts for filing upstream issues.")
parser.addoption(
"--save_mlir",
action="store_true",
default="False",
help="Pass option to save input MLIR",
)
parser.addoption(
"--save_vmfb",
action="store_true",
default="False",
help="Pass option to save IREE output .vmfb",
)
parser.addoption(
"--benchmark",
action="store_true",
default="False",
help="Pass option to benchmark and write results.csv",
)
parser.addoption(
"--save_temps",
action="store_true",
default="False",
help="Saves IREE reproduction artifacts for filing upstream issues.",
)

20
tank/model_utils.py
View File

@@ -1,10 +1,13 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
import torch
import numpy as np
import torchvision.models as models
from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel
from transformers import (
AutoModelForSequenceClassification,
BertTokenizer,
TFBertModel,
)
torch.manual_seed(0)
@@ -12,17 +15,13 @@ torch.manual_seed(0)
class HuggingFaceLanguage(torch.nn.Module):
def __init__(self, hf_model_name):
super().__init__()
self.model = AutoModelForSequenceClassification.from_pretrained(
hf_model_name, # The pretrained model.
num_labels=
2, # The number of output labels--2 for binary classification.
output_attentions=
False, # Whether the model returns attentions weights.
output_hidden_states=
False, # Whether the model returns all hidden-states.
num_labels=2, # The number of output labels--2 for binary classification.
output_attentions=False, # Whether the model returns attentions weights.
output_hidden_states=False, # Whether the model returns all hidden-states.
torchscript=True,
)
@@ -44,7 +43,6 @@ def get_hf_model(name):
class VisionModule(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
@@ -61,6 +59,7 @@ def get_vision_model(torch_model):
actual_out = model(test_input)
return model, test_input, actual_out
################################################################################
# Utility function for comparing two tensors (torch).
@@ -70,4 +69,3 @@ def compare_tensors(torch_tensor, numpy_tensor):
atol = 1e-03
torch_to_numpy = torch_tensor.detach().numpy()
return np.allclose(torch_to_numpy, numpy_tensor, rtol, atol)

42
tank/model_utils_tf.py
View File

@@ -1,9 +1,12 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
import tensorflow as tf
import numpy as np
from transformers import AutoModelForSequenceClassification, BertTokenizer, TFBertModel
from transformers import (
AutoModelForSequenceClassification,
BertTokenizer,
TFBertModel,
)
##################### Tensorflow Hugging Face LM Models ###################################
MAX_SEQUENCE_LENGTH = 512
@@ -13,20 +16,20 @@ BATCH_SIZE = 1
tf_bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
class TFHuggingFaceLanguage(tf.Module):
class TFHuggingFaceLanguage(tf.Module):
def __init__(self, hf_model_name):
super(TFHuggingFaceLanguage, self).__init__()
# Create a BERT trainer with the created network.
self.m = TFBertModel.from_pretrained(
hf_model_name, from_pt=True)
self.m = TFBertModel.from_pretrained(hf_model_name, from_pt=True)
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=tf_bert_input)
def forward(self, input_ids, attention_mask, token_type_ids):
@@ -36,17 +39,24 @@ class TFHuggingFaceLanguage(tf.Module):
def get_TFhf_model(name):
model = TFHuggingFaceLanguage(name)
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased")
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
encoded_input = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
test_input = (encoded_input["input_ids"], encoded_input["attention_mask"],
encoded_input["token_type_ids"])
tf.convert_to_tensor(encoded_input[key]), 0
)
test_input = (
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
)
actual_out = model.forward(*test_input)
return model, test_input, actual_out
@@ -58,5 +68,3 @@ def compare_tensors_tf(tf_tensor, numpy_tensor):
atol = 1e-03
tf_to_numpy = tf_tensor.numpy()
return np.allclose(tf_to_numpy, numpy_tensor, rtol, atol)

53
tank/pytorch/MiniLM-L12-H384-uncased/MiniLM-L12-H384-uncased_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_hf_model, compare_tensors
from shark.parser import shark_args
@@ -10,8 +10,8 @@ import pytest
torch.manual_seed(0)
class MiniLMModuleTester:
class MiniLMModuleTester:
def __init__(
self,
dynamic=False,
@@ -25,24 +25,29 @@ class MiniLMModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_hf_model("microsoft/MiniLM-L12-H384-uncased")
model, input, act_out = get_hf_model(
"microsoft/MiniLM-L12-H384-uncased"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(model, (input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True)
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class MiniLMModuleTest(unittest.TestCase):
class MiniLMModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = MiniLMModuleTester(self)
@@ -50,20 +55,24 @@ class MiniLMModuleTest(unittest.TestCase):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.xfail(reason="language models failing for dynamic case")
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
@@ -71,9 +80,9 @@ class MiniLMModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
@@ -81,14 +90,14 @@ class MiniLMModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

54
tank/pytorch/albert-base-v2/albert-base-v2_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_hf_model, compare_tensors
from shark.parser import shark_args
@@ -8,10 +8,10 @@ import unittest
import numpy as np
import pytest
#torch.manual_seed(0)
# torch.manual_seed(0)
class AlbertModuleTester:
def __init__(
self,
dynamic=False,
@@ -28,21 +28,24 @@ class AlbertModuleTester:
model, input, act_out = get_hf_model("albert-base-v2")
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(model, (input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True)
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class AlbertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = AlbertModuleTester(self)
self.module_tester.save_mlir = self.save_mlir
@@ -51,20 +54,26 @@ class AlbertModuleTest(unittest.TestCase):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
@@ -72,9 +81,9 @@ class AlbertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
@@ -82,13 +91,14 @@ class AlbertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

48
tank/pytorch/alexnet/alexnet_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class AlexnetModuleTester:
class AlexnetModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,69 +26,75 @@ class AlexnetModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.alexnet(pretrained=True))
model, input, act_out = get_vision_model(
models.alexnet(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class AlexnetModuleTest(unittest.TestCase):
class AlexnetModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = AlexnetModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

55
tank/pytorch/bert-base-uncased/bert-base-uncased_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_hf_model, compare_tensors
from shark.parser import shark_args
@@ -8,10 +8,10 @@ import unittest
import numpy as np
import pytest
#torch.manual_seed(0)
# torch.manual_seed(0)
class BertModuleTester:
def __init__(
self,
dynamic=False,
@@ -28,42 +28,51 @@ class BertModuleTester:
model, input, act_out = get_hf_model("bert-base-uncased")
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(model, (input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True)
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class BertModuleTest(unittest.TestCase):
class BertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = BertModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
@@ -71,9 +80,9 @@ class BertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
@@ -81,14 +90,14 @@ class BertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

59
tank/pytorch/distilbert-base-uncased/distilbert-base-uncased_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_hf_model, compare_tensors
from shark.parser import shark_args
@@ -8,10 +8,10 @@ import unittest
import numpy as np
import pytest
#torch.manual_seed(0)
# torch.manual_seed(0)
class DistilBertModuleTester:
def __init__(
self,
dynamic=False,
@@ -28,43 +28,54 @@ class DistilBertModuleTester:
model, input, act_out = get_hf_model("distilbert-base-uncased")
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(model, (input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True)
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class DistilBertModuleTest(unittest.TestCase):
class DistilBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = DistilBertModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.xfail(reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
@@ -72,9 +83,9 @@ class DistilBertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
@@ -82,14 +93,14 @@ class DistilBertModuleTest(unittest.TestCase):
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9554")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

49
tank/pytorch/resnet101/resnet101_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class Resnet101ModuleTester:
class Resnet101ModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,68 +26,75 @@ class Resnet101ModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.resnet101(pretrained=True))
model, input, act_out = get_vision_model(
models.resnet101(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class Resnet101ModuleTest(unittest.TestCase):
class Resnet101ModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = Resnet101ModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

48
tank/pytorch/resnet18/resnet18_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class Resnet18ModuleTester:
class Resnet18ModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,69 +26,75 @@ class Resnet18ModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.resnet18(pretrained=True))
model, input, act_out = get_vision_model(
models.resnet18(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class Resnet18ModuleTest(unittest.TestCase):
class Resnet18ModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = Resnet18ModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

49
tank/pytorch/resnet50/resnet50_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class Resnet50ModuleTester:
class Resnet50ModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,68 +26,75 @@ class Resnet50ModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.resnet50(pretrained=True))
model, input, act_out = get_vision_model(
models.resnet50(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class Resnet50ModuleTest(unittest.TestCase):
class Resnet50ModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = Resnet50ModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

48
tank/pytorch/squeezenet1_0/squeezenet1_0_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class SqueezenetModuleTester:
class SqueezenetModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,69 +26,75 @@ class SqueezenetModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.squeezenet1_0(pretrained=True))
model, input, act_out = get_vision_model(
models.squeezenet1_0(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class SqueezenetModuleTest(unittest.TestCase):
class SqueezenetModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = SqueezenetModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

11
tank/pytorch/v_diffusion/cfg_sample.py
View File

@@ -47,18 +47,23 @@ if os.path.exists(checkpoint):
model.load_state_dict(torch.load(checkpoint, map_location="cpu"))
model = model.to(device).eval().requires_grad_(False)
clip_model_name = model.clip_model if hasattr(model, "clip_model") else "ViT-B/16"
clip_model_name = (
model.clip_model if hasattr(model, "clip_model") else "ViT-B/16"
)
clip_model = clip.load(clip_model_name, jit=False, device=device)[0]
clip_model.eval().requires_grad_(False)
normalize = transforms.Normalize(
mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
mean=[0.48145466, 0.4578275, 0.40821073],
std=[0.26862954, 0.26130258, 0.27577711],
)
zero_embed = torch.zeros([1, clip_model.visual.output_dim], device=device)
target_embeds, weights = [zero_embed], []
txt, weight = parse_prompt(args.prompts[0])
target_embeds.append(clip_model.encode_text(clip.tokenize(txt).to(device)).float())
target_embeds.append(
clip_model.encode_text(clip.tokenize(txt).to(device)).float()
)
weights.append(weight)
weights = torch.tensor([1 - sum(weights), *weights], device=device)

48
tank/pytorch/wide_resnet50_2/wide_resnet50_2_pytorch_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils import get_vision_model, compare_tensors
from shark.parser import shark_args
@@ -11,8 +11,8 @@ import pytest
torch.manual_seed(0)
class WideResnet50ModuleTester:
class WideResnet50ModuleTester:
def __init__(
self,
dynamic=False,
@@ -26,69 +26,75 @@ class WideResnet50ModuleTester:
self.save_vmfb = save_vmfb
def create_and_check_module(self):
model, input, act_out = get_vision_model(models.wide_resnet50_2(pretrained=True))
model, input, act_out = get_vision_model(
models.wide_resnet50_2(pretrained=True)
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
shark_module = SharkInference(
model,
(input,),
device=self.device,
dynamic=self.dynamic,
model,
(input,),
device=self.device,
dynamic=self.dynamic,
)
shark_module.compile()
results = shark_module.forward((input,))
assert True == compare_tensors(act_out, results)
class WideResnet50ModuleTest(unittest.TestCase):
class WideResnet50ModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
self.save_vmfb = pytestconfig.getoption("save_vmfb")
def setUp(self):
self.module_tester = WideResnet50ModuleTester(self)
def test_module_static_cpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
def test_module_dynamic_cpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
self.module_tester.dynamic = False
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(check_device_drivers("gpu"), reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
self.module_tester.dynamic = True
self.module_tester.device = "gpu"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
self.module_tester.dynamic = False
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
)
check_device_drivers("vulkan"),
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
self.module_tester.dynamic = True
self.module_tester.device = "vulkan"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

32
tank/tf/automodelmaskedlm.py
View File

@@ -9,11 +9,11 @@ inputs_signature = [
class AutoModelMaskedLM(tf.Module):
def __init__(self, model_name):
super(AutoModelMaskedLM, self).__init__()
self.m = TFAutoModelForMaskedLM.from_pretrained(model_name,
output_attentions=False)
self.m = TFAutoModelForMaskedLM.from_pretrained(
model_name, output_attentions=False
)
self.m.predict = lambda x: self.m(input_ids=x)
@tf.function(input_signature=inputs_signature)
@@ -24,20 +24,26 @@ class AutoModelMaskedLM(tf.Module):
fail_models = ["microsoft/deberta-base", "google/rembert", "google/tapas-base"]
supported_models = [
"albert-base-v2", "bert-base-uncased", "camembert-base",
"dbmdz/convbert-base-turkish-cased", "distilbert-base-uncased",
"albert-base-v2",
"bert-base-uncased",
"camembert-base",
"dbmdz/convbert-base-turkish-cased",
"distilbert-base-uncased",
"google/electra-small-discriminator",
"hf-internal-testing/tiny-random-flaubert", "funnel-transformer/small",
"microsoft/layoutlm-base-uncased", "allenai/longformer-base-4096",
"google/mobilebert-uncased", "microsoft/mpnet-base", "roberta-base",
"xlm-roberta-base"
"hf-internal-testing/tiny-random-flaubert",
"funnel-transformer/small",
"microsoft/layoutlm-base-uncased",
"allenai/longformer-base-4096",
"google/mobilebert-uncased",
"microsoft/mpnet-base",
"roberta-base",
"xlm-roberta-base",
]
if __name__ == "__main__":
inputs = tf.random.uniform(shape=[1, 512],
maxval=3,
dtype=tf.int32,
seed=10)
inputs = tf.random.uniform(
shape=[1, 512], maxval=3, dtype=tf.int32, seed=10
)
for model_name in supported_models:
print(f"Running model: {model_name}")

49
tank/tf/bert_large_gen.py
View File

@@ -19,24 +19,26 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs,
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -46,15 +48,20 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input0: input_word_ids
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input1: input_mask
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32) # input3: labels
])
@tf.function(
input_signature=[
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input0: input_word_ids
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input1: input_mask
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels
]
)
def learn(self, input_word_ids, input_mask, segment_ids, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -77,12 +84,12 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# BertModule()
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["learn"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["learn"], import_only=True
)
# Save module as MLIR file in a directory
ARITFACTS_DIR = os.getcwd()
mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir")
with open(mlir_path, "wt") as output_file:
output_file.write(compiler_module.decode('utf-8'))
output_file.write(compiler_module.decode("utf-8"))
print(f"Wrote MLIR to path '{mlir_path}'")

67
tank/tf/bert_large_run.py
View File

@@ -21,24 +21,26 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs,
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -49,10 +51,12 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels
])
@tf.function(
input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
]
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -69,26 +73,29 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# BertModule()
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["learn"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["learn"], import_only=True
)
# Compile the model using IREE
backend = "dylib-llvm-aot"
args = [
"--iree-llvm-target-cpu-features=host",
"--iree-mhlo-demote-i64-to-i32=false",
"--iree-stream-resource-index-bits=64", "--iree-vm-target-index-bits=64"
"--iree-stream-resource-index-bits=64",
"--iree-vm-target-index-bits=64",
]
backend_config = "dylib"
#backend = "cuda"
#backend_config = "cuda"
#args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo")
#flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# backend = "cuda"
# backend_config = "cuda"
# args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(
compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo",
)
# flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# Save module as MLIR file in a directory
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
@@ -100,21 +107,23 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
learn_sample_input = [
predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))
np.random.randint(5, size=(BATCH_SIZE)),
]
warmup = 5
total_iter = 10
num_iter = total_iter - warmup
for i in range(10):
if (i == warmup - 1):
if i == warmup - 1:
start = time.time()
print(
BertCompiled.learn(predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))))
BertCompiled.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
end = time.time()
total_time = end - start
print("time: " + str(total_time))

38
tank/tf/bert_large_tf.py
View File

@@ -14,24 +14,26 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size,
num_layers=24,
hidden_size=1024,
num_attention_heads=16,
dict_outputs=dict_outputs,
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -42,10 +44,12 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels
])
@tf.function(
input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
]
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -64,7 +68,7 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
bert_model = BertModule()
warmup = 1
@@ -72,9 +76,11 @@ if __name__ == "__main__":
num_iter = total_iter - warmup
for i in range(total_iter):
print(
bert_model.learn(predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))))
if (i == warmup - 1):
bert_model.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
if i == warmup - 1:
start = time.time()
end = time.time()

46
tank/tf/bert_small_gen.py
View File

@@ -1,5 +1,6 @@
from iree import runtime as ireert
#from iree.tf.support import module_utils
# from iree.tf.support import module_utils
from iree.compiler import tf as tfc
from absl import app
@@ -19,22 +20,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=2,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -44,15 +45,20 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input0: input_word_ids
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input1: input_mask
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH],
dtype=tf.int32), #input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32) # input3: labels
])
@tf.function(
input_signature=[
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input0: input_word_ids
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input1: input_mask
tf.TensorSpec(
shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32
), # input2: segment_ids
tf.TensorSpec([BATCH_SIZE], tf.int32), # input3: labels
]
)
def learn(self, input_word_ids, input_mask, segment_ids, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -75,13 +81,13 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# BertModule()
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["learn"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["learn"], import_only=True
)
print(type(compiler_module))
# Save module as MLIR file in a directory
ARITFACTS_DIR = os.getcwd()
mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir")
with open(mlir_path, "wt") as output_file:
output_file.write(compiler_module.decode('utf-8'))
output_file.write(compiler_module.decode("utf-8"))
print(f"Wrote MLIR to path '{mlir_path}'")

61
tank/tf/bert_small_run.py
View File

@@ -21,22 +21,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=2,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -47,10 +47,12 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels
])
@tf.function(
input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
]
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -67,25 +69,28 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# BertModule()
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["learn"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["learn"], import_only=True
)
# Compile the model using IREE
backend = "dylib-llvm-aot"
args = [
"--iree-llvm-target-cpu-features=host",
"--iree-mhlo-demote-i64-to-i32=false", "--iree-flow-demote-i64-to-i32"
"--iree-mhlo-demote-i64-to-i32=false",
"--iree-flow-demote-i64-to-i32",
]
backend_config = "dylib"
#backend = "cuda"
#backend_config = "cuda"
#args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo")
#flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# backend = "cuda"
# backend_config = "cuda"
# args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(
compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo",
)
# flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# Save module as MLIR file in a directory
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
@@ -97,21 +102,23 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
learn_sample_input = [
predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))
np.random.randint(5, size=(BATCH_SIZE)),
]
warmup = 5
total_iter = 10
num_iter = total_iter - warmup
for i in range(10):
if (i == warmup - 1):
if i == warmup - 1:
start = time.time()
print(
BertCompiled.learn(predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))))
BertCompiled.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
end = time.time()
total_time = end - start
print("time: " + str(total_time))

32
tank/tf/bert_small_tf_run.py
View File

@@ -14,22 +14,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
dict_outputs = False
test_network = networks.BertEncoder(vocab_size=vocab_size,
num_layers=2,
dict_outputs=dict_outputs)
test_network = networks.BertEncoder(
vocab_size=vocab_size, num_layers=2, dict_outputs=dict_outputs
)
# Create a BERT trainer with the created network.
bert_trainer_model = bert_classifier.BertClassifier(
test_network, num_classes=NUM_CLASSES)
test_network, num_classes=NUM_CLASSES
)
bert_trainer_model.summary()
# Invoke the trainer model on the inputs. This causes the layer to be built.
@@ -40,10 +40,12 @@ class BertModule(tf.Module):
self.loss = tf.keras.losses.SparseCategoricalCrossentropy()
self.optimizer = tf.keras.optimizers.SGD(learning_rate=1e-2)
@tf.function(input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32) # labels
])
@tf.function(
input_signature=[
bert_input, # inputs
tf.TensorSpec(shape=[BATCH_SIZE], dtype=tf.int32), # labels
]
)
def learn(self, inputs, labels):
with tf.GradientTape() as tape:
# Capture the gradients from forward prop...
@@ -62,7 +64,7 @@ if __name__ == "__main__":
predict_sample_input = [
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH))
np.random.randint(5, size=(BATCH_SIZE, SEQUENCE_LENGTH)),
]
bert_model = BertModule()
warmup = 1
@@ -70,9 +72,11 @@ if __name__ == "__main__":
num_iter = total_iter - warmup
for i in range(total_iter):
print(
bert_model.learn(predict_sample_input,
np.random.randint(5, size=(BATCH_SIZE))))
if (i == warmup - 1):
bert_model.learn(
predict_sample_input, np.random.randint(5, size=(BATCH_SIZE))
)
)
if i == warmup - 1:
start = time.time()
end = time.time()

87
tank/tf/hf_masked_lm/albert-base-v2_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class AlbertBaseModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -42,71 +41,80 @@ class AlbertBaseModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"albert-base-v2",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "albert-base-v2", dynamic, device
)
class AlbertBaseModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = AlbertBaseModuleTester(self)
self.module_tester.save_temps=pytestconfig.getoption("save_temps")
self.module_tester.save_mlir=pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb=pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark=pytestconfig.getoption("benchmark")
self.module_tester.save_temps = pytestconfig.getoption("save_temps")
self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="https://github.com/google/iree/issues/9553")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.xfail(reason="https://github.com/google/iree/issues/9553")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -115,8 +123,7 @@ class AlbertBaseModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -124,11 +131,11 @@ class AlbertBaseModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -136,5 +143,5 @@ class AlbertBaseModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

75
tank/tf/hf_masked_lm/bert-base-uncased_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class BertBaseUncasedModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -42,36 +41,40 @@ class BertBaseUncasedModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"bert_base_uncased",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "bert_base_uncased", dynamic, device
)
class BertBaseUncasedModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = BertBaseUncasedModuleTester(self)
@@ -80,8 +83,9 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
@@ -89,24 +93,28 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -115,8 +123,7 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -124,11 +131,11 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -136,5 +143,5 @@ class BertBaseUncasedModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

81
tank/tf/hf_masked_lm/camembert-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class CamemBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -42,44 +41,51 @@ class CamemBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"camembert-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "camembert-base", dynamic, device
)
class CamemBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester=CamemBertModuleTester(self)
self.module_tester = CamemBertModuleTester(self)
self.module_tester.save_temps = pytestconfig.getoption("save_temps")
self.module_tester.save_vmfb = pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
@@ -87,24 +93,28 @@ class CamemBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -113,8 +123,7 @@ class CamemBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -122,11 +131,11 @@ class CamemBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -134,5 +143,5 @@ class CamemBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

80
tank/tf/hf_masked_lm/convbert-base-turkish-cased_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -11,24 +11,23 @@ import numpy as np
import tempfile
class ConvBertModuleTester:
class ConvBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
self.benchmark = benchmark
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model(
"dbmdz/convbert-base-turkish-cased")
"dbmdz/convbert-base-turkish-cased"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if self.save_temps == True:
@@ -44,36 +43,40 @@ class ConvBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"convbert-base-turkish-cased",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "convbert-base-turkish-cased", dynamic, device
)
class ConvBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = ConvBertModuleTester(self)
@@ -82,7 +85,9 @@ class ConvBertModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536.")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536."
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
@@ -90,24 +95,28 @@ class ConvBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -116,8 +125,7 @@ class ConvBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -125,11 +133,11 @@ class ConvBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -137,5 +145,5 @@ class ConvBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

99
tank/tf/hf_masked_lm/deberta-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class DebertaModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("microsoft/deberta-base")
shark_args.save_mlir = self.save_mlir
@@ -41,36 +40,40 @@ class DebertaModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"deberta-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "deberta-base", dynamic, device
)
class DebertaModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = DebertaModuleTester(self)
@@ -78,10 +81,11 @@ class DebertaModuleTest(unittest.TestCase):
self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail
@pytest.mark.skip(reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="deberta currently failing in the lowering passes."
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
@@ -89,52 +93,59 @@ class DebertaModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skip(
reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="deberta currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="deberta currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="deberta currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -142,5 +153,5 @@ class DebertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

75
tank/tf/hf_masked_lm/distilbert-base-uncased_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class DistilBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -41,35 +40,40 @@ class DistilBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"distilbert-base-uncased",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "distilbert-base-uncased", dynamic, device
)
class DistilBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = DistilBertModuleTester(self)
@@ -78,7 +82,9 @@ class DistilBertModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
@@ -86,24 +92,28 @@ class DistilBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -112,8 +122,7 @@ class DistilBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -121,11 +130,11 @@ class DistilBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -133,5 +142,5 @@ class DistilBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

86
tank/tf/hf_masked_lm/electra-small-discriminator_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,20 +12,21 @@ import tempfile
class ElectraModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("google/electra-small-discriminator")
model, input, act_out = get_causal_lm_model(
"google/electra-small-discriminator"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if self.save_temps == True:
@@ -39,71 +40,81 @@ class ElectraModuleTester:
np.save(f"{temp_dir}/input2.npy", input[1])
exp_out = act_out.numpy()
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"electra-small-discriminator",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "electra-small-discriminator", dynamic, device
)
class ElectraModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = ElectraModuleTester(self)
self.module_tester.save_temps = pytestconfig.getoption("save_temps")
self.module_tester.save_temps = pytestconfig.getoption("save_temps")
self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -112,8 +123,7 @@ class ElectraModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -121,11 +131,11 @@ class ElectraModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -133,5 +143,5 @@ class ElectraModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

71
tank/tf/hf_masked_lm/funnel-transformer_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class FunnelModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -41,35 +40,40 @@ class FunnelModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"funnel-transformer-small",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "funnel-transformer-small", dynamic, device
)
class FunnelModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = FunnelModuleTester(self)
@@ -85,15 +89,17 @@ class FunnelModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="funnel currently failing in the lowering passes.")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
@@ -101,9 +107,11 @@ class FunnelModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="funnel currently failing in the lowering passes.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -112,8 +120,7 @@ class FunnelModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="funnel currently failing in the lowering passes.")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -122,11 +129,11 @@ class FunnelModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="funnel currently failing in the lowering passes.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -134,5 +141,5 @@ class FunnelModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

82
tank/tf/hf_masked_lm/layoutlm-base-uncased_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,21 +12,21 @@ import tempfile
class LayoutLmModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model(
"microsoft/layoutlm-base-uncased")
"microsoft/layoutlm-base-uncased"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if self.save_temps == True:
@@ -42,35 +42,40 @@ class LayoutLmModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"layoutlm-base-uncased",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "layoutlm-base-uncased", dynamic, device
)
class LayoutLmModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = LayoutLmModuleTester(self)
@@ -78,32 +83,38 @@ class LayoutLmModuleTest(unittest.TestCase):
self.module_tester.save_mlir = pytestconfig.getoption("save_mlir")
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -112,8 +123,7 @@ class LayoutLmModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -121,11 +131,11 @@ class LayoutLmModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -133,5 +143,5 @@ class LayoutLmModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

92
tank/tf/hf_masked_lm/longformer-base-4096_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,20 +12,21 @@ import tempfile
class LongFormerModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("allenai/longformer-base-4096")
model, input, act_out = get_causal_lm_model(
"allenai/longformer-base-4096"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if self.save_temps == True:
@@ -41,35 +42,40 @@ class LongFormerModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"longformer-base-4096",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "longformer-base-4096", dynamic, device
)
class LongFormerModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = LongFormerModuleTester(self)
@@ -79,45 +85,52 @@ class LongFormerModuleTest(unittest.TestCase):
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.skip(
reason="longformer currently failing in the lowering passes.")
reason="longformer currently failing in the lowering passes."
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="longformer currently failing in the lowering passes.")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="longformer currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="longformer currently failing in the lowering passes.")
reason="longformer currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="longformer currently failing in the lowering passes.")
reason="longformer currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -125,13 +138,14 @@ class LongFormerModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="longformer currently failing in the lowering passes.")
reason="longformer currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -139,5 +153,5 @@ class LongFormerModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

19
tank/tf/hf_masked_lm/masked_lm.py
View File

@@ -17,21 +17,22 @@ inputs_signature = [
def preprocess_input(model_name, text="This is just used to compile the model"):
tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer(text,
padding="max_length",
return_tensors="tf",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
inputs = tokenizer(
text,
padding="max_length",
return_tensors="tf",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
return inputs
class MaskedLM(tf.Module):
def __init__(self, model_name):
super(MaskedLM, self).__init__()
self.m = TFAutoModelForMaskedLM.from_pretrained(model_name,
output_attentions=False,
num_labels=2)
self.m = TFAutoModelForMaskedLM.from_pretrained(
model_name, output_attentions=False, num_labels=2
)
self.m.predict = lambda x, y: self.m(input_ids=x, attention_mask=y)[0]
@tf.function(input_signature=inputs_signature)

75
tank/tf/hf_masked_lm/mobilebert-uncased_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,13 +12,12 @@ import tempfile
class MobileBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
@@ -41,35 +40,40 @@ class MobileBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"mobilebert-uncased",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "mobilebert-uncased", dynamic, device
)
class MobileBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = MobileBertModuleTester(self)
@@ -78,31 +82,37 @@ class MobileBertModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -111,8 +121,7 @@ class MobileBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -120,11 +129,11 @@ class MobileBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -132,5 +141,5 @@ class MobileBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

77
tank/tf/hf_masked_lm/mpnet-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class MpNetModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("microsoft/mpnet-base")
shark_args.save_mlir = self.save_mlir
@@ -41,35 +40,40 @@ class MpNetModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"mpnet-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "mpnet-base", dynamic, device
)
class MpNetModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = MpNetModuleTester(self)
@@ -78,31 +82,37 @@ class MpNetModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -111,8 +121,7 @@ class MpNetModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -120,11 +129,11 @@ class MpNetModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -132,5 +141,5 @@ class MpNetModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

96
tank/tf/hf_masked_lm/rembert_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class RemBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("google/rembert")
shark_args.save_mlir = self.save_mlir
@@ -41,35 +40,38 @@ class RemBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"rembert",
dynamic,
device)
shark_module.benchmark_all_csv((input), "rembert", dynamic, device)
class RemBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = RemBertModuleTester(self)
@@ -78,60 +80,68 @@ class RemBertModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.skip(reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="rembert currently failing in the lowering passes."
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skip(
reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="rembert currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
device = "vulkan"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="rembert currently failing in the lowering passes."
)
@pytest.mark.skip(
reason="rembert currently failing in the lowering passes."
)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -139,5 +149,5 @@ class RemBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

77
tank/tf/hf_masked_lm/roberta-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class RobertaModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("roberta-base")
shark_args.save_mlir = self.save_mlir
@@ -41,35 +40,40 @@ class RobertaModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"roberta-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "roberta-base", dynamic, device
)
class RobertaModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = RobertaModuleTester(self)
@@ -78,31 +82,37 @@ class RobertaModuleTest(unittest.TestCase):
self.module_tester.save_vmfb = pytestconfig.getoption("save_vmfb")
self.module_tester.benchmark = pytestconfig.getoption("benchmark")
@pytest.mark.xfail(reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536")
@pytest.mark.xfail(
reason="Upstream IREE issue, see https://github.com/google/iree/issues/9536"
)
def test_module_static_cpu(self):
dynamic = False
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -111,8 +121,7 @@ class RobertaModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -120,11 +129,11 @@ class RobertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -132,5 +141,5 @@ class RobertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

73
tank/tf/hf_masked_lm/tapas-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class TapasBaseModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("google/tapas-base")
shark_args.save_mlir = self.save_mlir
@@ -41,35 +40,40 @@ class TapasBaseModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"tapas-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "tapas-base", dynamic, device
)
class TapasBaseModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = TapasBaseModuleTester(self)
@@ -85,15 +89,17 @@ class TapasBaseModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="tapas currently failing in the lowering passes.")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
@@ -101,9 +107,11 @@ class TapasBaseModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="tapas currently failing in the lowering passes.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -112,8 +120,7 @@ class TapasBaseModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="tapas currently failing in the lowering passes.")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -122,11 +129,11 @@ class TapasBaseModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="tapas currently failing in the lowering passes.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -134,5 +141,5 @@ class TapasBaseModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

77
tank/tf/hf_masked_lm/tiny-random-flaubert_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,20 +12,21 @@ import tempfile
class FlauBertModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("hf-internal-testing/tiny-random-flaubert")
model, input, act_out = get_causal_lm_model(
"hf-internal-testing/tiny-random-flaubert"
)
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
if self.save_temps == True:
@@ -41,35 +42,40 @@ class FlauBertModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"tiny-random-flaubert",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "tiny-random-flaubert", dynamic, device
)
class FlauBertModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = FlauBertModuleTester(self)
@@ -84,24 +90,28 @@ class FlauBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -110,8 +120,7 @@ class FlauBertModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -119,11 +128,11 @@ class FlauBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -131,5 +140,5 @@ class FlauBertModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

73
tank/tf/hf_masked_lm/xlm-roberta-base_test.py
View File

@@ -1,6 +1,6 @@
from masked_lm import get_causal_lm_model
from tank.model_utils_tf import compare_tensors_tf
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from shark.shark_inference import SharkInference
from shark.parser import shark_args
@@ -12,18 +12,17 @@ import tempfile
class XLMRobertaModuleTester:
def __init__(
self,
save_temps=False,
save_mlir=False,
save_vmfb=False,
benchmark=False
benchmark=False,
):
self.save_temps = save_temps
self.save_mlir = save_mlir
self.save_vmfb = save_vmfb
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_causal_lm_model("xlm-roberta-base")
shark_args.save_mlir = self.save_mlir
@@ -41,35 +40,40 @@ class XLMRobertaModuleTester:
with open(f"{temp_dir}/expected_out.txt", "w") as out_file:
out_file.write(np.array2string(exp_out))
with ireec.tools.TempFileSaver(temp_dir):
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
else:
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark)
else:
shark_module = SharkInference(
model,
(input,),
device=device,
dynamic=dynamic,
jit_trace=True,
benchmark_mode=self.benchmark,
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
assert True == compare_tensors_tf(act_out, results)
if self.benchmark == True:
shark_module.benchmark_all_csv((input),
"xlm-roberta-base",
dynamic,
device)
shark_module.benchmark_all_csv(
(input), "xlm-roberta-base", dynamic, device
)
class XLMRobertaModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.module_tester = XLMRobertaModuleTester(self)
@@ -85,24 +89,28 @@ class XLMRobertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -111,8 +119,7 @@ class XLMRobertaModuleTest(unittest.TestCase):
@pytest.mark.xfail
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -120,11 +127,11 @@ class XLMRobertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -132,5 +139,5 @@ class XLMRobertaModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

17
tank/tf/huggingface_MiniLM_gen.py
View File

@@ -15,21 +15,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
# Create a BERT trainer with the created network.
self.m = TFBertModel.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", from_pt=True)
"microsoft/MiniLM-L12-H384-uncased", from_pt=True
)
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=bert_input)
def predict(self, input_word_ids, input_mask, segment_ids):
@@ -39,12 +40,12 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# BertModule()
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["predict"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["predict"], import_only=True
)
# Save module as MLIR file in a directory
ARITFACTS_DIR = os.getcwd()
mlir_path = os.path.join(ARITFACTS_DIR, "model.mlir")
with open(mlir_path, "wt") as output_file:
output_file.write(compiler_module.decode('utf-8'))
output_file.write(compiler_module.decode("utf-8"))
print(f"Wrote MLIR to path '{mlir_path}'")

60
tank/tf/huggingface_MiniLM_run.py
View File

@@ -16,21 +16,22 @@ BATCH_SIZE = 1
bert_input = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32)
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
class BertModule(tf.Module):
def __init__(self):
super(BertModule, self).__init__()
# Create a BERT trainer with the created network.
self.m = TFBertModel.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", from_pt=True)
"microsoft/MiniLM-L12-H384-uncased", from_pt=True
)
# Invoke the trainer model on the inputs. This causes the layer to be built.
self.m.predict = lambda x, y, z: self.m.call(
input_ids=x, attention_mask=y, token_type_ids=z, training=False)
input_ids=x, attention_mask=y, token_type_ids=z, training=False
)
@tf.function(input_signature=bert_input)
def predict(self, input_ids, attention_mask, token_type_ids):
@@ -40,36 +41,43 @@ class BertModule(tf.Module):
if __name__ == "__main__":
# Prepping Data
tokenizer = BertTokenizer.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased")
"microsoft/MiniLM-L12-H384-uncased"
)
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
encoded_input = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
tf.convert_to_tensor(encoded_input[key]), 0
)
# Compile the model using IREE
compiler_module = tfc.compile_module(BertModule(),
exported_names=["predict"],
import_only=True)
compiler_module = tfc.compile_module(
BertModule(), exported_names=["predict"], import_only=True
)
# Compile the model using IREE
backend = "dylib-llvm-aot"
args = [
"--iree-llvm-target-cpu-features=host",
"--iree-mhlo-demote-i64-to-i32=false", "--iree-flow-demote-i64-to-i32"
"--iree-mhlo-demote-i64-to-i32=false",
"--iree-flow-demote-i64-to-i32",
]
backend_config = "dylib"
#backend = "cuda"
#backend_config = "cuda"
#args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo")
#flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# backend = "cuda"
# backend_config = "cuda"
# args = ["--iree-cuda-llvm-target-arch=sm_80", "--iree-hal-cuda-disable-loop-nounroll-wa", "--iree-enable-fusion-with-reduction-ops"]
flatbuffer_blob = compile_str(
compiler_module,
target_backends=[backend],
extra_args=args,
input_type="mhlo",
)
# flatbuffer_blob = compile_str(compiler_module, target_backends=["dylib-llvm-aot"])
# Save module as MLIR file in a directory
vm_module = ireert.VmModule.from_flatbuffer(flatbuffer_blob)
@@ -78,7 +86,9 @@ if __name__ == "__main__":
ctx = ireert.SystemContext(config=config)
ctx.add_vm_module(vm_module)
BertCompiled = ctx.modules.module
result = BertCompiled.predict(encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"])
result = BertCompiled.predict(
encoded_input["input_ids"],
encoded_input["attention_mask"],
encoded_input["token_type_ids"],
)
print(result)

15
tank/tf/huggingface_MiniLM_tf.py
View File

@@ -1,18 +1,19 @@
import tensorflow as tf
from transformers import BertModel, BertTokenizer, TFBertModel
tf_model = TFBertModel.from_pretrained("microsoft/MiniLM-L12-H384-uncased",
from_pt=True)
tf_model = TFBertModel.from_pretrained(
"microsoft/MiniLM-L12-H384-uncased", from_pt=True
)
tokenizer = BertTokenizer.from_pretrained("microsoft/MiniLM-L12-H384-uncased")
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text,
padding='max_length',
truncation=True,
max_length=512)
encoded_input = tokenizer(
text, padding="max_length", truncation=True, max_length=512
)
for key in encoded_input:
encoded_input[key] = tf.expand_dims(
tf.convert_to_tensor(encoded_input[key]), 0)
tf.convert_to_tensor(encoded_input[key]), 0
)
output = tf_model(encoded_input)
print(output)

41
tank/tf/minilm_tf_test.py
View File

@@ -1,5 +1,5 @@
from shark.shark_inference import SharkInference
from shark.iree_utils import check_device_drivers
from shark.iree_utils._common import check_device_drivers
from tank.model_utils_tf import get_TFhf_model, compare_tensors_tf
import tensorflow as tf
@@ -9,14 +9,13 @@ import pytest
class MiniLMTFModuleTester:
def create_and_check_module(self, dynamic, device):
model, input, act_out = get_TFhf_model(
"microsoft/MiniLM-L12-H384-uncased")
shark_module = SharkInference(model, (input,),
device=device,
dynamic=dynamic,
jit_trace=True)
"microsoft/MiniLM-L12-H384-uncased"
)
shark_module = SharkInference(
model, (input,), device=device, dynamic=dynamic, jit_trace=True
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
results = shark_module.forward((input))
@@ -24,7 +23,6 @@ class MiniLMTFModuleTester:
class MiniLMTFModuleTest(unittest.TestCase):
def setUp(self):
self.module_tester = MiniLMTFModuleTester()
@@ -36,15 +34,17 @@ class MiniLMTFModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="TF testing temporarily unavailable.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
def test_module_dynamic_cpu(self):
dynamic = True
device = "cpu"
self.module_tester.create_and_check_module(dynamic, device)
@pytest.mark.skip(reason="TF testing temporarily unavailable.")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_static_gpu(self):
dynamic = False
device = "gpu"
@@ -52,9 +52,11 @@ class MiniLMTFModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="TF testing temporarily unavailable.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
@pytest.mark.skipif(check_device_drivers("gpu"),
reason="nvidia-smi not found")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("gpu"), reason="nvidia-smi not found"
)
def test_module_dynamic_gpu(self):
dynamic = True
device = "gpu"
@@ -63,8 +65,7 @@ class MiniLMTFModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="TF testing temporarily unavailable.")
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_static_vulkan(self):
dynamic = False
@@ -73,11 +74,11 @@ class MiniLMTFModuleTest(unittest.TestCase):
@pytest.mark.skip(reason="TF testing temporarily unavailable.")
@pytest.mark.xfail(
reason="Language models currently failing for dynamic case")
reason="Language models currently failing for dynamic case"
)
@pytest.mark.skipif(
check_device_drivers("vulkan"),
reason=
"vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases"
reason="vulkaninfo not found, install from https://github.com/KhronosGroup/MoltenVK/releases",
)
def test_module_dynamic_vulkan(self):
dynamic = True
@@ -85,5 +86,5 @@ class MiniLMTFModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module(dynamic, device)
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()

46
tank/tf/seq_classification.py
View File

@@ -6,12 +6,15 @@ from shark.parser import shark_args
import argparse
seq_parser = argparse.ArgumentParser(description='Shark Sequence Classification.')
seq_parser = argparse.ArgumentParser(
description="Shark Sequence Classification."
)
seq_parser.add_argument(
"--hf_model_name",
type=str,
default="bert-base-uncased",
help="Hugging face model to run sequence classification.")
help="Hugging face model to run sequence classification.",
)
seq_args, unknown = seq_parser.parse_known_args()
@@ -25,45 +28,56 @@ inputs_signature = [
tf.TensorSpec(shape=[BATCH_SIZE, MAX_SEQUENCE_LENGTH], dtype=tf.int32),
]
# For supported models please see here:
# For supported models please see here:
# https://huggingface.co/docs/transformers/model_doc/auto#transformers.TFAutoModelForSequenceClassification
def preprocess_input(text = "This is just used to compile the model"):
def preprocess_input(text="This is just used to compile the model"):
tokenizer = AutoTokenizer.from_pretrained(seq_args.hf_model_name)
inputs = tokenizer(text,
padding="max_length",
return_tensors="tf",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH)
inputs = tokenizer(
text,
padding="max_length",
return_tensors="tf",
truncation=True,
max_length=MAX_SEQUENCE_LENGTH,
)
return inputs
class SeqClassification(tf.Module):
def __init__(self, model_name):
super(SeqClassification, self).__init__()
self.m = TFAutoModelForSequenceClassification.from_pretrained(
model_name, output_attentions=False, num_labels=2)
model_name, output_attentions=False, num_labels=2
)
self.m.predict = lambda x, y: self.m(input_ids=x, attention_mask=y)[0]
@tf.function(input_signature=inputs_signature)
def forward(self, input_ids, attention_mask):
return tf.math.softmax(self.m.predict(input_ids, attention_mask),
axis=-1)
return tf.math.softmax(
self.m.predict(input_ids, attention_mask), axis=-1
)
if __name__ == "__main__":
inputs = preprocess_input()
shark_module = SharkInference(
SeqClassification(seq_args.hf_model_name),
(inputs["input_ids"], inputs["attention_mask"]))
(inputs["input_ids"], inputs["attention_mask"]),
)
shark_module.set_frontend("tensorflow")
shark_module.compile()
print(f"Model has been successfully compiled on {shark_args.device}")
while True:
input_text = input("Enter the text to classify (press q or nothing to exit): ")
input_text = input(
"Enter the text to classify (press q or nothing to exit): "
)
if not input_text or input_text == "q":
break
inputs = preprocess_input(input_text)
print(shark_module.forward((inputs["input_ids"], inputs["attention_mask"])))
print(
shark_module.forward(
(inputs["input_ids"], inputs["attention_mask"])
)
)

39
tank/tflite/albert.py
View File

@@ -13,26 +13,35 @@ def generate_inputs(input_details):
args = []
args.append(
np.random.randint(low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"]))
np.random.randint(
low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"],
)
)
args.append(
np.ones(shape=input_details[1]["shape"],
dtype=input_details[1]["dtype"]))
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(shape=input_details[2]["shape"],
dtype=input_details[2]["dtype"]))
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
if __name__ == '__main__':
my_shark_importer = SharkImporter(model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device="cpu",
dynamic=False,
jit_trace=True)
if __name__ == "__main__":
my_shark_importer = SharkImporter(
model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device="cpu",
dynamic=False,
jit_trace=True,
)
# Case1: Use default inputs
my_shark_importer.compile()
shark_results = my_shark_importer.forward()

40
tank/tflite/albert_lite_base/albert_lite_base_tflite_mlir_test.py
View File

@@ -6,13 +6,12 @@ from shark.parser import shark_args
class AlbertTfliteModuleTester:
def __init__(
self,
dynamic=False,
device="cpu",
save_mlir=False,
save_vmfb=False,
self,
dynamic=False,
device="cpu",
save_mlir=False,
save_vmfb=False,
):
self.dynamic = dynamic
self.device = device
@@ -22,19 +21,23 @@ class AlbertTfliteModuleTester:
def create_and_check_module(self):
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
self.shark_downloader = SharkDownloader(model_name="albert_lite_base",
tank_url="https://storage.googleapis.com/shark_tank",
local_tank_dir="./../gen_shark_tank/tflite",
model_type="tflite-tosa",
input_json="input.json",
input_type="int32")
self.shark_downloader = SharkDownloader(
model_name="albert_lite_base",
tank_url="https://storage.googleapis.com/shark_tank",
local_tank_dir="./../gen_shark_tank/tflite",
model_type="tflite-tosa",
input_json="input.json",
input_type="int32",
)
tflite_tosa_model = self.shark_downloader.get_mlir_file()
inputs = self.shark_downloader.get_inputs()
self.shark_module = SharkInference(tflite_tosa_model,
inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=True)
self.shark_module = SharkInference(
tflite_tosa_model,
inputs,
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
)
self.shark_module.set_frontend("tflite-tosa")
self.shark_module.compile()
self.shark_module.forward(inputs)
@@ -42,7 +45,6 @@ class AlbertTfliteModuleTester:
class AlbertTfliteModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
@@ -58,7 +60,7 @@ class AlbertTfliteModuleTest(unittest.TestCase):
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
unittest.main()
# module_tester = AlbertTfliteModuleTester()
# module_tester.create_and_check_module()

47
tank/tflite/albert_lite_base/albert_lite_base_tflite_test.py
View File

@@ -14,20 +14,27 @@ def generate_inputs(input_details):
args = []
args.append(
np.random.randint(low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"]))
np.random.randint(
low=0,
high=256,
size=input_details[0]["shape"],
dtype=input_details[0]["dtype"],
)
)
args.append(
np.ones(shape=input_details[1]["shape"],
dtype=input_details[1]["dtype"]))
np.ones(
shape=input_details[1]["shape"], dtype=input_details[1]["dtype"]
)
)
args.append(
np.zeros(shape=input_details[2]["shape"],
dtype=input_details[2]["dtype"]))
np.zeros(
shape=input_details[2]["shape"], dtype=input_details[2]["dtype"]
)
)
return args
class AlbertTfliteModuleTester:
class AlbertTfliteModuleTester:
def __init__(
self,
dynamic=False,
@@ -43,14 +50,16 @@ class AlbertTfliteModuleTester:
def create_and_check_module(self):
shark_args.save_mlir = self.save_mlir
shark_args.save_vmfb = self.save_vmfb
my_shark_importer = SharkImporter(model_name="albert_lite_base",
# model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
tank_url=None)
my_shark_importer = SharkImporter(
model_name="albert_lite_base",
# model_path=model_path,
model_type="tflite",
model_source_hub="tfhub",
device=self.device,
dynamic=self.dynamic,
jit_trace=True,
tank_url=None,
)
# Case1: Use default inputs
my_shark_importer.compile()
shark_results = my_shark_importer.forward()
@@ -63,7 +72,6 @@ class AlbertTfliteModuleTester:
class AlbertTfliteModuleTest(unittest.TestCase):
@pytest.fixture(autouse=True)
def configure(self, pytestconfig):
self.save_mlir = pytestconfig.getoption("save_mlir")
@@ -78,7 +86,8 @@ class AlbertTfliteModuleTest(unittest.TestCase):
self.module_tester.device = "cpu"
self.module_tester.create_and_check_module()
if __name__ == '__main__':
if __name__ == "__main__":
# module_tester = AlbertTfliteModuleTester()
# module_tester.create_and_check_module()
unittest.main()

3
tank/tflite/asr_conformer_test.py
View File

@@ -10,7 +10,6 @@ model_path = "https://tfhub.dev/neso613/lite-model/ASR_TFLite/pre_trained_models
# Failure is due to dynamic shapes:
# - Some improvements to tfl.strided_slice lowering are next steps
class AsrConformerTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(AsrConformerTest, self).__init__(model_path, *args, **kwargs)
@@ -18,5 +17,5 @@ class AsrConformerTest(test_util.TFLiteModelTest):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

15
tank/tflite/bird_classifier_test.py
View File

@@ -11,18 +11,21 @@ model_path = "https://tfhub.dev/google/lite-model/aiy/vision/classifier/birds_V1
class BirdClassifierTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(BirdClassifierTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(BirdClassifierTest, self).compare_results(iree_results,
tflite_results, details)
super(BirdClassifierTest, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all())
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()
)
def generate_inputs(self, input_details):
img_path = "https://github.com/google-coral/test_data/raw/master/bird.bmp"
img_path = (
"https://github.com/google-coral/test_data/raw/master/bird.bmp"
)
local_path = "/".join([self.workdir, "bird.bmp"])
urllib.request.urlretrieve(img_path, local_path)
@@ -35,5 +38,5 @@ class BirdClassifierTest(test_util.TFLiteModelTest):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

7
tank/tflite/cartoon_gan_test.py
View File

@@ -4,11 +4,12 @@
import absl.testing
import test_util
model_path = "https://tfhub.dev/sayakpaul/lite-model/cartoongan/dr/1?lite-format=tflite"
model_path = (
"https://tfhub.dev/sayakpaul/lite-model/cartoongan/dr/1?lite-format=tflite"
)
class CartoonGanTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(CartoonGanTest, self).__init__(model_path, *args, **kwargs)
@@ -16,5 +17,5 @@ class CartoonGanTest(test_util.TFLiteModelTest):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

8
tank/tflite/craft_text_test.py
View File

@@ -10,17 +10,17 @@ model_path = "https://tfhub.dev/tulasiram58827/lite-model/craft-text-detector/dr
# Failure: Resize lowering does not handle inferred dynamic shapes. Furthermore, the entire model
# requires dynamic shape support.
class CraftTextTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(CraftTextTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(CraftTextTest, self).compare_results(iree_results, tflite_results,
details)
super(CraftTextTest, self).compare_results(
iree_results, tflite_results, details
)
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

11
tank/tflite/deeplab_v3_test.py
View File

@@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2?lit
class DeepLabV3Test(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(DeepLabV3Test, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(DeepLabV3Test, self).compare_results(iree_results, tflite_results,
details)
super(DeepLabV3Test, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all())
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()
)
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

11
tank/tflite/densenet_test.py
View File

@@ -8,19 +8,20 @@ model_path = "https://tfhub.dev/tensorflow/lite-model/densenet/1/metadata/1?lite
class DenseNetTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(DenseNetTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(DenseNetTest, self).compare_results(iree_results, tflite_results,
details)
super(DenseNetTest, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-5).all())
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-5).all()
)
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

16
tank/tflite/east_text_detector_test.py
View File

@@ -8,28 +8,30 @@ model_path = "https://tfhub.dev/sayakpaul/lite-model/east-text-detector/dr/1?lit
class EastTextDetectorTest(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(EastTextDetectorTest, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(EastTextDetectorTest,
self).compare_results(iree_results, tflite_results, details)
super(EastTextDetectorTest, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all())
numpy.isclose(iree_results[0], tflite_results[0], atol=1e-3).all()
)
# The second return is extremely noisy as it is not a binary classification. To handle we
# check normalized correlation with an expectation of "close enough".
iree_norm = numpy.sqrt(iree_results[1] * iree_results[1])
tflite_norm = numpy.sqrt(tflite_results[1] * tflite_results[1])
correlation = numpy.average(iree_results[1] * tflite_results[1] /
iree_norm / tflite_norm)
correlation = numpy.average(
iree_results[1] * tflite_results[1] / iree_norm / tflite_norm
)
self.assertTrue(numpy.isclose(correlation, 1.0, atol=1e-2).all())
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

24
tank/tflite/efficientnet_lite0_int8_test.py
View File

@@ -10,23 +10,25 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/efficientnet_l
class EfficientnetLite0Int8Test(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(EfficientnetLite0Int8Test, self).__init__(model_path, *args,
**kwargs)
super(EfficientnetLite0Int8Test, self).__init__(
model_path, *args, **kwargs
)
def compare_results(self, iree_results, tflite_results, details):
super(EfficientnetLite0Int8Test,
self).compare_results(iree_results, tflite_results, details)
super(EfficientnetLite0Int8Test, self).compare_results(
iree_results, tflite_results, details
)
# Dequantize outputs.
zero_point = details[0]['quantization_parameters']['zero_points'][0]
scale = details[0]['quantization_parameters']['scales'][0]
zero_point = details[0]["quantization_parameters"]["zero_points"][0]
scale = details[0]["quantization_parameters"]["scales"][0]
dequantized_iree_results = (iree_results - zero_point) * scale
dequantized_tflite_results = (tflite_results - zero_point) * scale
self.assertTrue(
numpy.isclose(dequantized_iree_results,
dequantized_tflite_results,
atol=5e-3).all())
numpy.isclose(
dequantized_iree_results, dequantized_tflite_results, atol=5e-3
).all()
)
def generate_inputs(self, input_details):
return [imagenet_test_data.generate_input(self.workdir, input_details)]
@@ -35,5 +37,5 @@ class EfficientnetLite0Int8Test(test_util.TFLiteModelTest):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

13
tank/tflite/efficientnet_lite0_test.py
View File

@@ -10,25 +10,26 @@ model_path = "https://storage.googleapis.com/iree-model-artifacts/efficientnet_l
class EfficientnetLite0Test(test_util.TFLiteModelTest):
def __init__(self, *args, **kwargs):
super(EfficientnetLite0Test, self).__init__(model_path, *args, **kwargs)
def compare_results(self, iree_results, tflite_results, details):
super(EfficientnetLite0Test,
self).compare_results(iree_results, tflite_results, details)
super(EfficientnetLite0Test, self).compare_results(
iree_results, tflite_results, details
)
self.assertTrue(
numpy.isclose(iree_results, tflite_results, atol=1e-4).all())
numpy.isclose(iree_results, tflite_results, atol=1e-4).all()
)
def generate_inputs(self, input_details):
inputs = imagenet_test_data.generate_input(self.workdir, input_details)
# Normalize inputs to [-1, 1].
inputs = (inputs.astype('float32') / 127.5) - 1
inputs = (inputs.astype("float32") / 127.5) - 1
return [inputs]
def test_compile_tflite(self):
self.compile_and_execute()
if __name__ == '__main__':
if __name__ == "__main__":
absl.testing.absltest.main()

Some files were not shown because too many files have changed in this diff Show More