[4/17] Private ordered DLA scaffold and node ordering

modelopt/onnx/graph_surgery/dla_transforms/dla_cast_to_fp32.py

@@ -0,0 +1,184 @@
+# 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.
+
+"""Redirect eligible standard Cast targets to FP32 using ONNX type constraints.
+
+Shared type variables connect operands and outputs across operations. Fixed
+non-FP32 operands, integer or boolean controls, opaque operations and lexical
+captures prevent conversion of their connected type groups. Compatible Casts
+can convert together; their connected tensor metadata is updated as well.
+"""
+
+from __future__ import annotations
+
+import onnx
+from onnx import TensorProto
+
+from ...logging_config import logger
+from ..utils.graph_utils import _captured_names
+
+
+def _fp32_cast_outputs(model: onnx.ModelProto) -> set[str]:
+    """Find tensor type groups that can become FP32 without changing fixed operands."""
+    graph = model.graph
+    candidates = {
+        node.output[0]
+        for node in graph.node
+        if node.op_type == "Cast"
+        and node.domain in {"", "ai.onnx"}
+        and not node.overload
+        and node.output
+        and any(attr.name == "to" and attr.i != TensorProto.FLOAT for attr in node.attribute)
+    }
+    if not candidates:
+        return set()
+    parents: dict[str, str] = {}
+
+    def find(name: str) -> str:
+        """Find and compress a tensor's shared-type component."""
+        parents.setdefault(name, name)
+        while parents[name] != name:
+            parents[name] = parents[parents[name]]
+            name = parents[name]
+        return name
+
+    blocked: set[str] = set()
+    types = {
+        value.name: value.type.tensor_type.elem_type
+        for value in (*graph.input, *graph.value_info, *graph.output)
+        if value.type.HasField("tensor_type")
+    }
+    types.update((value.name, value.data_type) for value in graph.initializer)
+    fixed = {value.name for value in graph.input} | {value.name for value in graph.initializer}
+    fixed.update(value.values.name for value in graph.sparse_initializer)
+    types.update((value.values.name, value.values.data_type) for value in graph.sparse_initializer)
+    opsets = {opset.domain: opset.version for opset in model.opset_import}
+    opsets[""] = opsets.get("", opsets.get("ai.onnx", 1))
+    schemas = {}
+    for node in graph.node:
+        for attribute in node.attribute:
+            subgraphs = (
+                [attribute.g] if attribute.type == onnx.AttributeProto.GRAPH else attribute.graphs
+            )
+            for subgraph in subgraphs:
+                blocked.update(_captured_names(subgraph))
+        domain = "" if node.domain == "ai.onnx" else node.domain
+        key = (domain, node.op_type)
+        if key not in schemas:
+            try:
+                schemas[key] = onnx.defs.get_schema(node.op_type, opsets.get(domain, 1), domain)
+            except onnx.defs.SchemaError:
+                schemas[key] = None
+        schema = schemas[key]
+        if (
+            schema is None
+            or node.overload
+            or any(
+                parameter.option == onnx.defs.OpSchema.FormalParameterOption.Variadic
+                and not parameter.is_homogeneous
+                for parameter in (*schema.inputs, *schema.outputs)
+            )
+        ):
+            blocked.update(name for name in (*node.input, *node.output) if name)
+            continue
+        if (
+            domain == ""
+            and node.op_type == "Mod"
+            and schema.since_version < 28
+            and not any(
+                attribute.name == "fmod" and attribute.i == 1 for attribute in node.attribute
+            )
+        ):
+            # Older Mod schemas require integer inputs in the default attribute mode.
+            blocked.update(name for name in (*node.input, *node.output) if name)
+            continue
+        allowed = {
+            constraint.type_param_str: constraint.allowed_type_strs
+            for constraint in schema.type_constraints
+        }
+        groups: dict[str, list[str]] = {}
+        input_types: set[str] = set()
+        for values, parameters, is_output in (
+            (node.input, schema.inputs, False),
+            (node.output, schema.outputs, True),
+        ):
+            for index, name in enumerate(values):
+                if not name:
+                    continue
+                parameter = parameters[min(index, len(parameters) - 1)]
+                type_name = parameter.type_str
+                groups.setdefault(type_name, []).append(name)
+                if "tensor(float)" not in allowed.get(type_name, [type_name]):
+                    blocked.add(name)
+                if not is_output:
+                    input_types.add(type_name)
+                elif type_name not in input_types:
+                    fixed.add(name)
+                    if type_name == "tensor(float)":
+                        types[name] = TensorProto.FLOAT
+                    elif type_name in allowed and node.op_type != "Cast":
+                        # Independent output variables may still default to an input's dtype.
+                        blocked.update(name for name in node.input if name)
+        for names in groups.values():
+            for name in names[1:]:
+                parents[find(name)] = find(names[0])
+        if node.op_type == "Cast" and domain == "":
+            for attribute in node.attribute:
+                if attribute.name == "to":
+                    types[node.output[0]] = attribute.i
+        elif node.op_type == "Constant" and domain == "":
+            for attribute in node.attribute:
+                if attribute.name == "value":
+                    types[node.output[0]] = attribute.t.data_type
+        # Runtime attention inputs have implementation-specific precision requirements.
+        if node.op_type == "GroupQueryAttention":
+            blocked.update(name for name in node.input if name)
+    blocked.update(name for name in fixed - candidates if types.get(name) != TensorProto.FLOAT)
+    blocked_roots = {find(name) for name in blocked}
+    promoted_roots = {find(name) for name in candidates if find(name) not in blocked_roots}
+    return {name for name in parents if find(name) in promoted_roots}
+
+
+def _apply_cast_to_fp32(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Change every eligible Cast node's ``to`` attribute to FLOAT (float32).
+
+    Preserve connected type requirements, custom operations and subgraph captures.
+    Update intermediate and graph-output descriptors for converted type groups.
+    """
+    graph = model.graph
+    promoted = _fp32_cast_outputs(model)
+    value_info_by_name = {}
+    for value_info in graph.value_info:
+        value_info_by_name.setdefault(value_info.name, value_info)
+    outputs_by_name = {output.name: output for output in graph.output}
+    cnt = 0
+    for node in graph.node:
+        if node.op_type != "Cast":
+            continue
+        if not node.output or node.output[0] not in promoted:
+            continue
+        for attr in node.attribute:
+            if attr.name == "to" and attr.i != TensorProto.FLOAT:
+                attr.i = TensorProto.FLOAT
+                cnt += 1
+                break
+
+    for name in promoted:
+        for value_info in (value_info_by_name.get(name), outputs_by_name.get(name)):
+            if value_info is not None:
+                value_info.type.tensor_type.elem_type = TensorProto.FLOAT
+
+    logger.debug("cast_to_fp32: redirected %d Cast node(s) to float32.", cnt)
+    return model

modelopt/onnx/graph_surgery/dla_transforms/dla_reorder_nodes.py

@@ -0,0 +1,160 @@
+# 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.
+
+"""Final priority topological node reorder for DLA ONNX models."""
+
+from __future__ import annotations
+
+import heapq
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import onnx
+
+from ..utils.graph_utils import _node_inputs
+
+PRIORITIZED_OPS = {
+    "Constant",
+    "ConstantOfShape",
+    "QuantizeLinear",
+    "DequantizeLinear",
+}
+
+
+def _node_label(node: onnx.NodeProto, index: int) -> str:
+    """Identify a node in an error even when its optional name is absent."""
+    return node.name or f"<unnamed {node.op_type} #{index}>"
+
+
+def _build_edges(nodes: list[onnx.NodeProto]) -> tuple[list[set[int]], list[list[int]]]:
+    """Collect unique dependencies, including values captured by nested subgraphs."""
+    output_to_node: dict[str, int] = {}
+    for index, node in enumerate(nodes):
+        for output in node.output:
+            if output:
+                output_to_node[output] = index
+
+    dependencies = [set() for _ in nodes]
+    consumers: list[list[int]] = [[] for _ in nodes]
+    seen_edges: set[tuple[int, int]] = set()
+    for consumer_index, node in enumerate(nodes):
+        for input_name in _node_inputs(node):
+            producer_index = output_to_node.get(input_name)
+            if producer_index is None:
+                continue
+
+            edge = (producer_index, consumer_index)
+            if edge in seen_edges:
+                continue
+
+            seen_edges.add(edge)
+            dependencies[consumer_index].add(producer_index)
+            consumers[producer_index].append(consumer_index)
+
+    return dependencies, consumers
+
+
+def _unresolved_dependency_error(nodes: list[onnx.NodeProto], indegrees: list[int]) -> ValueError:
+    """Describe nodes left blocked after a topological traversal."""
+    remaining = [
+        _node_label(nodes[index], index) for index, indegree in enumerate(indegrees) if indegree > 0
+    ]
+    return ValueError(
+        "Graph contains a cycle or unresolved dependency involving: " + ", ".join(remaining[:10])
+    )
+
+
+def _compute_depths(
+    nodes: list[onnx.NodeProto],
+    dependencies: list[set[int]],
+    consumers: list[list[int]],
+) -> list[int]:
+    """Compute each node's longest dependency depth and reject cycles."""
+    depths = [0] * len(nodes)
+    indegrees = [len(node_dependencies) for node_dependencies in dependencies]
+    ready = [index for index, indegree in enumerate(indegrees) if indegree == 0]
+
+    processed = 0
+    read_index = 0
+    while read_index < len(ready):
+        index = ready[read_index]
+        read_index += 1
+        processed += 1
+
+        for consumer_index in consumers[index]:
+            depths[consumer_index] = max(depths[consumer_index], depths[index] + 1)
+            indegrees[consumer_index] -= 1
+            if indegrees[consumer_index] == 0:
+                ready.append(consumer_index)
+
+    if processed != len(nodes):
+        raise _unresolved_dependency_error(nodes, indegrees)
+
+    return depths
+
+
+def priority_topological_order(
+    graph: onnx.GraphProto,
+) -> tuple[list[int], list[int], list[int]]:
+    """Return node indices in DLA priority topological order.
+
+    The order always respects data dependencies. Among ready nodes, generic ops
+    keep the lowest priority, while Constant/QDQ-style support nodes are delayed
+    toward their first consumer depth so they appear nearer to their use.
+    """
+    nodes = list(graph.node)
+    dependencies, consumers = _build_edges(nodes)
+    depths = _compute_depths(nodes, dependencies, consumers)
+    sink_priority = max(depths, default=0) + 1
+
+    def node_priority(index: int) -> int:
+        """Delay support nodes toward their first consumer's dependency depth."""
+        if nodes[index].op_type not in PRIORITIZED_OPS:
+            return 0
+        if consumers[index]:
+            return depths[consumers[index][0]]
+        return sink_priority
+
+    indegrees = [len(node_dependencies) for node_dependencies in dependencies]
+    ready: list[tuple[int, int]] = []
+    for index, indegree in enumerate(indegrees):
+        if indegree == 0:
+            heapq.heappush(ready, (node_priority(index), index))
+
+    order: list[int] = []
+    while ready:
+        _, index = heapq.heappop(ready)
+        order.append(index)
+
+        for consumer_index in consumers[index]:
+            indegrees[consumer_index] -= 1
+            if indegrees[consumer_index] == 0:
+                heapq.heappush(ready, (node_priority(consumer_index), consumer_index))
+
+    if len(order) != len(nodes):
+        raise _unresolved_dependency_error(nodes, indegrees)
+
+    priorities = [node_priority(index) for index in range(len(nodes))]
+    return order, depths, priorities
+
+
+def _apply_reorder_nodes(model: onnx.ModelProto) -> onnx.ModelProto:
+    """Reorder ``graph.node`` in-place using DLA priority topological order."""
+    order, _, _ = priority_topological_order(model.graph)
+    nodes = list(model.graph.node)
+    sorted_nodes = [nodes[index] for index in order]
+    del model.graph.node[:]
+    model.graph.node.extend(sorted_nodes)
+    return model