From c1a7424c761302d9c2bb0d5a471e3376541f70ff Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 25 Nov 2025 22:51:39 +0300 Subject: [PATCH] ipex fix flux2 and cleanup --- modules/intel/ipex/__init__.py | 71 +++---- modules/intel/ipex/diffusers.py | 4 +- modules/intel/ipex/int_mm.py | 328 -------------------------------- 3 files changed, 31 insertions(+), 372 deletions(-) delete mode 100644 modules/intel/ipex/int_mm.py diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py index 726e5e742..6bd4cf9cd 100644 --- a/modules/intel/ipex/__init__.py +++ b/modules/intel/ipex/__init__.py @@ -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 diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py index cf585b05f..363f3a991 100644 --- a/modules/intel/ipex/diffusers.py +++ b/modules/intel/ipex/diffusers.py @@ -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 diff --git a/modules/intel/ipex/int_mm.py b/modules/intel/ipex/int_mm.py deleted file mode 100644 index 4256937ac..000000000 --- a/modules/intel/ipex/int_mm.py +++ /dev/null @@ -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