Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

ipex fix flux2 and cleanup

Browse Source
This commit is contained in:
Disty0
2025-11-25 22:51:39 +03:00
parent da0df35106
commit c1a7424c76
3 changed files with 31 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/intel/ipex/__init__.py
+29 -42
View File
@@ -42,48 +42,50 @@ def ipex_init(): # pylint: disable=too-many-statements
torch.cuda.init = torch.xpu.init
torch.cuda.is_available = torch.xpu.is_available
torch.cuda.is_initialized = torch.xpu.is_initialized
torch.cuda.is_current_stream_capturing = lambda: False
torch.cuda.stream = torch.xpu.stream
torch.cuda.Event = torch.xpu.Event
torch.cuda.Stream = torch.xpu.Stream
torch.Tensor.cuda = torch.Tensor.xpu
torch.Tensor.is_cuda = torch.Tensor.is_xpu
torch.nn.Module.cuda = torch.nn.Module.xpu
torch.cuda.Optional = torch.xpu.Optional
torch.cuda.__cached__ = getattr(torch.xpu, "__cached__", None)
torch.cuda.__loader__ = getattr(torch.xpu, "__loader__", None)
torch.cuda.streams = torch.xpu.streams
torch.cuda.Any = torch.xpu.Any
torch.cuda.__doc__ = torch.xpu.__doc__
torch.cuda.default_generators = torch.xpu.default_generators
torch.cuda._get_device_index = torch.xpu._get_device_index
torch.cuda.__path__ = torch.xpu.__path__
torch.cuda.set_stream = torch.xpu.set_stream
torch.cuda.torch = torch.xpu.torch
torch.cuda.Union = torch.xpu.Union
torch.cuda.__annotations__ = torch.xpu.__annotations__
torch.cuda.__package__ = getattr(torch.xpu, "__package__", None)
torch.cuda.__builtins__ = torch.xpu.__builtins__
torch.cuda._lazy_init = torch.xpu._lazy_init
torch.cuda.StreamContext = torch.xpu.StreamContext
torch.cuda._lazy_call = torch.xpu._lazy_call
torch.cuda.random = torch.xpu.random
torch.cuda._get_device_index = torch.xpu._get_device_index
torch.cuda._lazy_init = torch.xpu._lazy_init
torch.cuda._lazy_call = torch.xpu._lazy_call
torch.cuda._device = torch.xpu._device
torch.cuda.__name__ = torch.xpu.__name__
torch.cuda._device_t = torch.xpu._device_t
torch.cuda.is_current_stream_capturing = lambda: False
torch.cuda.__annotations__ = torch.xpu.__annotations__
torch.cuda.__builtins__ = torch.xpu.__builtins__
torch.cuda.__name__ = torch.xpu.__name__
torch.cuda.__spec__ = torch.xpu.__spec__
torch.cuda.__file__ = torch.xpu.__file__
# torch.cuda.is_current_stream_capturing = torch.xpu.is_current_stream_capturing
torch.cuda.__path__ = torch.xpu.__path__
torch.cuda.__doc__ = torch.xpu.__doc__
torch.cuda.__package__ = getattr(torch.xpu, "__package__", None)
torch.cuda.__cached__ = getattr(torch.xpu, "__cached__", None)
torch.cuda.__loader__ = getattr(torch.xpu, "__loader__", None)
torch.Tensor.cuda = torch.Tensor.xpu
torch.Tensor.is_cuda = torch.Tensor.is_xpu
torch.nn.Module.cuda = torch.nn.Module.xpu
if torch_version[0] < 2 or (torch_version[0] == 2 and torch_version[1] < 3):
torch.cuda.threading = torch.xpu.lazy_init.threading
torch.cuda.traceback = torch.xpu.lazy_init.traceback
torch.cuda._initialization_lock = torch.xpu.lazy_init._initialization_lock
torch.cuda._initialized = torch.xpu.lazy_init._initialized
torch.cuda._is_in_bad_fork = torch.xpu.lazy_init._is_in_bad_fork
torch.cuda._lazy_seed_tracker = torch.xpu.lazy_init._lazy_seed_tracker
torch.cuda._queued_calls = torch.xpu.lazy_init._queued_calls
torch.cuda._tls = torch.xpu.lazy_init._tls
torch.cuda.threading = torch.xpu.lazy_init.threading
torch.cuda.traceback = torch.xpu.lazy_init.traceback
torch.cuda._lazy_new = torch.xpu._lazy_new
torch.cuda.FloatTensor = torch.xpu.FloatTensor
@@ -111,14 +113,16 @@ def ipex_init(): # pylint: disable=too-many-statements
if has_ipex:
torch._C._cuda_getCurrentRawStream = ipex._C._getCurrentRawStream
else:
torch.cuda.threading = torch.xpu.threading
torch.cuda.traceback = torch.xpu.traceback
torch.cuda._initialization_lock = torch.xpu._initialization_lock
torch.cuda._initialized = torch.xpu._initialized
torch.cuda._is_in_bad_fork = torch.xpu._is_in_bad_fork
torch.cuda._lazy_seed_tracker = torch.xpu._lazy_seed_tracker
torch.cuda._queued_calls = torch.xpu._queued_calls
torch.cuda._tls = torch.xpu._tls
torch.cuda.threading = torch.xpu.threading
torch.cuda.traceback = torch.xpu.traceback
torch._C._cuda_getCurrentRawStream = torch._C._xpu_getCurrentRawStream
if torch_version[0] < 2 or (torch_version[0] == 2 and torch_version[1] < 5):
@@ -137,24 +141,6 @@ def ipex_init(): # pylint: disable=too-many-statements
torch.cuda.memory_summary = torch.xpu.memory_summary
torch.cuda.memory_snapshot = torch.xpu.memory_snapshot
if torch_version[0] < 2 or (torch_version[0] == 2 and torch_version[1] < 9):
# torch._int_mm via onednn quantized matmul is supported with torch 2.9
# ipex 2.7+ has the same torch._int_mm support as torch 2.9 but doesn't support torch.compile
# torch._int_mm directly uses onednn quantized matmul
# onednn qlinear is a wrapper around onednn quantized matmul
if hasattr(torch.ops, "onednn") and hasattr(torch.ops.onednn, "qlinear_pointwise"):
def onednn_mm(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
# supports int8, fp32, fp16, and bf16 matmul with accumulation using a different dtype
# int8 matmul with onednn is slower than 16 bit with dim_size < 4096
return torch.ops.onednn.qlinear_pointwise.default(x, 1.0, 0, y, torch.ones(1, device=y.device), torch.zeros(1, device=y.device), None, 1.0, 0, torch.float32, "none", [], "none")
torch._int_mm = onednn_mm
try:
# torch.compile fix
from .int_mm import qlinear_unary
torch._inductor.mkldnn_lowerings.register_onednn_fusion_ops.qlinear_unary = qlinear_unary
except Exception:
pass
# Memory:
if "linux" in sys.platform and "WSL2" in os.popen("uname -a").read():
torch.xpu.empty_cache = lambda: None
@@ -187,15 +173,16 @@ def ipex_init(): # pylint: disable=too-many-statements
# Fix functions with ipex:
# torch.xpu.mem_get_info always returns the total memory as free memory
torch.has_cuda = True
torch.version.cuda = "12.1"
torch.backends.cuda.is_built = lambda *args, **kwargs: True
torch._utils._get_available_device_type = lambda: "xpu"
torch.xpu.mem_get_info = lambda device=None: [(torch.xpu.get_device_properties(device).total_memory - torch.xpu.memory_reserved(device)), torch.xpu.get_device_properties(device).total_memory]
torch.cuda.mem_get_info = torch.xpu.mem_get_info
torch._utils._get_available_device_type = lambda: "xpu"
torch.has_cuda = True
torch.cuda.has_half = True
torch.cuda.is_bf16_supported = getattr(torch.xpu, "is_bf16_supported", lambda *args, **kwargs: True)
torch.cuda.is_fp16_supported = lambda *args, **kwargs: True
torch.backends.cuda.is_built = lambda *args, **kwargs: True
torch.version.cuda = "12.1"
torch.cuda.get_arch_list = getattr(torch.xpu, "get_arch_list", lambda: ["pvc", "dg2", "ats-m150"])
torch.cuda.get_device_capability = lambda *args, **kwargs: (12,1)
torch.cuda.ipc_collect = lambda *args, **kwargs: None
modules/intel/ipex/diffusers.py
+2 -2
View File
@@ -23,11 +23,10 @@ class FluxPosEmbed(torch.nn.Module):
self.axes_dim = axes_dim
def forward(self, ids: torch.Tensor) -> torch.Tensor:
n_axes = ids.shape[-1]
cos_out = []
sin_out = []
pos = ids.to(dtype=torch.float32)
for i in range(n_axes):
for i in range(ids.shape[-1]):
cos, sin = diffusers.models.embeddings.get_1d_rotary_pos_embed(
self.axes_dim[i],
pos[:, i],
@@ -121,6 +120,7 @@ def ipex_diffusers(device_supports_fp64=False):
diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed
diffusers.models.embeddings.apply_rotary_emb = apply_rotary_emb
diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
diffusers.models.transformers.transformer_flux2.Flux2PosEmbed = FluxPosEmbed
diffusers.models.transformers.transformer_lumina2.apply_rotary_emb = apply_rotary_emb
diffusers.models.transformers.transformer_hidream_image.rope = hidream_rope
diffusers.models.transformers.transformer_chroma.FluxPosEmbed = FluxPosEmbed
modules/intel/ipex/int_mm.py
-328
View File
@@ -1,328 +0,0 @@
import functools
import torch
import torch._inductor
from torch._inductor.select_algorithm import ExternKernelChoice, ChoiceCaller, autotune_select_algorithm, extern_kernels
from torch._inductor.utils import use_aten_gemm_kernels, use_cpp_gemm_template, use_max_autotune
from torch._inductor.codegen.cpp_gemm_template import CppGemmTemplate
from torch._inductor.codegen.cpp_utils import create_epilogue_with_attr
from torch._inductor.lowering import register_lowering, lowerings, view
from torch._inductor.kernel.mm_common import mm_args
from torch._inductor import ir, mkldnn_ir
from torch._inductor.ir import TensorBox
from torch._inductor.virtualized import ops, V
lowerings.pop(extern_kernels.qlinear_pointwise)
del extern_kernels.qlinear_pointwise
aten_mkldnn_qlinear_unary = ExternKernelChoice(
torch.ops.onednn.qlinear_pointwise,
"onednn::qlinear_pointwise",
has_out_variant=False,
kernel_creator=mkldnn_ir.QLinearPointwisePT2E.create,
)
@register_lowering(torch.ops.onednn.qlinear_pointwise, type_promotion_kind=None)
@register_lowering(torch.ops.onednn.qlinear_pointwise.default, type_promotion_kind=None)
def qlinear_unary(
x: TensorBox,
x_scale,
x_zp,
packed_weight: TensorBox,
w_scale: TensorBox,
w_zp: TensorBox,
bias: TensorBox,
o_scale,
o_zero_point,
output_dtype,
attr,
scalars,
algorithm,
layout=None,
):
assert packed_weight.get_dtype() is torch.int8, (
"Only int8 weights are supported by oneDNN qlinear."
)
x_size = x.get_size()
if len(x_size) > 2:
# GEMM template needs 2D input, normalize input shape here
x = view(x, [-1, x_size[-1]])
if not isinstance(x_scale, ir.TensorBox):
assert isinstance(x_scale, float)
x_scale = V.graph.add_tensor_constant(
torch.tensor(x_scale, dtype=torch.float32), name="x_scale"
)
else:
x_scale.realize()
if all(dim == 1 for dim in x_scale.get_size()):
# Corner-case discovered with LLaMA series.
# If all outer dims of x_scale are 1, make it a 0D tensor.
# Otherwise, epilogue creator will run into indexing issues.
x_scale = view(x_scale, [])
assert len(x_scale.get_size()) in [0, 1], "x_scale must be 0D or 1D"
if x_zp is None:
# If x_zp is None, x is int8 quantized per-tensor and its scale is not reshaped,
# then the codegened code would segfault if we don't create a tensor for x_zp.
# It's safe to do so since x is a symmetrically quantized int8 tensor.
# Moreover, oneDNN qlinear API doesn't accept None value for zp
x_zp = V.graph.add_tensor_constant(
torch.tensor(0, dtype=torch.int32), name="x_zp"
)
if not isinstance(x_zp, ir.TensorBox):
assert isinstance(x_zp, int)
x_zp = V.graph.add_tensor_constant(
torch.tensor(x_zp, dtype=torch.int32), name="x_zp"
)
else:
x_zp.realize()
assert x_zp.get_numel() == 1, "x_zp is incompatible with oneDNN qlinear"
# When channels less than 8, w_scale/w_zp is Pointwise instead of ConstantBuffer
# Refer to https://github.com/pytorch/pytorch/blob
# /f353d17755ed23b02924c962a86ff99a3405fe10/torch/_inductor/graph.py#L570-L577
if w_zp is None:
# If w_zp is None, then it's a dummy tensor created to denote the
# absence of a zero point, and thus w is int8 symmetrically quantized.
# Moreover, oneDNN qlinear API doesn't accept None value for zp
w_zp = V.graph.add_tensor_constant(
torch.tensor(0, dtype=torch.int32), name="w_zp"
)
w_scale.realize()
w_zp.realize()
if w_zp.get_dtype() != torch.int32 and isinstance(
ir.InputsKernel.unwrap_storage_for_input(w_zp),
ir.ConstantBuffer,
):
# W_zp might be a ConstantBuffer with int64, convert it to int32
w_zp_tensor = V.graph.constants[w_zp.get_name()].to(torch.int32)
w_zp = V.graph.add_tensor_constant(
torch.tensor(w_zp_tensor, dtype=torch.int32), name=w_zp.get_name()
)
bias_dtype = None if bias is None else bias.get_dtype()
choices: list[ChoiceCaller] = []
if use_max_autotune():
*_, layout, x, packed_weight = mm_args(
x, packed_weight, layout=layout, out_dtype=output_dtype
)
if (
# GEMM template currently only supports symmetrically quantized weights
isinstance(
ir.InputsKernel.unwrap_storage_for_input(w_zp),
ir.ConstantBuffer,
)
and torch.equal(
torch.zeros_like(V.graph.constants[w_zp.get_name()]),
V.graph.constants[w_zp.get_name()],
)
) and use_cpp_gemm_template(layout, x, packed_weight):
W_tensor = V.graph.constants[packed_weight.get_name()].to_dense()
weight_compens_tensor = torch.sum(W_tensor.to(torch.float), dim=0)
weight_compens = V.graph.add_tensor_constant(
weight_compens_tensor,
name=packed_weight.get_name() + "_BMatrixCompens",
)
def epilogue_creator(input_buffer):
# Epilogue to convert from s32 to f32 for u8s8f32
assert output_dtype in [
torch.float32,
torch.bfloat16,
torch.uint8,
torch.int8,
]
input_loader = input_buffer.make_loader()
weight_compens_loader = weight_compens.make_loader()
x_scale_loader = x_scale.make_loader()
w_scale_loader = w_scale.make_loader()
x_zp_loader = x_zp.make_loader()
nonlocal bias
bias_loader = None
if bias is not None:
bias_loader = bias.make_loader()
def inner_fn(index):
nonlocal bias
input = input_loader(index)
# MicroKernel Output is with int32
# cvt to FP32 before doing compensation
input = ops.to_dtype(input, torch.float32)
weight_compens_index = (index[-1],)
_x_scale = x_scale_loader(())
_x_zp = x_zp_loader(())
_w_scale = w_scale_loader(weight_compens_index)
_weight_compo = weight_compens_loader(weight_compens_index)
# Step 1: Compute s8s8->s32 or u8s8->s32 GEMM & then apply compensation
temp = ops.mul(
ops.mul(
input,
_x_scale,
),
_w_scale,
)
# NOTE: We will apply compensation even if the x_zp is 0 for int8 quantization.
# That's because when torch.compile is invoked for dynamic quantization,
# x might coincidentally have such values that x_zp might be zero despite
# asymmetric quantization.
# Besides, if x_zp is dummy for int8 x, or if x is statically quantized,
# we'd still perform that redundant compute to avoid making the code messy
# because we discovered that redundant computation of compensation did not
# lead to performance degradation with the input shapes tested.
temp = ops.sub(
temp,
ops.mul(
ops.mul(
ops.mul(
_x_scale,
_w_scale,
),
_x_zp,
),
_weight_compo,
),
)
# Step 2: add Bias if applicable
if bias is not None:
_bias = bias_loader(weight_compens_index)
nonlocal bias_dtype
assert bias_dtype in [torch.float32, torch.bfloat16]
if bias_dtype == torch.bfloat16:
_bias = ops.to_dtype(_bias, torch.float32)
temp = ops.add(temp, _bias)
return temp
output_buf = ir.Pointwise(
device=input_buffer.get_device(),
dtype=torch.float32, # Hardcode to FP32 for u8s8f32 & s8s8f32
inner_fn=inner_fn,
ranges=input_buffer.get_size(),
)
# Step 3: Doing the unary post op fusion
if attr != "none":
output_buf = create_epilogue_with_attr(
output_buf, attr, scalars=scalars, algorithm=algorithm
)
# Step 4: Cast output to Target Dtype
if output_dtype == torch.bfloat16:
output_cast_loader = output_buf.make_loader()
def inner_fn_cast_output_to_bf16(index):
input = output_cast_loader(index)
return ops.to_dtype(input, output_dtype)
output_buf = ir.Pointwise(
device=output_buf.get_device_or_error(),
dtype=output_dtype,
inner_fn=inner_fn_cast_output_to_bf16,
ranges=output_buf.get_size(),
)
elif output_dtype in [torch.uint8, torch.int8]:
from .lowering import _create_constants
requant_input_loader = output_buf.make_loader()
def inner_fn_requant(index, scale, zero_point):
input = requant_input_loader(index)
inv_scale, zero_point = _create_constants(
1.0 / scale, zero_point, dtype=torch.float32
)
val = ops.round(input * inv_scale) + zero_point
if output_dtype == torch.uint8:
qmin, qmax = _create_constants(
0, 255, dtype=torch.float32
)
else:
qmin, qmax = _create_constants(
-128, 127, dtype=torch.float32
)
clamped = ops.minimum(ops.maximum(val, qmin), qmax)
return ops.to_dtype(clamped, output_dtype)
output_buf = ir.Pointwise(
device=output_buf.get_device_or_error(),
dtype=output_dtype,
inner_fn=functools.partial(
inner_fn_requant,
scale=float(o_scale),
zero_point=int(o_zero_point),
),
ranges=output_buf.get_size(),
)
return output_buf
assert x.get_dtype() in [torch.uint8, torch.int8]
CppGemmTemplate.add_choices(
choices,
layout,
[x, x_scale, x_zp, packed_weight, w_scale, w_zp]
if bias is None
else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias],
has_bias=bias is not None,
epilogue_creator=epilogue_creator,
input_indices=[0, 3, 1, 2, 4, 5]
if bias is None
else [6, 0, 3, 1, 2, 4, 5],
)
if len(choices) == 0 or use_aten_gemm_kernels():
kwargs = dict(
output_scale=o_scale,
output_zero_point=o_zero_point,
output_dtype=output_dtype,
post_op_name=attr,
post_op_args=scalars,
post_op_algorithm=algorithm,
)
if bias is None:
kwargs["bias"] = None
choices.append(
aten_mkldnn_qlinear_unary.bind(
(x, x_scale, x_zp, packed_weight, w_scale, w_zp)
if bias is None
else (x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias),
layout,
**kwargs,
)
)
# this line is not needed and causes unnecessary errors
#assert packed_weight.get_name() in V.graph.constants
input_gen_fns = {
3: lambda x: V.graph.constants[x.get_name()], # packed weight
4: lambda x: V.graph.constants[x.get_name()], # weight scale
5: lambda x: V.graph.constants[x.get_name()], # weight zp
6: lambda x: V.graph.constants[x.get_name()], # bias
}
if isinstance(
ir.InputsKernel.unwrap_storage_for_input(x_scale),
ir.ConstantBuffer,
):
# x is statically quantized
input_gen_fns[1] = lambda x: V.graph.constants[x.get_name()]
if isinstance(
ir.InputsKernel.unwrap_storage_for_input(x_zp),
ir.ConstantBuffer,
):
input_gen_fns[2] = lambda x: V.graph.constants[x.get_name()]
result = autotune_select_algorithm(
"qlinear_unary",
choices,
[x, x_scale, x_zp, packed_weight, w_scale, w_zp]
if bias is None
else [x, x_scale, x_zp, packed_weight, w_scale, w_zp, bias],
layout,
input_gen_fns=input_gen_fns,
)
if len(x_size) > 2:
result = view(result, (*x_size[:-1], result.get_size()[-1]))
return result