已合并
[inductor]Disable fully static Welford optimization in scenarios with inplace memory update #44708
rain-666创建于 6 天前
[inductor]Disable fully static Welford optimization in scenarios with inplace memory update #44708
已合并
共 2 个文件变更+203-8
| @@ -1,5 +1,6 @@ | |||
| 1 | import torch | 1 | import torch |
| 2 | import torch.nn.functional as F | 2 | import torch.nn.functional as F |
| 3 | +from torch._inductor import config | ||
| 3 | from torch._inductor.utils import run_and_get_code | 4 | from torch._inductor.utils import run_and_get_code |
| 4 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests | 5 | from torch.testing._internal.common_utils import run_tests, parametrize, instantiate_parametrized_tests |
| 5 | from testutils import TestUtils | 6 | from testutils import TestUtils |
| @@ -131,6 +132,109 @@ class TestVarMean(TestUtils): | |||
| 131 | npu_config.enable_welford = previous | 132 | npu_config.enable_welford = previous |
| 132 | torch._dynamo.reset() | 133 | torch._dynamo.reset() |
| 133 | 134 | ||
| 135 | + def test_welford_full_static_mutation_hazard(self): | ||
| 136 | + previous = npu_config.enable_welford | ||
| 137 | + npu_config.enable_welford = True | ||
| 138 | + torch._dynamo.reset() | ||
| 139 | + try: | ||
| 140 | + def randn_fp16(shape, scale=0.1): | ||
| 141 | + return ( | ||
| 142 | + torch.randn(shape, device="npu", dtype=torch.float16) * scale | ||
| 143 | + ) | ||
| 144 | + | ||
| 145 | + ids = torch.randint( | ||
| 146 | + 0, 10000, (200, 1), device="npu", dtype=torch.int64 | ||
| 147 | + ) | ||
| 148 | + features = randn_fp16((200, 16)) | ||
| 149 | + left_weight = randn_fp16((16, 32)) | ||
| 150 | + left_bias = randn_fp16((32,)) | ||
| 151 | + right_weight = randn_fp16((16, 32)) | ||
| 152 | + right_bias = randn_fp16((32,)) | ||
| 153 | + norm_weight = torch.ones((32,), device="npu", dtype=torch.float16) | ||
| 154 | + norm_weight = norm_weight + randn_fp16((32,)) | ||
| 155 | + norm_bias = randn_fp16((32,)) | ||
| 156 | + side_weight = randn_fp16((32, 16)) | ||
| 157 | + side_bias = randn_fp16((16,)) | ||
| 158 | + | ||
| 159 | + def layer_norm_with_reused_input( | ||
| 160 | + ids, | ||
| 161 | + features, | ||
| 162 | + left_weight, | ||
| 163 | + left_bias, | ||
| 164 | + right_weight, | ||
| 165 | + right_bias, | ||
| 166 | + norm_weight, | ||
| 167 | + norm_bias, | ||
| 168 | + side_weight, | ||
| 169 | + side_bias, | ||
| 170 | + ): | ||
| 171 | + left = torch.addmm(left_bias, features, left_weight) | ||
| 172 | + right = torch.addmm(right_bias, features, right_weight) | ||
| 173 | + condition = torch.logical_or(ids == 9998, ids == 3).repeat(1, 32) | ||
| 174 | + pre_norm = torch.where(condition, left, right) | ||
| 175 | + variance, mean = torch.var_mean( | ||
| 176 | + pre_norm, dim=-1, correction=0, keepdim=True | ||
| 177 | + ) | ||
| 178 | + normalized = ( | ||
| 179 | + (pre_norm - mean) * torch.rsqrt(variance + 1e-6) * norm_weight | ||
| 180 | + + norm_bias | ||
| 181 | + ) | ||
| 182 | + projected = torch.addmm(side_bias, pre_norm, side_weight) | ||
| 183 | + side_sum = pre_norm + normalized.sum(dim=-1, keepdim=True) | ||
| 184 | + return pre_norm, normalized, projected, side_sum | ||
| 185 | + | ||
| 186 | + args = ( | ||
| 187 | + ids, | ||
| 188 | + features, | ||
| 189 | + left_weight, | ||
| 190 | + left_bias, | ||
| 191 | + right_weight, | ||
| 192 | + right_bias, | ||
| 193 | + norm_weight, | ||
| 194 | + norm_bias, | ||
| 195 | + side_weight, | ||
| 196 | + side_bias, | ||
| 197 | + ) | ||
| 198 | + expected = layer_norm_with_reused_input(*args) | ||
| 199 | + compiled = torch.compile( | ||
| 200 | + layer_norm_with_reused_input, | ||
| 201 | + backend="inductor", | ||
| 202 | + dynamic=False, | ||
| 203 | + options={"unroll_reductions_threshold": 1}, | ||
| 204 | + ) | ||
| 205 | + with config.patch("triton.codegen_upcast_to_fp32", False): | ||
| 206 | + actual, codes = run_and_get_code(compiled, *args) | ||
| 207 | + | ||
| 208 | + for output in (*expected, *actual): | ||
| 209 | + self.assertTrue(torch.isfinite(output).all().item()) | ||
| 210 | + | ||
| 211 | + for expected_output, actual_output in zip(expected, actual): | ||
| 212 | + self.assertEqual( | ||
| 213 | + expected_output, actual_output, atol=1e-1, rtol=1e-1 | ||
| 214 | + ) | ||
| 215 | + | ||
| 216 | + actual_variance, actual_mean = torch.var_mean( | ||
| 217 | + actual[0], dim=-1, correction=0, keepdim=True | ||
| 218 | + ) | ||
| 219 | + normalized_from_actual_input = ( | ||
| 220 | + (actual[0] - actual_mean) | ||
| 221 | + * torch.rsqrt(actual_variance + 1e-6) | ||
| 222 | + * norm_weight | ||
| 223 | + + norm_bias | ||
| 224 | + ) | ||
| 225 | + self.assertEqual( | ||
| 226 | + normalized_from_actual_input, actual[1], atol=1e-2, rtol=1e-2 | ||
| 227 | + ) | ||
| 228 | + | ||
| 229 | + code = "\n".join(codes) | ||
| 230 | + if npu_config.is_ascend950: | ||
| 231 | + self.assertIn("npu_kernel_type': 'simd'", code) | ||
| 232 | + self.assertIn("mutated_arg_names': ['in_out_ptr", code) | ||
| 233 | + self.assertNotIn("'vectorized_welford_axis':", code) | ||
| 234 | + finally: | ||
| 235 | + npu_config.enable_welford = previous | ||
| 236 | + torch._dynamo.reset() | ||
| 237 | + | ||
| 134 | 238 | ||
| 135 | def test_welford_simd_low_precision_codegen(self, dtype): | 239 | def test_welford_simd_low_precision_codegen(self, dtype): |
| 136 | if not npu_config.is_ascend950: | 240 | if not npu_config.is_ascend950: |
| @@ -1738,6 +1738,99 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1738 | if node not in (EnableReduction, DisableReduction): | 1738 | if node not in (EnableReduction, DisableReduction): |
| 1739 | yield node | 1739 | yield node |
| 1740 | 1740 | ||
| 1741 | + def _full_static_welford_mutation_hazards(self) -> OrderedSet[str]: | ||
| 1742 | + """Find mutated buffers whose values are used across the reduction.""" | ||
| 1743 | + reduction_buffers = OrderedSet() | ||
| 1744 | + post_reduction_buffers = OrderedSet() | ||
| 1745 | + welford_source_buffers = OrderedSet() | ||
| 1746 | + welford_post_buffers = OrderedSet() | ||
| 1747 | + alias_groups: list[OrderedSet[str]] = [] | ||
| 1748 | + mutated_buffers = OrderedSet() | ||
| 1749 | + inside_reduction = True | ||
| 1750 | + seen_welford = False | ||
| 1751 | + | ||
| 1752 | + for node in self.node_schedule: | ||
| 1753 | + if node is DisableReduction: | ||
| 1754 | + inside_reduction = False | ||
| 1755 | + continue | ||
| 1756 | + if node is EnableReduction: | ||
| 1757 | + inside_reduction = True | ||
| 1758 | + continue | ||
| 1759 | + | ||
| 1760 | + used_buffers = node.used_buffer_names() | ||
| 1761 | + if inside_reduction: | ||
| 1762 | + reduction_buffers.update(used_buffers) | ||
| 1763 | + else: | ||
| 1764 | + post_reduction_buffers.update(used_buffers) | ||
| 1765 | + | ||
| 1766 | + reduction = getattr(node.node, "data", None) | ||
| 1767 | + if getattr(reduction, "reduction_type", None) == "welford_reduce": | ||
| 1768 | + welford_source_buffers.update(used_buffers) | ||
| 1769 | + seen_welford = True | ||
| 1770 | + elif seen_welford: | ||
| 1771 | + welford_post_buffers.update(used_buffers) | ||
| 1772 | + | ||
| 1773 | + for output in node.get_outputs(): | ||
| 1774 | + mutations = OrderedSet(output.get_mutations()) | ||
| 1775 | + group = OrderedSet( | ||
| 1776 | + [output.get_name(), *output.get_aliases(), *mutations] | ||
| 1777 | + ) | ||
| 1778 | + if len(group) > 1: | ||
| 1779 | + alias_groups.append(group) | ||
| 1780 | + if mutations: | ||
| 1781 | + mutated_buffers.update(group) | ||
| 1782 | + | ||
| 1783 | + for output_name, input_name in self.inplace_update_buffers.items(): | ||
| 1784 | + group = OrderedSet([output_name, input_name]) | ||
| 1785 | + alias_groups.append(group) | ||
| 1786 | + mutated_buffers.update(group) | ||
| 1787 | + mutated_buffers.update(self.mutations) | ||
| 1788 | + | ||
| 1789 | + # Scheduler names may refer to either a logical output or the physical | ||
| 1790 | + # input it aliases. Expand both sets to the same transitive alias closure. | ||
| 1791 | + def expand_aliases(buffer_names: Iterable[str]) -> OrderedSet[str]: | ||
| 1792 | + expanded = OrderedSet(buffer_names) | ||
| 1793 | + changed = True | ||
| 1794 | + while changed: | ||
| 1795 | + changed = False | ||
| 1796 | + for group in alias_groups: | ||
| 1797 | + if expanded & group and group - expanded: | ||
| 1798 | + expanded.update(group) | ||
| 1799 | + changed = True | ||
| 1800 | + return expanded | ||
| 1801 | + | ||
| 1802 | + reduction_buffers = expand_aliases(reduction_buffers) | ||
| 1803 | + post_reduction_buffers = expand_aliases(post_reduction_buffers) | ||
| 1804 | + welford_source_buffers = expand_aliases(welford_source_buffers) | ||
| 1805 | + welford_post_buffers = expand_aliases(welford_post_buffers) | ||
| 1806 | + cross_reduction_buffers = ( | ||
| 1807 | + reduction_buffers & post_reduction_buffers | ||
| 1808 | + ) | (welford_source_buffers & welford_post_buffers) | ||
| 1809 | + mutated_buffers = expand_aliases(mutated_buffers) | ||
| 1810 | + hazards = cross_reduction_buffers & mutated_buffers | ||
| 1811 | + log.debug( | ||
| 1812 | + "full-static Welford mutation analysis: reduction=%s, post=%s, " | ||
| 1813 | + "welford_source=%s, welford_post=%s, cross=%s, mutated=%s, " | ||
| 1814 | + "hazards=%s", | ||
| 1815 | + sorted(reduction_buffers), | ||
| 1816 | + sorted(post_reduction_buffers), | ||
| 1817 | + sorted(welford_source_buffers), | ||
| 1818 | + sorted(welford_post_buffers), | ||
| 1819 | + sorted(cross_reduction_buffers), | ||
| 1820 | + sorted(mutated_buffers), | ||
| 1821 | + sorted(hazards), | ||
| 1822 | + ) | ||
| 1823 | + return hazards | ||
| 1824 | + | ||
| 1825 | + def finalize_indexing(self, indices: Sequence[sympy.Expr]) -> None: | ||
| 1826 | + super().finalize_indexing(indices) | ||
| 1827 | + if self.full_static_welford_reduction: | ||
| 1828 | + # The first scheduler pass has now finalized buffer reuse and | ||
| 1829 | + # in-place updates, while load codegen has not started yet. | ||
| 1830 | + self.full_static_welford_reduction = not ( | ||
| 1831 | + self._full_static_welford_mutation_hazards() | ||
| 1832 | + ) | ||
| 1833 | + | ||
| 1741 | def _iter_store_indices(self): | 1834 | def _iter_store_indices(self): |
| 1742 | """ | 1835 | """ |
| 1743 | Yield Store key/index pairs from scheduled nodes. | 1836 | Yield Store key/index pairs from scheduled nodes. |
| @@ -1887,14 +1980,12 @@ class NPUIndexTritonKernel(TritonKernel): | |||
| 1887 | axis = max(candidates, key=lambda candidate: candidate.sorted_order) | 1980 | axis = max(candidates, key=lambda candidate: candidate.sorted_order) |
| 1888 | axis.is_vectorized_split = True | 1981 | axis.is_vectorized_split = True |
| 1889 | self.vectorized_welford_axis = axis | 1982 | self.vectorized_welford_axis = axis |
| 1890 | - # A persistent SIMD tile covers every static reduction element. With | 1983 | + # A persistent SIMD tile covers every static reduction element. Loads |
| 1891 | - # no mutation, its loads can safely remain live for post-reduction use. | 1984 | + # may remain live unless a buffer used on both sides of the reduction |
| 1892 | - self.full_static_welford_reduction = ( | 1985 | + # is mutated by this kernel. |
| 1893 | - not self.mutations | 1986 | + self.full_static_welford_reduction = all( |
| 1894 | - and all( | 1987 | + int(self.range_tree_nodes[reduction_axis].length) > 0 |
| 1895 | - int(self.range_tree_nodes[reduction_axis].length) > 0 | 1988 | + for reduction_axis in self.reduction_axis_list() |
| 1896 | - for reduction_axis in self.reduction_axis_list() | ||
| 1897 | - ) | ||
| 1898 | ) | 1989 | ) |
| 1899 | 1990 | ||
| 1900 | def _static_welford_reduction_numel(self): | 1991 | def _static_welford_reduction_numel(self): |