已合并
fix(npu-inductor): int32 index overflow — promote only overflow addends (variant C) #43736
huyuchao创建于 8月4日
fix(npu-inductor): int32 index overflow — promote only overflow addends (variant C) #43736
已合并
共 4 个文件变更+915-52
| @@ -0,0 +1,440 @@ | |||
| 1 | +# Owner(s): ["module: tests"] | ||
| 2 | +import functools | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +import torch._dynamo as dynamo | ||
| 6 | +from torch._inductor.utils import run_and_get_code | ||
| 7 | +from torch.testing._internal.common_utils import ( | ||
| 8 | + run_tests, | ||
| 9 | + instantiate_parametrized_tests, | ||
| 10 | +) | ||
| 11 | +from testutils import TestUtils | ||
| 12 | + | ||
| 13 | +import torch_npu # noqa: F401 | ||
| 14 | +import torch_npu._inductor.triton_experimental.config as ncfg | ||
| 15 | + | ||
| 16 | + | ||
| 17 | +# Regression test for the int32 index-overflow bug in the triton_experimental | ||
| 18 | +# NPU backend (2026-08). | ||
| 19 | +# | ||
| 20 | +# Background: for a reduction kernel whose total element count | ||
| 21 | +# (numel * reduction_numel) exceeds int32_max, select_index_dtype() selects | ||
| 22 | +# int64 indexing. Before the fix, that signal was propagated to every index | ||
| 23 | +# construction (pid / arange tile / scalar odometer), which (a) let the | ||
| 24 | +# overflowing pointer addend (e.g. 268435456*x1) wrap to a negative address in | ||
| 25 | +# int32 → Ascend vector-core fault 507035, and (b) when the whole kernel was | ||
| 26 | +# upcast to int64 instead, doubled UB on big arange tiles → vector-core timeout | ||
| 27 | +# 507034. The fix keeps the big arange tiles int32 and widens EVERY axis | ||
| 28 | +# factor in the address expressions at the load/store use site | ||
| 29 | +# (coeff*var.to(tl.int64)) — unconditional, constructive correctness. | ||
| 30 | +# | ||
| 31 | +# This test pins the behavior at the 2^31 boundary: 8388609*256 == 2^31 + 257 | ||
| 32 | +# elements, just past int32_max, so the kernel must go int64-indexed without a | ||
| 33 | +# full tile upcast. | ||
| 34 | +def _free_hbm_bytes(): | ||
| 35 | + try: | ||
| 36 | + free, _total = torch.npu.mem_get_info() | ||
| 37 | + return free | ||
| 38 | + except Exception: | ||
| 39 | + return None | ||
| 40 | + | ||
| 41 | + | ||
| 42 | +def skipIfInsufficientHBM(min_free_bytes): | ||
| 43 | + """Skip a large test when the device's free HBM is verifiably below the | ||
| 44 | + requirement. The overflow corpus keeps the input, an eager reference, | ||
| 45 | + outputs and compile caches resident at once, so budget ~2x the tensor.""" | ||
| 46 | + def deco(fn): | ||
| 47 | + | ||
| 48 | + def wrapper(self): | ||
| 49 | + free = _free_hbm_bytes() | ||
| 50 | + if free is not None and free < min_free_bytes: | ||
| 51 | + self.skipTest( | ||
| 52 | + f"large test: needs >= {min_free_bytes >> 30} GiB free " | ||
| 53 | + f"HBM, only {free >> 30} GiB available" | ||
| 54 | + ) | ||
| 55 | + return fn(self) | ||
| 56 | + return wrapper | ||
| 57 | + return deco | ||
| 58 | + | ||
| 59 | + | ||
| 60 | +def _cut_autotune_runs(): | ||
| 61 | + """Test-speedup: cut the mspti autotune benchmark to 1 warmup + 1 active run. | ||
| 62 | + | ||
| 63 | + Guard-value note: the assertions of this suite are codegen-text and | ||
| 64 | + numeric-correctness checks, only ONE real kernel run is needed; the | ||
| 65 | + production 25-run-per-candidate bench ranks configs for performance, | ||
| 66 | + which the tests do not measure. The candidate set itself is left | ||
| 67 | + untouched (the with_index kernel at > 2^31 elements needs the | ||
| 68 | + production set: single-candidate pinning was tried and the largest | ||
| 69 | + tile fails UB compile while small-tile fallbacks produce wrong | ||
| 70 | + indices). Returns (old_warmup, old_active) for a try/finally restore. | ||
| 71 | + """ | ||
| 72 | + old = (ncfg.mspti_warmup, ncfg.mspti_active) | ||
| 73 | + ncfg.mspti_warmup = 1 | ||
| 74 | + ncfg.mspti_active = 1 | ||
| 75 | + return old | ||
| 76 | + | ||
| 77 | + | ||
| 78 | +# P0 drift pin (2026-08-28 audit): the audited call surface of the type-name | ||
| 79 | +# helpers in the PINNED torch's codegen, keyed by source text without line | ||
| 80 | +# numbers — a moved line does not false-alarm, an added/changed call site | ||
| 81 | +# does. A torch bump failing test_int64_type_callsite_surface_pin means the | ||
| 82 | +# int64 type-role assumptions need a re-audit (update this list only after | ||
| 83 | +# re-auditing). | ||
| 84 | +_EXPECTED_INT64_TYPE_CALLSURFACE = [ | ||
| 85 | + "simd: return self.dtype_to_str(self.get_index_dtype_as_torch_dtype())", | ||
| 86 | + "triton: asm_triton_type = triton_type(dtype)", | ||
| 87 | + "triton: cast_inputs.append(f\"{inp}.to({triton_type(input_dtypes[i])})\")", | ||
| 88 | + "triton: f\"({str(logical_index)}).to({self.dtype_to_str(index_dtype)})\"", | ||
| 89 | + "triton: f\"({var}).to({triton_type(dtype)})\",", | ||
| 90 | + "triton: f\".to({triton_type(result_dtype)})\"", | ||
| 91 | + "triton: f\"{result} = {result}.to({triton_compute_type(target_dtype)})\"", | ||
| 92 | + "triton: f\"{torch.iinfo(index_dtype).max}, {self.dtype_to_str(index_dtype)})\"", | ||
| 93 | + "triton: f\"{value}.to({triton_compute_type(dtype)})\",", | ||
| 94 | + "triton: f\"{value}.to({triton_store_type(store_dtype)})\",", | ||
| 95 | + "triton: line += f\".to({triton_type(dtype)})\"", | ||
| 96 | + "triton: out = f\"{out}.to({triton_type(out_dtype)})\"", | ||
| 97 | + "triton: out = f\"{out}.to({triton_type(upcast_compute_type(dtype))})\"", | ||
| 98 | + "triton: out = f\"{x}.to({triton_type(dtype)}, bitcast=True)\"", | ||
| 99 | + "triton: out_dtype = triton_compute_type(dtype)", | ||
| 100 | + "triton: out_dtype = triton_store_type(dtype)", | ||
| 101 | + "triton: result = f\"{result}.to({self.dtype_to_str(result_type)})\"", | ||
| 102 | + "triton: result = f\"{result}.to({triton_type(result_dtype)})\"", | ||
| 103 | + "triton: return f\"({arg}).to({triton_type(dtype)})\"", | ||
| 104 | + "triton: return f\"{result}.to({triton_type(dtype)})\"", | ||
| 105 | + "triton: return f\"{result}.to({triton_type(result_dtype)})\"", | ||
| 106 | + "triton: return triton_compute_type(upcast_acc_dtype(dtype))", | ||
| 107 | + "triton: return triton_type(dtype)", | ||
| 108 | + "triton: return triton_type(upcast_compute_type(dtype))", | ||
| 109 | + "triton: triton_type = triton_compute_type(dtype)", | ||
| 110 | + "triton: value = f\"{value}.to({triton_store_type(store_dtype)})\"", | ||
| 111 | + "triton: x = f\"{x}.to({triton_type(src_dtype)})\"", | ||
| 112 | + "triton: {result_var}_ws = ({ws_name} + {self.index_to_str(ws_offset)}).to(tl.pointer_type({triton_type(dtype)}))", | ||
| 113 | +] | ||
| 114 | + | ||
| 115 | + | ||
| 116 | +class TestTritonExperimentalInt32Overflow(TestUtils): | ||
| 117 | + | ||
| 118 | + | ||
| 119 | + def test_sum_over_int32_max_promotes_overflow_addend(self): | ||
| 120 | + # 2^31 + 257 elements (8.0 GiB fp32) — the minimal > int32_max case. | ||
| 121 | + x = torch.randn(8388609, 256, device=torch.device("npu")) | ||
| 122 | + ref = x.sum(dim=1) | ||
| 123 | + | ||
| 124 | + def fn(t): | ||
| 125 | + return t.sum(dim=1) | ||
| 126 | + | ||
| 127 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 128 | + y, codes = run_and_get_code(cf, x) | ||
| 129 | + | ||
| 130 | + # Numeric correctness at > 2^31 scale. | ||
| 131 | + self.assertTrue(torch.allclose(y, ref, atol=1e-3, rtol=1e-3)) | ||
| 132 | + # Every axis factor of the address expression is widened to int64 at | ||
| 133 | + # the load site (unconditional); lanes/tiles stay int32. | ||
| 134 | + self.assertIn("to(tl.int64)", codes[0]) | ||
R | |||
| 135 | + | ||
| 136 | + | ||
| 137 | + def test_dynamic_over_int32_max_unconditional_widen(self): | ||
| 138 | + # Dynamic axis length whose trace-time hint already exceeds int32_max. | ||
| 139 | + # The widening is unconditional (every axis factor in the address | ||
| 140 | + # expression), so there is no snapshot-derived promote decision and no | ||
| 141 | + # runtime guard contract: correctness holds for ANY runtime shape and | ||
| 142 | + # the hint boundary is irrelevant. This case used to exercise the | ||
| 143 | + # guard-skip path; it now pins that the dynamic >2^31 kernel simply | ||
| 144 | + # compiles and runs correctly. | ||
| 145 | + s0 = 2_150_000_000 # numel s0*2 = 4.3e9 > 2^31 (8.6 GiB fp16) | ||
| 146 | + x = torch.full((s0, 2), 1.0, device="npu", dtype=torch.float16) | ||
| 147 | + dynamo.mark_dynamic(x, 0) | ||
| 148 | + ref = x.sum(dim=1) | ||
| 149 | + | ||
| 150 | + def fn(t): | ||
| 151 | + return t.sum(dim=1) | ||
| 152 | + | ||
| 153 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 154 | + y = cf(x) | ||
| 155 | + torch.npu.synchronize() | ||
| 156 | + | ||
| 157 | + self.assertTrue(torch.allclose(y, ref)) | ||
| 158 | + | ||
| 159 | + | ||
| 160 | + def test_non_linearize_over_int32_max_keeps_tiles_int32(self): | ||
| 161 | + # With codegen_linearize=False, past 2^31 elements the | ||
| 162 | + # non-linearize structure cannot stay correct — xoffset = | ||
| 163 | + # pid.to(int64)*XBLOCK upcasts every arange tile to int64 via mixed | ||
| 164 | + # broadcast (UB doubling, 507034) and the results go wrong. The kernel | ||
| 165 | + # must force the linearize structure (i64 rides the scalar | ||
| 166 | + # group_base/real_block chain and the widened address factors, tiles | ||
| 167 | + # stay int32) regardless of the config, which then only governs | ||
| 168 | + # in-range kernels. | ||
| 169 | + import torch_npu._inductor.triton_experimental.codegen.triton as tmod | ||
| 170 | + | ||
| 171 | + orig = tmod.triton_codegen_linearize | ||
| 172 | + tmod.triton_codegen_linearize = False | ||
| 173 | + try: | ||
| 174 | + x = torch.full((8388609, 256), 1.0, device="npu", dtype=torch.float32) | ||
| 175 | + ref = x.sum(dim=1) | ||
| 176 | + | ||
| 177 | + def fn(t): | ||
| 178 | + return t.sum(dim=1) | ||
| 179 | + | ||
| 180 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 181 | + y, codes = run_and_get_code(cf, x) | ||
| 182 | + torch.npu.synchronize() | ||
| 183 | + | ||
| 184 | + # Numeric correctness at > 2^31 scale with linearize forced off. | ||
| 185 | + self.assertTrue(torch.allclose(y, ref, atol=1e-3, rtol=1e-3)) | ||
| 186 | + # Reverse-guard the whole-tile upcast: no arange/full tile may be | ||
| 187 | + # int64 AT ALL (as a dtype literal or via a cast). The pre-fix | ||
| 188 | + # all-tile upcast emitted tl.arange(..., tl.int64) with no ".to" | ||
| 189 | + # call, so asserting only on the cast would not catch it (507034). | ||
| 190 | + for line in codes[0].splitlines(): | ||
| 191 | + if "tl.arange" in line or "tl.full" in line: | ||
| 192 | + self.assertNotIn( | ||
| 193 | + "tl.int64", line, | ||
| 194 | + f"non-linearize tile upcast past 2^31: {line.strip()}", | ||
| 195 | + ) | ||
| 196 | + # Forward-guard: the overflowing addend — 256*x0 | ||
| 197 | + # has term_max 256*8388608 == 2^31 > int32_max — must carry the | ||
| 198 | + # int64 cast. Widening is unconditional (constructive correctness), | ||
| 199 | + # so other axis factors of the same index may be cast too; | ||
| 200 | + # exclusivity is deliberately NOT pinned — a wrong widening | ||
| 201 | + # selection can only cost speed, never correctness. | ||
| 202 | + cast_lines = [l for l in codes[0].splitlines() if ".to(tl.int64)" in l] | ||
| 203 | + self.assertTrue( | ||
| 204 | + cast_lines, | ||
| 205 | + "no int64 cast in the address expressions at all", | ||
| 206 | + ) | ||
| 207 | + self.assertRegex( | ||
| 208 | + "\n".join(cast_lines), | ||
| 209 | + r"\b256\*x0\.to\(tl\.int64\)", | ||
| 210 | + f"overflow addend cast missing: {cast_lines}", | ||
| 211 | + ) | ||
| 212 | + finally: | ||
| 213 | + tmod.triton_codegen_linearize = orig | ||
| 214 | + | ||
| 215 | + | ||
| 216 | + def test_max_with_index_over_int32_max_correct(self): | ||
| 217 | + # The upstream arg-reduction index accumulator follows | ||
| 218 | + # select_index_dtype() and is emitted as a full int64 tile past 2^31 | ||
| 219 | + # elements (tl.full(..., tl.int64) + *_with_index compare/select over | ||
| 220 | + # the whole [X, R] tile). Variant C keeps this upstream default — the | ||
| 221 | + # tile works on NPU, verified correct — and promotes ONLY the pointer | ||
| 222 | + # addend to int64 at the load use site. Note aten.argmax itself does | ||
| 223 | + # NOT reach this path on NPU: it falls back to eager (no triton kernel | ||
| 224 | + # is generated), so torch.max(dim=1) is the reachable form of the | ||
| 225 | + # arg-reduction code path. | ||
| 226 | + # | ||
| 227 | + # Test time budget: the kernel is > 2^31 logical elements (hard | ||
| 228 | + # requirement of the overflow guard) and the with_index accumulator is | ||
| 229 | + # a full int64 tile, so a single real run costs ~3s on AIV; the | ||
| 230 | + # production autotune would then benchmark 20+ candidates x 25 runs | ||
| 231 | + # (~15min) without adding guard value. The assertions below need only | ||
| 232 | + # ONE real run (codegen texts + numeric correctness are | ||
| 233 | + # config-independent), so cut the mspti bench to 1+1 and leave the | ||
| 234 | + # candidate set untouched. | ||
| 235 | + orig_warmup, orig_active = _cut_autotune_runs() | ||
| 236 | + try: | ||
| 237 | + x = torch.randn(8388609, 256, device="npu") | ||
| 238 | + ref_v, ref_i = x.max(dim=1) | ||
| 239 | + | ||
| 240 | + def fn(t): | ||
| 241 | + return t.max(dim=1) | ||
| 242 | + | ||
| 243 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 244 | + result, codes = run_and_get_code(cf, x) | ||
| 245 | + yv, yi = result | ||
| 246 | + torch.npu.synchronize() | ||
| 247 | + | ||
| 248 | + self.assertTrue(torch.allclose(yv, ref_v, atol=1e-3, rtol=1e-3)) | ||
| 249 | + self.assertTrue(torch.equal(yi, ref_i)) | ||
| 250 | + # The kernel must exist (compiled, not eager) and carry the overflow | ||
| 251 | + # addend cast (variant C at the load site). | ||
| 252 | + self.assertIn("with_index", codes[0]) | ||
| 253 | + self.assertIn(".to(tl.int64)", codes[0]) | ||
| 254 | + # R7-c trap pin: the with_index accumulator type and its | ||
| 255 | + # torch.iinfo(index_dtype).max sentinel fill must stay PAIRED — | ||
| 256 | + # narrowing dtype_to_str alone would emit the int64 max into a | ||
| 257 | + # tl.int32 full and break compilation. | ||
| 258 | + self.assertIn("9223372036854775807, tl.int64", codes[0]) | ||
| 259 | + finally: | ||
| 260 | + ncfg.mspti_warmup = orig_warmup | ||
| 261 | + ncfg.mspti_active = orig_active | ||
| 262 | + | ||
| 263 | + def test_mask_cmp_lhs_int64_narrow_protects_fp32(self): | ||
| 264 | + # With mask_cmp_fp32 on, the int32 narrow must | ||
| 265 | + # stay TERMINAL — falling through to the fp32 compare hangs the | ||
| 266 | + # kernel (triton-ascend lowers the fp32 mask compare into vector-core | ||
| 267 | + # work that never completes: 507034 vector-core timeout on a minimal | ||
| 268 | + # non-reduction kernel, >10 min no result on pointwise; verified | ||
| 269 | + # 2026-08-14). Pin the generated mask LHS texts: the default-off | ||
| 270 | + # layout byte-for-byte, and the fp32-on layout keeping the int64-narrow | ||
| 271 | + # case on the runnable int32 compare. These are pure-text assertions | ||
| 272 | + # (no compile) so they stay green while the fp32 path is broken. | ||
| 273 | + import torch_npu._inductor.triton_experimental.codegen.npu_header as nh | ||
| 274 | + import sympy | ||
| 275 | + | ||
| 276 | + off_cases = [ | ||
| 277 | + # (label, index_expr, numel, index_dtype, expected_lhs) | ||
| 278 | + ("int64-narrow", "x0index", 256, "tl.int64", "(x0index).to(tl.int32)"), | ||
| 279 | + ("int64-at-boundary", "x0index", 2**31, "tl.int64", "x0index"), | ||
| 280 | + ("int64-dynamic", "x0index", sympy.Symbol("s"), "tl.int64", "x0index"), | ||
| 281 | + ("int32", "x0index", 256, "tl.int32", "x0index"), | ||
| 282 | + ] | ||
| 283 | + orig = nh.npu_mask_cmp_fp32 | ||
| 284 | + try: | ||
| 285 | + nh.npu_mask_cmp_fp32 = False | ||
| 286 | + for label, e, n, d, want in off_cases: | ||
| 287 | + self.assertEqual( | ||
| 288 | + nh._mask_cmp_lhs(e, n, d), want, | ||
| 289 | + f"mask_cmp_fp32 off: {label}", | ||
| 290 | + ) | ||
| 291 | + # fp32 on: the int64 narrow keeps the int32 compare (protection); | ||
| 292 | + # the in-range int32 case is unchanged (the pre-fix fp32 cast). | ||
| 293 | + nh.npu_mask_cmp_fp32 = True | ||
| 294 | + self.assertEqual( | ||
| 295 | + nh._mask_cmp_lhs("x0index", 256, "tl.int64"), | ||
| 296 | + "(x0index).to(tl.int32)", | ||
| 297 | + "int64 narrow must not fall through to the hanging fp32 compare", | ||
| 298 | + ) | ||
| 299 | + self.assertEqual( | ||
| 300 | + nh._mask_cmp_lhs("x0index", 256, "tl.int32"), | ||
| 301 | + "(x0index).to(tl.float32)", | ||
| 302 | + ) | ||
| 303 | + finally: | ||
| 304 | + nh.npu_mask_cmp_fp32 = orig | ||
| 305 | + | ||
| 306 | + def test_non_linearize_in_range_reduction_grid_exact(self): | ||
| 307 | + # In-range non-linearize reduction: the launcher grid must be exactly | ||
| 308 | + # ceil(xnumel/XBLOCK). The heuristics' grid_0 defaults to the persistent | ||
| 309 | + # NPU_CU_COUNT (48) for reduction kernels, which is only valid under the | ||
| 310 | + # linearize structure's group dispatch — a non-linearize kernel | ||
| 311 | + # (xoffset = pid*XBLOCK, always-true xmask, no group folding) over-reads | ||
| 312 | + # past the input when 48 > ceil (MTE fault 507035, the pre-fix failure) | ||
| 313 | + # and silently drops tiles when 48 < ceil (uninitialized output rows, the | ||
| 314 | + # min(ceil, 48) clamp regression). Codegen flags | ||
| 315 | + # inductor_meta["npu_linearize"]=False; the heuristics then launches the | ||
| 316 | + # exact tile count instead of the fixed 48. | ||
| 317 | + import torch_npu._inductor.triton_experimental.codegen.triton as tmod | ||
| 318 | + | ||
| 319 | + orig = tmod.triton_codegen_linearize | ||
| 320 | + tmod.triton_codegen_linearize = False | ||
| 321 | + try: | ||
| 322 | + x = torch.randn(64, 128, 256, device="npu") | ||
| 323 | + ref = x.sum(dim=-1) | ||
| 324 | + | ||
| 325 | + def fn(t): | ||
| 326 | + return t.sum(dim=-1) | ||
| 327 | + | ||
| 328 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 329 | + y, codes = run_and_get_code(cf, x) | ||
| 330 | + torch.npu.synchronize() | ||
| 331 | + | ||
| 332 | + # Numeric correctness pins both failure modes: the 507035 over-read | ||
| 333 | + # (pre-fix) and the missing-tile clamp regression (wrong results). | ||
| 334 | + self.assertTrue(torch.allclose(y, ref, atol=1e-3, rtol=1e-3)) | ||
| 335 | + # Classic non-linearize structure, in-range so no int64 promotion. | ||
| 336 | + self.assertIn("xoffset = tl.program_id(0) * XBLOCK", codes[0]) | ||
| 337 | + self.assertNotIn("tl.int64", codes[0]) | ||
| 338 | + finally: | ||
| 339 | + tmod.triton_codegen_linearize = orig | ||
| 340 | + | ||
| 341 | + | ||
| 342 | + | ||
| 343 | + def test_expand_over_int32_blocks_odometer_i64(self): | ||
| 344 | + # Oversized-block-count dispatch audit: a static numel >= 2^31 must | ||
| 345 | + # never become | ||
| 346 | + # a triton literal. Three poison paths found and fixed by the | ||
| 347 | + # expand->sum >2^31 probe: (1) codegen_static_numels stomping the i64 | ||
| 348 | + # runtime arg with a bare literal (uint32 typing in [2^31, 2^32) -> | ||
| 349 | + # signedness errors in the div/mod dispatch chains), (2) per-axis | ||
| 350 | + # tl.constexpr numels in the linearize header, (3) the r-tree | ||
| 351 | + # tt.equal_to constants specialization (uint32->i64 vcast rejected by | ||
| 352 | + # BiShengIR). All three keep the numel on its i64 runtime arg, so the | ||
| 353 | + # whole block-dispatch scalar chain promotes to int64. This case has | ||
| 354 | + # > 2^31 total blocks via a stride-0 expand axis (numel 2.2e9, | ||
| 355 | + # storage ONE element) — cheap, no HBM skip needed. | ||
| 356 | + n = 2_200_000_000 # > 2^31, inside the uint32 window [2^31, 2^32) | ||
| 357 | + x = torch.ones((1,), device="npu").expand((n,)) | ||
| 358 | + ref = x.sum() | ||
| 359 | + | ||
| 360 | + def fn(t): | ||
| 361 | + return t.sum() | ||
| 362 | + | ||
| 363 | + cf = torch.compile(fn, options={"npu_backend": "triton_experimental"}) | ||
| 364 | + y, codes = run_and_get_code(cf, x) | ||
| 365 | + torch.npu.synchronize() | ||
| 366 | + self.assertTrue(torch.allclose(y, ref, rtol=1e-5)) | ||
| 367 | + # no bare >= 2^31 numel literal in any generated kernel (stomp or | ||
| 368 | + # constexpr forms) | ||
| 369 | + for code in codes: | ||
| 370 | + self.assertNotRegex(code, r"numel = 2\d{9}") | ||
| 371 | + self.assertNotRegex(code, r"numel : tl.constexpr = 2\d{9}") | ||
| 372 | + | ||
| 373 | + def test_dtype_role_isolation_pins(self): | ||
| 374 | + # P0 pin: the dtype-role separation is currently BY CONVENTION (three | ||
| 375 | + # mutually-counteracting global patches; see NPUTritonKernel. | ||
| 376 | + # dtype_to_str's SCOPE comment for the audited truth). These pins turn | ||
| 377 | + # silent drift (upstream sync, future edits) into loud failures. | ||
| 378 | + import torch._inductor.utils as inductor_utils | ||
| 379 | + | ||
| 380 | + import torch_npu._inductor.triton_experimental.codegen.triton as te_triton | ||
| 381 | + | ||
| 382 | + # (a) in-range kernel stays byte-clean int32: the widening must not | ||
| 383 | + # leak into kernels that do not need it | ||
| 384 | + x = torch.randn(1024, 64, device="npu") | ||
| 385 | + | ||
| 386 | + cf = torch.compile(lambda t: t.sum(1), options={"npu_backend": "triton_experimental"}) | ||
| 387 | + y, codes = run_and_get_code(cf, x) | ||
| 388 | + self.assertTrue(torch.allclose(y, x.sum(1), atol=1e-3, rtol=1e-3)) | ||
| 389 | + self.assertNotIn("tl.int64", codes[0]) | ||
| 390 | + | ||
| 391 | + # (b) compute types are NOT demoted: npu_triton_compute_type has no | ||
| 392 | + # int64 branch and bypasses the demotion mapping | ||
| 393 | + self.assertEqual(te_triton.npu_triton_compute_type(torch.int64), "tl.int64") | ||
| 394 | + | ||
| 395 | + # (c) seam policy pins: the mapping itself still demotes bare | ||
| 396 | + # triton_type calls (kept for default-backend coexistence), but the | ||
| 397 | + # two former seams are explicitly OPEN — int64 stores stay tl.int64 | ||
| 398 | + # (npu_triton_store_type) and non-int64 dtypes delegate upstream. | ||
| 399 | + # Any change in either direction must be a reviewed decision. | ||
| 400 | + # Requires activation (the compile above installs the patches). | ||
| 401 | + import torch._inductor.codegen.triton as up_codegen | ||
| 402 | + | ||
| 403 | + self.assertEqual(inductor_utils.triton_type(torch.int64), "tl.int32") | ||
| 404 | + self.assertEqual(up_codegen.triton_store_type(torch.int64), "tl.int64") | ||
| 405 | + self.assertEqual(up_codegen.triton_store_type(torch.bool), "tl.int8") | ||
| 406 | + self.assertEqual( | ||
| 407 | + up_codegen.triton_store_type(torch.float32), "tl.float32" | ||
| 408 | + ) | ||
| 409 | + | ||
| 410 | + def test_int64_type_callsite_surface_pin(self): | ||
| 411 | + # P0 pin: upstream drift detector, paired with | ||
| 412 | + # _EXPECTED_INT64_TYPE_CALLSURFACE above. Failing on a torch bump is | ||
| 413 | + # BY DESIGN: re-audit the int64 type-role surface, then update the | ||
| 414 | + # list. | ||
| 415 | + import inspect | ||
| 416 | + import re | ||
| 417 | + | ||
| 418 | + import torch._inductor.codegen.simd as up_simd | ||
| 419 | + import torch._inductor.codegen.triton as up_triton | ||
| 420 | + | ||
| 421 | + pat = re.compile(r"(dtype_to_str|triton_type|triton_compute_type|triton_store_type)\(") | ||
| 422 | + surface = set() | ||
| 423 | + for mod in (up_triton, up_simd): | ||
| 424 | + short = mod.__name__.rsplit(".", 1)[-1] | ||
| 425 | + for line in open(inspect.getsourcefile(mod)).read().splitlines(): | ||
| 426 | + s = line.strip() | ||
| 427 | + if pat.search(s) and not s.startswith(("def ", "#")): | ||
| 428 | + surface.add(f"{short}: {s}") | ||
| 429 | + self.assertEqual( | ||
| 430 | + sorted(surface), | ||
| 431 | + sorted(_EXPECTED_INT64_TYPE_CALLSURFACE), | ||
| 432 | + "int64 type-helper call surface drifted — re-audit the dtype " | ||
| 433 | + "roles before updating this snapshot", | ||
| 434 | + ) | ||
| 435 | + | ||
| 436 | + | ||
| 437 | +instantiate_parametrized_tests(TestTritonExperimentalInt32Overflow) | ||
| 438 | + | ||
| 439 | +if __name__ == "__main__": | ||
| 440 | + run_tests() | ||
| @@ -23,13 +23,69 @@ from torch._inductor.codegen.triton import texpr | |||
| 23 | npu_mask_cmp_fp32 = ncfg.mask_cmp_fp32 | 23 | npu_mask_cmp_fp32 = ncfg.mask_cmp_fp32 |
| 24 | 24 | ||
| 25 | 25 | ||
| 26 | -def _mask_cmp_lhs(index_expr: str) -> str: | 26 | +def _mask_cmp_lhs(index_expr: str, numel=None, index_dtype: str = "tl.int32") -> str: |
| 27 | - """Wrap a mask LHS so the `<` runs on the vector unit (fp32) when possible.""" | 27 | + """Wrap a mask LHS so the `<` runs on the vector unit when possible. |
| 28 | + | ||
| 29 | + int64 index vectors (kernels whose total element count exceeds 2^31) have | ||
| 30 | + no native int64 vector compare on Ascend, so a raw int64 ``<`` decays to a | ||
| 31 | + scalar loop. Per-axis indices stay below 2^31 -- only the final linear | ||
| 32 | + address combine needs int64 -- so cast the LHS back to int32 for the | ||
| 33 | + compare: exact below 2^31 and keeps the compare on the vector unit. The | ||
| 34 | + int32 cast is skipped when the axis itself could reach 2^31 (a wrapping | ||
| 35 | + mask would be silently wrong) or its length is dynamic (numel not a static | ||
| 36 | + int); those keep the correct (scalarized) int64 compare. | ||
| 37 | + | ||
| 38 | + The early return is deliberate PROTECTION, not an oversight: with | ||
| 39 | + mask_cmp_fp32 on, falling through to the fp32 compare | ||
| 40 | + hangs the kernel -- triton-ascend lowers the fp32 mask compare into | ||
| 41 | + vector-core work that never completes (507034 vector-core timeout on a | ||
| 42 | + minimal non-reduction kernel, and >10 min with no result on pointwise; | ||
| 43 | + verified 2026-08-14). The int32 narrow must therefore stay terminal; the | ||
| 44 | + fp32 path should be re-enabled only after the triton-ascend lowering is | ||
| 45 | + fixed, at which point the fall-through can be restored per the review. | ||
| 46 | + """ | ||
| 47 | + if index_dtype == "tl.int64" and isinstance(numel, (int, sympy.Integer)) and int(numel) < 2**31: | ||
| 48 | + return f"({index_expr}).to(tl.int32)" | ||
| 28 | if not npu_mask_cmp_fp32: | 49 | if not npu_mask_cmp_fp32: |
| 29 | return index_expr | 50 | return index_expr |
| 30 | return f"({index_expr}).to(tl.float32)" | 51 | return f"({index_expr}).to(tl.float32)" |
| 31 | 52 | ||
| 32 | 53 | ||
| 54 | +def _npu_emit_axis_numel(pre_loop, tree, node): | ||
| 55 | + """Emit the per-axis numel definition (oversized-block-count dispatch audit). Static values < 2^31 fold to | ||
| 56 | + a tl.constexpr (upstream behaviour). A static value >= 2^31 must NOT | ||
| 57 | + become a triton literal — values in [2^31, 2^32) type as uint32 and | ||
| 58 | + poison the div/mod dispatch chains with signedness errors, and larger | ||
| 59 | + values blow int32 block counts — so it aliases the tree's i64 runtime | ||
| 60 | + arg (guaranteed i64: an axis >= 2^31 forces total numel >= 2^31, which | ||
| 61 | + selects index_dtype=int64 and size_dtype types the arg i64). The alias | ||
| 62 | + divides out the sibling axes' product, exact by range-tree | ||
| 63 | + construction (tree numel == product of free-axis numels).""" | ||
| 64 | + length = node.length | ||
| 65 | + if not isinstance(length, (int, sympy.Integer)): | ||
| 66 | + return | ||
| 67 | + if int(length) < 2**31: | ||
| 68 | + pre_loop.writeline(f"{node.name}numel : tl.constexpr = {int(length)}") | ||
| 69 | + return | ||
| 70 | + sib_prod = sympy.Integer(1) | ||
| 71 | + for n in tree.nodes.values(): | ||
| 72 | + if n.name in tree.tree_node_mapping or n.name == node.name: | ||
| 73 | + continue | ||
| 74 | + sib_prod = sib_prod * sympy.sympify(n.length) | ||
| 75 | + if not (sib_prod.is_Integer and int(sib_prod) >= 1): | ||
| 76 | + raise RuntimeError( | ||
| 77 | + f"[triton_experimental] axis {node.name!r} has a static numel " | ||
| 78 | + f">= 2^31 ({length}) with non-static sibling axes; this " | ||
| 79 | + "combination needs a dedicated lowering (oversized-block-count dispatch)" | ||
| 80 | + ) | ||
| 81 | + if int(sib_prod) == 1: | ||
| 82 | + pre_loop.writeline(f"{node.name}numel = {tree.prefix}numel") | ||
| 83 | + else: | ||
| 84 | + pre_loop.writeline( | ||
| 85 | + f"{node.name}numel = {tree.prefix}numel // {int(sib_prod)}" | ||
| 86 | + ) | ||
| 87 | + | ||
| 88 | + | ||
| 33 | def _ordered_mapping_items(mapping): | 89 | def _ordered_mapping_items(mapping): |
| 34 | """Yield (name, expr) entries in dependency order: any name referenced by | 90 | """Yield (name, expr) entries in dependency order: any name referenced by |
| 35 | another mapping's expression is emitted before that mapping. The dict's | 91 | another mapping's expression is emitted before that mapping. The dict's |
| @@ -470,14 +526,12 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 470 | # Scalar odometer axis: no register tile. real_block==1 so the odometer | 526 | # Scalar odometer axis: no register tile. real_block==1 so the odometer |
| 471 | # walks one element/block (block count == numel) and the index is a scalar. | 527 | # walks one element/block (block count == numel) and the index is a scalar. |
| 472 | if node.name in _scalar_odo_names: | 528 | if node.name in _scalar_odo_names: |
| 473 | - if isinstance(node.length, (int, sympy.Integer)): | 529 | + _npu_emit_axis_numel(pre_loop, tree, node) |
| 474 | - pre_loop.writeline(f"{node.name}numel : tl.constexpr = {int(node.length)}") | ||
| 475 | pre_loop.writeline(f"real_block_{node.name} : tl.constexpr = 1") | 530 | pre_loop.writeline(f"real_block_{node.name} : tl.constexpr = 1") |
| 476 | continue | 531 | continue |
| 477 | # Unify candidates get real_block emitted after the tile is aligned below. | 532 | # Unify candidates get real_block emitted after the tile is aligned below. |
| 478 | if node.name in _unify_names: | 533 | if node.name in _unify_names: |
| 479 | - if isinstance(node.length, (int, sympy.Integer)): | 534 | + _npu_emit_axis_numel(pre_loop, tree, node) |
| 480 | - pre_loop.writeline(f"{node.name}numel : tl.constexpr = {int(node.length)}") | ||
| 481 | continue | 535 | continue |
| 482 | # Under greedy-via-unify every non-scalar-odo free axis is a unify | 536 | # Under greedy-via-unify every non-scalar-odo free axis is a unify |
| 483 | # candidate; the only other path was the legacy real_block=numel//divisor | 537 | # candidate; the only other path was the legacy real_block=numel//divisor |
| @@ -658,6 +712,20 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 658 | # offsets, emitted after real_block. odometer_opt (default ON): B1 hoists the | 712 | # offsets, emitted after real_block. odometer_opt (default ON): B1 hoists the |
| 659 | # cumulative block product and divides pid once (shorter div chain); B2 drops | 713 | # cumulative block product and divides pid once (shorter div chain); B2 drops |
| 660 | # provably single-block axes (static numel==1 -> offset 0, no cumprod term). | 714 | # provably single-block axes (static numel==1 -> offset 0, no cumprod term). |
| 715 | + # Oversized-block-count dispatch audit — static numels >= 2^31 must | ||
| 716 | + # never become triton | ||
| 717 | + # literals. The original "safe by promotion" audit was REFUTED by the | ||
| 718 | + # expand->sum >2^31 probe (2026-08-29): literals in [2^31, 2^32) type as | ||
| 719 | + # uint32 (signedness errors in the div/mod dispatch chains) and equal_to | ||
| 720 | + # specialization turns the uint32 into a BiShengIR-rejected uint32->i64 | ||
| 721 | + # vcast. Fixed at every emission point (codegen_static_numels, the | ||
| 722 | + # per-axis constexpr numels via _npu_emit_axis_numel, block counts, and | ||
| 723 | + # the r-tree constants specialization): such numels stay on their i64 | ||
| 724 | + # runtime arg (size_dtype=index_dtype in int64 mode), so the block | ||
| 725 | + # counts, cumblk products, (group_base + i) and % blocks promote to i64 | ||
| 726 | + # end-to-end through triton's own rules. Pinned by | ||
| 727 | + # test_expand_over_int32_blocks_odometer_i64 (> 2^31 total blocks via a | ||
| 728 | + # stride-0 expand axis — cheap, 1-element storage). | ||
| 661 | _odo_opt = ncfg.odometer_opt | 729 | _odo_opt = ncfg.odometer_opt |
| 662 | _free_nodes_ordered = [n for n in tree.nodes.values() | 730 | _free_nodes_ordered = [n for n in tree.nodes.values() |
| 663 | if n.name not in tree.tree_node_mapping] | 731 | if n.name not in tree.tree_node_mapping] |
| @@ -669,7 +737,16 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 669 | for node in _free_nodes_ordered: | 737 | for node in _free_nodes_ordered: |
| 670 | divisor_is_static = isinstance(node.divisor, (int, sympy.Integer)) | 738 | divisor_is_static = isinstance(node.divisor, (int, sympy.Integer)) |
| 671 | length_is_static = isinstance(node.length, (int, sympy.Integer)) | 739 | length_is_static = isinstance(node.length, (int, sympy.Integer)) |
| 672 | - blocks_is_constexpr = length_is_static and divisor_is_static | 740 | + # Oversized block counts: a static numel >= 2^31 is aliased to the i64 runtime arg by |
| 741 | + # _npu_emit_axis_numel, so its block count must be runtime too. | ||
| 742 | + blocks_is_constexpr = ( | ||
| 743 | + length_is_static | ||
| 744 | + and divisor_is_static | ||
| 745 | + and not ( | ||
| 746 | + isinstance(node.length, (int, sympy.Integer)) | ||
| 747 | + and int(node.length) >= 2**31 | ||
| 748 | + ) | ||
| 749 | + ) | ||
| 673 | if blocks_is_constexpr: | 750 | if blocks_is_constexpr: |
| 674 | pre_loop.writeline(f"{node.name}_blocks : tl.constexpr = ({node.name}numel + real_block_{node.name} - 1) // real_block_{node.name}") # noqa: B950 | 751 | pre_loop.writeline(f"{node.name}_blocks : tl.constexpr = ({node.name}numel + real_block_{node.name} - 1) // real_block_{node.name}") # noqa: B950 |
| 675 | else: | 752 | else: |
| @@ -708,6 +785,11 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 708 | tree.node_block_constexpr = {} | 785 | tree.node_block_constexpr = {} |
| 709 | tree.node_block_constexpr.update(node_arange_upper) | 786 | tree.node_block_constexpr.update(node_arange_upper) |
| 710 | 787 | ||
| 788 | + # Kernels whose total element count exceeds 2^31 index with tl.int64 (the | ||
| 789 | + # dtype_to_str override in triton.py restores index_dtype). Arange tiles stay | ||
| 790 | + # int32 (variant C): index_to_str in triton.py promotes only the overflow | ||
| 791 | + # addend (e.g. 268435456*x1) to int64 at the load/store use site, keeping the | ||
| 792 | + # big vector tiles int32 (all-int64 tiles double UB and hang -- 507034). | ||
| 711 | for node in tree.nodes.values(): | 793 | for node in tree.nodes.values(): |
| 712 | if node.name in tree.tree_node_mapping: | 794 | if node.name in tree.tree_node_mapping: |
| 713 | continue | 795 | continue |
| @@ -738,7 +820,8 @@ def _codegen_header_npu_for_tree(kernel, tree, code, outer_blocks=None): | |||
| 738 | line = kernel.iteration_ranges_scalar_code(tree, f"{node.name}offset") | 820 | line = kernel.iteration_ranges_scalar_code(tree, f"{node.name}offset") |
| 739 | header_code.writeline(f"{node.name}index = {line}") | 821 | header_code.writeline(f"{node.name}index = {line}") |
| 740 | header_code.writeline(f"{node.name} = {node.name}index") | 822 | header_code.writeline(f"{node.name} = {node.name}index") |
| 741 | - header_code.writeline(f"{node.name}mask = {_mask_cmp_lhs(f'{node.name}index')} < {node.name}numel") | 823 | + mask_lhs = _mask_cmp_lhs(f"{node.name}index", node.length, kernel.index_dtype) |
| 824 | + header_code.writeline(f"{node.name}mask = {mask_lhs} < {node.name}numel") | ||
| 742 | 825 | ||
| 743 | # Every free axis is now tiled flat (greedy-via-unify / scalar-odometer / | 826 | # Every free axis is now tiled flat (greedy-via-unify / scalar-odometer / |
| 744 | # static-constexpr): needs_inner_loop is False for all of them, so no axis | 827 | # static-constexpr): needs_inner_loop is False for all of them, so no axis |
| @@ -755,6 +755,87 @@ _TritonPrinter._print_Float = _npu_print_Float | |||
| 755 | _TritonPrinter._print_ToFloat = _npu_print_ToFloat | 755 | _TritonPrinter._print_ToFloat = _npu_print_ToFloat |
| 756 | 756 | ||
| 757 | 757 | ||
| 758 | +# Variant C int64 widening, structural (replaces the old re.sub text rewrite). | ||
| 759 | +# NpuWiden(x) is an unevaluated marker whose ONLY meaning is printer-level: it | ||
| 760 | +# renders as "x.to(tl.int64)" (same trick as ToFloat above). The cast is | ||
| 761 | +# attached to axis factors inside pointer-arithmetic expressions at the sympy | ||
| 762 | +# level (NPUTritonKernel.index_to_str), so no rendered-text matching is | ||
| 763 | +# involved and composite terms (c*x0*x1, FloorDiv(x, k)*c) widen correctly. | ||
| 764 | +class NpuWiden(sympy.Function): | ||
| 765 | + | ||
| 766 | + def eval(cls, arg): | ||
| 767 | + return None # stay unevaluated; survives Add/Mul rebalancing | ||
| 768 | + | ||
| 769 | + | ||
| 770 | +def _npu_print_NpuWiden(self, expr): | ||
| 771 | + # Parenthesize anything that is not an atom (same idiom as ToFloat); for | ||
| 772 | + # the plain monomial case the atom renders bare, byte-identical to the | ||
| 773 | + # old regex output ("268435456*x1.to(tl.int64)"). | ||
| 774 | + from sympy.printing.precedence import PRECEDENCE | ||
| 775 | + s = self.parenthesize(expr.args[0], PRECEDENCE["Atom"] - 0.5) | ||
| 776 | + return f"{s}.to(tl.int64)" | ||
| 777 | + | ||
| 778 | + | ||
| 779 | +_TritonPrinter._print_NpuWiden = _npu_print_NpuWiden | ||
| 780 | + | ||
| 781 | + | ||
| 782 | +# R3 hardening: the numel-level gate (select_index_dtype -> | ||
| 783 | +# can_use_32bit_indexing) only checks total numel and buffer storage sizes, | ||
| 784 | +# so address expressions carrying large index-arithmetic constants (upstream | ||
| 785 | +# #186057: model integer math, unfold/flatten offset algebra — constants NOT | ||
| 786 | +# reflected in any buffer's storage) can overflow int32 in a kernel the gate | ||
| 787 | +# typed as int32. Widening is cheap (perf-waived) and shape-independent, so | ||
| 788 | +# arm it on any static proximity hint instead of trying to prove overflow: | ||
| 789 | +# an integer literal of magnitude >= 2^30 anywhere in the expression | ||
| 790 | +# (addend, stride coefficient or mod/div base) can bust 2^31 once combined | ||
| 791 | +# with an axis extent of the same order; a STATIC axis within a whisker of | ||
| 792 | +# int32_max busts with even a small constant addend. Ordinary strides and | ||
| 793 | +# offsets sit far below both thresholds. | ||
| 794 | +_NPU_INT32_ARM_LITERAL = 2**30 | ||
| 795 | + | ||
| 796 | + | ||
| 797 | +def _npu_should_widen_address(index_dtype: str, expr: sympy.Expr, range_tree_nodes) -> bool: | ||
| 798 | + """Whether index_to_str widens this address expression's axis factors: | ||
| 799 | + int64-typed kernels always; int32-typed kernels only on an #186057-class | ||
| 800 | + hint (conservative arm — a false positive costs the waived i64 address | ||
| 801 | + tax, never correctness). | ||
| 802 | + | ||
| 803 | + Two hints: (a) any integer literal >= 2^30 (fast path); (b) an | ||
| 804 | + expression-level bound — substitute every axis symbol with its max lane | ||
| 805 | + value and evaluate against the trace-time hint (the upstream | ||
| 806 | + #186057/d630a2d direction, vendored as a boolean gate). Arm when the | ||
| 807 | + WHOLE expression can bust int32, which closes the dynamic | ||
| 808 | + near-max-axis + small-constant corner the literal/axis heuristics | ||
| 809 | + missed. Non-monotone terms (ModularIndexing) make a point substitution | ||
| 810 | + unsound as a maximum, so their presence arms unconditionally; no hint | ||
| 811 | + (unbacked) or any evaluation failure arms too — every tie breaks toward | ||
| 812 | + the fail-safe direction. A different runtime shape recompiles through | ||
| 813 | + the ordinary shape guards with a fresh hint, so the decision is | ||
| 814 | + re-evaluated per instance.""" | ||
| 815 | + if index_dtype == "tl.int64": | ||
| 816 | + return True | ||
| 817 | + if any(abs(int(lit)) >= _NPU_INT32_ARM_LITERAL for lit in expr.atoms(sympy.Integer)): | ||
| 818 | + return True | ||
| 819 | + if expr.has(ModularIndexing): | ||
| 820 | + return True | ||
| 821 | + subs = {} | ||
| 822 | + for sym in expr.free_symbols: | ||
| 823 | + node = range_tree_nodes.get(sym) | ||
| 824 | + if node is not None: | ||
| 825 | + subs[sym] = node.length - 1 | ||
| 826 | + if not subs: | ||
| 827 | + return False | ||
| 828 | + try: | ||
| 829 | + bound_expr = expr.subs(subs) | ||
| 830 | + if isinstance(bound_expr, (int, sympy.Integer)): | ||
| 831 | + bound = int(bound_expr) | ||
| 832 | + else: | ||
| 833 | + bound = int(V.graph.sizevars.guarding_hint_or_throw(bound_expr)) | ||
| 834 | + return abs(bound) > 2**31 - 1 | ||
| 835 | + except Exception: | ||
| 836 | + return True | ||
| 837 | + | ||
| 838 | + | ||
| 758 | # Tensor-dimension symbol kinds: any tensor actually indexed by a running kernel | 839 | # Tensor-dimension symbol kinds: any tensor actually indexed by a running kernel |
| 759 | # has every dim >= 1, so an expression built only from these is >= 1. | 840 | # has every dim >= 1, so an expression built only from these is >= 1. |
| 760 | _DIM_SYMT = (SymT.SIZE, SymT.PRECOMPUTED_SIZE, SymT.UNBACKED_INT) | 841 | _DIM_SYMT = (SymT.SIZE, SymT.PRECOMPUTED_SIZE, SymT.UNBACKED_INT) |
| @@ -982,6 +1063,69 @@ def npu_triton_compute_type(dtype): | |||
| 982 | torch._inductor.codegen.triton.triton_compute_type = npu_triton_compute_type | 1063 | torch._inductor.codegen.triton.triton_compute_type = npu_triton_compute_type |
| 983 | 1064 | ||
| 984 | 1065 | ||
| 1066 | +# P1-1 / R7 seams (audited 2026-08-29): int64 must not be silently demoted on | ||
| 1067 | +# the two routes that bypass NPUTritonKernel.dtype_to_str — stores of int64 | ||
| 1068 | +# tensors (value-correct until now only while stored values fit int32) and | ||
| 1069 | +# index-as-value casts (which actively NARROWED int64 values back to tl.int32 | ||
| 1070 | +# in int32-mode kernels). Both are re-opened with the same module-patch | ||
| 1071 | +# pattern as npu_triton_compute_type above; like it, they are process-global | ||
| 1072 | +# and thus shared with the default backend in mixed processes (accepted, | ||
| 1073 | +# consistent precedent). | ||
| 1074 | +from torch._inductor.codegen.triton import ( | ||
| 1075 | + TritonSymbols as _TritonSymbols, | ||
| 1076 | + triton_store_type as _upstream_triton_store_type, | ||
| 1077 | + triton_type as _upstream_triton_type, | ||
| 1078 | +) | ||
| 1079 | + | ||
| 1080 | + | ||
| 1081 | +def npu_triton_store_type(dtype): | ||
| 1082 | + """int64 stores stay tl.int64 (upstream routes through the demotion | ||
| 1083 | + mapping); every other dtype delegates to upstream unchanged.""" | ||
| 1084 | + if dtype == torch.int64: | ||
| 1085 | + return "tl.int64" | ||
| 1086 | + return _upstream_triton_store_type(dtype) | ||
| 1087 | + | ||
| 1088 | + | ||
| 1089 | +torch._inductor.codegen.triton.triton_store_type = npu_triton_store_type | ||
| 1090 | + | ||
| 1091 | + | ||
| 1092 | +def _npu_value_type_str(dtype): | ||
| 1093 | + if dtype == torch.int64: | ||
| 1094 | + return "tl.int64" | ||
| 1095 | + return _upstream_triton_type(dtype) | ||
| 1096 | + | ||
| 1097 | + | ||
| 1098 | + | ||
| 1099 | +def _npu_value_expr(cls, expr, dtype): | ||
| 1100 | + """Patched copy of upstream TritonSymbols.value_expr, kept in sync with | ||
| 1101 | + it; the ONLY change is the final cast routing through | ||
| 1102 | + _npu_value_type_str so int64 index values are not narrowed back to | ||
| 1103 | + tl.int32 by the demotion mapping. | ||
| 1104 | + | ||
| 1105 | + Like :meth:`index_expr`, but honors ``dtype`` by setting the kernel | ||
| 1106 | + index dtype before emitting, and casting the result if needed. | ||
| 1107 | + """ | ||
| 1108 | + real_index_dtype = V.kernel._index_dtype | ||
| 1109 | + V.kernel._index_dtype = ( | ||
| 1110 | + dtype if dtype in (torch.int32, torch.int64) else torch.int64 | ||
| 1111 | + ) | ||
| 1112 | + try: | ||
| 1113 | + var = cls.index_expr(expr, dtype) | ||
| 1114 | + finally: | ||
| 1115 | + V.kernel._index_dtype = real_index_dtype | ||
| 1116 | + if real_index_dtype != dtype or var.dtype != dtype: | ||
| 1117 | + var = V.kernel.cse.generate( | ||
| 1118 | + V.kernel.compute, | ||
| 1119 | + f"({var}).to({_npu_value_type_str(dtype)})", | ||
| 1120 | + dtype=dtype, | ||
| 1121 | + shape=var.shape, | ||
| 1122 | + ) | ||
| 1123 | + return var | ||
| 1124 | + | ||
| 1125 | + | ||
| 1126 | +_TritonSymbols.value_expr = _npu_value_expr | ||
| 1127 | + | ||
| 1128 | + | ||
| 985 | class _FlatMapExpr: | 1129 | class _FlatMapExpr: |
| 986 | """Render-only adapter for collapse_rowmajor_xtree: quacks like a sympy expr | 1130 | """Render-only adapter for collapse_rowmajor_xtree: quacks like a sympy expr |
| 987 | (``free_symbols`` + ``str()``) but emits plain ``(F // d) % L`` the printer can | 1131 | (``free_symbols`` + ``str()``) but emits plain ``(F // d) % L`` the printer can |
| @@ -1521,6 +1665,69 @@ class NPUTritonKernelOverrides(TritonKernelOverrides): | |||
| 1521 | class NPUTritonKernel(TritonKernel): | 1665 | class NPUTritonKernel(TritonKernel): |
| 1522 | overrides = NPUTritonKernelOverrides # type: ignore[assignment] | 1666 | overrides = NPUTritonKernelOverrides # type: ignore[assignment] |
| 1523 | 1667 | ||
| 1668 | + def dtype_to_str(self, dtype): | ||
| 1669 | + # The apply_npu_codegen_patches demotion (_triton_type_mapping | ||
| 1670 | + # ["tl.int64"] -> "tl.int32") forces index_dtype to "tl.int32" for every | ||
| 1671 | + # kernel, silently wrapping pointer offsets past 2^31 (e.g. | ||
| 1672 | + # 268435456*x1 for x1 >= 8 in a >2^31-element reduction) into negative | ||
| 1673 | + # addresses that fault on Ascend (507035 / vector core exception). | ||
| 1674 | + # Returning the upstream-native "tl.int64" here restores that signal. | ||
| 1675 | + # | ||
| 1676 | + # SCOPE (audited 2026-08-28; do not trust the old "only the index | ||
| 1677 | + # path is affected" claim): this override reaches EVERY | ||
| 1678 | + # dtype_to_str(torch.int64) call site — the index_dtype derivation | ||
| 1679 | + # (simd.py index_dtype property), the arg-reduction with_index trio | ||
| 1680 | + # (accumulator type AND its torch.iinfo(index_dtype).max sentinel, | ||
| 1681 | + # which MUST stay type-consistent — narrowing one side alone emits | ||
| 1682 | + # tl.full(size, 2**63-1, tl.int32) and breaks compilation), and the | ||
| 1683 | + # reduction intermediate-result cast (upstream | ||
| 1684 | + # reduction_collapse_dims) which is a DATA context. npu_triton_compute_type | ||
| 1685 | + # likewise has no int64->int32 branch and bypasses the mapping, so | ||
| 1686 | + # compute types are not demoted either. As of 2026-08-29 the backend | ||
| 1687 | + # is fully int64-open across ALL routes: the two former seams — | ||
| 1688 | + # value_expr (index-as-value casts) and triton_store_type (int64 | ||
| 1689 | + # tensor stores) — are explicitly re-opened via module patches next | ||
| 1690 | + # to npu_triton_compute_type (they previously stayed demoted through | ||
| 1691 | + # the mapping, correct only while stored values fit int32). The | ||
| 1692 | + # remaining demotion surface is the mapping's int64 entry itself, | ||
| 1693 | + # vestigial for this backend (dtype_to_str overridden, patched routes | ||
| 1694 | + # bypass it) and kept for default-backend coexistence in mixed | ||
| 1695 | + # processes; any NEW direct triton_type(torch.int64) call site is | ||
| 1696 | + # caught by the callsite surface pin in | ||
| 1697 | + # test_triton_experimental_int32_overflow. | ||
| 1698 | + if dtype == torch.int64: | ||
R 上游 arg reduction 的 index accumulator 也会走   | |||
| 1699 | + return "tl.int64" | ||
| 1700 | + return super().dtype_to_str(dtype) | ||
| 1701 | + | ||
| 1702 | + def codegen_static_numels(self, code): | ||
| 1703 | + # Oversized-block-count dispatch audit: upstream stomps static | ||
| 1704 | + # numels with bare | ||
| 1705 | + # literals ("r0_numel = 2200000000"). A literal in [2^31, 2^32) types | ||
| 1706 | + # as triton uint32 and poisons every downstream scalar chain with | ||
| 1707 | + # signedness errors (rsplit dispatch //, group-base %, found by the | ||
| 1708 | + # expand->sum >2^31 probe), and block counts past 2^31 wrap int32. | ||
| 1709 | + # Keep such numels on their runtime arg — typed i64 via | ||
| 1710 | + # size_dtype=index_dtype in int64 mode — so the whole dispatch chain | ||
| 1711 | + # promotes to int64 instead. Otherwise kept in sync with upstream | ||
| 1712 | + # TritonKernel.codegen_static_numels; the persistent-reduction branch | ||
| 1713 | + # raises because NPU disables persistent reductions. | ||
| 1714 | + for tree in self.range_trees: | ||
| 1715 | + if not tree.is_reduction or self.inside_reduction: | ||
| 1716 | + simplified_tree_numel = V.graph.sizevars.simplify(tree.numel) | ||
| 1717 | + if ( | ||
| 1718 | + isinstance(simplified_tree_numel, (sympy.Integer, int)) | ||
| 1719 | + and int(simplified_tree_numel) < 2**31 | ||
| 1720 | + ): | ||
| 1721 | + code.writeline(f"{tree.prefix}numel = {int(simplified_tree_numel)}") | ||
| 1722 | + if tree.is_reduction and self.persistent_reduction: | ||
| 1723 | + raise AssertionError( | ||
| 1724 | + "persistent reduction is disabled on NPU " | ||
| 1725 | + "(should_use_persistent_reduction -> False); extend this " | ||
| 1726 | + "override from upstream if that ever changes" | ||
| 1727 | + ) | ||
| 1728 | + if tree.prefix == "x" and self.no_x_dim: | ||
| 1729 | + code.writeline("XBLOCK: tl.constexpr = 1") | ||
| 1730 | + | ||
| 1524 | def should_use_persistent_reduction(self) -> bool: | 1731 | def should_use_persistent_reduction(self) -> bool: |
| 1525 | # NPU does not support persistent reduction; always use the looped path | 1732 | # NPU does not support persistent reduction; always use the looped path |
| 1526 | # so that the accumulator (tl.full) is properly initialized. | 1733 | # so that the accumulator (tl.full) is properly initialized. |
| @@ -1790,7 +1997,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 1790 | """ | 1997 | """ |
| 1791 | if not ncfg.select_extract_slice: | 1998 | if not ncfg.select_extract_slice: |
| 1792 | return | 1999 | return |
| 1793 | - if not triton_codegen_linearize: | 2000 | + if not self._npu_linearize: |
| 1794 | return | 2001 | return |
| 1795 | if self.inside_reduction: | 2002 | if self.inside_reduction: |
| 1796 | return | 2003 | return |
| @@ -1837,9 +2044,60 @@ class NPUTritonKernel(TritonKernel): | |||
| 1837 | # needs to prove dim >= 1. Mirror the upstream list path. | 2044 | # needs to prove dim >= 1. Mirror the upstream list path. |
| 1838 | if isinstance(index, list): | 2045 | if isinstance(index, list): |
| 1839 | return f"[{', '.join(map(self.index_to_str, index))}]" | 2046 | return f"[{', '.join(map(self.index_to_str, index))}]" |
| 1840 | - return self.kexpr(self.rename_indexing(_drop_size_clamp(index))) | 2047 | + expr = self.rename_indexing(_drop_size_clamp(index)) |
| 2048 | + if _npu_should_widen_address(self.index_dtype, expr, self.range_tree_nodes): | ||
| 2049 | + # Constructive correctness: widen EVERY axis factor in the address | ||
| 2050 | + # expression, unconditionally — no term analysis, no bounds, no | ||
| 2051 | + # runtime guards. Every product/sum then computes in int64 (triton | ||
| 2052 | + # promotes on any i64 operand), so no address can wrap regardless | ||
| 2053 | + # of term shape (composite strides, FloorDiv/Mod, negative | ||
| 2054 | + # monomials) and correctness holds for ANY runtime shape: there is | ||
| 2055 | + # no snapshot to distrust and nothing to recompile. Tile headers | ||
| 2056 | + # keep lanes int32 (linearize bypasses the upstream | ||
| 2057 | + # arange/full/pid casts); the i64 footprint is bounded by | ||
| 2058 | + # (#free axes x XBLOCK) per address site and absorbed by the | ||
| 2059 | + # autotuner's small-tile fallback when a kernel sits near the UB | ||
| 2060 | + # ceiling. Dropping provably-redundant casts is deliberately NOT | ||
| 2061 | + # this layer's job: a wrong selection may only cost speed, never | ||
| 2062 | + # correctness, so it belongs to an optional demotion pass on top | ||
| 2063 | + # (upstream #91028/d630a2d ValueRange semantics), never here. | ||
| 2064 | + for axis in expr.free_symbols & set(self.range_tree_nodes): | ||
| 2065 | + expr = expr.subs(axis, NpuWiden(axis)) | ||
| 2066 | + return self.kexpr(expr) | ||
| 1841 | 2067 | ||
| 1842 | def __init__(self, *args, **kwargs): | 2068 | def __init__(self, *args, **kwargs): |
| 2069 | + # Placeholder before super().__init__() — initialize_range_tree() and | ||
| 2070 | + # the pre-codegen header pass (both run inside super().__init__) gate | ||
| 2071 | + # on this flag. The cross-boundary test MUST match select_index_dtype | ||
| 2072 | + # exactly (same features instance, cached): past 2^31 elements the | ||
| 2073 | + # non-linearize structure cannot stay correct (xoffset = | ||
| 2074 | + # pid.to(i64)*XBLOCK upcasts every arange tile via mixed broadcast — | ||
| 2075 | + # UB doubling 507034 and wrong results), so force the | ||
| 2076 | + # linearize structure: the i64 rides the scalar group_base/real_block | ||
| 2077 | + # chain, tiles stay int32. codegen_linearize then only governs | ||
| 2078 | + # in-range kernels. | ||
| 2079 | + features = kwargs.get("features") | ||
| 2080 | + if features is None and len(args) > 1: | ||
| 2081 | + features = args[1] | ||
| 2082 | + force_linearize = False | ||
| 2083 | + if features is not None: | ||
| 2084 | + try: | ||
| 2085 | + force_linearize = ( | ||
| 2086 | + self.dtype_to_str(features.select_index_dtype()) == "tl.int64" | ||
| 2087 | + ) | ||
| 2088 | + except Exception: | ||
| 2089 | + # Fail SAFE, not silent: linearize is the structure this | ||
| 2090 | + # backend deems correct past 2^31, so an undecidable kernel | ||
| 2091 | + # must not fall back to the non-linearize path on a swallowed | ||
| 2092 | + # exception (config semantics may not silently revert to the | ||
| 2093 | + # unsafe path). An in-range kernel hitting this only pays the | ||
| 2094 | + # linearize structure — never correctness. | ||
| 2095 | + log.warning( | ||
| 2096 | + "select_index_dtype raised; forcing linearize (safe side)", | ||
| 2097 | + exc_info=True, | ||
| 2098 | + ) | ||
| 2099 | + force_linearize = True | ||
| 2100 | + self._npu_linearize: bool = bool(triton_codegen_linearize) or force_linearize | ||
| 1843 | super().__init__(*args, **kwargs) | 2101 | super().__init__(*args, **kwargs) |
| 1844 | self._axis_split_subs: Dict[sympy.Symbol, sympy.Expr] = {} | 2102 | self._axis_split_subs: Dict[sympy.Symbol, sympy.Expr] = {} |
| 1845 | # r-axis cross-core split (OUTER reduction): when True, codegen_kernel | 2103 | # r-axis cross-core split (OUTER reduction): when True, codegen_kernel |
| @@ -1859,7 +2117,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 1859 | # detection replaces the old regex text classification. | 2117 | # detection replaces the old regex text classification. |
| 1860 | self._npu_select_lane_loads: Dict[str, Dict[str, Any]] = {} | 2118 | self._npu_select_lane_loads: Dict[str, Dict[str, Any]] = {} |
| 1861 | 2119 | ||
| 1862 | - if triton_codegen_linearize: | 2120 | + if self._npu_linearize: |
| 1863 | for tree in self.range_trees: | 2121 | for tree in self.range_trees: |
| 1864 | if not hasattr(tree, 'tree_node_mapping'): | 2122 | if not hasattr(tree, 'tree_node_mapping'): |
| 1865 | tree.tree_node_mapping = {} | 2123 | tree.tree_node_mapping = {} |
| @@ -1943,7 +2201,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 1943 | 2201 | ||
| 1944 | def prepare_indexing(self, index): | 2202 | def prepare_indexing(self, index): |
| 1945 | index = super().prepare_indexing(index) | 2203 | index = super().prepare_indexing(index) |
| 1946 | - if triton_codegen_linearize: | 2204 | + if self._npu_linearize: |
| 1947 | index = self._maybe_split_fused_axes(index) | 2205 | index = self._maybe_split_fused_axes(index) |
| 1948 | index = self._maybe_split_strided_axis(index) | 2206 | index = self._maybe_split_strided_axis(index) |
| 1949 | index = self._simplify_compound_indexing(index) | 2207 | index = self._simplify_compound_indexing(index) |
| @@ -2923,7 +3181,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 2923 | def initialize_range_tree(self, pid_cache): | 3181 | def initialize_range_tree(self, pid_cache): |
| 2924 | """Override to add tree_node_mapping for linearize mode.""" | 3182 | """Override to add tree_node_mapping for linearize mode.""" |
| 2925 | super().initialize_range_tree(pid_cache) | 3183 | super().initialize_range_tree(pid_cache) |
| 2926 | - if triton_codegen_linearize: | 3184 | + if self._npu_linearize: |
| 2927 | for tree in self.range_trees: | 3185 | for tree in self.range_trees: |
| 2928 | if not hasattr(tree, 'tree_node_mapping'): | 3186 | if not hasattr(tree, 'tree_node_mapping'): |
| 2929 | tree.tree_node_mapping = {} | 3187 | tree.tree_node_mapping = {} |
| @@ -2958,7 +3216,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 2958 | return imports.getvalue() | 3216 | return imports.getvalue() |
| 2959 | 3217 | ||
| 2960 | def triton_tensor_ndim(self): | 3218 | def triton_tensor_ndim(self): |
| 2961 | - if triton_codegen_linearize and getattr(self, '_linearize_applied', False): | 3219 | + if self._npu_linearize and getattr(self, '_linearize_applied', False): |
| 2962 | ndim = 0 | 3220 | ndim = 0 |
| 2963 | for tree in self.range_trees: | 3221 | for tree in self.range_trees: |
| 2964 | if tree.tensor_dim is not None: | 3222 | if tree.tensor_dim is not None: |
| @@ -2981,7 +3239,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 2981 | return sum(int(tree.tensor_dim is not None) for tree in self.range_trees) | 3239 | return sum(int(tree.tensor_dim is not None) for tree in self.range_trees) |
| 2982 | 3240 | ||
| 2983 | def dense_size_str(self): | 3241 | def dense_size_str(self): |
| 2984 | - if triton_codegen_linearize: | 3242 | + if self._npu_linearize: |
| 2985 | ndim = self.triton_tensor_ndim() | 3243 | ndim = self.triton_tensor_ndim() |
| 2986 | if ndim == 0: | 3244 | if ndim == 0: |
| 2987 | return "[]" | 3245 | return "[]" |
| @@ -2996,7 +3254,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 2996 | return f"[{', '.join(sizes)}]" | 3254 | return f"[{', '.join(sizes)}]" |
| 2997 | 3255 | ||
| 2998 | def reduction_resize(self, value): | 3256 | def reduction_resize(self, value): |
| 2999 | - if triton_codegen_linearize: | 3257 | + if self._npu_linearize: |
| 3000 | if not self.no_x_dim and self.inside_reduction: | 3258 | if not self.no_x_dim and self.inside_reduction: |
| 3001 | ndims = self.triton_tensor_ndim() | 3259 | ndims = self.triton_tensor_ndim() |
| 3002 | if ndims <= 1: | 3260 | if ndims <= 1: |
| @@ -3031,7 +3289,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3031 | # covers sum → upstream emits "[:, None]", broadcasting the OUTER store to [XBLOCK, | 3289 | # covers sum → upstream emits "[:, None]", broadcasting the OUTER store to [XBLOCK, |
| 3032 | # XBLOCK] (~100x). Mirror the permuted-slot logic, deviating ONLY in permuted-linearize. | 3290 | # XBLOCK] (~100x). Mirror the permuted-slot logic, deviating ONLY in permuted-linearize. |
| 3033 | if ( | 3291 | if ( |
| 3034 | - triton_codegen_linearize | 3292 | + self._npu_linearize |
| 3035 | and not self.no_x_dim | 3293 | and not self.no_x_dim |
| 3036 | and self.inside_reduction | 3294 | and self.inside_reduction |
| 3037 | and getattr(self, "_npu_tile_permuted", False) | 3295 | and getattr(self, "_npu_tile_permuted", False) |
| @@ -3058,7 +3316,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3058 | # arange slice and mask. Upstream _combine_contiguous_dims() merges [x0(100), | 3316 | # arange slice and mask. Upstream _combine_contiguous_dims() merges [x0(100), |
| 3059 | # x1(4)] into one flat x2(400), which has no per-node shape and breaks the | 3317 | # x1(4)] into one flat x2(400), which has no per-node shape and breaks the |
| 3060 | # broadcast. Disable entirely in linearize mode. | 3318 | # broadcast. Disable entirely in linearize mode. |
| 3061 | - if triton_codegen_linearize: | 3319 | + if self._npu_linearize: |
| 3062 | return index | 3320 | return index |
| 3063 | return super()._combine_contiguous_dims(index, tree) | 3321 | return super()._combine_contiguous_dims(index, tree) |
| 3064 | 3322 | ||
| @@ -3075,7 +3333,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3075 | self._npu_prepared_load_index = ( | 3333 | self._npu_prepared_load_index = ( |
| 3076 | result.index if isinstance(result, IndexingOptions) else None | 3334 | result.index if isinstance(result, IndexingOptions) else None |
| 3077 | ) | 3335 | ) |
| 3078 | - if not triton_codegen_linearize: | 3336 | + if not self._npu_linearize: |
| 3079 | return result | 3337 | return result |
| 3080 | 3338 | ||
| 3081 | # Linearize mode emits per-node masks (x0mask, x1mask, …) and xmask = their AND, | 3339 | # Linearize mode emits per-node masks (x0mask, x1mask, …) and xmask = their AND, |
| @@ -3218,7 +3476,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3218 | ) | 3476 | ) |
| 3219 | 3477 | ||
| 3220 | def iteration_ranges_get_pid(self, entry: IterationRangesRoot) -> str: | 3478 | def iteration_ranges_get_pid(self, entry: IterationRangesRoot) -> str: |
| 3221 | - if not triton_codegen_linearize: | 3479 | + if not self._npu_linearize: |
| 3222 | return super().iteration_ranges_get_pid(entry) | 3480 | return super().iteration_ranges_get_pid(entry) |
| 3223 | 3481 | ||
| 3224 | assert entry.grid_dim is not None | 3482 | assert entry.grid_dim is not None |
| @@ -3226,8 +3484,6 @@ class NPUTritonKernel(TritonKernel): | |||
| 3226 | # all dimensions — yz grid overflow handling is not needed here. | 3484 | # all dimensions — yz grid overflow handling is not needed here. |
| 3227 | key = "(group_base + i)" | 3485 | key = "(group_base + i)" |
| 3228 | pid = entry.pid_cache.get(key, key) | 3486 | pid = entry.pid_cache.get(key, key) |
| 3229 | - if self.index_dtype != "tl.int32": | ||
| 3230 | - return f"{pid}.to({self.index_dtype})" | ||
| 3231 | return pid | 3487 | return pid |
| 3232 | 3488 | ||
| 3233 | def codegen_range_tree(self): | 3489 | def codegen_range_tree(self): |
| @@ -3236,7 +3492,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3236 | For non-r dimensions, calls our custom codegen_header_npu instead of the | 3492 | For non-r dimensions, calls our custom codegen_header_npu instead of the |
| 3237 | default iteration_ranges_codegen_header. | 3493 | default iteration_ranges_codegen_header. |
| 3238 | """ | 3494 | """ |
| 3239 | - if not triton_codegen_linearize: | 3495 | + if not self._npu_linearize: |
| 3240 | return super().codegen_range_tree() | 3496 | return super().codegen_range_tree() |
| 3241 | 3497 | ||
| 3242 | # npu_header owns the large linearize header generator; import lazily to | 3498 | # npu_header owns the large linearize header generator; import lazily to |
| @@ -3276,9 +3532,9 @@ class NPUTritonKernel(TritonKernel): | |||
| 3276 | total_ndim = max(orig_ndim, npu_ndim, tree.tensor_dim + 1) | 3532 | total_ndim = max(orig_ndim, npu_ndim, tree.tensor_dim + 1) |
| 3277 | sizes = ["None"] * total_ndim | 3533 | sizes = ["None"] * total_ndim |
| 3278 | sizes[tree.tensor_dim] = ":" | 3534 | sizes[tree.tensor_dim] = ":" |
| 3279 | - index_dtype = self.index_dtype | 3535 | + # Arange tiles stay int32 (variant C): the pointer arithmetic |
| 3280 | - suffix = f".to({index_dtype})" if index_dtype != "tl.int32" else "" | 3536 | + # promotes only the overflow addend to int64 at the use site. |
| 3281 | - ranges_code = f"tl.arange(0, {tree.prefix.upper()}BLOCK)[{', '.join(sizes)}]{suffix}" | 3537 | + ranges_code = f"tl.arange(0, {tree.prefix.upper()}BLOCK)[{', '.join(sizes)}]" |
| 3282 | self.body.writeline( | 3538 | self.body.writeline( |
| 3283 | f"{tree.prefix}base = {ranges_code}" | 3539 | f"{tree.prefix}base = {ranges_code}" |
| 3284 | ) | 3540 | ) |
| @@ -3299,7 +3555,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3299 | assignments (x0 = x0index) are emitted by _codegen_header_npu_for_tree. | 3555 | assignments (x0 = x0index) are emitted by _codegen_header_npu_for_tree. |
| 3300 | Only r-tree (loop) entries use the default body/indexing_code path. | 3556 | Only r-tree (loop) entries use the default body/indexing_code path. |
| 3301 | """ | 3557 | """ |
| 3302 | - if not triton_codegen_linearize: | 3558 | + if not self._npu_linearize: |
| 3303 | return super().codegen_iteration_ranges_entry(entry) | 3559 | return super().codegen_iteration_ranges_entry(entry) |
| 3304 | 3560 | ||
| 3305 | if not entry.root.is_reduction: | 3561 | if not entry.root.is_reduction: |
| @@ -3344,7 +3600,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3344 | transpose of the in-loop ``r0_base`` (the latter built from tensor_dim) — | 3600 | transpose of the in-loop ``r0_base`` (the latter built from tensor_dim) — |
| 3345 | a loop-carried-type conflict when a free X-axis takes the inner-loop path. | 3601 | a loop-carried-type conflict when a free X-axis takes the inner-loop path. |
| 3346 | """ | 3602 | """ |
| 3347 | - if not triton_codegen_linearize: | 3603 | + if not self._npu_linearize: |
| 3348 | return super().iteration_ranges_ranges_code(entry) | 3604 | return super().iteration_ranges_ranges_code(entry) |
| 3349 | 3605 | ||
| 3350 | assert entry.tensor_dim is not None | 3606 | assert entry.tensor_dim is not None |
| @@ -3370,23 +3626,23 @@ class NPUTritonKernel(TritonKernel): | |||
| 3370 | else: | 3626 | else: |
| 3371 | effective_dim += 1 | 3627 | effective_dim += 1 |
| 3372 | size = self.indexing_size_str(effective_dim) | 3628 | size = self.indexing_size_str(effective_dim) |
| 3373 | - index_dtype = self.index_dtype | 3629 | + # Arange tiles stay int32 (variant C); see codegen_range_tree. |
| 3374 | - convert = f".to({index_dtype})" if index_dtype != "tl.int32" else "" | 3630 | + return f"tl.arange(0, {entry.prefix.upper()}BLOCK){size}" |
| 3375 | - return f"tl.arange(0, {entry.prefix.upper()}BLOCK){size}{convert}" | ||
| 3376 | 3631 | ||
| 3377 | def iteration_ranges_scalar_code(self, entry: IterationRangesRoot, value) -> str: | 3632 | def iteration_ranges_scalar_code(self, entry: IterationRangesRoot, value) -> str: |
| 3378 | """ | 3633 | """ |
| 3379 | Linearize mode: override scalar_code. | 3634 | Linearize mode: override scalar_code. |
| 3380 | """ | 3635 | """ |
| 3381 | - if not triton_codegen_linearize: | 3636 | + if not self._npu_linearize: |
| 3382 | return super().iteration_ranges_scalar_code(entry, value) | 3637 | return super().iteration_ranges_scalar_code(entry, value) |
| 3383 | 3638 | ||
| 3384 | - index_dtype = self.index_dtype | 3639 | + # Scalar odometer offsets stay int32 (variant C): the int64 cast for |
| 3640 | + # overflow addends rides on the scalar/latent tile at the use site. | ||
| 3385 | ndim = self.triton_tensor_ndim() | 3641 | ndim = self.triton_tensor_ndim() |
| 3386 | size = [1] * ndim | 3642 | size = [1] * ndim |
| 3387 | if self.no_x_dim: | 3643 | if self.no_x_dim: |
| 3388 | - return f"tl.full([1, 1], {value}, {index_dtype})" | 3644 | + return f"tl.full([1, 1], {value}, tl.int32)" |
| 3389 | - return f"tl.full({size}, {value}, {index_dtype})" | 3645 | + return f"tl.full({size}, {value}, tl.int32)" |
| 3390 | 3646 | ||
| 3391 | def _npu_rsplit_rprefix(self) -> str: | 3647 | def _npu_rsplit_rprefix(self) -> str: |
| 3392 | """Reduction-tree prefix (e.g. 'r0_') for the rsplit partial kernel.""" | 3648 | """Reduction-tree prefix (e.g. 'r0_') for the rsplit partial kernel.""" |
| @@ -3661,7 +3917,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3661 | # Collect block hints for linearize mode | 3917 | # Collect block hints for linearize mode |
| 3662 | block_hints = {} | 3918 | block_hints = {} |
| 3663 | axis_hints = [] | 3919 | axis_hints = [] |
| 3664 | - if triton_codegen_linearize: | 3920 | + if self._npu_linearize: |
| 3665 | for tree in self.range_trees: | 3921 | for tree in self.range_trees: |
| 3666 | if tree.prefix != 'r': | 3922 | if tree.prefix != 'r': |
| 3667 | block_hints[f"{tree.prefix.upper()}BLOCK_HINT"] = tree.get_block_hint() | 3923 | block_hints[f"{tree.prefix.upper()}BLOCK_HINT"] = tree.get_block_hint() |
| @@ -3733,7 +3989,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3733 | # is no longer emitted as a runtime arg: the greedy-via-unify tile | 3989 | # is no longer emitted as a runtime arg: the greedy-via-unify tile |
| 3734 | # scheme computes real_block from XBLOCK + size hints and never | 3990 | # scheme computes real_block from XBLOCK + size hints and never |
| 3735 | # references <node>divisor in the body, so the arg was dead.) | 3991 | # references <node>divisor in the body, so the arg was dead.) |
| 3736 | - if tree.prefix != 'r' and triton_codegen_linearize: | 3992 | + if tree.prefix != 'r' and self._npu_linearize: |
| 3737 | tree_node_mapping = getattr(tree, 'tree_node_mapping', {}) | 3993 | tree_node_mapping = getattr(tree, 'tree_node_mapping', {}) |
| 3738 | for node in tree.nodes.values(): | 3994 | for node in tree.nodes.values(): |
| 3739 | if node.name in tree_node_mapping: | 3995 | if node.name in tree_node_mapping: |
| @@ -3769,13 +4025,13 @@ class NPUTritonKernel(TritonKernel): | |||
| 3769 | "mix_mode": "aiv", # NPU: force vector kernel generation | 4025 | "mix_mode": "aiv", # NPU: force vector kernel generation |
| 3770 | } | 4026 | } |
| 3771 | triton_meta["configs"] = [config_of(signature)] | 4027 | triton_meta["configs"] = [config_of(signature)] |
| 3772 | - if triton_codegen_linearize: | 4028 | + if self._npu_linearize: |
| 3773 | triton_meta['block_hints'] = block_hints | 4029 | triton_meta['block_hints'] = block_hints |
| 3774 | triton_meta['axis_hints'] = axis_hints | 4030 | triton_meta['axis_hints'] = axis_hints |
| 3775 | 4031 | ||
| 3776 | optimize_mem = V.graph.is_inference or V.graph.is_backward | 4032 | optimize_mem = V.graph.is_inference or V.graph.is_backward |
| 3777 | npu_num_x_nodes = 0 | 4033 | npu_num_x_nodes = 0 |
| 3778 | - if triton_codegen_linearize: | 4034 | + if self._npu_linearize: |
| 3779 | for tree in self.range_trees: | 4035 | for tree in self.range_trees: |
| 3780 | if not tree.is_reduction and not tree.is_loop: | 4036 | if not tree.is_reduction and not tree.is_loop: |
| 3781 | tree_node_mapping = getattr(tree, 'tree_node_mapping', {}) | 4037 | tree_node_mapping = getattr(tree, 'tree_node_mapping', {}) |
| @@ -3812,12 +4068,19 @@ class NPUTritonKernel(TritonKernel): | |||
| 3812 | # Mark static r-tree numel as constants so NPU compiler can prove | 4068 | # Mark static r-tree numel as constants so NPU compiler can prove |
| 3813 | # r0_mask = (r0_index < r0_numel) is always true when R0_BLOCK == r0_numel, | 4069 | # r0_mask = (r0_index < r0_numel) is always true when R0_BLOCK == r0_numel, |
| 3814 | # eliminating the scalar select/boundary-check path in the ttadapter. | 4070 | # eliminating the scalar select/boundary-check path in the ttadapter. |
| 3815 | - if triton_codegen_linearize: | 4071 | + if self._npu_linearize: |
| 3816 | for tree in self.range_trees: | 4072 | for tree in self.range_trees: |
| 3817 | if tree.is_reduction and isinstance(tree.numel, (int, sympy.Integer)): | 4073 | if tree.is_reduction and isinstance(tree.numel, (int, sympy.Integer)): |
| 3818 | numel_name = f"{tree.prefix}numel" | 4074 | numel_name = f"{tree.prefix}numel" |
| 3819 | if any(getattr(s, 'name', None) == numel_name for s in signature): | 4075 | if any(getattr(s, 'name', None) == numel_name for s in signature): |
| 3820 | - triton_meta["constants"][numel_name] = int(tree.numel) | 4076 | + # Oversized block counts: a constant >= 2^31 specializes the i64 arg to a |
| 4077 | + # value triton types as uint32 (in [2^31, 2^32)), and | ||
| 4078 | + # the resulting uint32->i64 vcast is rejected by | ||
| 4079 | + # BiShengIR. Keep such numels on their i64 runtime | ||
| 4080 | + # arg; the mask-elimination this enables is | ||
| 4081 | + # irrelevant at >2^31 reduction sizes anyway. | ||
| 4082 | + if int(tree.numel) < 2**31: | ||
| 4083 | + triton_meta["constants"][numel_name] = int(tree.numel) | ||
| 3821 | 4084 | ||
| 3822 | self.triton_meta = triton_meta | 4085 | self.triton_meta = triton_meta |
| 3823 | 4086 | ||
| @@ -3839,11 +4102,22 @@ class NPUTritonKernel(TritonKernel): | |||
| 3839 | # codegen_body() has populated pre_loop_code, so the recipe reads finished | 4102 | # codegen_body() has populated pre_loop_code, so the recipe reads finished |
| 3840 | # block-count lines. A5 only (other chips keep group dispatch). Must run before the | 4103 | # block-count lines. A5 only (other chips keep group dispatch). Must run before the |
| 3841 | # heuristics decorator serializes inductor_meta. | 4104 | # heuristics decorator serializes inductor_meta. |
| 3842 | - if triton_codegen_linearize and device_props.is_a5(): | 4105 | + if self._npu_linearize and device_props.is_a5(): |
| 3843 | _recipe = self._npu_build_grid_recipe() | 4106 | _recipe = self._npu_build_grid_recipe() |
| 3844 | if _recipe is not None: | 4107 | if _recipe is not None: |
| 3845 | inductor_meta["npu_dispatch_recipe"] = _recipe | 4108 | inductor_meta["npu_dispatch_recipe"] = _recipe |
| 3846 | 4109 | ||
| 4110 | + # Linearize flag for the launcher grid: the heuristics grid_0 defaults to | ||
| 4111 | + # NPU_CU_COUNT (persistent 48-core dispatch) for reduction kernels, which | ||
| 4112 | + # is only valid under the linearize structure's group logic. A | ||
| 4113 | + # non-linearize kernel emits xoffset = pid*XBLOCK with an always-true | ||
| 4114 | + # xmask (xnumel % XBLOCK == 0) and one tile per pid, so the grid must be | ||
| 4115 | + # exactly ceil(xnumel/XBLOCK): larger reads past the input (MTE fault | ||
| 4116 | + # 507035, verified at XBLOCK=512/grid=48 on a 8192-xnumel sum), smaller | ||
| 4117 | + # silently drops tiles. The heuristics launches the exact tile count | ||
| 4118 | + # when this flag is False. | ||
| 4119 | + inductor_meta["npu_linearize"] = self._npu_linearize | ||
| 4120 | + | ||
| 3847 | for helper in self.helper_functions: | 4121 | for helper in self.helper_functions: |
| 3848 | code.writeline("") | 4122 | code.writeline("") |
| 3849 | code.splice(helper) | 4123 | code.splice(helper) |
| @@ -3886,7 +4160,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 3886 | self.codegen_static_numels(code) | 4160 | self.codegen_static_numels(code) |
| 3887 | for old, new in self.args.aliases(): | 4161 | for old, new in self.args.aliases(): |
| 3888 | code.writeline(f"{old} = {new}") | 4162 | code.writeline(f"{old} = {new}") |
| 3889 | - if triton_codegen_linearize: | 4163 | + if self._npu_linearize: |
| 3890 | # Emit the intra-core block->core dispatch prologue (pre_loop | 4164 | # Emit the intra-core block->core dispatch prologue (pre_loop |
| 3891 | # hoists, total_blocks, group_size/group_base), then the body in | 4165 | # hoists, total_blocks, group_size/group_base), then the body in |
| 3892 | # the per-core "for i" loop. | 4166 | # the per-core "for i" loop. |
| @@ -4330,7 +4604,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 4330 | return | 4604 | return |
| 4331 | if self.inside_reduction: | 4605 | if self.inside_reduction: |
| 4332 | return | 4606 | return |
| 4333 | - if not (triton_codegen_linearize and getattr(self, "_linearize_applied", False)): | 4607 | + if not (self._npu_linearize and getattr(self, "_linearize_applied", False)): |
| 4334 | return | 4608 | return |
| 4335 | targets = getattr(self, "_npu_select_lane_loads", None) | 4609 | targets = getattr(self, "_npu_select_lane_loads", None) |
| 4336 | if not targets: | 4610 | if not targets: |
| @@ -4555,7 +4829,7 @@ class NPUTritonKernel(TritonKernel): | |||
| 4555 | numel/divisor args in linearize mode. | 4829 | numel/divisor args in linearize mode. |
| 4556 | In 2.7.1 the method is add_numel_to_call_args (no _and_grid suffix). | 4830 | In 2.7.1 the method is add_numel_to_call_args (no _and_grid suffix). |
| 4557 | """ | 4831 | """ |
| 4558 | - if not triton_codegen_linearize: | 4832 | + if not self._npu_linearize: |
| 4559 | # Non-linearize: upstream behavior | 4833 | # Non-linearize: upstream behavior |
| 4560 | for tree in self.range_trees: | 4834 | for tree in self.range_trees: |
| 4561 | expr = tree.numel if isinstance(tree.numel, (sympy.Integer, sympy.Symbol)) else V.graph.wrapper_code.generate_numel_expr(name, tree) # noqa: B950 | 4835 | expr = tree.numel if isinstance(tree.numel, (sympy.Integer, sympy.Symbol)) else V.graph.wrapper_code.generate_numel_expr(name, tree) # noqa: B950 |
| @@ -4970,7 +5244,7 @@ class NPUTritonScheduling(TritonScheduling): | |||
| 4970 | for kernel in kernels: | 5244 | for kernel in kernels: |
| 4971 | self.codegen_node_schedule_with_kernel(node_schedule, kernel) | 5245 | self.codegen_node_schedule_with_kernel(node_schedule, kernel) |
| 4972 | 5246 | ||
| 4973 | - if triton_codegen_linearize: | 5247 | + if kernels and kernels[0]._npu_linearize: |
| 4974 | self._apply_linearize(kernels[0] if len(kernels) == 1 else None, node_schedule) | 5248 | self._apply_linearize(kernels[0] if len(kernels) == 1 else None, node_schedule) |
| 4975 | 5249 | ||
| 4976 | # r-axis cross-core split for OUTER reductions: when the x-axis core | 5250 | # r-axis cross-core split for OUTER reductions: when the x-axis core |
| @@ -5125,8 +5399,16 @@ class NPUTritonScheduling(TritonScheduling): | |||
| 5125 | **partial.inductor_meta_common(), | 5399 | **partial.inductor_meta_common(), |
| 5126 | } | 5400 | } |
| 5127 | 5401 | ||
| 5402 | + # Oversized block counts: a literal >= 2^31 types as uint32 in [2^31, 2^32) and poisons | ||
| 5403 | + # the div/mod dispatch chain with signedness errors; alias the i64 | ||
| 5404 | + # runtime arg instead (x_total_hint is the sole x axis here). | ||
| 5405 | + x0numel_def = ( | ||
| 5406 | + "x0numel = xnumel" | ||
| 5407 | + if int(x_total_hint) >= 2**31 | ||
| 5408 | + else f"x0numel = {int(x_total_hint)}" | ||
| 5409 | + ) | ||
| 5128 | pre_loop_lines = [ | 5410 | pre_loop_lines = [ |
| 5129 | - f"x0numel = {int(x_total_hint)}", | 5411 | + x0numel_def, |
| 5130 | "real_block_x0 = x0numel if x0numel <= XBLOCK else XBLOCK", | 5412 | "real_block_x0 = x0numel if x0numel <= XBLOCK else XBLOCK", |
| 5131 | "x0_blocks = (x0numel + real_block_x0 - 1) // real_block_x0", | 5413 | "x0_blocks = (x0numel + real_block_x0 - 1) // real_block_x0", |
| 5132 | ] | 5414 | ] |
| @@ -6989,7 +7271,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 6989 | """ | 7271 | """ |
| 6990 | Override to track index_vars and var_ranges per node (needed for linearize mode). | 7272 | Override to track index_vars and var_ranges per node (needed for linearize mode). |
| 6991 | """ | 7273 | """ |
| 6992 | - if triton_codegen_linearize: | 7274 | + if kernel._npu_linearize: |
| 6993 | kernel.var_ranges_per_node = [] | 7275 | kernel.var_ranges_per_node = [] |
| 6994 | kernel.index_vars_per_node = [] | 7276 | kernel.index_vars_per_node = [] |
| 6995 | 7277 | ||
| @@ -7025,7 +7307,7 @@ def {combine_name}(in_ptr0, out_ptr0, xnumel, r0_numel, XBLOCK : tl.constexpr, R | |||
| 7025 | indexing_dtype_strength_reduction(node._body) | 7307 | indexing_dtype_strength_reduction(node._body) |
| 7026 | index_vars = kernel.split_and_set_ranges(node.get_ranges()) | 7308 | index_vars = kernel.split_and_set_ranges(node.get_ranges()) |
| 7027 | 7309 | ||
| 7028 | - if triton_codegen_linearize: | 7310 | + if kernel._npu_linearize: |
| 7029 | kernel.var_ranges_per_node.append(node.get_ranges()) | 7311 | kernel.var_ranges_per_node.append(node.get_ranges()) |
| 7030 | kernel.index_vars_per_node.append(index_vars) | 7312 | kernel.index_vars_per_node.append(index_vars) |
| 7031 | 7313 | ||
| @@ -7119,6 +7401,14 @@ def apply_npu_codegen_patches(): | |||
| 7119 | _patch_zero_dim_cpu_tensor_for_npu() | 7401 | _patch_zero_dim_cpu_tensor_for_npu() |
| 7120 | 7402 | ||
| 7121 | # NPU AI Vector Core does not natively support int64 arithmetic — demote to int32. | 7403 | # NPU AI Vector Core does not natively support int64 arithmetic — demote to int32. |
| 7404 | + # Mapping ownership (R7 audit): this entry is process-global and shared | ||
| 7405 | + # with the default backend in mixed processes. Within triton_experimental | ||
| 7406 | + # it is now VESTIGIAL for int64 — dtype_to_str is overridden, and the | ||
| 7407 | + # store/value routes are patched (see npu_triton_store_type / | ||
| 7408 | + # _npu_value_expr); it survives only so the default backend's own | ||
| 7409 | + # consumers keep their behaviour. New direct triton_type(torch.int64) | ||
| 7410 | + # call sites surface through the callsite pin in | ||
| 7411 | + # test_triton_experimental_int32_overflow. | ||
| 7122 | import torch._inductor.utils as _inductor_utils | 7412 | import torch._inductor.utils as _inductor_utils |
| 7123 | _inductor_utils._triton_type_mapping["tl.int64"] = "tl.int32" | 7413 | _inductor_utils._triton_type_mapping["tl.int64"] = "tl.int32" |
| 7124 | # Also patch triton_compute_type which torch_npu overrides with its own | 7414 | # Also patch triton_compute_type which torch_npu overrides with its own |
| @@ -197,6 +197,17 @@ def disable_pointwise_autotuning(inductor_meta): | |||
| 197 | return not inductor_meta.get("autotune_pointwise", True) | 197 | return not inductor_meta.get("autotune_pointwise", True) |
| 198 | 198 | ||
| 199 | 199 | ||
| 200 | +def _is_constexpr_signature_error(e: BaseException) -> bool: | ||
| 201 | + """The ONE ValueError that is a skippable autotune candidate: triton's | ||
| 202 | + ASTSource validates every Config constexpr against the kernel signature | ||
| 203 | + and raises "'<name>' is not in list" for fold candidates whose kernel | ||
| 204 | + body folds inline (no declared constexpr). A blanket ValueError catch | ||
| 205 | + would swallow real codegen bugs, so only this exact shape is | ||
| 206 | + candidate-level; both the serial and the parallel precompile paths must | ||
| 207 | + discriminate identically.""" | ||
| 208 | + return isinstance(e, ValueError) and "is not in list" in str(e) | ||
| 209 | + | ||
| 210 | + | ||
| 200 | def _fmt_config(cfg): | 211 | def _fmt_config(cfg): |
| 201 | try: | 212 | try: |
| 202 | return (f"kwargs={dict(cfg.kwargs)} warps={getattr(cfg, 'num_warps', None)} " | 213 | return (f"kwargs={dict(cfg.kwargs)} warps={getattr(cfg, 'num_warps', None)} " |
| @@ -435,7 +446,12 @@ class NPUTritonCompileResult(TritonCompileResult): | |||
| 435 | self.inductor_meta, def_args, fn.arg_names | 446 | self.inductor_meta, def_args, fn.arg_names |
| 436 | ) | 447 | ) |
| 437 | is_unsplit_scalar_reduction = ( | 448 | is_unsplit_scalar_reduction = ( |
| 438 | - npu_num_x_nodes == 0 | 449 | + # Non-linearize bodies index xoffset = pid*XBLOCK directly (no |
| 450 | + # group-dispatch odometer folding), so a grid of 1 would only run | ||
| 451 | + # the first tile — they must fall through to the exact-grid branch | ||
| 452 | + # below regardless of the x-node count. | ||
| 453 | + self.inductor_meta.get("npu_linearize", True) | ||
| 454 | + and npu_num_x_nodes == 0 | ||
| 439 | and grid_type == "Grid1D" | 455 | and grid_type == "Grid1D" |
| 440 | and "R0_BLOCK" in set(fn.arg_names) | 456 | and "R0_BLOCK" in set(fn.arg_names) |
| 441 | and not npu_rsplit_partial | 457 | and not npu_rsplit_partial |
| @@ -484,6 +500,20 @@ class NPUTritonCompileResult(TritonCompileResult): | |||
| 484 | # max/min run only when xnumel changes, still correct for any value. The cache | 500 | # max/min run only when xnumel changes, still correct for any value. The cache |
| 485 | # cell is bound as a hidden default param (below) so the lookup is LOAD_FAST. | 501 | # cell is bound as a hidden default param (below) so the lookup is LOAD_FAST. |
| 486 | grid_0_is_memoized = True | 502 | grid_0_is_memoized = True |
| 503 | + elif not self.inductor_meta.get("npu_linearize", True) and grid_type == "Grid1D" and "xnumel" in def_args: | ||
| 504 | + # Non-linearize structure: xoffset = pid*XBLOCK with an always-true | ||
| 505 | + # xmask (xnumel % XBLOCK == 0) and no group-dispatch folding, so | ||
| 506 | + # the grid must be EXACTLY ceil(xnumel/XBLOCK) — one tile per pid. | ||
| 507 | + # Larger grids read past the input (MTE fault 507035, vector core | ||
| 508 | + # exception — codegen/triton.py sets inductor_meta["npu_linearize"] | ||
| 509 | + # =False for these kernels); smaller ones silently drop tiles | ||
| 510 | + # (uninitialized output rows, no fault — a min(ceil, NPU_CU_COUNT) | ||
| 511 | + # clamp regressed sum(64,128,256) at XBLOCK=128 where ceil=64 > 48). | ||
| 512 | + # Unlike the simple-1D pointwise branch there is no num_x_nodes | ||
| 513 | + # guarantee that ceil <= NPU_CU_COUNT, so no upper clamp: reductions | ||
| 514 | + # time-slice tiles over the cores instead. | ||
| 515 | + grid_0_expr = f"max(1, (xnumel + {xblock_val} - 1) // {xblock_val})" | ||
| 516 | + grid_0_is_memoized = True | ||
| 487 | else: | 517 | else: |
| 488 | grid_0_expr = str(NPU_CU_COUNT) | 518 | grid_0_expr = str(NPU_CU_COUNT) |
| 489 | grid_0_is_memoized = False | 519 | grid_0_is_memoized = False |
| @@ -743,7 +773,11 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 743 | compile_failed.append((c, e)) | 773 | compile_failed.append((c, e)) |
| 744 | last_exc = e | 774 | last_exc = e |
| 745 | except Exception as e: | 775 | except Exception as e: |
| 746 | - if isinstance(e, OutOfResources): | 776 | + # ValueError with the constexpr-signature shape ("'X' is |
| 777 | + # not in list"): fold candidates whose kernel body folds | ||
| 778 | + # inline must be skipped, not fail the whole compile. Any | ||
| 779 | + # OTHER ValueError is a real codegen bug and must raise. | ||
| 780 | + if isinstance(e, OutOfResources) or _is_constexpr_signature_error(e): | ||
| 747 | log.debug(" [COMPILE FAIL] %s -> %s: %s", _fmt_config(c), type(e).__name__, e) # noqa: G200 | 781 | log.debug(" [COMPILE FAIL] %s -> %s: %s", _fmt_config(c), type(e).__name__, e) # noqa: G200 |
| 748 | compile_failed.append((c, e)) | 782 | compile_failed.append((c, e)) |
| 749 | last_exc = e | 783 | last_exc = e |
| @@ -756,6 +790,13 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 756 | return self._precompile_config(cfg), None | 790 | return self._precompile_config(cfg), None |
| 757 | except (MLIRCompilationError, CompilationError, OutOfResources) as e: | 791 | except (MLIRCompilationError, CompilationError, OutOfResources) as e: |
| 758 | return None, e | 792 | return None, e |
| 793 | + except ValueError as e: | ||
| 794 | + # Same discrimination as the serial loop: only the | ||
| 795 | + # constexpr-signature shape is candidate-level, every other | ||
| 796 | + # ValueError is a real codegen bug and propagates. | ||
| 797 | + if _is_constexpr_signature_error(e): | ||
| 798 | + return None, e | ||
| 799 | + raise | ||
| 759 | 800 | ||
| 760 | max_workers = min(compile_threads, len(configs)) | 801 | max_workers = min(compile_threads, len(configs)) |
| 761 | with ThreadPoolExecutor(max_workers=max_workers) as executor: | 802 | with ThreadPoolExecutor(max_workers=max_workers) as executor: |
| @@ -918,7 +959,16 @@ class NPUCachingAutotuner(CachingAutotuner): | |||
| 918 | compile_meta = copy.deepcopy(self.triton_meta) | 959 | compile_meta = copy.deepcopy(self.triton_meta) |
| 919 | 960 | ||
| 920 | cfg_kwargs = cfg.kwargs | 961 | cfg_kwargs = cfg.kwargs |
| 921 | - compile_meta["constants"].update(cfg_kwargs) | 962 | + # Only kernel-declared constexprs may enter ASTSource constants. A |
| 963 | + # cache-loaded Config (read_best -> _load_cached_autotuning) can carry | ||
| 964 | + # foreign bookkeeping keys inside kwargs; triton's ast_to_ttir validates | ||
| 965 | + # every constants key against the kernel signature and rejects the rest | ||
| 966 | + # ("'X' is not in list"), killing every config at once. Filter at the | ||
| 967 | + # boundary so a dirty Config degrades to its constexpr payload only. | ||
| 968 | + constexpr_names = {self.fn.arg_names[i] for i in self.fn.constexprs} | ||
| 969 | + compile_meta["constants"].update( | ||
| 970 | + {k: v for k, v in cfg_kwargs.items() if k in constexpr_names} | ||
| 971 | + ) | ||
| 922 | for i in self.fn.constexprs: | 972 | for i in self.fn.constexprs: |
| 923 | arg_name = self.fn.arg_names[i] | 973 | arg_name = self.fn.arg_names[i] |
| 924 | if arg_name not in compile_meta["constants"] and arg_name in ( | 974 | if arg_name not in compile_meta["constants"] and arg_name in ( |
这条断言对旧的全量 upcast 同样会通过，无法防住 507034。测试应确认 cast 只落在会溢出的具体地址 addend 上，同时反向禁止完整的 int64 arange/full；尤其应补 >2^31 argmax 和
codegen_linearize=False的覆盖，确保不会重新生成满 int64 tile。