diff --git a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py index 9bd576ad0e5..0a4aaa3f538 100644 --- a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py +++ b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py @@ -26,8 +26,8 @@ records (``cl_to_nchw``: triton::convolution gets ``output_strides``, the other paths a ``contiguous()`` that Inductor may fuse). -Suites: ``shapes`` (dense and kernel == stride convolutions at representative -shapes), ``envelope`` (each axis of what the +Suites: ``shapes`` (dense, kernel == stride and transposed convolutions at +representative shapes), ``envelope`` (each axis of what the template supports varied around base shapes, in every precision and input layout) and ``boundary`` (every geometry the router takes at every corner of its limits, ``implicit_gemm.boundary_calls()``). The router's @@ -47,12 +47,14 @@ import torch import torch.nn.functional as F from executorch.backends.cuda.autotune.launch_params import autotune_launch_params +from executorch.backends.cuda.cuda_backend import _conv_transpose_as_conv from executorch.backends.cuda.triton.kernels.conv.conv import ( convolution, force_conv_candidate, ) from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm from executorch.backends.cuda.triton.kernels.conv.core.common import conv_allows_tf32 +from torch.fx.experimental.proxy_tensor import make_fx @dataclass(frozen=True) @@ -70,6 +72,7 @@ class Case: # The padding is a separate pad op before a padding-0 convolution, as the # model writes it; both paths run that pad. explicit_pad: bool = False + transposed: bool = False @dataclass(frozen=True) @@ -80,7 +83,17 @@ class Pad: after: tuple[int, ...] -def _case(label, shape, out_channels, kernel, stride, padding, bias, precision): +def _case( + label, + shape, + out_channels, + kernel, + stride, + padding, + bias, + precision, + transposed=False, +): rank = len(shape) - 2 def expand(v): @@ -102,6 +115,7 @@ def expand(v): bias, precision, explicit_pad=explicit, + transposed=transposed, ) @@ -367,12 +381,26 @@ def expand(v): ] +# Transposed 1-D convolutions, likewise: kernel 4, stride 2, upsampling +# sequences of 4 to 32 steps. +TRANSPOSED_SHAPES = [ + ("1d_transposed", "1024 ch T4", (1, 1024, 4), 1024, 4, 2, 1, True), + ("1d_transposed", "512 ch T8", (1, 512, 8), 512, 4, 2, 1, True), + ("1d_transposed", "1024 ch T16", (1, 1024, 16), 1024, 4, 2, 1, True), + ("1d_transposed", "512 ch T32", (1, 512, 32), 512, 4, 2, 1, True), +] + + def shape_cases(precision, layout, groups=None): return [ Case( - **{**asdict(_case(f"{group} {label}", *rest, precision)), "layout": layout} + **{ + **asdict(_case(f"{group} {label}", *rest, precision, transposed)), + "layout": layout, + } ) - for group, label, *rest in SHAPES + for table, transposed in ((SHAPES, False), (TRANSPOSED_SHAPES, True)) + for group, label, *rest in table if groups is None or group in groups ] @@ -530,7 +558,7 @@ def _args(case): list(case.before), list(case.after), [1] * rank, - False, + case.transposed, [0] * rank, 1, ) @@ -573,6 +601,10 @@ def aten_conv(x, w, b, case): """What the CUDA backend gives Inductor: a separate pad for explicit or asymmetric padding, a 1-D convolution as a 2-D one, and a contiguous output where the graph records one.""" + if case.transposed: + conv = {1: F.conv_transpose1d, 2: F.conv_transpose2d, 3: F.conv_transpose3d} + out = conv[len(case.kernel)](x, w, b, case.stride, case.before) + return out.contiguous() if case.layout == "cl_to_nchw" else out conv = {1: _conv1d, 2: F.conv2d, 3: F.conv3d}[len(case.kernel)] if case.explicit_pad or case.before != case.after: out = conv(_pad(x, case), w, b, case.stride) @@ -583,14 +615,21 @@ def aten_conv(x, w, b, case): def inductor_runner(x, w, b, case): torch._dynamo.reset() + fn = lambda x, w, b: aten_conv(x, w, b, case) # noqa: E731 + if case.transposed: + # The backend rewrites a 1-D or 2-D transposed convolution as a + # convolution before Inductor sees it. + fn = make_fx( + lambda x, w, b: _conv_transpose_as_conv( + x, w, b, list(case.stride), list(case.before) + ) + )(x, w, b) with torch._inductor.config.patch( max_autotune=True, max_autotune_conv_backends="TRITON", max_autotune_gemm_backends="TRITON", ): - compiled = torch.compile( - lambda x, w, b: aten_conv(x, w, b, case), dynamic=False, fullgraph=True - ) + compiled = torch.compile(fn, dynamic=False, fullgraph=True) compiled(x, w, b) return lambda: compiled(x, w, b) @@ -618,12 +657,12 @@ def _tensors(case): channels = case.shape[1] fan_in = channels * torch.Size(case.kernel).numel() x = torch.randn(case.shape, device="cuda", dtype=dtype) - w = ( - torch.randn( - case.out_channels, channels, *case.kernel, device="cuda", dtype=dtype - ) - / fan_in**0.5 + w_shape = ( + (channels, case.out_channels) + if case.transposed + else (case.out_channels, channels) ) + w = torch.randn(*w_shape, *case.kernel, device="cuda", dtype=dtype) / fan_in**0.5 b = ( torch.randn(case.out_channels, device="cuda", dtype=dtype) if case.bias @@ -665,7 +704,7 @@ def run_case(case, flush, threshold): } out = ours() nbytes = sum(t.numel() * t.element_size() for t in (x, w, b, out) if t is not None) - flops = 2 * out.numel() * w[0].numel() + flops = 2 * out.numel() * x.shape[1] * torch.Size(case.kernel).numel() speedup = results["inductor"][0] / results["ours"][0] significant = speedup >= threshold and results["ours"][2] < results["inductor"][1] return { diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py index 5fc5392b8ad..284509b5400 100644 --- a/backends/cuda/tests/test_conv_implicit_gemm.py +++ b/backends/cuda/tests/test_conv_implicit_gemm.py @@ -27,10 +27,17 @@ _TF32_TOLERANCE = 5e-3 -def _reference(x, w, b, stride, before, after): +def _reference(x, w, b, stride, before, after, output_padding=None): + """The convolution in FP64; transposed if ``output_padding`` is given, as + the full transposed convolution cropped by ``before`` and ``after``.""" + bias = None if b is None else b.double() + if output_padding is not None: + conv = {3: F.conv_transpose1d, 4: F.conv_transpose2d, 5: F.conv_transpose3d} + full = conv[x.dim()](x.double(), w.double(), bias, stride, 0, output_padding) + crop = [slice(s, n - a) for s, n, a in zip(before, full.shape[2:], after)] + return full[(..., *crop)] pad = [p for pair in zip(reversed(before), reversed(after)) for p in pair] conv = {3: F.conv1d, 4: F.conv2d, 5: F.conv3d}[x.dim()] - bias = None if b is None else b.double() return conv(F.pad(x.double(), pad), w.double(), bias, stride) @@ -58,7 +65,18 @@ def setUp(self): torch.manual_seed(0) self.tensor_cores = torch.cuda.get_device_capability() >= (8, 0) - def convolution(self, x, w, b, stride, before, after, output_strides=None): + def convolution( + self, + x, + w, + b, + stride, + before, + after, + output_strides=None, + output_padding=None, + ): + """triton::convolution; transposed if ``output_padding`` is given.""" from executorch.backends.cuda.triton.kernels.conv.conv import convolution rank = x.dim() - 2 @@ -70,19 +88,29 @@ def convolution(self, x, w, b, stride, before, after, output_strides=None): list(before), list(after), [1] * rank, - False, - [0] * rank, + output_padding is not None, + [0] * rank if output_padding is None else list(output_padding), 1, output_strides, ) def check( - self, x, w, b, stride, before, after, every_config=False, output_strides=None + self, + x, + w, + b, + stride, + before, + after, + every_config=False, + output_strides=None, + output_padding=None, ): """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 fake implementation's, and ``output_strides`` if given); the tensor-core - schedule runs FP32 operands in TF32.""" + schedule runs FP32 operands in TF32. Transposed if ``output_padding`` is + given.""" from executorch.backends.cuda.triton.kernels.conv.conv import ( force_conv_candidate, ) @@ -99,8 +127,8 @@ def check( tuple(before), tuple(after), (1,) * rank, - False, - (0,) * rank, + output_padding is not None, + (0,) * rank if output_padding is None else tuple(output_padding), 1, ) call = ConvCall( @@ -108,7 +136,7 @@ def check( ) 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) + expected = _reference(x, w, b, stride, before, after, output_padding) layout = out_shape_and_stride(call) with FakeTensorMode() as mode: fake = self.convolution( @@ -117,6 +145,7 @@ def check( before, after, output_strides, + output_padding, ) self.assertEqual((tuple(fake.shape), fake.stride()), layout) if output_strides is not None: @@ -133,7 +162,7 @@ def check( 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, output_strides + x, w, b, stride, before, after, output_strides, output_padding ) self.assertEqual((tuple(out.shape), out.stride()), layout) torch.testing.assert_close( @@ -145,13 +174,20 @@ def check( msg=lambda m, s=schedule, c=config: f"{s} {c}: {m}", ) - def tensors(self, shape, out_channels, kernel, dtype=torch.float32, bias=True): + def tensors( + self, + shape, + out_channels, + kernel, + dtype=torch.float32, + bias=True, + transposed=False, + ): 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 - ) + w_shape = (channels, out_channels) if transposed else (out_channels, channels) + w = torch.randn(*w_shape, *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 @@ -272,6 +308,55 @@ def test_output_layouts(self): x, w, b, stride, padding, padding, output_strides=strides ) + def test_transposed(self): + geometries = ( + # input, C_out, kernel, stride, before, after, output_padding: a + # kernel-4 stride-2 1-D upsample, then edge cases; an asymmetric crop + # only the canonical op can express; 16+ channels walk K tap-major. + ((1, 32, 8), 16, (4,), (2,), (1,), (1,), (0,)), + ((2, 8, 5, 6), 12, (3, 3), (2, 2), (1, 1), (1, 1), (1, 1)), + ( + (1, 6, 3, 4, 5), + 10, + (2, 3, 3), + (2, 1, 3), + (0, 1, 1), + (0, 1, 1), + (1, 0, 2), + ), + ((1, 20, 7), 9, (5,), (3,), (3,), (1,), (0,)), + ((1, 4, 9), 4, (3,), (1,), (1,), (1,), (0,)), + ((2, 24, 6, 7), 16, (4, 4), (2, 2), (1, 1), (1, 1), (0, 0)), + ) + cases = itertools.product( + geometries, + (torch.float32, torch.float16, torch.bfloat16), + ("contiguous", "channels_last"), + ) + for geometry, dtype, layout in cases: + shape, c_out, kernel, stride, before, after, output_padding = geometry + x, w, b = self.tensors(shape, c_out, kernel, dtype=dtype, transposed=True) + if layout == "channels_last": + x = _channels_last(x) + with self.subTest(shape=shape, kernel=kernel, dtype=dtype, layout=layout): + self.check( + x, w, b, stride, before, after, output_padding=output_padding + ) + x, w, b = self.tensors((1, 32, 8), 16, (4,), transposed=True) + with self.subTest(every_config=True): + self.check(x, w, b, (2,), (1,), (1,), True, output_padding=(0,)) + with self.subTest(output_strides=(256, 1, 16)): + self.check( + x, + w, + b, + (2,), + (1,), + (1,), + output_strides=(256, 1, 16), + output_padding=(0,), + ) + 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)) @@ -476,6 +561,19 @@ def forward(self, x): return self.time(y) +class _Upsample(nn.Module): + """A transposed 1-D upsample, and a transposed + 3-D one.""" + + def __init__(self): + super().__init__() + self.up1d = nn.ConvTranspose1d(32, 32, 4, stride=2, padding=1) + self.up3d = nn.ConvTranspose3d(8, 4, 3, stride=2, padding=1, output_padding=1) + + def forward(self, x, y): + return self.up1d(x), self.up3d(y) + + class ImplicitGemmLoweringTest(unittest.TestCase): def setUp(self): if not torch.cuda.is_available() or torch.version.hip is not None: @@ -605,3 +703,19 @@ def test_conv3d_block_lowers_and_runs(self): ) with mock.patch.object(routing, "TEMPLATES", take_all): self.lower_and_run(model, inputs, 3, 3e-2) + + def test_transposed_lowers_and_runs(self): + torch.manual_seed(0) + model = _Upsample().to("cuda").eval() + inputs = ( + torch.randn(1, 32, 8, device="cuda"), + torch.randn(1, 8, 3, 4, 5, device="cuda"), + ) + # Inductor has no transposed template on CUDA. to_edge keeps both as + # aten.convolution; to_edge_transform_and_lower keeps conv_transpose1d, + # which the backend rewrites as a convolution, and conv_transpose3d. + self.lower_and_run(model, inputs, 2, 1e-4, route="required", to_backend=True) + output = self.lower_and_run( + model, inputs, 1, self.fp32_tolerance(), route="required" + ) + self.assertIn("transposed convolution is not benchmarked faster", output) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index a31d0767f9f..66e59730f23 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/backends/cuda/tests/test_conv_routing.py @@ -1308,3 +1308,84 @@ def test_explicit_pad_stays_its_own_node(self): ) rank = len(shape) - 2 self.assertEqual(conv.args[4:6], ([0] * rank, [0] * rank)) + + def test_transposed(self): + # (model, input shape, dtype, whether each API's graph needs the + # template): to_edge keeps aten.convolution, which Inductor cannot lower + # transposed on CUDA; to_edge_transform_and_lower keeps + # conv_transpose1d/2d, which the backend rewrites as convolutions, and + # conv_transpose3d, which it cannot lower. + cases = ( + ( + "k4 s2 upsample", + _dense( + nn.ConvTranspose1d, + 1024, + 1024, + 4, + torch.float32, + stride=2, + padding=1, + ), + (1, 1024, 4), + torch.float32, + {"to_edge_transform_and_lower": False, "to_edge_to_backend": True}, + ), + ( + "3-channel 2-D", + _dense( + nn.ConvTranspose2d, 3, 64, 3, torch.bfloat16, stride=2, padding=1 + ), + (1, 3, 56, 56), + torch.bfloat16, + {"to_edge_transform_and_lower": False, "to_edge_to_backend": True}, + ), + ( + "3-D", + _dense( + nn.ConvTranspose3d, + 8, + 4, + 3, + torch.float32, + stride=2, + output_padding=1, + ), + (1, 8, 3, 4, 5), + torch.float32, + {"to_edge_transform_and_lower": True, "to_edge_to_backend": True}, + ), + ) + for row, model, shape, dtype, required in cases: + graphs = _backend_graphs( + model, lambda shape=shape, dtype=dtype: torch.randn(shape, dtype=dtype) + ) + for api, graph_module in graphs.items(): + with self.subTest(row=row, api=api): + targets, logs = self.run_pass(graph_module) + if required[api]: + (node,) = [ + n + for n in graph_module.graph.nodes + if n.target == torch.ops.triton.convolution.default + ] + self.assertIn("required, implicit_gemm template", logs) + self.assertTrue(node.args[7]) + self.assertEqual(node.args[-1], list(node.meta["val"].stride())) + else: + self.assertNotIn(torch.ops.triton.convolution.default, targets) + self.assertIn("transposed convolution is not benchmarked", logs) + + def test_unsupported_transposed_is_kept(self): + for model, reason in ( + ( + _dense(nn.ConvTranspose1d, 8, 8, 3, torch.float32, dilation=2), + "implicit GEMM: dilation [2]", + ), + ( + _dense(nn.ConvTranspose1d, 8, 8, 3, torch.float32, groups=2), + "implicit GEMM: grouped convolution", + ), + ): + with self.subTest(reason=reason): + self.assertKept(model, lambda: torch.randn(1, 8, 9), reason=reason) diff --git a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py index 4273941a31e..546da6c7e64 100644 --- a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py +++ b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py @@ -4,7 +4,8 @@ # 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) 1-D, 2-D and 3-D forward convolution. +"""Implicit-GEMM template: dense (groups 1) 1-D, 2-D and 3-D forward and +transposed 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 @@ -21,6 +22,11 @@ - flat: one index over (channel, tap), decoded per element, which wastes no lanes on inputs with few channels, such as 3. +A transposed convolution is the same GEMM over its output pixels: output index +``o`` reads, at tap ``k``, input index ``(o + padding_before - k) / stride`` +when that is an integer inside the input, and its ``[C_in, C_out, ...]`` +weight is read with its first two strides swapped. + Two schedules write the same output layout, accumulating in FP32 and adding the bias in FP32 before the single store: @@ -70,6 +76,18 @@ def _output_size(call: ConvCall) -> list[int]: x, w, p = call.input, call.weight, call.params + if p.transposed: + return [ + (i - 1) * s - b - a + k + o + for i, k, s, b, a, o in zip( + x.shape[2:], + w.shape[2:], + p.stride, + p.padding_before, + p.padding_after, + p.output_padding, + ) + ] return [ (i + b + a - k) // s + 1 for i, k, s, b, a in zip( @@ -78,13 +96,16 @@ def _output_size(call: ConvCall) -> list[int]: ] +def _out_channels(call: ConvCall) -> int: + return call.weight.shape[1 if call.params.transposed else 0] + + def _geometry_reason(call: ConvCall) -> Optional[str]: x, w, p = call.input, call.weight, call.params - 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]: + in_channels = w.shape[0] if p.transposed else w.shape[1] + if p.groups != 1 or in_channels != x.shape[1]: return ( f"implicit GEMM: grouped convolution (groups {p.groups}, input " f"{list(x.shape)}, weight {list(w.shape)})" @@ -105,7 +126,9 @@ def supports(call: ConvCall) -> Optional[str]: 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()): + if b is not None and ( + tuple(b.shape) != (_out_channels(call),) or not b.is_contiguous() + ): return "implicit GEMM: bias is not a contiguous [C_out] tensor" return None @@ -181,6 +204,8 @@ def _family(call: ConvCall) -> Optional[RoutedFamily]: def worth_routing(call: ConvCall) -> Optional[str]: """Why a call ``supports`` accepts is left to the regular lowering, or None.""" x, w, p = call.input, call.weight, call.params + if p.transposed: + return "implicit GEMM: a transposed convolution is not benchmarked faster" if all(k == 1 for k in w.shape[2:]): return "implicit GEMM: a 1x1 convolution keeps the regular lowering" tensor_cores = x.dtype != torch.float32 or conv_allows_tf32() @@ -307,6 +332,16 @@ def _bound(ok, i, SIZE: tl.constexpr, LOW: tl.constexpr, HIGH: tl.constexpr): return ok +@triton.jit +def _source(o, k, PAD: tl.constexpr, STRIDE: tl.constexpr, SIZE: tl.constexpr): + """The input index a transposed convolution reads for output index ``o`` and + tap ``k``, and whether it exists: ``o + PAD - k`` must be a multiple of the + stride inside the input.""" + q = o + PAD - k + i = q // STRIDE + return i, (q >= 0) & (i * STRIDE == q) & (i < SIZE) + + @triton.jit def _dot(a, b, acc, FFMA: tl.constexpr, CHANNEL_MAJOR: tl.constexpr): if CHANNEL_MAJOR: @@ -390,6 +425,7 @@ def _kernel( INDEX_64: tl.constexpr, PIXEL_TILES: tl.constexpr, CHANNEL_MAJOR: tl.constexpr, + TRANSPOSED: tl.constexpr, TAP_MAJOR: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, @@ -437,7 +473,8 @@ def _kernel( 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 + x_images = x_ptr + n * SX_N + x_rows = x_images + 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) @@ -453,14 +490,20 @@ def _kernel( 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) + if TRANSPOSED: + i_d, ok_d = _source(od, kd, PAD_D, STRIDE_D, D) + i_h, ok_h = _source(oh, kh, PAD_H, STRIDE_H, H) + i_w, ok_w = _source(ow, kw, PAD_W, STRIDE_W, W) + ok = m_mask & ok_d & ok_h & ok_w + taps = x_images + i_d * SX_D + i_h * SX_H + i_w * SX_W + else: + 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) + taps = x_rows + (kd * SX_D + kh * SX_H + kw * SX_W) c_ok = c < C a = tl.load( - x_rows[:, None] - + (kd * SX_D + kh * SX_H + kw * SX_W) - + c[None, :] * SX_C, + taps[:, None] + c[None, :] * SX_C, mask=ok[:, None] & c_ok[None, :], other=0.0, ) @@ -484,21 +527,26 @@ def _kernel( 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, - ) + if TRANSPOSED: + i_d, ok_d = _source(od[:, None], kd[None, :], PAD_D, STRIDE_D, D) + i_h, ok_h = _source(oh[:, None], kh[None, :], PAD_H, STRIDE_H, H) + i_w, ok_w = _source(ow[:, None], kw[None, :], PAD_W, STRIDE_W, W) + ok = m_mask[:, None] & k_ok[None, :] & ok_d & ok_h & ok_w + # With pixel tiles the image (and its base pointer) is a scalar. + images = x_images if PIXEL_TILES else x_images[:, None] + taps = images + i_d * SX_D + i_h * SX_H + i_w * SX_W + else: + 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) + taps = x_rows[:, None] + (kd * SX_D + kh * SX_H + kw * SX_W)[None, :] + a = tl.load(taps + (c * SX_C)[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], @@ -604,7 +652,12 @@ def run_schedule( """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] + co = out.shape[1] + # A transposed weight is [C_in, C_out, ...]: the kernel indexes it by + # (output channel, input channel) strides. + w_strides = weight.stride()[:2] + if params.transposed: + w_strides = w_strides[::-1] out_size = _as_3d(out.shape[2:], 1) tap_major = c >= TAP_MAJOR_MIN_CHANNELS options = {} @@ -648,7 +701,7 @@ def grid(meta): *_as_3d(params.padding_before, 0), *x.stride()[:2], *_as_3d(x.stride()[2:], 0), - *weight.stride()[:2], + *w_strides, *_as_3d(weight.stride()[2:], 0), *out.stride()[:2], *out_strides, @@ -658,6 +711,7 @@ def grid(meta): or n * pixels * co > _INT32_MAX, PIXEL_TILES=pixel_tiles, CHANNEL_MAJOR=name == "igemm_tc" and out.stride(1) != 1, + TRANSPOSED=params.transposed, **options, )