From 69fbaee9d56e8c60fa591f994eb06d58242b8cc6 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 13 Nov 2026 23:58:34 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/BUCK | 1 + backends/cuda/tests/targets.bzl | 24 + .../cuda/tests/test_conv_implicit_gemm.py | 323 +++++++++++ backends/cuda/triton/kernels/conv/conv.py | 11 +- .../cuda/triton/kernels/conv/core/common.py | 21 +- .../cuda/triton/kernels/conv/core/direct.py | 9 +- .../triton/kernels/conv/core/implicit_gemm.py | 510 ++++++++++++++++++ 7 files changed, 890 insertions(+), 9 deletions(-) create mode 100644 backends/cuda/tests/test_conv_implicit_gemm.py create mode 100644 backends/cuda/triton/kernels/conv/core/implicit_gemm.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 1da2346bc15..7a895f691f0 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -174,6 +174,7 @@ fbcode_target( "triton/kernels/conv/core/__init__.py", "triton/kernels/conv/core/common.py", "triton/kernels/conv/core/direct.py", + "triton/kernels/conv/core/implicit_gemm.py", "triton/kernels/conv/routing.py", "triton/kernels/fused_moe.py", "triton/kernels/int4_quantized_gemm.py", diff --git a/backends/cuda/tests/targets.bzl b/backends/cuda/tests/targets.bzl index 163b610bae3..f447e7163cd 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -121,6 +121,30 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_conv_implicit_gemm", + srcs = [ + "test_conv_implicit_gemm.py", + ], + visibility = [ + "//executorch/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/cuda:autotune", + "//executorch/backends/cuda:cuda_backend", + "//executorch/backends/cuda:cuda_partitioner", + "//executorch/backends/cuda:triton_kernels", + "//executorch/exir:lib", + "//executorch/export:export_utils", + ], + keep_gpu_sections = True, + remote_execution = re_test_utils.remote_execution( + platform = "gpu-remote-execution", + subplatform = "A100", + ), + ) + runtime.python_library( name = "autotune_test_utils", srcs = [ diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py new file mode 100644 index 00000000000..0aac36316db --- /dev/null +++ b/backends/cuda/tests/test_conv_implicit_gemm.py @@ -0,0 +1,323 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""triton::convolution on dense convolutions: every schedule and K order of +the implicit-GEMM template against an FP64 reference.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch +import torch.nn.functional as F + +_TOLERANCE = {torch.float32: 2e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2} +_TF32_TOLERANCE = 5e-3 + + +def _reference(x, w, b, stride, before, after): + pad = [p for pair in zip(reversed(before), reversed(after)) for p in pair] + conv = F.conv2d if x.dim() == 4 else F.conv3d + bias = None if b is None else b.double() + return conv(F.pad(x.double(), pad), w.double(), bias, stride) + + +def _channels_last(t): + return t.to( + memory_format=torch.channels_last if t.dim() == 4 else torch.channels_last_3d + ) + + +class ImplicitGemmScheduleTest(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available() or torch.version.hip is not None: + self.skipTest("CUDA required") + torch.manual_seed(0) + self.tensor_cores = torch.cuda.get_device_capability() >= (8, 0) + + def convolution(self, x, w, b, stride, before, after): + from executorch.backends.cuda.triton.kernels.conv.conv import convolution + + rank = x.dim() - 2 + return convolution( + x, + w, + b, + list(stride), + list(before), + list(after), + [1] * rank, + False, + [0] * rank, + 1, + ) + + def check(self, x, w, b, stride, before, after, every_config=False): + """Runs every schedule and K order on the call, with the first config of + each or every config, and checks the result and the layout; the + tensor-core schedule runs FP32 operands in TF32.""" + from executorch.backends.cuda.triton.kernels.conv.conv import ( + force_conv_candidate, + ) + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + out_shape_and_stride, + ) + + rank = x.dim() - 2 + params = ConvParams( + tuple(stride), + tuple(before), + tuple(after), + (1,) * rank, + False, + (0,) * rank, + 1, + ) + call = ConvCall(x, w, b, params) + tf32 = mock.patch.object(implicit_gemm, "conv_allows_tf32", return_value=True) + self.assertIsNone(implicit_gemm.supports(call)) + expected = _reference(x, w, b, stride, before, after) + layout = out_shape_and_stride(call) + schedules = ["igemm_ffma"] + (["igemm_tc"] if self.tensor_cores else []) + for (schedule, _), configs in implicit_gemm.CONFIGS.items(): + if schedule not in schedules: + continue + with tf32: + self.assertIsNone(implicit_gemm.schedule_supports(schedule, call)) + tolerance = _TOLERANCE[x.dtype] + if schedule == "igemm_tc" and x.dtype == torch.float32: + tolerance = _TF32_TOLERANCE + for config in configs if every_config else configs[:1]: + with tf32, force_conv_candidate(schedule, config): + out = self.convolution(x, w, b, stride, before, after) + self.assertEqual((tuple(out.shape), out.stride()), layout) + torch.testing.assert_close( + out.double(), + expected.to(x.dtype).double(), + atol=tolerance * max(1.0, expected.abs().max().item()), + rtol=tolerance, + equal_nan=True, + msg=lambda m, s=schedule, c=config: f"{s} {c}: {m}", + ) + + def tensors(self, shape, out_channels, kernel, dtype=torch.float32, bias=True): + channels = shape[1] + x = torch.randn(shape, device="cuda", dtype=dtype) + fan_in = channels * torch.Size(kernel).numel() + w = torch.randn(out_channels, channels, *kernel, device="cuda", dtype=dtype) / ( + fan_in**0.5 + ) + b = torch.randn(out_channels, device="cuda", dtype=dtype) if bias else None + return x, w, b + + def test_every_config_of_every_schedule(self): + for shape, out_channels, kernel, stride, padding in ( + ((2, 32, 13, 11), 40, (3, 3), (1, 1), (1, 1)), + ((2, 3, 19, 17), 24, (7, 7), (2, 2), (3, 3)), + ): + x, w, b = self.tensors(shape, out_channels, kernel) + with self.subTest(shape=shape): + self.check(x, w, b, stride, padding, padding, every_config=True) + + def check_geometries(self, dtype, layouts=("contiguous", "channels_last")): + """Every schedule on each geometry in ``dtype`` and each of ``layouts``: + a contiguous or channels-last input, or a channels-last weight.""" + geometries = ( + # input, C_out, kernel, stride, before, after: representative + # geometries at small sizes, then edge cases. + ((6, 3, 2, 14, 14), 40, (2, 14, 14), (2, 14, 14), (0,) * 3, (0,) * 3), + ((1, 24, 3, 6, 10), 33, (1, 2, 2), (1, 2, 2), (0,) * 3, (0,) * 3), + ((1, 3, 56, 42), 72, (14, 14), (14, 14), (0, 0), (0, 0)), + ((2, 3, 32, 48), 48, (16, 16), (16, 16), (0, 0), (0, 0)), + ((3, 3, 29, 31), 40, (7, 7), (4, 4), (3, 3), (3, 3)), + ((2, 3, 30, 26), 64, (7, 7), (2, 2), (3, 3), (3, 3)), + ((2, 32, 11, 13), 32, (3, 3), (1, 1), (1, 1), (1, 1)), + ((1, 48, 12, 9), 64, (3, 3), (2, 2), (1, 1), (1, 1)), + ((2, 40, 7, 9), 24, (1, 1), (1, 1), (0, 0), (0, 0)), + ((2, 40, 7, 9), 24, (1, 1), (2, 2), (0, 0), (0, 0)), + ((1, 3, 41, 39), 32, (7, 7), (2, 2), (0, 0), (0, 0)), + ((2, 3, 46, 36), 24, (7, 7), (1, 1), (0, 0), (0, 0)), + ((1, 32, 39, 39), 32, (5, 5), (2, 2), (0, 0), (0, 0)), + ((1, 32, 8, 8), 32, (3, 3), (2, 2), (0, 0), (0, 0)), + ((1, 5, 10, 12), 7, (3, 5), (2, 1), (0, 2), (1, 0)), + ((1, 20, 9, 7), 3, (2, 4), (1, 3), (2, 0), (0, 5)), + ((1, 17, 4, 9, 8), 18, (3, 3, 3), (1, 2, 1), (2, 1, 1), (0, 1, 1)), + ((1, 2, 1, 1), 1, (1, 1), (1, 1), (0, 0), (0, 0)), + ) + cases = [ + (geometry, layout, layout == "weight_cl" or i % 2 == 0) + for i, geometry in enumerate(geometries) + for layout in layouts + ] + for (shape, c_out, kernel, stride, before, after), layout, bias in cases: + x, w, b = self.tensors(shape, c_out, kernel, dtype=dtype, bias=bias) + if layout == "channels_last": + x = _channels_last(x) + if layout == "weight_cl": + w = _channels_last(w) + with self.subTest(shape=shape, kernel=kernel, stride=stride, layout=layout): + self.check(x, w, b, stride, before, after) + + def test_geometries_fp32(self): + self.check_geometries(torch.float32, ("contiguous",)) + + def test_geometries_fp32_channels_last(self): + self.check_geometries(torch.float32, ("channels_last",)) + + def test_geometries_fp32_channels_last_weight(self): + self.check_geometries(torch.float32, ("weight_cl",)) + + def test_geometries_fp16(self): + self.check_geometries(torch.float16) + + def test_geometries_bf16(self): + self.check_geometries(torch.bfloat16) + + def test_sliced_inputs_with_storage_offset(self): + base = torch.randn(2, 20, 15, 18, device="cuda") + _, w, b = self.tensors((1, 16, 1, 1), 24, (3, 3)) + for x in ( + base[:, 2:18, 1:-1, 2:], + _channels_last(base)[:, 4:, 3:, :-1], + base[:, 2:18, :, ::2], + ): + with self.subTest(stride=x.stride(), offset=x.storage_offset()): + self.assertGreater(x.storage_offset(), 0) + self.check(x, w, b, (1, 1), (1, 1), (1, 1)) + + def test_bias_is_added_in_float32(self): + for dtype in (torch.float32, torch.float16, torch.bfloat16): + x, w, _ = self.tensors((1, 6, 6, 6), 8, (3, 3), dtype=dtype) + conv = _reference(x, w, None, (1, 1), (1, 1), (1, 1)) + finfo = torch.finfo(dtype) + biases = { + "zero": torch.zeros(8), + "tiny": torch.full((8,), finfo.tiny), + "huge": torch.full((8,), finfo.max / 2), + "near max": torch.full((8,), finfo.max * 0.999), + "cancels": -conv[0, :, 3, 3].float().cpu(), + "nan": torch.full((8,), float("nan")), + "inf": torch.tensor([float("inf"), -float("inf")] * 4), + } + for name, bias in biases.items(): + with self.subTest(dtype=dtype, bias=name): + self.check(x, w, bias.to("cuda", dtype), (1, 1), (1, 1), (1, 1)) + + def test_capability_and_tf32_gate_the_tensor_core_schedule(self): + from executorch.backends.cuda.triton.kernels.conv.conv import ( + force_conv_candidate, + ) + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + Schedule, + ) + + self.assertEqual(implicit_gemm.SCHEDULES[0], Schedule("igemm_ffma", 0)) + params = ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 1) + x, w, b = self.tensors((1, 8, 6, 6), 8, (3, 3)) + fp32, bf16 = ( + ConvCall(x, w, b, params), + ConvCall(x.bfloat16(), w.bfloat16(), b.bfloat16(), params), + ) + expected = { + # (export GPU capability, TF32 allowed): (FP32, BF16) run igemm_tc + ((7, 5), False): (False, False), + ((7, 5), True): (False, False), + ((8, 0), False): (False, True), + ((8, 6), True): (True, True), + } + for (capability, tf32), allowed in expected.items(): + with self.subTest(capability=capability, tf32=tf32), mock.patch.object( + torch.cuda, "get_device_capability", return_value=capability + ), mock.patch.object(implicit_gemm, "conv_allows_tf32", return_value=tf32): + for call in (fp32, bf16): + self.assertIsNone( + implicit_gemm.schedule_supports("igemm_ffma", call) + ) + got = tuple( + implicit_gemm.schedule_supports("igemm_tc", call) is None + for call in (fp32, bf16) + ) + self.assertEqual(got, allowed) + with mock.patch.object( + torch.cuda, "get_device_capability", return_value=(8, 0) + ), mock.patch.object( + implicit_gemm, "conv_allows_tf32", return_value=False + ), force_conv_candidate( + "igemm_tc" + ), self.assertRaisesRegex( + ValueError, "TF32" + ): + self.convolution(x, w, b, (1, 1), (1, 1), (1, 1)) + + def test_timed_schedule_choice(self): + from executorch.backends.cuda.autotune.launch_params import ( + autotune_launch_params, + clear_launch_param_cache, + ) + + if not self.tensor_cores: + self.skipTest("needs tensor cores") + clear_launch_param_cache() + x, w, b = self.tensors((2, 16, 20, 20), 32, (3, 3), dtype=torch.bfloat16) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), autotune_launch_params() as stats: + out = self.convolution(x, w, b, (1, 1), (1, 1), (1, 1)) + stream.synchronize() + self.assertEqual(stats.measured, 1) + self.assertEqual(set(stats.picks[0].times_us), {"igemm_ffma", "igemm_tc"}) + torch.testing.assert_close( + out.double(), + _reference(x, w, b, (1, 1), (1, 1), (1, 1)), + atol=2e-2, + rtol=2e-2, + ) + + def test_64_bit_indexing(self): + from executorch.backends.cuda.triton.kernels.conv.conv import ( + force_conv_candidate, + ) + + # 2^31 + 2^17 input elements: offsets past the int32 range. + shape = (1, 2, 32769, 32768) + if torch.cuda.mem_get_info()[0] < 16 * 2**30: + self.skipTest("needs 16 GiB of free device memory") + w = torch.randn(4, 2, 3, 3, device="cuda", dtype=torch.float16) + for memory_format in (torch.contiguous_format, torch.channels_last): + x = torch.randn(shape, device="cuda", dtype=torch.float16) + x = x.contiguous(memory_format=memory_format) + with self.subTest(memory_format=memory_format), force_conv_candidate( + "igemm_ffma" + ): + out = self.convolution(x, w, None, (1, 1), (1, 1), (1, 1)) + top = _reference(x[:, :, :4], w, None, (1, 1), (1, 1), (0, 1)) + bottom = _reference(x[:, :, -4:], w, None, (1, 1), (0, 1), (1, 1)) + for got, expected in ( + (out[:, :, :3], top[:, :, :3]), + (out[:, :, -3:], bottom), + ): + torch.testing.assert_close( + got, expected.to(torch.float16), atol=3e-3, rtol=3e-3 + ) + del x, out + + +class Tf32AuthorityTest(unittest.TestCase): + def test_follows_the_inductor_conv_template_setting(self): + from executorch.backends.cuda.triton.kernels.conv.core import common + from torch._dynamo import device_interface + + for allowed in (True, False): + interface = SimpleNamespace(conv_allow_tf32=lambda allowed=allowed: allowed) + with self.subTest(allowed=allowed), mock.patch.object( + device_interface, "get_interface_for_device", return_value=interface + ): + self.assertEqual(common.conv_allows_tf32(), allowed) diff --git a/backends/cuda/triton/kernels/conv/conv.py b/backends/cuda/triton/kernels/conv/conv.py index aabbc138679..66709311cd0 100644 --- a/backends/cuda/triton/kernels/conv/conv.py +++ b/backends/cuda/triton/kernels/conv/conv.py @@ -20,7 +20,7 @@ import torch import triton -from executorch.backends.cuda.triton.kernels.conv.core import direct +from executorch.backends.cuda.triton.kernels.conv.core import direct, implicit_gemm from executorch.backends.cuda.triton.kernels.conv.core.common import ( check_params, ConvCall, @@ -53,6 +53,15 @@ class ConvTemplate: direct.run_schedule, direct.worth_routing, ), + ConvTemplate( + "implicit_gemm", + implicit_gemm.supports, + tuple(s.name for s in implicit_gemm.SCHEDULES), + implicit_gemm.schedule_supports, + implicit_gemm.launch, + implicit_gemm.run_schedule, + implicit_gemm.worth_routing, + ), ) _FORCED: contextvars.ContextVar[Optional[tuple[str, Optional[triton.Config]]]] = ( diff --git a/backends/cuda/triton/kernels/conv/core/common.py b/backends/cuda/triton/kernels/conv/core/common.py index 1ce053cc21f..909aa82d687 100644 --- a/backends/cuda/triton/kernels/conv/core/common.py +++ b/backends/cuda/triton/kernels/conv/core/common.py @@ -26,7 +26,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Callable, Optional, Sequence +from typing import Any, Callable, NamedTuple, Optional, Sequence import torch @@ -47,6 +47,25 @@ } +class Schedule(NamedTuple): + name: str + # The compute capability the export GPU needs; 0: every NVIDIA GPU. + min_capability: int + + +def conv_allows_tf32() -> bool: + """Whether FP32 convolutions may compute in TF32: the setting Inductor's + Triton conv templates follow, so a convolution keeps its precision whether + it is routed or not.""" + from torch._dynamo.device_interface import get_interface_for_device + + interface = get_interface_for_device("cuda") + # Older PyTorch releases read the cuDNN-wide setting directly. + if not hasattr(interface, "conv_allow_tf32"): + return torch.backends.cudnn.fp32_precision == "tf32" + return bool(interface.conv_allow_tf32()) + + class ConvBindingError(ValueError): """The arguments of a conv call cannot be bound to the canonical ABI.""" diff --git a/backends/cuda/triton/kernels/conv/core/direct.py b/backends/cuda/triton/kernels/conv/core/direct.py index c3c1781da81..d74023937d7 100644 --- a/backends/cuda/triton/kernels/conv/core/direct.py +++ b/backends/cuda/triton/kernels/conv/core/direct.py @@ -39,7 +39,7 @@ from __future__ import annotations -from typing import NamedTuple, Optional +from typing import Optional import torch import triton @@ -52,6 +52,7 @@ ConvCall, ConvParams, max_offset, + Schedule, SUPPORTED_DTYPES, ) from executorch.backends.cuda.triton.routing import device_supports @@ -66,12 +67,6 @@ _INT32_MAX = 2**31 - 1 -class Schedule(NamedTuple): - name: str - # The compute capability the export GPU needs; 0: every NVIDIA GPU. - min_capability: int - - SCHEDULES = ( Schedule("dw_flat", 0), Schedule("dw_row8", 0), diff --git a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py new file mode 100644 index 00000000000..7335f6e1b5e --- /dev/null +++ b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py @@ -0,0 +1,510 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Implicit-GEMM template: dense (groups 1) forward convolution. + +The convolution is a GEMM of the output pixels (``M = N * OD * OH * OW``) by the +output channels, reducing over ``K = C_in * KD * KH * KW``; input tiles are +gathered from the input as the GEMM walks K. The kernel is written once for +(D, H, W): a 2-D convolution has a unit depth, which, like every shape, stride +and geometry argument, is constexpr and compiles away; so does every bounds +check the geometry cannot fail (all of them when kernel == stride and +there is no padding). Tiles are autotuned. K is walked in +one of two orders: + +- tap-major: for each kernel tap, blocks of input channels, sharing the bounds + check of the tap. Used when the input has at least ``TAP_MAJOR_MIN_CHANNELS`` + channels. +- flat: one index over (channel, tap), decoded per element, which wastes no + lanes on inputs with few channels, such as 3. + +Two schedules write the same output layout, accumulating in FP32 and adding +the bias in FP32 before the single store: + +- ``igemm_ffma``: FP32 FMA on operands converted to FP32. The base schedule, + correct on every NVIDIA GPU. +- ``igemm_tc``: tensor cores on FP16/BF16 operands, or FP32 operands rounded to + TF32 when TF32 is allowed (``conv_allows_tf32``). Needs sm_80 on the export + GPU. For a channel-first output it accumulates the + transposed tile (output channels by pixels), whose rows are contiguous in + memory, so its store is coalesced. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl +from executorch.backends.cuda.autotune.launch_params import ( + autotune_launch_param, + InvalidLaunchParam, +) +from executorch.backends.cuda.triton.kernels.conv.core.common import ( + conv_allows_tf32, + ConvCall, + ConvParams, + max_offset, + Schedule, + SUPPORTED_DTYPES, +) +from executorch.backends.cuda.triton.routing import device_supports +from torch.library import wrap_triton + +SCHEDULES = (Schedule("igemm_ffma", 0), Schedule("igemm_tc", 80)) +TAP_MAJOR_MIN_CHANNELS = 16 +# Output pixels per image from which an output with dense spatial dims is +# tiled per image (always for one image): the partial last tile of each image +# wastes at most 1/8 of a 128-row tile. +PIXEL_TILES_MIN_PIXELS = 1024 +_INT32_MAX = 2**31 - 1 + + +def _output_size(call: ConvCall) -> list[int]: + x, w, p = call.input, call.weight, call.params + return [ + (i + b + a - k) // s + 1 + for i, k, s, b, a in zip( + x.shape[2:], w.shape[2:], p.stride, p.padding_before, p.padding_after + ) + ] + + +def _geometry_reason(call: ConvCall) -> Optional[str]: + x, w, p = call.input, call.weight, call.params + if p.rank not in (2, 3): + return f"implicit GEMM: {p.rank}-D convolution, needs 2-D or 3-D" + if p.transposed: + return "implicit GEMM: transposed convolution" + if any(d != 1 for d in p.dilation): + return f"implicit GEMM: dilation {list(p.dilation)}" + if p.groups != 1 or w.shape[1] != x.shape[1]: + return ( + f"implicit GEMM: grouped convolution (groups {p.groups}, input " + f"{list(x.shape)}, weight {list(w.shape)})" + ) + if min(*x.shape, *w.shape, *_output_size(call)) <= 0: + return f"implicit GEMM: empty input or output (input {list(x.shape)})" + return None + + +def supports(call: ConvCall) -> Optional[str]: + """Why this template cannot run ``call``, or None. Assumes valid canonical + arguments (``check_params``).""" + reason = _geometry_reason(call) + if reason is not None: + return reason + x, w, b = call.input, call.weight, call.bias + if x.dtype not in SUPPORTED_DTYPES: + return f"implicit GEMM: dtype {x.dtype}" + if w.dtype != x.dtype or (b is not None and b.dtype != x.dtype): + return "implicit GEMM: input, weight and bias dtypes differ" + if b is not None and (tuple(b.shape) != (w.shape[0],) or not b.is_contiguous()): + return "implicit GEMM: bias is not a contiguous [C_out] tensor" + return None + + +def worth_routing(call: ConvCall) -> Optional[str]: + """Why a call ``supports`` accepts is left to the regular lowering, or None.""" + return "implicit GEMM: not benchmarked" + + +def schedule_supports(name: str, call: ConvCall) -> Optional[str]: + """Why schedule ``name`` cannot run ``call`` (which ``supports``), or None.""" + schedule = next((s for s in SCHEDULES if s.name == name), None) + if schedule is None: + return f"unknown schedule {name}" + if not device_supports(schedule.min_capability, call.input.device): + return f"{name} needs sm_{schedule.min_capability}" + if ( + name == "igemm_tc" + and call.input.dtype == torch.float32 + and not conv_allows_tf32() + ): + return "igemm_tc needs TF32 for FP32 operands, which is not allowed" + return None + + +@triton.jit +def _bound(ok, i, SIZE: tl.constexpr, LOW: tl.constexpr, HIGH: tl.constexpr): + if LOW: + ok = ok & (i >= 0) + if HIGH: + ok = ok & (i < SIZE) + return ok + + +@triton.jit +def _dot(a, b, acc, FFMA: tl.constexpr, CHANNEL_MAJOR: tl.constexpr): + if CHANNEL_MAJOR: + a, b = tl.trans(b), tl.trans(a) + if FFMA: + return tl.dot(a.to(tl.float32), b.to(tl.float32), acc, input_precision="ieee") + if a.dtype == tl.float32: + return tl.dot(a, b, acc, input_precision="tf32") + return tl.dot(a, b, acc) + + +@triton.jit +def _store( + out_ptr, + acc, + b_ptr, + out_rows, + m_mask, + offs_n, + n_mask, + SO_C: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNEL_MAJOR: tl.constexpr, +): + if HAS_BIAS: + bias = tl.load(b_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + if CHANNEL_MAJOR: + if HAS_BIAS: + acc += bias[:, None] + out_ptrs = out_ptr + offs_n[:, None] * SO_C + out_rows[None, :] + mask = n_mask[:, None] & m_mask[None, :] + else: + if HAS_BIAS: + acc += bias[None, :] + out_ptrs = out_ptr + out_rows[:, None] + offs_n[None, :] * SO_C + mask = m_mask[:, None] & n_mask[None, :] + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + D: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + CO: tl.constexpr, + OD: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KD: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_D: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_D: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_D: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_O: tl.constexpr, + SW_C: tl.constexpr, + SW_D: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_D: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + FFMA: tl.constexpr, + INDEX_64: tl.constexpr, + PIXEL_TILES: tl.constexpr, + CHANNEL_MAJOR: tl.constexpr, + TAP_MAJOR: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_M: tl.constexpr, +): + PIXELS: tl.constexpr = OD * OH * OW + TAPS: tl.constexpr = KD * KH * KW + # Which bounds checks the geometry can fail. + LOW_D: tl.constexpr = PAD_D > 0 + LOW_H: tl.constexpr = PAD_H > 0 + LOW_W: tl.constexpr = PAD_W > 0 + HIGH_D: tl.constexpr = (OD - 1) * STRIDE_D - PAD_D + KD > D + HIGH_H: tl.constexpr = (OH - 1) * STRIDE_H - PAD_H + KH > H + HIGH_W: tl.constexpr = (OW - 1) * STRIDE_W - PAD_W + KW > W + # With pixel tiles, an M tile holds consecutive output pixels of one image, + # so the store into a channel-first output is provably contiguous. + PIXEL_BLOCKS: tl.constexpr = (PIXELS + BLOCK_M - 1) // BLOCK_M + GRID_M: tl.constexpr = ( + N * PIXEL_BLOCKS if PIXEL_TILES else (N * PIXELS + BLOCK_M - 1) // BLOCK_M + ) + GRID_N: tl.constexpr = (CO + BLOCK_N - 1) // BLOCK_N + + pid = tl.program_id(0) + if INDEX_64: + pid = pid.to(tl.int64) + first_m = pid // (GROUP_M * GRID_N) * GROUP_M + rows = tl.minimum(GRID_M - first_m, GROUP_M) + pid_m = first_m + pid % (GROUP_M * GRID_N) % rows + pid_n = pid % (GROUP_M * GRID_N) // rows + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + n_mask = offs_n < CO + if PIXEL_TILES: + n = pid_m // PIXEL_BLOCKS + pixel = pid_m % PIXEL_BLOCKS * BLOCK_M + tl.arange(0, BLOCK_M) + m_mask = pixel < PIXELS + else: + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + m_mask = offs_m < N * PIXELS + n = offs_m // PIXELS + pixel = offs_m % PIXELS + ow = pixel % OW + oh = pixel // OW % OH + od = pixel // (OW * OH) + d0 = od * STRIDE_D - PAD_D + h0 = oh * STRIDE_H - PAD_H + w0 = ow * STRIDE_W - PAD_W + x_rows = x_ptr + n * SX_N + d0 * SX_D + h0 * SX_H + w0 * SX_W + w_cols = w_ptr + offs_n * SW_O + if CHANNEL_MAJOR: + acc = tl.zeros((BLOCK_N, BLOCK_M), dtype=tl.float32) + else: + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + if TAP_MAJOR: + C_BLOCKS: tl.constexpr = (C + BLOCK_K - 1) // BLOCK_K + for step in range(TAPS * C_BLOCKS): + tap = step // C_BLOCKS + c = (step % C_BLOCKS) * BLOCK_K + tl.arange(0, BLOCK_K) + if INDEX_64: + c = c.to(tl.int64) + kd = tap // (KH * KW) + kh = tap // KW % KH + kw = tap % KW + ok = _bound(m_mask, d0 + kd, D, LOW_D, HIGH_D) + ok = _bound(ok, h0 + kh, H, LOW_H, HIGH_H) + ok = _bound(ok, w0 + kw, W, LOW_W, HIGH_W) + c_ok = c < C + a = tl.load( + x_rows[:, None] + + (kd * SX_D + kh * SX_H + kw * SX_W) + + c[None, :] * SX_C, + mask=ok[:, None] & c_ok[None, :], + other=0.0, + ) + b = tl.load( + w_cols[None, :] + + (kd * SW_D + kh * SW_H + kw * SW_W) + + c[:, None] * SW_C, + mask=c_ok[:, None] & n_mask[None, :], + other=0.0, + ) + acc = _dot(a, b, acc, FFMA, CHANNEL_MAJOR) + else: + K: tl.constexpr = C * TAPS + for k0 in range(0, K, BLOCK_K): + k = k0 + tl.arange(0, BLOCK_K) + if INDEX_64: + k = k.to(tl.int64) + c = k // TAPS + tap = k % TAPS + kd = tap // (KH * KW) + kh = tap // KW % KH + kw = tap % KW + k_ok = k < K + ok = _bound( + m_mask[:, None] & k_ok[None, :], + d0[:, None] + kd[None, :], + D, + LOW_D, + HIGH_D, + ) + ok = _bound(ok, h0[:, None] + kh[None, :], H, LOW_H, HIGH_H) + ok = _bound(ok, w0[:, None] + kw[None, :], W, LOW_W, HIGH_W) + a = tl.load( + x_rows[:, None] + + (c * SX_C + kd * SX_D + kh * SX_H + kw * SX_W)[None, :], + mask=ok, + other=0.0, + ) + b = tl.load( + w_cols[None, :] + + (c * SW_C + kd * SW_D + kh * SW_H + kw * SW_W)[:, None], + mask=k_ok[:, None] & n_mask[None, :], + other=0.0, + ) + acc = _dot(a, b, acc, FFMA, CHANNEL_MAJOR) + if PIXEL_TILES: + out_rows = n * SO_N + pixel + else: + out_rows = n * SO_N + od * SO_D + oh * SO_H + ow * SO_W + _store( + out_ptr, + acc, + b_ptr, + out_rows, + m_mask, + offs_n, + n_mask, + SO_C, + HAS_BIAS, + CHANNEL_MAJOR, + ) + + +def _configs(tiles, tap_major: bool) -> list[triton.Config]: + return [ + triton.Config( + { + "TAP_MAJOR": tap_major, + "BLOCK_M": m, + "BLOCK_N": n, + "BLOCK_K": k, + "GROUP_M": 8, + }, + num_warps=warps, + num_stages=stages, + ) + for m, n, k, warps, stages in tiles + ] + + +# (BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) +_FFMA_TILES = ( + (64, 64, 16, 4, 3), + (128, 64, 16, 8, 3), + (64, 128, 16, 8, 3), + (128, 128, 16, 8, 2), + (64, 64, 32, 8, 3), + (64, 32, 32, 4, 3), + (128, 32, 16, 4, 3), +) +_TC_TILES = ( + (64, 64, 32, 4, 4), + (128, 64, 32, 4, 4), + (64, 128, 32, 4, 4), + (128, 128, 32, 8, 3), + (128, 32, 32, 4, 4), + (64, 32, 32, 4, 3), + (64, 32, 64, 4, 3), + (32, 64, 64, 4, 3), +) +CONFIGS = { + (schedule, tap_major): _configs(tiles, tap_major) + for schedule, tiles in (("igemm_ffma", _FFMA_TILES), ("igemm_tc", _TC_TILES)) + for tap_major in (True, False) +} +_TUNING_KEY = [ + "N", + "C", + "D", + "H", + "W", + "CO", + "KD", + "KH", + "KW", + "STRIDE_D", + "STRIDE_H", + "STRIDE_W", + "SX_C", + "SO_C", +] +_KERNELS = { + key: triton.autotune(configs=configs, key=_TUNING_KEY)(_kernel) + for key, configs in CONFIGS.items() +} + + +def _as_3d(values, fill: int) -> tuple[int, ...]: + return (fill,) * (3 - len(values)) + tuple(values) + + +def run_schedule( + name: str, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + params: ConvParams, + config: Optional[triton.Config] = None, +) -> None: + """Launches schedule ``name``; ``config`` replaces its autotuned config + (and the K order, if it sets ``TAP_MAJOR``).""" + n, c = x.shape[:2] + co = weight.shape[0] + out_size = _as_3d(out.shape[2:], 1) + tap_major = c >= TAP_MAJOR_MIN_CHANNELS + options = {} + if config is not None: + tap_major = config.kwargs.get("TAP_MAJOR", tap_major) + options = { + **config.kwargs, + "TAP_MAJOR": tap_major, + "num_warps": config.num_warps, + "num_stages": config.num_stages, + } + kernel = _KERNELS[(name, tap_major)] + if config is not None: + kernel = kernel.fn + pixels = out_size[0] * out_size[1] * out_size[2] + out_strides = _as_3d(out.stride()[2:], 0) + dense = (out_size[1] * out_size[2], out_size[2], 1) + pixel_tiles = (n == 1 or pixels >= PIXEL_TILES_MIN_PIXELS) and all( + size == 1 or s == d for size, s, d in zip(out_size, out_strides, dense) + ) + + def grid(meta): + if pixel_tiles: + rows = n * triton.cdiv(pixels, meta["BLOCK_M"]) + else: + rows = triton.cdiv(n * pixels, meta["BLOCK_M"]) + return (rows * triton.cdiv(co, meta["BLOCK_N"]),) + + wrap_triton(kernel)[grid]( + x, + weight, + x if bias is None else bias, + out, + n, + c, + *_as_3d(x.shape[2:], 1), + co, + *out_size, + *_as_3d(weight.shape[2:], 1), + *_as_3d(params.stride, 1), + *_as_3d(params.padding_before, 0), + *x.stride()[:2], + *_as_3d(x.stride()[2:], 0), + *weight.stride()[:2], + *_as_3d(weight.stride()[2:], 0), + *out.stride()[:2], + *out_strides, + HAS_BIAS=bias is not None, + FFMA=name == "igemm_ffma", + INDEX_64=max(max_offset(x), max_offset(out), max_offset(weight)) > _INT32_MAX + or n * pixels * co > _INT32_MAX, + PIXEL_TILES=pixel_tiles, + CHANNEL_MAJOR=name == "igemm_tc" and out.stride(1) != 1, + **options, + ) + + +@autotune_launch_param("schedule", tuple(s.name for s in SCHEDULES)) +def launch( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + out: torch.Tensor, + params: ConvParams, + *, + schedule: str, +) -> None: + """Runs the convolution into ``out``; during AOTI compile the schedule is timed.""" + reason = schedule_supports(schedule, ConvCall(x, weight, bias, params)) + if reason is not None: + raise InvalidLaunchParam(reason) + run_schedule(schedule, x, weight, bias, out, params)