[5/17] DLA QDQ and constant preprocessing

modelopt/onnx/graph_surgery/dla_transforms/dla_constants_to_initializers.py

@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+
+"""Fold ``Constant`` ops into graph ``initializer`` entries."""
+
+from __future__ import annotations
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+from ...logging_config import logger
+from ..utils.graph_utils import add_unique_initializers
+
+
+def _tensor_from_constant_node(node: onnx.NodeProto) -> onnx.TensorProto:
+    """Build a dense ``TensorProto`` initializer from a ``Constant`` node's attributes."""
+    if not node.output:
+        raise ValueError(f"Constant node {node.name!r} has no output")
+    out_name = node.output[0]
+
+    for attr in node.attribute:
+        if attr.name == "value":
+            arr = numpy_helper.to_array(attr.t)
+            return numpy_helper.from_array(arr, name=out_name)
+        if attr.name == "value_float":
+            arr = np.asarray(attr.f, dtype=np.float32).reshape(())
+            return numpy_helper.from_array(arr, name=out_name)
+        if attr.name == "value_floats":
+            arr = np.asarray(list(attr.floats), dtype=np.float32)
+            return numpy_helper.from_array(arr, name=out_name)
+        if attr.name == "value_int":
+            arr = np.asarray(np.int64(attr.i), dtype=np.int64).reshape(())
+            return numpy_helper.from_array(arr, name=out_name)
+        if attr.name == "value_ints":
+            arr = np.asarray(list(attr.ints), dtype=np.int64)
+            return numpy_helper.from_array(arr, name=out_name)
+        if attr.name == "sparse_value":
+            sparse = attr.sparse_tensor
+            values = numpy_helper.to_array(sparse.values)
+            indices = numpy_helper.to_array(sparse.indices)
+            arr = np.zeros(sparse.dims, dtype=values.dtype)
+            if indices.ndim == 1:
+                arr.reshape(-1)[indices] = values
+            else:
+                arr[tuple(indices.T)] = values
+            return numpy_helper.from_array(arr, name=out_name)
+
+    attr_names = [a.name for a in node.attribute]
+    raise ValueError(
+        f"Constant node {node.name!r} has no supported value attribute "
+        f"(expected value, value_float(s), value_int(s), or sparse_value); got {attr_names}"
+    )
+
+
+def _add_legacy_initializer_inputs(model: onnx.ModelProto) -> None:
+    """Declare initializer defaults as graph inputs for ONNX IR versions through 3."""
+    if model.ir_version > 3:
+        return
+    graph = model.graph
+    input_names = {value.name for value in graph.input}
+    for init in graph.initializer:
+        if init.name not in input_names:
+            graph.input.append(
+                helper.make_tensor_value_info(init.name, init.data_type, list(init.dims))
+            )
+            input_names.add(init.name)
+
+
+def _constants_to_initializers_from_model(
+    model: onnx.ModelProto,
+    *,
+    verbose: bool,
+) -> int:
+    """Fold supported Constant nodes in place and return their conversion count."""
+    graph = model.graph
+    unique: set[str] = {i.name for i in graph.initializer}
+    nodes_to_remove: list[onnx.NodeProto] = []
+
+    for node in list(graph.node):
+        if node.op_type != "Constant" or node.domain not in {"", "ai.onnx"} or node.overload:
+            continue
+        if not node.output:
+            continue
+        try:
+            init = _tensor_from_constant_node(node)
+        except ValueError as exc:
+            if verbose:
+                logger.warning("%s", exc)
+            continue
+        before = len(unique)
+        add_unique_initializers(graph, unique, [init])
+        if len(unique) > before:
+            nodes_to_remove.append(node)
+
+    for node in nodes_to_remove:
+        graph.node.remove(node)
+    if nodes_to_remove:
+        _add_legacy_initializer_inputs(model)
+
+    if verbose:
+        logger.debug("Converted %d Constant nodes to initializers", len(nodes_to_remove))
+    return len(nodes_to_remove)

modelopt/onnx/graph_surgery/dla_transforms/dla_handle_qdq.py

