Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions backends/arm/quantizer/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,12 @@
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""Expose quantizer APIs and load optional quantized kernels.

Import the public quantizer classes and configuration helpers for Arm
backends. Attempt to load portable and quantized libraries; fall back to a
log message if unavailable.
"""

from .quantization_config import QuantizationConfig # noqa # usort: skip
from .arm_quantizer import ( # noqa
Expand Down
38 changes: 31 additions & 7 deletions backends/arm/quantizer/arm_quantizer_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,12 @@
# LICENSE file in the root directory of this source tree.

# pyre-unsafe
"""Provide utilities for quantization annotations.

#
# Utility functions for TOSAQuantizer
#
Use these helpers to check and mark annotation state when working with
``QuantizationAnnotation`` entries in FX node metadata.

"""

from typing import cast

Expand All @@ -20,15 +22,31 @@


def is_annotated(node: Node) -> bool:
"""Given a node return whether the node is annotated."""
"""Return True if the node is annotated.

Args:
node (Node): FX node to inspect.

Returns:
bool: True if ``Q_ANNOTATION_KEY`` exists and ``_annotated`` is set.

"""
return (
Q_ANNOTATION_KEY in node.meta
and cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])._annotated
)


def is_output_annotated(node: Node) -> bool:
"""Given a node, return whether the output of the node is annotated."""
"""Return True if the node's output is annotated.

Args:
node (Node): FX node to inspect.

Returns:
bool: True if annotated and an output qspec is present.

"""
if Q_ANNOTATION_KEY in node.meta:
annotation = cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])
return annotation._annotated and annotation.output_qspec is not None
Expand All @@ -37,8 +55,14 @@ def is_output_annotated(node: Node) -> bool:


def mark_node_as_annotated(node: Node) -> None:
"""Marks node as annotated. If needed, an empty QuantizationAnnotation is added
to the quantization_annotation node meta entry.
"""Mark a node as annotated.

Create an empty ``QuantizationAnnotation`` on the node when missing and set
its ``_annotated`` flag to True.

Args:
node (Node): FX node to update.

"""
if Q_ANNOTATION_KEY not in node.meta:
node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation()
Expand Down
Loading