mirror of https://github.com/llvm/torch-mlir
145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
# Based on code Copyright (c) Advanced Micro Devices, Inc.
|
|
#
|
|
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
# See https://llvm.org/LICENSE.txt for license information.
|
|
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
# Also available under a BSD-style license. See LICENSE.
|
|
|
|
# RUN: %PYTHON %s --output %t
|
|
|
|
from pathlib import Path
|
|
|
|
import logging
|
|
import shutil
|
|
import sys
|
|
import subprocess
|
|
import unittest
|
|
import unittest.mock
|
|
|
|
import onnx
|
|
|
|
from torch_mlir.tools.import_onnx import __main__
|
|
|
|
# For ONNX models
|
|
|
|
import numpy
|
|
from onnx import numpy_helper, TensorProto
|
|
from onnx.helper import (
|
|
make_model, make_node, make_graph,
|
|
make_tensor_value_info)
|
|
from onnx.external_data_helper import convert_model_to_external_data
|
|
from onnx.checker import check_model
|
|
|
|
# Accept the output path on the command line or default to a sibling
|
|
# to this file. We have to pop this off explicitly or else unittest
|
|
# won't understand.
|
|
if len(sys.argv) > 1 and sys.argv[1] == "--output":
|
|
OUTPUT_PATH = Path(sys.argv[2])
|
|
del sys.argv[1:3]
|
|
else:
|
|
OUTPUT_PATH = Path(__file__).resolve().parent / "output"
|
|
|
|
OUTPUT_PATH.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def const_model() -> onnx.ModelProto:
|
|
# Note: data_path must be relative to model_file
|
|
|
|
const = make_node(
|
|
'Constant', [], ['c_shape'], 'const',
|
|
value=numpy_helper.from_array(numpy.array([4], dtype=numpy.int64)))
|
|
cofshape = make_node(
|
|
'ConstantOfShape', ['c_shape'], ['c_out'], 'cofshape',
|
|
value=numpy_helper.from_array(numpy.array([1], dtype=numpy.int64)))
|
|
|
|
outval = make_tensor_value_info('c_out', TensorProto.INT64, [None])
|
|
graph = make_graph([const, cofshape], 'constgraph', [], [outval])
|
|
|
|
onnx_model = make_model(graph)
|
|
check_model(onnx_model)
|
|
return onnx_model
|
|
|
|
|
|
def linear_model() -> onnx.ModelProto:
|
|
# initializers
|
|
k_dim = 32
|
|
value = numpy.arange(k_dim).reshape([k_dim, 1])
|
|
value = numpy.asarray(value, dtype=numpy.float32)
|
|
A = numpy_helper.from_array(value, name='A')
|
|
|
|
value = numpy.array([0.4], dtype=numpy.float32).reshape([1, 1])
|
|
C = numpy_helper.from_array(value, name='C')
|
|
|
|
# the part which does not change
|
|
X = make_tensor_value_info('X', TensorProto.FLOAT, [1, k_dim])
|
|
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None, None])
|
|
node1 = make_node('MatMul', ['X', 'A'], ['AX'])
|
|
node2 = make_node('Add', ['AX', 'C'], ['Y'])
|
|
graph = make_graph([node1, node2], 'lr', [X], [Y], [A, C])
|
|
onnx_model = make_model(graph)
|
|
check_model(onnx_model)
|
|
return onnx_model
|
|
|
|
|
|
ALL_MODELS = [
|
|
const_model,
|
|
linear_model
|
|
]
|
|
|
|
|
|
class CommandLineTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.test_dir = OUTPUT_PATH / "command-line"
|
|
shutil.rmtree(cls.test_dir, ignore_errors=True)
|
|
cls.test_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def get_run_path(self, model_name: str) -> Path:
|
|
run_path = CommandLineTest.test_dir / model_name
|
|
run_path.mkdir(exist_ok=True)
|
|
return run_path
|
|
|
|
def run_model_intern(self, onnx_model: onnx.ModelProto, model_name: str):
|
|
run_path = self.get_run_path(model_name)
|
|
model_file = run_path / f"{model_name}-i.onnx"
|
|
mlir_file = run_path / f"{model_name}-i.torch.mlir"
|
|
onnx.save(onnx_model, model_file)
|
|
args = __main__.parse_arguments([
|
|
str(model_file), "-o", str(mlir_file)])
|
|
__main__.main(args)
|
|
|
|
def run_model_extern(self, onnx_model: onnx.ModelProto, model_name: str):
|
|
run_path = self.get_run_path(model_name)
|
|
model_file = run_path / f"{model_name}-e.onnx"
|
|
mlir_file = run_path / f"{model_name}-e.torch.mlir"
|
|
data_dir_name = f"{model_name}-data"
|
|
model_data_dir = run_path / data_dir_name
|
|
model_data_dir.mkdir(exist_ok=True)
|
|
convert_model_to_external_data(
|
|
onnx_model, all_tensors_to_one_file=True,
|
|
location=data_dir_name + "/data.bin",
|
|
size_threshold=48,
|
|
convert_attribute=True)
|
|
onnx.save(onnx_model, model_file)
|
|
temp_dir = run_path / "temp"
|
|
temp_dir.mkdir(exist_ok=True)
|
|
args = __main__.parse_arguments([
|
|
str(model_file), "-o", str(mlir_file), "--keep-temps", "--temp-dir",
|
|
str(temp_dir), "--data-dir", str(run_path)])
|
|
__main__.main(args)
|
|
|
|
def test_all(self):
|
|
for model_func in ALL_MODELS:
|
|
model_name = model_func.__name__
|
|
model = model_func()
|
|
with self.subTest(f"model {model_name}", model_name=model_name):
|
|
with self.subTest("Internal data"):
|
|
self.run_model_intern(model, model_name)
|
|
with self.subTest("External data"):
|
|
self.run_model_extern(model, model_name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
unittest.main()
|