ComfyUI/tests-unit/comfy_api_test/torch_compile_test.py
2026-07-13 14:39:10 +09:00

51 lines
1.5 KiB
Python

import torch
from comfy_api.torch_helpers import torch_compile
def test_windows_rocm_default_mode_drops_injected_guard_options(monkeypatch):
monkeypatch.setattr(torch_compile.sys, "platform", "win32")
monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
result = torch_compile.normalize_torch_compile_kwargs(
{
"backend": "inductor",
"mode": "default",
"options": {"guard_filter_fn": object()},
}
)
assert result["mode"] is None
assert result["options"] is None
def test_windows_rocm_custom_options_disable_cudagraphs(monkeypatch):
monkeypatch.setattr(torch_compile.sys, "platform", "win32")
monkeypatch.setattr(torch.version, "hip", "7.15", raising=False)
result = torch_compile.normalize_torch_compile_kwargs(
{
"backend": "inductor",
"mode": "default",
"options": {"max_autotune": True},
}
)
assert result["mode"] is None
assert result["options"] == {
"max_autotune": True,
"triton.cudagraphs": False,
"triton.cudagraph_trees": False,
}
def test_non_rocm_compile_options_are_unchanged(monkeypatch):
monkeypatch.setattr(torch_compile.sys, "platform", "linux")
compile_kwargs = {
"backend": "inductor",
"mode": "default",
"options": {"max_autotune": True},
}
assert torch_compile.normalize_torch_compile_kwargs(compile_kwargs) == compile_kwargs