2189 lines
79 KiB
Python
2189 lines
79 KiB
Python
import array
|
|
import enum
|
|
import functools
|
|
import logging
|
|
import operator
|
|
import struct
|
|
import sys
|
|
from typing import List, NamedTuple, Optional, Tuple
|
|
|
|
import torch
|
|
|
|
|
|
# TODO: Add type annotations
|
|
# TODO: Check tensor types for ops
|
|
|
|
|
|
LOG = logging.getLogger("nnapi_serialize")
|
|
|
|
|
|
class NNAPI_OperandCode:
|
|
FLOAT32 = 0
|
|
INT32 = 1
|
|
UINT32 = 2
|
|
TENSOR_FLOAT32 = 3
|
|
TENSOR_INT32 = 4
|
|
TENSOR_QUANT8_ASYMM = 5
|
|
BOOL = 6
|
|
TENSOR_QUANT16_SYMM = 7
|
|
TENSOR_FLOAT16 = 8
|
|
TENSOR_BOOL8 = 9
|
|
FLOAT16 = 10
|
|
TENSOR_QUANT8_SYMM_PER_CHANNEL = 11
|
|
TENSOR_QUANT16_ASYMM = 12
|
|
|
|
|
|
class NNAPI_OperationCode:
|
|
ADD = 0
|
|
AVERAGE_POOL_2D = 1
|
|
CONCATENATION = 2
|
|
CONV_2D = 3
|
|
DEPTHWISE_CONV_2D = 4
|
|
DEPTH_TO_SPACE = 5
|
|
DEQUANTIZE = 6
|
|
EMBEDDING_LOOKUP = 7
|
|
FLOOR = 8
|
|
FULLY_CONNECTED = 9
|
|
HASHTABLE_LOOKUP = 10
|
|
L2_NORMALIZATION = 11
|
|
L2_POOL_2D = 12
|
|
LOCAL_RESPONSE_NORMALIZATION = 13
|
|
LOGISTIC = 14
|
|
LSH_PROJECTION = 15
|
|
LSTM = 16
|
|
MAX_POOL_2D = 17
|
|
MUL = 18
|
|
RELU = 19
|
|
RELU1 = 20
|
|
RELU6 = 21
|
|
RESHAPE = 22
|
|
RESIZE_BILINEAR = 23
|
|
RNN = 24
|
|
SOFTMAX = 25
|
|
SPACE_TO_DEPTH = 26
|
|
SVDF = 27
|
|
TANH = 28
|
|
BATCH_TO_SPACE_ND = 29
|
|
DIV = 30
|
|
MEAN = 31
|
|
PAD = 32
|
|
SPACE_TO_BATCH_ND = 33
|
|
SQUEEZE = 34
|
|
STRIDED_SLICE = 35
|
|
SUB = 36
|
|
TRANSPOSE = 37
|
|
ABS = 38
|
|
ARGMAX = 39
|
|
ARGMIN = 40
|
|
AXIS_ALIGNED_BBOX_TRANSFORM = 41
|
|
BIDIRECTIONAL_SEQUENCE_LSTM = 42
|
|
BIDIRECTIONAL_SEQUENCE_RNN = 43
|
|
BOX_WITH_NMS_LIMIT = 44
|
|
CAST = 45
|
|
CHANNEL_SHUFFLE = 46
|
|
DETECTION_POSTPROCESSING = 47
|
|
EQUAL = 48
|
|
EXP = 49
|
|
EXPAND_DIMS = 50
|
|
GATHER = 51
|
|
GENERATE_PROPOSALS = 52
|
|
GREATER = 53
|
|
GREATER_EQUAL = 54
|
|
GROUPED_CONV_2D = 55
|
|
HEATMAP_MAX_KEYPOINT = 56
|
|
INSTANCE_NORMALIZATION = 57
|
|
LESS = 58
|
|
LESS_EQUAL = 59
|
|
LOG = 60
|
|
LOGICAL_AND = 61
|
|
LOGICAL_NOT = 62
|
|
LOGICAL_OR = 63
|
|
LOG_SOFTMAX = 64
|
|
MAXIMUM = 65
|
|
MINIMUM = 66
|
|
NEG = 67
|
|
NOT_EQUAL = 68
|
|
PAD_V2 = 69
|
|
POW = 70
|
|
PRELU = 71
|
|
QUANTIZE = 72
|
|
QUANTIZED_16BIT_LSTM = 73
|
|
RANDOM_MULTINOMIAL = 74
|
|
REDUCE_ALL = 75
|
|
REDUCE_ANY = 76
|
|
REDUCE_MAX = 77
|
|
REDUCE_MIN = 78
|
|
REDUCE_PROD = 79
|
|
REDUCE_SUM = 80
|
|
ROI_ALIGN = 81
|
|
ROI_POOLING = 82
|
|
RSQRT = 83
|
|
SELECT = 84
|
|
SIN = 85
|
|
SLICE = 86
|
|
SPLIT = 87
|
|
SQRT = 88
|
|
TILE = 89
|
|
TOPK_V2 = 90
|
|
TRANSPOSE_CONV_2D = 91
|
|
UNIDIRECTIONAL_SEQUENCE_LSTM = 92
|
|
UNIDIRECTIONAL_SEQUENCE_RNN = 93
|
|
RESIZE_NEAREST_NEIGHBOR = 94
|
|
|
|
|
|
class NNAPI_FuseCode:
|
|
FUSED_NONE = 0
|
|
FUSED_RELU = 1
|
|
FUSED_RELU1 = 2
|
|
FUSED_RELU6 = 3
|
|
|
|
|
|
class OperandValueSourceType:
|
|
IMMEDIATE = 0
|
|
NUMBERED_BUFFER = 2
|
|
NUMBERED_MEMORY = 3
|
|
|
|
|
|
# Scalar types that appear explicitly in models.
|
|
# These must be kept in sync with
|
|
# AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS.
|
|
# TODO: Expose these directly to Python to avoid maintaining this list.
|
|
class TorchScalarTypes(enum.Enum):
|
|
QUINT8 = 13
|
|
|
|
|
|
def approx_equal(lhs, rhs, tolerance=1e-6):
|
|
return abs(lhs - rhs) <= tolerance * min(lhs, rhs)
|
|
|
|
|
|
def tensor_size(op_type, dims):
|
|
ITEM_SIZES = {
|
|
NNAPI_OperandCode.TENSOR_FLOAT32: 4,
|
|
NNAPI_OperandCode.TENSOR_INT32: 4,
|
|
NNAPI_OperandCode.TENSOR_QUANT8_ASYMM: 1,
|
|
NNAPI_OperandCode.TENSOR_QUANT16_SYMM: 2,
|
|
NNAPI_OperandCode.TENSOR_QUANT16_ASYMM: 2,
|
|
}
|
|
size = ITEM_SIZES[op_type]
|
|
for d in dims:
|
|
size *= d
|
|
return size
|
|
|
|
|
|
def change_element(tup, index, value):
|
|
ls = list(tup)
|
|
ls[index] = value
|
|
return tuple(ls)
|
|
|
|
|
|
class ConvPoolArgs2d(NamedTuple):
|
|
"""Configuration arguments for a convolution."""
|
|
|
|
kernel_h: int
|
|
kernel_w: int
|
|
stride_h: int
|
|
stride_w: int
|
|
pad_t: int
|
|
pad_b: int
|
|
pad_l: int
|
|
pad_r: int
|
|
dilation_h: int
|
|
dilation_w: int
|
|
group: int
|
|
|
|
|
|
class DimOrder(enum.Enum):
|
|
PRESUMED_CONTIGUOUS = 0
|
|
CHANNELS_LAST = 1
|
|
SCALAR_OR_VECTOR = 2
|
|
UNKNOWN_CONSTANT = 999
|
|
|
|
|
|
class Operand(NamedTuple):
|
|
"""Represenation of an NNAPI operand."""
|
|
|
|
# NNAPI operand type. One of NNAPI_OperandCode.
|
|
# TODO: Make this an enum.
|
|
op_type: int
|
|
|
|
# This is always the PyTorch shape, which is NCHW for feature maps.
|
|
# The actual NNAPI operand might have a transposed shape.
|
|
# we use 0 for load time dynamic shapes & -1 for runtime dynamic shapes
|
|
shape: Tuple[int, ...]
|
|
|
|
# Specifies how the shape of the operand that we define in NNAPI
|
|
# relates to the shape we track above.
|
|
# - PRESUMED_CONTIGUOUS: physical NNAPI operand will exactly match
|
|
# the shape of the PyTorch tensor.
|
|
# - CHANNELS_LAST: The PyTorch tensor is expected to be NCHW, and
|
|
# the NNAPI operand will be represented explicitly as NHWC.
|
|
dim_order: DimOrder
|
|
|
|
# Quantization params
|
|
scale: float
|
|
zero_point: int
|
|
|
|
def use_nchw(self):
|
|
if self.dim_order is DimOrder.PRESUMED_CONTIGUOUS:
|
|
return True
|
|
if self.dim_order is DimOrder.CHANNELS_LAST:
|
|
return False
|
|
raise Exception("Unknown dim order")
|
|
|
|
|
|
def broadcast_shapes(shape1, shape2):
|
|
assert len(shape1) > 0
|
|
assert len(shape2) > 0
|
|
s1 = list(shape1)
|
|
s2 = list(shape2)
|
|
# TODO: Support non-equal-rank broadcast where semantics match.
|
|
# This can be tricky for NHWC tensors because dimension orders
|
|
# don't match between PT and NNAPI, even though semantics match.
|
|
if len(s1) > len(s2):
|
|
# s2 = [1] * (len(s1) - len(s2)) + s2
|
|
raise Exception("Non-equal-rank broadcast is not supported yet.")
|
|
if len(s2) > len(s1):
|
|
# s3 = [1] * (len(s2) - len(s1)) + s1
|
|
raise Exception("Non-equal-rank broadcast is not supported yet.")
|
|
ret = []
|
|
for d1, d2 in zip(s1, s2):
|
|
if d1 == 1:
|
|
ret.append(d2)
|
|
elif d2 == 1:
|
|
ret.append(d1)
|
|
elif d1 == d2:
|
|
ret.append(d1)
|
|
else:
|
|
raise Exception(f"Cannot broadcast shapes: {shape1} and {shape2}")
|
|
return tuple(ret)
|
|
|
|
|
|
def get_conv_pool_shape(image_shape, args, out_ch, transpose):
|
|
batch, in_c, in_h, in_w = image_shape
|
|
|
|
# TODO: Handle dilation
|
|
if args.dilation_h != 1 or args.dilation_w != 1:
|
|
raise Exception("Dilation not supported yet.")
|
|
|
|
if transpose:
|
|
out_h = (in_h - 1) * args.stride_h + args.kernel_h - args.pad_t - args.pad_b
|
|
out_w = (in_w - 1) * args.stride_w + args.kernel_w - args.pad_l - args.pad_l
|
|
else:
|
|
out_h = (in_h - args.kernel_h + args.pad_t + args.pad_b) // args.stride_h + 1
|
|
out_w = (in_w - args.kernel_w + args.pad_l + args.pad_r) // args.stride_w + 1
|
|
|
|
# Handle variable-sized tensors.
|
|
if in_h == 0:
|
|
out_h = 0
|
|
if in_w == 0:
|
|
out_w = 0
|
|
|
|
out_shape = (batch, out_ch, out_h, out_w)
|
|
return out_shape
|
|
|
|
|
|
def fix_shape(shape, dim_order):
|
|
# Return the actual shape that an operand should have in NNAPI,
|
|
# given a PyTorch shape and dimension order. This is where we
|
|
# convert from PyTorch's "always NCHW" shape to explicit NHWC.
|
|
if dim_order is DimOrder.PRESUMED_CONTIGUOUS:
|
|
return shape
|
|
if dim_order is DimOrder.CHANNELS_LAST:
|
|
return tuple([shape[0]] + list(shape[2:]) + [shape[1]])
|
|
if dim_order is DimOrder.SCALAR_OR_VECTOR:
|
|
assert len(shape) == 0 or len(shape) == 1
|
|
return shape
|
|
if dim_order is DimOrder.UNKNOWN_CONSTANT:
|
|
# XXX think this through
|
|
return shape
|
|
raise Exception(f"Bad dim_order: {dim_order!r}.")
|
|
|
|
|
|
def reverse_map_dim(dim_order, d):
|
|
# Return the original PyTorch dimension position for a given dimension.
|
|
# d should be the dimension that NNAPI will see.
|
|
# reverse_map_dim(PRESUMED_CONTIGUOUS, x) == x
|
|
# reverse_map_dim(CHANNELS_LAST, 3) == 1
|
|
if dim_order in (DimOrder.PRESUMED_CONTIGUOUS, DimOrder.SCALAR_OR_VECTOR):
|
|
return d
|
|
assert dim_order is DimOrder.CHANNELS_LAST
|
|
return [0, 2, 3, 1][d]
|
|
|
|
|
|
def flex_name(op_id, dim):
|
|
# Return the local variable name for the computed flexible size
|
|
# for a given op and dimension.
|
|
return f"s_{op_id}_{dim}"
|
|
|
|
|
|
class _NnapiSerializer:
|
|
def __init__(self, config, use_int16_for_qint16=False):
|
|
self.operands = []
|
|
self.values = []
|
|
self.operations = []
|
|
self.value_data = []
|
|
self.operation_args = []
|
|
self.inputs = []
|
|
self.outputs = []
|
|
self.flexible_shape_computation_lines = []
|
|
|
|
self.modules = {}
|
|
self.constants = {}
|
|
self.tensor_sequences = {}
|
|
self.jitval_operand_map = {}
|
|
self.cached_immediates = {}
|
|
self.used_weights = []
|
|
self.weight_offset = 0
|
|
self.use_int16_for_qint16 = use_int16_for_qint16
|
|
|
|
if config is None:
|
|
config = {}
|
|
|
|
def get_next_operand_id(self):
|
|
return len(self.operands)
|
|
|
|
# Add a tensor operand corresponding to a JIT Value.
|
|
# Returns the NNAPI operand ID. Can be looked up later with
|
|
# get_tensor_operand_by_jitval.
|
|
def add_tensor_operand(self, jitval, oper):
|
|
assert isinstance(oper, Operand)
|
|
if jitval in self.jitval_operand_map:
|
|
raise Exception(f"Duplicate tensor: {jitval!r}")
|
|
|
|
operand_id = self.get_next_operand_id()
|
|
self.operands.append(oper)
|
|
self.jitval_operand_map[jitval] = operand_id
|
|
return operand_id
|
|
|
|
# Add a tensor operand that does not correspond to a JIT Value.
|
|
# Useful for cases where multiple NNAPI operands are required
|
|
# to implement one JIT IR node. Returns the NNAPI operand ID.
|
|
def add_anonymous_tensor_operand(self, oper):
|
|
assert isinstance(oper, Operand)
|
|
operand_id = self.get_next_operand_id()
|
|
self.operands.append(oper)
|
|
return operand_id
|
|
|
|
def torch_tensor_to_operand(self, tensor, dim_order):
|
|
dtype = str(tensor.dtype).replace("torch.", "")
|
|
scale = 0.0
|
|
zero_point = 0
|
|
if dtype == "float32":
|
|
op_type = NNAPI_OperandCode.TENSOR_FLOAT32
|
|
elif dtype == "int32":
|
|
op_type = NNAPI_OperandCode.TENSOR_INT32
|
|
elif dtype == "quint8":
|
|
op_type = NNAPI_OperandCode.TENSOR_QUANT8_ASYMM
|
|
scale = tensor.q_scale()
|
|
zero_point = tensor.q_zero_point()
|
|
elif dtype == "qint32":
|
|
op_type = NNAPI_OperandCode.TENSOR_INT32
|
|
scale = tensor.q_scale()
|
|
zero_point = tensor.q_zero_point()
|
|
assert zero_point == 0
|
|
elif dtype == "int16":
|
|
if self.use_int16_for_qint16:
|
|
nnapi_dtype = getattr(tensor, "nnapi_dtype", None)
|
|
op_codes = (
|
|
NNAPI_OperandCode.TENSOR_QUANT16_SYMM,
|
|
NNAPI_OperandCode.TENSOR_QUANT16_ASYMM,
|
|
)
|
|
if nnapi_dtype in op_codes:
|
|
op_type = nnapi_dtype
|
|
scale = tensor.nnapi_scale
|
|
zero_point = tensor.nnapi_zero_point
|
|
else:
|
|
raise Exception(
|
|
f"`nnapi_type` needs to be one of {op_codes} for `int16`"
|
|
)
|
|
else:
|
|
raise Exception(
|
|
"`int16` isn't supported. If you're trying to represent NNAPI"
|
|
" qint16 with Pytorch int16, set `use_int16_for_qint16 = True`"
|
|
)
|
|
else:
|
|
raise Exception(f"Can't handle input with dtype '{tensor.dtype}'")
|
|
return Operand(
|
|
shape=tuple(tensor.shape),
|
|
op_type=op_type,
|
|
dim_order=dim_order,
|
|
scale=scale,
|
|
zero_point=zero_point,
|
|
)
|
|
|
|
def add_tensor_operand_for_input(self, arg_idx, jitval, tensor):
|
|
dim_order = (
|
|
DimOrder.CHANNELS_LAST
|
|
if getattr(tensor, "nnapi_nhwc", False)
|
|
else DimOrder.PRESUMED_CONTIGUOUS
|
|
)
|
|
toper = self.torch_tensor_to_operand(tensor, dim_order)
|
|
operand_id = self.add_tensor_operand(jitval, toper)
|
|
self.inputs.append(operand_id)
|
|
for dim, size in enumerate(tensor.shape):
|
|
if size == 0:
|
|
self.compute_operand_shape(
|
|
operand_id, dim, f"args[{arg_idx}].shape[{dim}]"
|
|
)
|
|
return operand_id
|
|
|
|
def add_tensor_operand_for_weight(
|
|
self, tensor, dim_order=DimOrder.UNKNOWN_CONSTANT
|
|
):
|
|
toper = self.torch_tensor_to_operand(tensor, dim_order)
|
|
operand_id = len(self.operands)
|
|
self.operands.append(toper)
|
|
tsize = tensor_size(toper.op_type, toper.shape)
|
|
psize = ((tsize - 1) | 0x3) + 1
|
|
self.values.append((operand_id, OperandValueSourceType.NUMBERED_BUFFER))
|
|
buf_num = len(self.used_weights)
|
|
offset = 0
|
|
self.value_data.append(struct.pack("iii", buf_num, offset, tsize))
|
|
# For NHWC NNAPI op, lay out data in the same dim order by permuting torch tensor
|
|
if dim_order == DimOrder.CHANNELS_LAST:
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
self.used_weights.append(tensor)
|
|
return operand_id
|
|
|
|
def add_immediate_operand(self, code, value, dims):
|
|
assert isinstance(dims, tuple)
|
|
cache_key = (code, value)
|
|
if cache_key not in self.cached_immediates:
|
|
operand_id = len(self.operands)
|
|
self.operands.append(Operand(code, dims, DimOrder.SCALAR_OR_VECTOR, 0.0, 0))
|
|
self.values.append((operand_id, OperandValueSourceType.IMMEDIATE))
|
|
self.value_data.append(value)
|
|
self.cached_immediates[cache_key] = operand_id
|
|
return self.cached_immediates[cache_key]
|
|
|
|
def add_immediate_int_scalar(self, value):
|
|
return self.add_immediate_operand(
|
|
NNAPI_OperandCode.INT32, struct.pack("i", value), ()
|
|
)
|
|
|
|
def add_immediate_float_scalar(self, value):
|
|
return self.add_immediate_operand(
|
|
NNAPI_OperandCode.FLOAT32, struct.pack("f", value), ()
|
|
)
|
|
|
|
def add_immediate_bool_scalar(self, value):
|
|
return self.add_immediate_operand(
|
|
NNAPI_OperandCode.BOOL, b"\x01" if value else b"\x00", ()
|
|
)
|
|
|
|
def add_immediate_int_vector(self, value):
|
|
return self.add_immediate_operand(
|
|
NNAPI_OperandCode.TENSOR_INT32,
|
|
array.array("i", value).tobytes(),
|
|
(len(value),),
|
|
)
|
|
|
|
def has_operand_for_jitval(self, jitval):
|
|
return jitval in self.jitval_operand_map
|
|
|
|
def get_tensor_operand_by_jitval(self, jitval):
|
|
operand_id = self.jitval_operand_map[jitval]
|
|
return (operand_id, self.operands[operand_id])
|
|
|
|
def get_tensor_operand_by_jitval_fixed_size(self, jitval):
|
|
op_id, oper = self.get_tensor_operand_by_jitval(jitval)
|
|
for s in oper.shape:
|
|
if s == 0:
|
|
# TODO: Improve this error message, possibly after converting
|
|
# many callsites to support flexible size.
|
|
raise Exception("Flexible size is not supported for this operand.")
|
|
if s < 0:
|
|
# runtime flex
|
|
LOG.warning("Operand %s has runtime flex shape", oper)
|
|
return op_id, oper
|
|
|
|
def get_tensor_operand_or_constant(
|
|
self, jitval, dim_order=DimOrder.PRESUMED_CONTIGUOUS
|
|
):
|
|
operand_id = self.jitval_operand_map.get(jitval)
|
|
if operand_id is None:
|
|
_, value = self.get_constant_value(jitval, "TensorType")
|
|
operand_id = self.add_tensor_operand_for_weight(value, dim_order)
|
|
return (operand_id, self.operands[operand_id])
|
|
|
|
def get_tensor_operand_for_weight(self, jitval):
|
|
_, value = self.get_constant_value(jitval, "TensorType")
|
|
operand_id = self.add_tensor_operand_for_weight(value)
|
|
return (operand_id, self.operands[operand_id])
|
|
|
|
def add_operation(self, opcode, inputs, outputs):
|
|
self.operations.append((opcode, len(inputs), len(outputs)))
|
|
self.operation_args.extend(inputs + outputs)
|
|
|
|
def add_tensor_sequence(self, jitval, values):
|
|
assert jitval not in self.tensor_sequences
|
|
self.tensor_sequences[jitval] = values
|
|
|
|
def add_constant_value(self, jitval, ctype, value):
|
|
assert jitval not in self.constants
|
|
self.constants[jitval] = (ctype, value)
|
|
|
|
def get_constant_value(self, jitval, typekind=None):
|
|
record = self.constants.get(jitval)
|
|
if record is None:
|
|
raise Exception(f"Could not find constant value for '{jitval!r}'.")
|
|
ctype, _ = record
|
|
if typekind is not None and ctype.kind() != typekind:
|
|
raise Exception(
|
|
f"Expected constant value of type {typekind}, but got {ctype.kind()} for value '{jitval!r}'"
|
|
)
|
|
return record
|
|
|
|
def operand_to_template_torchscript(self, op_id, oper, shape=None):
|
|
"""Return a TorchScript expression to build a template for a given operand."""
|
|
if shape is None:
|
|
shape = oper.shape
|
|
else:
|
|
assert len(shape) == len(oper.shape)
|
|
|
|
shape_parts = ["("]
|
|
for d, s in enumerate(shape):
|
|
if s > 0:
|
|
# Fixed shape dimension: just add the value.
|
|
shape_parts.append(str(s))
|
|
elif s == 0:
|
|
# Load time flexible shape dimension: it should have been computed in a variable.
|
|
shape_parts.append(flex_name(op_id, d))
|
|
elif s == -1:
|
|
# Runtime flexible shape
|
|
shape_parts.append("0")
|
|
else:
|
|
raise Exception("Unknown dim value, dimensions should be >= -1")
|
|
shape_parts.append(",")
|
|
shape_parts.append(")")
|
|
shape_code = "".join(shape_parts)
|
|
if oper.op_type == NNAPI_OperandCode.TENSOR_FLOAT32:
|
|
return f"torch.zeros({shape_code}, dtype=torch.float32)"
|
|
elif oper.op_type == NNAPI_OperandCode.TENSOR_INT32:
|
|
return f"torch.zeros({shape_code}, dtype=torch.int32)"
|
|
elif oper.op_type == NNAPI_OperandCode.TENSOR_QUANT8_ASYMM:
|
|
return (
|
|
f"torch.quantize_per_tensor("
|
|
f"torch.zeros(1), scale={oper.scale}, zero_point={oper.zero_point}, dtype=torch.quint8)"
|
|
f".expand({shape_code}).contiguous()"
|
|
)
|
|
elif oper.op_type in (
|
|
NNAPI_OperandCode.TENSOR_QUANT16_ASYMM,
|
|
NNAPI_OperandCode.TENSOR_QUANT16_SYMM,
|
|
):
|
|
if self.use_int16_for_qint16:
|
|
return f"torch.zeros({shape_code}, dtype=torch.int16)"
|
|
else:
|
|
raise Exception(
|
|
"`int16` isn't supported. If you're trying to represent NNAPI"
|
|
" qint16 with Pytorch int16, set `use_int16_for_qint16 = True`"
|
|
)
|
|
|
|
raise Exception(f"Unsupported output operand type: {oper.op_type}")
|
|
|
|
def forward_operand_shape(self, out_op_id, out_dim, in_op_id, in_dim):
|
|
self.compute_operand_shape(out_op_id, out_dim, flex_name(in_op_id, in_dim))
|
|
|
|
def compute_operand_shape(self, op_id, dim, expr):
|
|
self.flexible_shape_computation_lines.append(
|
|
f"{flex_name(op_id, dim)} = {expr}"
|
|
)
|
|
|
|
def transpose_to_nhwc(self, in_id, oper):
|
|
if oper.shape[2:] != (1, 1):
|
|
raise Exception("Automatic transpose only supported for H,W == 1,1")
|
|
|
|
out_oper = oper._replace(dim_order=DimOrder.CHANNELS_LAST)
|
|
|
|
inputs = [None] * 2
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_vector([0, 2, 3, 1])
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_anonymous_tensor_operand(out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.TRANSPOSE, inputs, outputs)
|
|
|
|
return outputs[0], out_oper
|
|
|
|
# Transpose inputs as necessary to allow broadcasting.
|
|
def transpose_for_broadcast(self, in0_id, in0_oper, in1_id, in1_oper):
|
|
if in0_oper.dim_order == in1_oper.dim_order:
|
|
return in0_id, in0_oper, in1_id, in1_oper
|
|
|
|
# Assume NHWC is preferred if there is a mismatch.
|
|
orders = (in0_oper.dim_order, in1_oper.dim_order)
|
|
if orders == (DimOrder.PRESUMED_CONTIGUOUS, DimOrder.CHANNELS_LAST):
|
|
return self.transpose_to_nhwc(in0_id, in0_oper) + (in1_id, in1_oper)
|
|
if orders == (DimOrder.CHANNELS_LAST, DimOrder.PRESUMED_CONTIGUOUS):
|
|
return (in0_id, in0_oper) + self.transpose_to_nhwc(in1_id, in1_oper)
|
|
|
|
raise Exception(
|
|
f"Automatic transpose not supported for dim_orders: {in0_oper.dim_order!r}, {in1_oper.dim_order!r}"
|
|
)
|
|
|
|
def get_size_arg(self, jitval):
|
|
ctype, value = self.get_constant_value(jitval)
|
|
if ctype.kind() == "ListType":
|
|
assert ctype.getElementType().kind() == "IntType"
|
|
return value
|
|
raise Exception(f"Can't handle size arg of type '{ctype!r}' for '{jitval!r}'")
|
|
|
|
def get_conv_pool_args_2d_from_pack(self, kernel_size, packed_config):
|
|
pc = [i.item() for i in packed_config]
|
|
assert pc[0] == 2
|
|
strides = [pc[1], pc[2]]
|
|
paddings = [pc[3], pc[4]]
|
|
dilations = [pc[5], pc[6]]
|
|
output_padding = [pc[7], pc[8]]
|
|
group_num = pc[9]
|
|
|
|
assert len(pc) == 11
|
|
assert output_padding == [0, 0]
|
|
|
|
return self.get_conv_pool_args_2d_common(
|
|
kernel_size, strides, paddings, dilations, group_num
|
|
)
|
|
|
|
def get_conv_pool_args_2d_from_jit(
|
|
self, kernel_size, stride, padding, dilation=None, group=None
|
|
):
|
|
strides = self.get_size_arg(stride)
|
|
paddings = self.get_size_arg(padding)
|
|
if dilation is None:
|
|
dilations = [1, 1]
|
|
else:
|
|
dilations = self.get_size_arg(dilation)
|
|
if group is not None:
|
|
_, group_num = self.get_constant_value(group, "IntType")
|
|
else:
|
|
group_num = None
|
|
return self.get_conv_pool_args_2d_common(
|
|
kernel_size, strides, paddings, dilations, group_num
|
|
)
|
|
|
|
def get_conv_pool_args_2d_common(
|
|
self, kernel_size, strides, paddings, dilations, group_num
|
|
):
|
|
kernels = list(kernel_size)
|
|
|
|
assert len(kernels) == 2
|
|
assert len(strides) == 2
|
|
assert len(paddings) == 2
|
|
assert len(dilations) == 2
|
|
|
|
# NNAPI uses 4 values for padding.
|
|
ph, pw = paddings
|
|
real_paddings = [ph, ph, pw, pw]
|
|
|
|
return ConvPoolArgs2d(
|
|
*(kernels + strides + real_paddings + dilations + [group_num])
|
|
)
|
|
|
|
def serialize_model(self, model, inputs, return_shapes=None):
|
|
self.add_immediate_bool_scalar(False)
|
|
self.add_immediate_bool_scalar(True)
|
|
|
|
inp_dim_orders = []
|
|
out_dim_orders = []
|
|
|
|
self_jitval = next(model.graph.inputs())
|
|
self.add_constant_value(self_jitval, self_jitval.type(), model)
|
|
|
|
for arg_idx, (input_value, input_tensor) in enumerate(
|
|
zip(list(model.graph.inputs())[1:], inputs)
|
|
):
|
|
op_id = self.add_tensor_operand_for_input(
|
|
arg_idx, input_value, input_tensor
|
|
)
|
|
inp_dim_orders.append(self.operands[op_id].dim_order.value)
|
|
|
|
for idx, node in enumerate(model.graph.nodes()):
|
|
LOG.debug("Processing node #%d: %r", idx, node)
|
|
self.add_node(node)
|
|
|
|
retn = model.graph.return_node()
|
|
assert retn.inputsSize() == 1
|
|
assert retn.outputsSize() == 0
|
|
retn_input = retn.inputsAt(0)
|
|
template_return_lines = ["return ["]
|
|
if retn_input.type().kind() == "TensorType":
|
|
return_values = [retn_input]
|
|
retval_count = -1
|
|
elif retn_input.type().kind() == "TupleType":
|
|
return_values = self.tensor_sequences[retn_input]
|
|
retval_count = len(return_values)
|
|
else:
|
|
raise Exception(f"Unsupported return type: {retn_input.type()}")
|
|
|
|
if return_shapes is not None:
|
|
assert len(return_shapes) == len(return_values)
|
|
for i, v in enumerate(return_values):
|
|
op_id = self.jitval_operand_map[v]
|
|
self.outputs.append(op_id)
|
|
out_dim_orders.append(self.operands[op_id].dim_order.value)
|
|
shape = return_shapes[i] if return_shapes else None
|
|
template_return_lines.append(
|
|
self.operand_to_template_torchscript(op_id, self.operands[op_id], shape)
|
|
+ ","
|
|
)
|
|
template_return_lines.append("]")
|
|
|
|
model = []
|
|
|
|
version = 1
|
|
header = struct.pack(
|
|
"iiiiii",
|
|
version,
|
|
len(self.operands),
|
|
len(self.values),
|
|
len(self.operations),
|
|
len(self.inputs),
|
|
len(self.outputs),
|
|
)
|
|
model.append(header)
|
|
|
|
serialized_values, serialized_value_data = self.serialize_values()
|
|
|
|
model.extend(
|
|
struct.pack("iifi", t, len(d), s, z) for (t, d, _m, s, z) in self.operands
|
|
)
|
|
model.extend(serialized_values)
|
|
model.extend(struct.pack("iii", *x) for x in self.operations)
|
|
|
|
# Compact the model so we can get its length so far.
|
|
model = [b"".join(model)]
|
|
model_offset = len(model[0])
|
|
# Model offset is the index into the model (in 32-bit words, not bytes)
|
|
# of the next dimension we're about to serialize. If it's 0,
|
|
# generate code to mutate it before passing to NNAPI.
|
|
assert model_offset % 4 == 0
|
|
model_offset = int(model_offset / 4)
|
|
|
|
for op_id, (_, dims, dim_order, _, _) in enumerate(self.operands):
|
|
shape = fix_shape(dims, dim_order)
|
|
for d, s in enumerate(shape):
|
|
if s == 0:
|
|
pt_d = reverse_map_dim(dim_order, d)
|
|
self.flexible_shape_computation_lines.append(
|
|
f"ser_model[{model_offset}] = {flex_name(op_id, pt_d)}"
|
|
)
|
|
model_offset += 1
|
|
|
|
# convert runtime flex shape from -1 to 0
|
|
shape = tuple(d if d != -1 else 0 for d in shape)
|
|
model.append(self.serialize_ints(shape))
|
|
|
|
model.extend(serialized_value_data)
|
|
model.append(self.serialize_ints(self.operation_args))
|
|
model.append(self.serialize_ints(self.inputs))
|
|
model.append(self.serialize_ints(self.outputs))
|
|
|
|
self.flexible_shape_computation_lines.extend(template_return_lines)
|
|
|
|
return (
|
|
array.array("i", b"".join(model)),
|
|
self.used_weights,
|
|
inp_dim_orders,
|
|
out_dim_orders,
|
|
self.flexible_shape_computation_lines,
|
|
retval_count,
|
|
)
|
|
|
|
def serialize_values(self):
|
|
serialized_values = []
|
|
serialized_value_data = []
|
|
assert len(self.values) == len(self.value_data)
|
|
for (op_index, source_type), data in zip(self.values, self.value_data):
|
|
source_length = len(data)
|
|
|
|
# Pad with 0 bytes out to a multiple of 4 for alignment.
|
|
physical_length = ((source_length - 1) | 0x3) + 1
|
|
padded_data = data + (b"\0" * (physical_length - source_length))
|
|
|
|
serialized_values.append(
|
|
struct.pack("iii", op_index, source_type, source_length)
|
|
)
|
|
serialized_value_data.append(padded_data)
|
|
|
|
return serialized_values, serialized_value_data
|
|
|
|
@staticmethod
|
|
def serialize_ints(ints):
|
|
return array.array("i", ints).tobytes()
|
|
|
|
ADDER_MAP = {
|
|
"prim::GetAttr": lambda self, node: self.add_getattr(node),
|
|
"prim::Constant": lambda self, node: self.add_constant_node(node),
|
|
"prim::ListConstruct": lambda self, node: self.add_list_construct(node),
|
|
"prim::TupleConstruct": lambda self, node: self.add_tuple_construct(node),
|
|
"aten::unsqueeze": lambda self, node: self.add_unsqueeze(node),
|
|
"aten::to": lambda self, node: self.add_to(node),
|
|
"aten::detach": lambda self, node: self._identity(node),
|
|
"aten::reshape": lambda self, node: self.add_reshape(node),
|
|
"aten::flatten": lambda self, node: self.add_flatten(node),
|
|
"aten::slice": lambda self, node: self.add_slice(node),
|
|
"aten::size": lambda self, node: self.add_size(node),
|
|
"aten::cat": lambda self, node: self.add_cat(node),
|
|
"aten::mean": lambda self, node: self.add_mean(node),
|
|
"aten::quantize_per_tensor": lambda self, node: self.add_quantize(node),
|
|
"aten::dequantize": lambda self, node: self.add_dequantize(node),
|
|
"aten::add": lambda self, node: self.add_add_sub_op(
|
|
node, NNAPI_OperationCode.ADD, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"aten::sub": lambda self, node: self.add_add_sub_op(
|
|
node, NNAPI_OperationCode.SUB, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"aten::mul": lambda self, node: self.add_pointwise_simple_binary_broadcast_op(
|
|
node, NNAPI_OperationCode.MUL, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"aten::div": lambda self, node: self.add_pointwise_simple_binary_broadcast_op(
|
|
node, NNAPI_OperationCode.DIV, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"aten::relu": lambda self, node: self.add_pointwise_simple_unary_op(
|
|
node, NNAPI_OperationCode.RELU
|
|
),
|
|
"aten::sigmoid": lambda self, node: self.add_pointwise_simple_unary_op(
|
|
node, NNAPI_OperationCode.LOGISTIC
|
|
),
|
|
"aten::softmax": lambda self, node: self.add_softmax(node),
|
|
"aten::hardtanh": lambda self, node: self.add_hardtanh(node),
|
|
"aten::avg_pool2d": lambda self, node: self.add_avg_pool2d(node),
|
|
"aten::max_pool2d": lambda self, node: self.add_pool2d_node(
|
|
node, NNAPI_OperationCode.MAX_POOL_2D
|
|
),
|
|
"aten::adaptive_avg_pool2d": lambda self, node: self.add_adaptive_avg_pool2d(
|
|
node
|
|
),
|
|
"aten::upsample_nearest2d": lambda self, node: self.add_upsample_nearest2d(
|
|
node
|
|
),
|
|
"aten::prelu": lambda self, node: self.add_prelu_op(node),
|
|
"aten::addmm": lambda self, node: self.add_addmm(node),
|
|
"aten::linear": lambda self, node: self.add_linear(node),
|
|
"aten::_convolution": lambda self, node: self.add_conv_underscore(node),
|
|
"aten::conv2d": lambda self, node: self.add_conv2d(node),
|
|
"aten::log_softmax": lambda self, node: self.add_log_softmax(node),
|
|
"quantized::linear": lambda self, node: self.add_qlinear(node),
|
|
"quantized::conv2d": lambda self, node: self.add_qconv2d(
|
|
node, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"quantized::conv2d_relu": lambda self, node: self.add_qconv2d(
|
|
node, NNAPI_FuseCode.FUSED_RELU
|
|
),
|
|
"quantized::conv_transpose2d": lambda self, node: self.add_qconv2d(
|
|
node, NNAPI_FuseCode.FUSED_NONE, transpose=True
|
|
),
|
|
"quantized::add": lambda self, node: self.add_qadd(
|
|
node, NNAPI_OperationCode.ADD, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
"quantized::add_relu": lambda self, node: self.add_qadd(
|
|
node, NNAPI_OperationCode.ADD, NNAPI_FuseCode.FUSED_RELU
|
|
),
|
|
"quantized::mul": lambda self, node: self.add_qadd(
|
|
node, NNAPI_OperationCode.MUL, NNAPI_FuseCode.FUSED_NONE
|
|
),
|
|
}
|
|
|
|
def add_node(self, node):
|
|
adder = self.ADDER_MAP.get(node.kind())
|
|
if not adder:
|
|
raise Exception(f"Unsupported node kind ({node.kind()!r}) in node {node!r}")
|
|
adder(self, node)
|
|
|
|
def _identity(self, node):
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
jitval = node.outputsAt(0)
|
|
self.jitval_operand_map[jitval] = in_id
|
|
|
|
def add_getattr(self, node):
|
|
assert node.inputsSize() == 1
|
|
assert node.outputsSize() == 1
|
|
obj_ctype, obj = self.get_constant_value(node.inputsAt(0))
|
|
assert str(obj_ctype).startswith("__torch__.")
|
|
name = node.s("name")
|
|
value = getattr(obj, name)
|
|
output = node.outputsAt(0)
|
|
ctype = output.type()
|
|
self.add_constant_value(output, ctype, value)
|
|
|
|
def add_constant_node(self, node):
|
|
assert node.inputsSize() == 0
|
|
assert node.outputsSize() == 1
|
|
output = node.outputsAt(0)
|
|
ctype = output.type()
|
|
value = output.toIValue()
|
|
self.add_constant_value(output, ctype, value)
|
|
|
|
def add_list_construct(self, node):
|
|
assert node.outputsSize() == 1
|
|
output = node.outputsAt(0)
|
|
ctype = output.type()
|
|
const_vals: Optional[List] = []
|
|
tensors: Optional[List] = []
|
|
for inp in node.inputs():
|
|
if const_vals is not None and inp in self.constants:
|
|
_, val = self.get_constant_value(inp)
|
|
const_vals.append(val)
|
|
else:
|
|
const_vals = None
|
|
if tensors is not None and inp.type().kind() == "TensorType":
|
|
tensors.append(inp)
|
|
else:
|
|
tensors = None
|
|
|
|
if const_vals is not None:
|
|
# NOTE: Now that TorchScript supports list constants,
|
|
# this code path might not be used anymore.
|
|
self.add_constant_value(output, ctype, const_vals)
|
|
if tensors is not None:
|
|
self.add_tensor_sequence(output, tensors)
|
|
if const_vals is None and tensors is None:
|
|
raise Exception(
|
|
f"Unable to handle ListConstruct node. Neither all constants nor all tensors. {node!r}"
|
|
)
|
|
|
|
def add_tuple_construct(self, node):
|
|
assert node.outputsSize() == 1
|
|
output = node.outputsAt(0)
|
|
values = list(node.inputs())
|
|
self.add_tensor_sequence(output, values)
|
|
|
|
def add_unsqueeze(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
|
|
_, dim = self.get_constant_value(node.inputsAt(1), "IntType")
|
|
assert in_oper.dim_order == DimOrder.PRESUMED_CONTIGUOUS
|
|
|
|
real_dim = dim if dim >= 0 else dim + len(in_oper.shape) + 1
|
|
out_shape_list = list(in_oper.shape)
|
|
out_shape_list.insert(real_dim, 1)
|
|
out_shape = tuple(out_shape_list)
|
|
out_oper = in_oper._replace(shape=out_shape)
|
|
|
|
inputs = [None] * 2
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_scalar(dim)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.EXPAND_DIMS, inputs, outputs)
|
|
|
|
def add_to(self, node):
|
|
# Handle to("cpu") / to("gpu") case
|
|
self._identity(node)
|
|
|
|
def add_reshape(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
|
|
shape_ctype, shape = self.get_constant_value(node.inputsAt(1))
|
|
assert shape_ctype.kind() == "ListType"
|
|
assert shape_ctype.getElementType().kind() == "IntType"
|
|
is_trivial_reshape = len(shape) == 2 and shape[1] == -1
|
|
|
|
if in_oper.dim_order != DimOrder.PRESUMED_CONTIGUOUS and not is_trivial_reshape:
|
|
raise Exception(
|
|
"Currently, reshape is only supported on NHWC tensors if the target size is [X, -1]."
|
|
)
|
|
|
|
# Bit of a hack here. Use a real tensor to infer the output shape.
|
|
out_shape = torch.zeros(1).expand(in_oper.shape).reshape(shape).shape
|
|
out_oper = in_oper._replace(
|
|
shape=out_shape, dim_order=DimOrder.PRESUMED_CONTIGUOUS
|
|
)
|
|
|
|
inputs = [None] * 2
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_vector(shape)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.RESHAPE, inputs, outputs)
|
|
|
|
def add_flatten(self, node):
|
|
assert node.inputsSize() == 3
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
|
|
start_ctype, start_dim = self.get_constant_value(node.inputsAt(1), "IntType")
|
|
end_ctype, end_dim = self.get_constant_value(node.inputsAt(2), "IntType")
|
|
|
|
# channels last with channels == 1 or (height & width both 1)
|
|
is_trivial_flatten = len(in_oper.shape) == 4 and (
|
|
in_oper.shape[1] == 1 or (in_oper.shape[2] == 1 and in_oper.shape[3] == 1)
|
|
)
|
|
if in_oper.dim_order != DimOrder.PRESUMED_CONTIGUOUS and not is_trivial_flatten:
|
|
raise Exception(
|
|
"Currently, flatten is not supported on NHWC tensors unless C=1 or H=W=1"
|
|
)
|
|
|
|
if start_dim < 0:
|
|
start_dim += len(in_oper.shape)
|
|
if end_dim < 0:
|
|
end_dim += len(in_oper.shape)
|
|
|
|
out_shape = (
|
|
in_oper.shape[:start_dim]
|
|
+ (functools.reduce(operator.mul, in_oper.shape[start_dim : end_dim + 1]),)
|
|
+ in_oper.shape[end_dim + 1 :]
|
|
)
|
|
|
|
if any(dim == 0 for dim in in_oper.shape[start_dim : end_dim + 1]):
|
|
raise Exception("Flattening flexible dims is not supported yet")
|
|
non_flattened_dims = in_oper.shape[:start_dim] + in_oper.shape[end_dim + 1 :]
|
|
if non_flattened_dims.count(0) > 1:
|
|
raise Exception("Only 1 dim can be flexible")
|
|
|
|
out_oper = in_oper._replace(
|
|
shape=out_shape, dim_order=DimOrder.PRESUMED_CONTIGUOUS
|
|
)
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
for idx, dim in enumerate(out_shape):
|
|
if dim == 0:
|
|
self.forward_operand_shape(out_id, idx, in_id, in_oper.shape.index(0))
|
|
|
|
inputs_1 = tuple(dim if dim != 0 else -1 for dim in out_shape)
|
|
inputs = [None] * 2
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_vector(inputs_1)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.RESHAPE, inputs, outputs)
|
|
|
|
def add_slice(self, node):
|
|
assert node.inputsSize() == 5
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
_, dim_value = self.get_constant_value(node.inputsAt(1))
|
|
_, start_value = self.get_constant_value(node.inputsAt(2))
|
|
_, stop_value = self.get_constant_value(node.inputsAt(3))
|
|
_, step_value = self.get_constant_value(node.inputsAt(4))
|
|
|
|
if start_value is None:
|
|
start_value = 0
|
|
if stop_value is None:
|
|
stop_value = sys.maxsize
|
|
|
|
if start_value < 0:
|
|
start_value += in_oper.shape[dim_value]
|
|
elif start_value == sys.maxsize:
|
|
start_value = 0
|
|
|
|
if start_value == 0 and stop_value == sys.maxsize:
|
|
self._identity(node)
|
|
return
|
|
|
|
if in_oper.shape[dim_value] == 0:
|
|
raise Exception("Unable to slice with flexible shape")
|
|
|
|
if stop_value < 0:
|
|
stop_value += in_oper.shape[dim_value]
|
|
elif stop_value == sys.maxsize:
|
|
stop_value = in_oper.shape[dim_value]
|
|
|
|
if start_value >= stop_value:
|
|
raise Exception("Slice start value should be less than stop value")
|
|
|
|
out_len = (stop_value - start_value) // step_value
|
|
out_shape = tuple(
|
|
out_len if i == dim_value else dim for i, dim in enumerate(in_oper.shape)
|
|
)
|
|
out_id = self.add_tensor_operand(
|
|
node.outputsAt(0), in_oper._replace(shape=out_shape)
|
|
)
|
|
|
|
# flex inputs
|
|
end_mask = 0
|
|
for idx, dim in enumerate(out_shape):
|
|
if dim == 0:
|
|
self.forward_operand_shape(out_id, idx, in_id, idx)
|
|
end_mask |= 1 << idx
|
|
|
|
inputs = [None] * 7
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_vector(
|
|
[start_value if i == dim_value else 0 for i in range(len(in_oper.shape))]
|
|
)
|
|
inputs[2] = self.add_immediate_int_vector(
|
|
[
|
|
stop_value if i == dim_value else dim
|
|
for i, dim in enumerate(in_oper.shape)
|
|
]
|
|
)
|
|
inputs[3] = self.add_immediate_int_vector(
|
|
[step_value if i == dim_value else 1 for i in range(len(in_oper.shape))]
|
|
)
|
|
inputs[4] = self.add_immediate_int_scalar(0) # begin mask
|
|
inputs[5] = self.add_immediate_int_scalar(end_mask)
|
|
inputs[6] = self.add_immediate_int_scalar(0) # shrink axis mas
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.STRIDED_SLICE, inputs, outputs)
|
|
|
|
def add_size(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
_, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
_, value = self.constants[node.inputsAt(1)]
|
|
res = in_oper.shape[value]
|
|
output = node.outputsAt(0)
|
|
self.add_constant_value(output, output.type(), res)
|
|
|
|
def add_cat(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
tensors = self.tensor_sequences[node.inputsAt(0)]
|
|
_, dim = self.get_constant_value(node.inputsAt(1), "IntType")
|
|
|
|
assert len(tensors) > 0
|
|
in_ids = []
|
|
out_oper = None
|
|
out_dim_size = 0
|
|
for inp in tensors:
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(inp)
|
|
if out_oper is None:
|
|
out_shape = change_element(in_oper.shape, dim, -1)
|
|
out_oper = in_oper._replace(shape=out_shape)
|
|
assert in_oper.op_type == out_oper.op_type
|
|
assert in_oper.dim_order == out_oper.dim_order
|
|
assert change_element(in_oper.shape, dim, -1) == change_element(
|
|
out_oper.shape, dim, -1
|
|
)
|
|
# TODO: Possibly check scale and zero point.
|
|
in_ids.append(in_id)
|
|
# TODO: Possibly support variable-sized inputs.
|
|
out_dim_size += in_oper.shape[dim]
|
|
|
|
assert out_oper is not None
|
|
out_oper = out_oper._replace(
|
|
shape=change_element(out_oper.shape, dim, out_dim_size)
|
|
)
|
|
|
|
if in_oper.dim_order == DimOrder.CHANNELS_LAST: # type: ignore[possibly-undefined]
|
|
assert len(out_oper.shape) == 4
|
|
nnapi_dim = [0, 3, 1, 2][dim]
|
|
else:
|
|
nnapi_dim = dim
|
|
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
for idx, d in enumerate(out_oper.shape):
|
|
if d == 0:
|
|
if idx == dim:
|
|
shape = " + ".join(flex_name(ip_id, dim) for ip_id in in_ids)
|
|
self.compute_operand_shape(out_id, idx, shape)
|
|
else:
|
|
self.forward_operand_shape(out_id, idx, in_ids[0], idx)
|
|
|
|
inputs = in_ids + [self.add_immediate_int_scalar(nnapi_dim)]
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.CONCATENATION, inputs, outputs)
|
|
|
|
def add_mean(self, node):
|
|
assert node.inputsSize() == 4
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
dim_ctype, dim = self.get_constant_value(node.inputsAt(1))
|
|
assert dim_ctype.kind() == "ListType"
|
|
assert dim_ctype.getElementType().kind() == "IntType"
|
|
_, keep_dim = self.get_constant_value(node.inputsAt(2), "BoolType")
|
|
# Expect None for dtype
|
|
self.get_constant_value(node.inputsAt(3), "NoneType")
|
|
|
|
if in_oper.dim_order == DimOrder.CHANNELS_LAST:
|
|
assert len(in_oper.shape) == 4
|
|
nnapi_dim = [[0, 3, 1, 2][d] for d in dim]
|
|
else:
|
|
nnapi_dim = dim
|
|
|
|
collapsed_dims = set()
|
|
for d in dim:
|
|
if d < 0:
|
|
d += len(in_oper.shape)
|
|
collapsed_dims.add(d)
|
|
|
|
if in_oper.dim_order == DimOrder.CHANNELS_LAST and not keep_dim:
|
|
assert collapsed_dims.issuperset({2, 3})
|
|
out_dim_order = DimOrder.PRESUMED_CONTIGUOUS
|
|
else:
|
|
out_dim_order = in_oper.dim_order
|
|
|
|
out_shape = []
|
|
for i, s in enumerate(in_oper.shape):
|
|
if i not in collapsed_dims:
|
|
out_shape.append(s)
|
|
elif keep_dim:
|
|
out_shape.append(1)
|
|
|
|
out_oper = in_oper._replace(shape=out_shape, dim_order=out_dim_order)
|
|
|
|
inputs = [None] * 3
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_int_vector(nnapi_dim)
|
|
inputs[2] = self.add_immediate_int_scalar(keep_dim)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.MEAN, inputs, outputs)
|
|
|
|
def add_quantize(self, node):
|
|
assert node.inputsSize() == 4
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
if in_oper.dim_order != DimOrder.CHANNELS_LAST:
|
|
raise Exception(
|
|
"Most hardware backends prefer NHWC quantized tensors. "
|
|
"Try setting `t.nnapi_nhwc = True` on your tensor inputs. "
|
|
)
|
|
_, scale = self.get_constant_value(node.inputsAt(1), "FloatType")
|
|
_, zero_point = self.get_constant_value(node.inputsAt(2), "IntType")
|
|
_, scalar_type = self.get_constant_value(node.inputsAt(3), "IntType")
|
|
if scalar_type != TorchScalarTypes.QUINT8.value:
|
|
raise Exception(
|
|
"PyTorch NNAPI export only supports quantized tensors "
|
|
"with the quint8 dtype."
|
|
)
|
|
op_type = NNAPI_OperandCode.TENSOR_QUANT8_ASYMM
|
|
|
|
out_oper = in_oper._replace(
|
|
op_type=op_type,
|
|
scale=scale,
|
|
zero_point=zero_point,
|
|
)
|
|
|
|
inputs = [None] * 1
|
|
inputs[0] = in_id
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.QUANTIZE, inputs, outputs)
|
|
|
|
def add_dequantize(self, node):
|
|
assert node.inputsSize() == 1
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
out_oper = in_oper._replace(
|
|
op_type=NNAPI_OperandCode.TENSOR_FLOAT32,
|
|
scale=0.0,
|
|
zero_point=0,
|
|
)
|
|
|
|
inputs = [None] * 1
|
|
inputs[0] = in_id
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.DEQUANTIZE, inputs, outputs)
|
|
|
|
def add_pointwise_simple_unary_op(self, node, opcode):
|
|
assert node.inputsSize() == 1
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
|
|
out_oper = in_oper
|
|
if opcode == NNAPI_OperationCode.LOGISTIC:
|
|
# NNAPI docs: For ANEURALNETWORKS_TENSOR_QUANT8_ASYMM, the scale
|
|
# must be 1.f / 256 and the zeroPoint must be 0.
|
|
# https://fburl.com/h52stoog
|
|
if in_oper.op_type == NNAPI_OperandCode.TENSOR_QUANT8_ASYMM:
|
|
out_oper = in_oper._replace(zero_point=0, scale=1.0 / 256)
|
|
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
for idx, dim in enumerate(in_oper.shape):
|
|
if dim == 0:
|
|
self.forward_operand_shape(out_id, idx, in_id, idx)
|
|
|
|
inputs = [None] * 1
|
|
inputs[0] = in_id
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(opcode, inputs, outputs)
|
|
|
|
def _do_add_binary(self, node, opcode, fuse_code, *, qparams=None): # noqa: D401
|
|
"""Helper for pointwise binary broadcast ops with superfluous extra args."""
|
|
assert node.outputsSize() == 1
|
|
|
|
assert node.inputsAt(0).type().kind() == "TensorType"
|
|
assert node.inputsAt(1).type().kind() == "TensorType"
|
|
|
|
if self.has_operand_for_jitval(node.inputsAt(0)):
|
|
in0_id, in0_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
in1_id, in1_oper = self.get_tensor_operand_or_constant(
|
|
node.inputsAt(1), in0_oper.dim_order
|
|
)
|
|
elif self.has_operand_for_jitval(node.inputsAt(1)):
|
|
in1_id, in1_oper = self.get_tensor_operand_by_jitval(node.inputsAt(1))
|
|
in0_id, in0_oper = self.get_tensor_operand_or_constant(
|
|
node.inputsAt(0), in1_oper.dim_order
|
|
)
|
|
else:
|
|
raise Exception(f"Can't do a NNAPI binary op: {opcode} on two constants")
|
|
|
|
assert in0_oper.op_type == in1_oper.op_type
|
|
in0_id, in0_oper, in1_id, in1_oper = self.transpose_for_broadcast(
|
|
in0_id, in0_oper, in1_id, in1_oper
|
|
)
|
|
# NOTE: PyTorch and NNAPI have the same broadcast semantics.
|
|
out_shape = broadcast_shapes(in0_oper.shape, in1_oper.shape)
|
|
out_oper = in0_oper._replace(shape=out_shape)
|
|
if qparams is not None:
|
|
scale, zp = qparams
|
|
out_oper = out_oper._replace(scale=scale, zero_point=zp)
|
|
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
for idx, (d0, d1) in enumerate(zip(in0_oper.shape, in1_oper.shape)):
|
|
if d0 == 1 and d1 == 0:
|
|
self.forward_operand_shape(out_id, idx, in1_id, idx)
|
|
elif d0 == 0 and d1 == 1:
|
|
self.forward_operand_shape(out_id, idx, in0_id, idx)
|
|
elif d0 == 0 and d1 == 0:
|
|
self.flexible_shape_computation_lines.append(
|
|
f"assert {flex_name(in0_id, idx)} == {flex_name(in1_id, idx)}"
|
|
)
|
|
self.forward_operand_shape(out_id, idx, in0_id, idx)
|
|
|
|
inputs = [None] * 3
|
|
inputs[0] = in0_id
|
|
inputs[1] = in1_id
|
|
inputs[2] = self.add_immediate_int_scalar(fuse_code)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(opcode, inputs, outputs)
|
|
|
|
def add_pointwise_simple_binary_broadcast_op(self, node, opcode, fuse_code):
|
|
assert node.inputsSize() == 2
|
|
self._do_add_binary(node, opcode, fuse_code)
|
|
|
|
def add_add_sub_op(self, node, opcode, fuse_code):
|
|
assert node.inputsSize() == 3
|
|
|
|
_, alpha = self.get_constant_value(node.inputsAt(2), "IntType")
|
|
if alpha != 1:
|
|
raise Exception("NNAPI does not support add/sub with alpha.")
|
|
|
|
self._do_add_binary(node, opcode, fuse_code)
|
|
|
|
def add_qadd(self, node, opcode, fuse_code):
|
|
assert node.inputsSize() == 4
|
|
|
|
_, scale = self.get_constant_value(node.inputsAt(2), "FloatType")
|
|
_, zero_point = self.get_constant_value(node.inputsAt(3), "IntType")
|
|
|
|
self._do_add_binary(node, opcode, fuse_code, qparams=(scale, zero_point))
|
|
|
|
def add_softmax(self, node):
|
|
assert node.inputsSize() == 3
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
|
|
_, softmax_dim = self.get_constant_value(node.inputsAt(1), "IntType")
|
|
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), in_oper)
|
|
for dim, size in enumerate(in_oper.shape):
|
|
if size == 0:
|
|
self.forward_operand_shape(out_id, dim, in_id, dim)
|
|
|
|
inputs = [None] * 3
|
|
inputs[0] = in_id
|
|
inputs[1] = self.add_immediate_float_scalar(
|
|
1.0
|
|
) # positive scaling factor of exponent, beta
|
|
inputs[2] = self.add_immediate_int_scalar(softmax_dim)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.SOFTMAX, inputs, outputs)
|
|
|
|
def add_hardtanh(self, node):
|
|
assert node.inputsSize() == 3
|
|
assert node.outputsSize() == 1
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval_fixed_size(node.inputsAt(0))
|
|
_, min_val = self.get_constant_value(node.inputsAt(1), "FloatType")
|
|
_, max_val = self.get_constant_value(node.inputsAt(2), "FloatType")
|
|
|
|
op_map = {
|
|
(-1, 1): NNAPI_OperationCode.RELU1,
|
|
(0, 6): NNAPI_OperationCode.RELU6, # noqa: E201
|
|
}
|
|
|
|
opcode = op_map.get((min_val, max_val))
|
|
if opcode is None:
|
|
raise Exception("NNAPI only supports hardtanh with args (-1, 1) or (0, 6).")
|
|
|
|
inputs = [None] * 1
|
|
inputs[0] = in_id
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), in_oper)
|
|
|
|
self.add_operation(opcode, inputs, outputs)
|
|
|
|
def add_prelu_op(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
assert node.inputsAt(0).type().kind() == "TensorType"
|
|
assert node.inputsAt(1).type().kind() == "TensorType"
|
|
|
|
in_id, in_oper = self.get_tensor_operand_by_jitval(node.inputsAt(0))
|
|
w_id, w_oper = self.get_tensor_operand_for_weight(node.inputsAt(1))
|
|
assert len(w_oper.shape) == 1
|
|
assert w_oper.shape[0] > 0
|
|
if w_oper.shape[0] > 1:
|
|
if in_oper.use_nchw():
|
|
# TODO: Support this by adding trailing 1 dims.
|
|
raise Exception(
|
|
"Per-channel PReLU only supports channels_last right now."
|
|
)
|
|
|
|
out_id = self.add_tensor_operand(node.outputsAt(0), in_oper)
|
|
for dim, size in enumerate(in_oper.shape):
|
|
if size > 0:
|
|
pass
|
|
elif dim <= 1:
|
|
raise Exception("PReLU requires fixed size for dim 0 and dim 1.")
|
|
else:
|
|
self.forward_operand_shape(out_id, dim, in_id, dim)
|
|
|
|
inputs = [None] * 2
|
|
inputs[0] = in_id
|
|
inputs[1] = w_id
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.PRELU, inputs, outputs)
|
|
|
|
def add_pool2d_node(self, node, opcode):
|
|
assert node.inputsSize() == 6
|
|
assert node.outputsSize() == 1
|
|
image, kernel, stride, padding, dilation, ceil_mode = node.inputs()
|
|
|
|
stride = stride or kernel
|
|
|
|
# TODO: Validate ceil_mode semantics.
|
|
|
|
args = self.get_conv_pool_args_2d_from_jit(
|
|
self.get_size_arg(kernel), stride, padding, dilation
|
|
)
|
|
if args.dilation_h != 1 or args.dilation_w != 1:
|
|
raise Exception("NNAPI does not support dilated pooling.")
|
|
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval_fixed_size(image)
|
|
assert len(image_oper.shape) == 4
|
|
|
|
out_shape = get_conv_pool_shape(
|
|
image_oper.shape, args, image_oper.shape[1], False
|
|
)
|
|
use_nchw = image_oper.use_nchw()
|
|
|
|
inputs = [None] * 11
|
|
inputs[0] = image_id
|
|
inputs[1] = self.add_immediate_int_scalar(args.pad_l)
|
|
inputs[2] = self.add_immediate_int_scalar(args.pad_r)
|
|
inputs[3] = self.add_immediate_int_scalar(args.pad_t)
|
|
inputs[4] = self.add_immediate_int_scalar(args.pad_b)
|
|
inputs[5] = self.add_immediate_int_scalar(args.stride_w)
|
|
inputs[6] = self.add_immediate_int_scalar(args.stride_h)
|
|
inputs[7] = self.add_immediate_int_scalar(args.kernel_w)
|
|
inputs[8] = self.add_immediate_int_scalar(args.kernel_h)
|
|
inputs[9] = self.add_immediate_int_scalar(NNAPI_FuseCode.FUSED_NONE)
|
|
inputs[10] = self.add_immediate_bool_scalar(use_nchw)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(
|
|
node.outputsAt(0), image_oper._replace(shape=out_shape)
|
|
)
|
|
|
|
self.add_operation(opcode, inputs, outputs)
|
|
|
|
def add_avg_pool2d(self, node):
|
|
assert node.inputsSize() == 7
|
|
assert node.outputsSize() == 1
|
|
(
|
|
image,
|
|
kernel,
|
|
stride,
|
|
padding,
|
|
ceil_mode,
|
|
count_include_pad,
|
|
divisor_override,
|
|
) = node.inputs()
|
|
|
|
_, count_include_pad_value = self.get_constant_value(count_include_pad)
|
|
_, divisor_override_value = self.get_constant_value(divisor_override)
|
|
if not count_include_pad_value or divisor_override_value:
|
|
raise Exception(
|
|
"NNAPI doesn't support count_include_pad=False or divisor_override"
|
|
)
|
|
|
|
args = self.get_conv_pool_args_2d_from_jit(
|
|
self.get_size_arg(kernel), stride, padding
|
|
)
|
|
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval(image)
|
|
assert len(image_oper.shape) == 4
|
|
|
|
out_shape = get_conv_pool_shape(
|
|
image_oper.shape, args, image_oper.shape[1], False
|
|
)
|
|
use_nchw = image_oper.use_nchw()
|
|
|
|
inputs = [None] * 11
|
|
inputs[0] = image_id
|
|
inputs[1] = self.add_immediate_int_scalar(args.pad_l)
|
|
inputs[2] = self.add_immediate_int_scalar(args.pad_r)
|
|
inputs[3] = self.add_immediate_int_scalar(args.pad_t)
|
|
inputs[4] = self.add_immediate_int_scalar(args.pad_b)
|
|
inputs[5] = self.add_immediate_int_scalar(args.stride_w)
|
|
inputs[6] = self.add_immediate_int_scalar(args.stride_h)
|
|
inputs[7] = self.add_immediate_int_scalar(args.kernel_w)
|
|
inputs[8] = self.add_immediate_int_scalar(args.kernel_h)
|
|
inputs[9] = self.add_immediate_int_scalar(NNAPI_FuseCode.FUSED_NONE)
|
|
inputs[10] = self.add_immediate_bool_scalar(use_nchw)
|
|
|
|
outputs = [None] * 1
|
|
out_id = self.add_tensor_operand(
|
|
node.outputsAt(0), image_oper._replace(shape=out_shape)
|
|
)
|
|
self._handle_conv_pool_flexible_input(out_id, image, args, False)
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.AVERAGE_POOL_2D, inputs, outputs)
|
|
|
|
def add_adaptive_avg_pool2d(self, node):
|
|
assert node.inputsSize() == 2
|
|
assert node.outputsSize() == 1
|
|
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval_fixed_size(
|
|
node.inputsAt(0)
|
|
)
|
|
assert len(image_oper.shape) == 4
|
|
|
|
size_ctype, size_arg = self.get_constant_value(node.inputsAt(1))
|
|
assert size_ctype.kind() == "ListType"
|
|
assert size_ctype.getElementType().kind() == "IntType"
|
|
if size_arg != [1, 1]:
|
|
raise Exception(
|
|
"NNAPI only supports adaptive_avg_pool2d with output size (1, 1)."
|
|
)
|
|
|
|
out_shape = image_oper.shape[0:2] + tuple(size_arg)
|
|
use_nchw = image_oper.use_nchw()
|
|
|
|
inputs = [None] * 11
|
|
inputs[0] = image_id
|
|
inputs[1] = self.add_immediate_int_scalar(0)
|
|
inputs[2] = self.add_immediate_int_scalar(0)
|
|
inputs[3] = self.add_immediate_int_scalar(0)
|
|
inputs[4] = self.add_immediate_int_scalar(0)
|
|
inputs[5] = self.add_immediate_int_scalar(1)
|
|
inputs[6] = self.add_immediate_int_scalar(1)
|
|
inputs[7] = self.add_immediate_int_scalar(image_oper.shape[3])
|
|
inputs[8] = self.add_immediate_int_scalar(image_oper.shape[2])
|
|
inputs[9] = self.add_immediate_int_scalar(NNAPI_FuseCode.FUSED_NONE)
|
|
inputs[10] = self.add_immediate_bool_scalar(use_nchw)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(
|
|
node.outputsAt(0), image_oper._replace(shape=out_shape)
|
|
)
|
|
|
|
self.add_operation(NNAPI_OperationCode.AVERAGE_POOL_2D, inputs, outputs)
|
|
|
|
def add_upsample_nearest2d(self, node):
|
|
assert node.inputsSize() == 3 or node.inputsSize() == 4
|
|
assert node.outputsSize() == 1
|
|
if node.inputsSize() == 3:
|
|
image, size_jit, scale_jit = node.inputs()
|
|
else:
|
|
image, size_jit, scale_h_jit, scale_w_jit = node.inputs()
|
|
size_ctype, size_arg = self.get_constant_value(size_jit)
|
|
|
|
if node.inputsSize() == 3:
|
|
scale_ctype, scale_arg = self.get_constant_value(scale_jit) # type: ignore[possibly-undefined]
|
|
else:
|
|
scale_h_ctype, scale_h_arg = self.get_constant_value(scale_h_jit) # type: ignore[possibly-undefined]
|
|
scale_w_ctype, scale_w_arg = self.get_constant_value(scale_w_jit) # type: ignore[possibly-undefined]
|
|
|
|
# The only way for the 4-argument overload of upsample_nearest2d to
|
|
# have been added to the graph without error is if the scale_h and
|
|
# scale_w arguments are None
|
|
assert scale_h_ctype.kind() == "NoneType"
|
|
assert scale_w_ctype.kind() == "NoneType"
|
|
|
|
scale_ctype = scale_h_ctype
|
|
scale_arg = scale_h_arg
|
|
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval(image)
|
|
assert len(image_oper.shape) == 4
|
|
|
|
if size_ctype.kind() != "NoneType" and scale_ctype.kind() != "NoneType":
|
|
raise Exception("Size and scale cannot both be non-None.")
|
|
elif size_ctype.kind() != "NoneType":
|
|
assert size_ctype.kind() == "ListType"
|
|
assert size_ctype.getElementType().kind() == "IntType"
|
|
assert scale_ctype.kind() == "NoneType"
|
|
assert scale_arg is None
|
|
assert isinstance(size_arg, list)
|
|
assert size_arg
|
|
assert all(isinstance(val, int) for val in size_arg)
|
|
if len(size_arg) == 1:
|
|
size_arg = size_arg * 2
|
|
assert len(size_arg) == 2
|
|
out_h = size_arg[0]
|
|
out_w = size_arg[1]
|
|
arg_h = self.add_immediate_int_scalar(out_h)
|
|
arg_w = self.add_immediate_int_scalar(out_w)
|
|
elif scale_ctype.kind() != "NoneType":
|
|
assert scale_ctype.kind() == "ListType"
|
|
assert scale_ctype.getElementType().kind() == "FloatType"
|
|
assert size_ctype.kind() == "NoneType"
|
|
assert size_arg is None
|
|
assert isinstance(scale_arg, list)
|
|
assert scale_arg
|
|
assert all(isinstance(val, float) for val in scale_arg)
|
|
if len(scale_arg) == 1:
|
|
scale_arg = scale_arg * 2
|
|
assert len(scale_arg) == 2
|
|
out_h = int(scale_arg[0] * image_oper.shape[2])
|
|
out_w = int(scale_arg[1] * image_oper.shape[3])
|
|
arg_h = self.add_immediate_float_scalar(scale_arg[0])
|
|
arg_w = self.add_immediate_float_scalar(scale_arg[1])
|
|
else:
|
|
raise Exception("Size and scale cannot both be None.")
|
|
|
|
out_shape = (image_oper.shape[0], image_oper.shape[1], out_h, out_w)
|
|
use_nchw = image_oper.use_nchw()
|
|
out_id = self.add_tensor_operand(
|
|
node.outputsAt(0), image_oper._replace(shape=out_shape)
|
|
)
|
|
|
|
if image_oper.shape[0] == 0 or image_oper.shape[1] == 0:
|
|
raise Exception("Flexible batch or channels not supported")
|
|
|
|
# Handle variable input size
|
|
for dim in (2, 3): # h, w indices
|
|
if image_oper.shape[dim] == 0:
|
|
if size_ctype.kind() != "NoneType":
|
|
self.compute_operand_shape(out_id, dim, size_arg[dim - 2])
|
|
elif scale_ctype.kind() != "NoneType":
|
|
self.compute_operand_shape(
|
|
out_id,
|
|
dim,
|
|
f"int({scale_arg[dim - 2]} * {flex_name(image_id, dim)})",
|
|
)
|
|
else:
|
|
raise Exception("Size and scale cannot both be None.")
|
|
|
|
inputs = [None] * 4
|
|
inputs[0] = image_id
|
|
inputs[1] = arg_w
|
|
inputs[2] = arg_h
|
|
inputs[3] = self.add_immediate_bool_scalar(use_nchw)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.RESIZE_NEAREST_NEIGHBOR, inputs, outputs)
|
|
|
|
def add_addmm(self, node):
|
|
assert node.inputsSize() == 5
|
|
assert node.outputsSize() == 1
|
|
jit_bias, jit_input, jit_weight, jit_beta, jit_alpha = node.inputs()
|
|
|
|
for jitval in (jit_beta, jit_alpha):
|
|
scale_ctype, scale_value = self.get_constant_value(jitval)
|
|
assert scale_ctype.kind() in ("IntType", "FloatType")
|
|
if scale_value != 1:
|
|
raise Exception(
|
|
"NNAPI Fully-Connected does not support alpha and beta."
|
|
)
|
|
|
|
self.add_addmm_or_linear(node, True, jit_input, jit_weight, jit_bias)
|
|
|
|
def add_linear(self, node):
|
|
assert node.inputsSize() == 3
|
|
assert node.outputsSize() == 1
|
|
jit_input, jit_weight, jit_bias = node.inputs()
|
|
|
|
self.add_addmm_or_linear(node, False, jit_input, jit_weight, jit_bias)
|
|
|
|
def add_addmm_or_linear(
|
|
self, node, transpose_weight, jit_input, jit_weight, jit_bias
|
|
):
|
|
input_id, input_oper = self.get_tensor_operand_by_jitval(jit_input)
|
|
bias_id, bias_oper = self.get_tensor_operand_for_weight(jit_bias)
|
|
|
|
assert len(input_oper.shape) == 2
|
|
assert len(bias_oper.shape) == 1
|
|
|
|
# TODO: Transform at load time to share weights with CPU model.
|
|
_, weight_tensor = self.get_constant_value(jit_weight, "TensorType")
|
|
assert len(weight_tensor.shape) == 2
|
|
if transpose_weight:
|
|
nnapi_weight_tensor = weight_tensor.t().contiguous()
|
|
else:
|
|
nnapi_weight_tensor = weight_tensor.contiguous()
|
|
weight_id = self.add_tensor_operand_for_weight(nnapi_weight_tensor)
|
|
weight_oper = self.operands[weight_id]
|
|
|
|
out_shape = (input_oper.shape[0], weight_oper.shape[0])
|
|
out_id = self.add_tensor_operand(
|
|
node.outputsAt(0), input_oper._replace(shape=out_shape)
|
|
)
|
|
|
|
if input_oper.shape[0] == 0:
|
|
self.forward_operand_shape(out_id, 0, input_id, 0)
|
|
|
|
inputs = [None] * 4
|
|
inputs[0] = input_id
|
|
inputs[1] = weight_id
|
|
inputs[2] = bias_id
|
|
inputs[3] = self.add_immediate_int_scalar(NNAPI_FuseCode.FUSED_NONE)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = out_id
|
|
|
|
self.add_operation(NNAPI_OperationCode.FULLY_CONNECTED, inputs, outputs)
|
|
|
|
def add_qlinear(self, node):
|
|
assert node.inputsSize() == 4
|
|
assert node.outputsSize() == 1
|
|
(
|
|
jit_input,
|
|
jit_packed_weight,
|
|
jit_scale,
|
|
jit_zero_point,
|
|
) = node.inputs()
|
|
|
|
input_id, input_oper = self.get_tensor_operand_by_jitval_fixed_size(jit_input)
|
|
# TODO: Support automatic reshape
|
|
assert len(input_oper.shape) == 2
|
|
|
|
_, out_scale = self.get_constant_value(jit_scale, "FloatType")
|
|
_, out_zero_point = self.get_constant_value(jit_zero_point, "IntType")
|
|
weight_ctype, packed_weight = self.get_constant_value(jit_packed_weight)
|
|
assert weight_ctype.name() == "LinearPackedParamsBase"
|
|
raw_weight, raw_bias = packed_weight.__getstate__()[0]
|
|
assert raw_bias is not None
|
|
|
|
assert len(raw_weight.shape) == 2
|
|
assert len(raw_bias.shape) == 1
|
|
assert raw_bias.shape[0] == raw_weight.shape[0]
|
|
assert raw_weight.shape[1] == input_oper.shape[1]
|
|
|
|
assert raw_weight.qscheme() == torch.per_tensor_affine
|
|
if raw_weight.dtype == torch.quint8:
|
|
unsigned_weight = raw_weight
|
|
else:
|
|
assert raw_weight.dtype == torch.qint8
|
|
unsigned_weight = torch._make_per_tensor_quantized_tensor(
|
|
(raw_weight.int_repr().int() + 128).to(torch.uint8),
|
|
scale=raw_weight.q_scale(),
|
|
zero_point=raw_weight.q_zero_point() + 128,
|
|
)
|
|
weight_scale = unsigned_weight.q_scale()
|
|
bias_scale = input_oper.scale * weight_scale
|
|
int_bias = torch.quantize_per_tensor(raw_bias, bias_scale, 0, torch.qint32)
|
|
bias_id = self.add_tensor_operand_for_weight(int_bias)
|
|
|
|
multiplier = input_oper.scale * weight_scale / out_scale
|
|
assert multiplier > 0
|
|
if multiplier >= 1:
|
|
raise Exception(
|
|
"Quantized convolution multiplier is greater than 1. "
|
|
"This is supported by NNAPI, but not by most hardware backends. "
|
|
"Try training a model without quantization-aware training. "
|
|
)
|
|
|
|
# TODO: Transform at load time to share weights with CPU model.
|
|
nnapi_weight_tensor = unsigned_weight.contiguous()
|
|
weight_id = self.add_tensor_operand_for_weight(nnapi_weight_tensor)
|
|
weight_oper = self.operands[weight_id]
|
|
|
|
out_shape = (input_oper.shape[0], weight_oper.shape[0])
|
|
out_oper = input_oper._replace(
|
|
shape=out_shape,
|
|
scale=out_scale,
|
|
zero_point=out_zero_point,
|
|
)
|
|
|
|
inputs = [None] * 4
|
|
inputs[0] = input_id
|
|
inputs[1] = weight_id
|
|
inputs[2] = bias_id
|
|
inputs[3] = self.add_immediate_int_scalar(NNAPI_FuseCode.FUSED_NONE)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(node.outputsAt(0), out_oper)
|
|
|
|
self.add_operation(NNAPI_OperationCode.FULLY_CONNECTED, inputs, outputs)
|
|
|
|
def get_optional_bias(self, jit_bias, weight_tensor, transpose=False):
|
|
ctype, value = self.get_constant_value(jit_bias)
|
|
if ctype.kind() == "NoneType":
|
|
bias_idx = 1 if transpose else 0
|
|
nnapi_bias_tensor = torch.zeros(
|
|
weight_tensor.size()[bias_idx], dtype=weight_tensor.dtype
|
|
)
|
|
bias_id = self.add_tensor_operand_for_weight(nnapi_bias_tensor)
|
|
bias_oper = self.operands[bias_id]
|
|
return bias_id, bias_oper
|
|
else:
|
|
return self.get_tensor_operand_for_weight(jit_bias)
|
|
|
|
def add_conv2d(self, node):
|
|
assert node.inputsSize() == 7
|
|
assert node.outputsSize() == 1
|
|
|
|
(
|
|
jit_image,
|
|
jit_weight,
|
|
jit_bias,
|
|
jit_stride,
|
|
jit_pad,
|
|
jit_dilation,
|
|
jit_groups,
|
|
) = node.inputs()
|
|
|
|
_, weight_tensor = self.get_constant_value(jit_weight, "TensorType")
|
|
bias_id, bias_oper = self.get_optional_bias(jit_bias, weight_tensor)
|
|
args = self.get_conv_pool_args_2d_from_jit(
|
|
weight_tensor.shape[2:4], jit_stride, jit_pad, jit_dilation, jit_groups
|
|
)
|
|
|
|
return self.add_conv2d_common(
|
|
node.outputsAt(0),
|
|
0.0,
|
|
0,
|
|
jit_image,
|
|
weight_tensor,
|
|
bias_id,
|
|
args,
|
|
False, # transpose
|
|
NNAPI_FuseCode.FUSED_NONE,
|
|
)
|
|
|
|
def add_conv_underscore(self, node):
|
|
assert node.inputsSize() == 13
|
|
assert node.outputsSize() == 1
|
|
|
|
(
|
|
jit_image,
|
|
jit_weight,
|
|
jit_bias,
|
|
jit_stride,
|
|
jit_pad,
|
|
jit_dilation,
|
|
jit_transpose,
|
|
_,
|
|
jit_groups,
|
|
_,
|
|
_,
|
|
_,
|
|
_,
|
|
) = node.inputs()
|
|
|
|
_, weight_tensor = self.get_constant_value(jit_weight, "TensorType")
|
|
_, transpose = self.get_constant_value(jit_transpose)
|
|
bias_id, bias_oper = self.get_optional_bias(jit_bias, weight_tensor, transpose)
|
|
args = self.get_conv_pool_args_2d_from_jit(
|
|
weight_tensor.shape[2:4], jit_stride, jit_pad, jit_dilation, jit_groups
|
|
)
|
|
|
|
return self.add_conv2d_common(
|
|
node.outputsAt(0),
|
|
0.0,
|
|
0,
|
|
jit_image,
|
|
weight_tensor,
|
|
bias_id,
|
|
args,
|
|
transpose,
|
|
NNAPI_FuseCode.FUSED_NONE,
|
|
)
|
|
|
|
def add_log_softmax(self, node):
|
|
assert node.inputsSize() == 3
|
|
assert node.outputsSize() == 1
|
|
|
|
(jit_input, jit_dim, jit_half_to_float) = node.inputs()
|
|
input_id, input_oper = self.get_tensor_operand_by_jitval_fixed_size(jit_input)
|
|
_, dim = self.get_constant_value(jit_dim, "IntType")
|
|
|
|
out_shape = input_oper.shape
|
|
|
|
inputs = [None] * 3
|
|
inputs[0] = input_id
|
|
# specifying 1 as the scaling factor for the exponent, beta
|
|
inputs[1] = self.add_immediate_float_scalar(1)
|
|
inputs[2] = self.add_immediate_int_scalar(dim)
|
|
|
|
outputs = [None] * 1
|
|
outputs[0] = self.add_tensor_operand(
|
|
node.outputsAt(0), input_oper._replace(shape=out_shape)
|
|
)
|
|
self.add_operation(NNAPI_OperationCode.LOG_SOFTMAX, inputs, outputs)
|
|
|
|
def add_qconv2d(self, node, fuse_code, transpose=False):
|
|
assert node.inputsSize() == 4
|
|
assert node.outputsSize() == 1
|
|
|
|
(
|
|
jit_image,
|
|
jit_packed_weight,
|
|
jit_scale,
|
|
jit_zero_point,
|
|
) = node.inputs()
|
|
|
|
_, out_scale = self.get_constant_value(jit_scale, "FloatType")
|
|
_, out_zero_point = self.get_constant_value(jit_zero_point, "IntType")
|
|
weight_ctype, packed_weight = self.get_constant_value(jit_packed_weight)
|
|
assert weight_ctype.name() == "Conv2dPackedParamsBase"
|
|
(
|
|
pack_version,
|
|
tensors,
|
|
opt_tensors,
|
|
) = packed_weight.__getstate__()[0]
|
|
assert pack_version == "2"
|
|
packed_config, raw_weight = tensors
|
|
(raw_bias,) = opt_tensors
|
|
assert raw_bias is not None
|
|
args = self.get_conv_pool_args_2d_from_pack(
|
|
raw_weight.shape[2:4], packed_config
|
|
)
|
|
|
|
assert raw_weight.qscheme() == torch.per_tensor_affine
|
|
if raw_weight.dtype == torch.quint8:
|
|
unsigned_weight = raw_weight
|
|
else:
|
|
assert raw_weight.dtype == torch.qint8
|
|
unsigned_weight = torch._make_per_tensor_quantized_tensor(
|
|
(raw_weight.int_repr().int() + 128).to(torch.uint8),
|
|
scale=raw_weight.q_scale(),
|
|
zero_point=raw_weight.q_zero_point() + 128,
|
|
)
|
|
weight_scale = unsigned_weight.q_scale()
|
|
_, image_oper = self.get_tensor_operand_by_jitval(jit_image)
|
|
bias_scale = image_oper.scale * weight_scale
|
|
int_bias = torch.quantize_per_tensor(raw_bias, bias_scale, 0, torch.qint32)
|
|
bias_id = self.add_tensor_operand_for_weight(int_bias)
|
|
|
|
multiplier = image_oper.scale * weight_scale / out_scale
|
|
assert multiplier > 0
|
|
if multiplier >= 1:
|
|
raise Exception(
|
|
"Quantized convolution multiplier is greater than 1. "
|
|
"This is supported by NNAPI, but not by most hardware backends. "
|
|
"Try training a model without quantization-aware training. "
|
|
)
|
|
|
|
return self.add_conv2d_common(
|
|
node.outputsAt(0),
|
|
out_scale,
|
|
out_zero_point,
|
|
jit_image,
|
|
unsigned_weight,
|
|
bias_id,
|
|
args,
|
|
transpose,
|
|
fuse_code,
|
|
)
|
|
|
|
def add_conv2d_common(
|
|
self,
|
|
jit_out,
|
|
out_scale,
|
|
out_zero_point,
|
|
jit_image,
|
|
weight_tensor,
|
|
bias_id,
|
|
args,
|
|
transpose,
|
|
fuse_code,
|
|
):
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval(jit_image)
|
|
in_c = image_oper.shape[1]
|
|
|
|
if args.group == 1:
|
|
# Full convolution
|
|
depthwise = False
|
|
if transpose:
|
|
weight_permutation = (1, 2, 3, 0)
|
|
else:
|
|
weight_permutation = (0, 2, 3, 1)
|
|
elif args.group == in_c:
|
|
# Depthwise convolution
|
|
depthwise = True
|
|
weight_permutation = (1, 2, 3, 0)
|
|
else:
|
|
raise Exception("Group convolution not supported yet.")
|
|
|
|
# TODO: Transform at load time to share weights with CPU model.
|
|
nnapi_weight_tensor = weight_tensor.permute(*weight_permutation).contiguous()
|
|
weight_id = self.add_tensor_operand_for_weight(nnapi_weight_tensor)
|
|
weight_oper = self.operands[weight_id]
|
|
|
|
bias_oper = self.operands[bias_id]
|
|
|
|
if image_oper.op_type == NNAPI_OperandCode.TENSOR_FLOAT32:
|
|
assert weight_oper.op_type == NNAPI_OperandCode.TENSOR_FLOAT32
|
|
assert bias_oper.op_type == NNAPI_OperandCode.TENSOR_FLOAT32
|
|
elif image_oper.op_type == NNAPI_OperandCode.TENSOR_QUANT8_ASYMM:
|
|
assert weight_oper.op_type == NNAPI_OperandCode.TENSOR_QUANT8_ASYMM
|
|
assert bias_oper.op_type == NNAPI_OperandCode.TENSOR_INT32
|
|
assert approx_equal(image_oper.scale * weight_oper.scale, bias_oper.scale)
|
|
assert bias_oper.zero_point == 0
|
|
else:
|
|
raise Exception(f"Unsupported input type for conv2d: {image_oper.op_type}")
|
|
|
|
assert len(image_oper.shape) == 4
|
|
assert len(weight_oper.shape) == 4
|
|
assert len(bias_oper.shape) == 1
|
|
|
|
if depthwise:
|
|
# Depthwise convolution
|
|
one, kern_h, kern_w, out_c = weight_oper.shape
|
|
assert one == 1
|
|
assert out_c % in_c == 0
|
|
channel_multiplier = out_c // in_c
|
|
assert channel_multiplier == 1 # Don't support multiplier
|
|
assert out_c == in_c
|
|
else:
|
|
# Full convolution
|
|
out_c, kern_h, kern_w, kern_d = weight_oper.shape
|
|
assert kern_d == in_c
|
|
|
|
assert out_c == bias_oper.shape[0]
|
|
|
|
use_nchw = image_oper.use_nchw()
|
|
|
|
if depthwise:
|
|
num_args = 12
|
|
opcode = NNAPI_OperationCode.DEPTHWISE_CONV_2D
|
|
else:
|
|
num_args = 11
|
|
if transpose:
|
|
opcode = NNAPI_OperationCode.TRANSPOSE_CONV_2D
|
|
else:
|
|
opcode = NNAPI_OperationCode.CONV_2D
|
|
|
|
inputs = [None] * num_args
|
|
inputs[0] = image_id
|
|
inputs[1] = weight_id
|
|
inputs[2] = bias_id
|
|
inputs[3] = self.add_immediate_int_scalar(args.pad_l)
|
|
inputs[4] = self.add_immediate_int_scalar(args.pad_r)
|
|
inputs[5] = self.add_immediate_int_scalar(args.pad_t)
|
|
inputs[6] = self.add_immediate_int_scalar(args.pad_b)
|
|
inputs[7] = self.add_immediate_int_scalar(args.stride_w)
|
|
inputs[8] = self.add_immediate_int_scalar(args.stride_h)
|
|
if depthwise:
|
|
inputs[9] = self.add_immediate_int_scalar(1)
|
|
inputs[10] = self.add_immediate_int_scalar(fuse_code)
|
|
inputs[11] = self.add_immediate_bool_scalar(use_nchw)
|
|
else:
|
|
inputs[9] = self.add_immediate_int_scalar(fuse_code)
|
|
inputs[10] = self.add_immediate_bool_scalar(use_nchw)
|
|
|
|
outputs = [None] * 1
|
|
out_shape = get_conv_pool_shape(image_oper.shape, args, out_c, transpose)
|
|
out_oper = image_oper._replace(
|
|
shape=out_shape,
|
|
scale=out_scale,
|
|
zero_point=out_zero_point,
|
|
)
|
|
out_id = self.add_tensor_operand(jit_out, out_oper)
|
|
self._handle_conv_pool_flexible_input(out_id, jit_image, args, transpose)
|
|
|
|
outputs[0] = out_id
|
|
self.add_operation(opcode, inputs, outputs)
|
|
|
|
def _handle_conv_pool_flexible_input(self, out_id, jit_image, args, transpose):
|
|
image_id, image_oper = self.get_tensor_operand_by_jitval(jit_image)
|
|
batch, in_ch, in_h, in_w = image_oper.shape
|
|
|
|
if batch == 0:
|
|
self.forward_operand_shape(out_id, 0, image_id, 0)
|
|
if in_ch == 0:
|
|
raise Exception("Input channels can't be flexible")
|
|
# H & W
|
|
if transpose:
|
|
if in_h == 0:
|
|
self.compute_operand_shape(
|
|
out_id,
|
|
2,
|
|
f"({flex_name(image_id, 2)} - 1) * {args.stride_h} + {args.kernel_h} - {args.pad_t} - {args.pad_b}",
|
|
)
|
|
if in_w == 0:
|
|
self.compute_operand_shape(
|
|
out_id,
|
|
3,
|
|
f"({flex_name(image_id, 3)} - 1) * {args.stride_w} + {args.kernel_w} - {args.pad_l} - {args.pad_r}",
|
|
)
|
|
else:
|
|
if in_h == 0:
|
|
self.compute_operand_shape(
|
|
out_id,
|
|
2,
|
|
f"({flex_name(image_id, 2)} - {args.kernel_h} + {args.pad_t} + {args.pad_b}) // {args.stride_h} + 1",
|
|
)
|
|
if in_w == 0:
|
|
self.compute_operand_shape(
|
|
out_id,
|
|
3,
|
|
f"({flex_name(image_id, 3)} - {args.kernel_w} + {args.pad_l} + {args.pad_r}) // {args.stride_w} + 1",
|
|
)
|
|
|
|
|
|
def serialize_model(
|
|
module, inputs, *, config=None, return_shapes=None, use_int16_for_qint16=False
|
|
):
|
|
"""Convert to NNAPI and serialize torchscript module.
|
|
|
|
Parameters:
|
|
module: Torchscript module to convert
|
|
inputs: Tensors used to specify input details for NNAPI
|
|
config (optional): Optional config to attach to module
|
|
return_shapes (optional): Specify shape of outputs if
|
|
your module uses runtime flexible shapes to set output
|
|
buffer size for NNAPI
|
|
use_int16_for_qint16 (optional): Use Pytorch int16 to represent NNAPI qint16 values
|
|
"""
|
|
return _NnapiSerializer(config, use_int16_for_qint16).serialize_model(
|
|
module, inputs, return_shapes
|
|
)
|