已合并
[refactor]decomp_2.10.0 #40576
yeh创建于 7月9日
[refactor]decomp_2.10.0 #40576
已合并
共 4 个文件变更+11-228
| @@ -17,7 +17,6 @@ from .utils import patch_has_triton, patch_device_supports_tma, patch_is_gpu, ge | |||
| 17 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system | 17 | from .codegen.common import register_device_op_overrides_npu, patch_cache_base_get_system |
| 18 | from .shape_handling import NPUShapeHandling, patch_shape_handling | 18 | from .shape_handling import NPUShapeHandling, patch_shape_handling |
| 19 | from ._npu_meta_registration import npu_patch_meta | 19 | from ._npu_meta_registration import npu_patch_meta |
| 20 | - | ||
| 21 | # 顶层 patch：所有 inductor backend（triton / mlir / dvm / ascendc）都需要的 NPU 设备级patch， | 20 | # 顶层 patch：所有 inductor backend（triton / mlir / dvm / ascendc）都需要的 NPU 设备级patch， |
| 22 | # 与 codegen 后端选择无关，在任何 backend loader 之前无条件执行 | 21 | # 与 codegen 后端选择无关，在任何 backend loader 之前无条件执行 |
| 23 | npu_patch_meta() | 22 | npu_patch_meta() |
| @@ -124,7 +123,7 @@ def _load_triton_backend(): | |||
| 124 | patch_get_optimization_cflags, | 123 | patch_get_optimization_cflags, |
| 125 | ) | 124 | ) |
| 126 | from .codegen.cpp_utils import patch_device_to_aten | 125 | from .codegen.cpp_utils import patch_device_to_aten |
| 127 | - from .decomposition import _register_npu_inductor_decompositions | 126 | + from .decomposition import _register_triton_decompositions |
| 128 | from .dependencies import patch_extract_read_writes | 127 | from .dependencies import patch_extract_read_writes |
| 129 | from .fx_passes import patch_pattern_mm_plus_mm, register_fav3_partition_pass | 128 | from .fx_passes import patch_pattern_mm_plus_mm, register_fav3_partition_pass |
| 130 | from .fx_passes.graph_match_pass import ( | 129 | from .fx_passes.graph_match_pass import ( |
| @@ -186,7 +185,7 @@ def _load_triton_backend(): | |||
| 186 | _register_npu_inductor_fallbacks, | 185 | _register_npu_inductor_fallbacks, |
| 187 | ) | 186 | ) |
| 188 | 187 | ||
| 189 | - _register_npu_inductor_decompositions(backend="triton") | 188 | + _register_triton_decompositions() |
| 190 | 189 | ||
| 191 | if npu_config.enable_full_lowering_fallback.strip() == "allfallback": | 190 | if npu_config.enable_full_lowering_fallback.strip() == "allfallback": |
| 192 | _enable_full_lowering_fallback() | 191 | _enable_full_lowering_fallback() |
| @@ -1,210 +0,0 @@ | |||
| 1 | -import functools | ||
| 2 | -from typing import Optional, Tuple | ||
| 3 | - | ||
| 4 | -import torch | ||
| 5 | -import torch.nn.functional as F | ||
| 6 | -from torch._inductor import decomposition as inductor_decomp | ||
| 7 | -from torch._C import DispatchKey | ||
| 8 | -from torch import Tensor | ||
| 9 | - | ||
| 10 | -from torch._decomp import ( | ||
| 11 | - remove_decompositions, | ||
| 12 | -) | ||
| 13 | - | ||
| 14 | -from .. import config as anir_config | ||
| 15 | - | ||
| 16 | -aten = torch.ops.aten | ||
| 17 | -npu = torch.ops.npu | ||
| 18 | - | ||
| 19 | -remove_decompositions(inductor_decomp.decompositions, anir_config.decomps_to_exclude_npu) | ||
| 20 | - | ||
| 21 | -# Batch_norm_decomposition function registered to fix dynamic shape dynamo tracing issue. | ||
| 22 | - | ||
| 23 | - | ||
| 24 | -def batch_norm_decomposition( | ||
| 25 | - input: Tensor, | ||
| 26 | - weight: Optional[Tensor], | ||
| 27 | - bias: Optional[Tensor], | ||
| 28 | - running_mean: Optional[Tensor], | ||
| 29 | - running_var: Optional[Tensor], | ||
| 30 | - training: bool, | ||
| 31 | - momentum: float, | ||
| 32 | - eps: float, | ||
| 33 | - cudnn_enabled: bool, | ||
| 34 | -) -> Tensor: | ||
| 35 | - if input.numel() == 0: | ||
| 36 | - out = input.clone() | ||
| 37 | - if weight is not None: | ||
| 38 | - out *= weight[0] | ||
| 39 | - if bias is not None: | ||
| 40 | - out += bias[0] | ||
| 41 | - return out | ||
| 42 | - return aten._batch_norm_impl_index.default( | ||
| 43 | - input, | ||
| 44 | - weight, | ||
| 45 | - bias, | ||
| 46 | - running_mean, | ||
| 47 | - running_var, | ||
| 48 | - training, | ||
| 49 | - momentum, | ||
| 50 | - eps, | ||
| 51 | - cudnn_enabled, | ||
| 52 | - )[0] | ||
| 53 | - | ||
| 54 | -def npu_convolution_backward( | ||
| 55 | - grad_output, | ||
| 56 | - input, | ||
| 57 | - weight, | ||
| 58 | - bias_sizes, | ||
| 59 | - stride, | ||
| 60 | - padding, | ||
| 61 | - dilation, | ||
| 62 | - transposed, | ||
| 63 | - output_padding, | ||
| 64 | - groups, | ||
| 65 | - output_mask, | ||
| 66 | -): | ||
| 67 | - if not output_mask[2]: | ||
| 68 | - return NotImplemented | ||
| 69 | - grad_bias = torch.ops.aten.sum(grad_output, [0] + list(range(2, grad_output.dim()))) | ||
| 70 | - grad_inp, grad_weight, _ = torch.ops.aten.convolution_backward( | ||
| 71 | - grad_output, | ||
| 72 | - input, | ||
| 73 | - weight, | ||
| 74 | - bias_sizes, | ||
| 75 | - stride, | ||
| 76 | - padding, | ||
| 77 | - dilation, | ||
| 78 | - transposed, | ||
| 79 | - output_padding, | ||
| 80 | - groups, | ||
| 81 | - [output_mask[0], output_mask[1], False], | ||
| 82 | - ) | ||
| 83 | - return (grad_inp, grad_weight, grad_bias) | ||
| 84 | - | ||
| 85 | -def npu__softmax_backward_data( | ||
| 86 | - grad_output: torch.Tensor, | ||
| 87 | - output: torch.Tensor, | ||
| 88 | - dim: int, | ||
| 89 | - input_dtype: torch.dtype, | ||
| 90 | -) -> torch.Tensor: | ||
| 91 | - new_grad_output = grad_output * output | ||
| 92 | - sum_new_grad = torch.sum(new_grad_output, dim=dim, keepdim=True) | ||
| 93 | - grad_input = new_grad_output - output * sum_new_grad | ||
| 94 | - # grad_input = inductor_prims.fma(-output, sum_new_grad, new_grad_output) | ||
| 95 | - | ||
| 96 | - # CPU kernel doesn't respect input_dtype, but following check doesn't work for meta tensor | ||
| 97 | - # if grad_output.device == torch.device("cpu"): | ||
| 98 | - # return grad_input.contiguous() | ||
| 99 | - | ||
| 100 | - if grad_output.dtype != input_dtype: | ||
| 101 | - grad_input = grad_input.to(input_dtype) | ||
| 102 | - return grad_input.contiguous() | ||
| 103 | - | ||
| 104 | -def npu_rms_norm( | ||
| 105 | - x: torch.Tensor, | ||
| 106 | - weight: torch.Tensor, | ||
| 107 | - epsilon=1e-6 | ||
| 108 | -) -> torch.Tensor: | ||
| 109 | - dtype = x.dtype | ||
| 110 | - x = x.float() | ||
| 111 | - rsqrt = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + epsilon) | ||
| 112 | - output = (x * rsqrt * weight).to(dtype) | ||
| 113 | - return output, rsqrt | ||
| 114 | - | ||
| 115 | -def npu_rms_norm_backward(grad_output: torch.Tensor, | ||
| 116 | - x: torch.Tensor, | ||
| 117 | - weight: torch.Tensor, | ||
| 118 | - rsqrt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | ||
| 119 | - dx = (grad_output * weight - x * rsqrt * (grad_output * weight * x * rsqrt).mean(-1, keepdim=True)) * rsqrt | ||
| 120 | - dgamma = (grad_output * x * rsqrt).sum(0, keepdim=False) | ||
| 121 | - return dx, dgamma | ||
| 122 | - | ||
| 123 | -def npu_swiglu(x, dim=-1): | ||
| 124 | - x = torch.chunk(x, 2, dim=dim) | ||
| 125 | - return F.silu(x[0]) * x[1] | ||
| 126 | - | ||
| 127 | -def npu_swiglu_backward(grad_output, x, dim=-1): | ||
| 128 | - x0, x1 = torch.chunk(x, 2, dim=dim) | ||
| 129 | - | ||
| 130 | - # 计算 x0 的梯度 | ||
| 131 | - sigmoid_x0 = torch.sigmoid(x0) | ||
| 132 | - silu_grad = sigmoid_x0 * (1 + x0 * (1 - sigmoid_x0)) # SiLU 的导数 | ||
| 133 | - grad_x0 = grad_output * x1 * silu_grad | ||
| 134 | - | ||
| 135 | - # 计算 x1 的梯度 | ||
| 136 | - grad_x1 = grad_output * F.silu(x0) | ||
| 137 | - grad_x = torch.cat([grad_x0, grad_x1], dim=dim) | ||
| 138 | - return grad_x | ||
| 139 | - | ||
| 140 | -def _rotate_half(x: Tensor) -> Tensor: | ||
| 141 | - x1, x2 = torch.chunk(x, 2, dim=-1) | ||
| 142 | - return torch.cat((-x2, x1), dim=-1) | ||
| 143 | - | ||
| 144 | -def npu_rotary_mul(t, cos_, sin_): | ||
| 145 | - t = (t * cos_) + (_rotate_half(t) * sin_) | ||
| 146 | - return t | ||
| 147 | - | ||
| 148 | -def npu_rotary_mul_backward(grad_output, t, cos_, sin_): | ||
| 149 | - rotated_t = _rotate_half(t) | ||
| 150 | - grad_t = cos_ * grad_output | ||
| 151 | - grad_rotated_part = grad_output * sin_ | ||
| 152 | - a, b = torch.chunk(grad_rotated_part, 2, dim=-1) | ||
| 153 | - grad_rotated_t = torch.cat((b, -a), dim=-1) | ||
| 154 | - grad_t = grad_t + grad_rotated_t | ||
| 155 | - | ||
| 156 | - grad_cos = t * grad_output | ||
| 157 | - grad_sin = rotated_t * grad_output | ||
| 158 | - | ||
| 159 | - return grad_t, grad_cos, grad_sin | ||
| 160 | - | ||
| 161 | -def gelu(a, approximate: str = "none"): | ||
| 162 | - """ | ||
| 163 | - Reference implementation of torch.nn.functional.gelu | ||
| 164 | - """ | ||
| 165 | - M_SQRT2 = 1.41421356237309504880 | ||
| 166 | - M_2_SQRTPI = 1.12837916709551257390 | ||
| 167 | - kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 | ||
| 168 | - kKappa = 0.044715 | ||
| 169 | - a_cube = a * a * a | ||
| 170 | - inner = kBeta * (a + kKappa * a_cube) | ||
| 171 | - return 0.5 * a * (1 + torch.tanh(inner)) | ||
| 172 | - | ||
| 173 | -def gelu_backward(grad: Tensor, self: Tensor, approximate: str = "none"): | ||
| 174 | - M_SQRT2 = 1.41421356237309504880 | ||
| 175 | - M_SQRT1_2 = 0.70710678118654752440 | ||
| 176 | - M_2_SQRTPI = 1.12837916709551257390 | ||
| 177 | - kBeta = M_SQRT2 * M_2_SQRTPI * 0.5 | ||
| 178 | - kKappa = 0.044715 | ||
| 179 | - x_sq = self * self | ||
| 180 | - x_cube = x_sq * self | ||
| 181 | - inner = kBeta * (self + kKappa * x_cube) | ||
| 182 | - tanh_inner = torch.tanh(inner) | ||
| 183 | - | ||
| 184 | - left = 0.5 * self | ||
| 185 | - right = 1.0 + tanh_inner | ||
| 186 | - | ||
| 187 | - left_derivative = 0.5 * right | ||
| 188 | - | ||
| 189 | - tanh_derivative = (tanh_inner * tanh_inner) * -1.0 + 1.0 | ||
| 190 | - inner_derivative = kBeta * (1.0 + 3.0 * kKappa * x_sq) | ||
| 191 | - right_derivative = left * tanh_derivative * inner_derivative | ||
| 192 | - | ||
| 193 | - return grad * (left_derivative + right_derivative) | ||
| 194 | - | ||
| 195 | - | ||
| 196 | -def expm1(x): | ||
| 197 | - tensor = torch.exp(x) - torch.ones_like(x) | ||
| 198 | - return tensor | ||
| 199 | - | ||
| 200 | -inductor_decomp.register_decomposition(torch.ops.aten.expm1)(expm1) | ||
| 201 | -inductor_decomp.register_decomposition(torch.ops.aten.convolution_backward)(npu_convolution_backward) | ||
| 202 | -inductor_decomp.register_decomposition(torch.ops.aten._softmax_backward_data.default)(npu__softmax_backward_data) | ||
| 203 | -inductor_decomp.register_decomposition(torch.ops.aten.gelu.default)(gelu) | ||
| 204 | -inductor_decomp.register_decomposition(torch.ops.aten.gelu_backward.default)(gelu_backward) | ||
| 205 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_rms_norm.default)(npu_rms_norm) | ||
| 206 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_rms_norm_backward.default)(npu_rms_norm_backward) | ||
| 207 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_swiglu.default)(npu_swiglu) | ||
| 208 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_swiglu_backward.default)(npu_swiglu_backward) | ||
| 209 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_rotary_mul.default)(npu_rotary_mul) | ||
| 210 | -# inductor_decomp.register_decomposition(torch.ops.npu.npu_rotary_mul_backward.default)(npu_rotary_mul_backward) | ||
| @@ -104,8 +104,8 @@ class NewNpuInterface(NpuInterface): | |||
| 104 | register_interface_for_device("npu", NewNpuInterface) | 104 | register_interface_for_device("npu", NewNpuInterface) |
| 105 | 105 | ||
| 106 | ## npu patch | 106 | ## npu patch |
| 107 | -from torch_npu._inductor.decomposition import _register_npu_inductor_decompositions | 107 | +from torch_npu._inductor.decomposition import _register_mlir_dvm_decompositions |
| 108 | -_register_npu_inductor_decompositions(backend="mlir_dvm") | 108 | +_register_mlir_dvm_decompositions() |
| 109 | from torch._C import DispatchKey | 109 | from torch._C import DispatchKey |
| 110 | from torch._prims_common.wrappers import out_wrapper | 110 | from torch._prims_common.wrappers import out_wrapper |
| 111 | 111 | ||
| @@ -4,10 +4,9 @@ import torch | |||
| 4 | import torch._ops | 4 | import torch._ops |
| 5 | from torch import Tensor | 5 | from torch import Tensor |
| 6 | from torch._inductor import decomposition as inductor_decomp | 6 | from torch._inductor import decomposition as inductor_decomp |
| 7 | -from torch._inductor.decomposition import decompositions, pw_cast_for_opmath, register_decomposition | 7 | +from torch._inductor.decomposition import decompositions, register_decomposition |
| 8 | from torch._C import DispatchKey | 8 | from torch._C import DispatchKey |
| 9 | from torch._decomp import remove_decompositions | 9 | from torch._decomp import remove_decompositions |
| 10 | -from torch._prims_common.wrappers import out_wrapper | ||
| 11 | import torch.nn.functional as F | 10 | import torch.nn.functional as F |
| 12 | 11 | ||
| 13 | from .lowering_common import add_overload | 12 | from .lowering_common import add_overload |
| @@ -164,10 +163,12 @@ def _register_mlir_dvm_decompositions(): | |||
| 164 | output = (x * rsqrt * weight).to(dtype) | 163 | output = (x * rsqrt * weight).to(dtype) |
| 165 | return output, rsqrt | 164 | return output, rsqrt |
| 166 | 165 | ||
| 167 | - def npu_rms_norm_backward(grad_output: torch.Tensor, | 166 | + def npu_rms_norm_backward( |
| 168 | - x: torch.Tensor, | 167 | + grad_output: torch.Tensor, |
| 169 | - weight: torch.Tensor, | 168 | + x: torch.Tensor, |
| 170 | - rsqrt: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | 169 | + weight: torch.Tensor, |
| 170 | + rsqrt: torch.Tensor | ||
| 171 | + ) -> Tuple[torch.Tensor, torch.Tensor]: | ||
| 171 | dx = (grad_output * weight - x * rsqrt * (grad_output * weight * x * rsqrt).mean(-1, keepdim=True)) * rsqrt | 172 | dx = (grad_output * weight - x * rsqrt * (grad_output * weight * x * rsqrt).mean(-1, keepdim=True)) * rsqrt |
| 172 | dgamma = (grad_output * x * rsqrt).sum(0, keepdim=False) | 173 | dgamma = (grad_output * x * rsqrt).sum(0, keepdim=False) |
| 173 | return dx, dgamma | 174 | return dx, dgamma |
| @@ -254,10 +255,3 @@ def _register_mlir_dvm_decompositions(): | |||
| 254 | # register_decomposition(torch.ops.npu.npu_swiglu_backward.default)(npu_swiglu_backward) | 255 | # register_decomposition(torch.ops.npu.npu_swiglu_backward.default)(npu_swiglu_backward) |
| 255 | # register_decomposition(torch.ops.npu.npu_rotary_mul.default)(npu_rotary_mul) | 256 | # register_decomposition(torch.ops.npu.npu_rotary_mul.default)(npu_rotary_mul) |
| 256 | # register_decomposition(torch.ops.npu.npu_rotary_mul_backward.default)(npu_rotary_mul_backward) | 257 | # register_decomposition(torch.ops.npu.npu_rotary_mul_backward.default)(npu_rotary_mul_backward) |
| 257 | - | ||
| 258 | - | ||
| 259 | -def _register_npu_inductor_decompositions(backend=None): | ||
| 260 | - if backend == "triton": | ||
| 261 | - _register_triton_decompositions() | ||
| 262 | - elif backend in ["mlir_dvm"]: | ||
| 263 | - _register_mlir_dvm_decompositions() | ||