From 821aefab9f590115398fa934dbf46061f721f726 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Fri, 13 Nov 2026 23:58:40 -0800 Subject: [PATCH] Update [ghstack-poisoned] --- .../conv/benchmark_implicit_gemm.py | 582 ++++++++++++++++++ .../cuda/tests/test_conv_implicit_gemm.py | 126 +++- backends/cuda/tests/test_conv_routing.py | 390 +++++++++++- .../triton/kernels/conv/core/implicit_gemm.py | 180 +++++- 4 files changed, 1274 insertions(+), 4 deletions(-) create mode 100644 backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py diff --git a/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py new file mode 100644 index 00000000000..ea2abe1aedf --- /dev/null +++ b/backends/cuda/benchmarks/conv/benchmark_implicit_gemm.py @@ -0,0 +1,582 @@ +#!/usr/bin/env python3 +# 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. + +""" +Benchmark the implicit-GEMM template of triton::convolution on dense +convolutions against Inductor's Triton conv template (what the CUDA backend +lowers a convolution to otherwise) and cuDNN. + +Timing, the compile options of the two Triton paths and "significantly +faster" are those of ``benchmark_depthwise.py``: cold L2, CUDA graph replay, +median and 10th-90th percentile spread, a median speedup of at least +``--threshold`` with our 90th percentile below Inductor's 10th. triton::convolution +runs with the schedule its compile would pick, for the targets of this GPU and +the TF32 setting the backend reads (``conv_allows_tf32``). Each case also +records the kernels Inductor launched, which shows the path it took (e.g. a +1x1 convolution as a matrix multiplication). + +Precisions: ``fp32`` (TF32 off everywhere), ``tf32`` (TF32 allowed; needs a +PyTorch whose Inductor conv templates can use TF32), ``fp16`` and ``bf16``. + +Suites: ``shapes`` (dense and kernel == stride 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 +``implicit_gemm.worth_routing`` takes the range these suites show to be +significantly faster. + +Usage: + python benchmark_implicit_gemm.py --precision bf16 + python benchmark_implicit_gemm.py --suite envelope --shard 0/8 --json out.jsonl +""" + +import argparse +import json +import statistics +from dataclasses import asdict, dataclass + +import torch +import torch.nn.functional as F +from executorch.backends.cuda.autotune.launch_params import autotune_launch_params +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 + + +@dataclass(frozen=True) +class Case: + label: str + shape: tuple[int, ...] + out_channels: int + kernel: tuple[int, ...] + stride: tuple[int, ...] + before: tuple[int, ...] + after: tuple[int, ...] + bias: bool + precision: str = "fp32" + layout: str = "nchw" + + +def _case(label, shape, out_channels, kernel, stride, padding, bias, precision): + rank = len(shape) - 2 + + def expand(v): + return tuple(v) if isinstance(v, tuple) else (v,) * rank + + padding = expand(padding) + return Case( + label, + shape, + out_channels, + expand(kernel), + expand(stride), + padding, + padding, + bias, + precision, + ) + + +# (group, label, input shape, C_out, kernel, stride, padding, bias) +SHAPES = [ + ( + "kernel_stride_3d", + "2x14x14 x256", + (256, 3, 2, 14, 14), + 1280, + (2, 14, 14), + (2, 14, 14), + 0, + False, + ), + ( + "kernel_stride_3d", + "2x14x14 x324", + (324, 3, 2, 14, 14), + 1280, + (2, 14, 14), + (2, 14, 14), + 0, + False, + ), + ( + "kernel_stride_3d", + "2x14x14 x2048", + (2048, 3, 2, 14, 14), + 1280, + (2, 14, 14), + (2, 14, 14), + 0, + False, + ), + ( + "kernel_stride_3d", + "2x16x16 x256", + (256, 3, 2, 16, 16), + 1024, + (2, 16, 16), + (2, 16, 16), + 0, + True, + ), + ( + "kernel_stride_3d", + "2x16x16 x392", + (392, 3, 2, 16, 16), + 1024, + (2, 16, 16), + (2, 16, 16), + 0, + True, + ), + ( + "kernel_stride_3d", + "2x16x16 x512", + (512, 3, 2, 16, 16), + 1024, + (2, 16, 16), + (2, 16, 16), + 0, + True, + ), + ( + "kernel_stride_3d", + "2x16x16 x768", + (768, 3, 2, 16, 16), + 1024, + (2, 16, 16), + (2, 16, 16), + 0, + True, + ), + ( + "kernel_stride_3d", + "1x2x2, 48 ch", + (1, 48, 1, 14, 28), + 3072, + (1, 2, 2), + (1, 2, 2), + 0, + True, + ), + ( + "kernel_stride_3d", + "1x4x4, 2x544x736", + (1, 3, 2, 544, 736), + 256, + (1, 4, 4), + (1, 4, 4), + 0, + True, + ), + ("kernel_stride_2d", "14x14 224x224", (1, 3, 224, 224), 1152, 14, 14, 0, True), + ("kernel_stride_2d", "14x14 448x448", (1, 3, 448, 448), 1152, 14, 14, 0, True), + ("kernel_stride_2d", "14x14 448x448 x3", (3, 3, 448, 448), 1024, 14, 14, 0, True), + ("kernel_stride_2d", "16x16 512x512", (1, 3, 512, 512), 768, 16, 16, 0, True), + ("kernel_stride_2d", "16x16 224x224 x2", (2, 3, 224, 224), 768, 16, 16, 0, False), + ("kernel_stride_2d", "16x16 256x256", (1, 3, 256, 256), 768, 16, 16, 0, True), + ("3ch_7x7", "7x7 s4 224x224 x3", (3, 3, 224, 224), 256, 7, 4, 3, True), + ("3ch_7x7", "7x7 s2 480x640", (1, 3, 480, 640), 64, 7, 2, 3, False), + ("3ch_7x7", "7x7 s2 84x84 x2", (2, 3, 84, 84), 64, 7, 2, 3, False), + ("3ch_7x7", "7x7 s2 84x84 x5", (5, 3, 84, 84), 64, 7, 2, 3, False), + ("3x3", "64 ch 120x160", (1, 64, 120, 160), 64, 3, 1, 1, False), + ("3x3", "128 ch 60x80", (1, 128, 60, 80), 128, 3, 1, 1, False), + ("3x3", "256 ch 30x40", (1, 256, 30, 40), 256, 3, 1, 1, False), + ("3x3", "512 ch 15x20", (1, 512, 15, 20), 512, 3, 1, 1, False), + ("3x3", "64 ch 21x21 x2", (2, 64, 21, 21), 64, 3, 1, 1, False), + ("3x3", "128 ch 11x11 x2", (2, 128, 11, 11), 128, 3, 1, 1, False), + ("3x3", "256 ch 6x6 x2", (2, 256, 6, 6), 256, 3, 1, 1, False), + ("3x3", "512 ch 3x3 x2", (2, 512, 3, 3), 512, 3, 1, 1, False), + ("3x3", "64 ch 21x21 x5", (5, 64, 21, 21), 64, 3, 1, 1, False), + ("3x3", "512 ch 3x3 x5", (5, 512, 3, 3), 512, 3, 1, 1, False), + ("3x3_s2", "64-128 ch 120x160", (1, 64, 120, 160), 128, 3, 2, 1, False), + ("3x3_s2", "128-256 ch 60x80", (1, 128, 60, 80), 256, 3, 2, 1, False), + ("3x3_s2", "256-512 ch 30x40", (1, 256, 30, 40), 512, 3, 2, 1, False), + ("3x3_s2", "64-128 ch 21x21 x2", (2, 64, 21, 21), 128, 3, 2, 1, False), + ("3x3_s2", "128-256 ch 11x11 x2", (2, 128, 11, 11), 256, 3, 2, 1, False), + ("3x3_s2", "256-512 ch 6x6 x2", (2, 256, 6, 6), 512, 3, 2, 1, False), + ("3x3_s2", "256-512 ch 56x56 x3", (3, 256, 56, 56), 512, 3, 2, 1, True), + ("3x3_s2", "512-1024 ch 28x28 x3", (3, 512, 28, 28), 1024, 3, 2, 1, True), + ("3x3_s2", "1024-2048 ch 14x14 x3", (3, 1024, 14, 14), 2048, 3, 2, 1, True), + ("1x1", "512 ch 15x20", (1, 512, 15, 20), 512, 1, 1, 0, True), + ("1x1", "512-32 ch 3x3 x2", (2, 512, 3, 3), 32, 1, 1, 0, True), + ("1x1", "512-32 ch 3x3 x5", (5, 512, 3, 3), 32, 1, 1, 0, True), + ("1x1", "640-1920 ch 14x28", (1, 640, 14, 28), 1920, 1, 1, 0, True), + ("1x1", "640 ch 14x28", (1, 640, 14, 28), 640, 1, 1, 0, True), + ("1x1", "640-1920 ch 8x8 x2", (2, 640, 8, 8), 1920, 1, 1, 0, True), + ("1x1_s2", "64-128 ch 120x160", (1, 64, 120, 160), 128, 1, 2, 0, False), + ("1x1_s2", "128-256 ch 60x80", (1, 128, 60, 80), 256, 1, 2, 0, False), + ("1x1_s2", "256-512 ch 30x40", (1, 256, 30, 40), 512, 1, 2, 0, False), + ("1x1_s2", "64-128 ch 21x21 x2", (2, 64, 21, 21), 128, 1, 2, 0, False), + ("1x1_s2", "128-256 ch 11x11 x2", (2, 128, 11, 11), 256, 1, 2, 0, False), + ("1x1_s2", "256-512 ch 6x6 x2", (2, 256, 6, 6), 512, 1, 2, 0, False), + ("32ch_s2", "3-32 ch 7x7 84x84", (1, 3, 84, 84), 32, 7, 2, 0, True), + ("32ch_s2", "32 ch 5x5 39x39", (1, 32, 39, 39), 32, 5, 2, 0, True), + ("32ch_s2", "32 ch 3x3 18x18", (1, 32, 18, 18), 32, 3, 2, 0, True), + ("32ch_s2", "32 ch 3x3 8x8", (1, 32, 8, 8), 32, 3, 2, 0, True), + ("32ch_s2", "3-32 ch 7x7 256x256 x2", (2, 3, 256, 256), 32, 7, 2, 0, True), + ("32ch_s2", "32 ch 5x5 125x125 x2", (2, 32, 125, 125), 32, 5, 2, 0, True), + ("32ch_s2", "32 ch 3x3 61x61 x2", (2, 32, 61, 61), 32, 3, 2, 0, True), + ("32ch_s2", "32 ch 3x3 30x30 x2", (2, 32, 30, 30), 32, 3, 2, 0, True), +] + + +def shape_cases(precision, layout): + return [ + Case( + **{**asdict(_case(f"{group} {label}", *rest, precision)), "layout": layout} + ) + for group, label, *rest in SHAPES + ] + + +def _envelope_geometries(): + """(axis, input shape, C_out, kernel, stride, padding) around the routed + range (fewer than 16 input channels), 1x1 convolutions, and a few dense + convolutions with more input channels.""" + geometries = [] + for channels in (1, 2, 3, 4, 6, 8, 12, 15, 16, 32): + geometries.append(("channels", (2, channels, 56, 56), 64, 3, 1, 1)) + geometries.append(("channels", (1, channels, 224, 224), 128, 7, 2, 3)) + for kernel, stride, padding in ( + (2, 2, 0), + (3, 2, 0), + (3, 1, 2), + (4, 4, 0), + (5, 2, 2), + (7, 2, 3), + (7, 4, 3), + (8, 8, 0), + (11, 4, 2), + (14, 14, 0), + (16, 16, 0), + (32, 32, 0), + ((1, 7), 1, (0, 3)), + ((3, 5), (2, 1), (1, 2)), + ): + for channels in (3, 12): + geometries.append( + ("kernel", (2, channels, 224, 224), 128, kernel, stride, padding) + ) + for out_channels in (1, 8, 16, 32, 48, 96, 192, 384, 768, 1152, 1536, 3072): + geometries.append(("out_channels", (1, 3, 112, 112), out_channels, 3, 2, 1)) + geometries.append(("out_channels", (1, 3, 224, 224), out_channels, 16, 16, 0)) + for spatial in ( + (1, 1), + (3, 3), + (7, 9), + (33, 47), + (1, 1024), + (1024, 1), + (512, 512), + (1024, 1024), + (2048, 2048), + ): + geometries.append(("spatial", (1, 3, *spatial), 64, 3, 1, 1)) + for batch in (1, 4, 16, 64, 256): + geometries.append(("batch", (batch, 3, 56, 56), 64, 3, 1, 1)) + for shape, out_channels, kernel, stride, padding in ( + ((1, 3, 2, 14, 14), 1280, (2, 14, 14), (2, 14, 14), 0), + ((64, 3, 2, 14, 14), 1280, (2, 14, 14), (2, 14, 14), 0), + ((4096, 3, 2, 16, 16), 1024, (2, 16, 16), (2, 16, 16), 0), + ((1, 3, 8, 256, 256), 256, (1, 4, 4), (1, 4, 4), 0), + ((1, 12, 4, 64, 64), 1024, (1, 2, 2), (1, 2, 2), 0), + ((2, 3, 9, 64, 64), 64, 3, 1, 1), + ((1, 3, 16, 112, 112), 64, (3, 7, 7), (1, 2, 2), (1, 3, 3)), + ((1, 8, 5, 40, 40), 32, (3, 3, 3), (2, 2, 2), 1), + ): + geometries.append(("3-D", shape, out_channels, kernel, stride, padding)) + for channels, out_channels in ((64, 32), (64, 512), (640, 1920)): + for spatial in ((3, 3), (15, 20), (60, 80)): + for stride in (1, 2): + geometries.append( + ("1x1", (2, channels, *spatial), out_channels, 1, stride, 0) + ) + for channels in (16, 64, 256): + geometries.append(("dense", (2, channels, 56, 56), channels, 3, 1, 1)) + return geometries + + +def envelope_cases(precisions): + """``_envelope_geometries`` in each precision, with NCHW and channels-last + inputs, and in FP32 also with a channels-last weight.""" + cases = [] + for label, shape, out_channels, kernel, stride, padding in _envelope_geometries(): + for precision in precisions: + layouts = ["nchw", "channels_last"] + if precision == "fp32": + layouts.append("weight_cl") + for layout in layouts: + case = _case( + label, shape, out_channels, kernel, stride, padding, True, precision + ) + cases.append(Case(**{**asdict(case), "layout": layout})) + return cases + + +def boundary_cases(precision, layout): + """``implicit_gemm.boundary_calls()``: every routed geometry at every corner + of its limits.""" + return [ + Case( + "boundary", + shape, + out_channels, + kernel, + stride, + padding, + padding, + True, + precision, + layout, + ) + for shape, out_channels, kernel, stride, padding in ( + implicit_gemm.boundary_calls() + ) + ] + + +_CALLS_PER_GRAPH = 10 +_REPLAYS = 20 + + +def time_us(fn, flush): + """Per-call device times (us) of ``fn`` replayed in a CUDA graph.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + events = [ + ( + torch.cuda.Event(enable_timing=True, external=True), + torch.cuda.Event(enable_timing=True, external=True), + ) + for _ in range(_CALLS_PER_GRAPH) + ] + with torch.cuda.stream(stream): + fn() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for start, end in events: + if flush is not None: + flush.zero_() + start.record() + fn() + end.record() + graph.replay() + times = [] + for _ in range(_REPLAYS): + graph.replay() + events[-1][1].synchronize() + times += [s.elapsed_time(e) * 1000.0 for s, e in events] + torch.cuda.current_stream().wait_stream(stream) + return times + + +def summarize(times): + deciles = statistics.quantiles(times, n=10) + return statistics.median(times), deciles[0], deciles[-1] + + +def _args(case): + rank = len(case.kernel) + return ( + list(case.stride), + list(case.before), + list(case.after), + [1] * rank, + False, + [0] * rank, + 1, + ) + + +def triton_runner(x, w, b, case): + args = _args(case) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream), autotune_launch_params() as stats: + convolution(x, w, b, *args) + stream.synchronize() + schedule = stats.picks[0].choice if stats.picks else "igemm_ffma" + + def run(): + with force_conv_candidate(schedule): + return convolution(x, w, b, *args) + + return run, schedule + + +def aten_conv(x, w, b, case): + """What the exported graph holds: a separate pad for asymmetric padding.""" + conv = {1: F.conv1d, 2: F.conv2d, 3: F.conv3d}[len(case.kernel)] + if case.before == case.after: + return conv(x, w, b, case.stride, case.before) + pad = [p for pair in zip(reversed(case.before), reversed(case.after)) for p in pair] + return conv(F.pad(x, pad), w, b, case.stride) + + +def inductor_runner(x, w, b, case): + torch._dynamo.reset() + 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(x, w, b) + return lambda: compiled(x, w, b) + + +def kernel_names(fn): + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as prof: + fn() + torch.cuda.synchronize() + return sorted( + { + e.name + for e in prof.events() + if e.device_type == torch.autograd.DeviceType.CUDA + } + ) + + +def _tensors(case): + torch.manual_seed(0) + dtype = {"fp32": torch.float32, "tf32": torch.float32, "fp16": torch.float16}.get( + case.precision, torch.bfloat16 + ) + 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 + ) + b = ( + torch.randn(case.out_channels, device="cuda", dtype=dtype) + if case.bias + else None + ) + memory_format = torch.channels_last if x.dim() == 4 else torch.channels_last_3d + if case.layout == "channels_last": + x = x.to(memory_format=memory_format) + if case.layout == "weight_cl": + w = w.to(memory_format=memory_format) + return x, w, b + + +def run_case(case, flush, threshold): + tf32 = case.precision == "tf32" + torch.backends.cudnn.allow_tf32 = tf32 + torch.backends.cuda.matmul.allow_tf32 = tf32 + if conv_allows_tf32() != tf32: + raise RuntimeError( + f"{case.precision}: this PyTorch's Inductor conv templates " + f"{'cannot' if tf32 else 'must not'} use TF32" + ) + x, w, b = _tensors(case) + ours, schedule = triton_runner(x, w, b, case) + inductor = inductor_runner(x, w, b, case) + reference = aten_conv( + x.double(), w.double(), None if b is None else b.double(), case + ) + error = (ours().double() - reference).abs().max().item() / max( + 1.0, reference.abs().max().item() + ) + results = { + name: summarize(time_us(fn, flush)) + for name, fn in ( + ("ours", ours), + ("inductor", inductor), + ("cudnn", lambda: aten_conv(x, w, b, case)), + ) + } + 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() + speedup = results["inductor"][0] / results["ours"][0] + significant = speedup >= threshold and results["ours"][2] < results["inductor"][1] + return { + **asdict(case), + "schedule": schedule, + "inductor_kernels": kernel_names(inductor), + "max_rel_error": error, + "bytes": nbytes, + "flops": flops, + "speedup_vs_inductor": speedup, + "speedup_vs_cudnn": results["cudnn"][0] / results["ours"][0], + "significantly_faster": significant, + **{ + f"{name}_us": dict(zip(("median", "p10", "p90"), r)) + for name, r in results.items() + }, + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument( + "--suite", default="shapes", choices=["shapes", "envelope", "boundary"] + ) + parser.add_argument( + "--precision", + default="fp32", + help="fp32, tf32, fp16 or bf16; the envelope suite takes a comma-separated list", + ) + parser.add_argument("--layout", default="nchw", choices=["nchw", "channels_last"]) + parser.add_argument("--shard", default="0/1", help="run cases i::n of the suite") + parser.add_argument("--warm", action="store_true", help="do not flush L2") + parser.add_argument("--threshold", type=float, default=1.25) + parser.add_argument("--json", help="append one JSON object per case") + args = parser.parse_args() + + cases = { + "shapes": lambda: shape_cases(args.precision, args.layout), + "envelope": lambda: envelope_cases(args.precision.split(",")), + "boundary": lambda: boundary_cases(args.precision, args.layout), + }[args.suite]() + index, count = map(int, args.shard.split("/")) + l2_bytes = torch.cuda.get_device_properties(0).L2_cache_size + flush = ( + None + if args.warm + else torch.empty(2 * l2_bytes, dtype=torch.uint8, device="cuda") + ) + print( + f"{torch.cuda.get_device_name()} {'warm' if args.warm else 'cold'} L2, " + "median [p10, p90] us" + ) + for case in cases[index::count]: + r = run_case(case, flush, args.threshold) + times = " ".join( + f"{t['median']:8.1f} [{t['p10']:7.1f},{t['p90']:8.1f}]" + for t in (r["ours_us"], r["inductor_us"], r["cudnn_us"]) + ) + print( + f"{case.label:32s} {case.precision:4s} {case.layout:13s} " + f"{str(case.shape):22s} {case.out_channels:5d} {str(case.kernel):12s} " + f"{str(case.stride):12s} {str(case.before) + '/' + str(case.after):20s} " + f"{times} {r['speedup_vs_inductor']:6.2f}x {r['speedup_vs_cudnn']:6.2f}x " + f"{r['flops'] / r['ours_us']['median'] / 1e6:7.2f} {r['schedule']:10s} " + f"{'yes' if r['significantly_faster'] else 'no'}", + flush=True, + ) + if args.json: + with open(args.json, "a") as f: + f.write(json.dumps(r) + "\n") + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/tests/test_conv_implicit_gemm.py b/backends/cuda/tests/test_conv_implicit_gemm.py index 0aac36316db..a727595715a 100644 --- a/backends/cuda/tests/test_conv_implicit_gemm.py +++ b/backends/cuda/tests/test_conv_implicit_gemm.py @@ -5,15 +5,21 @@ # 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.""" +the implicit-GEMM template against an FP64 reference, and models lowered +through the CUDA backend.""" +import glob +import os +import tempfile import unittest from types import SimpleNamespace from unittest import mock import torch +import torch.nn as nn import torch.nn.functional as F +_PASS_LOGGER = "executorch.backends.cuda.triton.replacement_pass" _TOLERANCE = {torch.float32: 2e-5, torch.float16: 2e-3, torch.bfloat16: 1.6e-2} _TF32_TOLERANCE = 5e-3 @@ -321,3 +327,121 @@ def test_follows_the_inductor_conv_template_setting(self): device_interface, "get_interface_for_device", return_value=interface ): self.assertEqual(common.conv_allows_tf32(), allowed) + + +class _KernelEqualsStride(nn.Module): + """A 2-D and a 3-D convolution with kernel == stride.""" + + def __init__(self, dtype): + super().__init__() + self.conv2d = nn.Conv2d(3, 64, 14, stride=14, padding="valid", dtype=dtype) + self.conv3d = nn.Conv3d( + 3, 64, (2, 14, 14), stride=(2, 14, 14), bias=False, dtype=dtype + ) + + def forward(self, x, y): + return self.conv2d(x), self.conv3d(y.view(-1, 3, 2, 14, 14)) + + +class _ThreeChannelBlock(nn.Module): + """3-input-channel 7x7 stride-2 convolutions, padded and unpadded, and + the 1x1 convolutions left to Inductor.""" + + def __init__(self, dtype): + super().__init__() + self.padded = nn.Conv2d(3, 32, 7, stride=2, padding=3, bias=False, dtype=dtype) + self.unpadded = nn.Conv2d(3, 32, 7, stride=2, dtype=dtype) + self.down = nn.Conv2d(32, 64, 1, stride=2, bias=False, dtype=dtype) + self.proj = nn.Conv2d(64, 16, 1, dtype=dtype) + + def forward(self, x): + return self.proj(self.down(F.relu(self.padded(x)))), self.unpadded(x) + + +class ImplicitGemmLoweringTest(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available() or torch.version.hip is not None: + self.skipTest("CUDA required") + + def lower_and_run(self, model, inputs, routed, tolerance): + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + + with torch.no_grad(): + program = torch.export.export(model, inputs, strict=True) + with self.assertLogs(_PASS_LOGGER, "INFO") as logs: + edge = to_edge_transform_and_lower( + program, + partitioner=[ + CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + ], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + output = "\n".join(logs.output) + self.assertEqual(output.count("(optional, implicit_gemm template)"), routed) + self.assertIn(f"Replaced {routed} nodes with Triton kernels", output) + + from executorch.export.utils import is_fbcode + + if is_fbcode(): + # Running the program is skipped in fbcode. + return output + from executorch.runtime import Runtime + + et_program = edge.to_executorch() + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + et_program.write_to_file(f) + et_program.write_tensor_data_to_file(directory) + data_path = glob.glob(os.path.join(directory, "*.ptd"))[0] + method = Runtime.get().load_program(path, data_path=data_path) + outputs = method.load_method("forward").execute([i.cpu() for i in inputs]) + with torch.no_grad(), torch.backends.cudnn.flags( + enabled=True, allow_tf32=False + ): + expected = model(*inputs) + for got, want in zip(outputs, expected): + torch.testing.assert_close( + got.float(), want.float().cpu(), atol=tolerance, rtol=tolerance + ) + return output + + def fp32_tolerance(self): + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + conv_allows_tf32, + ) + + return _TF32_TOLERANCE if conv_allows_tf32() else 1e-4 + + def test_kernel_equals_stride_lowers_and_runs(self): + for dtype, tolerance in ( + (torch.float32, self.fp32_tolerance()), + (torch.bfloat16, 3e-2), + ): + torch.manual_seed(0) + model = _KernelEqualsStride(dtype).to("cuda").eval() + inputs = ( + torch.randn(1, 3, 224, 224, device="cuda", dtype=dtype), + torch.randn(256, 3 * 2 * 14 * 14, device="cuda", dtype=dtype), + ) + with self.subTest(dtype=dtype): + self.lower_and_run(model, inputs, routed=2, tolerance=tolerance) + + def test_three_channel_block_lowers_and_runs(self): + for dtype, tolerance in ( + (torch.float32, self.fp32_tolerance()), + (torch.bfloat16, 3e-2), + ): + torch.manual_seed(0) + model = _ThreeChannelBlock(dtype).to("cuda").eval() + inputs = (torch.randn(2, 3, 84, 84, device="cuda", dtype=dtype),) + with self.subTest(dtype=dtype): + output = self.lower_and_run( + model, inputs, routed=2, tolerance=tolerance + ) + # The 1x1 convolutions keep Inductor's lowering. + self.assertEqual(output.count("Skipping conv replacement"), 2) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index 32bb525b486..345b8383bd2 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/backends/cuda/tests/test_conv_routing.py @@ -14,11 +14,13 @@ import contextlib import dataclasses import importlib.util +import math import unittest from unittest import mock import torch import torch.nn as nn +from torch._subclasses.fake_tensor import FakeTensorMode _PASS_LOGGER = "executorch.backends.cuda.triton.replacement_pass" _APIS = ("to_edge_transform_and_lower", "to_edge_to_backend") @@ -122,7 +124,12 @@ def run_pass(self, graph_module): return _targets(graph_module), "\n".join(logs.output) def assertRouted( - self, model, *inputs, template, canonical_args=None, device="cuda" + self, + model, + *inputs, + template, + canonical_args=None, + device="cuda", ): for api, graph_module in _backend_graphs(model, *inputs, device=device).items(): with self.subTest(api=api): @@ -494,7 +501,7 @@ def test_other_convolutions_are_kept(self): class Dense(nn.Module): def __init__(self): super().__init__() - self.conv = nn.Conv2d(8, 4, 3) + self.conv = nn.Conv2d(8, 4, 1) def forward(self, x): return self.conv(x) @@ -578,3 +585,382 @@ def test_benchmarked_range(self): self.assertKept( model, lambda shape=shape: torch.randn(shape), reason=reason ) + + +def _dense(conv, in_channels, out_channels, kernel, dtype, view=None, **kwargs): + class Dense(nn.Module): + def __init__(self): + super().__init__() + self.conv = conv(in_channels, out_channels, kernel, dtype=dtype, **kwargs) + + def forward(self, x): + return self.conv(x if view is None else x.view(view)) + + return Dense + + +# Benchmarked convolutions in the patterns models write them: (row, model, +# input shape). The 3-input-channel ones are routed; the rest +# keep the regular lowering, and the reason. +_ROUTED_ROWS = ( + ( + "2x14x14 kernel == stride", + lambda dtype: _dense( + nn.Conv3d, + 3, + 1280, + (2, 14, 14), + dtype, + view=(-1, 3, 2, 14, 14), + stride=(2, 14, 14), + bias=False, + ), + (256, 3 * 2 * 14 * 14), + ), + ( + "2x16x16 kernel == stride", + lambda dtype: _dense( + nn.Conv3d, + 3, + 1024, + (2, 16, 16), + dtype, + view=(-1, 3, 2, 16, 16), + stride=(2, 16, 16), + ), + (512, 3 * 2 * 16 * 16), + ), + ( + "1x4x4 kernel == stride", + lambda dtype: _dense(nn.Conv3d, 3, 256, (1, 4, 4), dtype, stride=(1, 4, 4)), + (1, 3, 2, 544, 736), + ), + ( + "14x14 kernel == stride", + lambda dtype: _dense(nn.Conv2d, 3, 1152, 14, dtype, stride=14, padding="valid"), + (1, 3, 224, 224), + ), + ( + "16x16 kernel == stride", + lambda dtype: _dense(nn.Conv2d, 3, 768, 16, dtype, stride=16, bias=False), + (2, 3, 224, 224), + ), + ( + "7x7 s4, 3 channels", + lambda dtype: _dense(nn.Conv2d, 3, 256, 7, dtype, stride=4, padding=3), + (3, 3, 224, 224), + ), + ( + "7x7 s2 p3, 3 channels", + lambda dtype: _dense( + nn.Conv2d, 3, 64, 7, dtype, stride=2, padding=3, bias=False + ), + (1, 3, 480, 640), + ), + ( + "7x7 s2, 3 to 32 channels", + lambda dtype: _dense(nn.Conv2d, 3, 32, 7, dtype, stride=2), + (1, 3, 84, 84), + ), +) +_KEPT_ROWS = ( + ( + "1x2x2 kernel == stride, 48 channels", + lambda dtype: _dense(nn.Conv3d, 48, 3072, (1, 2, 2), dtype, stride=(1, 2, 2)), + (1, 48, 1, 14, 28), + "48 input channels", + ), + ( + "3x3", + lambda dtype: _dense(nn.Conv2d, 64, 64, 3, dtype, padding=1, bias=False), + (1, 64, 120, 160), + "64 input channels", + ), + ( + "3x3 s2", + lambda dtype: _dense( + nn.Conv2d, 256, 512, 3, dtype, stride=2, padding=1, bias=False + ), + (1, 256, 30, 40), + "256 input channels", + ), + ( + "1x1", + lambda dtype: _dense(nn.Conv2d, 512, 512, 1, dtype), + (1, 512, 15, 20), + "1x1 convolution", + ), + ( + "1x1 s2", + lambda dtype: _dense(nn.Conv2d, 64, 128, 1, dtype, stride=2, bias=False), + (1, 64, 120, 160), + "1x1 convolution", + ), + ( + "5x5 s2, 32 channels", + lambda dtype: _dense(nn.Conv2d, 32, 32, 5, dtype, stride=2), + (1, 32, 39, 39), + "32 input channels", + ), +) + + +def _steps_past(families, shape, c_out, kernel, stride, padding): + """The calls one step past each limit a benchmarked corner sits on.""" + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + family = families[(kernel, stride, padding)] + batch, channels, *extent = shape + yield (batch, channels - 1, *extent), c_out, kernel, stride, padding + yield (batch, channels + 1, *extent), c_out, kernel, stride, padding + per_image = implicit_gemm._outputs(1, extent, kernel, stride, padding) + if batch > 1 and (batch - 1) * per_image < implicit_gemm.MIN_ROUTED_OUTPUT_PIXELS: + yield (batch - 1, channels, *extent), c_out, kernel, stride, padding + pixels = batch * math.prod(extent) + if ( + batch == family.max_batch + or pixels + math.prod(extent) > implicit_gemm.MAX_ROUTED_INPUT_PIXELS + ): + yield (batch + 1, channels, *extent), c_out, kernel, stride, padding + for i, e in enumerate(extent): + grown = [*extent[:i], e + 1, *extent[i + 1 :]] + if ( + e == family.max_spatial[i] + or batch * math.prod(grown) > implicit_gemm.MAX_ROUTED_INPUT_PIXELS + ): + yield (batch, channels, *grown), c_out, kernel, stride, padding + outputs = implicit_gemm._outputs(batch, extent, kernel, stride, padding) + most = min( + implicit_gemm.MAX_ROUTED_OUT_CHANNELS, + implicit_gemm.MAX_ROUTED_OUTPUT_ELEMENTS // outputs, + ) + if c_out == most: + yield shape, c_out + 1, kernel, stride, padding + if c_out == implicit_gemm.MIN_ROUTED_OUT_CHANNELS > 1: + yield shape, c_out - 1, kernel, stride, padding + for geometry in ( + (kernel, tuple(s + 1 for s in stride), padding), + (tuple(k + 2 for k in kernel), stride, padding), + (kernel, stride, tuple(p + 1 for p in padding)), + ): + if geometry not in families: + yield (shape, c_out, *geometry) + + +class ImplicitGemmRoutingTest(_PassTest): + """Which dense convolutions the implicit-GEMM template takes.""" + + def setUp(self): + super().setUp() + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + # Exported on an A100, FP32 without TF32 unless a test allows it. + for patch in ( + mock.patch.object(torch.cuda, "get_device_capability", return_value=(8, 0)), + mock.patch.object(implicit_gemm, "conv_allows_tf32", return_value=False), + ): + patch.start() + self.addCleanup(patch.stop) + + def tf32(self, allowed): + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + return mock.patch.object( + implicit_gemm, "conv_allows_tf32", return_value=allowed + ) + + def test_benchmarked_rows(self): + for dtype in (torch.float32, torch.bfloat16): + for row, model, shape in _ROUTED_ROWS: + with self.subTest(row=row, dtype=dtype): + self.assertRouted( + model(dtype), + lambda shape=shape, dtype=dtype: torch.randn( + shape, dtype=dtype + ), + template="implicit_gemm", + ) + for row, model, shape, reason in _KEPT_ROWS: + with self.subTest(row=row, dtype=dtype): + self.assertKept( + model(dtype), + lambda shape=shape, dtype=dtype: torch.randn( + shape, dtype=dtype + ), + reason=reason, + ) + + def test_canonical_arguments(self): + # padding="valid" binds to zero padding on both sides. + self.assertRouted( + _ROUTED_ROWS[3][1](torch.float32), + lambda: torch.randn(1, 3, 224, 224), + template="implicit_gemm", + canonical_args=([14, 14], [0, 0], [0, 0], [1, 1], False, [0, 0], 1), + ) + + def test_precision_without_tensor_cores_is_kept(self): + _, model, shape = _ROUTED_ROWS[6] + with mock.patch.object( + torch.cuda, "get_device_capability", return_value=(7, 5) + ): + for dtype, tf32 in ((torch.bfloat16, False), (torch.float32, True)): + with self.subTest(dtype=dtype, tf32=tf32), self.tf32(tf32): + self.assertKept( + model(dtype), + lambda dtype=dtype: torch.randn(shape, dtype=dtype), + reason="without tensor cores", + ) + self.assertRouted( + model(torch.float32), + lambda: torch.randn(shape), + template="implicit_gemm", + ) + + def worth_routing(self, shape, out_channels, kernel, stride, before, after, **kw): + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + ) + + with FakeTensorMode(): + x = torch.empty(shape, dtype=kw["dtype"], device="cuda") + w = torch.empty( + out_channels, shape[1], *kernel, dtype=kw["dtype"], device="cuda" + ) + x, w = kw.get("layout", lambda x, w: (x, w))(x, w) + rank = len(kernel) + params = ConvParams( + tuple(stride), + tuple(before), + tuple(after), + (1,) * rank, + False, + (0,) * rank, + 1, + ) + with self.tf32(kw["tf32"]): + return implicit_gemm.worth_routing(ConvCall(x, w, None, params)) + + def precisions(self): + """(dtype, TF32 allowed) of the FMA and tensor-core cases.""" + return ( + (torch.float32, False), + (torch.bfloat16, False), + (torch.float16, False), + (torch.float32, True), + ) + + def test_benchmarked_calls_are_routed(self): + # Every call the benchmark's boundary suite measures. + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + for dtype, tf32 in self.precisions(): + for shape, c_out, kernel, stride, padding in implicit_gemm.boundary_calls(): + with self.subTest( + shape=shape, c_out=c_out, kernel=kernel, dtype=dtype, tf32=tf32 + ): + self.assertIsNone( + self.worth_routing( + shape, + c_out, + kernel, + stride, + padding, + padding, + dtype=dtype, + tf32=tf32, + ) + ) + + def test_one_step_past_the_benchmarked_calls_is_kept(self): + """Each benchmarked corner moved one step past a limit it sits on: one + input channel more or less, one image, pixel or output channel more, + one image fewer below the output-pixel minimum, or a stride, kernel or + padding outside the benchmarked geometries.""" + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + families = { + geometry: family + for family in implicit_gemm.ROUTED_FAMILIES + for geometry in family.geometries + } + + for dtype, tf32 in self.precisions(): + for call in implicit_gemm.boundary_calls(): + padding = call[4] + outside = [(*step, step[-1]) for step in _steps_past(families, *call)] + # Asymmetric padding: one more element after. + outside.append((*call[:4], padding, tuple(p + 1 for p in padding))) + for shape, c_out, kernel, stride, before, after in outside: + with self.subTest( + call=call, outside=(shape, c_out, kernel, stride, before, after) + ): + self.assertIsNotNone( + self.worth_routing( + shape, + c_out, + kernel, + stride, + before, + after, + dtype=dtype, + tf32=tf32, + ) + ) + + def test_only_benchmarked_layouts_are_routed(self): + """The benchmark measures dense NCHW and channels-last inputs with + contiguous weights; sliced or padded inputs and channels-last or sliced + weights keep the regular lowering.""" + from executorch.backends.cuda.triton.kernels.conv.core import implicit_gemm + + def channels_last(x, w): + fmt = torch.channels_last if x.dim() == 4 else torch.channels_last_3d + return x.contiguous(memory_format=fmt), w + + def sliced_input(x, w): + return x.as_strided(x.shape, x.stride(), 1), w + + def padded_input(x, w): + n, c, *spatial = x.shape + plane = math.prod(spatial) + 1 + strides = (c * plane, plane, *x.stride()[2:]) + return x.as_strided(x.shape, strides), w + + def channels_last_weight(x, w): + fmt = torch.channels_last if w.dim() == 4 else torch.channels_last_3d + return x, w.contiguous(memory_format=fmt) + + def sliced_weight(x, w): + return x, w.as_strided(w.shape, w.stride(), 1) + + for dtype, tf32 in self.precisions(): + for call in implicit_gemm.boundary_calls(): + shape, c_out, kernel, stride, padding = call + for layout, routed in ( + (channels_last, True), + (sliced_input, False), + (padded_input, False), + (channels_last_weight, False), + (sliced_weight, False), + ): + with self.subTest( + call=call, layout=layout.__name__, dtype=dtype, tf32=tf32 + ): + reason = self.worth_routing( + shape, + c_out, + kernel, + stride, + padding, + padding, + dtype=dtype, + tf32=tf32, + layout=layout, + ) + if routed: + self.assertIsNone(reason) + else: + self.assertIn("not a benchmarked layout", reason) diff --git a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py index 7335f6e1b5e..a2c9b7a2a20 100644 --- a/backends/cuda/triton/kernels/conv/core/implicit_gemm.py +++ b/backends/cuda/triton/kernels/conv/core/implicit_gemm.py @@ -35,6 +35,9 @@ from __future__ import annotations +import itertools +import math +from dataclasses import dataclass from typing import Optional import torch @@ -108,9 +111,184 @@ def supports(call: ConvCall) -> Optional[str]: return None +@dataclass(frozen=True) +class RoutedFamily: + """Geometries benchmarked together: every (kernel, stride, padding) listed, + with symmetric padding, up to ``max_batch`` images of at most + ``max_spatial`` pixels per dimension.""" + + name: str + geometries: tuple[tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]], ...] + max_batch: int + max_spatial: tuple[int, ...] + + +# The range where benchmarks/conv/benchmark_implicit_gemm.py finds every +# measured call significantly faster than Inductor's Triton conv template: +# 3-channel inputs, which the flat K order walks without the empty +# lanes of the template's channel blocks, in the geometry families below, with +# a dense NCHW or channels-last input and a contiguous weight (the layouts the +# benchmark measures). Its ``boundary`` suite measures ``boundary_calls()``: +# each geometry at every corner of these limits. +ROUTED_CHANNELS = 3 +MIN_ROUTED_OUT_CHANNELS = 32 +MAX_ROUTED_OUT_CHANNELS = 1280 +MIN_ROUTED_OUTPUT_PIXELS = 256 +MAX_ROUTED_INPUT_PIXELS = 1024 * 1024 +MAX_ROUTED_OUTPUT_ELEMENTS = 1 << 24 +ROUTED_FAMILIES = ( + RoutedFamily( + "window", + tuple( + ((k, k), (s, s), (p, p)) + for k in (3, 5, 7) + for s in (1, 2, 4) + if s <= k + for p in (0, k // 2) + ), + 256, + (1024, 1024), + ), + RoutedFamily( + "kernel_stride_2d", + tuple(((k, k), (k, k), (0, 0)) for k in (2, 4, 8, 14, 16, 32)), + 256, + (1024, 1024), + ), + RoutedFamily( + "kernel_stride_3d", + tuple((k, k, (0, 0, 0)) for k in ((1, 4, 4), (2, 14, 14), (2, 16, 16))), + 2048, + (4, 1024, 1024), + ), +) + + +def _benchmarked_layout(x: torch.Tensor, w: torch.Tensor) -> bool: + """A dense, unsliced NCHW or channels-last input and a contiguous weight.""" + channels_last = torch.channels_last if x.dim() == 4 else torch.channels_last_3d + dense_input = x.is_contiguous() or x.is_contiguous(memory_format=channels_last) + return ( + dense_input + and x.storage_offset() == 0 + and w.is_contiguous() + and w.storage_offset() == 0 + ) + + +def _family(call: ConvCall) -> Optional[RoutedFamily]: + p = call.params + if p.padding_before != p.padding_after: + return None + geometry = (tuple(call.weight.shape[2:]), p.stride, p.padding_before) + return next((f for f in ROUTED_FAMILIES if geometry in f.geometries), 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" + x, w, p = call.input, call.weight, call.params + 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() + if tensor_cores and not device_supports(SCHEDULES[1].min_capability, x.device): + return ( + f"implicit GEMM: {x.dtype} without tensor cores on the export GPU " + "is not benchmarked" + ) + if x.shape[1] != ROUTED_CHANNELS: + return f"implicit GEMM: {x.shape[1]} input channels are not benchmarked faster" + if not _benchmarked_layout(x, w): + return ( + f"implicit GEMM: input strides {list(x.stride())} (offset " + f"{x.storage_offset()}) or weight strides {list(w.stride())} are not " + "a benchmarked layout" + ) + family = _family(call) + if family is None: + return ( + f"implicit GEMM: kernel {list(w.shape[2:])}, stride {list(p.stride)}, " + f"padding {list(p.padding_before)}/{list(p.padding_after)} is not a " + "benchmarked geometry" + ) + outputs = x.shape[0] * math.prod(_output_size(call)) + if ( + x.shape[0] > family.max_batch + or any(e > m for e, m in zip(x.shape[2:], family.max_spatial)) + or x.shape[0] * math.prod(x.shape[2:]) > MAX_ROUTED_INPUT_PIXELS + or not MIN_ROUTED_OUT_CHANNELS <= w.shape[0] <= MAX_ROUTED_OUT_CHANNELS + or outputs < MIN_ROUTED_OUTPUT_PIXELS + or outputs * w.shape[0] > MAX_ROUTED_OUTPUT_ELEMENTS + ): + return ( + f"implicit GEMM: input {list(x.shape)} with {w.shape[0]} output " + f"channels is outside the benchmarked {family.name} range" + ) + return None + + +def _outputs(batch, extent, kernel, stride, padding) -> int: + return batch * math.prod( + (e + 2 * p - k) // s + 1 for e, k, s, p in zip(extent, kernel, stride, padding) + ) + + +def boundary_calls() -> ( + list[tuple[tuple[int, ...], int, tuple[int, ...], tuple[int, ...], tuple[int, ...]]] +): + """(input shape, C_out, kernel, stride, padding) of every routed geometry at + every corner of its limits: each spatial dimension at its smallest (one + output) or largest extent with the fewest images that give + ``MIN_ROUTED_OUTPUT_PIXELS`` outputs, and the most images of the smallest + extent; extents halved until the input fits ``MAX_ROUTED_INPUT_PIXELS`` and + the output fits ``MIN_ROUTED_OUT_CHANNELS``; with the fewest and the most + output channels the limits allow.""" + calls = [] + for family in ROUTED_FAMILIES: + for kernel, stride, padding in family.geometries: + smallest = tuple(k - 2 * p for k, p in zip(kernel, padding)) + corners = [ + (None, e) for e in itertools.product(*zip(smallest, family.max_spatial)) + ] + corners.append((family.max_batch, smallest)) + for batch, extent in corners: + extent = list(extent) + while True: + images = batch + per_image = _outputs(1, extent, kernel, stride, padding) + if images is None: + images = -(-MIN_ROUTED_OUTPUT_PIXELS // per_image) + images = min( + images, + family.max_batch, + MAX_ROUTED_INPUT_PIXELS // math.prod(extent), + ) + outputs = images * per_image + fits = ( + math.prod(extent) <= MAX_ROUTED_INPUT_PIXELS + and MAX_ROUTED_OUTPUT_ELEMENTS // max(outputs, 1) + >= MIN_ROUTED_OUT_CHANNELS + ) + if fits or list(extent) == list(smallest): + break + i = max(range(len(extent)), key=lambda d: extent[d] - smallest[d]) + extent[i] = max(smallest[i], extent[i] // 2) + batch = images + if outputs < MIN_ROUTED_OUTPUT_PIXELS: + continue + most = min( + MAX_ROUTED_OUT_CHANNELS, MAX_ROUTED_OUTPUT_ELEMENTS // outputs + ) + for out_channels in (MIN_ROUTED_OUT_CHANNELS, most): + calls.append( + ( + (batch, ROUTED_CHANNELS, *extent), + out_channels, + kernel, + stride, + padding, + ) + ) + return list(dict.fromkeys(calls)) def schedule_supports(name: str, call: ConvCall) -> Optional[str]: