From a6309ad64515dbf688c71eeedca6f116fba8e2da Mon Sep 17 00:00:00 2001
From: Yifan Zhao <yifanz16@illinois.edu>
Date: Wed, 27 Jan 2021 22:03:49 -0600
Subject: [PATCH] Improved torch2hpvm setup script
---
hpvm/projects/torch2hpvm/.gitignore | 3 +-
hpvm/projects/torch2hpvm/env.yaml | 10 --
hpvm/projects/torch2hpvm/main.py | 107 ------------------
hpvm/projects/torch2hpvm/setup.py | 16 +--
.../torch2hpvm/torch2hpvm/__init__.py | 2 +
.../torch2hpvm/torch2hpvm/__main__.py | 55 +++++++++
.../projects/torch2hpvm/torch2hpvm/compile.py | 52 +++++++++
7 files changed, 118 insertions(+), 127 deletions(-)
delete mode 100644 hpvm/projects/torch2hpvm/env.yaml
delete mode 100644 hpvm/projects/torch2hpvm/main.py
create mode 100644 hpvm/projects/torch2hpvm/torch2hpvm/__main__.py
create mode 100644 hpvm/projects/torch2hpvm/torch2hpvm/compile.py
diff --git a/hpvm/projects/torch2hpvm/.gitignore b/hpvm/projects/torch2hpvm/.gitignore
index 2f4db75f1d..25aacffde0 100644
--- a/hpvm/projects/torch2hpvm/.gitignore
+++ b/hpvm/projects/torch2hpvm/.gitignore
@@ -1,4 +1,3 @@
build/
dist/
-onnx2hpvm.egg-info/
-.ipynb_checkpoints/
+*.egg-info/
diff --git a/hpvm/projects/torch2hpvm/env.yaml b/hpvm/projects/torch2hpvm/env.yaml
deleted file mode 100644
index d04bb1b503..0000000000
--- a/hpvm/projects/torch2hpvm/env.yaml
+++ /dev/null
@@ -1,10 +0,0 @@
-name: onnxfront
-channels:
- - defaults
-dependencies:
- - pip
- - python=3.8.5=h7579374_1
- - jinja2
- - networkx
- - pip:
- - onnx==1.8.0
diff --git a/hpvm/projects/torch2hpvm/main.py b/hpvm/projects/torch2hpvm/main.py
deleted file mode 100644
index 7dc23d4873..0000000000
--- a/hpvm/projects/torch2hpvm/main.py
+++ /dev/null
@@ -1,107 +0,0 @@
-from pathlib import Path
-from typing import Optional
-
-import onnx
-
-
-def check_version(model, new_version):
- try:
- opset = model.opset_import[0].version if model.opset_import else 1
- except AttributeError:
- opset = 1 # default opset version set to 1 if not specified
- if opset != new_version:
- from onnx import version_converter
-
- try:
- converted_model = version_converter.convert_version(model, new_version)
- return converted_model
- except RuntimeError as e:
- raise RuntimeError(
- f"Current version {opset} of ONNX model not supported!\n"
- f"Conversion failed with message below: \n{e}"
- )
- return model
-
-
-def compile(
- onnx_file: Path,
- output_dir: Path,
- input_size: int,
- prefix: Optional[str],
- batch_size: Optional[int],
- opset: Optional[int],
- hpvmc: bool,
-):
- from frontend.graph_builder import DFG
- from frontend.codegen_tensor import TensorCodeGen
- from frontend.codegen_hpvm import HpvmCodeGen
-
- model = onnx.load(onnx_file)
- if opset is not None:
- model = check_version(model, opset)
- model = onnx.shape_inference.infer_shapes(model)
- dfg = DFG(model.graph)
- if hpvmc:
- hpvm_code_gen = HpvmCodeGen(dfg, output_dir, input_size, batch_size, prefix)
- hpvm_code_gen.compile()
- else:
- tensor_code_gen = TensorCodeGen(dfg, output_dir, input_size, batch_size, prefix)
- tensor_code_gen.compile()
- dfg.dump_weights(output_dir)
-
-
-def parse_args():
- import argparse
-
- parser = argparse.ArgumentParser(description="ONNX to HPVM-C")
- parser.add_argument("onnx_file", type=Path, help="Path to input ONNX file")
- parser.add_argument(
- "output_dir",
- type=Path,
- help="Output folder where source file and weight files are generated",
- )
- parser.add_argument(
- "input_size", type=int, help="Size of input dataset",
- )
- parser.add_argument(
- "-p",
- "--prefix",
- type=str,
- help="Prefix in generated code; will be attached before name of weight/input files. "
- "Defaults to output_dir.",
- )
- parser.add_argument(
- "-b",
- "--batch-size",
- type=int,
- help="Batch size to be used in the generated code. "
- "Defaults to input size (i.e., not using batch).",
- )
- parser.add_argument("--opset", type=int, help="ONNX opset version (enforced)")
- parser.add_argument(
- "-c",
- "--compile-mode",
- type=str,
- choices=["tensor", "hpvmc"],
- default="hpvmc",
- help="""Output mode.
-tensor: HPVM Tensor Runtime;
-hpvmc: HPVM C Interface. Default value is hpvmc.""",
- )
-
- args = parser.parse_args()
- args.hpvmc = args.compile_mode == "hpvmc"
- delattr(args, "compile_mode")
- return args
-
-
-def main():
- import os
-
- args = parse_args()
- os.makedirs(args.output_dir, exist_ok=True)
- compile(**vars(args))
-
-
-if __name__ == "__main__":
- main()
diff --git a/hpvm/projects/torch2hpvm/setup.py b/hpvm/projects/torch2hpvm/setup.py
index 37dcf5f039..fd21f8b59d 100644
--- a/hpvm/projects/torch2hpvm/setup.py
+++ b/hpvm/projects/torch2hpvm/setup.py
@@ -1,12 +1,12 @@
-
from setuptools import setup
setup(
- name='torch2hpvm',
- version='1.0',
- description='PyTorch frontend for HPVM',
- author='Yuanjing Shi, Yifan Zhao',
- author_email='ys26@illinois.edu, yifanz16@illinois.edu',
- packages=['torch2hpvm'],
- install_requires=[],
+ name="torch2hpvm",
+ version="1.0",
+ description="PyTorch frontend for HPVM",
+ author="Yuanjing Shi, Yifan Zhao",
+ author_email="ys26@illinois.edu, yifanz16@illinois.edu",
+ packages=["torch2hpvm"],
+ install_requires=["jinja2>=2.11", "networkx>=2.5", "onnx>=1.8.0"],
+ entry_points={"console_scripts": ["torch2hpvm=torch2hpvm:main"]},
)
diff --git a/hpvm/projects/torch2hpvm/torch2hpvm/__init__.py b/hpvm/projects/torch2hpvm/torch2hpvm/__init__.py
index e69de29bb2..8e7de3cf9a 100644
--- a/hpvm/projects/torch2hpvm/torch2hpvm/__init__.py
+++ b/hpvm/projects/torch2hpvm/torch2hpvm/__init__.py
@@ -0,0 +1,2 @@
+from .compile import compile
+from .__main__ import main
diff --git a/hpvm/projects/torch2hpvm/torch2hpvm/__main__.py b/hpvm/projects/torch2hpvm/torch2hpvm/__main__.py
new file mode 100644
index 0000000000..8eae267f47
--- /dev/null
+++ b/hpvm/projects/torch2hpvm/torch2hpvm/__main__.py
@@ -0,0 +1,55 @@
+import os
+from pathlib import Path
+
+from .compile import compile
+
+
+def parse_args():
+ import argparse
+
+ parser = argparse.ArgumentParser(description="ONNX to HPVM-C")
+ parser.add_argument("onnx_file", type=Path, help="Path to input ONNX file")
+ parser.add_argument(
+ "output_dir",
+ type=Path,
+ help="Output folder where source file and weight files are generated",
+ )
+ parser.add_argument(
+ "input_size", type=int, help="Size of input dataset",
+ )
+ parser.add_argument(
+ "-p",
+ "--prefix",
+ type=str,
+ help="Prefix in generated code; will be attached before name of weight/input files. "
+ "Defaults to output_dir.",
+ )
+ parser.add_argument(
+ "-b",
+ "--batch-size",
+ type=int,
+ help="Batch size to be used in the generated code. "
+ "Defaults to input size (i.e., not using batch).",
+ )
+ parser.add_argument("--opset", type=int, help="ONNX opset version (enforced)")
+ parser.add_argument(
+ "-c",
+ "--compile-mode",
+ type=str,
+ choices=["tensor", "hpvmc"],
+ default="hpvmc",
+ help="""Output mode.
+tensor: HPVM Tensor Runtime;
+hpvmc: HPVM C Interface. Default value is hpvmc.""",
+ )
+
+ args = parser.parse_args()
+ args.hpvmc = args.compile_mode == "hpvmc"
+ delattr(args, "compile_mode")
+ return args
+
+
+def main():
+ args = parse_args()
+ os.makedirs(args.output_dir, exist_ok=True)
+ compile(**vars(args))
diff --git a/hpvm/projects/torch2hpvm/torch2hpvm/compile.py b/hpvm/projects/torch2hpvm/torch2hpvm/compile.py
new file mode 100644
index 0000000000..694e5f6231
--- /dev/null
+++ b/hpvm/projects/torch2hpvm/torch2hpvm/compile.py
@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Optional, Union
+
+import onnx
+from onnx import version_converter
+
+from .codegen_hpvm import HpvmCodeGen
+from .codegen_tensor import TensorCodeGen
+from .graph_builder import DFG
+
+PathLike = Union[Path, str]
+
+
+def check_version(model, new_version):
+ try:
+ opset = model.opset_import[0].version if model.opset_import else 1
+ except AttributeError:
+ opset = 1 # default opset version set to 1 if not specified
+ if opset != new_version:
+
+ try:
+ converted_model = version_converter.convert_version(model, new_version)
+ return converted_model
+ except RuntimeError as e:
+ raise RuntimeError(
+ f"Current version {opset} of ONNX model not supported!\n"
+ f"Conversion failed with message below: \n{e}"
+ )
+ return model
+
+
+def compile(
+ onnx_file: Path,
+ output_dir: Path,
+ input_size: int,
+ prefix: Optional[str],
+ batch_size: Optional[int],
+ opset: Optional[int],
+ hpvmc: bool,
+):
+ model = onnx.load(onnx_file)
+ if opset is not None:
+ model = check_version(model, opset)
+ model = onnx.shape_inference.infer_shapes(model)
+ dfg = DFG(model.graph)
+ if hpvmc:
+ hpvm_code_gen = HpvmCodeGen(dfg, output_dir, input_size, batch_size, prefix)
+ hpvm_code_gen.compile()
+ else:
+ tensor_code_gen = TensorCodeGen(dfg, output_dir, input_size, batch_size, prefix)
+ tensor_code_gen.compile()
+ dfg.dump_weights(output_dir)
--
GitLab