From 5729bbdfc230eaca3fd32bb26926d084068b3f61 Mon Sep 17 00:00:00 2001 From: Bin Bao Date: Tue, 9 Jun 2026 06:02:21 -0700 Subject: [PATCH 1/2] [inductor] Unify threads-per-wave (warp_size) extraction (#181112) Route threads-per-wave through a single source of truth (DeviceProperties.warp_size, from torch.cuda.get_device_properties) and thread it through the config-generation helpers, replacing _NUM_THREADS_PER_WARP, cc_warp_size, and scattered `warp_size or 32` / `torch.version.hip` branches. Behavior change: _num_warps halves num_warps only when warp_size == 64 (AMD CDNA/gfx9), not on every HIP device; AMD RDNA (wave32) now follows the NVIDIA path. Authored with: Claude Pull Request resolved: https://github.com/pytorch/pytorch/pull/181112 Approved by: https://github.com/jansel (cherry picked from commit 521f9d784733cb2f4571804b0f6e4687fb7a91be) --- test/inductor/test_max_autotune.py | 3 +- test/inductor/test_torchinductor.py | 4 +- .../test_torchinductor_dynamic_shapes.py | 10 ++- test/inductor/test_triton_heuristics.py | 71 +++++++++++++++ torch/_inductor/choices.py | 2 +- .../codegen/cuda/device_op_overrides.py | 8 +- torch/_inductor/codegen/triton.py | 3 +- torch/_inductor/runtime/hints.py | 18 +++- torch/_inductor/runtime/triton_compat.py | 11 --- torch/_inductor/runtime/triton_heuristics.py | 89 +++++++++++++------ torch/_inductor/scheduler.py | 2 +- torch/_inductor/template_heuristics/triton.py | 1 + torch/_inductor/utils.py | 8 +- 13 files changed, 175 insertions(+), 55 deletions(-) diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py index 58f1ae439abdf..2e1db50e4ade0 100644 --- a/test/inductor/test_max_autotune.py +++ b/test/inductor/test_max_autotune.py @@ -40,6 +40,7 @@ from torch._inductor.graph import GraphLowering from torch._inductor.ir import Buffer, ChoiceCaller, FixedLayout, FlexibleLayout from torch._inductor.kernel.mm_plus_mm import aten_mm_plus_mm +from torch._inductor.runtime.hints import DeviceProperties from torch._inductor.runtime.triton_heuristics import CachingAutotuner, pointwise from torch._inductor.scheduler import Scheduler from torch._inductor.select_algorithm import ( @@ -184,7 +185,7 @@ def test_max_autotune_includes_max_autotune_pointwise_configs(self): Verify that `max_autotune` includes all pointwise configs from `max_autotune_pointwise` for 1D, 2D, and 3D pointwise kernels. """ - triton_meta = {"device": object()} + triton_meta = {"device": DeviceProperties.create(torch.device(GPU_TYPE, 0))} inductor_meta_common = {"autotune_pointwise": False} for size_hints in ( diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py index 8cfe86ebe7897..3f5b18106b9c0 100644 --- a/test/inductor/test_torchinductor.py +++ b/test/inductor/test_torchinductor.py @@ -17271,7 +17271,7 @@ def test_has_constant_mask(self, block_multiple, ynumel_exceed_ygrid_size): ) @config.patch("triton.persistent_reductions", True) def test_has_constant_mask_small_persistent_reduction(self, rnumel): - from torch._inductor.runtime.hints import DeviceProperties + from torch._inductor.runtime.hints import get_warp_size def fn(x): return x.sum(dim=-1) @@ -17281,7 +17281,7 @@ def fn(x): code = run_and_get_triton_code(opt_fn, x) device = torch.device(GPU_TYPE, 0) - warp_size = DeviceProperties.create(device).warp_size or 32 + warp_size = get_warp_size(device) rblock = 1 while rblock < rnumel: diff --git a/test/inductor/test_torchinductor_dynamic_shapes.py b/test/inductor/test_torchinductor_dynamic_shapes.py index 1144c1f9c4553..8d6adb5844720 100644 --- a/test/inductor/test_torchinductor_dynamic_shapes.py +++ b/test/inductor/test_torchinductor_dynamic_shapes.py @@ -643,6 +643,7 @@ def f(x, w): torch.compile(fullgraph=True)(f)(x, w).sum().backward() self.assertEqual(orig_w, w.grad) + @onlyOn(GPU_TYPE) @torch._dynamo.config.patch( capture_scalar_outputs=True, capture_dynamic_output_shape_ops=True ) @@ -653,6 +654,7 @@ def test_embedding_backward_dynamic_shapes_large_grid(self, device): On CUDA: uses num_blocks only (not num_blocks * num_warps * warp_size). On ROCm: uses num_blocks * num_warps * warp_size (total threads limit). """ + from torch._inductor.runtime.hints import get_warp_size from torch._inductor.runtime.triton_heuristics import ( _check_max_grid_x, _num_warps, @@ -660,13 +662,15 @@ def test_embedding_backward_dynamic_shapes_large_grid(self, device): size_hints = {"x": 600_000_000} x = 64 - num_warps = _num_warps(8) + warp_size = get_warp_size(torch.device(device)) + num_warps = _num_warps(8, warp_size=warp_size) - result_x, result_num_blocks = _check_max_grid_x(size_hints, x, num_warps) + result_x, result_num_blocks = _check_max_grid_x( + size_hints, x, num_warps, warp_size=warp_size + ) max_grid_x = 2147483647 if torch.version.hip: - warp_size = 64 # TODO: query warp size once #129663 is merged # ROCm limits total threads (num_blocks * num_warps * warp_size) self.assertLessEqual( result_num_blocks * num_warps * warp_size, diff --git a/test/inductor/test_triton_heuristics.py b/test/inductor/test_triton_heuristics.py index 5910b68b06419..8a646d0f471bd 100644 --- a/test/inductor/test_triton_heuristics.py +++ b/test/inductor/test_triton_heuristics.py @@ -45,6 +45,8 @@ ) from torch._inductor.runtime.triton_helpers import math as tl_math from torch._inductor.runtime.triton_heuristics import ( + _check_max_grid_x, + _num_warps, autotune_hints_to_configs, CachingAutotuner, template, @@ -209,6 +211,8 @@ def mock_triton_config( num_stages=None, num_elements_per_warp=None, min_elem_per_thread=None, + *, + warp_size=32, ): seen_num_elements_per_warp.add(num_elements_per_warp) return None @@ -759,6 +763,73 @@ def test_grid2d_yz_overflow_large_ynumel(self): self.assertGreaterEqual(y * z, 131070) + +class TestWarpSizeUnification(TestCase): + """Tests for the unified warp_size threading through config helpers.""" + + def test_warp_size_or_default(self): + none_props = DeviceProperties( + type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=None + ) + self.assertEqual(none_props.warp_size_or_default, 32) + + w32 = DeviceProperties( + type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=32 + ) + self.assertEqual(w32.warp_size_or_default, 32) + + w64 = DeviceProperties( + type="hip", index=0, multi_processor_count=80, cc=80, warp_size=64 + ) + self.assertEqual(w64.warp_size_or_default, 64) + + def test_num_warps_halves_on_wave64(self): + # wave64 (AMD CDNA/gfx9) halves the range so total threads match wave32. + self.assertEqual(_num_warps(8, max_num_warps=8, warp_size=64), 4) + + def test_num_warps_does_not_halve_on_wave32(self): + # wave32 (NVIDIA and AMD RDNA) must not halve. + self.assertEqual(_num_warps(8, max_num_warps=8, warp_size=32), 8) + # Default warp_size is 32; confirm the default path matches wave32. + self.assertEqual(_num_warps(8, max_num_warps=8), 8) + + def test_check_max_grid_x_respects_warp_size_hip(self): + # _check_max_grid_x only uses warp_size on the HIP path, where the + # bound is on total threads (num_blocks * num_warps * warp_size). + # Force the HIP branch so the assertion exercises that path even on + # a non-HIP build. + size_hints = {"x": 2**32} + with patch.object(torch.version, "hip", "6.0.0"): + x32, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=32) + x64, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=64) + # Doubling warp_size halves the per-block thread budget, so x must + # double to keep total threads under 2**31 - 1. + self.assertEqual(x64, 2 * x32) + self.assertLessEqual(x32, TRITON_MAX_BLOCK["X"]) + self.assertLessEqual(x64, TRITON_MAX_BLOCK["X"]) + + def test_check_max_grid_x_ignores_warp_size_on_nvidia(self): + # NVIDIA bounds num_blocks directly, so warp_size must not change x. + size_hints = {"x": 2**32} + with patch.object(torch.version, "hip", None): + x32, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=32) + x64, _ = _check_max_grid_x(size_hints, x=1, num_warps=8, warp_size=64) + self.assertEqual(x32, x64) + + def test_triton_config_uses_warp_size_for_min_elem(self): + # min_elem_per_thread * warp_size * num_warps drives the floor for + # block_size; doubling warp_size should at least double the XBLOCK. + size_hints = {"x": 4096} + cfg32 = triton_config( + size_hints, 128, num_warps=1, min_elem_per_thread=4, warp_size=32 + ) + cfg64 = triton_config( + size_hints, 128, num_warps=1, min_elem_per_thread=4, warp_size=64 + ) + self.assertEqual(cfg64.kwargs["XBLOCK"], 2 * cfg32.kwargs["XBLOCK"]) + + + if __name__ == "__main__": if IS_LINUX and HAS_GPU: run_tests() diff --git a/torch/_inductor/choices.py b/torch/_inductor/choices.py index deccd38641cc4..2265ad4befe97 100644 --- a/torch/_inductor/choices.py +++ b/torch/_inductor/choices.py @@ -459,7 +459,7 @@ def reduction_split_factor( so we will do the reduction in two phases.""" props = DeviceProperties.create(device) num_sm = props.multi_processor_count - warp_size = props.warp_size if props.warp_size is not None else 32 + warp_size = props.warp_size_or_default max_threads_per_sm = ( props.max_threads_per_multi_processor if props.max_threads_per_multi_processor is not None diff --git a/torch/_inductor/codegen/cuda/device_op_overrides.py b/torch/_inductor/codegen/cuda/device_op_overrides.py index 8d0f210fd6c25..e39c25e084fe2 100644 --- a/torch/_inductor/codegen/cuda/device_op_overrides.py +++ b/torch/_inductor/codegen/cuda/device_op_overrides.py @@ -2,6 +2,7 @@ import torch +from ...runtime.hints import get_warp_size from ..common import ( DeviceOpOverrides, register_device_op_overrides, @@ -124,10 +125,9 @@ def kernel_driver(self) -> str: """ if torch.version.hip is not None: # Adjusting the warp size to GPU supported wavefront size on AMD GPU - prop = torch.cuda.get_device_properties(torch.cuda.current_device()) - source_codes = source_codes.replace( - "32*numWarps", str(prop.warp_size) + "*numWarps" - ) + device = torch.device("cuda", torch.cuda.current_device()) + warp_size = get_warp_size(device) + source_codes = source_codes.replace("32*numWarps", f"{warp_size}*numWarps") return source_codes def tma_descriptor_helpers(self) -> str: diff --git a/torch/_inductor/codegen/triton.py b/torch/_inductor/codegen/triton.py index 58a71b789f030..7521df27614dd 100644 --- a/torch/_inductor/codegen/triton.py +++ b/torch/_inductor/codegen/triton.py @@ -53,6 +53,7 @@ from ..runtime.hints import ( AutotuneHint, DeviceProperties, + get_warp_size, ReductionHint, TRITON_MAX_BLOCK, TRITON_MAX_RSPLIT, @@ -6173,7 +6174,7 @@ def _has_constant_mask(self, tree: IterationRangesRoot) -> bool: # faults on AMD hardware. Keep the dynamic mask so that all # hardware stays correct. device = V.graph.get_current_device_or_throw() - warp_size = DeviceProperties.create(device).warp_size or 32 + warp_size = get_warp_size(device) if isinstance(max_block, int) and max_block < warp_size: return False elif tree.prefix == "x" and self.no_x_dim: diff --git a/torch/_inductor/runtime/hints.py b/torch/_inductor/runtime/hints.py index af91cdece3bd3..c4e82b7f8dedf 100644 --- a/torch/_inductor/runtime/hints.py +++ b/torch/_inductor/runtime/hints.py @@ -114,9 +114,6 @@ def AttrsDescriptorWrapper( ) -_NUM_THREADS_PER_WARP = 32 - - class HeuristicType(Enum): PERSISTENT_REDUCTION = auto() POINTWISE = auto() @@ -150,6 +147,10 @@ class DeviceProperties(typing.NamedTuple): max_threads_per_block: int | None = None warp_size: int | None = None + @property + def warp_size_or_default(self) -> int: + return self.warp_size if self.warp_size is not None else 32 + @classmethod @functools.cache def create(cls, device) -> DeviceProperties: @@ -187,6 +188,17 @@ def create(cls, device) -> DeviceProperties: ) +def get_warp_size(device) -> int: + """Return the wave/warp size in threads for the given device. + + Reads from torch.cuda.get_device_properties(device).warp_size via the + cached DeviceProperties.create(). Correct on both AMD (64 for CDNA/gfx9, + 32 for RDNA/gfx10+) and NVIDIA (always 32). Falls back to 32 only when + the field is unavailable. + """ + return DeviceProperties.create(device).warp_size_or_default + + class HalideInputSpec(typing.NamedTuple): ctype: str name: str diff --git a/torch/_inductor/runtime/triton_compat.py b/torch/_inductor/runtime/triton_compat.py index dc899a76de5b4..cd3d53f139c03 100644 --- a/torch/_inductor/runtime/triton_compat.py +++ b/torch/_inductor/runtime/triton_compat.py @@ -148,16 +148,6 @@ class JITFunction: # type: ignore[no-redef] HAS_TRITON = False -def cc_warp_size(cc: str | int) -> int: - if torch.version.hip: - if "gfx9" in str(cc): - return 64 - else: - return 32 - else: - return 32 - - try: autograd_profiler = torch.autograd.profiler except AttributeError: # Compile workers only have a mock version of torch @@ -180,7 +170,6 @@ class autograd_profiler: # type: ignore[no-redef] "libdevice", "math", "triton", - "cc_warp_size", "knobs", "triton_key", ] diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py index 1b6dae341aae6..e7dd10b0a99ad 100644 --- a/torch/_inductor/runtime/triton_heuristics.py +++ b/torch/_inductor/runtime/triton_heuristics.py @@ -42,7 +42,6 @@ from .benchmarking import benchmarker from .coordinate_descent_tuner import CoordescTuner from .hints import ( - _NUM_THREADS_PER_WARP, AutotuneHint, DeviceProperties, HeuristicType, @@ -75,7 +74,6 @@ from .triton_compat import ( ASTSource, autograd_profiler, - cc_warp_size, CompiledKernel, Config, GPUTarget, @@ -187,13 +185,13 @@ def autotune_hints_to_configs( (1, block_size // 4, 1), (1, 1, block_size // 4), ) + warp_size = device_props.warp_size_or_default configs.extend( triton_config( size_hints, *xyz, - num_elements_per_warp=( - device_props.warp_size if device_props.warp_size else 32 - ), + num_elements_per_warp=warp_size, + warp_size=warp_size, ) for xyz in xyz_options ) @@ -561,7 +559,7 @@ def _dynamic_scale_rblock(self): assert device_prop.max_threads_per_multi_processor assert device_prop.multi_processor_count seen_config_hashes: OrderedSet[Hashable] | None = None - warp_size = device_prop.warp_size or 32 + warp_size = device_prop.warp_size_or_default for result in self.compile_results: triton_config = result.config compiled_binary = result.kernel @@ -878,7 +876,7 @@ def _precompile_config(self, cfg: Config) -> CompileResult[_KernelType]: target = GPUTarget( compile_meta["device_type"], arch, - cc_warp_size(compile_meta["cc"]), + self.device_props.warp_size_or_default, ) options = self._create_compile_options(cfg, compile_meta) @@ -2614,10 +2612,19 @@ def check_max_block(cfg: dict[str, int]): ) -def _num_warps(num_warps, max_num_warps=8, min_num_warps=2, register_intensive=False): - # On AMD GPU each warp has 64 lanes which is double the size on NV GPU, - # therefore using half the number of warps here correspondingly. - if torch.version.hip: +def _num_warps( + num_warps, + max_num_warps=8, + min_num_warps=2, + register_intensive=False, + *, + warp_size: int = 32, +): + # On wave64 AMD GPUs (CDNA / gfx9) each warp has 64 lanes, double NVIDIA + # and RDNA, so use half the number of warps to keep total thread count + # comparable. RDNA (wave32) follows the NVIDIA path. + if warp_size == 64: + max_num_warps = (max_num_warps + 1) // 2 min_num_warps = (min_num_warps + 1) // 2 # persistent reduction is register intensive @@ -2626,13 +2633,10 @@ def _num_warps(num_warps, max_num_warps=8, min_num_warps=2, register_intensive=F return next_power_of_2(min(max(num_warps, min_num_warps), max_num_warps)) -def _check_max_grid_x(size_hints, x, num_warps): +def _check_max_grid_x(size_hints, x, num_warps, *, warp_size: int = 32): # Check if maxGridSize is exceeded - if so then must scale XBLOCK further max_grid_x = 2147483647 max_block_x = TRITON_MAX_BLOCK["X"] - warp_size = ( - 64 if torch.version.hip else 32 - ) # TODO: query warp size once #129663 is merged num_blocks = (size_hints["x"] + x - 1) // x if torch.version.hip: @@ -2669,6 +2673,8 @@ def triton_config( matrix_instr=None, waves_per_eu=None, kpack=None, + *, + warp_size: int = 32, ) -> Config: """ Construct a pointwise triton config with some adjustment heuristics @@ -2728,7 +2734,9 @@ def triton_config( # Calculate num_warps if they are not hard passed to config if num_warps is None: num_warps = _num_warps( - conditional_product(x, y, z) // num_elements_per_warp, min_num_warps=1 + conditional_product(x, y, z) // num_elements_per_warp, + min_num_warps=1, + warp_size=warp_size, ) # we are going to arrive at 2 warps only if bs was too small due to # numel being too small. However to workaround some ptx bugs we still @@ -2745,11 +2753,11 @@ def triton_config( # Increase x to satisfy min_elem_per_thread requirements. block_size = max( conditional_product(x, y, z), - min_elem_per_thread * _NUM_THREADS_PER_WARP * num_warps, + min_elem_per_thread * warp_size * num_warps, ) x *= math.ceil(block_size / conditional_product(x, y, z)) - x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps) + x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size) x = min(x, size_hints["x"]) cfg = {"XBLOCK": x} @@ -2821,6 +2829,8 @@ def triton_config_reduction( dynamic_scale_rblock=True, reduction_hint=None, min_num_warps=None, + *, + warp_size: int = 32, ) -> Config: """ Construct a reduction triton config with some adjustment heuristics @@ -2862,10 +2872,13 @@ def total_numel() -> int: _num_warps_func = _num_warps num_warps = _num_warps_func( - num_warps, max_num_warps=max_num_warps, register_intensive=register_intensive + num_warps, + max_num_warps=max_num_warps, + register_intensive=register_intensive, + warp_size=warp_size, ) - x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps) + x, _num_blocks = _check_max_grid_x(size_hints, x, num_warps, warp_size=warp_size) for prefix in sorted(rnumels): while total_numel() > target: @@ -3023,7 +3036,15 @@ def _handle_combo_kernel_per_subkernel_blocks( def triton_config_tiled_reduction( - size_hints, x, y, r, num_stages=1, register_intensive=False, waves_per_eu=None + size_hints, + x, + y, + r, + num_stages=1, + register_intensive=False, + waves_per_eu=None, + *, + warp_size: int = 32, ): """ Construct a tile reduction triton config with some adjustment @@ -3054,9 +3075,12 @@ def total_numel() -> int: y *= 2 cfg = _get_config({"x": x, "y": y, **rnumels}) - num_warps = _num_warps(total_numel() // 256, min_num_warps=1) + num_warps = _num_warps(total_numel() // 256, min_num_warps=1, warp_size=warp_size) num_warps = _num_warps( - num_warps, max_num_warps=16, register_intensive=register_intensive + num_warps, + max_num_warps=16, + register_intensive=register_intensive, + warp_size=warp_size, ) check_config(cfg, xnumel=size_hints["x"], ynumel=size_hints["y"]) check_max_block(cfg) @@ -3156,8 +3180,9 @@ def pointwise( triton_meta["device"], ) + warp_size = triton_meta["device"].warp_size_or_default triton_config_with_settings = functools.partial( - triton_config, min_elem_per_thread=min_elem_per_thread + triton_config, min_elem_per_thread=min_elem_per_thread, warp_size=warp_size ) configs = None @@ -3389,6 +3414,7 @@ def _reduction_configs( ) device_major = triton_meta["device"].major + warp_size = triton_meta["device"].warp_size_or_default # Prefer smaller MAX_R0_BLOCK for Blackwell MAX_R0_BLOCK = 1024 if device_major is not None and device_major >= 10 else 2048 if size_hints["x"] >= 1024 and loads_and_red >= 10: @@ -3442,6 +3468,7 @@ def make_config( num_stages=num_stages, register_intensive=register_intensive, waves_per_eu=waves_per_eu, + warp_size=warp_size, ) else: # For other cases, use the original function @@ -3455,6 +3482,7 @@ def make_config( waves_per_eu=waves_per_eu, dynamic_scale_rblock=dynamic_scale_rblock, reduction_hint=reduction_hint, + warp_size=warp_size, ) def outer_config_opt(): @@ -3653,6 +3681,8 @@ def adapt_config_for_tiling( register_intensive=False, persistent_reduction=False, waves_per_eu=None, + *, + warp_size: int = 32, ) -> Config: """ Create an adapted configuration based on tiling scores, @@ -3672,6 +3702,7 @@ def adapt_config_for_tiling( num_stages=num_stages, register_intensive=register_intensive, waves_per_eu=waves_per_eu, + warp_size=warp_size, ) @@ -3906,6 +3937,7 @@ def _persistent_reduction_configs( rnumel = get_total_reduction_numel(size_hints) MAX_PERSISTENT_BLOCK_NUMEL = 4096 + warp_size = triton_meta["device"].warp_size_or_default if triton_meta.get("native_matmul"): if len(size_hints) == 3: @@ -3938,6 +3970,7 @@ def _persistent_reduction_configs( rnumel, register_intensive=True, reduction_hint=reduction_hint, + warp_size=warp_size, ) for xblock in xblock_vals if xblock == 1 @@ -3956,7 +3989,11 @@ def _persistent_reduction_configs( ) configs.append( triton_config_tiled_reduction( - size_hints, block_sizes["x"], block_sizes["y"], rnumel + size_hints, + block_sizes["x"], + block_sizes["y"], + rnumel, + warp_size=warp_size, ) ) @@ -3965,6 +4002,7 @@ def _persistent_reduction_configs( size_hints, 2 * (256 // rnumel) if rnumel <= 256 else 1, rnumel, + warp_size=warp_size, ) ] @@ -4000,6 +4038,7 @@ def _persistent_reduction_configs( num_warps=num_warps, min_num_warps=min_num_warps, reduction_hint=reduction_hint, + warp_size=warp_size, ) ] diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py index 590cf68a481da..0b0c43b3af311 100644 --- a/torch/_inductor/scheduler.py +++ b/torch/_inductor/scheduler.py @@ -2908,7 +2908,7 @@ def _occupancy_before_and_after_fusion( return 1, 1 # Can't calculate, allow fusion assert num_warps - threads_per_block = num_warps * (device_props.warp_size or 32) + threads_per_block = num_warps * device_props.warp_size_or_default regs_per_block_unfused = unfused_n_regs * threads_per_block regs_per_block_fused = fused_n_regs * threads_per_block diff --git a/torch/_inductor/template_heuristics/triton.py b/torch/_inductor/template_heuristics/triton.py index 641c8568fb298..987800dea4f56 100644 --- a/torch/_inductor/template_heuristics/triton.py +++ b/torch/_inductor/template_heuristics/triton.py @@ -30,6 +30,7 @@ ) from ..kernel.mm_plus_mm import mm_plus_mm_template from ..kernel_inputs import KernelInputs, MMKernelInputs +from ..runtime.hints import DeviceProperties from ..utils import ( get_backend_num_stages, get_default_kpack, diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py index 7dc7476f2c9ad..ae25b3355540a 100644 --- a/torch/_inductor/utils.py +++ b/torch/_inductor/utils.py @@ -2944,9 +2944,11 @@ def get_gpu_shared_memory() -> int: def get_max_numwarps() -> int: if torch.cuda.is_available(): - warp_size = torch.cuda.get_device_properties().warp_size - # pyrefly: ignore [missing-attribute] - max_threads_per_block = torch.cuda.get_device_properties().max_threads_per_block + device = torch.device("cuda", torch.cuda.current_device()) + props = DeviceProperties.create(device) + warp_size = props.warp_size_or_default + max_threads_per_block = props.max_threads_per_block + assert max_threads_per_block is not None else: # Defaults warp_size = 32 From 0cc2280e44da48851d2b7aaf85d9e87fa2a3c796 Mon Sep 17 00:00:00 2001 From: "Nichols A. Romero" Date: Mon, 15 Jun 2026 14:34:04 +0000 Subject: [PATCH 2/2] [ROCm][Inductor] Make missing GPU warp-size metadata explicit (#183014) Upstream PR #181112 (`[inductor] Unify threads-per-wave (warp_size) extraction`) has now landed and did much of the original broad warp-size plumbing for this PR: it centralized `DeviceProperties.warp_size`, threaded `warp_size` through Inductor's Triton config helpers, and made HIP wave32 devices follow the non-wave64 path. After rebasing on top of that work, this PR is now focused on the remaining stricter ROCm-safe behavior: - Preserve missing `DeviceProperties.warp_size` as `None` instead of synthesizing `32` for CUDA/HIP devices. - Make `DeviceProperties.warp_size_or_default` raise for missing CUDA/HIP warp-size metadata, while preserving the explicit non-GPU fallback. - Skip optional autotune/rblock/fusion/split-reduction heuristics when required warp-size metadata is unavailable instead of guessing. - Keep CUDA launcher generation on the fast `32` path, but use the device warp size for HIP launcher block dimensions. - Use queried device warp size in template register-spill pruning, with graceful fallback on non-CUDA/XPU builds. - Update the relevant Inductor tests for the stricter missing CUDA/HIP metadata behavior. - Rebased cleaned stack onto `upstream/main` at `19791183fec`. - Rebuilt PyTorch from scratch with `/home/niromero/docker_workspace/framework_scripts/pytorch/build.sh`. - Ran impacted UTs: - `python test/inductor/test_triton_heuristics.py TestWarpSizeUnification TestTritonHeuristics.test_autotune_hints_to_configs TestFastLauncherDeviceSupport` - `python test/inductor/test_max_autotune.py TestMaxAutotune.test_max_autotune_includes_max_autotune_pointwise_configs` - `python test/inductor/test_torchinductor_dynamic_shapes.py -k test_embedding_backward_dynamic_shapes_large_grid` - `python test/inductor/test_torchinductor.py -k test_reduction_config_limit` Made with [Cursor](https://cursor.com) Pull Request resolved: https://github.com/pytorch/pytorch/pull/183014 Approved by: https://github.com/jansel, https://github.com/glen-amd, https://github.com/jeffdaily, https://github.com/eellison (cherry picked from commit 9183299236b498a69e999def9667a4c477c112fe) --- test/inductor/test_max_autotune.py | 13 ++++++++++++- test/inductor/test_torchinductor.py | 4 +++- test/inductor/test_triton_heuristics.py | 10 +++++++++- .../codegen/cuda/device_op_overrides.py | 14 +++++++++----- torch/_inductor/runtime/hints.py | 16 ++++++++++------ torch/_inductor/runtime/triton_heuristics.py | 18 +++++++++++++----- torch/_inductor/scheduler.py | 8 +++++--- torch/_inductor/template_heuristics/triton.py | 15 +++++++++++++-- torch/_inductor/utils.py | 12 ++++++------ 9 files changed, 80 insertions(+), 30 deletions(-) diff --git a/test/inductor/test_max_autotune.py b/test/inductor/test_max_autotune.py index 2e1db50e4ade0..dbcd721227a13 100644 --- a/test/inductor/test_max_autotune.py +++ b/test/inductor/test_max_autotune.py @@ -185,7 +185,18 @@ def test_max_autotune_includes_max_autotune_pointwise_configs(self): Verify that `max_autotune` includes all pointwise configs from `max_autotune_pointwise` for 1D, 2D, and 3D pointwise kernels. """ - triton_meta = {"device": DeviceProperties.create(torch.device(GPU_TYPE, 0))} + # Fake device properties for this unit test; no CUDA state is needed. + triton_meta = { + "device": DeviceProperties( + type="cuda", + index=0, + multi_processor_count=1, + cc=80, + major=8, + max_threads_per_block=1024, + warp_size=32, + ) + } inductor_meta_common = {"autotune_pointwise": False} for size_hints in ( diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py index 3f5b18106b9c0..8ba98bba4001c 100644 --- a/test/inductor/test_torchinductor.py +++ b/test/inductor/test_torchinductor.py @@ -2212,13 +2212,15 @@ def test_reduction_config_limit(self): triton reduction configs for large size hints. #128826 introduced a scaling XBLOCK feature to resolve the issue in reduction configs which may exceed the maxGridSize """ + from torch._inductor.runtime.hints import get_warp_size from torch._inductor.runtime.runtime_utils import next_power_of_2 from torch._inductor.runtime.triton_heuristics import triton_config_reduction + warp_size = get_warp_size(torch.device(self.device)) size_hints = {"x": 67108864, "r0_": 8192} for _ in range(4): size_hints["x"] = next_power_of_2(size_hints["x"]) - triton_config_reduction(size_hints, 1, 2048, 1, 8) + triton_config_reduction(size_hints, 1, 2048, 1, 8, warp_size=warp_size) def test_prod(self): def fn(a): diff --git a/test/inductor/test_triton_heuristics.py b/test/inductor/test_triton_heuristics.py index 8a646d0f471bd..20617345d8e60 100644 --- a/test/inductor/test_triton_heuristics.py +++ b/test/inductor/test_triton_heuristics.py @@ -771,7 +771,15 @@ def test_warp_size_or_default(self): none_props = DeviceProperties( type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=None ) - self.assertEqual(none_props.warp_size_or_default, 32) + with self.assertRaisesRegex( + RuntimeError, "cuda device properties must report warp_size" + ): + none_props.warp_size_or_default + + cpu_props = DeviceProperties( + type="cpu", index=0, multi_processor_count=80, cc=80, warp_size=None + ) + self.assertEqual(cpu_props.warp_size_or_default, 32) w32 = DeviceProperties( type="cuda", index=0, multi_processor_count=80, cc=80, warp_size=32 diff --git a/torch/_inductor/codegen/cuda/device_op_overrides.py b/torch/_inductor/codegen/cuda/device_op_overrides.py index e39c25e084fe2..dde9b0924cc8c 100644 --- a/torch/_inductor/codegen/cuda/device_op_overrides.py +++ b/torch/_inductor/codegen/cuda/device_op_overrides.py @@ -51,6 +51,10 @@ def kernel_header(self) -> str: return source_codes def kernel_driver(self) -> str: + """Return C++ host-side helpers (loadKernel, launchKernel, CUDA_DRIVER_CHECK) + embedded in AOTI-generated wrapper code.""" + # NVIDIA devices have a warp size of 32, while AMD devices can have a + # warp size of 32 or 64 depending on the architecture. source_codes = """ #define CUDA_DRIVER_CHECK(EXPR) \\ do { \\ @@ -119,16 +123,16 @@ def kernel_driver(self) -> str: void* args[], cudaStream_t stream) { CUDA_DRIVER_CHECK(cuLaunchKernel( - func, gridX, gridY, gridZ, 32*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr + func, gridX, gridY, gridZ, __WARP_SIZE__*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr )); } """ - if torch.version.hip is not None: - # Adjusting the warp size to GPU supported wavefront size on AMD GPU + if torch.version.hip is not None and torch.cuda.is_available(): device = torch.device("cuda", torch.cuda.current_device()) warp_size = get_warp_size(device) - source_codes = source_codes.replace("32*numWarps", f"{warp_size}*numWarps") - return source_codes + else: + warp_size = 32 + return source_codes.replace("__WARP_SIZE__", str(warp_size)) def tma_descriptor_helpers(self) -> str: """ diff --git a/torch/_inductor/runtime/hints.py b/torch/_inductor/runtime/hints.py index c4e82b7f8dedf..1552a0764aa08 100644 --- a/torch/_inductor/runtime/hints.py +++ b/torch/_inductor/runtime/hints.py @@ -149,7 +149,11 @@ class DeviceProperties(typing.NamedTuple): @property def warp_size_or_default(self) -> int: - return self.warp_size if self.warp_size is not None else 32 + if self.warp_size is not None: + return self.warp_size + if self.type in ("cuda", "hip"): + raise RuntimeError(f"{self.type} device properties must report warp_size") + return 32 @classmethod @functools.cache @@ -184,17 +188,17 @@ def create(cls, device) -> DeviceProperties: props, "max_threads_per_multi_processor", None ), max_threads_per_block=getattr(props, "max_threads_per_block", 1024), - warp_size=getattr(props, "warp_size", 32 if device_type != "cpu" else None), + warp_size=getattr(props, "warp_size", None), ) def get_warp_size(device) -> int: """Return the wave/warp size in threads for the given device. - Reads from torch.cuda.get_device_properties(device).warp_size via the - cached DeviceProperties.create(). Correct on both AMD (64 for CDNA/gfx9, - 32 for RDNA/gfx10+) and NVIDIA (always 32). Falls back to 32 only when - the field is unavailable. + Reads from torch.cuda.get_device_properties(device).warp_size via the cached + DeviceProperties.create(). Correct on both AMD (64 for CDNA/gfx9, 32 for + RDNA/gfx10+) and NVIDIA (always 32). Missing cuda/hip warp_size metadata is + treated as an error rather than silently falling back. """ return DeviceProperties.create(device).warp_size_or_default diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py index e7dd10b0a99ad..94084382be1e1 100644 --- a/torch/_inductor/runtime/triton_heuristics.py +++ b/torch/_inductor/runtime/triton_heuristics.py @@ -175,6 +175,13 @@ def autotune_hints_to_configs( configs: list[Config] = [] for hint in hints: if hint == AutotuneHint.ONE_ELEMENT_PER_THREAD: + if device_props.warp_size is None: + log.debug( + "Skipping %s autotune hint because device %s does not report warp_size", + AutotuneHint.ONE_ELEMENT_PER_THREAD, + device_props.type, + ) + continue if len(size_hints) == 1: xyz_options = ((block_size // 4, None, None),) elif len(size_hints) == 2: @@ -554,12 +561,13 @@ def _dynamic_scale_rblock(self): and device_prop.major and (device_prop.major >= 8 or torch.version.hip) and device_prop.regs_per_multiprocessor is not None + and device_prop.warp_size is not None ): assert device_prop.regs_per_multiprocessor assert device_prop.max_threads_per_multi_processor assert device_prop.multi_processor_count seen_config_hashes: OrderedSet[Hashable] | None = None - warp_size = device_prop.warp_size_or_default + warp_size = device_prop.warp_size for result in self.compile_results: triton_config = result.config compiled_binary = result.kernel @@ -598,14 +606,14 @@ def _dynamic_scale_rblock(self): nreg_per_warp = nreg * warp_size nreg_per_block = nreg_per_warp * triton_config.num_warps - # Previously we set max_blocks_per_sm to 'max_threads_per_multi_processo / (32 * num_warps)' + # Previously we set max_blocks_per_sm to 'max_threads_per_multi_processor / (warp_size * num_warps)' # The formula below is a tighter upper bound since we have the assumption that # nreg > device_prop.regs_per_multiprocessor // device_prop.max_threads_per_multi_processor # due to the if condition above and: # regs_per_multiprocessor / nreg_per_block - # = regs_per_multiprocessor / (nreg * 32 * num_warps) - # < regs_per_multiprocessor / ((regs_per_multiprocessor / max_threads_per_multi_processor) * 32 * num_warps) - # = max_threads_per_multi_processor / (32 * num_warps) + # = regs_per_multiprocessor / (nreg * warp_size * num_warps) + # < regs_per_multiprocessor / ((regs_per_multiprocessor / max_threads_per_multi_processor) * warp_size * num_warps) + # = max_threads_per_multi_processor / (warp_size * num_warps) # Using a tighter upper bound can reveal more optimization opportunities. max_blocks_per_sm = max( device_prop.regs_per_multiprocessor // nreg_per_block, 1 diff --git a/torch/_inductor/scheduler.py b/torch/_inductor/scheduler.py index 0b0c43b3af311..1a59e9214ae7d 100644 --- a/torch/_inductor/scheduler.py +++ b/torch/_inductor/scheduler.py @@ -2904,11 +2904,13 @@ def _occupancy_before_and_after_fusion( # # Need device info to calculate occupancy regs_per_sm = device_props.regs_per_multiprocessor - if regs_per_sm is None: + warp_size = device_props.warp_size + if regs_per_sm is None or warp_size is None: return 1, 1 # Can't calculate, allow fusion - assert num_warps - threads_per_block = num_warps * device_props.warp_size_or_default + if not num_warps: + raise AssertionError("expected num_warps to be truthy") + threads_per_block = num_warps * warp_size regs_per_block_unfused = unfused_n_regs * threads_per_block regs_per_block_fused = fused_n_regs * threads_per_block diff --git a/torch/_inductor/template_heuristics/triton.py b/torch/_inductor/template_heuristics/triton.py index 987800dea4f56..ae0998c7b5759 100644 --- a/torch/_inductor/template_heuristics/triton.py +++ b/torch/_inductor/template_heuristics/triton.py @@ -1020,11 +1020,22 @@ def _prune_reg_spill_configs( self, configs: list[BaseConfig], ) -> list[BaseConfig]: + try: + device = torch.cuda.current_device() + props = torch.cuda.get_device_properties(device) + warp_size = props.warp_size + except (RuntimeError, AttributeError, AssertionError): + return configs + pruned_configs = [] + # NVIDIA-oriented approximation for the per-thread register limit. + # TODO: generalize this threshold for ROCm devices. + NUM_REG = 255 for gemm_config in configs: - NUM_REG = 255 acc_regs = math.ceil( - gemm_config.block_m * gemm_config.block_n / (gemm_config.num_warps * 32) + gemm_config.block_m + * gemm_config.block_n + / (gemm_config.num_warps * warp_size) ) # Lower bound for register spillage, if exceeds the kernel will certainly spill if acc_regs > NUM_REG: diff --git a/torch/_inductor/utils.py b/torch/_inductor/utils.py index ae25b3355540a..b613819215670 100644 --- a/torch/_inductor/utils.py +++ b/torch/_inductor/utils.py @@ -2948,12 +2948,12 @@ def get_max_numwarps() -> int: props = DeviceProperties.create(device) warp_size = props.warp_size_or_default max_threads_per_block = props.max_threads_per_block - assert max_threads_per_block is not None - else: - # Defaults - warp_size = 32 - max_threads_per_block = 1024 - return max_threads_per_block // warp_size + if max_threads_per_block is None: + raise AssertionError("expected max_threads_per_block to be set") + return max_threads_per_block // warp_size + + log.debug("CUDA is not available; defaulting max num warps to 32") + return 32 def is_welford_reduction(reduction_type: str) -> bool: