IPEX 4GB alloc detection and log driver version

installer.py

@@ -678,13 +678,11 @@ def install_ipex(torch_command):
         os.environ.setdefault('IGC_EnableDPEmulation', '1') # FP64 Emulation
     if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', None) is None:
         # XPU PyTorch doesn't support Flash Atten or Memory Atten yet so Battlemage goes OOM without this
-        # Also force enabled because of FP64 emulation making auto-detect based on the GPU invalid
         os.environ.setdefault('IPEX_FORCE_ATTENTION_SLICE', '1')
 
     if "linux" in sys.platform:
         torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.5.1+cxx11.abi torchvision==0.20.1+cxx11.abi intel-extension-for-pytorch==2.5.10+xpu oneccl_bind_pt==2.5.0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/')
-        # torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/test/xpu') # test wheels are stable previews, significantly slower than IPEX
-        # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.1')
+        # os.environ.setdefault('TENSORFLOW_PACKAGE', 'tensorflow==2.15.1 intel-extension-for-tensorflow[xpu]==2.15.0.2')
     else:
         torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+xpu torchvision==0.21.0+xpu --index-url https://download.pytorch.org/whl/test/xpu')
 

modules/devices.py

@@ -82,7 +82,12 @@ def get_backend(shared_cmd_opts):
 def get_gpu_info():
     def get_driver():
         import subprocess
-        if torch.cuda.is_available() and torch.version.cuda:
+        if torch.xpu.is_available():
+            try:
+                return torch.xpu.get_device_properties(torch.xpu.current_device()).driver_version
+            except Exception:
+                return ''
+        elif torch.cuda.is_available() and torch.version.cuda:
             try:
                 result = subprocess.run('nvidia-smi --query-gpu=driver_version --format=csv,noheader', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                 version = result.stdout.decode(encoding="utf8", errors="ignore").strip()
@@ -121,6 +126,7 @@ def get_gpu_info():
                 return {
                     'device': f'{torch.xpu.get_device_name(torch.xpu.current_device())} n={torch.xpu.device_count()}',
                     'ipex': get_package_version('intel-extension-for-pytorch'),
+                    'driver': get_driver(),
                 }
             elif backend == 'cuda' or backend == 'zluda':
                 return {

modules/intel/ipex/__init__.py

@@ -206,10 +206,10 @@ def ipex_init(): # pylint: disable=too-many-statements
             torch.cuda.ipc_collect = lambda *args, **kwargs: None
             torch.cuda.utilization = lambda *args, **kwargs: 0
 
-            ipex_hijacks(legacy=legacy)
+            device_supports_fp64, can_allocate_plus_4gb = ipex_hijacks(legacy=legacy)
             try:
                 from .diffusers import ipex_diffusers
-                ipex_diffusers()
+                ipex_diffusers(device_supports_fp64=device_supports_fp64, can_allocate_plus_4gb=can_allocate_plus_4gb)
             except Exception: # pylint: disable=broad-exception-caught
                 pass
             torch.cuda.is_xpu_hijacked = True

modules/intel/ipex/diffusers.py

@@ -6,7 +6,6 @@ from diffusers.models.attention_processor import Attention
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
-device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties("xpu").has_fp64
 attention_slice_rate = float(os.environ.get('IPEX_ATTENTION_SLICE_RATE', 4))
 
 
@@ -342,10 +341,10 @@ class AttnProcessor:
         return hidden_states
 
 
-def ipex_diffusers():
+def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False):
     diffusers.utils.torch_utils.fourier_filter = fourier_filter
     #ARC GPUs can't allocate more than 4GB to a single block:
-    if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '1' or (not device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'):
+    if not can_allocate_plus_4gb:
         diffusers.models.attention_processor.SlicedAttnProcessor = SlicedAttnProcessor
         diffusers.models.attention_processor.AttnProcessor = AttnProcessor
     if not device_supports_fp64:

modules/intel/ipex/hijacks.py

@@ -6,6 +6,16 @@ import numpy as np
 from modules import devices, errors
 
 device_supports_fp64 = torch.xpu.has_fp64_dtype() if hasattr(torch.xpu, "has_fp64_dtype") else torch.xpu.get_device_properties("xpu").has_fp64
+if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0' and (torch.xpu.get_device_properties("xpu").total_memory / 1024 / 1024 / 1024) > 4.1:
+    try:
+        x = torch.ones((33000,33000), dtype=torch.float32, device="xpu")
+        del x
+        torch.xpu.empty_cache()
+        can_allocate_plus_4gb = True
+    except Exception:
+        can_allocate_plus_4gb = False
+else:
+    can_allocate_plus_4gb = bool(os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1')
 
 # pylint: disable=protected-access, missing-function-docstring, line-too-long, unnecessary-lambda, no-else-return
 
@@ -75,7 +85,7 @@ def as_tensor(data, dtype=None, device=None):
         return original_as_tensor(data, dtype=dtype, device=device)
 
 
-if os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '-1' or (device_supports_fp64 and os.environ.get('IPEX_FORCE_ATTENTION_SLICE', '0') == '0'):
+if can_allocate_plus_4gb:
     original_torch_bmm = torch.bmm
     original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
 else:
@@ -190,6 +200,7 @@ def functional_pad(input, pad, mode='constant', value=None):
 original_torch_tensor = torch.tensor
 @wraps(torch.tensor)
 def torch_tensor(data, *args, dtype=None, device=None, **kwargs):
+    global device_supports_fp64
     if check_device(device):
         device = return_xpu(device)
     if not device_supports_fp64:
@@ -313,6 +324,7 @@ def torch_load(f, map_location=None, *args, **kwargs):
 
 # Hijack Functions:
 def ipex_hijacks(legacy=True):
+    global device_supports_fp64, can_allocate_plus_4gb
     if legacy and float(torch.__version__[:3]) < 2.5:
         torch.nn.functional.interpolate = interpolate
     torch.tensor = torch_tensor
@@ -350,3 +362,4 @@ def ipex_hijacks(legacy=True):
     if not device_supports_fp64:
         torch.from_numpy = from_numpy
         torch.as_tensor = as_tensor
+    return device_supports_fp64, can_allocate_plus_4gb