IPEX fix dtype errors when GPU supports 64 bit

modules/intel/ipex/__init__.py

@@ -165,11 +165,6 @@ def ipex_init(): # pylint: disable=too-many-statements
 
  ipex_hijacks()
  if not torch.xpu.has_fp64_dtype():
- try:
- from .attention import attention_init
- attention_init()
- except Exception: # pylint: disable=broad-exception-caught
- pass
  try:
  from .diffusers import ipex_diffusers
  ipex_diffusers()

modules/intel/ipex/attention.py

@@ -4,11 +4,8 @@
 # pylint: disable=protected-access, missing-function-docstring, line-too-long
 
 original_torch_bmm = torch.bmm
-def torch_bmm(input, mat2, *, out=None):
- if input.dtype != mat2.dtype:
- mat2 = mat2.to(input.dtype)
-
- #ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+def torch_bmm_32_bit(input, mat2, *, out=None):
+ # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
  batch_size_attention, input_tokens, mat2_shape = input.shape[0], input.shape[1], mat2.shape[2]
  block_multiply = input.element_size()
  slice_block_size = input_tokens * mat2_shape / 1024 / 1024 * block_multiply
@@ -17,7 +14,7 @@ def torch_bmm(input, mat2, *, out=None):
  split_slice_size = batch_size_attention
  if block_size > 4:
  do_split = True
- #Find something divisible with the input_tokens
+ # Find something divisible with the input_tokens
  while (split_slice_size * slice_block_size) > 4:
  split_slice_size = split_slice_size // 2
  if split_slice_size <= 1:
@@ -30,7 +27,7 @@ def torch_bmm(input, mat2, *, out=None):
  if split_slice_size * slice_block_size > 4:
  slice_block_size2 = split_slice_size * mat2_shape / 1024 / 1024 * block_multiply
  do_split_2 = True
- #Find something divisible with the input_tokens
+ # Find something divisible with the input_tokens
  while (split_2_slice_size * slice_block_size2) > 4:
  split_2_slice_size = split_2_slice_size // 2
  if split_2_slice_size <= 1:
@@ -64,8 +61,8 @@ def torch_bmm(input, mat2, *, out=None):
  return hidden_states
 
 original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
-def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
- #ARC GPUs can't allocate more than 4GB to a single block, Slice it:
+def scaled_dot_product_attention_32_bit(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
+ # ARC GPUs can't allocate more than 4GB to a single block, Slice it:
  if len(query.shape) == 3:
  batch_size_attention, query_tokens, shape_four = query.shape
  shape_one = 1
@@ -74,19 +71,14 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
  shape_one, batch_size_attention, query_tokens, shape_four = query.shape
  no_shape_one = False
 
- if query.dtype != key.dtype:
- key = key.to(dtype=query.dtype)
- if query.dtype != value.dtype:
- value = value.to(dtype=query.dtype)
-
  block_multiply = query.element_size()
  slice_block_size = shape_one * query_tokens * shape_four / 1024 / 1024 * block_multiply
  block_size = batch_size_attention * slice_block_size
 
  split_slice_size = batch_size_attention
  if block_size > 6:
  do_split = True
- #Find something divisible with the shape_one
+ # Find something divisible with the shape_one
  while (split_slice_size * slice_block_size) > 4:
  split_slice_size = split_slice_size // 2
  if split_slice_size <= 1:
@@ -99,7 +91,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
  if split_slice_size * slice_block_size > 6:
  slice_block_size2 = shape_one * split_slice_size * shape_four / 1024 / 1024 * block_multiply
  do_split_2 = True
- #Find something divisible with the batch_size_attention
+ # Find something divisible with the batch_size_attention
  while (split_2_slice_size * slice_block_size2) > 4:
  split_2_slice_size = split_2_slice_size // 2
  if split_2_slice_size <= 1:
@@ -155,8 +147,3 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
  query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal
  )
  return hidden_states
-
-def attention_init():
- #ARC GPUs can't allocate more than 4GB to a single block:
- torch.bmm = torch_bmm
- torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention

modules/intel/ipex/hijacks.py

@@ -93,6 +93,31 @@ def linalg_solve(A, B, *args, **kwargs): # pylint: disable=invalid-name
  else:
  return original_linalg_solve(A, B, *args, **kwargs)
 
+if torch.xpu.has_fp64_dtype():
+ original_torch_bmm = torch.bmm
+ original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
+else:
+ # 64 bit attention workarounds for Alchemist:
+ try:
+ from .attention import torch_bmm_32_bit as original_torch_bmm
+ from .attention import scaled_dot_product_attention_32_bit as original_scaled_dot_product_attention
+ except Exception: # pylint: disable=broad-exception-caught
+ original_torch_bmm = torch.bmm
+ original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
+
+# dtype errors:
+def torch_bmm(input, mat2, *, out=None):
+ if input.dtype != mat2.dtype:
+ mat2 = mat2.to(input.dtype)
+ return original_torch_bmm(input, mat2, out=out)
+
+def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
+ if query.dtype != key.dtype:
+ key = key.to(dtype=query.dtype)
+ if query.dtype != value.dtype:
+ value = value.to(dtype=query.dtype)
+ return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal)
+
 @property
 def is_cuda(self):
  return self.device.type == 'xpu'
@@ -184,11 +209,16 @@ def ipex_hijacks():
  lambda orig_func, *args, **kwargs: True)
 
  # Functions that make compile mad with CondFunc:
- torch.utils.data.dataloader._MultiProcessingDataLoaderIter._shutdown_workers = _shutdown_workers
  torch.nn.DataParallel = DummyDataParallel
+ torch.utils.data.dataloader._MultiProcessingDataLoaderIter._shutdown_workers = _shutdown_workers
+
  torch.autocast = ipex_autocast
- torch.cat = torch_cat
- torch.linalg.solve = linalg_solve
+ torch.backends.cuda.sdp_kernel = return_null_context
  torch.UntypedStorage.is_cuda = is_cuda
+
  torch.nn.functional.interpolate = interpolate
- torch.backends.cuda.sdp_kernel = return_null_context
+ torch.linalg.solve = linalg_solve
+
+ torch.bmm = torch_bmm
+ torch.cat = torch_cat
+ torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention