[DirectML EP] MatMul with rank-5 inputs fails the whole session at initialization (80070057) instead of falling back #33549

Describe the issue

With the DirectML EP, a session containing a MatMul whose inputs have rank 5 fails during initialization instead of falling back to the CPU EP for that node:

Exception during initialization: ...\onnxruntime\core\providers\dml\DmlExecutionProvider\src\AbiCustomRegistry.cpp(519)\onnxruntime.dll!...: Exception(3) 80070057 The parameter is incorrect.

This is the limitation discussed in #26529 (DML_GEMM_OPERATOR_DESC only supports ranks 2–4), which was closed as stale. To answer the open questions there: it happens with and without broadcasting, and it is reproducible with a single-node model, outside of Chromium/WebNN:

The whole session is lost because of one node, so the model cannot use the GPU at all.

Real-world impact

The widely used ONNX export of LaMa inpainting (Carve/LaMa-ONNX, lama_fp32.onnx) implements its FFT as MatMul(A[N,N], B[b,c,W,H,1]) (144 such nodes), so it does not run on DirectML at all. Rewriting each of them as Unsqueeze(MatMul(Squeeze(B), Transpose(A)), -1) (rank 4, same results) makes it run on a GTX 1650 Ti at 0.4 s instead of 3.8 s on the CPU. The rewrite script, for anyone hitting the same problem: https://github.com/mschneider85/silberkorn/blob/main/crates/sk-ai/models/lama_dml.py

Suggestion

Either flatten the leading batch dimensions in the DML EP before handing the GEMM to DirectML (always possible when the batch dimensions are not broadcast; with a broadcast 2-D operand the batch dimensions can be folded into the rows, as above), or not claim MatMul nodes of rank > 4 so that they fall back to the CPU EP instead of failing the session.

To reproduce

import numpy as np, onnx, onnxruntime as ort
from onnx import helper, TensorProto, numpy_helper

def model(a_shape, b_shape, a_first):
    b = helper.make_tensor_value_info("B", TensorProto.FLOAT, b_shape)
    y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, None)
    a = numpy_helper.from_array(np.ones(a_shape, np.float32), "A")
    node = helper.make_node("MatMul", ["A", "B"] if a_first else ["B", "A"], ["Y"])
    g = helper.make_graph([node], "m", [b], [y], [a])
    return helper.make_model(g, opset_imports=[helper.make_opsetid("", 17)], ir_version=8)

for name, a, b, first in [("2d_5d", [4, 4], [1, 2, 3, 4, 1], True),
                          ("5d_2d", [4, 4], [1, 2, 3, 1, 4], False),
                          ("5d_5d", [1, 2, 3, 4, 4], [1, 2, 3, 4, 1], True),
                          ("2d_4d", [4, 4], [2, 3, 4, 1], True)]:
    path = f"matmul_{name}.onnx"
    onnx.save(model(a, b, first), path)
    opts = ort.SessionOptions()
    opts.enable_mem_pattern = False
    opts.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
    try:
        s = ort.InferenceSession(path, opts, providers=["DmlExecutionProvider"])
        s.run(None, {"B": np.ones(b, np.float32)})
        print(name, "ok")
    except Exception as e:
        print(name, "failed:", str(e).splitlines()[0])

(We observed it through the C API of the DirectML NuGet build via the Rust ort crate; the models above are the ones we tested.)

Urgency

Not urgent for us (we ship the rewritten model), but every model with a rank-5 MatMul silently loses the GPU.

Platform

Windows

OS Version

Windows 11, build 26200

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.24.4 (Microsoft.ML.OnnxRuntime.DirectML NuGet), DirectML 1.15.4

ONNX Runtime API

C

Architecture

X64

Execution Provider

DirectML

Execution Provider Library Version

DirectML 1.15.4, NVIDIA GeForce GTX 1650 Ti, driver 581.95