[release/2.12] Backport Inductor warp-size handling

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,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": object()}
+        # 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 (

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):
@@ -17271,7 +17273,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 +17283,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:

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,20 +654,23 @@ 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,
         )
 
         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,

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,81 @@ 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
+        )
+        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
+        )
+        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()

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

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,
@@ -50,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 {                                               \\
@@ -118,17 +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
-            prop = torch.cuda.get_device_properties(torch.cuda.current_device())
-            source_codes = source_codes.replace(
-                "32*numWarps", str(prop.warp_size) + "*numWarps"
-            )
-        return source_codes
+        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)
+        else:
+            warp_size = 32
+        return source_codes.replace("__WARP_SIZE__", str(warp_size))
 
     def tma_descriptor_helpers(self) -> str:
         """

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:

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,14 @@ class DeviceProperties(typing.NamedTuple):
     max_threads_per_block: int | None = None
     warp_size: int | None = None
 
+    @property
+    def warp_size_or_default(self) -> int:
+        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
     def create(cls, device) -> DeviceProperties:
@@ -183,10 +188,21 @@ 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). Missing cuda/hip warp_size metadata is
+    treated as an error rather than silently falling back.
+    """
+    return DeviceProperties.create(device).warp_size_or_default
+
+
 class HalideInputSpec(typing.NamedTuple):
     ctype: str
     name: str

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",
 ]