[IR] Add support for quant_parameter_tensor_names field

[justinchuby]

Support quantization_annotation in graph inputs, node in/out and graph outputs.

Two design decisions made are

1. Make ir.Value carry the quant_parameter_tensor_names information. This is similar to ValueInfoProto where in proto we store a list of proto messages whose keys point tensor names. But the information really belongs to individual values.
2. quantization_annotation is deserialized into the Value's meta field under the quant_parameter_tensor_names key. Values that are stored under this key will be serialized as quantization annotations. I chose to add a value in meta instead of creating a new property in value to avoid over complicating the preperties in Value.

Example usage

>>> from onnxscript import ir
>>> model = ir.load("l_1_n_12_z_384_i_1536.onnx")
>>> model.graph.node("MVAU_rtl_0").outputs[0]
Value('MVAU_rtl_0_out0', type=Tensor(FLOAT), shape=[1,128,384], producer=MVAU_rtl_0, index=0)
>>> model.graph.node("MVAU_rtl_0").outputs[0].meta
MetadataStore({'quant_parameter_tensor_names': {'finn_datatype': 'INT22'}}, invalid_keys=set())
>>> ir.save(model, "model_with_quant_params.textproto")

[codecov]

❌ 114 Tests Failed:

Tests completed	Failed	Passed	Skipped
9749	114	9635	1942

View the top 3 failed test(s) by shortest run time

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0573_test_logsoftmax_axis_0

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_logsoftmax_axis_0'

The above exception was the direct cause of the following exception:
.nox\test_torch_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_logsoftmax_axis_0' (e=No module named 'tests.onnx_backend_test_code.test_logsoftmax_axis_0') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_logsoftmax_axis_0.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_logsoftmax_axis_0.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT
E   from onnxscript.onnx_opset import opset13
E   
E   @script()
E   def bck_test_logsoftmax_axis_0(x: FLOAT[3,4,5]) -> (FLOAT[3,4,5]):
E       y = opset13.LogSoftmax(x, axis=0)
E       return y

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_0778_test_pow_types_int64_float32

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_pow_types_int64_float32'

The above exception was the direct cause of the following exception:
.nox\test_torch_nightly\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_pow_types_int64_float32' (e=No module named 'tests.onnx_backend_test_code.test_pow_types_int64_float32') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_pow_types_int64_float32.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_pow_types_int64_float32.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset15
E   
E   @script()
E   def bck_test_pow_types_int64_float32(x: INT64[3], y: FLOAT[3]) -> (INT64[3]):
E       z = opset15.Pow(x, y)
E       return z

onnxscript.backend.onnx_export_test.TestOnnxBackEnd::test_export2python_produces_correct_onnx_script_model_1101_test_shape_end_negative_1

Stack Traces | 0.003s run time

onnxscript\backend\onnx_export_test.py:137: in extract_functions
    mod = importlib.import_module(import_name)
C:\hostedtoolcache\windows\Python\3.11.9\x64\Lib\importlib\__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
E   ModuleNotFoundError: No module named 'tests.onnx_backend_test_code.test_shape_end_negative_1'

The above exception was the direct cause of the following exception:
.nox\test\Lib\site-packages\parameterized\parameterized.py:620: in standalone_func
    return func(*(a + p.args), **p.kwargs, **kw)
onnxscript\backend\onnx_export_test.py:271: in test_export2python_produces_correct_onnx_script_model
    functions = extract_functions(backend_test.name, code, self.test_folder)
onnxscript\backend\onnx_export_test.py:139: in extract_functions
    raise AssertionError(
E   AssertionError: Unable to import 'tests.onnx_backend_test_code.test_shape_end_negative_1' (e=No module named 'tests.onnx_backend_test_code.test_shape_end_negative_1') (file: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_negative_1.py', absolute path: 'D:\\a\\onnxscript\\onnxscript\\tests\\onnx_backend_test_code\\test_shape_end_negative_1.py', current folder: D:\a\onnxscript\onnxscript
E   ---- CONTENT --
E   import numpy
E   from onnx import TensorProto
E   from onnx.helper import make_tensor
E   from onnxscript import script, external_tensor
E   from onnxscript.values import Opset
E   from onnxscript.onnx_types import FLOAT, INT64
E   from onnxscript.onnx_opset import opset21
E   
E   @script()
E   def bck_test_shape_end_negative_1(x: FLOAT[3,4,5]) -> (INT64[2]):
E       y = opset21.Shape(x, end=-1)
E       return y

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

[Copilot]

PR Overview

This PR introduces support for quantization annotations by adding a new field ("quant_parameter_tensor_names") in the Value metadata. Key changes include:

• Updating the GraphProtocol documentation to outline the new quantization annotation behavior.
• Enhancing the deserialization logic by adding support for quantization annotations in graphs.
• Adding new serialization/deserialization utilities (_deserialize_quantization_annotation, _maybe_add_quantization_annotation, and _serialize_tensor_annotation_into) to incorporate quantization annotation data.

Reviewed Changes

File	Description
onnxscript/ir/_protocols.py	Updated docstring to reflect the new quantization annotation support.
onnxscript/ir/serde.py	Extended graph/node deserialization and serialization to handle quantization annotations.

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

onnxscript/ir/serde.py:1273

• Ensure that there are unit tests covering the new quantization annotation paths in both serialization and deserialization for initializer and non-initializer values.

if input_.name not in from_.initializers:

PR Overview

This PR adds support for the quant_parameter_tensor_names field for quantization annotations in the IR. Key changes include:

• Adding a new test suite to verify the correct serialization and deserialization of quantization annotations.
• Extending the serde module with new functions to deserialize and serialize quantization annotations.
• Updating the GraphProtocol documentation to reflect the new quantization_annotation support.

Reviewed Changes

File	Description
onnxscript/ir/serde_test.py	Introduces unit tests for verifying quantization annotation serialization/deserialization.
onnxscript/ir/serde.py	Implements the logic to correctly handle quant_parameter_tensor_names in graphs.
onnxscript/ir/_protocols.py	Updates the documentation to reflect the new quantization_annotation field.

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (2)

onnxscript/ir/serde.py:741

• Instead of using an assert to check that the tensor name in the annotation matches the value's name, consider adding explicit error handling (e.g., raising a descriptive exception) to handle any mismatches gracefully in production.

assert proto.tensor_name == value.name

onnxscript/ir/serde.py:1284

• [nitpick] Verify that using the condition 'if input_.name not in from_.initializers' reliably prevents double-adding quantization annotations for inputs. It may be beneficial to ensure that 'from_.initializers' is always set for values requiring annotations.

if input_.name not in from_.initializers:

[titaiwangms]

I am not so familiar to quantization_annotation in ONNX, and can't find more interpretation in this PR. Based on what I found in ONNX, it belongs to graph(https://github.com/search?q=repo%3Aonnx%2Fonnx%20quantization_annotation&type=code)? I am not sure if there is a better place for it in IR. cc @gramalingam

Approving this to unblock for the time being, also this does not seem like an irreversible change if needed.

[justinchuby]

@titaiwangms Thanks. I updated the pr field to explain the design choices. Essentially the info belong to individual values but is stored as a list in the GraphProto like ValueInfoProtos.