Skip to content
Snippets Groups Projects
Commit 89a1cc7c authored by Yifan Zhao's avatar Yifan Zhao
Browse files

Frontend prints metadata for tuning (WIP)

parent 38b4d263
No related branches found
No related tags found
No related merge requests found
[
{
"name": "11",
"speedup": 1.0,
"applies_to": [
"Conv2D",
"MatMul"
]
},
{
"name": "12",
"speedup": 1.5,
"applies_to": [
"Conv2D",
"MatMul"
]
},
{
"name": "151",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "152",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "153",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "154",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "155",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "156",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "157",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "158",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "159",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "160",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "161",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "162",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "163",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "164",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "165",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "166",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "167",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "168",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "261",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "262",
"speedup": 3.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "263",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "264",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "265",
"speedup": 2.25,
"applies_to": [
"Conv2D"
]
},
{
"name": "266",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "267",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "268",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
},
{
"name": "269",
"speedup": 2.0,
"applies_to": [
"Conv2D"
]
}
]
...@@ -25,6 +25,8 @@ ModelTy = Union[PathLike, onnx.ModelProto, Module] ...@@ -25,6 +25,8 @@ ModelTy = Union[PathLike, onnx.ModelProto, Module]
# Path to a pair of HPVM bin format file (input, labels), or PyTorch Dataset # Path to a pair of HPVM bin format file (input, labels), or PyTorch Dataset
DatasetTy = Union[BinDataset, Dataset] DatasetTy = Union[BinDataset, Dataset]
def_approx_knobs_file = Path(__file__).parent / "approxknobs.json"
class ModelExporter: class ModelExporter:
tuneset_name = "tune_input.bin", "tune_labels.bin" tuneset_name = "tune_input.bin", "tune_labels.bin"
...@@ -74,9 +76,30 @@ class ModelExporter: ...@@ -74,9 +76,30 @@ class ModelExporter:
input_, labels = self.testset_name input_, labels = self.testset_name
self._dump_dataset(self.test_dataset, self.weight_dir / input_, self.weight_dir / labels) self._dump_dataset(self.test_dataset, self.weight_dir / input_, self.weight_dir / labels)
def export_metadata(
self,
output: PathLike = None, approx_knobs_file: PathLike = def_approx_knobs_file
):
import json
from collections import defaultdict
with Path(approx_knobs_file).open() as f:
knobs = json.load(f)
ty_knobs = defaultdict(list)
for k in knobs:
for ty in k.pop("applies_to"):
ty_knobs[ty].append((k["name"], k["speedup"]))
knobs_used = set()
for node in self.dfg.traverse_order:
knobs = ty_knobs.get(node.op_type, [])
flops = node.get_flops()
knobs_used.update(knobs)
print(f"{node.name} ({node.op_type}) -> {knobs}, {flops}")
def export_all(self, output: PathLike = None, batch_size: Optional[int] = None): def export_all(self, output: PathLike = None, batch_size: Optional[int] = None):
default_codefile = self.output_dir / self.source_file_name default_codefile = self.output_dir / self.source_file_name
self.export_source_code(output or default_codefile, batch_size) self.export_source_code(output or default_codefile, batch_size)
self.export_metadata()
self.export_weights() self.export_weights()
self.export_datasets() self.export_datasets()
...@@ -195,3 +218,7 @@ def torch_to_onnx( ...@@ -195,3 +218,7 @@ def torch_to_onnx(
output_names=["output"], # the model's output names output_names=["output"], # the model's output names
strip_doc_string=False, strip_doc_string=False,
) )
def get_full_typename(o: object) -> str:
return o.__module__ + "." + o.__class__.__qualname__
...@@ -195,7 +195,7 @@ class DFG(object): ...@@ -195,7 +195,7 @@ class DFG(object):
attrs["input_shapes"] = predec_shapes[:2] attrs["input_shapes"] = predec_shapes[:2]
mul_node = g.MatMulNode(node.name, **attrs) mul_node = g.MatMulNode(node.name, **attrs)
if node.op_type == "Gemm": if node.op_type == "Gemm":
mul_node.gemm_transpose(node, predec) mul_node.gemm_transpose(predec)
if len(predec) == 2: if len(predec) == 2:
return mul_node return mul_node
# Split into mul followed by an addition # Split into mul followed by an addition
......
...@@ -5,8 +5,6 @@ from typing import List, Optional, Sequence, Tuple ...@@ -5,8 +5,6 @@ from typing import List, Optional, Sequence, Tuple
import numpy as np import numpy as np
import onnx import onnx
from .onnx_attr import node_attr_to_dict
ShapeT = Optional[List[int]] ShapeT = Optional[List[int]]
...@@ -24,12 +22,7 @@ class DFGNode(abc.ABC): ...@@ -24,12 +22,7 @@ class DFGNode(abc.ABC):
self.name = name self.name = name
self.input_shapes = input_shapes self.input_shapes = input_shapes
self.output_shape = output_shape self.output_shape = output_shape
sin = " x ".join(str(sh) if sh else "??" for sh in input_shapes) self.attrs = kwargs
sout = output_shape if output_shape else "??"
if sin:
print(f"{name}: {sin} -> {sout})")
else:
print(f"{name}: {sout}")
def codegen(self) -> Tuple[str, list]: def codegen(self) -> Tuple[str, list]:
return "", [] return "", []
...@@ -37,8 +30,16 @@ class DFGNode(abc.ABC): ...@@ -37,8 +30,16 @@ class DFGNode(abc.ABC):
def hpvm_codegen(self) -> Tuple[str, list]: def hpvm_codegen(self) -> Tuple[str, list]:
return "", [] return "", []
def get_flops(self) -> int:
return np.prod(self.output_shape) if self.output_shape else 0
def __repr__(self) -> str: def __repr__(self) -> str:
return f"{self.op_type}({self.name})" sin = " x ".join(str(sh) if sh else "??" for sh in self.input_shapes)
sout = self.output_shape if self.output_shape else "??"
if sin:
return f"{self.name}({self.op_type}): {sin} -> {sout} ({self.get_flops()})"
else:
return f"{self.name}({self.op_type}): {sout}"
class TensorNode(DFGNode, abc.ABC): class TensorNode(DFGNode, abc.ABC):
...@@ -148,6 +149,14 @@ class Conv2DNode(DFGNode): ...@@ -148,6 +149,14 @@ class Conv2DNode(DFGNode):
[self.pads, self.pads, self.sh, self.sw], [self.pads, self.pads, self.sh, self.sw],
) )
def get_flops(self) -> int:
_, kshape = self.input_shapes
if not self.output_shape or not kshape:
return 0
_, _, h, w = self.output_shape
c1, c2, kh, kw = kshape
return int(c1 * c2 * h * w * kh * kw / (self.sh * self.sw))
class _Pool2DNode(DFGNode, abc.ABC): class _Pool2DNode(DFGNode, abc.ABC):
"""Common super class of Average pooling and Max pooling.""" """Common super class of Average pooling and Max pooling."""
...@@ -188,6 +197,10 @@ class _Pool2DNode(DFGNode, abc.ABC): ...@@ -188,6 +197,10 @@ class _Pool2DNode(DFGNode, abc.ABC):
[*self.kernel_shape, *self.pads, *self.strides], [*self.kernel_shape, *self.pads, *self.strides],
) )
def get_flops(self) -> int:
input0 = self.input_shapes[0]
return np.prod(input0) if input0 else 0
class MaxPool2DNode(_Pool2DNode): class MaxPool2DNode(_Pool2DNode):
pool_type = "0" pool_type = "0"
...@@ -218,27 +231,36 @@ class MatMulNode(DFGNode): ...@@ -218,27 +231,36 @@ class MatMulNode(DFGNode):
def hpvm_codegen(self): def hpvm_codegen(self):
return "__hpvm__tensor_mul", [] return "__hpvm__tensor_mul", []
@staticmethod def gemm_transpose(self, predec):
def gemm_transpose(onnx_gemm_node, predec):
"""Find and transpose weights of the onnx gemm node. """Find and transpose weights of the onnx gemm node.
This way we transpose the constant weight instead of exporting This way we transpose the constant weight instead of exporting
a transpose node (which doesn't yet exist in HPVM). a transpose node (which doesn't yet exist in HPVM).
""" """
def _transpose(weight): def _transpose(idx: int):
weight = predec[idx]
if not isinstance(weight, WeightTensor): if not isinstance(weight, WeightTensor):
raise ValueError( raise ValueError(
f"Cannot transpose non-const {weight} (transpose op needed)" f"Cannot transpose non-const {weight} (transpose op needed)"
) )
weight.transpose_() weight.transpose_()
self.input_shapes[idx] = weight.output_shape
# Some tensors may need transposing # Some tensors may need transposing
attrs = node_attr_to_dict(onnx_gemm_node) if self.attrs.get("transA", False):
if attrs.get("transA", False): _transpose(0)
_transpose(predec[0]) if self.attrs.get("transB", False):
if attrs.get("transB", False): _transpose(1)
_transpose(predec[1])
def get_flops(self) -> int:
ishape, wshape = self.input_shapes
if not ishape or not wshape:
return 0
input_len = np.prod(ishape)
_, _, len_, k = wshape
assert input_len == len_
return input_len * k
class SoftMaxNode(DFGNode): class SoftMaxNode(DFGNode):
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment