#########################################################
# Copyright (C) 2024 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Christopher Rodrigues
#########################################################
"""
Functions to work with node rankings for mixed precision search.
"""
from typing import Set
from afe.core.configs import Opt, QuantizationPrecision
from afe.ir.defines import NodeName, TensorValue
from afe.ir.net import AwesomeNet
from afe.ir.node import node_is_sima_ir, node_is_placeholder, node_is_awesomenet
from afe.ir.tensor_type import TensorType, ScalarType
[docs]
def rank_int16_nodes(analysis_result: dict[NodeName, float]) -> list[list[NodeName]]:
"""
Rank nodes by their sensitivity.
:param analysis_result: Estimated sensitivity of nodes.
Larger values indicate greater sensitivity. Noise metrics such as MSE
can be used as sensitivity. The reciprocal of PSNR can be used as sensitivity.
:return: Sets of node names ordered from high sensitivity to low sensitivity.
Nodes in the same set have equal sensitivity.
"""
ranked_list: list[tuple[NodeName, float]] = list(analysis_result.items())
ranked_list.sort(key=lambda x: x[1], reverse=True)
# Group nodes that have equal sensitivity
grouped_list: list[list[NodeName]] = []
last_value = None
for name, value in ranked_list:
if last_value is not None and value == last_value:
# Append to previous group
grouped_list[-1].append(name)
else:
# Create new group
grouped_list.append([name])
last_value = value
return grouped_list
[docs]
def annotate_int16_nodes(net: AwesomeNet, node_names: Set[NodeName]) -> None:
"""
Annotate nodes in the list so that they will be quantized with int16 precision.
:param net: Net to annotate
:param node_names: Nodes to be given int16 precision.
"""
for node in net.nodes.values():
if node_is_awesomenet(node):
annotate_int16_nodes(node.ir, node_names)
elif node_is_sima_ir(node) and node.name in node_names:
node.ir.quant_config.quantization_precision = Opt(QuantizationPrecision.INT_16)
# If input node is a placeholder node then precision in the placeholder node should also be set to INT16
input_node_names = node.input_node_names
for input_node_name in input_node_names:
input_node = net.nodes[input_node_name]
if node_is_placeholder(input_node):
if input_node.ir.quant_config.quantization_precision.is_empty():
input_node.ir.quant_config.quantization_precision = Opt(QuantizationPrecision.INT_16)
[docs]
def count_int16_nodes(net: AwesomeNet) -> tuple[int, int]:
"""
Count the number of SiMa IR nodes where the output type is int16.
Only SiMaIR nodes (representing single operators) are counted.
:param net: Network to analyze
:return: Number of int16 SiMaIR nodes and total number of SiMaIR nodes
"""
int16_count = 0
total_count = 0
for node in net.iter_nodes_recursive():
if node_is_sima_ir(node):
total_count += 1
match node.get_type().output:
case TensorValue(TensorType(scalar=ScalarType.int16)):
int16_count += 1
return int16_count, total_count