@@ -0,0 +1,156 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+
+"""Wrap every QuantizeLinear / DequantizeLinear node with Unsqueeze → op → Squeeze.
+
+Both Q and DQ are handled independently regardless of whether their data input
+is a static initializer or a dynamic tensor.  Scale and zero-point inputs are
+left untouched.  The ``axis`` attribute is shifted by ``4 - orig_rank`` to
+keep per-channel quantization referencing the correct dimension.
+Blocked quantization is left unchanged because its scales track the data rank.
+
+Pattern applied to each node::
+
+    data [*orig_shape]
+    → Unsqueeze  → data_4d [1,…,1, *orig_shape]
+    → Q / DQ     → out_4d  [1,…,1, *orig_shape]
+    → Squeeze    → out     [*orig_shape]
+"""
+
+from __future__ import annotations
+
+import onnx
+from onnx import helper
+
+from ...logging_config import logger
+from ..utils.graph_utils import (
+    GraphCache,
+    add_value_info,
+    axes_for_rank,
+    get_node_attr_i,
+    pad4d,
+    set_node_attr_i,
+)
+from .dla_constants_to_initializers import _add_legacy_initializer_inputs
+from .dla_graph_cleanup import _set_axes
+from .dla_remove_qdq import _qdq_output_dtype
+
+_axes_for_rank = axes_for_rank
+
+
+def _apply_handle_qdq(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Wrap eligible Q / DQ nodes with Unsqueeze → node → Squeeze."""
+    graph = model.graph
+    cache = GraphCache(graph)
+    versions = {opset.domain: opset.version for opset in model.opset_import}
+    opset = versions.get("", versions.get("ai.onnx", 1))
+
+    replacements: dict[int, list[onnx.NodeProto]] = {}
+    cnt = [0]
+
+    for node_pos, node in enumerate(list(graph.node)):
+        if node.domain not in {"", "ai.onnx", "com.microsoft"} or node.overload:
+            continue
+        if node.op_type not in ("QuantizeLinear", "DequantizeLinear"):
+            continue
+        if not node.input or not node.output:
+            continue
+        if get_node_attr_i(node, "block_size", 0) > 0:
+            continue
+
+        data_name = node.input[0]
+        orig_shape: list | None = cache.get_shape(data_name)
+        if orig_shape is None:
+            info = cache.vi_map.get(data_name)
+            if info is not None and info.type.tensor_type.HasField("shape"):
+                orig_shape = [
+                    dim.dim_value if dim.HasField("dim_value") else dim.dim_param or None
+                    for dim in info.type.tensor_type.shape.dim
+                ]
+
+        if orig_shape is None:
+            logger.debug("%s %r: data input shape unknown; skipping.", node.op_type, node.name)
+            continue
+        if len(orig_shape) not in (1, 2, 3):
+            continue  # already 4-D, nothing to do
+
+        orig_rank = len(orig_shape)
+        delta = 4 - orig_rank
+        axes = _axes_for_rank(orig_rank)
+        pfx = node.name or f"qdq_{cnt[0]}"
+        c = cnt[0]
+        cnt[0] += 1
+
+        # ── Unsqueeze data input to 4D ────────────────────────────────────────
+        in4d = cache.unique_name(f"{pfx}_in4d_{c}")
+        unsq = helper.make_node(
+            "Unsqueeze",
+            inputs=[data_name],
+            outputs=[in4d],
+            name=cache.unique_name(f"{pfx}_unsq_{c}"),
+        )
+        _set_axes(model, cache, unsq, axes, f"{pfx}_unsq_axes_{c}")
+        in_dtype = cache.get_dtype(data_name) or onnx.TensorProto.FLOAT
+        add_value_info(graph, in4d, in_dtype, pad4d(orig_shape))
+        node.input[0] = in4d
+
+        # ── Update axis attribute (ONNX Q/DQ default: 1) ─────────────────────
+        has_axis = any(a.name == "axis" for a in node.attribute)
+        if (opset >= 13 or node.domain == "com.microsoft") and (
+            has_axis or (node.domain != "com.microsoft" and len(orig_shape) > 1)
+        ):
+            orig_axis = get_node_attr_i(node, "axis", 1)
+            if orig_axis >= 0:
+                set_node_attr_i(node, "axis", orig_axis + delta)
+
+        # ── Squeeze output back to original rank ──────────────────────────────
+        orig_out = node.output[0]
+        out4d = cache.unique_name(f"{pfx}_out4d_{c}")
+        out_dtype = cache.get_dtype(orig_out) or _qdq_output_dtype(node, cache)
+        if out_dtype:
+            add_value_info(graph, out4d, out_dtype, pad4d(orig_shape))
+        node.output[0] = out4d
+
+        c2 = cnt[0]
+        cnt[0] += 1
+        sq = helper.make_node(
+            "Squeeze",
+            inputs=[out4d],
+            outputs=[orig_out],
+            name=cache.unique_name(f"{pfx}_sq_{c2}"),
+        )
+        _set_axes(model, cache, sq, axes, f"{pfx}_sq_axes_{c2}")
+
+        replacements[node_pos] = [unsq, node, sq]
+        logger.debug(
+            "%s %r: wrapped Unsqueeze→op→Squeeze (rank %d→4, axis delta=%d).",
+            node.op_type,
+            node.name,
+            orig_rank,
+            delta,
+        )
+
+    if replacements:
+        nodes = [
+            replacement
+            for index, node in enumerate(graph.node)
+            for replacement in replacements.get(index, [node])
+        ]
+        del graph.node[:]
+        graph.node.extend(nodes)
+        _add_legacy_initializer_inputs(model)
+
+    logger.debug("handle_qdq: wrapped %d Q/DQ node(s) with Unsqueeze/Squeeze.", len(replacements))
+    return model