From d50e6764bdbb04f36720d6c3e4ce25874d50df40 Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 15 Oct 2026 12:46:10 +0000 Subject: [PATCH 1/7] [DLA_transforms_8] [8/17] DLA MatMul and Gemm convolution conversion Signed-off-by: mgohil --- .../dla_matmul_to_transpose_conv_transpose.py | 483 ++++++++++++++++++ .../onnx/graph_surgery/make_dla_compatible.py | 8 + ..._dla_matmul_to_transpose_conv_transpose.py | 442 ++++++++++++++++ 3 files changed, 933 insertions(+) create mode 100644 modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py create mode 100644 tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py new file mode 100644 index 00000000000..830ac2e34aa --- /dev/null +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -0,0 +1,483 @@ +# 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. + +"""Replace MatMul/Gemm nodes with an Unsqueeze-Transpose-Conv-Transpose-Squeeze subgraph. + +For nodes that pass eligibility (see below), the full Conv conversion is applied. +For **ineligible** nodes, a lightweight 4-D promotion fallback is used instead: +dynamic inputs are Unsqueezed to 4-D, static initializers are padded with leading +1-dims in-place, and a Squeeze is inserted after the node to restore the original +output rank. This ensures all MatMul/Gemm tensors are 4-D after this transform. + +Eligibility (checked via :func:`check_to_apply_transpose`): +- The **second input** (weight) must be a constant initializer or a Const→DequantizeLinear chain. + Dynamic / graph-input weights are NOT eligible for Conv conversion. +- Weight shape must be 2-D, 3-D with ``shape[0]==1``, or 4-D with ``shape[0]==1, shape[1]==1``. +- The **first input** (activation) shape must be known and be 2-D, 3-D, or 4-D with + ``shape[0]==1``. + +Shape contract after transform (per activation rank): + + +--------------+---------------+--------------------------------------+--------------+-------------+ + | Orig rank | Unsqueeze | After input Transpose (→ Conv input) | Conv output | Squeeze | + +==============+===============+======================================+==============+=============+ + | 2D [M,N] | [0,1]→[1,1,M,N]| perm[0,3,2,1]→[1,N,M,1] | [1,K,M,1] | [0,1]→[M,K] | + | 3D [B,M,N] | [0]→[1,B,M,N] | perm[0,3,2,1]→[1,N,M,B] | [1,K,M,B] | [0]→[B,M,K] | + | 4D [1,B,M,N] | none | perm[0,3,2,1]→[1,N,M,B] | [1,K,M,B] | none | + +--------------+---------------+--------------------------------------+--------------+-------------+ + +Output Transpose always uses perm ``[0,3,2,1]``: ``[1,K,M,B] → [1,B,M,K]``. +""" + +import numpy as np +from onnx import helper, numpy_helper + +from ...logging_config import logger +from ..utils.graph_utils import ( + GraphCache, + add_unique_initializers, + batch_replace_nodes, + check_to_apply_transpose, + is_const_dq_input, + pad4d, +) + + +def _apply_matmul_to_transpose_conv_transpose(model, *, defer_rank_conversion=False): + """Convert MatMul/Gemm; optionally let the pipeline promote tensor ranks later.""" + graph = model.graph + cache = GraphCache(graph) + unique_initializer: set = set() + + # shape map snapshot required by check_to_apply_transpose + tensor_name_dim_map: dict = {} + for vi in list(graph.value_info) + list(graph.input) + list(graph.output): + s = cache.get_shape(vi.name) + if s is not None: + tensor_name_dim_map[vi.name] = s + + node_idx = {id(n): i for i, n in enumerate(graph.node)} + replacements_batch: dict[int, tuple[list, list]] = {} + fallback_batch: dict[int, tuple[list, list]] = {} + for candidate in graph.node: + if candidate.op_type not in {"MatMul", "Gemm"}: + continue + if candidate.op_type == "Gemm": + if any( + attribute.name in {"alpha", "beta"} and attribute.f != 1.0 + for attribute in candidate.attribute + ): + raise ValueError( + f"{candidate.name!r}: Gemm alpha/beta other than 1 are unsupported" + ) + if len(candidate.input) > 2 and candidate.input[2]: + bias_shape = cache.get_shape(candidate.input[2]) + if bias_shape is None or len(bias_shape) != 1: + raise ValueError(f"Bias for {candidate.name!r} must have a known 1-D shape") + + def update_initializers( + graph, name, current_node, nodes_to_add, trans_b=False, is_const_dq_init=False + ): + """Reshape weight initializer to a 4-D Conv kernel. + + If the initializer is shared by other consumers besides the current + rewrite, a clone (new initializer + cloned DQ if applicable) is + created so other consumers keep the original layout; this MatMul/Gemm + gets rewired to the clone. When the initializer has a single consumer + it's rewritten in place as before. + + Maps: + 2D [C,K] → [K,C,1,1] (or [C,K,1,1] when transB) + 3D [C,K,M] → [M,K,C,1] (or [C,K,M,1] when transB) + 4D [N,C,K,M] → [M,K,C,N] perm (or [C,K,M,N] when transB) + """ + for initializer in graph.initializer: + if initializer.name != name: + continue + axis = None + dq_node = None + shape = None + if is_const_dq_init: + dq_node = next( + ( + n + for n in graph.node + if n.op_type == "DequantizeLinear" and n.output[0] == current_node.input[1] + ), + None, + ) + if dq_node is not None: + axis = next((attr.i for attr in dq_node.attribute if attr.name == "axis"), None) + + # 2D [C,K] + if len(initializer.dims) == 2: + c, k = initializer.dims[0], initializer.dims[1] + init_arr = numpy_helper.to_array(initializer) + if not trans_b: + init_arr = init_arr.T + shape = (k, c, 1, 1) + if axis is not None: + axis = 1 - axis + else: + shape = (c, k, 1, 1) + + # 3D [C,K,M] + elif len(initializer.dims) == 3: + c, k, m = initializer.dims[0], initializer.dims[1], initializer.dims[2] + init_arr = numpy_helper.to_array(initializer) + if not trans_b: + init_arr = init_arr.T + shape = (m, k, c, 1) + if axis is not None: + axis = 2 - axis + else: + shape = (c, k, m, 1) + + # 4D [N,C,K,M] + elif len(initializer.dims) == 4: + n, c, k, m = ( + initializer.dims[0], + initializer.dims[1], + initializer.dims[2], + initializer.dims[3], + ) + init_arr = numpy_helper.to_array(initializer) + if not trans_b: + init_arr = np.transpose(init_arr, (3, 2, 1, 0)) + shape = (m, k, c, n) + if axis is not None: + axis = 3 - axis + else: + init_arr = np.transpose(init_arr, (1, 2, 3, 0)) + shape = (c, k, m, n) + + reshaped_arr = np.reshape(init_arr, shape) + if len(cache.get_consumers(name)) > 1 or ( + dq_node is not None + and ( + len(cache.get_consumers(dq_node.output[0])) > 1 + or dq_node.output[0] in cache.graph_output_names + ) + ): + clone_idx = len(replacements_batch) + new_init_name = cache.unique_name(f"{name}_conv4d_{clone_idx}") + new_init_proto = numpy_helper.from_array(reshaped_arr, new_init_name) + add_unique_initializers(graph, unique_initializer, [new_init_proto]) + if dq_node is not None: + new_dq_out = f"{dq_node.output[0]}_conv4d_{clone_idx}" + dq_clone = helper.make_node( + "DequantizeLinear", + inputs=[new_init_name, *list(dq_node.input[1:])], + outputs=[new_dq_out], + name=f"{dq_node.name or name}_conv4d_{clone_idx}", + ) + dq_clone.attribute.extend(dq_node.attribute) + if axis is not None: + for attr in dq_clone.attribute: + if attr.name == "axis": + attr.i = axis + nodes_to_add.append(dq_clone) + current_node.input[1] = new_dq_out + else: + current_node.input[1] = new_init_name + return + + if dq_node is not None: + for attr in dq_node.attribute: + if attr.name == "axis": + attr.i = axis + dq_node.name = f"{dq_node.name}_axis_updated" + + reshaped_arr = np.reshape(init_arr, shape) + initializer.CopyFrom(numpy_helper.from_array(reshaped_arr, initializer.name)) + + # ── Separate eligible and ineligible nodes in one pass ─────────────────── + eligible_nodes = [] + ineligible_nodes = [] + for node in graph.node: + if node.op_type not in ("MatMul", "Gemm"): + continue + if check_to_apply_transpose(node, model, tensor_name_dim_map): + eligible_nodes.append(node) + else: + ineligible_nodes.append(node) + + # ── Fallback: promote inputs/output of ineligible MatMul/Gemm to 4D ───── + # For each ineligible node: + # * Dynamic inputs → Unsqueeze to 4D. + # * Static initializer inputs (including DQ-wrapped) → expand dims in-place. + # * Output → Squeeze back to original rank when the original was not 4D. + for node in ineligible_nodes: + if defer_rank_conversion and node.input[1] not in cache.graph_input_names: + continue + cnt_fb = len(fallback_batch) + pfx = node.name or f"mm_fallback_{cnt_fb}" + initializer_names = cache.init_names + pre: list = [] + post: list = [] + + for inp_idx in range(min(2, len(node.input))): + inp = node.input[inp_idx] + if defer_rank_conversion and (inp_idx != 1 or inp not in cache.graph_input_names): + continue + if not inp: + continue + in_shape = cache.get_shape(inp) + if in_shape is None or len(in_shape) >= 4: + continue + rank = len(in_shape) + # Check whether this is a static initializer (possibly DQ-wrapped) + dq_src = is_const_dq_input(initializer_names, inp, graph) + init_name = dq_src or (inp if inp in initializer_names else None) + if init_name is not None: + if dq_src is not None: + # Preserve the quantization layout when a static DQ weight needs promotion. + clone = helper.make_node( + "Unsqueeze", + [inp, ""], + [cache.unique_name(f"{pfx}_in{inp_idx}_4d_{cnt_fb}")], + name=cache.unique_name(pfx + "_weight_unsqueeze"), + ) + axes_name = cache.unique_name(f"{pfx}_in{inp_idx}_unsq_axes_{cnt_fb}") + graph.initializer.append( + numpy_helper.from_array(np.arange(4 - rank, dtype=np.int64), axes_name) + ) + clone.input[1] = axes_name + pre.append(clone) + node.input[inp_idx] = clone.output[0] + continue + # Static weight: expand initializer dims in-place with 1-padding + for init in graph.initializer: + if init.name == init_name: + arr = numpy_helper.to_array(init) + if arr.ndim < 4: + new_shape = pad4d(list(arr.shape)) + if ( + len(cache.get_consumers(init.name)) > 1 + or init.name in cache.graph_output_names + ): + clone = type(init)() + clone.CopyFrom(init) + clone.name = cache.unique_name(init.name + "_4d") + del clone.dims[:] + clone.dims.extend(new_shape) + graph.initializer.append(clone) + node.input[inp_idx] = clone.name + else: + del init.dims[:] + init.dims.extend(new_shape) + break + else: + # Dynamic input: Unsqueeze to 4D + axes = list(range(4 - rank)) + axes_init = numpy_helper.from_array( + np.array(axes, dtype=np.int64), + name=f"{pfx}_in{inp_idx}_unsq_axes_{cnt_fb}", + ) + add_unique_initializers(graph, unique_initializer, [axes_init]) + inp_4d = f"{pfx}_in{inp_idx}_4d_{cnt_fb}" + pre.append( + helper.make_node( + "Unsqueeze", + inputs=[inp, axes_init.name], + outputs=[inp_4d], + name=f"{pfx}_in{inp_idx}_unsq_{cnt_fb}", + ) + ) + node.input[inp_idx] = inp_4d + + # Output: Squeeze back to original rank + out_name = node.output[0] if node.output else None + if out_name and not defer_rank_conversion: + out_shape = cache.get_shape(out_name) + if out_shape is not None and len(out_shape) < 4: + temp_out = f"{pfx}_out_4d_{cnt_fb}" + sq_axes = list(range(4 - len(out_shape))) + sq_axes_init = numpy_helper.from_array( + np.array(sq_axes, dtype=np.int64), + name=f"{pfx}_out_sq_axes_{cnt_fb}", + ) + add_unique_initializers(graph, unique_initializer, [sq_axes_init]) + post.append( + helper.make_node( + "Squeeze", + inputs=[temp_out, sq_axes_init.name], + outputs=[out_name], + name=f"{pfx}_out_sq_{cnt_fb}", + ) + ) + node.output[0] = temp_out + + if pre or post: + fallback_batch[node_idx[id(node)]] = ([node], [*pre, node, *post]) + + for node in eligible_nodes: + cnt = len(replacements_batch) + pfx = node.name or f"matmul_{cnt}" + nodes_to_add = [] + initializer_names = cache.init_names + + # ── Gemm transA / transB attributes ────────────────────────────────── + trans_a = False + trans_b = False + if node.op_type == "Gemm": + for attr in node.attribute: + if attr.name == "transA": + trans_a = bool(attr.i) + elif attr.name == "transB": + trans_b = bool(attr.i) + + # ── Activation rank ─────────────────────────────────────────────────── + input_shape = cache.get_shape(node.input[0]) + if input_shape is None: + raise ValueError(f"Shape unknown for activation tensor {node.input[0]!r}") + orig_rank = len(input_shape) + output_shape = cache.get_shape(node.output[0]) + output_rank = len(output_shape) if output_shape is not None else orig_rank + + # ── Step 1: Unsqueeze activation to 4D ─────────────────────────────── + # 2D [M,N] → axes [0,1] → [1,1,M,N] + # 3D [B,M,N] → axes [0] → [1,B,M,N] + activation_4d = node.input[0] + if orig_rank < 4 and not defer_rank_conversion: + axes = [0, 1] if orig_rank == 2 else [0] + axes_init = numpy_helper.from_array( + np.array(axes, dtype=np.int64), name=f"{pfx}_unsq_axes_{cnt}" + ) + add_unique_initializers(graph, unique_initializer, [axes_init]) + activation_4d = f"{pfx}_unsqueeze4d_{cnt}" + nodes_to_add.append( + helper.make_node( + "Unsqueeze", + inputs=[node.input[0], axes_init.name], + outputs=[activation_4d], + name=f"{pfx}_unsqueeze4d_node_{cnt}", + ) + ) + + # ── Step 2: Transpose activation (guaranteed 4D at this point) ─────── + # transA=False: [1,B,M,N] → perm[0,3,2,1] → [1,N,M,B] (C_in=N, H=M, W=B) + # transA=True : [1,B,N,M] → perm[0,2,3,1] → [1,N,M,B] (same result) + perm_a = [0, 2, 3, 1] if trans_a else [0, 3, 2, 1] + transposed_a = f"{pfx}_tr_before0_{cnt}" + nodes_to_add.append( + helper.make_node( + "Transpose", + inputs=[activation_4d], + outputs=[transposed_a], + name=f"{pfx}_transpose_before0_{cnt}", + perm=perm_a, + ) + ) + + # ── Step 3: Prepare weight ──────────────────────────────────────────── + dq_init_name = is_const_dq_input(initializer_names, node.input[1], graph) + is_const_dq = dq_init_name is not None + init_name = ( + dq_init_name + if is_const_dq + else (node.input[1] if node.input[1] in initializer_names else None) + ) + + if init_name is not None: + # Static weight: reshape initializer to 4-D conv kernel format + update_initializers(graph, init_name, node, nodes_to_add, trans_b, is_const_dq) + conv_input1 = node.input[1] # tensor name unchanged; dims updated + else: + # Dynamic weight: add a Transpose node before Conv + # (check_to_apply_transpose normally prevents reaching here) + perm_b = [2, 3, 0, 1] if trans_b else [3, 2, 0, 1] + transposed_b = f"{pfx}_tr_before1_{cnt}" + nodes_to_add.append( + helper.make_node( + "Transpose", + inputs=[node.input[1]], + outputs=[transposed_b], + name=f"{pfx}_transpose_before1_{cnt}", + perm=perm_b, + ) + ) + conv_input1 = transposed_b + + # ── Step 4: Build Conv inputs (include optional Gemm bias) ─────────── + conv_inputs = [transposed_a, conv_input1] + if len(node.input) >= 3 and node.input[2]: + bias_shape = cache.get_shape(node.input[2]) + if bias_shape is not None and len(bias_shape) != 1: + raise ValueError( + f"[matmul_to_conv] Bias for {node.name!r} must be 1-D, got {bias_shape}" + ) + conv_inputs.append(node.input[2]) + + # ── Step 5: Conv ────────────────────────────────────────────────────── + conv_out = f"{pfx}_conv_out_{cnt}" + nodes_to_add.append( + helper.make_node( + "Conv", + inputs=conv_inputs, + outputs=[conv_out], + name=f"{pfx}_conv_{cnt}", + ) + ) + + # ── Step 6: Output Transpose ────────────────────────────────────────── + # [1,K,M,B] → perm[0,3,2,1] → [1,B,M,K] + # Write to a temp name when Squeeze follows; otherwise directly to node.output[0]. + tr_after_out = ( + f"{pfx}_tr_after_out_{cnt}" + if output_rank < 4 and not defer_rank_conversion + else node.output[0] + ) + nodes_to_add.append( + helper.make_node( + "Transpose", + inputs=[conv_out], + outputs=[tr_after_out], + name=f"{pfx}_transpose_after_{cnt}", + perm=[0, 3, 2, 1], + ) + ) + + # ── Step 7: Squeeze back to original rank ───────────────────────────── + # 2D case: [1,1,M,K] → squeeze [0,1] → [M,K] + # 3D case: [1,B,M,K] → squeeze [0] → [B,M,K] + if output_rank < 4 and not defer_rank_conversion: + sq_axes = list(range(4 - output_rank)) + sq_axes_init = numpy_helper.from_array( + np.array(sq_axes, dtype=np.int64), name=f"{pfx}_sq_axes_{cnt}" + ) + add_unique_initializers(graph, unique_initializer, [sq_axes_init]) + nodes_to_add.append( + helper.make_node( + "Squeeze", + inputs=[tr_after_out, sq_axes_init.name], + outputs=list(node.output), + name=f"{pfx}_squeeze_after_{cnt}", + ) + ) + + replacements_batch[node_idx[id(node)]] = ([node], nodes_to_add) + + # ── Apply all replacements ──────────────────────────────────────────────── + all_batch = {**replacements_batch, **fallback_batch} + batch_replace_nodes(graph, all_batch) + + logger.debug( + "Replaced %d MatMul/Gemm nodes with Transpose-Conv-Transpose nodes; " + "promoted %d ineligible MatMul/Gemm nodes to 4D.", + len(replacements_batch), + len(fallback_batch), + ) + return model diff --git a/modelopt/onnx/graph_surgery/make_dla_compatible.py b/modelopt/onnx/graph_surgery/make_dla_compatible.py index 100480015ed..787a17767d3 100644 --- a/modelopt/onnx/graph_surgery/make_dla_compatible.py +++ b/modelopt/onnx/graph_surgery/make_dla_compatible.py @@ -24,6 +24,7 @@ from __future__ import annotations import time +from functools import partial from typing import TYPE_CHECKING import onnx @@ -32,6 +33,9 @@ from .dla_transforms.dla_cast_to_fp32 import _apply_cast_to_fp32 from .dla_transforms.dla_constants_to_initializers import _constants_to_initializers_from_model from .dla_transforms.dla_graph_cleanup import _remove_unused_initializers +from .dla_transforms.dla_matmul_to_transpose_conv_transpose import ( + _apply_matmul_to_transpose_conv_transpose, +) from .dla_transforms.dla_remove_deqlin import _apply_remove_deqlin from .dla_transforms.dla_remove_qdq import _remove_qdq_from_model from .dla_transforms.dla_reorder_nodes import _apply_reorder_nodes @@ -77,6 +81,10 @@ def _apply_constants_to_initializers(model, **options): ("ReshapeTransposeReshape", _apply_reshape_transpose_reshape), ("ReshapeQDQReduceSum", _apply_reshape_qdq_reduce_sum), ("ReshapeReduceSum", _apply_reshape_reduce_sum), + ( + "MatMulToTransposeConvTranspose", + partial(_apply_matmul_to_transpose_conv_transpose, defer_rank_conversion=True), + ), ("ShapeInference", infer_shapes), ("RemoveUnusedInitializers", _remove_unused_initializers), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py new file mode 100644 index 00000000000..0a65faee3fd --- /dev/null +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -0,0 +1,442 @@ +# 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. + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Unit tests for :mod:`modelopt.onnx.graph_surgery.dla_transforms.dla_matmul_to_transpose_conv_transpose`. + +"""ORT CPU parity tests for MatMul/Gemm → Unsqueeze-Transpose-Conv-Transpose-Squeeze. + +Test matrix +----------- +Converts: + 2D activation [M,N] x W_init[N,K] → Y[M,K] + 3D activation [B,M,N] x W_init[N,K] → Y[B,M,K] + 4D activation [1,B,M,N] x W_init[N,K] → Y[1,B,M,K] (no Unsqueeze/Squeeze) + Gemm transB=1 A[M,N] x W_init[K,N]^T → Y[M,K] + Gemm transA=1 A[N,M]^T x W_init[N,K] → Y[M,K] + Gemm bias A[M,N] x W_init[N,K] + B[K] → Y[M,K] + DQ weight A[M,N] x dequant(W_q_int8[N,K]) → Y[M,K] (init→DQ chain) + +Does NOT convert: + Dynamic weight A[M,N] x W_graph_input[N,K] → MatMul left unchanged + NCHW non-unary A[2,3,M,N] x W_init[2,3,N,K] → MatMul left unchanged (4D weight, N/C > 1) + +Each test saves a base + transformed ONNX pair under +``scratch_space/matmul_conv_test_models/`` for offline inspection. +""" + +from __future__ import annotations + +import copy +from pathlib import Path + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from modelopt.onnx.graph_surgery.dla_transforms.dla_matmul_to_transpose_conv_transpose import ( + _apply_matmul_to_transpose_conv_transpose as _apply_matmul_to_conv, +) + +pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") + +_REPO_ROOT = Path(__file__).resolve().parents[5] +_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "matmul_conv_test_models" +_OPSET_VERSION = 17 +_MAX_IR_VERSION_FOR_ORT = 9 + +# Small but non-trivial dims: activation [M,N], weight [N,K], batch B +M, N, K, B = 8, 16, 12, 3 + + +# ── Shared test utilities ───────────────────────────────────────────────────────────────────────── + + +def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + return numpy_helper.from_array(arr, name=name) + + +def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + return helper.make_tensor_value_info(name, elem_type, shape) + + +def _clamp_ir_version(model: onnx.ModelProto) -> None: + model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) + + +def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: + import onnxruntime as ort # type: ignore[import-not-found] + + _clamp_ir_version(model) + sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} + + +def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + a64 = np.asarray(a, dtype=np.float64).ravel() + b64 = np.asarray(b, dtype=np.float64).ravel() + na, nb = float(np.linalg.norm(a64)), float(np.linalg.norm(b64)) + if na == 0.0 and nb == 0.0: + return 1.0 + if na == 0.0 or nb == 0.0: + return 0.0 + return float(np.dot(a64, b64) / (na * nb)) + + +def _assert_outputs_match(ref: dict, got: dict) -> None: + assert ref.keys() == got.keys() + for name in ref: + a, b = ref[name], got[name] + assert a.shape == b.shape, f"{name}: shape {a.shape} vs {b.shape}" + cos = _cosine_similarity(a, b) + assert cos >= 1.0 - 1e-5, f"{name}: cosine similarity {cos:.8f} (expected ~1.0)" + np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=f"{name}: value mismatch") + + +def _save_test_pair(base: onnx.ModelProto, transformed: onnx.ModelProto, tag: str) -> None: + _SAVED_MODELS_DIR.mkdir(parents=True, exist_ok=True) + onnx.save(base, str(_SAVED_MODELS_DIR / f"matmul_base_{tag}.onnx")) + import contextlib + + with contextlib.suppress(Exception): + transformed = onnx.shape_inference.infer_shapes(copy.deepcopy(transformed)) + onnx.save(transformed, str(_SAVED_MODELS_DIR / f"matmul_transformed_{tag}.onnx")) + + +def _make_model(nodes, graph_name, inputs, outputs, initializers=None): + """Build, shape-infer, and checker-validate a simple ONNX model.""" + graph = helper.make_graph(nodes, graph_name, inputs, outputs, initializer=initializers or []) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) + _clamp_ir_version(model) + try: + from onnxruntime.tools.symbolic_shape_infer import ( + SymbolicShapeInference, # type: ignore[import-untyped] + ) + + model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model + except Exception: + pass + onnx.checker.check_model(model) + return model + + +_RNG = np.random.default_rng(42) + + +# ── Model builders ──────────────────────────────────────────────────────────────────────────────── + + +def _make_matmul_2d(): + """MatMul: A[M,N] x W[N,K] → Y[M,K]. W is a float32 initializer.""" + w = _RNG.standard_normal((N, K)).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="mm")], + "mm_2d", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((M, N)).astype(np.float32)} + return model, feeds + + +def _make_matmul_3d(): + """MatMul: A[B,M,N] x W[N,K] → Y[B,M,K]. W is a float32 initializer.""" + w = _RNG.standard_normal((N, K)).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="mm")], + "mm_3d", + [_vi("A", TensorProto.FLOAT, [B, M, N])], + [_vi("Y", TensorProto.FLOAT, [B, M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((B, M, N)).astype(np.float32)} + return model, feeds + + +def _make_matmul_4d(): + """MatMul: A[1,B,M,N] x W[N,K] → Y[1,B,M,K]. W is a float32 initializer.""" + w = _RNG.standard_normal((N, K)).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="mm")], + "mm_4d", + [_vi("A", TensorProto.FLOAT, [1, B, M, N])], + [_vi("Y", TensorProto.FLOAT, [1, B, M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((1, B, M, N)).astype(np.float32)} + return model, feeds + + +def _make_gemm_transb(): + """Gemm transB=1: A[M,N] x W[K,N]^T → Y[M,K]. W is a float32 initializer.""" + w = _RNG.standard_normal((K, N)).astype(np.float32) + model = _make_model( + [helper.make_node("Gemm", ["A", "W"], ["Y"], name="gemm", transB=1)], + "gemm_transb", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((M, N)).astype(np.float32)} + return model, feeds + + +def _make_gemm_transa(): + """Gemm transA=1: A[N,M]^T x W[N,K] → Y[M,K]. W is a float32 initializer.""" + w = _RNG.standard_normal((N, K)).astype(np.float32) + model = _make_model( + [helper.make_node("Gemm", ["A", "W"], ["Y"], name="gemm", transA=1)], + "gemm_transa", + [_vi("A", TensorProto.FLOAT, [N, M])], # A is [N,M]; transA makes it [M,N] + [_vi("Y", TensorProto.FLOAT, [M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((N, M)).astype(np.float32)} + return model, feeds + + +def _make_gemm_bias(): + """Gemm: A[M,N] x W[N,K] + Bias[K] → Y[M,K]. W and Bias are float32 initializers.""" + w = _RNG.standard_normal((N, K)).astype(np.float32) + bias = _RNG.standard_normal((K,)).astype(np.float32) + model = _make_model( + [helper.make_node("Gemm", ["A", "W", "Bias"], ["Y"], name="gemm")], + "gemm_bias", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + [_init("W", w), _init("Bias", bias)], + ) + feeds = {"A": _RNG.standard_normal((M, N)).astype(np.float32)} + return model, feeds + + +def _make_matmul_dq_weight(): + """MatMul: A[M,N] x dequant(W_q[N,K]) → Y[M,K]. Per-tensor int8 weight via DQ chain.""" + w_q = _RNG.integers(-127, 127, (N, K)).astype(np.int8) + scale = np.float32(0.05) + zp = np.int8(0) + nodes = [ + helper.make_node("DequantizeLinear", ["W_q", "W_scale", "W_zp"], ["W_dq"], name="dq_W"), + helper.make_node("MatMul", ["A", "W_dq"], ["Y"], name="mm"), + ] + model = _make_model( + nodes, + "mm_dq", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + [_init("W_q", w_q), _init("W_scale", np.array(scale)), _init("W_zp", np.array(zp))], + ) + feeds = {"A": _RNG.standard_normal((M, N)).astype(np.float32)} + return model, feeds + + +def _make_matmul_dynamic_weight(): + """MatMul: A[M,N] x W[N,K] → Y[M,K]. W is a graph input — should NOT be converted.""" + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="mm")], + "mm_dyn", + [_vi("A", TensorProto.FLOAT, [M, N]), _vi("W", TensorProto.FLOAT, [N, K])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + ) + feeds = { + "A": _RNG.standard_normal((M, N)).astype(np.float32), + "W": _RNG.standard_normal((N, K)).astype(np.float32), + } + return model, feeds + + +def _make_matmul_nchw_non_unary(): + """Batch MatMul: A[2,3,M,N] x W[2,3,N,K] → Y[2,3,M,K]. + + W is a 4D initializer with shape[0]=2 (N > 1) — check_to_apply_transpose returns False, + so the MatMul must remain unchanged. + """ + c1, c2 = 2, 3 + w = _RNG.standard_normal((c1, c2, N, K)).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="mm")], + "mm_nchw", + [_vi("A", TensorProto.FLOAT, [c1, c2, M, N])], + [_vi("Y", TensorProto.FLOAT, [c1, c2, M, K])], + [_init("W", w)], + ) + feeds = {"A": _RNG.standard_normal((c1, c2, M, N)).astype(np.float32)} + return model, feeds + + +# ── Helper: apply transform and return (ref_outputs, got_outputs, transformed_model) ───────────── + + +def _apply_and_run(model, feeds): + ref = _run_ort(model, feeds) + transformed = copy.deepcopy(model) + transformed = _apply_matmul_to_conv(transformed) + _clamp_ir_version(transformed) + got = _run_ort(transformed, feeds) + return ref, got, transformed + + +# ── Tests: cases that SHOULD convert ───────────────────────────────────────────────────────────── + + +def test_matmul_2d_activation_converts(): + """2D activation → Unsqueeze[0,1] + Conv + Squeeze[0,1]; output must match MatMul.""" + model, feeds = _make_matmul_2d() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node), ( + "MatMul should be replaced" + ) + assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 2D input" + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 2D input" + + _save_test_pair(model, tr, "2d") + _assert_outputs_match(ref, got) + + +def test_matmul_3d_activation_converts(): + """3D activation [B,M,N] → Unsqueeze[0] + Conv + Squeeze[0]; output must match MatMul.""" + model, feeds = _make_matmul_3d() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 3D input" + assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 3D input" + + _save_test_pair(model, tr, "3d") + _assert_outputs_match(ref, got) + + +def test_matmul_4d_activation_converts(): + """4D activation [1,B,M,N] → Conv only, NO Unsqueeze/Squeeze; output must match MatMul.""" + model, feeds = _make_matmul_4d() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 + assert not any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + "No Unsqueeze expected for 4D input" + ) + assert not any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "No Squeeze expected for 4D input" + ) + + _save_test_pair(model, tr, "4d") + _assert_outputs_match(ref, got) + + +def test_gemm_transb_converts(): + """Gemm transB=1 (W[K,N] → used without extra transpose): output must match Gemm.""" + model, feeds = _make_gemm_transb() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + + _save_test_pair(model, tr, "gemm_transb") + _assert_outputs_match(ref, got) + + +def test_gemm_transa_converts(): + """Gemm transA=1 (A[N,M] used transposed): output must match Gemm.""" + model, feeds = _make_gemm_transa() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + + _save_test_pair(model, tr, "gemm_transa") + _assert_outputs_match(ref, got) + + +def test_gemm_with_bias_converts(): + """Gemm with 1D bias: bias forwarded to Conv as third input; output must match Gemm.""" + model, feeds = _make_gemm_bias() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + conv_nodes = [n for n in tr.graph.node if n.op_type == "Conv"] + assert len(conv_nodes) == 1 + assert len(conv_nodes[0].input) == 3, "Conv should carry bias as third input" + + _save_test_pair(model, tr, "gemm_bias") + _assert_outputs_match(ref, got) + + +def test_matmul_dq_weight_converts(): + """Int8 init→DQ weight: W_q reshaped to [K,N,1,1], DQ node preserved, Conv uses DQ output.""" + model, feeds = _make_matmul_dq_weight() + ref, got, tr = _apply_and_run(model, feeds) + + assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) + assert any(n.op_type == "DequantizeLinear" for n in tr.graph.node), "DQ node should remain" + assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 + + # Verify W_q was reshaped to 4D [K,N,1,1] + w_q_init = next((i for i in tr.graph.initializer if i.name == "W_q"), None) + assert w_q_init is not None + assert list(w_q_init.dims) == [K, N, 1, 1], ( + f"W_q dims: {list(w_q_init.dims)}, expected [{K},{N},1,1]" + ) + + _save_test_pair(model, tr, "dq_weight") + _assert_outputs_match(ref, got) + + +# ── Tests: cases that should NOT convert ───────────────────────────────────────────────────────── + + +def test_matmul_dynamic_weight_not_converted(): + """Dynamic weight: MatMul remains; both inputs and output promoted to 4D; ORT parity.""" + model, feeds = _make_matmul_dynamic_weight() + ref, got, tr = _apply_and_run(model, feeds) + + matmul_nodes = [n for n in tr.graph.node if n.op_type == "MatMul"] + assert len(matmul_nodes) == 1, "MatMul must remain when weight is a graph input" + assert not any(n.op_type == "Conv" for n in tr.graph.node), "No Conv should be added" + # Fallback: both 2D inputs promoted → Unsqueeze nodes inserted + assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( + "Unsqueeze expected for 4D promotion of dynamic inputs" + ) + # Fallback: 2D output squeezed back + assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( + "Squeeze expected to restore original output rank" + ) + + _save_test_pair(model, tr, "dynamic_weight") + _assert_outputs_match(ref, got) + + +def test_matmul_nchw_non_unary_not_converted(): + """Batch MatMul [2,3,M,N] x W[2,3,N,K]: ineligible for Conv; inputs/output promoted to 4D.""" + model, feeds = _make_matmul_nchw_non_unary() + ref, got, tr = _apply_and_run(model, feeds) + + matmul_nodes = [n for n in tr.graph.node if n.op_type == "MatMul"] + assert len(matmul_nodes) == 1, "MatMul must remain when weight 4D shape[0] > 1" + assert not any(n.op_type == "Conv" for n in tr.graph.node), "No Conv should be added" + # Activation [2,3,M,N] is already 4D — no Unsqueeze needed for it + # Weight W[2,3,N,K] is already 4D — no init expansion needed + # Output [2,3,M,K] is already 4D — no Squeeze needed + assert not any(n.op_type in ("Unsqueeze", "Squeeze") for n in tr.graph.node), ( + "No Unsqueeze/Squeeze expected when all tensors are already 4D" + ) + + _save_test_pair(model, tr, "nchw_non_unary") + _assert_outputs_match(ref, got) From d55b4a62a2b6297de31678323dad4584dfa244b4 Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 21:01:16 +0000 Subject: [PATCH 2/7] Preserve DLA matrix ranks and shared weight layouts Signed-off-by: mgohil --- .../dla_matmul_to_transpose_conv_transpose.py | 409 ++++++++--------- ..._dla_matmul_to_transpose_conv_transpose.py | 427 ++++++++++++++++-- 2 files changed, 607 insertions(+), 229 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index 830ac2e34aa..91d80384f39 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -15,11 +15,10 @@ """Replace MatMul/Gemm nodes with an Unsqueeze-Transpose-Conv-Transpose-Squeeze subgraph. -For nodes that pass eligibility (see below), the full Conv conversion is applied. -For **ineligible** nodes, a lightweight 4-D promotion fallback is used instead: -dynamic inputs are Unsqueezed to 4-D, static initializers are padded with leading -1-dims in-place, and a Squeeze is inserted after the node to restore the original -output rank. This ensures all MatMul/Gemm tensors are 4-D after this transform. +Eligible nodes use Conv with rank adapters at the current stage boundary. +Ineligible Gemm nodes remain unchanged. Ineligible MatMul matrix operands use +four-dimensional promotion when requested; vector promotion is unsupported and +can only be deferred to a later stage. Rank adapters require ONNX opset 13+. Eligibility (checked via :func:`check_to_apply_transpose`): - The **second input** (weight) must be a constant initializer or a Const→DequantizeLinear chain. @@ -42,7 +41,7 @@ """ import numpy as np -from onnx import helper, numpy_helper +from onnx import TensorProto, helper, numpy_helper from ...logging_config import logger from ..utils.graph_utils import ( @@ -50,15 +49,46 @@ add_unique_initializers, batch_replace_nodes, check_to_apply_transpose, + get_node_attr_i, is_const_dq_input, pad4d, + set_node_attr_i, ) +from .dla_5d_reshape_to_4d import _known_non_scalar_control +from .dla_constants_to_initializers import _add_legacy_initializer_inputs +from .dla_graph_cleanup import _subgraph_captures def _apply_matmul_to_transpose_conv_transpose(model, *, defer_rank_conversion=False): - """Convert MatMul/Gemm; optionally let the pipeline promote tensor ranks later.""" + """Convert eligible matrices and optionally defer fallback rank promotion.""" graph = model.graph cache = GraphCache(graph) + captured = _subgraph_captures(graph) + local_functions = { + (function.domain, function.name, function.overload) for function in model.functions + } + opsets = {entry.domain: entry.version for entry in model.opset_import} + + def tensor_rank(name): + """Read rank metadata without requiring concrete dimension values.""" + info = cache.vi_map.get(name) + if ( + info is not None + and info.type.HasField("tensor_type") + and info.type.tensor_type.HasField("shape") + ): + return len(info.type.tensor_type.shape.dim) + shape = cache.get_shape(name) + return len(shape) if shape is not None else None + + candidates = [ + node + for node in graph.node + if node.op_type in {"MatMul", "Gemm"} + and node.domain in {"", "ai.onnx"} + and not node.overload + and (node.domain, node.op_type, node.overload) not in local_functions + ] unique_initializer: set = set() # shape map snapshot required by check_to_apply_transpose @@ -71,9 +101,13 @@ def _apply_matmul_to_transpose_conv_transpose(model, *, defer_rank_conversion=Fa node_idx = {id(n): i for i, n in enumerate(graph.node)} replacements_batch: dict[int, tuple[list, list]] = {} fallback_batch: dict[int, tuple[list, list]] = {} - for candidate in graph.node: - if candidate.op_type not in {"MatMul", "Gemm"}: - continue + for candidate in candidates: + if ( + candidate.op_type == "MatMul" + and not defer_rank_conversion + and any(tensor_rank(name) == 1 for name in candidate.input[:2]) + ): + raise ValueError(f"{candidate.name!r}: vector MatMul rank promotion is unsupported") if candidate.op_type == "Gemm": if any( attribute.name in {"alpha", "beta"} and attribute.f != 1.0 @@ -87,141 +121,143 @@ def _apply_matmul_to_transpose_conv_transpose(model, *, defer_rank_conversion=Fa if bias_shape is None or len(bias_shape) != 1: raise ValueError(f"Bias for {candidate.name!r} must have a known 1-D shape") - def update_initializers( - graph, name, current_node, nodes_to_add, trans_b=False, is_const_dq_init=False - ): - """Reshape weight initializer to a 4-D Conv kernel. - - If the initializer is shared by other consumers besides the current - rewrite, a clone (new initializer + cloned DQ if applicable) is - created so other consumers keep the original layout; this MatMul/Gemm - gets rewired to the clone. When the initializer has a single consumer - it's rewritten in place as before. - - Maps: - 2D [C,K] → [K,C,1,1] (or [C,K,1,1] when transB) - 3D [C,K,M] → [M,K,C,1] (or [C,K,M,1] when transB) - 4D [N,C,K,M] → [M,K,C,N] perm (or [C,K,M,N] when transB) - """ - for initializer in graph.initializer: - if initializer.name != name: - continue - axis = None - dq_node = None - shape = None - if is_const_dq_init: - dq_node = next( - ( - n - for n in graph.node - if n.op_type == "DequantizeLinear" and n.output[0] == current_node.input[1] - ), - None, - ) - if dq_node is not None: - axis = next((attr.i for attr in dq_node.attribute if attr.name == "axis"), None) - - # 2D [C,K] - if len(initializer.dims) == 2: - c, k = initializer.dims[0], initializer.dims[1] - init_arr = numpy_helper.to_array(initializer) - if not trans_b: - init_arr = init_arr.T - shape = (k, c, 1, 1) - if axis is not None: - axis = 1 - axis - else: - shape = (c, k, 1, 1) - - # 3D [C,K,M] - elif len(initializer.dims) == 3: - c, k, m = initializer.dims[0], initializer.dims[1], initializer.dims[2] - init_arr = numpy_helper.to_array(initializer) - if not trans_b: - init_arr = init_arr.T - shape = (m, k, c, 1) - if axis is not None: - axis = 2 - axis - else: - shape = (c, k, m, 1) - - # 4D [N,C,K,M] - elif len(initializer.dims) == 4: - n, c, k, m = ( - initializer.dims[0], - initializer.dims[1], - initializer.dims[2], - initializer.dims[3], - ) - init_arr = numpy_helper.to_array(initializer) - if not trans_b: - init_arr = np.transpose(init_arr, (3, 2, 1, 0)) - shape = (m, k, c, n) - if axis is not None: - axis = 3 - axis - else: - init_arr = np.transpose(init_arr, (1, 2, 3, 0)) - shape = (c, k, m, n) - - reshaped_arr = np.reshape(init_arr, shape) - if len(cache.get_consumers(name)) > 1 or ( + def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node=None): + """Allocate the Conv layout while retaining every shared, public or captured weight view.""" + initializer = cache.get_init(name) + if initializer is None: + raise ValueError(f"Missing weight initializer {name!r}") + dims = list(initializer.dims) + rank = len(dims) + perm = None if trans_b else list(reversed(range(rank))) + if trans_b and rank == 4: + perm = [1, 2, 3, 0] + shape = [dims[index] for index in perm] if perm is not None else dims + shape = [*shape, *([1] * (4 - rank))] + axis = None + if dq_node is not None: + axis = next( + (attribute.i for attribute in dq_node.attribute if attribute.name == "axis"), None + ) + if axis is None and any( + _known_non_scalar_control(cache, control) + or (cache.get_init(control) is None and cache.get_shape(control) is None) + for control in dq_node.input[1:] + if control + ): + axis = 1 + if axis is not None: + axis %= rank + if perm is not None: + axis = perm.index(axis) + shared = ( + len(cache.get_consumers(name)) > 1 + or name in cache.graph_output_names | captured + or ( dq_node is not None and ( len(cache.get_consumers(dq_node.output[0])) > 1 - or dq_node.output[0] in cache.graph_output_names + or dq_node.output[0] in cache.graph_output_names | captured ) - ): - clone_idx = len(replacements_batch) - new_init_name = cache.unique_name(f"{name}_conv4d_{clone_idx}") - new_init_proto = numpy_helper.from_array(reshaped_arr, new_init_name) - add_unique_initializers(graph, unique_initializer, [new_init_proto]) - if dq_node is not None: - new_dq_out = f"{dq_node.output[0]}_conv4d_{clone_idx}" - dq_clone = helper.make_node( - "DequantizeLinear", - inputs=[new_init_name, *list(dq_node.input[1:])], - outputs=[new_dq_out], - name=f"{dq_node.name or name}_conv4d_{clone_idx}", - ) - dq_clone.attribute.extend(dq_node.attribute) - if axis is not None: - for attr in dq_clone.attribute: - if attr.name == "axis": - attr.i = axis - nodes_to_add.append(dq_clone) - current_node.input[1] = new_dq_out - else: - current_node.input[1] = new_init_name - return - + ) + ) + clone_idx = len(replacements_batch) + target_name = cache.unique_name(f"{name}_conv4d_{clone_idx}") if shared else name + if perm is None: + target = graph.initializer.add() if shared else initializer + if shared: + target.CopyFrom(initializer) + target.name = target_name + del target.dims[:] + target.dims.extend(shape) + else: + array = np.transpose(numpy_helper.to_array(initializer), perm).reshape(shape) + converted = numpy_helper.from_array(array, target_name) + if shared: + add_unique_initializers(graph, unique_initializer, [converted]) + else: + initializer.CopyFrom(converted) + if shared: if dq_node is not None: - for attr in dq_node.attribute: - if attr.name == "axis": - attr.i = axis - dq_node.name = f"{dq_node.name}_axis_updated" - - reshaped_arr = np.reshape(init_arr, shape) - initializer.CopyFrom(numpy_helper.from_array(reshaped_arr, initializer.name)) + dq_clone = type(dq_node)() + dq_clone.CopyFrom(dq_node) + dq_clone.name = cache.unique_name(f"{dq_node.name or name}_conv4d_{clone_idx}") + dq_clone.input[0] = target_name + dq_clone.output[0] = cache.unique_name(f"{dq_node.output[0]}_conv4d_{clone_idx}") + if axis is not None: + set_node_attr_i(dq_clone, "axis", axis) + nodes_to_add.append(dq_clone) + current_node.input[1] = dq_clone.output[0] + else: + current_node.input[1] = target_name + elif dq_node is not None and axis is not None: + set_node_attr_i(dq_node, "axis", axis) + dq_node.name = cache.unique_name(f"{dq_node.name}_axis_updated") - # ── Separate eligible and ineligible nodes in one pass ─────────────────── eligible_nodes = [] ineligible_nodes = [] - for node in graph.node: - if node.op_type not in ("MatMul", "Gemm"): - continue - if check_to_apply_transpose(node, model, tensor_name_dim_map): - eligible_nodes.append(node) + for node in candidates: + producer = cache.get_producer(node.input[1]) + dq_node = ( + producer if producer is not None and producer.op_type == "DequantizeLinear" else None + ) + known_dq = dq_node is None or ( + not dq_node.overload + and (dq_node.domain, dq_node.op_type, dq_node.overload) not in local_functions + and ( + dq_node.domain in {"", "ai.onnx"} + or (dq_node.domain == "com.microsoft" and opsets.get(dq_node.domain) == 1) + ) + and not get_node_attr_i(dq_node, "block_size", 0) + ) + weight_name = dq_node.input[0] if dq_node is not None else node.input[1] + weight = cache.get_init(weight_name) + dtype = cache.get_dtype(node.input[0]) + activation_shape = cache.get_shape(node.input[0]) + if ( + not known_dq + or weight is None + or weight_name in cache.graph_input_names + or any(tensor_rank(name) == 1 for name in node.input[:2]) + or ( + dtype is not None + and dtype not in {TensorProto.FLOAT, TensorProto.FLOAT16, TensorProto.DOUBLE} + ) + or any(dim == 0 for dim in weight.dims) + or (activation_shape is not None and 0 in activation_shape) + ): + ineligible_nodes.append(node) + elif check_to_apply_transpose(node, model, tensor_name_dim_map): + output_rank = tensor_rank(node.output[0]) + if output_rank is None: + output_rank = ( + max(len(tensor_name_dim_map[node.input[0]]), len(weight.dims)) + if node.op_type == "MatMul" + else 2 + ) + eligible_nodes.append((node, weight_name, dq_node, output_rank)) else: ineligible_nodes.append(node) + active_fallback = [ + node + for node in ineligible_nodes + if node.op_type == "MatMul" + and not any(tensor_rank(name) == 1 for name in node.input[:2]) + and (not defer_rank_conversion or node.input[1] in cache.graph_input_names) + ] + needs_adapters = any( + len(tensor_name_dim_map[node.input[0]]) < 4 or rank < 4 + for node, _, _, rank in eligible_nodes + ) or any( + any( + tensor_rank(name) is not None and tensor_rank(name) < 4 + for name in [*node.input[:2], node.output[0]] + ) + for node in active_fallback + ) + if needs_adapters and opsets.get("", opsets.get("ai.onnx", 1)) < 13: + raise ValueError("DLA matrix rank adapters require ONNX opset 13 or newer") - # ── Fallback: promote inputs/output of ineligible MatMul/Gemm to 4D ───── - # For each ineligible node: - # * Dynamic inputs → Unsqueeze to 4D. - # * Static initializer inputs (including DQ-wrapped) → expand dims in-place. - # * Output → Squeeze back to original rank when the original was not 4D. - for node in ineligible_nodes: - if defer_rank_conversion and node.input[1] not in cache.graph_input_names: - continue + for node in active_fallback: cnt_fb = len(fallback_batch) pfx = node.name or f"mm_fallback_{cnt_fb}" initializer_names = cache.init_names @@ -234,13 +270,14 @@ def update_initializers( continue if not inp: continue - in_shape = cache.get_shape(inp) - if in_shape is None or len(in_shape) >= 4: + rank = tensor_rank(inp) + if rank is None or rank >= 4: continue - rank = len(in_shape) # Check whether this is a static initializer (possibly DQ-wrapped) dq_src = is_const_dq_input(initializer_names, inp, graph) - init_name = dq_src or (inp if inp in initializer_names else None) + init_name = dq_src or ( + inp if inp in initializer_names and inp not in cache.graph_input_names else None + ) if init_name is not None: if dq_src is not None: # Preserve the quantization layout when a static DQ weight needs promotion. @@ -261,19 +298,17 @@ def update_initializers( # Static weight: expand initializer dims in-place with 1-padding for init in graph.initializer: if init.name == init_name: - arr = numpy_helper.to_array(init) - if arr.ndim < 4: - new_shape = pad4d(list(arr.shape)) + if len(init.dims) < 4: + new_shape = pad4d(list(init.dims)) if ( len(cache.get_consumers(init.name)) > 1 - or init.name in cache.graph_output_names + or init.name in cache.graph_output_names | captured ): - clone = type(init)() + clone = graph.initializer.add() clone.CopyFrom(init) clone.name = cache.unique_name(init.name + "_4d") del clone.dims[:] clone.dims.extend(new_shape) - graph.initializer.append(clone) node.input[inp_idx] = clone.name else: del init.dims[:] @@ -284,30 +319,30 @@ def update_initializers( axes = list(range(4 - rank)) axes_init = numpy_helper.from_array( np.array(axes, dtype=np.int64), - name=f"{pfx}_in{inp_idx}_unsq_axes_{cnt_fb}", + name=cache.unique_name(f"{pfx}_in{inp_idx}_unsq_axes_{cnt_fb}"), ) add_unique_initializers(graph, unique_initializer, [axes_init]) - inp_4d = f"{pfx}_in{inp_idx}_4d_{cnt_fb}" + inp_4d = cache.unique_name(f"{pfx}_in{inp_idx}_4d_{cnt_fb}") pre.append( helper.make_node( "Unsqueeze", inputs=[inp, axes_init.name], outputs=[inp_4d], - name=f"{pfx}_in{inp_idx}_unsq_{cnt_fb}", + name=cache.unique_name(f"{pfx}_in{inp_idx}_unsq_{cnt_fb}"), ) ) node.input[inp_idx] = inp_4d # Output: Squeeze back to original rank out_name = node.output[0] if node.output else None - if out_name and not defer_rank_conversion: - out_shape = cache.get_shape(out_name) - if out_shape is not None and len(out_shape) < 4: - temp_out = f"{pfx}_out_4d_{cnt_fb}" - sq_axes = list(range(4 - len(out_shape))) + if out_name: + out_rank = tensor_rank(out_name) + if out_rank is not None and out_rank < 4: + temp_out = cache.unique_name(f"{pfx}_out_4d_{cnt_fb}") + sq_axes = list(range(4 - out_rank)) sq_axes_init = numpy_helper.from_array( np.array(sq_axes, dtype=np.int64), - name=f"{pfx}_out_sq_axes_{cnt_fb}", + name=cache.unique_name(f"{pfx}_out_sq_axes_{cnt_fb}"), ) add_unique_initializers(graph, unique_initializer, [sq_axes_init]) post.append( @@ -315,7 +350,7 @@ def update_initializers( "Squeeze", inputs=[temp_out, sq_axes_init.name], outputs=[out_name], - name=f"{pfx}_out_sq_{cnt_fb}", + name=cache.unique_name(f"{pfx}_out_sq_{cnt_fb}"), ) ) node.output[0] = temp_out @@ -323,11 +358,10 @@ def update_initializers( if pre or post: fallback_batch[node_idx[id(node)]] = ([node], [*pre, node, *post]) - for node in eligible_nodes: + for node, weight_name, weight_dq, output_rank in eligible_nodes: cnt = len(replacements_batch) pfx = node.name or f"matmul_{cnt}" nodes_to_add = [] - initializer_names = cache.init_names # ── Gemm transA / transB attributes ────────────────────────────────── trans_a = False @@ -344,26 +378,24 @@ def update_initializers( if input_shape is None: raise ValueError(f"Shape unknown for activation tensor {node.input[0]!r}") orig_rank = len(input_shape) - output_shape = cache.get_shape(node.output[0]) - output_rank = len(output_shape) if output_shape is not None else orig_rank # ── Step 1: Unsqueeze activation to 4D ─────────────────────────────── # 2D [M,N] → axes [0,1] → [1,1,M,N] # 3D [B,M,N] → axes [0] → [1,B,M,N] activation_4d = node.input[0] - if orig_rank < 4 and not defer_rank_conversion: + if orig_rank < 4: axes = [0, 1] if orig_rank == 2 else [0] axes_init = numpy_helper.from_array( - np.array(axes, dtype=np.int64), name=f"{pfx}_unsq_axes_{cnt}" + np.array(axes, dtype=np.int64), name=cache.unique_name(f"{pfx}_unsq_axes_{cnt}") ) add_unique_initializers(graph, unique_initializer, [axes_init]) - activation_4d = f"{pfx}_unsqueeze4d_{cnt}" + activation_4d = cache.unique_name(f"{pfx}_unsqueeze4d_{cnt}") nodes_to_add.append( helper.make_node( "Unsqueeze", inputs=[node.input[0], axes_init.name], outputs=[activation_4d], - name=f"{pfx}_unsqueeze4d_node_{cnt}", + name=cache.unique_name(f"{pfx}_unsqueeze4d_node_{cnt}"), ) ) @@ -371,47 +403,21 @@ def update_initializers( # transA=False: [1,B,M,N] → perm[0,3,2,1] → [1,N,M,B] (C_in=N, H=M, W=B) # transA=True : [1,B,N,M] → perm[0,2,3,1] → [1,N,M,B] (same result) perm_a = [0, 2, 3, 1] if trans_a else [0, 3, 2, 1] - transposed_a = f"{pfx}_tr_before0_{cnt}" + transposed_a = cache.unique_name(f"{pfx}_tr_before0_{cnt}") nodes_to_add.append( helper.make_node( "Transpose", inputs=[activation_4d], outputs=[transposed_a], - name=f"{pfx}_transpose_before0_{cnt}", + name=cache.unique_name(f"{pfx}_transpose_before0_{cnt}"), perm=perm_a, ) ) # ── Step 3: Prepare weight ──────────────────────────────────────────── - dq_init_name = is_const_dq_input(initializer_names, node.input[1], graph) - is_const_dq = dq_init_name is not None - init_name = ( - dq_init_name - if is_const_dq - else (node.input[1] if node.input[1] in initializer_names else None) - ) - - if init_name is not None: - # Static weight: reshape initializer to 4-D conv kernel format - update_initializers(graph, init_name, node, nodes_to_add, trans_b, is_const_dq) - conv_input1 = node.input[1] # tensor name unchanged; dims updated - else: - # Dynamic weight: add a Transpose node before Conv - # (check_to_apply_transpose normally prevents reaching here) - perm_b = [2, 3, 0, 1] if trans_b else [3, 2, 0, 1] - transposed_b = f"{pfx}_tr_before1_{cnt}" - nodes_to_add.append( - helper.make_node( - "Transpose", - inputs=[node.input[1]], - outputs=[transposed_b], - name=f"{pfx}_transpose_before1_{cnt}", - perm=perm_b, - ) - ) - conv_input1 = transposed_b + update_initializers(weight_name, node, nodes_to_add, trans_b, weight_dq) + conv_input1 = node.input[1] - # ── Step 4: Build Conv inputs (include optional Gemm bias) ─────────── conv_inputs = [transposed_a, conv_input1] if len(node.input) >= 3 and node.input[2]: bias_shape = cache.get_shape(node.input[2]) @@ -422,13 +428,13 @@ def update_initializers( conv_inputs.append(node.input[2]) # ── Step 5: Conv ────────────────────────────────────────────────────── - conv_out = f"{pfx}_conv_out_{cnt}" + conv_out = cache.unique_name(f"{pfx}_conv_out_{cnt}") nodes_to_add.append( helper.make_node( "Conv", inputs=conv_inputs, outputs=[conv_out], - name=f"{pfx}_conv_{cnt}", + name=cache.unique_name(f"{pfx}_conv_{cnt}"), ) ) @@ -436,16 +442,14 @@ def update_initializers( # [1,K,M,B] → perm[0,3,2,1] → [1,B,M,K] # Write to a temp name when Squeeze follows; otherwise directly to node.output[0]. tr_after_out = ( - f"{pfx}_tr_after_out_{cnt}" - if output_rank < 4 and not defer_rank_conversion - else node.output[0] + cache.unique_name(f"{pfx}_tr_after_out_{cnt}") if output_rank < 4 else node.output[0] ) nodes_to_add.append( helper.make_node( "Transpose", inputs=[conv_out], outputs=[tr_after_out], - name=f"{pfx}_transpose_after_{cnt}", + name=cache.unique_name(f"{pfx}_transpose_after_{cnt}"), perm=[0, 3, 2, 1], ) ) @@ -453,10 +457,10 @@ def update_initializers( # ── Step 7: Squeeze back to original rank ───────────────────────────── # 2D case: [1,1,M,K] → squeeze [0,1] → [M,K] # 3D case: [1,B,M,K] → squeeze [0] → [B,M,K] - if output_rank < 4 and not defer_rank_conversion: + if output_rank < 4: sq_axes = list(range(4 - output_rank)) sq_axes_init = numpy_helper.from_array( - np.array(sq_axes, dtype=np.int64), name=f"{pfx}_sq_axes_{cnt}" + np.array(sq_axes, dtype=np.int64), name=cache.unique_name(f"{pfx}_sq_axes_{cnt}") ) add_unique_initializers(graph, unique_initializer, [sq_axes_init]) nodes_to_add.append( @@ -464,7 +468,7 @@ def update_initializers( "Squeeze", inputs=[tr_after_out, sq_axes_init.name], outputs=list(node.output), - name=f"{pfx}_squeeze_after_{cnt}", + name=cache.unique_name(f"{pfx}_squeeze_after_{cnt}"), ) ) @@ -473,10 +477,11 @@ def update_initializers( # ── Apply all replacements ──────────────────────────────────────────────── all_batch = {**replacements_batch, **fallback_batch} batch_replace_nodes(graph, all_batch) + _add_legacy_initializer_inputs(model) logger.debug( "Replaced %d MatMul/Gemm nodes with Transpose-Conv-Transpose nodes; " - "promoted %d ineligible MatMul/Gemm nodes to 4D.", + "promoted %d ineligible MatMul nodes to 4D.", len(replacements_batch), len(fallback_batch), ) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index 0a65faee3fd..fd4b0e6fef6 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -35,28 +35,25 @@ Dynamic weight A[M,N] x W_graph_input[N,K] → MatMul left unchanged NCHW non-unary A[2,3,M,N] x W_init[2,3,N,K] → MatMul left unchanged (4D weight, N/C > 1) -Each test saves a base + transformed ONNX pair under -``scratch_space/matmul_conv_test_models/`` for offline inspection. """ from __future__ import annotations import copy -from pathlib import Path import numpy as np import onnx import pytest from onnx import TensorProto, helper, numpy_helper +from onnx.reference import ReferenceEvaluator from modelopt.onnx.graph_surgery.dla_transforms.dla_matmul_to_transpose_conv_transpose import ( _apply_matmul_to_transpose_conv_transpose as _apply_matmul_to_conv, ) +from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible -pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") +ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") -_REPO_ROOT = Path(__file__).resolve().parents[5] -_SAVED_MODELS_DIR = _REPO_ROOT / "scratch_space" / "matmul_conv_test_models" _OPSET_VERSION = 17 _MAX_IR_VERSION_FOR_ORT = 9 @@ -68,26 +65,29 @@ def _init(name: str, arr: np.ndarray) -> onnx.TensorProto: + """Create a named tensor initializer.""" return numpy_helper.from_array(arr, name=name) def _vi(name: str, elem_type: int, shape: list) -> onnx.ValueInfoProto: + """Create typed shape metadata.""" return helper.make_tensor_value_info(name, elem_type, shape) def _clamp_ir_version(model: onnx.ModelProto) -> None: + """Use an IR version supported by the test runtime.""" model.ir_version = min(model.ir_version, _MAX_IR_VERSION_FOR_ORT) def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: - import onnxruntime as ort # type: ignore[import-not-found] - + """Execute all graph outputs with the optional CPU runtime.""" _clamp_ir_version(model) sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: + """Measure the directional agreement of numerical outputs.""" a64 = np.asarray(a, dtype=np.float64).ravel() b64 = np.asarray(b, dtype=np.float64).ravel() na, nb = float(np.linalg.norm(a64)), float(np.linalg.norm(b64)) @@ -99,6 +99,7 @@ def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float: def _assert_outputs_match(ref: dict, got: dict) -> None: + """Check output names, shapes and numerical agreement.""" assert ref.keys() == got.keys() for name in ref: a, b = ref[name], got[name] @@ -108,22 +109,13 @@ def _assert_outputs_match(ref: dict, got: dict) -> None: np.testing.assert_allclose(a, b, rtol=1e-4, atol=1e-4, err_msg=f"{name}: value mismatch") -def _save_test_pair(base: onnx.ModelProto, transformed: onnx.ModelProto, tag: str) -> None: - _SAVED_MODELS_DIR.mkdir(parents=True, exist_ok=True) - onnx.save(base, str(_SAVED_MODELS_DIR / f"matmul_base_{tag}.onnx")) - import contextlib - - with contextlib.suppress(Exception): - transformed = onnx.shape_inference.infer_shapes(copy.deepcopy(transformed)) - onnx.save(transformed, str(_SAVED_MODELS_DIR / f"matmul_transformed_{tag}.onnx")) - - def _make_model(nodes, graph_name, inputs, outputs, initializers=None): """Build, shape-infer, and checker-validate a simple ONNX model.""" graph = helper.make_graph(nodes, graph_name, inputs, outputs, initializer=initializers or []) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", _OPSET_VERSION)]) _clamp_ir_version(model) try: + # Optional runtime tools may be absent without preventing test collection. from onnxruntime.tools.symbolic_shape_infer import ( SymbolicShapeInference, # type: ignore[import-untyped] ) @@ -284,6 +276,7 @@ def _make_matmul_nchw_non_unary(): def _apply_and_run(model, feeds): + """Execute the original and converted graph for comparison.""" ref = _run_ort(model, feeds) transformed = copy.deepcopy(model) transformed = _apply_matmul_to_conv(transformed) @@ -307,7 +300,6 @@ def test_matmul_2d_activation_converts(): assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 2D input" assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 2D input" - _save_test_pair(model, tr, "2d") _assert_outputs_match(ref, got) @@ -321,7 +313,6 @@ def test_matmul_3d_activation_converts(): assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 3D input" assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 3D input" - _save_test_pair(model, tr, "3d") _assert_outputs_match(ref, got) @@ -339,7 +330,6 @@ def test_matmul_4d_activation_converts(): "No Squeeze expected for 4D input" ) - _save_test_pair(model, tr, "4d") _assert_outputs_match(ref, got) @@ -350,7 +340,6 @@ def test_gemm_transb_converts(): assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - _save_test_pair(model, tr, "gemm_transb") _assert_outputs_match(ref, got) @@ -361,7 +350,6 @@ def test_gemm_transa_converts(): assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - _save_test_pair(model, tr, "gemm_transa") _assert_outputs_match(ref, got) @@ -375,7 +363,6 @@ def test_gemm_with_bias_converts(): assert len(conv_nodes) == 1 assert len(conv_nodes[0].input) == 3, "Conv should carry bias as third input" - _save_test_pair(model, tr, "gemm_bias") _assert_outputs_match(ref, got) @@ -395,7 +382,6 @@ def test_matmul_dq_weight_converts(): f"W_q dims: {list(w_q_init.dims)}, expected [{K},{N},1,1]" ) - _save_test_pair(model, tr, "dq_weight") _assert_outputs_match(ref, got) @@ -419,7 +405,6 @@ def test_matmul_dynamic_weight_not_converted(): "Squeeze expected to restore original output rank" ) - _save_test_pair(model, tr, "dynamic_weight") _assert_outputs_match(ref, got) @@ -438,5 +423,393 @@ def test_matmul_nchw_non_unary_not_converted(): "No Unsqueeze/Squeeze expected when all tensors are already 4D" ) - _save_test_pair(model, tr, "nchw_non_unary") _assert_outputs_match(ref, got) + + +@pytest.mark.parametrize("trans_a", [False, True]) +@pytest.mark.parametrize("trans_b", [False, True]) +@pytest.mark.parametrize("defer", [False, True]) +def test_dynamic_gemm_preserves_two_dimensional_contract(trans_a, trans_b, defer): + """Leave dynamic Gemm operands and transpose semantics valid in both modes.""" + a_shape = [N, M] if trans_a else [M, N] + b_shape = [K, N] if trans_b else [N, K] + model = _make_model( + [ + helper.make_node( + "Gemm", ["A", "W"], ["Y"], name="gemm", transA=int(trans_a), transB=int(trans_b) + ) + ], + "dynamic_gemm", + [_vi("A", TensorProto.FLOAT, a_shape), _vi("W", TensorProto.FLOAT, b_shape)], + [_vi("Y", TensorProto.FLOAT, [M, K])], + ) + feeds = { + "A": _RNG.standard_normal(a_shape).astype(np.float32), + "W": _RNG.standard_normal(b_shape).astype(np.float32), + } + expected = _run_ort(model, feeds) + before = model.SerializeToString() + _apply_matmul_to_conv(model, defer_rank_conversion=defer) + onnx.checker.check_model(model, full_check=True) + assert model.SerializeToString() == before + _assert_outputs_match(expected, _run_ort(model, feeds)) + + +@pytest.mark.parametrize("rank", [2, 3]) +def test_eligible_deferred_conversion_keeps_rank_adapters(rank): + """Keep eligible Conv rewrites valid before any later rank-promotion stage.""" + model, feeds = (_make_matmul_2d if rank == 2 else _make_matmul_3d)() + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=True) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_internal_pipeline_executes_eligible_two_dimensional_matmul(): + """Run the registered deferred conversion through the actual internal sequence.""" + model, feeds = _make_matmul_2d() + expected = _run_ort(model, feeds) + transformed = _transform_make_dla_compatible(copy.deepcopy(model), verbose=False) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("axis", [None, -1, 1]) +@pytest.mark.parametrize("shared", [False, True]) +def test_per_axis_dq_weight_layout_is_preserved(axis, shared): + """Remap implicit and negative quantization axes while preserving shared DQ outputs.""" + model, feeds = _make_matmul_dq_weight() + for tensor in model.graph.initializer: + if tensor.name == "W_scale": + tensor.CopyFrom(_init("W_scale", np.linspace(0.01, 0.02, K, dtype=np.float32))) + elif tensor.name == "W_zp": + tensor.CopyFrom(_init("W_zp", np.zeros(K, np.int8))) + if axis is not None: + model.graph.node[0].attribute.append(helper.make_attribute("axis", axis)) + if shared: + model.graph.output.append(_vi("W_dq", TensorProto.FLOAT, [N, K])) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_public_weight_initializer_keeps_original_layout(quantized): + """Retain public weight tensors while allocating a Conv kernel for this consumer.""" + model, feeds = (_make_matmul_dq_weight if quantized else _make_matmul_2d)() + name = "W_q" if quantized else "W" + model.graph.output.append( + _vi(name, TensorProto.INT8 if quantized else TensorProto.FLOAT, [N, K]) + ) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("defer", [False, True]) +def test_overridable_weight_initializer_uses_runtime_feed(quantized, defer): + """Preserve runtime weight overrides instead of converting their initializer defaults.""" + model, feeds = (_make_matmul_dq_weight if quantized else _make_matmul_2d)() + name = "W_q" if quantized else "W" + tensor = next(tensor for tensor in model.graph.initializer if tensor.name == name) + model.graph.input.append(_vi(name, tensor.data_type, list(tensor.dims))) + feeds[name] = -numpy_helper.to_array(tensor) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=defer) + onnx.checker.check_model(transformed, full_check=True) + assert not any(node.op_type == "Conv" for node in transformed.graph.node) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("op_type", ["MatMul", "Gemm"]) +def test_custom_matrix_function_is_not_lowered_as_builtin(op_type): + """Preserve custom matrix functions whose result differs from the standard operation.""" + function = helper.make_function( + "custom", + op_type, + ["A", "W"], + ["Y"], + [ + helper.make_node("MatMul", ["A", "W"], ["product"]), + helper.make_node("Neg", ["product"], ["Y"]), + ], + opset_imports=[helper.make_opsetid("", 17)], + ) + model = helper.make_model( + helper.make_graph( + [helper.make_node(op_type, ["A", "W"], ["Y"], name="custom_matrix", domain="custom")], + "custom_matrix", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [M, K])], + initializer=[_init("W", _RNG.standard_normal([N, K]).astype(np.float32))], + ), + opset_imports=[helper.make_opsetid("", 17), helper.make_opsetid("custom", 1)], + functions=[function], + ) + _clamp_ir_version(model) + onnx.checker.check_model(model, full_check=True) + feeds = {"A": _RNG.standard_normal([M, N]).astype(np.float32)} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + _apply_matmul_to_conv(model) + assert model.SerializeToString() == before + _assert_outputs_match(expected, _run_ort(model, feeds)) + + +@pytest.mark.parametrize("collision", ["activation", "cloned-dq"]) +def test_generated_matrix_names_preserve_unrelated_outputs(collision): + """Allocate fresh generated names across activation and cloned DQ branches.""" + model, feeds = (_make_matmul_dq_weight if collision == "cloned-dq" else _make_matmul_2d)() + if collision == "cloned-dq": + model.graph.output.append(_vi("W_dq", TensorProto.FLOAT, [N, K])) + name = "W_dq_conv4d_0" if collision == "cloned-dq" else "mm_tr_before0_0" + model.graph.node.insert(0, helper.make_node("Identity", ["A"], [name], name="existing_alias")) + model.graph.output.append(_vi(name, TensorProto.FLOAT, [M, N])) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_integer_matmul_remains_a_valid_matmul(): + """Keep integer matrix multiplication on its supported operator and preserve exact values.""" + data = np.arange(6, dtype=np.int64).reshape(2, 3) + weight = np.arange(12, dtype=np.int64).reshape(3, 4) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="integer_mm")], + "integer_mm", + [_vi("A", TensorProto.INT64, [2, 3])], + [_vi("Y", TensorProto.INT64, [2, 4])], + [_init("W", weight)], + ) + expected = ReferenceEvaluator(model).run(None, {"A": data})[0] + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + np.testing.assert_array_equal( + ReferenceEvaluator(transformed).run(None, {"A": data})[0], expected + ) + + +def test_deferred_dynamic_matmul_restores_output_rank(): + """Keep deferred graph-input weight promotion valid at the current stage boundary.""" + model, feeds = _make_matmul_dynamic_weight() + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=True) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("symbolic", [False, True]) +def test_batched_weight_fallback_preserves_values_without_decoding(symbolic): + """Promote batched matrix layouts using ranks even when dimension values are symbolic.""" + weight = _RNG.standard_normal([2, N, K]).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="batched")], + "batched_fallback", + [_vi("A", TensorProto.FLOAT, [2, M, N])], + [_vi("Y", TensorProto.FLOAT, [2, M, K])], + [_init("W", weight)], + ) + if symbolic: + model.graph.input[0].CopyFrom(_vi("A", TensorProto.FLOAT, [2, "rows", N])) + model.graph.output[0].CopyFrom(_vi("Y", TensorProto.FLOAT, [2, "rows", K])) + onnx.checker.check_model(model, full_check=True) + feeds = {"A": _RNG.standard_normal([2, M, N]).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + assert ( + next(tensor for tensor in transformed.graph.initializer if tensor.name == "W").raw_data + == next(tensor for tensor in model.graph.initializer if tensor.name == "W").raw_data + ) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("weight_rank", [3, 4]) +def test_batched_weight_output_rank_survives_symbolic_metadata(weight_rank): + """Restore the declared batch rank independently of concrete output dimension values.""" + weight = _RNG.standard_normal([*([1] * (weight_rank - 2)), N, K]).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="batched")], + "batched_output", + [_vi("A", TensorProto.FLOAT, [M, N])], + [_vi("Y", TensorProto.FLOAT, [*([1] * (weight_rank - 2)), M, K])], + [_init("W", weight)], + ) + model.graph.output[0].CopyFrom( + _vi("Y", TensorProto.FLOAT, [f"dim_{index}" for index in range(weight_rank)]) + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"A": _RNG.standard_normal([M, N]).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_captured_weight_layout_is_preserved(quantized): + """Clone converted weights when nested control flow still reads their original layout.""" + model, feeds = (_make_matmul_dq_weight if quantized else _make_matmul_2d)() + name = "W_dq" if quantized else "W" + branch = helper.make_graph( + [helper.make_node("Identity", [name], ["branch_out"])], + "branch", + [], + [_vi("branch_out", TensorProto.FLOAT, [N, K])], + ) + model.graph.input.append(_vi("condition", TensorProto.BOOL, [])) + model.graph.node.append( + helper.make_node( + "If", + ["condition"], + ["captured_weight"], + name="capture", + then_branch=branch, + else_branch=copy.deepcopy(branch), + ) + ) + model.graph.output.append(_vi("captured_weight", TensorProto.FLOAT, [N, K])) + feeds["condition"] = np.asarray(True) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + actual = _run_ort(transformed, feeds) + np.testing.assert_array_equal(actual["captured_weight"], expected["captured_weight"]) + _assert_outputs_match(expected, actual) + + +def test_blocked_weight_dequantization_retains_its_layout(): + """Keep blocked scale/data coordinates intact while promoting only the DQ result.""" + model, feeds = _make_matmul_dq_weight() + model.opset_import[0].version = 21 + model.graph.node[0].attribute.extend( + [helper.make_attribute("axis", 1), helper.make_attribute("block_size", 4)] + ) + for tensor in model.graph.initializer: + if tensor.name == "W_scale": + tensor.CopyFrom(_init("W_scale", np.full([N, K // 4], 0.05, np.float32))) + elif tensor.name == "W_zp": + tensor.CopyFrom(_init("W_zp", np.zeros([N, K // 4], np.int8))) + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + for actual, reference in zip(ReferenceEvaluator(transformed).run(None, feeds), expected): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("domain", ["custom", "com.microsoft"]) +def test_opaque_weight_function_is_not_reinterpreted_as_dq(domain): + """Preserve model-local weight functions with nonpointwise semantics.""" + model, feeds = _make_matmul_dq_weight() + function = helper.make_function( + domain, + "DequantizeLinear", + ["data", "scale", "zero"], + ["output"], + [ + helper.make_node("Cast", ["data"], ["cast"], to=TensorProto.FLOAT), + helper.make_node("Softmax", ["cast"], ["output"], axis=1), + ], + opset_imports=[helper.make_opsetid("", 17)], + ) + model.functions.append(function) + model.opset_import.append(helper.make_opsetid(domain, 1)) + model.graph.node[0].domain = domain + onnx.checker.check_model(model, full_check=True) + expected = ReferenceEvaluator(model).run(None, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + for actual, reference in zip(ReferenceEvaluator(transformed).run(None, feeds), expected): + np.testing.assert_allclose(actual, reference, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize(("vector_a", "vector_b"), [(True, False), (False, True), (True, True)]) +@pytest.mark.parametrize("defer", [False, True]) +def test_vector_matmul_is_refused_or_left_for_later_conversion(vector_a, vector_b, defer): + """Avoid applying matrix-only rank promotion to vector multiplication semantics.""" + a_shape = [N] if vector_a else [M, N] + w_shape = [N] if vector_b else [N, K] + output_shape = [] if vector_a and vector_b else [K] if vector_a else [M] + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="vector")], + "vector", + [_vi("A", TensorProto.FLOAT, a_shape)], + [_vi("Y", TensorProto.FLOAT, output_shape)], + [_init("W", _RNG.standard_normal(w_shape).astype(np.float32))], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"A": _RNG.standard_normal(a_shape).astype(np.float32)} + expected = _run_ort(model, feeds) + before = model.SerializeToString() + if defer: + _apply_matmul_to_conv(model, defer_rank_conversion=True) + assert model.SerializeToString() == before + _assert_outputs_match(expected, _run_ort(model, feeds)) + else: + with pytest.raises(ValueError, match=r"[Uu]nsupported|[Vv]ector"): + _apply_matmul_to_conv(model) + assert model.SerializeToString() == before + + +@pytest.mark.parametrize("empty", ["rows", "outputs"]) +def test_empty_matrix_dimensions_remain_valid(empty): + """Keep empty matrix results on MatMul instead of creating an unsupported Conv kernel.""" + rows = 0 if empty == "rows" else M + outputs = 0 if empty == "outputs" else K + weight = _RNG.standard_normal([N, outputs]).astype(np.float32) + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="empty")], + "empty", + [_vi("A", TensorProto.FLOAT, [rows, N])], + [_vi("Y", TensorProto.FLOAT, [rows, outputs])], + [_init("W", weight)], + ) + onnx.checker.check_model(model, full_check=True) + feeds = {"A": _RNG.standard_normal([rows, N]).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +def test_legacy_rank_adapter_schema_refuses_before_mutation(): + """Refuse old Unsqueeze/Squeeze schemas before adding input-form rank adapters.""" + model, feeds = _make_matmul_2d() + model.opset_import[0].version = 11 + onnx.checker.check_model(model, full_check=True) + _run_ort(model, feeds) + before = model.SerializeToString() + with pytest.raises(ValueError, match="opset"): + _apply_matmul_to_conv(model) + assert model.SerializeToString() == before + + +def test_legacy_ir_registers_new_initializer_inputs(): + """Keep IR3 valid when rank adapters introduce new constant axes tensors.""" + model, feeds = _make_matmul_2d() + model.ir_version = 3 + for tensor in model.graph.initializer: + model.graph.input.append(_vi(tensor.name, tensor.data_type, list(tensor.dims))) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + assert {tensor.name for tensor in transformed.graph.initializer} <= { + value.name for value in transformed.graph.input + } + _assert_outputs_match(expected, _run_ort(transformed, feeds)) From a8333540ab86c7d4d345d9dbcb0d640c2e431c0c Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 21:05:13 +0000 Subject: [PATCH 3/7] Record DLA matrix adapter shapes for later rank passes Signed-off-by: mgohil --- .../dla_matmul_to_transpose_conv_transpose.py | 21 +++++++++++++++++++ ..._dla_matmul_to_transpose_conv_transpose.py | 10 +++++++++ 2 files changed, 31 insertions(+) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index 91d80384f39..9904f25a1c6 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -476,6 +476,27 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node # ── Apply all replacements ──────────────────────────────────────────────── all_batch = {**replacements_batch, **fallback_batch} + # Later rank passes consume the original shape snapshot without another inference step. + for _, replacements in all_batch.values(): + for adapter in replacements: + if adapter.op_type not in {"Unsqueeze", "Squeeze"}: + continue + original, padded = ( + (adapter.input[0], adapter.output[0]) + if adapter.op_type == "Unsqueeze" + else (adapter.output[0], adapter.input[0]) + ) + info = cache.vi_map.get(original) + if info is not None and info.type.tensor_type.HasField("shape"): + promoted = graph.value_info.add() + promoted.CopyFrom(info) + promoted.name = padded + dims = promoted.type.tensor_type.shape.dim + original_dims = list(dims) + del dims[:] + for _ in range(4 - len(original_dims)): + dims.add().dim_value = 1 + dims.extend(original_dims) batch_replace_nodes(graph, all_batch) _add_legacy_initializer_inputs(model) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index fd4b0e6fef6..d04784586c1 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -51,6 +51,7 @@ _apply_matmul_to_transpose_conv_transpose as _apply_matmul_to_conv, ) from modelopt.onnx.graph_surgery.make_dla_compatible import _transform_make_dla_compatible +from modelopt.onnx.graph_surgery.utils.graph_utils import GraphCache ort = pytest.importorskip("onnxruntime", reason="onnxruntime required for ORT parity checks") @@ -461,6 +462,15 @@ def test_eligible_deferred_conversion_keeps_rank_adapters(rank): model, feeds = (_make_matmul_2d if rank == 2 else _make_matmul_3d)() expected = _run_ort(model, feeds) transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=True) + cache = GraphCache(transformed.graph) + for adapter in transformed.graph.node: + if adapter.op_type == "Unsqueeze": + shape = cache.get_shape(adapter.output[0]) + elif adapter.op_type == "Squeeze": + shape = cache.get_shape(adapter.input[0]) + else: + continue + assert shape is not None and len(shape) == 4 onnx.checker.check_model(transformed, full_check=True) _assert_outputs_match(expected, _run_ort(transformed, feeds)) From 7a3a9c49d116031cd8b7c6a235b65353367b6abd Mon Sep 17 00:00:00 2001 From: mgohil Date: Mon, 26 Oct 2026 21:26:49 +0000 Subject: [PATCH 4/7] Preserve DLA fallback ranks and inferred weight metadata Signed-off-by: mgohil --- .../dla_matmul_to_transpose_conv_transpose.py | 26 ++++- ..._dla_matmul_to_transpose_conv_transpose.py | 95 ++++++++++++++++++- 2 files changed, 116 insertions(+), 5 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index 9904f25a1c6..02519dc6815 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -121,6 +121,15 @@ def tensor_rank(name): if bias_shape is None or len(bias_shape) != 1: raise ValueError(f"Bias for {candidate.name!r} must have a known 1-D shape") + def update_weight_metadata(name, shape): + """Keep private weight annotations consistent with an in-place layout change.""" + info = cache.vi_map.get(name) + if info is not None: + dims = info.type.tensor_type.shape.dim + del dims[:] + for size in shape: + dims.add().dim_value = size + def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node=None): """Allocate the Conv layout while retaining every shared, public or captured weight view.""" initializer = cache.get_init(name) @@ -176,6 +185,10 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node add_unique_initializers(graph, unique_initializer, [converted]) else: initializer.CopyFrom(converted) + if not shared: + update_weight_metadata(name, shape) + if dq_node is not None: + update_weight_metadata(dq_node.output[0], shape) if shared: if dq_node is not None: dq_clone = type(dq_node)() @@ -226,7 +239,14 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node or (activation_shape is not None and 0 in activation_shape) ): ineligible_nodes.append(node) - elif check_to_apply_transpose(node, model, tensor_name_dim_map): + elif ( + activation_shape is not None + and ( + len(activation_shape) in (2, 3) + or (len(activation_shape) == 4 and activation_shape[0] == 1) + ) + and check_to_apply_transpose(node, model, tensor_name_dim_map) + ): output_rank = tensor_rank(node.output[0]) if output_rank is None: output_rank = ( @@ -241,7 +261,8 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node node for node in ineligible_nodes if node.op_type == "MatMul" - and not any(tensor_rank(name) == 1 for name in node.input[:2]) + and all(tensor_rank(name) is not None and tensor_rank(name) >= 2 for name in node.input[:2]) + and tensor_rank(node.output[0]) is not None and (not defer_rank_conversion or node.input[1] in cache.graph_input_names) ] needs_adapters = any( @@ -313,6 +334,7 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node else: del init.dims[:] init.dims.extend(new_shape) + update_weight_metadata(init.name, new_shape) break else: # Dynamic input: Unsqueeze to 4D diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index d04784586c1..6aacaf9d66f 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -83,7 +83,12 @@ def _clamp_ir_version(model: onnx.ModelProto) -> None: def _run_ort(model: onnx.ModelProto, feeds: dict) -> dict: """Execute all graph outputs with the optional CPU runtime.""" _clamp_ir_version(model) - sess = ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"]) + options = ort.SessionOptions() + # Compare graph-surgery semantics before runtime quantization fusions. + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL + sess = ort.InferenceSession( + model.SerializeToString(), options, providers=["CPUExecutionProvider"] + ) return {o.name: sess.run([o.name], feeds)[0] for o in model.graph.output} @@ -121,8 +126,8 @@ def _make_model(nodes, graph_name, inputs, outputs, initializers=None): SymbolicShapeInference, # type: ignore[import-untyped] ) - model = SymbolicShapeInference.infer_shapes(model, auto_merge_symbolic_dims=True) or model - except Exception: + model = SymbolicShapeInference.infer_shapes(model, auto_merge=True) or model + except ImportError: pass onnx.checker.check_model(model) return model @@ -823,3 +828,87 @@ def test_legacy_ir_registers_new_initializer_inputs(): value.name for value in transformed.graph.input } _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("layout", ["symbolic_2d", "symbolic_3d", "batch_4d", "rank_5d"]) +@pytest.mark.parametrize("defer", [False, True]) +def test_ineligible_activation_layout_keeps_matmul_semantics(layout, defer): + """Route valid dynamic or unsupported Conv activation layouts through matrix fallback.""" + shape = { + "symbolic_2d": [M, N], + "symbolic_3d": [B, M, N], + "batch_4d": [2, 3, M, N], + "rank_5d": [2, 1, 3, M, N], + }[layout] + metadata = list(shape) + if layout.startswith("symbolic"): + metadata[0] = "batch" + model = _make_model( + [helper.make_node("MatMul", ["A", "W"], ["Y"], name="matrix")], + "activation_fallback", + [_vi("A", TensorProto.FLOAT, metadata)], + [_vi("Y", TensorProto.FLOAT, [*metadata[:-1], K])], + [_init("W", _RNG.standard_normal((N, K)).astype(np.float32))], + ) + feeds = {"A": _RNG.standard_normal(shape).astype(np.float32)} + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=defer) + assert any(node.op_type == "MatMul" for node in transformed.graph.node) + onnx.checker.check_model(transformed, full_check=True) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("defer", [False, True]) +def test_unknown_intermediate_result_rank_is_preserved(defer): + """Leave a dynamic matrix operation unchanged when its result rank is absent.""" + model, feeds = _make_matmul_dynamic_weight() + node = next(node for node in model.graph.node if node.op_type == "MatMul") + node.output[0] = "product" + model.graph.node.append(helper.make_node("Identity", ["product"], ["Y"])) + retained = [info for info in model.graph.value_info if info.name != "product"] + del model.graph.value_info[:] + model.graph.value_info.extend(retained) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + before = model.SerializeToString() + _apply_matmul_to_conv(model, defer_rank_conversion=defer) + onnx.checker.check_model(model, full_check=True) + assert model.SerializeToString() == before + _assert_outputs_match(expected, _run_ort(model, feeds)) + + +def test_test_fixture_inference_supplies_intermediate_metadata(): + """Exercise the fixture inference rather than silently accepting keyword failures.""" + model = _make_model( + [ + helper.make_node("Identity", ["X"], ["middle"]), + helper.make_node("Identity", ["middle"], ["Y"]), + ], + "inferred_fixture", + [_vi("X", TensorProto.FLOAT, [2, 3])], + [_vi("Y", TensorProto.FLOAT, [2, 3])], + ) + assert GraphCache(model.graph).get_shape("middle") == [2, 3] + + +@pytest.mark.parametrize("per_axis", [False, True]) +def test_dq_matrix_rewrite_matches_independent_reference(per_axis): + """Check quantized-weight rewrites against independent unfused execution.""" + model, feeds = _make_matmul_dq_weight() + # The reference evaluator implements these INT8 controls starting at opset 19. + model.opset_import[0].version = 21 + if per_axis: + for tensor in model.graph.initializer: + if tensor.name == "W_scale": + tensor.CopyFrom(_init("W_scale", np.linspace(0.01, 0.02, K, dtype=np.float32))) + elif tensor.name == "W_zp": + tensor.CopyFrom(_init("W_zp", np.zeros(K, np.int8))) + expected = ReferenceEvaluator(model).run(None, feeds)[0] + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + onnx.checker.check_model(transformed, full_check=True) + for actual in [ + ReferenceEvaluator(transformed).run(None, feeds)[0], + _run_ort(model, feeds)["Y"], + _run_ort(transformed, feeds)["Y"], + ]: + np.testing.assert_allclose(actual, expected, rtol=1e-4, atol=1e-4) From 7ba65e1d3f63c1751038a5904d621f962afa78ad Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 07:01:20 +0000 Subject: [PATCH 5/7] Preserve unsupported Gemm and fixed activation contracts Signed-off-by: mgohil --- .../dla_matmul_to_transpose_conv_transpose.py | 34 ++++----- ..._dla_matmul_to_transpose_conv_transpose.py | 73 +++++++++++++++++++ 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py index 02519dc6815..b75a2acd9fd 100644 --- a/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py +++ b/modelopt/onnx/graph_surgery/dla_transforms/dla_matmul_to_transpose_conv_transpose.py @@ -108,18 +108,6 @@ def tensor_rank(name): and any(tensor_rank(name) == 1 for name in candidate.input[:2]) ): raise ValueError(f"{candidate.name!r}: vector MatMul rank promotion is unsupported") - if candidate.op_type == "Gemm": - if any( - attribute.name in {"alpha", "beta"} and attribute.f != 1.0 - for attribute in candidate.attribute - ): - raise ValueError( - f"{candidate.name!r}: Gemm alpha/beta other than 1 are unsupported" - ) - if len(candidate.input) > 2 and candidate.input[2]: - bias_shape = cache.get_shape(candidate.input[2]) - if bias_shape is None or len(bias_shape) != 1: - raise ValueError(f"Bias for {candidate.name!r} must have a known 1-D shape") def update_weight_metadata(name, shape): """Keep private weight annotations consistent with an in-place layout change.""" @@ -226,8 +214,22 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node weight = cache.get_init(weight_name) dtype = cache.get_dtype(node.input[0]) activation_shape = cache.get_shape(node.input[0]) + unsupported_gemm = False + if node.op_type == "Gemm": + unsupported_gemm = any( + attribute.name in {"alpha", "beta"} and attribute.f != 1.0 + for attribute in node.attribute + ) + if len(node.input) > 2 and node.input[2]: + channels = ( + weight.dims[-2 if get_node_attr_i(node, "transB", 0) else -1] + if weight is not None and len(weight.dims) >= 2 + else None + ) + unsupported_gemm |= channels is None or cache.get_shape(node.input[2]) != [channels] if ( - not known_dq + unsupported_gemm + or not known_dq or weight is None or weight_name in cache.graph_input_names or any(tensor_rank(name) == 1 for name in node.input[:2]) @@ -241,6 +243,7 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node ineligible_nodes.append(node) elif ( activation_shape is not None + and node.input[0] in tensor_name_dim_map and ( len(activation_shape) in (2, 3) or (len(activation_shape) == 4 and activation_shape[0] == 1) @@ -442,11 +445,6 @@ def update_initializers(name, current_node, nodes_to_add, trans_b=False, dq_node conv_inputs = [transposed_a, conv_input1] if len(node.input) >= 3 and node.input[2]: - bias_shape = cache.get_shape(node.input[2]) - if bias_shape is not None and len(bias_shape) != 1: - raise ValueError( - f"[matmul_to_conv] Bias for {node.name!r} must be 1-D, got {bias_shape}" - ) conv_inputs.append(node.input[2]) # ── Step 5: Conv ────────────────────────────────────────────────────── diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index 6aacaf9d66f..079ef89390a 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -359,6 +359,79 @@ def test_gemm_transa_converts(): _assert_outputs_match(ref, got) +@pytest.mark.parametrize("trans_b", [False, True]) +@pytest.mark.parametrize( + "form", ["alpha", "beta", "row", "matrix", "scalar", "singleton", "symbolic", "dynamic"] +) +def test_ineligible_gemm_preserves_native_broadcast_and_attributes(form, trans_b): + """Leave valid unsupported Gemm forms intact, including broadcast biases and dynamic weights.""" + bias_shape = {"row": [1, 4], "matrix": [2, 4], "scalar": [], "singleton": [1]}.get(form, [4]) + bias = np.ones(bias_shape, np.float32) + weight = np.arange(12, dtype=np.float32).reshape(3, 4) + if trans_b: + weight = weight.T.copy() + attrs = {"transB": int(trans_b)} + if form in {"alpha", "dynamic"}: + attrs["alpha"] = 2.0 + if form == "beta": + attrs["beta"] = 0.5 + inputs = [_vi("A", TensorProto.FLOAT, [2, 3])] + feeds = {"A": np.arange(6, dtype=np.float32).reshape(2, 3)} + inits = [] + if form == "dynamic": + inputs.append(_vi("W", TensorProto.FLOAT, list(weight.shape))) + feeds["W"] = weight + else: + inits.append(_init("W", weight)) + if form == "symbolic": + inputs.append(_vi("Bias", TensorProto.FLOAT, ["channels"])) + feeds["Bias"] = bias + else: + inits.append(_init("Bias", bias)) + model = helper.make_model( + helper.make_graph( + [helper.make_node("Gemm", ["A", "W", "Bias"], ["Y"], **attrs)], + "native_gemm", + inputs, + [_vi("Y", TensorProto.FLOAT, [2, 4])], + inits, + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir_version(model) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, feeds) + before = model.SerializeToString() + transformed = _apply_matmul_to_conv(copy.deepcopy(model)) + assert transformed.SerializeToString() == before + _assert_outputs_match(expected, _run_ort(transformed, feeds)) + + +@pytest.mark.parametrize("operation", ["MatMul", "Gemm"]) +@pytest.mark.parametrize("defer", [False, True]) +def test_initializer_activation_without_value_info_preserves_values(operation, defer): + """Fixed activations need not have graph-input or intermediate metadata to remain valid.""" + model = helper.make_model( + helper.make_graph( + [helper.make_node(operation, ["A", "W"], ["Y"])], + "fixed_activation", + [], + [_vi("Y", TensorProto.FLOAT, [2, 4])], + [ + _init("A", np.arange(6, dtype=np.float32).reshape(2, 3)), + _init("W", np.arange(12, dtype=np.float32).reshape(3, 4)), + ], + ), + opset_imports=[helper.make_opsetid("", _OPSET_VERSION)], + ) + _clamp_ir_version(model) + onnx.checker.check_model(model, full_check=True) + expected = _run_ort(model, {}) + converted = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=defer) + onnx.checker.check_model(converted, full_check=True) + _assert_outputs_match(expected, _run_ort(converted, {})) + + def test_gemm_with_bias_converts(): """Gemm with 1D bias: bias forwarded to Conv as third input; output must match Gemm.""" model, feeds = _make_gemm_bias() From 009083bd2da1b7d99faadec51e7f22f659e8dd6a Mon Sep 17 00:00:00 2001 From: mgohil Date: Tue, 27 Oct 2026 08:21:33 +0000 Subject: [PATCH 6/7] Skip inference-specific DLA fixture when optional ORT tools are absent Signed-off-by: mgohil --- .../test_dla_matmul_to_transpose_conv_transpose.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index 079ef89390a..b9a6dfb4cf3 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -952,6 +952,10 @@ def test_unknown_intermediate_result_rank_is_preserved(defer): def test_test_fixture_inference_supplies_intermediate_metadata(): """Exercise the fixture inference rather than silently accepting keyword failures.""" + pytest.importorskip( + "onnxruntime.tools.symbolic_shape_infer", + reason="This fixture requires ORT symbolic shape inference", + ) model = _make_model( [ helper.make_node("Identity", ["X"], ["middle"]), From 33834cbdb7bea4032e04275ef809418098d4eb3a Mon Sep 17 00:00:00 2001 From: mgohil Date: Thu, 29 Oct 2026 07:57:41 +0000 Subject: [PATCH 7/7] Consolidate matrix conversion regression cases Signed-off-by: mgohil --- ..._dla_matmul_to_transpose_conv_transpose.py | 118 +++++------------- 1 file changed, 31 insertions(+), 87 deletions(-) diff --git a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py index b9a6dfb4cf3..d8b941b7a3c 100644 --- a/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py +++ b/tests/unit/onnx/graph_surgery/dla_transforms/test_dla_matmul_to_transpose_conv_transpose.py @@ -294,69 +294,28 @@ def _apply_and_run(model, feeds): # ── Tests: cases that SHOULD convert ───────────────────────────────────────────────────────────── -def test_matmul_2d_activation_converts(): - """2D activation → Unsqueeze[0,1] + Conv + Squeeze[0,1]; output must match MatMul.""" - model, feeds = _make_matmul_2d() - ref, got, tr = _apply_and_run(model, feeds) - - assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node), ( - "MatMul should be replaced" - ) - assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 - assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 2D input" - assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 2D input" - - _assert_outputs_match(ref, got) - - -def test_matmul_3d_activation_converts(): - """3D activation [B,M,N] → Unsqueeze[0] + Conv + Squeeze[0]; output must match MatMul.""" - model, feeds = _make_matmul_3d() - ref, got, tr = _apply_and_run(model, feeds) - - assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 - assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), "Unsqueeze expected for 3D input" - assert any(n.op_type == "Squeeze" for n in tr.graph.node), "Squeeze expected for 3D input" - - _assert_outputs_match(ref, got) - - -def test_matmul_4d_activation_converts(): - """4D activation [1,B,M,N] → Conv only, NO Unsqueeze/Squeeze; output must match MatMul.""" - model, feeds = _make_matmul_4d() - ref, got, tr = _apply_and_run(model, feeds) - - assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - assert sum(1 for n in tr.graph.node if n.op_type == "Conv") == 1 - assert not any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( - "No Unsqueeze expected for 4D input" - ) - assert not any(n.op_type == "Squeeze" for n in tr.graph.node), ( - "No Squeeze expected for 4D input" - ) - - _assert_outputs_match(ref, got) - - -def test_gemm_transb_converts(): - """Gemm transB=1 (W[K,N] → used without extra transpose): output must match Gemm.""" - model, feeds = _make_gemm_transb() - ref, got, tr = _apply_and_run(model, feeds) - - assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - - _assert_outputs_match(ref, got) - - -def test_gemm_transa_converts(): - """Gemm transA=1 (A[N,M] used transposed): output must match Gemm.""" - model, feeds = _make_gemm_transa() - ref, got, tr = _apply_and_run(model, feeds) +@pytest.mark.parametrize( + ("builder", "adapters"), + [(_make_matmul_2d, True), (_make_matmul_3d, True), (_make_matmul_4d, False)], +) +def test_matmul_activation_converts(builder, adapters): + """Convert each supported activation rank with only the required boundary adapters.""" + model, feeds = builder() + expected, actual, transformed = _apply_and_run(model, feeds) + assert not any(node.op_type in {"MatMul", "Gemm"} for node in transformed.graph.node) + assert sum(node.op_type == "Conv" for node in transformed.graph.node) == 1 + for operation in ["Unsqueeze", "Squeeze"]: + assert any(node.op_type == operation for node in transformed.graph.node) == adapters + _assert_outputs_match(expected, actual) - assert not any(n.op_type in ("MatMul", "Gemm") for n in tr.graph.node) - _assert_outputs_match(ref, got) +@pytest.mark.parametrize("builder", [_make_gemm_transb, _make_gemm_transa]) +def test_gemm_transpose_converts(builder): + """Preserve values with either supported Gemm transpose convention.""" + model, feeds = builder() + expected, actual, transformed = _apply_and_run(model, feeds) + assert not any(node.op_type in {"MatMul", "Gemm"} for node in transformed.graph.node) + _assert_outputs_match(expected, actual) @pytest.mark.parametrize("trans_b", [False, True]) @@ -467,24 +426,18 @@ def test_matmul_dq_weight_converts(): # ── Tests: cases that should NOT convert ───────────────────────────────────────────────────────── -def test_matmul_dynamic_weight_not_converted(): - """Dynamic weight: MatMul remains; both inputs and output promoted to 4D; ORT parity.""" +@pytest.mark.parametrize("defer", [False, True]) +def test_matmul_dynamic_weight_not_converted(defer): + """Keep dynamic-weight MatMul valid with immediate or deferred rank conversion.""" model, feeds = _make_matmul_dynamic_weight() - ref, got, tr = _apply_and_run(model, feeds) - - matmul_nodes = [n for n in tr.graph.node if n.op_type == "MatMul"] - assert len(matmul_nodes) == 1, "MatMul must remain when weight is a graph input" - assert not any(n.op_type == "Conv" for n in tr.graph.node), "No Conv should be added" - # Fallback: both 2D inputs promoted → Unsqueeze nodes inserted - assert any(n.op_type == "Unsqueeze" for n in tr.graph.node), ( - "Unsqueeze expected for 4D promotion of dynamic inputs" - ) - # Fallback: 2D output squeezed back - assert any(n.op_type == "Squeeze" for n in tr.graph.node), ( - "Squeeze expected to restore original output rank" - ) - - _assert_outputs_match(ref, got) + expected = _run_ort(model, feeds) + transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=defer) + onnx.checker.check_model(transformed, full_check=True) + assert sum(node.op_type == "MatMul" for node in transformed.graph.node) == 1 + assert not any(node.op_type == "Conv" for node in transformed.graph.node) + assert any(node.op_type == "Unsqueeze" for node in transformed.graph.node) + assert any(node.op_type == "Squeeze" for node in transformed.graph.node) + _assert_outputs_match(expected, _run_ort(transformed, feeds)) def test_matmul_nchw_non_unary_not_converted(): @@ -686,15 +639,6 @@ def test_integer_matmul_remains_a_valid_matmul(): ) -def test_deferred_dynamic_matmul_restores_output_rank(): - """Keep deferred graph-input weight promotion valid at the current stage boundary.""" - model, feeds = _make_matmul_dynamic_weight() - expected = _run_ort(model, feeds) - transformed = _apply_matmul_to_conv(copy.deepcopy(model), defer_rank_conversion=True) - onnx.checker.check_model(transformed, full_check=True) - _assert_outputs_match(expected, _run_ort(transformed, feeds)) - - @pytest.mark.parametrize("symbolic", [False, True]) def test_batched_weight_fallback_preserves_values_without_decoding(symbolic): """Promote batched matrix layouts using ranks even when dimension values are symbolic."""