已合并
fix(inductor): deliver auto_blockify_size via the Config.extra_options slot #44846
AllenGuan创建于 8月18日
fix(inductor): deliver auto_blockify_size via the Config.extra_options slot #44846
已合并
共 2 个文件变更+186-17
| @@ -33,6 +33,8 @@ from testutils import TestUtils | |||
| 33 | import torch_npu._inductor.triton_experimental.npu_triton_heuristics as _heur | 33 | import torch_npu._inductor.triton_experimental.npu_triton_heuristics as _heur |
| 34 | from torch_npu._inductor.triton_experimental.npu_triton_heuristics import ( | 34 | from torch_npu._inductor.triton_experimental.npu_triton_heuristics import ( |
| 35 | _pw1d_formula_configs, _red_formula_configs, | 35 | _pw1d_formula_configs, _red_formula_configs, |
| 36 | + npu_triton_config, | ||
| 37 | + _npu_config_key, _npu_unique_configs, _npu_hash_configs, | ||
| 36 | _NPU_PTR_ELEM_BYTES, _NPU_UB_CAPACITY_BYTES, _NPU_UB_OVERHEAD_FACTOR) | 38 | _NPU_PTR_ELEM_BYTES, _NPU_UB_CAPACITY_BYTES, _NPU_UB_OVERHEAD_FACTOR) |
| 37 | from torch_npu._inductor.triton_experimental import config as ncfg | 39 | from torch_npu._inductor.triton_experimental import config as ncfg |
| 38 | 40 | ||
| @@ -43,11 +45,21 @@ HOLE_R = (512, 1024, 2048, 4096) | |||
| 43 | 45 | ||
| 44 | # L0 helpers: call the pure generators directly with synthesized dicts. | 46 | # L0 helpers: call the pure generators directly with synthesized dicts. |
| 45 | def _norm(c): | 47 | def _norm(c): |
| 46 | - """Config object OR serialized {kwargs, num_warps, num_stages} dict -> tuple.""" | 48 | + """Config object OR serialized {kwargs, num_warps, num_stages} dict -> tuple. |
| 47 | - kw = c.kwargs if hasattr(c, "kwargs") else c["kwargs"] | 49 | + |
| 50 | + Mirrors the backend identity key (_npu_config_key): kwargs + num_warps + | ||
| 51 | + num_stages + the backend extra_options slot, so auto_blockify variants of | ||
| 52 | + one tile shape count as distinct configs.""" | ||
| 53 | + is_obj = hasattr(c, "kwargs") | ||
| 54 | + kw = c.kwargs if is_obj else c["kwargs"] | ||
| 48 | nw = c.num_warps if hasattr(c, "num_warps") else c["num_warps"] | 55 | nw = c.num_warps if hasattr(c, "num_warps") else c["num_warps"] |
| 49 | ns = c.num_stages if hasattr(c, "num_stages") else c["num_stages"] | 56 | ns = c.num_stages if hasattr(c, "num_stages") else c["num_stages"] |
| 50 | - return (tuple(sorted(kw.items())), nw, ns) | 57 | + extra = (getattr(c, "extra_options", None) if is_obj else c.get("extra_options")) or {} |
| 58 | + return ( | ||
| 59 | + tuple(sorted(kw.items())) | ||
| 60 | + + (("num_warps", nw), ("num_stages", ns)) | ||
| 61 | + + tuple(sorted(extra.items())) | ||
| 62 | + ) | ||
| 51 | 63 | ||
| 52 | 64 | ||
| 53 | def _norm_list(cfgs): | 65 | def _norm_list(cfgs): |
| @@ -147,15 +159,20 @@ class TestAutotuneFormula(TestCase): | |||
| 147 | 159 | ||
| 148 | def test_huge_kernel_auto_blockify_on(self): | 160 | def test_huge_kernel_auto_blockify_on(self): |
| 149 | # all_blocks_parallel=True (default): huge numel appends auto_blockify_size | 161 | # all_blocks_parallel=True (default): huge numel appends auto_blockify_size |
| 150 | - # {2,4,8} on the cap XBLOCK. | 162 | + # {2,4,8} on the cap XBLOCK. The value rides the backend Config.extra_options |
| 163 | + # slot (NOT cfg.kwargs, whose keys must all be kernel-signature constexprs). | ||
| 151 | self.assertTrue(ncfg.all_blocks_parallel) | 164 | self.assertTrue(ncfg.all_blocks_parallel) |
| 152 | cfgs = pw1d(2_000_000_000, "fp32", 1) | 165 | cfgs = pw1d(2_000_000_000, "fp32", 1) |
| 153 | - blockify = sorted(c.kwargs["auto_blockify_size"] | 166 | + blockify = sorted(c.extra_options["auto_blockify_size"] |
| 154 | - for c in cfgs if "auto_blockify_size" in c.kwargs) | 167 | + for c in cfgs |
| 168 | + if getattr(c, "extra_options", None)) | ||
| 155 | self.assertEqual(blockify, [2, 4, 8]) | 169 | self.assertEqual(blockify, [2, 4, 8]) |
| 156 | cap_xb = max(c.kwargs["XBLOCK"] for c in cfgs) | 170 | cap_xb = max(c.kwargs["XBLOCK"] for c in cfgs) |
| 157 | self.assertTrue(all(c.kwargs["XBLOCK"] == cap_xb | 171 | self.assertTrue(all(c.kwargs["XBLOCK"] == cap_xb |
| 158 | - for c in cfgs if "auto_blockify_size" in c.kwargs)) | 172 | + for c in cfgs if getattr(c, "extra_options", None))) |
| 173 | + # The slot separation is the contract: no backend option may leak into | ||
| 174 | + # kwargs (constants channel). | ||
| 175 | + self.assertFalse(any("auto_blockify_size" in c.kwargs for c in cfgs)) | ||
| 159 | 176 | ||
| 160 | def test_numel_one(self): | 177 | def test_numel_one(self): |
| 161 | # numel < align: bracket floors at 1, cap skipped -> single [1]. | 178 | # numel < align: bracket floors at 1, cap skipped -> single [1]. |
| @@ -341,6 +358,106 @@ class TestAutotuneFormula(TestCase): | |||
| 341 | instantiate_parametrized_tests(TestAutotuneFormula) | 358 | instantiate_parametrized_tests(TestAutotuneFormula) |
| 342 | 359 | ||
| 343 | 360 | ||
| 361 | +# L0: auto_blockify slot + gate contract. auto_blockify_size is a triton-ascend | ||
| 362 | +# compile option (NPUOptions field), NOT a kernel constexpr: it must ride the | ||
| 363 | +# upstream third-party backend slot Config.extra_options (AutotuneCache.save | ||
| 364 | +# serializes it at autotune_cache.py:325, read_best restores it at :674/:694/:701) | ||
| 365 | +# and reach triton.compile via the options dict. If it ever lands in | ||
| 366 | +# Config.kwargs, the constants merge in _precompile_config feeds it to | ||
| 367 | +# ast_to_ttir and every compile of the config dies with | ||
| 368 | +# "ValueError: 'auto_blockify_size' is not in list". The gate semantics are the | ||
| 369 | +# pre-slot-era ones, unchanged: all_blocks_parallel=True AND | ||
| 370 | +# ceildiv(numel, cap_XBLOCK) > 65535 -> append {2,4,8} on the cap tile. | ||
| 371 | +class TestAutoBlockify(TestCase): | ||
| 372 | + | ||
| 373 | + HUGE = 2_000_000_000 # fp32/num_load=1: cap 12288, grid 162760 >> 65535 | ||
| 374 | + | ||
| 375 | + def _ab_cfgs(self, numel, dtype="fp32", num_load=1): | ||
| 376 | + cfgs = pw1d(numel, dtype, num_load) | ||
| 377 | + with_ab = [c for c in cfgs if getattr(c, "extra_options", None)] | ||
| 378 | + without_ab = [c for c in cfgs if getattr(c, "extra_options", None) is None] | ||
| 379 | + return cfgs, with_ab, without_ab | ||
| 380 | + | ||
| 381 | + def test_slot_placement_when_on(self): | ||
| 382 | + # ON case: {2,4,8} ride extra_options, exactly one key per config, | ||
| 383 | + # kwargs stay pure constexprs, and only the cap tile carries the slot. | ||
| 384 | + cfgs, with_ab, without_ab = self._ab_cfgs(self.HUGE) | ||
| 385 | + self.assertEqual( | ||
| 386 | + sorted(c.extra_options["auto_blockify_size"] for c in with_ab), [2, 4, 8]) | ||
| 387 | + for c in with_ab: | ||
| 388 | + self.assertEqual(set(c.extra_options), {"auto_blockify_size"}) | ||
| 389 | + self.assertFalse(any("auto_blockify_size" in c.kwargs for c in cfgs)) | ||
| 390 | + self.assertEqual(len(with_ab) + len(without_ab), len(cfgs)) | ||
| 391 | + cap = max(c.kwargs["XBLOCK"] for c in cfgs) | ||
| 392 | + self.assertTrue(all(c.kwargs["XBLOCK"] == cap for c in with_ab)) | ||
| 393 | + | ||
| 394 | + def test_gate_off_below_grid_threshold(self): | ||
| 395 | + # ceildiv(numel, cap) == 65535 exactly: OFF (the gate is strictly >). | ||
| 396 | + align = _pw1d_align("fp32") | ||
| 397 | + hi = _pw1d_hi(self.HUGE, "fp32", 1) | ||
| 398 | + cap = (hi // align) * align | ||
| 399 | + _, with_ab, _ = self._ab_cfgs(65535 * cap) | ||
| 400 | + self.assertEqual(with_ab, []) | ||
| 401 | + | ||
| 402 | + def test_gate_on_just_above_grid_threshold(self): | ||
| 403 | + # One block over the 65535 coreDim limit: ON -- the boundary the | ||
| 404 | + # mobilevit_s bs=128 crash kernel sits exactly on (grid 65536). | ||
| 405 | + align = _pw1d_align("fp32") | ||
| 406 | + hi = _pw1d_hi(self.HUGE, "fp32", 1) | ||
| 407 | + cap = (hi // align) * align | ||
| 408 | + _, with_ab, _ = self._ab_cfgs(65536 * cap) | ||
| 409 | + self.assertEqual( | ||
| 410 | + sorted(c.extra_options["auto_blockify_size"] for c in with_ab), [2, 4, 8]) | ||
| 411 | + | ||
| 412 | + def test_gate_small_kernel_off(self): | ||
| 413 | + # grid far below 65535: never generates backend-option candidates. | ||
| 414 | + _, with_ab, _ = self._ab_cfgs(1_000_000, "fp32", 2) | ||
| 415 | + self.assertEqual(with_ab, []) | ||
| 416 | + | ||
| 417 | + def test_gate_all_blocks_parallel_off(self): | ||
| 418 | + # The feature switch (config all_blocks_parallel, default True) kills | ||
| 419 | + # candidate generation even for a huge grid. | ||
| 420 | + saved = ncfg.all_blocks_parallel | ||
| 421 | + try: | ||
| 422 | + ncfg.all_blocks_parallel = False | ||
| 423 | + _, with_ab, _ = self._ab_cfgs(self.HUGE) | ||
| 424 | + self.assertEqual(with_ab, []) | ||
| 425 | + finally: | ||
| 426 | + ncfg.all_blocks_parallel = saved | ||
| 427 | + | ||
| 428 | + def test_identity_distinguishes_backend_options(self): | ||
| 429 | + # Upstream's dedup/hash key (kwargs+num_warps+num_stages, see | ||
| 430 | + # runtime_utils.triton_config_to_hashable / triton_heuristics.hash_configs) | ||
| 431 | + # cannot see extra_options: the {2,4,8} candidates share one tile shape | ||
| 432 | + # and would collapse into one. _npu_config_key/_npu_unique_configs/ | ||
| 433 | + # _npu_hash_configs (used by cached_autotune) must keep them distinct. | ||
| 434 | + size_hints = {"x": self.HUGE} | ||
| 435 | + plain = npu_triton_config(size_hints, 4096) | ||
| 436 | + ab = [npu_triton_config(size_hints, 4096, auto_blockify_size=v) | ||
| 437 | + for v in (2, 4, 8)] | ||
| 438 | + keys = [_npu_config_key(c) for c in [plain] + ab] | ||
| 439 | + self.assertEqual(len(keys), len(set(keys))) | ||
| 440 | + self.assertEqual(len(_npu_unique_configs([plain] + ab)), 4) | ||
| 441 | + self.assertEqual(len(_npu_unique_configs([plain, plain])), 1) | ||
| 442 | + self.assertNotEqual(_npu_hash_configs([plain]), _npu_hash_configs([ab[0]])) | ||
| 443 | + | ||
| 444 | + def test_norm_roundtrip_carries_extra_options(self): | ||
| 445 | + # The serialized config dict (the shape AutotuneCache.save writes / | ||
| 446 | + # read_best restores) and the live Config must normalize to the same | ||
| 447 | + # identity -- losing extra_options must change the identity. | ||
| 448 | + size_hints = {"x": self.HUGE} | ||
| 449 | + cfg = npu_triton_config(size_hints, 4096, auto_blockify_size=4) | ||
| 450 | + d = { | ||
| 451 | + "kwargs": dict(cfg.kwargs), | ||
| 452 | + "num_warps": cfg.num_warps, | ||
| 453 | + "num_stages": cfg.num_stages, | ||
| 454 | + "extra_options": dict(cfg.extra_options), | ||
| 455 | + } | ||
| 456 | + self.assertEqual(_norm(cfg), _norm(d)) | ||
| 457 | + d.pop("extra_options") | ||
| 458 | + self.assertNotEqual(_norm(cfg), _norm(d)) | ||
| 459 | + | ||
| 460 | + | ||
| 344 | # L0: broad-grid minimal contract for the REDUCTION generator only (varied dtype / | 461 | # L0: broad-grid minimal contract for the REDUCTION generator only (varied dtype / |
| 345 | # num_load / num_reduction that the fixed-param invariant tests above do not cover) | 462 | # num_load / num_reduction that the fixed-param invariant tests above do not cover) |
| 346 | # -- non-empty, well-formed, positive, in-bounds, num_stages==1. | 463 | # -- non-empty, well-formed, positive, in-bounds, num_stages==1. |
| @@ -22,6 +22,7 @@ | |||
| 22 | 22 | ||
| 23 | import copy | 23 | import copy |
| 24 | import functools | 24 | import functools |
| 25 | +import hashlib | ||
| 25 | import logging | 26 | import logging |
| 26 | import math | 27 | import math |
| 27 | import operator | 28 | import operator |
| @@ -48,9 +49,7 @@ from torch._inductor.runtime.triton_heuristics import ( | |||
| 48 | CachingAutotuner, | 49 | CachingAutotuner, |
| 49 | TritonCompileResult, | 50 | TritonCompileResult, |
| 50 | autotune_hints_to_configs, | 51 | autotune_hints_to_configs, |
| 51 | - unique_configs, | ||
| 52 | triton_config_reduction, | 52 | triton_config_reduction, |
| 53 | - hash_configs, | ||
| 54 | get_first_attr, | 53 | get_first_attr, |
| 55 | ) | 54 | ) |
| 56 | from torch._inductor.runtime.autotune_cache import AutotuneCache | 55 | from torch._inductor.runtime.autotune_cache import AutotuneCache |
| @@ -854,8 +853,9 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 854 | config fails to compile (typically UB overflow on giant fused kernels). | 853 | config fails to compile (typically UB overflow on giant fused kernels). |
| 855 | 854 | ||
| 856 | We rebuild Configs from one of the failed configs as a template — that | 855 | We rebuild Configs from one of the failed configs as a template — that |
| 857 | - preserves auxiliary kwargs like ``auto_blockify_size`` | 856 | + preserves backend compile options riding the Config.extra_options slot |
| 858 | - — but override the *_BLOCK kwargs with powers of two in [1, 256]. | 857 | + (``auto_blockify_size``) — but override the *_BLOCK kwargs with powers |
| 858 | + of two in [1, 256]. | ||
| 859 | Reduction R*_BLOCK kwargs are left untouched (UB overflow is dominated | 859 | Reduction R*_BLOCK kwargs are left untouched (UB overflow is dominated |
| 860 | by the pointwise tile; reduction tiles are sized differently). | 860 | by the pointwise tile; reduction tiles are sized differently). |
| 861 | """ | 861 | """ |
| @@ -898,11 +898,15 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 898 | if key in seen: | 898 | if key in seen: |
| 899 | continue | 899 | continue |
| 900 | seen.add(key) | 900 | seen.add(key) |
| 901 | - candidates.append(Config( | 901 | + fallback_cfg = Config( |
| 902 | cfg_kwargs, | 902 | cfg_kwargs, |
| 903 | num_warps=getattr(template, "num_warps", 8), | 903 | num_warps=getattr(template, "num_warps", 8), |
| 904 | num_stages=getattr(template, "num_stages", 1), | 904 | num_stages=getattr(template, "num_stages", 1), |
| 905 | - )) | 905 | + ) |
| 906 | + template_extra = getattr(template, "extra_options", None) | ||
| 907 | + if template_extra: | ||
| 908 | + fallback_cfg.extra_options = template_extra | ||
| 909 | + candidates.append(fallback_cfg) | ||
| 906 | 910 | ||
| 907 | results = [] | 911 | results = [] |
| 908 | last_exc = None | 912 | last_exc = None |
| @@ -986,6 +990,11 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 986 | "target": target, | 990 | "target": target, |
| 987 | "options": options, | 991 | "options": options, |
| 988 | } | 992 | } |
| 993 | + # Backend compile options (auto_blockify_size & co.) ride the upstream | ||
| 994 | + # third-party Config.extra_options slot; forward them through the | ||
| 995 | + # options channel. NPUOptions.parse_options maps the keys it knows onto | ||
| 996 | + # its fields and silently drops the rest. | ||
| 997 | + options.update(getattr(cfg, "extra_options", None) or {}) | ||
| 989 | 998 | ||
| 990 | try: | 999 | try: |
| 991 | binary = triton.compile(*compile_args, **compile_kwargs) | 1000 | binary = triton.compile(*compile_args, **compile_kwargs) |
| @@ -1295,6 +1304,41 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 1295 | return times[len(times) // 2] | 1304 | return times[len(times) // 2] |
| 1296 | 1305 | ||
| 1297 | 1306 | ||
| 1307 | +def _npu_config_key(cfg: Config) -> tuple: | ||
| 1308 | + """Identity key for config dedup and autotune-cache hashing: upstream's | ||
| 1309 | + triton_config_to_hashable/hash_configs key ONLY (kwargs, num_warps, | ||
| 1310 | + num_stages), which would collapse candidates that differ solely in the | ||
| 1311 | + backend extra_options slot (auto_blockify_size {2,4,8} all share the same | ||
| 1312 | + tile kwargs). Extend the key so each backend-option candidate stays a | ||
| 1313 | + distinct autotune candidate; save and read_best both go through this file, | ||
| 1314 | + so the hash stays self-consistent end to end.""" | ||
| 1315 | + items = sorted(cfg.kwargs.items()) | ||
| 1316 | + items.append(("num_warps", cfg.num_warps)) | ||
| 1317 | + items.append(("num_stages", cfg.num_stages)) | ||
| 1318 | + extra = getattr(cfg, "extra_options", None) | ||
| 1319 | + if extra: | ||
| 1320 | + items.extend(sorted(extra.items())) | ||
| 1321 | + return tuple(items) | ||
| 1322 | + | ||
| 1323 | + | ||
| 1324 | +def _npu_unique_configs(configs): | ||
| 1325 | + seen = set() | ||
| 1326 | + out = [] | ||
| 1327 | + for cfg in configs: | ||
| 1328 | + key = _npu_config_key(cfg) | ||
| 1329 | + if key in seen: | ||
| 1330 | + continue | ||
| 1331 | + seen.add(key) | ||
| 1332 | + out.append(cfg) | ||
| 1333 | + return out | ||
| 1334 | + | ||
| 1335 | + | ||
| 1336 | +def _npu_hash_configs(configs): | ||
| 1337 | + hasher = hashlib.sha256() | ||
| 1338 | + for cfg in configs: | ||
| 1339 | + hasher.update(f"{_npu_config_key(cfg)}\n".encode()) | ||
| 1340 | + return hasher.hexdigest() | ||
| 1341 | + | ||
| 1298 | 1342 | ||
| 1299 | def cached_autotune( | 1343 | def cached_autotune( |
| 1300 | size_hints: Optional[List[int]], | 1344 | size_hints: Optional[List[int]], |
| @@ -1309,7 +1353,7 @@ def cached_autotune( | |||
| 1309 | Override of cached_autotune to redirect to NPU autotuner subclass. | 1353 | Override of cached_autotune to redirect to NPU autotuner subclass. |
| 1310 | In 2.7.1, uses AutotuneCache instead of manual file caching. | 1354 | In 2.7.1, uses AutotuneCache instead of manual file caching. |
| 1311 | """ | 1355 | """ |
| 1312 | - configs = unique_configs(configs) | 1356 | + configs = _npu_unique_configs(configs) |
| 1313 | if len(configs) != 1 and not filename: | 1357 | if len(configs) != 1 and not filename: |
| 1314 | raise ValueError("[triton_experimental] cached_autotune requires a filename when multiple configs are given") | 1358 | raise ValueError("[triton_experimental] cached_autotune requires a filename when multiple configs are given") |
| 1315 | inductor_meta = {} if inductor_meta is None else inductor_meta | 1359 | inductor_meta = {} if inductor_meta is None else inductor_meta |
| @@ -1323,7 +1367,7 @@ def cached_autotune( | |||
| 1323 | and (len(configs) > 1 or inductor_meta.get("coordinate_descent_tuning")) | 1367 | and (len(configs) > 1 or inductor_meta.get("coordinate_descent_tuning")) |
| 1324 | and not os.environ.get("TRITON_INTERPRET", "0") == "1" | 1368 | and not os.environ.get("TRITON_INTERPRET", "0") == "1" |
| 1325 | ): | 1369 | ): |
| 1326 | - configs_hash = hash_configs(configs) | 1370 | + configs_hash = _npu_hash_configs(configs) |
| 1327 | autotune_cache = AutotuneCache.create(inductor_meta, filename, configs_hash) | 1371 | autotune_cache = AutotuneCache.create(inductor_meta, filename, configs_hash) |
| 1328 | if autotune_cache and (best_config := autotune_cache.read_best(inductor_meta, configs)): | 1372 | if autotune_cache and (best_config := autotune_cache.read_best(inductor_meta, configs)): |
| 1329 | configs = [best_config] | 1373 | configs = [best_config] |
| @@ -1381,9 +1425,17 @@ def npu_triton_config( | |||
| 1381 | cfg["YBLOCK"] = min(y, size_hints["y"]) | 1425 | cfg["YBLOCK"] = min(y, size_hints["y"]) |
| 1382 | if z: | 1426 | if z: |
| 1383 | cfg["ZBLOCK"] = min(z, size_hints["z"]) | 1427 | cfg["ZBLOCK"] = min(z, size_hints["z"]) |
| 1428 | + config = Config(cfg, num_warps=8, num_stages=num_stages) | ||
| 1384 | if auto_blockify_size is not None: | 1429 | if auto_blockify_size is not None: |
| 1385 | - cfg["auto_blockify_size"] = auto_blockify_size | 1430 | + # Backend compile option, NOT a kernel constexpr: it must never enter |
| 1386 | - return Config(cfg, num_warps=8, num_stages=num_stages) | 1431 | + # cfg.kwargs (whose keys all flow into ASTSource constants and must be |
| 1432 | + # kernel signature params). Ride the upstream third-party backend slot | ||
| 1433 | + # Config.extra_options instead -- AutotuneCache.save serializes it and | ||
| 1434 | + # read_best restores it -- and reach triton.compile via the options | ||
| 1435 | + # dict (NPUOptions.auto_blockify_size), same channel triton-ascend's | ||
| 1436 | + # own autotune examples use at JIT-launch level. | ||
| 1437 | + config.extra_options = {"auto_blockify_size": auto_blockify_size} | ||
| 1438 | + return config | ||
| 1387 | 1439 | ||
| 1388 | 1440 | ||
| 1389 | autotune_enhance = ncfg.autotune_enhance | 1441 | autotune_enhance = ncfg.autotune_enhance |