wehub-resource-sync
194 lines
7.9 KiB
Python
194 lines
7.9 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
"""Quantized Neural Network (QNN) Operators"""
|
|
|
|
import tvm
|
|
from tvm import te, tirx, topi
|
|
|
|
SQNN_DISABLE = 0
|
|
SQNN_INT8 = 1
|
|
SQNN_UINT8 = 2
|
|
SQNN_INT32 = 3
|
|
|
|
SQNN_DTYPE_TO_CODE = {
|
|
"disable": SQNN_DISABLE,
|
|
"int8": SQNN_INT8,
|
|
"uint8": SQNN_UINT8,
|
|
"int32": SQNN_INT32,
|
|
}
|
|
|
|
SQNN_CODE_TO_DTYPE = {v: k for k, v in SQNN_DTYPE_TO_CODE.items()}
|
|
|
|
|
|
@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
|
|
def simulated_quantize(data, out_dtype, output_scale=None, output_zero_point=None, axis=-1):
|
|
"""Simulated QNN quantize operator that mimics QNN outputs without changing datatype.
|
|
The benefit of this operator over true QNN quantize is that this operator allows dynamic
|
|
datatype selection and can operate on both per-channel and scalar scales and zero points while
|
|
QNN quantize requires both of these to be fixed at compile time.
|
|
|
|
Parameters
|
|
----------
|
|
data: tvm.te.Tensor
|
|
An N-D input tensor to the operator.
|
|
|
|
out_dtype: tvm.te.Tensor
|
|
A scalar variable that indicates which datatype to simulate quantization with. Use
|
|
SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
|
|
value.
|
|
|
|
output_scale: tvm.te.Tensor, optional
|
|
A scalar tensor representing the scale to use when quantizing to integer datatypes.
|
|
When it contains more than a single value, N must match the number of channels in data.
|
|
|
|
output_zero_point: tvm.te.Tensor, optional
|
|
A 1-D tensor representing the zero point to use when quantizing to integer datatypes.
|
|
When it contains more than a single value, N must match the number of channels in data.
|
|
|
|
axis: int, optional
|
|
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
|
|
|
|
"""
|
|
|
|
# When disabled, just pass through the input values.
|
|
def _compute_pass_through(value, *indices):
|
|
return value[indices]
|
|
|
|
# Simulate quantization for arbitrary integer datatypes. The computation for all datatypes is:
|
|
# Q_output = clip((round(input_tensor/output_scale) + output_zero_point),
|
|
# out_dtype::min,
|
|
# out_dtype::max)
|
|
def _compute_intn(dtype, value, *indices):
|
|
assert output_scale is not None and output_zero_point is not None
|
|
const_min = tvm.tirx.min_value(dtype)
|
|
const_max = tvm.tirx.max_value(dtype)
|
|
# Use indexmod to handle both scalar and per-channel QNN parameters.
|
|
scale_idx = tirx.indexmod(indices[axis], topi.shape(output_scale)[0])
|
|
zp_idx = tirx.indexmod(indices[axis], topi.shape(output_zero_point)[0])
|
|
return te.max(
|
|
te.min(
|
|
te.round(value[indices] / output_scale[scale_idx]) + output_zero_point[zp_idx],
|
|
const_max,
|
|
),
|
|
const_min,
|
|
)
|
|
|
|
# Use an if chain to dynamically return the proper quantization based on the input datatype.
|
|
# This allows the op to compile once but apply different quantization approaches
|
|
# using a variable datatype input.
|
|
def _dispatch_sim_quantize(value):
|
|
pass_through_value = te.compute(
|
|
data.shape, lambda *indices: _compute_pass_through(value, *indices)
|
|
)
|
|
int8_value = te.compute(
|
|
data.shape,
|
|
lambda *indices: tirx.if_then_else(
|
|
out_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
|
|
_compute_intn("int8", value, *indices),
|
|
pass_through_value[indices],
|
|
),
|
|
)
|
|
uint8_value = te.compute(
|
|
data.shape,
|
|
lambda *indices: tirx.if_then_else(
|
|
out_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
|
|
_compute_intn("uint8", value, *indices),
|
|
int8_value[indices],
|
|
),
|
|
)
|
|
int32_value = te.compute(
|
|
data.shape,
|
|
lambda *indices: tirx.if_then_else(
|
|
out_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
|
|
_compute_intn("int32", value, *indices),
|
|
uint8_value[indices],
|
|
),
|
|
)
|
|
|
|
return int32_value
|
|
|
|
return te.compute(data.shape, lambda *indices: _dispatch_sim_quantize(data)[indices])
|
|
|
|
|
|
@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
|
|
def simulated_dequantize(data, in_dtype, input_scale=None, input_zero_point=None, axis=-1):
|
|
"""Simulated QNN dequantize operator that mimics QNN outputs without changing datatype.
|
|
The benefit of this operator over true QNN dequantize is that this operator allows dynamic
|
|
datatype selection and can operate on both per-channel and scalar scales and zero points while
|
|
QNN dequantize requires both of these to be fixed at compile time.
|
|
|
|
Parameters
|
|
----------
|
|
data: tvm.te.Tensor
|
|
An N-D input tensor to the operator.
|
|
|
|
in_dtype: tvm.te.Tensor
|
|
A scalar variable that indicates which datatype to simulate dequantization with. Use
|
|
SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
|
|
value.
|
|
|
|
input_scale: tvm.te.Tensor, optional
|
|
A scalar tensor representing the scale to use when dequantizing from integer datatypes.
|
|
When it contains more than a single value, N must match the number of channels in data.
|
|
|
|
input_zero_point: tvm.te.Tensor, optional
|
|
A 1-D tensor representing the zero point to use when dequantizing from integer datatypes.
|
|
When it contains more than a single value, N must match the number of channels in data.
|
|
|
|
axis: int, optional
|
|
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
|
|
|
|
"""
|
|
|
|
# When disabled simply return the input tensor.
|
|
def _compute_pass_through(value, *indices):
|
|
return value[indices]
|
|
|
|
# Simulate dequantization for arbitrary integer datatypes. The computation for all datatypes is:
|
|
# DQ_output = (input - zero_point) * scale
|
|
def _compute_intn(value, *indices):
|
|
assert input_scale is not None and input_zero_point is not None
|
|
# Use indexmod to handle both scalar and per-channel QNN parameters.
|
|
scale_idx = tirx.indexmod(indices[axis], topi.shape(input_scale)[0])
|
|
zp_idx = tirx.indexmod(indices[axis], topi.shape(input_zero_point)[0])
|
|
return (value[indices] - input_zero_point[zp_idx]) * input_scale[scale_idx]
|
|
|
|
# Use an if chain to dynamically return the proper dequantization based on the input datatype.
|
|
# This allows the op to compile once but apply different quantization approaches
|
|
# using a variable datatype input.
|
|
def _dispatch_sim_dequantize(value):
|
|
pass_through_value = te.compute(
|
|
data.shape, lambda *indices: _compute_pass_through(value, *indices)
|
|
)
|
|
intn_condition = tvm.te.any(
|
|
in_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
|
|
in_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
|
|
in_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
|
|
)
|
|
intn_value = te.compute(
|
|
data.shape,
|
|
lambda *indices: tirx.if_then_else(
|
|
intn_condition,
|
|
_compute_intn(value, *indices),
|
|
pass_through_value[indices],
|
|
),
|
|
)
|
|
|
|
return intn_value
|
|
|
|
return te.compute(data.shape, lambda *indices: _dispatch_sim_dequantize(data)[indices])
|