已合并
test(test_pytorch_onnx_onnxruntime.py) fix testcase test_quantized_conv3d/test_quantized_conv3d_relu #36631
test(test_pytorch_onnx_onnxruntime.py) fix testcase test_quantized_conv3d/test_quantized_conv3d_relu #36631
已合并
yuanqi1104创建于 5月25日
2 个文件变更+125-3
@@ -3,9 +3,12 @@ from __future__ import annotations
3 3 
4import functools4import functools
5import os5import os
6+import platform
6import random7import random
8+import re
7import sys9import sys
8import unittest10import unittest
11+from contextlib import contextmanager
9from enum import auto, Enum12from enum import auto, Enum
10from typing import Optional13from typing import Optional
11 14 
@@ -64,6 +67,96 @@ skipIfQuantizationBackendQNNPack = _skipper(
64)67)
65 68 
66 69 
70+def skipIfOneDnnVersionLessThan(major, minor, patch):
71+ """Skips test if oneDNN is not enabled OR version is less than specified.
72+ 
73+ Usage:
74+ @skipIfOneDnnVersionLessThan(3, 9, 0)
75+ def test_xxx(self): ...
76+ """
77+ 
78+ @functools.lru_cache(maxsize=10)
79+ def get_mkl_dnn_info():
80+ """Extracts oneDNN enablement status and version from torch config.
81+ 
82+ Returns:
83+ A tuple: (is_enabled: bool, version: tuple)
84+ """
85+ config_str = torch.__config__.show()
86+ is_enabled = False
87+ version = (0, 0, 0)
88+ 
89+ # 1. Check if oneDNN is enabled in build settings (Strict match for USE_MKLDNN=1)
90+ # Looking for "USE_MKLDNN=1" in the Build settings line
91+ build_settings_match = re.search(r"USE_MKLDNN=1", config_str)
92+ if build_settings_match:
93+ is_enabled = True
94+ 
95+ # 2. If enabled, extract the version (Strict match for the version string)
96+ # Pattern: Intel(R) MKL-DNN vX.Y.Z ...
97+ version_pattern = r"Intel\(R\)\s+MKL-DNN\s+v(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)"
98+ 
99+ for line in config_str.splitlines():
100+ version_match = re.search(version_pattern, line)
101+ if version_match:
102+ major = int(version_match.group('major'))
103+ minor = int(version_match.group('minor'))
104+ patch = int(version_match.group('patch'))
105+ version = (major, minor, patch)
106+ break
107+ 
108+ return is_enabled, version
109+ 
110+ def condition_fn():
111+ is_enabled, current_version = get_mkl_dnn_info()
112+ 
113+ # Condition to skip:
114+ # 1. If oneDNN is not enabled, skip.
115+ # 2. If enabled but version is too low, skip.
116+ if not is_enabled:
117+ return True
118+ return current_version < (major, minor, patch)
119+ 
120+ reason = f"Requires oneDNN enabled and version >= {major}.{minor}.{patch}"
121+ return _skipper(condition_fn, reason)
122+ 
123+def useBackendOnednnOnArm(func=None, *, force=False):
124+ @contextmanager
125+ def switch_to_onednn_if_needed():
126+ is_arm = platform.machine().lower() in ("aarch64", "arm64", "armv7l", "arm")
127+ 
128+ if not is_arm:
129+ yield
130+ return
131+ 
132+ supported = torch.backends.quantized.supported_engines
133+ if "onednn" not in supported:
134+ yield
135+ return
136+ 
137+ old_engine = torch.backends.quantized.engine
138+ if old_engine == "onednn":
139+ yield
140+ return
141+ 
142+ try:
143+ torch.backends.quantized.engine = "onednn"
144+ yield
145+ finally:
146+ torch.backends.quantized.engine = old_engine
147+ 
148+ def decorator(fn):
149+ @functools.wraps(fn)
150+ def wrapper(*args, **kwargs):
151+ with switch_to_onednn_if_needed():
152+ return fn(*args, **kwargs)
153+ return wrapper
154+ 
155+ if func is None:
156+ return decorator
157+ else:
158+ return decorator(func)
159+ 
67# skips tests for all versions below min_opset_version.160# skips tests for all versions below min_opset_version.
68# add this wrapper to prevent running the test for opset_versions161# add this wrapper to prevent running the test for opset_versions
69# smaller than `min_opset_version`.162# smaller than `min_opset_version`.
@@ -1,10 +1,21 @@
1diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py1diff --git a/test/onnx/test_pytorch_onnx_onnxruntime.py b/test/onnx/test_pytorch_onnx_onnxruntime.py
2-index f993808..82c6bcf 1006442+index f9938084067..13f474258de 100644
3--- a/test/onnx/test_pytorch_onnx_onnxruntime.py3--- a/test/onnx/test_pytorch_onnx_onnxruntime.py
4+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py4+++ b/test/onnx/test_pytorch_onnx_onnxruntime.py
5-@@ -39,6 +39,15 @@ from pytorch_test_common import (5+@@ -30,15 +30,26 @@ from pytorch_test_common import (
6+ RNN_SEQUENCE_LENGTH,
7+ skipDtypeChecking,
8+ skipIfQuantizationBackendQNNPack,
9++ skipIfOneDnnVersionLessThan,
10+ skipIfUnsupportedMaxOpsetVersion,
11+ skipIfUnsupportedMinOpsetVersion,
12+ skipIfUnsupportedOpsetVersion,
13+ skipScriptTest,
14+ skipShapeChecking,
15+ skipTraceTest,
16++ useBackendOnednnOnArm
6 )17 )
7- 18+ 
8 import torch19 import torch
9+import torch_npu20+import torch_npu
10+# from torch_npu.contrib import transfer_to_npu21+# from torch_npu.contrib import transfer_to_npu
@@ -18,3 +29,21 @@ index f993808..82c6bcf 100644
18 from torch import Tensor29 from torch import Tensor
19 from torch.nn.utils import rnn as rnn_utils30 from torch.nn.utils import rnn as rnn_utils
20 from torch.onnx import errors, verification31 from torch.onnx import errors, verification
32+@@ -12809,6 +12820,8 @@ class TestONNXRuntime(onnx_test_common._TestONNXRuntime):
33+ 
34+ @skipIfUnsupportedMinOpsetVersion(10)
35+ @skipIfQuantizationBackendQNNPack
36++ @skipIfOneDnnVersionLessThan(3, 9, 0)
37++ @useBackendOnednnOnArm
38+ def test_quantized_conv3d(self):
39+ model = torch.ao.nn.quantized.Conv3d(16, 33, [2, 3, 4], stride=[3, 1, 2])
40+ # Manually initialize model weight and bias to random numbers.
41+@@ -12859,6 +12872,8 @@ class TestONNXRuntime(onnx_test_common._TestONNXRuntime):
42+ 
43+ @skipIfUnsupportedMinOpsetVersion(10)
44+ @skipIfQuantizationBackendQNNPack
45++ @skipIfOneDnnVersionLessThan(3, 9, 0)
46++ @useBackendOnednnOnArm
47+ def test_quantized_conv3d_relu(self):
48+ model = torch.ao.nn.intrinsic.quantized.ConvReLU3d(
49+ 16, 33, [2, 3, 4], stride=[3, 1, 2]