File size: 4,121 Bytes
fd32dda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | import os
from typing import Any, Dict, Tuple
import torch
def configure_torch_runtime(
torch_num_threads: int = 0,
torch_interop_threads: int = 0,
allow_tf32: bool = True,
matmul_precision: str = "high",
strict_determinism: bool = False,
) -> None:
cpu_count = os.cpu_count() or 1
n_threads = int(torch_num_threads)
if n_threads <= 0:
n_threads = cpu_count
try:
torch.set_num_threads(max(1, n_threads))
except Exception:
pass
interop = int(torch_interop_threads)
if interop <= 0:
interop = min(8, max(1, cpu_count // 2))
try:
torch.set_num_interop_threads(max(1, interop))
except Exception:
pass
try:
torch.set_float32_matmul_precision(str(matmul_precision))
except Exception:
pass
if hasattr(torch.backends, "cuda") and torch.cuda.is_available():
try:
torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
if bool(strict_determinism):
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.use_deterministic_algorithms(True, warn_only=True)
else:
torch.backends.cudnn.benchmark = True
except Exception:
pass
def _try_directml() -> Any:
try:
import torch_directml # type: ignore
return torch_directml.device()
except Exception:
return None
def resolve_device(device_spec: str = "auto", preference: str = "cuda,npu,xpu,dml,mps,cpu") -> Tuple[Any, Dict[str, str]]:
spec = str(device_spec or "auto").strip().lower()
if spec and spec != "auto":
if spec == "dml":
dml = _try_directml()
if dml is None:
raise RuntimeError("Requested device 'dml' but torch_directml is not available.")
return dml, {"requested": spec, "resolved": "dml"}
return torch.device(spec), {"requested": spec, "resolved": spec}
order = [x.strip().lower() for x in str(preference or "").split(",") if x.strip()]
if not order:
order = ["cuda", "npu", "xpu", "dml", "mps", "cpu"]
for kind in order:
if kind == "cuda":
try:
if torch.cuda.is_available():
return torch.device("cuda"), {"requested": "auto", "resolved": "cuda"}
except Exception:
pass
elif kind == "xpu":
try:
if hasattr(torch, "xpu") and torch.xpu.is_available(): # type: ignore[attr-defined]
return torch.device("xpu"), {"requested": "auto", "resolved": "xpu"}
except Exception:
pass
elif kind == "mps":
try:
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps"), {"requested": "auto", "resolved": "mps"}
except Exception:
pass
elif kind == "npu":
try:
if hasattr(torch, "npu") and torch.npu.is_available(): # type: ignore[attr-defined]
return torch.device("npu"), {"requested": "auto", "resolved": "npu"}
except Exception:
pass
try:
import torch_npu # type: ignore # noqa: F401
if hasattr(torch, "npu") and torch.npu.is_available(): # type: ignore[attr-defined]
return torch.device("npu"), {"requested": "auto", "resolved": "npu"}
except Exception:
pass
elif kind == "dml":
dml = _try_directml()
if dml is not None:
return dml, {"requested": "auto", "resolved": "dml"}
elif kind == "cpu":
return torch.device("cpu"), {"requested": "auto", "resolved": "cpu"}
return torch.device("cpu"), {"requested": "auto", "resolved": "cpu"}
|