58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
|
"""GPU-accelerated compressed execution path for XIC.
|
||
|
|
|
||
|
|
has_gpu() probes for CUDA via torch. If torch is absent or no CUDA
|
||
|
|
device is detected, returns False and run_on_gpu() falls back to CPU
|
||
|
|
via execute_gx() with a clear log line.
|
||
|
|
"""
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
def has_gpu() -> bool:
|
||
|
|
"""Check if CUDA GPU is available via torch.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
True if torch is installed and CUDA device is detected, False otherwise
|
||
|
|
"""
|
||
|
|
try:
|
||
|
|
import torch
|
||
|
|
return torch.cuda.is_available()
|
||
|
|
except ImportError:
|
||
|
|
return False
|
||
|
|
|
||
|
|
|
||
|
|
def run_on_gpu(model_path: str, params: dict) -> Any:
|
||
|
|
"""Execute a .gx model with optional GPU acceleration.
|
||
|
|
|
||
|
|
If GPU is available (torch + CUDA), logs device info and runs on GPU.
|
||
|
|
If GPU is not available, logs fallback and runs on CPU via execute_gx().
|
||
|
|
|
||
|
|
Args:
|
||
|
|
model_path: Path to .gx model file
|
||
|
|
params: Execution parameters dict (trace, profile, etc.)
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
ExecutionContext from execute_gx()
|
||
|
|
"""
|
||
|
|
from runtime_executor.runner import execute_gx
|
||
|
|
|
||
|
|
if has_gpu():
|
||
|
|
try:
|
||
|
|
import torch
|
||
|
|
device_name = torch.cuda.get_device_name(0)
|
||
|
|
print(f"[XIC-GPU] Device: {device_name}")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"[XIC-GPU] Warning: Could not get device name: {e}")
|
||
|
|
|
||
|
|
return execute_gx(
|
||
|
|
model_path,
|
||
|
|
trace=params.get("trace", False),
|
||
|
|
profile=params.get("profile", False),
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
print("[XIC-GPU] No CUDA device — executing on CPU")
|
||
|
|
return execute_gx(
|
||
|
|
model_path,
|
||
|
|
trace=params.get("trace", False),
|
||
|
|
profile=params.get("profile", False),
|
||
|
|
)
|