| 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
|
|
|
| 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():
|
| 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():
|
| return torch.device("npu"), {"requested": "auto", "resolved": "npu"}
|
| except Exception:
|
| pass
|
| try:
|
| import torch_npu
|
|
|
| if hasattr(torch, "npu") and torch.npu.is_available():
|
| 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"}
|
|
|
