已开启
fix(inductor): fix triton_experimental UT failures #44742
liuyutong创建于 28 天前
fix(inductor): fix triton_experimental UT failures #44742
已开启
共 3 个文件变更+54-0
| @@ -751,8 +751,16 @@ def _npu_print_ToFloat(self, expr): | |||
| 751 | s = self.parenthesize(expr.args[0], PRECEDENCE["Atom"] - 0.5) | 751 | s = self.parenthesize(expr.args[0], PRECEDENCE["Atom"] - 0.5) |
| 752 | return f"{s}.to(tl.float32)" | 752 | return f"{s}.to(tl.float32)" |
| 753 | 753 | ||
| 754 | +def _npu_print_RoundToInt(self, expr): | ||
| 755 | + assert len(expr.args) == 1 | ||
| 756 | + return ( | ||
| 757 | + f"libdevice.nearbyint({self._print(expr.args[0])})" | ||
| 758 | + f".to({V.kernel.index_dtype})" | ||
| 759 | + ) | ||
| 760 | + | ||
| 754 | _TritonPrinter._print_Float = _npu_print_Float | 761 | _TritonPrinter._print_Float = _npu_print_Float |
| 755 | _TritonPrinter._print_ToFloat = _npu_print_ToFloat | 762 | _TritonPrinter._print_ToFloat = _npu_print_ToFloat |
| 763 | +_TritonPrinter._print_RoundToInt = _npu_print_RoundToInt | ||
| 756 | 764 | ||
| 757 | 765 | ||
| 758 | # Tensor-dimension symbol kinds: any tensor actually indexed by a running kernel | 766 | # Tensor-dimension symbol kinds: any tensor actually indexed by a running kernel |
| @@ -131,3 +131,16 @@ KEEP_UPSTREAM_LOWERING = [control_deps] + [ | |||
| 131 | ) | 131 | ) |
| 132 | if hasattr(aten, _name) | 132 | if hasattr(aten, _name) |
| 133 | ] | 133 | ] |
| 134 | + | ||
| 135 | +# torch.ops._inductor_test.realize is defined by torch._inductor.test_operators | ||
| 136 | +# (imported by inductor scheduler tests such as test_forced_buffer_realize). Its | ||
| 137 | +# upstream lowering is pure IR bookkeeping (x.realize() + clone) and emits no | ||
| 138 | +# device kernel, so it must NOT be clobbered into a fallback: the CPU fallback | ||
| 139 | +# breaks forced buffer realization and inflates ir_nodes_pre_fusion. Importing | ||
| 140 | +# test_operators just defines the _inductor_test torch.library -- no side effect. | ||
| 141 | +try: | ||
| 142 | + from torch._inductor import test_operators as _test_operators # noqa: F401 | ||
| 143 | + | ||
| 144 | + KEEP_UPSTREAM_LOWERING.append(torch.ops._inductor_test.realize) | ||
| 145 | +except ImportError: | ||
| 146 | + pass | ||
| @@ -84,6 +84,34 @@ def _disable_pad_mm_pass(): | |||
| 84 | inductor_config.shape_padding = False | 84 | inductor_config.shape_padding = False |
| 85 | 85 | ||
| 86 | 86 | ||
| 87 | +def _register_l1_loss_backward_decomposition(): | ||
| 88 | + """Register an inductor decomposition for ``npu.l1_loss_backward``. | ||
| 89 | + | ||
| 90 | + torch_npu's l1_loss autograd formula emits ``torch.ops.npu.l1_loss_backward`` | ||
| 91 | + into the backward FX graph (torch core 2.13 has no aten.l1_loss_backward op to | ||
| 92 | + decompose), but the triton_experimental backend has neither a lowering nor a | ||
| 93 | + decomposition for it -> MissingOperatorWithoutDecomp at backward compile time | ||
| 94 | + (test_inductor_sequence_nr). Decompose it into the exact math of the NPU | ||
| 95 | + kernel (grad * sign(self - target), / numel for Mean reduction) so Inductor | ||
| 96 | + fuses it into the surrounding backward pointwise chain instead of erroring | ||
| 97 | + out. Mirrors torch._decomp's own mse_loss_backward decomposition.""" | ||
| 98 | + import torch | ||
| 99 | + from torch._inductor.decomposition import decompositions as _ind_decomps | ||
| 100 | + | ||
| 101 | + def l1_loss_backward(grad_output, self, target, reduction): | ||
| 102 | + diff = self - target | ||
| 103 | + # sign(diff) = (diff > 0) - (diff < 0): avoids aten.sign, which is not in | ||
| 104 | + # this backend's GENERATE_LIST (would fall back to aclnn and break the | ||
| 105 | + # fusion). gt/lt/convert_element_type/sub are all generatable pointwise. | ||
| 106 | + sign = (diff > 0).to(diff.dtype) - (diff < 0).to(diff.dtype) | ||
| 107 | + grad_input = sign * grad_output | ||
| 108 | + if reduction == 1: # Reduction::Mean | ||
| 109 | + return grad_input / self.numel() | ||
| 110 | + return grad_input | ||
| 111 | + | ||
| 112 | + _ind_decomps[torch.ops.npu.l1_loss_backward.default] = l1_loss_backward | ||
| 113 | + | ||
| 114 | + | ||
| 87 | def apply_npu_overrides(): | 115 | def apply_npu_overrides(): |
| 88 | # config-flag flips + FX passes that don't touch the decomposition table. | 116 | # config-flag flips + FX passes that don't touch the decomposition table. |
| 89 | _override_disable_pointwise_autotuning() | 117 | _override_disable_pointwise_autotuning() |
| @@ -99,6 +127,11 @@ def apply_npu_overrides(): | |||
| 99 | # (_register_triton_experimental_decompositions) and are invoked directly by the | 127 | # (_register_triton_experimental_decompositions) and are invoked directly by the |
| 100 | # loader _load_triton_experimental_backend before _activate(), mirroring the | 128 | # loader _load_triton_experimental_backend before _activate(), mirroring the |
| 101 | # other backends' _register_*_decompositions entry points. | 129 | # other backends' _register_*_decompositions entry points. |
| 130 | + # | ||
| 131 | + # Exception: the npu.l1_loss_backward decomposition is a backend-specific gap | ||
| 132 | + # (no NPU lowering exists for this op), so it is registered here to keep the | ||
| 133 | + # change physically inside triton_experimental. | ||
| 134 | + _register_l1_loss_backward_decomposition() | ||
| 102 | 135 | ||
| 103 | # Triton codegen / scheduler monkeypatches (codegen/triton.py). | 136 | # Triton codegen / scheduler monkeypatches (codegen/triton.py). |
| 104 | apply_npu_codegen_patches() | 137 | apply_npu_codegen_patches() |