mirror of
https://github.com/mit-han-lab/tinyengine.git
synced 2025-05-10 17:31:24 +08:00
88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
import warnings
|
|
|
|
from numpy import iterable
|
|
|
|
from .basic_utils import basicOperator, deep_copy_dicts, overwrite_dicts
|
|
|
|
__all__ = ["div"]
|
|
|
|
default_params = {
|
|
# op related
|
|
"op": "DIV",
|
|
"input_idx": None,
|
|
"input2_idx": None,
|
|
"output_idx": None,
|
|
# tensor related
|
|
"input_size": None,
|
|
"input2": None,
|
|
"input_dtype": "int8",
|
|
"input2_dtype": "int8",
|
|
"output_dtype": "int8",
|
|
# quantization related
|
|
"weight_value": None,
|
|
"bias": None,
|
|
"input_zero_point": None,
|
|
"output_zero_point": None,
|
|
"input_scale": None,
|
|
"output_scale": None,
|
|
"multiplier": None,
|
|
"shift": None,
|
|
# input of scale from some conv2d
|
|
"scale_from_add": None,
|
|
}
|
|
|
|
|
|
class div(basicOperator):
|
|
div_const_cnt = 0
|
|
|
|
def __init__(self, params: dict) -> None:
|
|
self.params = deep_copy_dicts(default_params)
|
|
overwrite_dicts(self.params, params)
|
|
super().__init__()
|
|
# handle input/output tensors in HWC format
|
|
self._add_input(self.params["input_idx"], self.params["input_dtype"], self.params["input_size"], 1, 1)
|
|
self._add_input(self.params["input2_idx"], self.params["input2_dtype"], self.params["input_size"], 1, 1)
|
|
self._add_output(self.params["output_idx"], self.params["output_dtype"], self.params["input_size"], 1, 1)
|
|
|
|
if None in default_params:
|
|
warnings.warn(f"parameters are not all set for op {self.params['op']}")
|
|
|
|
def generate_inference_str(self):
|
|
params = self.params
|
|
string = ""
|
|
|
|
if params["input_dtype"] == "float32":
|
|
if self.params["scale_from_add"] is not None:
|
|
scale_divisor = f"{self.params['scale_from_add']}"
|
|
string += (
|
|
f"fptr3 = {self._getBufferstrCast(params['output_buf_add'], params['output_buf_add_offset'])}; "
|
|
+ f"fptr2 = {self._getBufferstrCast(params['input_buf_add'], params['input_buf_add_offset'])};\n"
|
|
)
|
|
string += (
|
|
f"for(int i = 0; i < {self.params['input_size']}; i++) *fptr3++ = *fptr2++ / {scale_divisor};\n"
|
|
)
|
|
return string
|
|
elif "constant" in self.params["input2_idx"]:
|
|
string += f"const float fptr{div.div_const_cnt}[] = " + " {"
|
|
if iterable(self.params["input2"]):
|
|
for v in self.params["input2"]:
|
|
string += f"{str(v)},"
|
|
else:
|
|
string += f"{str(self.params['input2'])},"
|
|
string += "};\n"
|
|
input2_str = f"fptr{div.div_const_cnt}"
|
|
div.div_const_cnt += 1
|
|
else:
|
|
input2_str = f"{self._getBufferstrCast(params['input2_buf_add'], params['input2_buf_add_offset'])}"
|
|
|
|
string += (
|
|
f"div_fp({self.params['input_size']},"
|
|
+ f"{self._getBufferstrCast(params['input_buf_add'], params['input_buf_add_offset'])},"
|
|
+ f"{input2_str},"
|
|
+ f"{self._getBufferstrCast(params['output_buf_add'], params['output_buf_add_offset'])});\n"
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return string
|