Batch-transpose fusion into `FusedMatMul` sets `transBatchA` on inputs of different ranks, which ONNX Runtime can't compute

オープン
  • 何か問題がありますか？
#3,156 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 4 日以内に返信

まだ誰も着手していません。

評価

難易度
3/5
見積もり時間
1〜2日
初心者へのやさしさ
72/100
issue の種類
バグ
明瞭さ
明確に書かれている
活発さ
活発
技術スタック
python
領域
backend, machine-learning

調査の方向性

Start by finding the ORT fusion rules that fold Transpose into FusedMatMul as transBatchA or transBatchB, then trace how they inspect input shapes and ranks. Reproduce the reported case with rank 3 x and rank 2 w, and check that the fusion is skipped when the ranks differ. Done means the optimized graph runs in ONNX Runtime and matches the unfused result.

索引モデルが issue の本文から書いたものです。

説明

The rules that fold a Transpose into FusedMatMul as transBatchA / transBatchB don't check the other input's rank. ONNX Runtime's FusedMatMul only batch-transposes when both inputs have the same rank, of at least 3, so the fused graph fails at runtime. A projection of a transposed activation, Transpose(x, [1, 0, 2]) @ w with x rank 3 and w rank 2, is enough.

Repro (onnxscript 0.7.2, onnxruntime 1.31.1):

"""onnxscript folds a batch `Transpose` into `FusedMatMul` (`transBatchA`) when the other input has a lower rank, which
ONNX Runtime can't compute.

`Div(MatMul(Transpose(x, perm=[1, 0, 2]), w), 2)` with `x` rank 3 and `w` rank 2 (a projection of a transposed
activation). onnxscript's ORT fusions turn it into `FusedMatMul(x, w, transBatchA=1, alpha=0.5)`.
"""

import numpy as np
import onnx_ir as ir
import onnxruntime as ort
from onnx import TensorProto, helper
from onnxscript.rewriter.ort_fusions import optimize_for_ort

x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [4, 2, 8])
w = helper.make_tensor("w", TensorProto.FLOAT, [8, 16], np.random.default_rng(0).standard_normal((8, 16)).astype(np.float32).ravel())
two = helper.make_tensor("two", TensorProto.FLOAT, [], [2.0])
nodes = [
    helper.make_node("Transpose", ["x"], ["xt"], perm=[1, 0, 2]),
    helper.make_node("MatMul", ["xt", "w"], ["product"]),
    helper.make_node("Div", ["product", "two"], ["y"]),
]
y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 4, 16])
proto = helper.make_model(helper.make_graph(nodes, "g", [x], [y], [w, two]), opset_imports=[helper.make_opsetid("", 20)], ir_version=10)

feeds = {"x": np.random.default_rng(1).standard_normal((4, 2, 8)).astype(np.float32)}
expected = ort.InferenceSession(proto.SerializeToString(), providers=["CPUExecutionProvider"]).run(None, feeds)[0]

model, _ = optimize_for_ort(ir.serde.deserialize_model(proto))
for node in model.graph.all_nodes():
    print(node.op_type, {k: a.value for k, a in node.attributes.items()}, [str(v.shape) for v in node.inputs if v is not None])
try:
    got = ort.InferenceSession(ir.to_proto(model).SerializeToString(), providers=["CPUExecutionProvider"]).run(None, feeds)[0]
    print("fused graph max diff:", np.abs(got - expected).max())
except Exception as error:
    print("ONNX Runtime:", str(error).split("Status Message: ")[-1].strip())

FusedMatMul {'alpha': 0.5, 'transBatchA': 1} ['[4,2,8]', '[8,16]']
ONNX Runtime: .../matmul_helper.h:73 ... left_num_dims > 2 && left_num_dims == right_num_dims was false. Two inputs should have same rank and rank >= 3 if transBatchA or transBatchB is true

The unfused graph runs and gives [2, 4, 16].

主要言語
Python
スター
457
フォーク
149
平均マージ
3日 20時間
マージ済み PR（30日）
22

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

microsoft/onnxscript のほかの issue

microsoft/onnxscript の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。