已合并
add testcase test_guard_control for dynamo tests. #34962
rmch创建于 5月7日
add testcase test_guard_control for dynamo tests. #34962
已合并
共 1 个文件变更+555-0
| @@ -0,0 +1,555 @@ | |||
| 1 | +# Owner(s): ["module: dynamo"] | ||
| 2 | +import copy | ||
| 3 | +import functools | ||
| 4 | +import unittest | ||
| 5 | +from collections.abc import Callable | ||
| 6 | +from dataclasses import dataclass | ||
| 7 | +from typing import Any | ||
| 8 | + | ||
| 9 | +import torch | ||
| 10 | +import torch.nn as nn | ||
| 11 | +from torch._dynamo.test_case import TestCase | ||
| 12 | +from torch.overrides import TorchFunctionMode | ||
| 13 | + | ||
| 14 | + | ||
| 15 | +FilterFn = Callable[[list], list[bool]] | ||
| 16 | + | ||
| 17 | +_DICT_GUARD_TYPES = frozenset( | ||
| 18 | + { | ||
| 19 | + "DICT_VERSION", | ||
| 20 | + "DICT_KEYS", | ||
| 21 | + "DICT_KEYS_MATCH", | ||
| 22 | + "DICT_CONTAINS", | ||
| 23 | + } | ||
| 24 | +) | ||
| 25 | + | ||
| 26 | +_OPTIONAL_TYPE_GUARD_TYPES = frozenset( | ||
| 27 | + { | ||
| 28 | + "TYPE_MATCH", | ||
| 29 | + "OPTIONAL_TENSOR", | ||
| 30 | + } | ||
| 31 | +) | ||
| 32 | + | ||
| 33 | +_HASATTR_GUARD_TYPES = frozenset( | ||
| 34 | + { | ||
| 35 | + "HASATTR", | ||
| 36 | + "NOT_PRESENT_IN_GENERIC_DICT", | ||
| 37 | + } | ||
| 38 | +) | ||
| 39 | + | ||
| 40 | +_RUNTIME_STATE_GUARD_TYPES = frozenset( | ||
| 41 | + { | ||
| 42 | + "GRAD_MODE", | ||
| 43 | + "TORCH_FUNCTION_STATE", | ||
| 44 | + "GLOBAL_STATE", | ||
| 45 | + "DEFAULT_DEVICE", | ||
| 46 | + "DETERMINISTIC_ALGORITHMS", | ||
| 47 | + "AUTOCAST_STATE", | ||
| 48 | + "FSDP_TRAINING_STATE", | ||
| 49 | + } | ||
| 50 | +) | ||
| 51 | + | ||
| 52 | +_ORIGINAL_TORCH_COMPILE = None | ||
| 53 | +_TORCH_COMPILE_STACK = [] | ||
| 54 | + | ||
| 55 | + | ||
| 56 | +def _entry_type(entry) -> str: | ||
| 57 | + return getattr(entry, "guard_type", "") or "" | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +def make_npu_guard_filter( | ||
| 61 | + *, | ||
| 62 | + disable_dict_version: bool = False, | ||
| 63 | + disable_optional_type: bool = False, | ||
| 64 | + disable_hasattr: bool = False, | ||
| 65 | + disable_runtime_state: bool = False, | ||
| 66 | +) -> FilterFn: | ||
| 67 | + def _filter(entries) -> list[bool]: | ||
| 68 | + keep = [True] * len(entries) | ||
| 69 | + for i, e in enumerate(entries): | ||
| 70 | + t = _entry_type(e) | ||
| 71 | + if disable_dict_version and t in _DICT_GUARD_TYPES: | ||
| 72 | + keep[i] = False | ||
| 73 | + continue | ||
| 74 | + if disable_optional_type and t in _OPTIONAL_TYPE_GUARD_TYPES: | ||
| 75 | + keep[i] = False | ||
| 76 | + continue | ||
| 77 | + if disable_hasattr and t in _HASATTR_GUARD_TYPES: | ||
| 78 | + keep[i] = False | ||
| 79 | + continue | ||
| 80 | + if disable_runtime_state and t in _RUNTIME_STATE_GUARD_TYPES: | ||
| 81 | + keep[i] = False | ||
| 82 | + continue | ||
| 83 | + return keep | ||
| 84 | + | ||
| 85 | + return _filter | ||
| 86 | + | ||
| 87 | + | ||
| 88 | +def _supports_dynamo_guard_filter_config() -> bool: | ||
| 89 | + import torch._dynamo.config as cfg | ||
| 90 | + | ||
| 91 | + return hasattr(cfg, "_config") and "guard_filter_fn" in cfg._config | ||
| 92 | + | ||
| 93 | + | ||
| 94 | +def _set_wrapped_torch_compile(filter_fn: FilterFn | None) -> None: | ||
| 95 | + global _ORIGINAL_TORCH_COMPILE | ||
| 96 | + if _ORIGINAL_TORCH_COMPILE is None: | ||
| 97 | + _ORIGINAL_TORCH_COMPILE = torch.compile | ||
| 98 | + | ||
| 99 | + if filter_fn is None: | ||
| 100 | + torch.compile = _ORIGINAL_TORCH_COMPILE | ||
| 101 | + return | ||
| 102 | + | ||
| 103 | + | ||
| 104 | + def wrapped_compile(model=None, *args, **kwargs): | ||
| 105 | + options = kwargs.get("options") | ||
| 106 | + if options is None: | ||
| 107 | + options = {} | ||
| 108 | + else: | ||
| 109 | + options = dict(options) | ||
| 110 | + options.setdefault("guard_filter_fn", filter_fn) | ||
| 111 | + kwargs["options"] = options | ||
| 112 | + return _ORIGINAL_TORCH_COMPILE(model, *args, **kwargs) | ||
| 113 | + | ||
| 114 | + torch.compile = wrapped_compile | ||
| 115 | + | ||
| 116 | + | ||
| 117 | +def _install_global_guard_filter(filter_fn: FilterFn): | ||
| 118 | + if _supports_dynamo_guard_filter_config(): | ||
| 119 | + import torch._dynamo.config as cfg | ||
| 120 | + | ||
| 121 | + prev = getattr(cfg, "guard_filter_fn", None) | ||
| 122 | + cfg.guard_filter_fn = filter_fn | ||
| 123 | + return ("dynamo_config", prev) | ||
| 124 | + | ||
| 125 | + previous = _TORCH_COMPILE_STACK[-1] if _TORCH_COMPILE_STACK else None | ||
| 126 | + _TORCH_COMPILE_STACK.append(filter_fn) | ||
| 127 | + _set_wrapped_torch_compile(filter_fn) | ||
| 128 | + return ("torch_compile", previous) | ||
| 129 | + | ||
| 130 | + | ||
| 131 | +def _restore_global_guard_filter(state) -> None: | ||
| 132 | + kind, prev = state | ||
| 133 | + if kind == "dynamo_config": | ||
| 134 | + import torch._dynamo.config as cfg | ||
| 135 | + | ||
| 136 | + cfg.guard_filter_fn = prev | ||
| 137 | + return | ||
| 138 | + | ||
| 139 | + if _TORCH_COMPILE_STACK: | ||
| 140 | + _TORCH_COMPILE_STACK.pop() | ||
| 141 | + if _TORCH_COMPILE_STACK: | ||
| 142 | + _set_wrapped_torch_compile(_TORCH_COMPILE_STACK[-1]) | ||
| 143 | + else: | ||
| 144 | + _set_wrapped_torch_compile(None) | ||
| 145 | + | ||
| 146 | + | ||
| 147 | +def _get_installed_guard_filter(): | ||
| 148 | + if _supports_dynamo_guard_filter_config(): | ||
| 149 | + import torch._dynamo.config as cfg | ||
| 150 | + | ||
| 151 | + return getattr(cfg, "guard_filter_fn", None) | ||
| 152 | + | ||
| 153 | + compile_fn = torch.compile | ||
| 154 | + closure = getattr(compile_fn, "__closure__", None) or () | ||
| 155 | + for cell in closure: | ||
| 156 | + try: | ||
| 157 | + value = cell.cell_contents | ||
| 158 | + except ValueError: | ||
| 159 | + continue | ||
| 160 | + if callable(value) and getattr(value, "__name__", "") == "_filter": | ||
| 161 | + return value | ||
| 162 | + return None | ||
| 163 | + | ||
| 164 | + | ||
| 165 | +class NpuGuardPolicy: | ||
| 166 | + def __init__( | ||
| 167 | + self, | ||
| 168 | + *, | ||
| 169 | + filter_fn: FilterFn | None = None, | ||
| 170 | + disable_dict_version: bool = False, | ||
| 171 | + disable_optional_type: bool = False, | ||
| 172 | + disable_hasattr: bool = False, | ||
| 173 | + disable_runtime_state: bool = False, | ||
| 174 | + ): | ||
| 175 | + switches_used = any( | ||
| 176 | + [ | ||
| 177 | + disable_dict_version, | ||
| 178 | + disable_optional_type, | ||
| 179 | + disable_hasattr, | ||
| 180 | + disable_runtime_state, | ||
| 181 | + ] | ||
| 182 | + ) | ||
| 183 | + if filter_fn is not None and switches_used: | ||
| 184 | + raise ValueError("filter_fn and disable_* switches are mutually exclusive") | ||
| 185 | + | ||
| 186 | + if filter_fn is not None: | ||
| 187 | + self._fn = filter_fn | ||
| 188 | + else: | ||
| 189 | + self._fn = make_npu_guard_filter( | ||
| 190 | + disable_dict_version=disable_dict_version, | ||
| 191 | + disable_optional_type=disable_optional_type, | ||
| 192 | + disable_hasattr=disable_hasattr, | ||
| 193 | + disable_runtime_state=disable_runtime_state, | ||
| 194 | + ) | ||
| 195 | + self._prev = None | ||
| 196 | + self._active = False | ||
| 197 | + | ||
| 198 | + def __enter__(self) -> "NpuGuardPolicy": | ||
| 199 | + self._prev = _install_global_guard_filter(self._fn) | ||
| 200 | + self._active = True | ||
| 201 | + return self | ||
| 202 | + | ||
| 203 | + def __exit__(self, exc_type, exc, tb) -> None: | ||
| 204 | + if not self._active: | ||
| 205 | + return | ||
| 206 | + _restore_global_guard_filter(self._prev) | ||
| 207 | + self._prev = None | ||
| 208 | + self._active = False | ||
| 209 | + | ||
| 210 | + | ||
| 211 | + | ||
| 212 | +class FakeEntry: | ||
| 213 | + guard_type: str | ||
| 214 | + name: str = "" | ||
| 215 | + is_global: bool = False | ||
| 216 | + has_value: bool = False | ||
| 217 | + value: Any = None | ||
| 218 | + | ||
| 219 | + | ||
| 220 | +def _entries(*types: str): | ||
| 221 | + return [FakeEntry(guard_type=t, name=f"L[{i}]") for i, t in enumerate(types)] | ||
| 222 | + | ||
| 223 | + | ||
| 224 | +def _make_mlp(seed: int = 0) -> nn.Module: | ||
| 225 | + torch.manual_seed(seed) | ||
| 226 | + return nn.Sequential(nn.Linear(8, 8), nn.ReLU(), nn.Linear(8, 4)) | ||
| 227 | + | ||
| 228 | + | ||
| 229 | +def _recompile_count() -> int: | ||
| 230 | + from torch._dynamo.utils import counters | ||
| 231 | + | ||
| 232 | + return sum(counters["recompiles"].values()) | ||
| 233 | + | ||
| 234 | + | ||
| 235 | +def _reset_dynamo() -> None: | ||
| 236 | + from torch._dynamo.utils import counters | ||
| 237 | + | ||
| 238 | + torch._dynamo.reset() | ||
| 239 | + counters.clear() | ||
| 240 | + | ||
| 241 | + | ||
| 242 | +class GlobalTorchFunctionMode(TorchFunctionMode): | ||
| 243 | + def __torch_function__(self, func, types, args=(), kwargs=None): | ||
| 244 | + if kwargs is None: | ||
| 245 | + kwargs = {} | ||
| 246 | + return func(*args, **kwargs) | ||
| 247 | + | ||
| 248 | + | ||
| 249 | +class TestDictDimension(TestCase): | ||
| 250 | + def test_drops_dict_family(self): | ||
| 251 | + f = make_npu_guard_filter(disable_dict_version=True) | ||
| 252 | + es = _entries( | ||
| 253 | + "DICT_VERSION", "DICT_KEYS", "DICT_CONTAINS", "TENSOR_MATCH", "TYPE_MATCH" | ||
| 254 | + ) | ||
| 255 | + self.assertEqual(f(es), [False, False, False, True, True]) | ||
| 256 | + | ||
| 257 | + def test_inactive_keeps_all(self): | ||
| 258 | + f = make_npu_guard_filter() | ||
| 259 | + es = _entries("DICT_VERSION", "TENSOR_MATCH") | ||
| 260 | + self.assertEqual(f(es), [True, True]) | ||
| 261 | + | ||
| 262 | + | ||
| 263 | +class TestOptionalTypeDimension(TestCase): | ||
| 264 | + def test_drops_type_family(self): | ||
| 265 | + f = make_npu_guard_filter(disable_optional_type=True) | ||
| 266 | + es = _entries("TYPE_MATCH", "OPTIONAL_TENSOR", "TENSOR_MATCH", "DICT_VERSION") | ||
| 267 | + self.assertEqual(f(es), [False, False, True, True]) | ||
| 268 | + | ||
| 269 | + | ||
| 270 | +class TestHasattrDimension(TestCase): | ||
| 271 | + def test_drops_hasattr_family(self): | ||
| 272 | + f = make_npu_guard_filter(disable_hasattr=True) | ||
| 273 | + es = _entries("HASATTR", "NOT_PRESENT_IN_GENERIC_DICT", "TENSOR_MATCH") | ||
| 274 | + self.assertEqual(f(es), [False, False, True]) | ||
| 275 | + | ||
| 276 | + | ||
| 277 | +class TestRuntimeStateDimension(TestCase): | ||
| 278 | + def test_drops_runtime_state_family(self): | ||
| 279 | + f = make_npu_guard_filter(disable_runtime_state=True) | ||
| 280 | + es = _entries( | ||
| 281 | + "GRAD_MODE", | ||
| 282 | + "TORCH_FUNCTION_STATE", | ||
| 283 | + "GLOBAL_STATE", | ||
| 284 | + "DEFAULT_DEVICE", | ||
| 285 | + "DETERMINISTIC_ALGORITHMS", | ||
| 286 | + "AUTOCAST_STATE", | ||
| 287 | + "FSDP_TRAINING_STATE", | ||
| 288 | + "TENSOR_MATCH", | ||
| 289 | + ) | ||
| 290 | + self.assertEqual(f(es), [False] * 7 + [True]) | ||
| 291 | + | ||
| 292 | + | ||
| 293 | +class TestUpstreamGuardFilterCoverage(TestCase): | ||
| 294 | + def setUp(self): | ||
| 295 | + super().setUp() | ||
| 296 | + _reset_dynamo() | ||
| 297 | + | ||
| 298 | + def test_guard_filter_fn_by_id(self): | ||
| 299 | + def guard_filter_fn(entries): | ||
| 300 | + return [entry.guard_type != "ID_MATCH" for entry in entries] | ||
| 301 | + | ||
| 302 | + | ||
| 303 | + def fn(x): | ||
| 304 | + return id(x) | ||
| 305 | + | ||
| 306 | + inputs = (torch.randn(3, 2),) | ||
| 307 | + fn(*inputs) | ||
| 308 | + | ||
| 309 | + inputs_1 = (torch.randn(3, 2),) | ||
| 310 | + with torch.compiler.set_stance("fail_on_recompile"): | ||
| 311 | + self.assertEqual(fn(*inputs_1), id(inputs[0])) | ||
| 312 | + | ||
| 313 | + def test_guard_filter_fn_by_is_global(self): | ||
| 314 | + def guard_filter_fn(entries): | ||
| 315 | + return [not entry.is_global for entry in entries] | ||
| 316 | + | ||
| 317 | + global GLOBAL_INT | ||
| 318 | + | ||
| 319 | + | ||
| 320 | + def fn(x): | ||
| 321 | + return x + GLOBAL_INT | ||
| 322 | + | ||
| 323 | + GLOBAL_INT = 1 | ||
| 324 | + fn(torch.randn(3, 2)) | ||
| 325 | + | ||
| 326 | + GLOBAL_INT = 2 | ||
| 327 | + inputs = (torch.randn(3, 2),) | ||
| 328 | + with torch.compiler.set_stance("fail_on_recompile"): | ||
| 329 | + self.assertTrue(torch.equal(fn(*inputs), inputs[0] + 1)) | ||
| 330 | + | ||
| 331 | + def test_guard_filter_fn_by_name_and_value(self): | ||
| 332 | + def guard_filter_fn(entries): | ||
| 333 | + return [ | ||
| 334 | + not (entry.name == "y" and entry.value is None) for entry in entries | ||
| 335 | + ] | ||
| 336 | + | ||
| 337 | + | ||
| 338 | + def fn(x, y): | ||
| 339 | + if y is not None: | ||
| 340 | + x += y | ||
| 341 | + return x | ||
| 342 | + | ||
| 343 | + fn(torch.randn(3, 2), None) | ||
| 344 | + | ||
| 345 | + inputs = (torch.randn(3, 2), torch.tensor(1)) | ||
| 346 | + with torch.compiler.set_stance("fail_on_recompile"): | ||
| 347 | + self.assertTrue(torch.equal(fn(*inputs), inputs[0])) | ||
| 348 | + | ||
| 349 | + def test_guard_filter_inbuilt_nn_modules(self): | ||
| 350 | + class Mod(torch.nn.Module): | ||
| 351 | + def __init__(self): | ||
| 352 | + super().__init__() | ||
| 353 | + self.norm = torch.nn.LayerNorm(8) | ||
| 354 | + | ||
| 355 | + def forward(self, x): | ||
| 356 | + return self.norm(x) | ||
| 357 | + | ||
| 358 | + mod = Mod() | ||
| 359 | + opt_mod = torch.compile( | ||
| 360 | + mod, | ||
| 361 | + options={ | ||
| 362 | + "guard_filter_fn": torch.compiler.skip_guard_on_inbuilt_nn_modules_unsafe | ||
| 363 | + }, | ||
| 364 | + ) | ||
| 365 | + | ||
| 366 | + x = torch.rand(4, 8) | ||
| 367 | + opt_mod(x) | ||
| 368 | + | ||
| 369 | + mod.norm.eps = 1e-02 | ||
| 370 | + with unittest.mock.patch("torch._dynamo.config.error_on_recompile", True): | ||
| 371 | + opt_mod(x) | ||
| 372 | + | ||
| 373 | + | ||
| 374 | +class TestCombinedSwitches(TestCase): | ||
| 375 | + def test_union_drop(self): | ||
| 376 | + f = make_npu_guard_filter( | ||
| 377 | + disable_dict_version=True, | ||
| 378 | + disable_optional_type=True, | ||
| 379 | + disable_hasattr=True, | ||
| 380 | + disable_runtime_state=True, | ||
| 381 | + ) | ||
| 382 | + es = _entries( | ||
| 383 | + "DICT_VERSION", | ||
| 384 | + "TYPE_MATCH", | ||
| 385 | + "HASATTR", | ||
| 386 | + "GRAD_MODE", | ||
| 387 | + "TENSOR_MATCH", | ||
| 388 | + "ID_MATCH", | ||
| 389 | + ) | ||
| 390 | + self.assertEqual(f(es), [False, False, False, False, True, True]) | ||
| 391 | + | ||
| 392 | + def test_returns_list_of_correct_length(self): | ||
| 393 | + f = make_npu_guard_filter(disable_dict_version=True) | ||
| 394 | + es = _entries("DICT_VERSION", "TENSOR_MATCH", "TYPE_MATCH") | ||
| 395 | + keep = f(es) | ||
| 396 | + self.assertIsInstance(keep, list) | ||
| 397 | + self.assertEqual(len(keep), len(es)) | ||
| 398 | + | ||
| 399 | + | ||
| 400 | +class TestPolicy(TestCase): | ||
| 401 | + def setUp(self): | ||
| 402 | + super().setUp() | ||
| 403 | + _reset_dynamo() | ||
| 404 | + | ||
| 405 | + def test_installs_and_restores(self): | ||
| 406 | + prev = _get_installed_guard_filter() | ||
| 407 | + with NpuGuardPolicy(disable_dict_version=True): | ||
| 408 | + self.assertIsNotNone(_get_installed_guard_filter()) | ||
| 409 | + self.assertNotEqual(_get_installed_guard_filter(), prev) | ||
| 410 | + self.assertEqual(_get_installed_guard_filter(), prev) | ||
| 411 | + | ||
| 412 | + def test_restores_on_exception(self): | ||
| 413 | + prev = _get_installed_guard_filter() | ||
| 414 | + with self.assertRaises(RuntimeError): | ||
| 415 | + with NpuGuardPolicy(disable_runtime_state=True): | ||
| 416 | + raise RuntimeError("boom") | ||
| 417 | + self.assertEqual(_get_installed_guard_filter(), prev) | ||
| 418 | + | ||
| 419 | + def test_nesting(self): | ||
| 420 | + prev = _get_installed_guard_filter() | ||
| 421 | + with NpuGuardPolicy(disable_dict_version=True): | ||
| 422 | + outer_fn = _get_installed_guard_filter() | ||
| 423 | + with NpuGuardPolicy(disable_runtime_state=True): | ||
| 424 | + self.assertNotEqual(_get_installed_guard_filter(), outer_fn) | ||
| 425 | + self.assertEqual(_get_installed_guard_filter(), outer_fn) | ||
| 426 | + self.assertEqual(_get_installed_guard_filter(), prev) | ||
| 427 | + | ||
| 428 | + def test_with_custom_filter_fn(self): | ||
| 429 | + called = {"n": 0} | ||
| 430 | + | ||
| 431 | + def my_filter(entries): | ||
| 432 | + called["n"] += 1 | ||
| 433 | + return [True] * len(entries) | ||
| 434 | + | ||
| 435 | + with NpuGuardPolicy(filter_fn=my_filter): | ||
| 436 | + c = torch.compile(_make_mlp()) | ||
| 437 | + c(torch.randn(2, 8)) | ||
| 438 | + self.assertGreaterEqual(called["n"], 1) | ||
| 439 | + | ||
| 440 | + def test_filter_fn_and_switches_are_exclusive(self): | ||
| 441 | + with self.assertRaises(ValueError): | ||
| 442 | + NpuGuardPolicy( | ||
| 443 | + filter_fn=lambda e: [True] * len(e), disable_dict_version=True | ||
| 444 | + ) | ||
| 445 | + | ||
| 446 | + | ||
| 447 | +class TestParity(TestCase): | ||
| 448 | + def setUp(self): | ||
| 449 | + super().setUp() | ||
| 450 | + _reset_dynamo() | ||
| 451 | + | ||
| 452 | + def _parity(self, **switches): | ||
| 453 | + torch.manual_seed(0) | ||
| 454 | + m = _make_mlp().eval() | ||
| 455 | + x = torch.randn(2, 8) | ||
| 456 | + with torch.no_grad(): | ||
| 457 | + ref = m(x) | ||
| 458 | + | ||
| 459 | + m2 = copy.deepcopy(m) | ||
| 460 | + c = torch.compile( | ||
| 461 | + m2, options={"guard_filter_fn": make_npu_guard_filter(**switches)} | ||
| 462 | + ) | ||
| 463 | + with torch.no_grad(): | ||
| 464 | + got = c(x) | ||
| 465 | + self.assertTrue(torch.allclose(got, ref, atol=1e-5, rtol=1e-5)) | ||
| 466 | + | ||
| 467 | + def test_parity_dict(self): | ||
| 468 | + self._parity(disable_dict_version=True) | ||
| 469 | + | ||
| 470 | + def test_parity_optional_type(self): | ||
| 471 | + self._parity(disable_optional_type=True) | ||
| 472 | + | ||
| 473 | + def test_parity_hasattr(self): | ||
| 474 | + self._parity(disable_hasattr=True) | ||
| 475 | + | ||
| 476 | + def test_parity_runtime_state(self): | ||
| 477 | + self._parity(disable_runtime_state=True) | ||
| 478 | + | ||
| 479 | + def test_parity_all(self): | ||
| 480 | + self._parity( | ||
| 481 | + disable_dict_version=True, | ||
| 482 | + disable_optional_type=True, | ||
| 483 | + disable_hasattr=True, | ||
| 484 | + disable_runtime_state=True, | ||
| 485 | + ) | ||
| 486 | + | ||
| 487 | + | ||
| 488 | +class TestRuntimeStateRecompileBehavior(TestCase): | ||
| 489 | + def setUp(self): | ||
| 490 | + super().setUp() | ||
| 491 | + _reset_dynamo() | ||
| 492 | + | ||
| 493 | + def test_no_recompile_across_grad_mode_change(self): | ||
| 494 | + def foo(x): | ||
| 495 | + return x + 1 | ||
| 496 | + | ||
| 497 | + x = torch.randn(3, 2) | ||
| 498 | + compiled_fn = torch.compile( | ||
| 499 | + foo, | ||
| 500 | + options={ | ||
| 501 | + "guard_filter_fn": make_npu_guard_filter(disable_runtime_state=True) | ||
| 502 | + }, | ||
| 503 | + ) | ||
| 504 | + | ||
| 505 | + with torch.no_grad(): | ||
| 506 | + compiled_fn(x) | ||
| 507 | + | ||
| 508 | + with torch.enable_grad(), torch.compiler.set_stance("fail_on_recompile"): | ||
| 509 | + self.assertTrue(torch.equal(compiled_fn(x), foo(x))) | ||
| 510 | + | ||
| 511 | + def test_no_recompile_across_torch_function_mode_change(self): | ||
| 512 | + def foo(x): | ||
| 513 | + return x + 1 | ||
| 514 | + | ||
| 515 | + x = torch.randn(3, 2) | ||
| 516 | + with GlobalTorchFunctionMode(): | ||
| 517 | + compiled_fn = torch.compile( | ||
| 518 | + foo, | ||
| 519 | + options={ | ||
| 520 | + "guard_filter_fn": make_npu_guard_filter(disable_runtime_state=True) | ||
| 521 | + }, | ||
| 522 | + ) | ||
| 523 | + compiled_fn(x) | ||
| 524 | + | ||
| 525 | + with torch.compiler.set_stance("fail_on_recompile"): | ||
| 526 | + self.assertTrue(torch.equal(compiled_fn(x), foo(x))) | ||
| 527 | + | ||
| 528 | + | ||
| 529 | +class TestRecompileCount(TestCase): | ||
| 530 | + def setUp(self): | ||
| 531 | + super().setUp() | ||
| 532 | + _reset_dynamo() | ||
| 533 | + | ||
| 534 | + def test_no_recompile_on_repeated_same_shape(self): | ||
| 535 | + m = _make_mlp() | ||
| 536 | + c = torch.compile( | ||
| 537 | + m, | ||
| 538 | + options={ | ||
| 539 | + "guard_filter_fn": make_npu_guard_filter( | ||
| 540 | + disable_dict_version=True, | ||
| 541 | + disable_optional_type=True, | ||
| 542 | + disable_hasattr=True, | ||
| 543 | + disable_runtime_state=True, | ||
| 544 | + ) | ||
| 545 | + }, | ||
| 546 | + ) | ||
| 547 | + for _ in range(5): | ||
| 548 | + c(torch.randn(2, 8)) | ||
| 549 | + self.assertEqual(_recompile_count(), 0) | ||
| 550 | + | ||
| 551 | + | ||
| 552 | +if __name__ == "__main__": | ||
| 553 | + from torch._dynamo.test_case import run_tests | ||
| 554 | + | ||
| 555 | + run_tests() | ||