diff --git a/.ci/scripts/test_cuda_clean_runtime.sh b/.ci/scripts/test_cuda_clean_runtime.sh new file mode 100755 index 00000000000..9a231c764fc --- /dev/null +++ b/.ci/scripts/test_cuda_clean_runtime.sh @@ -0,0 +1,118 @@ +#!/bin/bash +# 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. + +# Runs an exported CUDA program twice with nothing but the prebuilt +# executor_runner and its runtime libraries: empty HOME, TMPDIR, working +# directory and Triton / Inductor / CUDA cache directories, no Python or +# compiler on PATH and no other environment variables. Checks that +# - the runner links no PyTorch, Python or Triton library, +# - neither run starts another program or opens a file of a Python, PyTorch or +# Triton installation, a compiler, or a cache, +# - neither run leaves a file behind (the runtime removes the shared library +# it extracts from the .pte), and +# - both runs write the same output, which matches the eager reference. +# +# Usage: test_cuda_clean_runtime.sh RUNNER MODEL_DIR MODEL_NAME [DTYPE] +# MODEL_DIR holds MODEL_NAME.pte, aoti_cuda_blob.ptd and the +# MODEL_NAME_input_.bin / MODEL_NAME_expected_0.bin files written by +# `python -m examples.cuda.scripts.export --save_example_io`. DTYPE is the +# numpy dtype of the output (default float32). + +set -euo pipefail + +runner="$(realpath "$1")" +model_dir="$(realpath "$2")" +model="$3" +dtype="${4:-float32}" + +fail() { + echo "::error::$*" + exit 1 +} + +strace_bin="$(command -v strace || true)" +if [ -z "${strace_bin}" ]; then + conda install -y -c conda-forge strace + strace_bin="$(command -v strace)" +fi + +work="$(mktemp -d)" +mkdir -p "${work}"/{home,tmp,cwd,out,bin} "${work}"/cache/{xdg,triton,inductor,cuda} +inputs="$(find "${model_dir}" -maxdepth 1 -name "${model}_input_*.bin" | sort -V | paste -sd, -)" + +ldd "${runner}" | tee "${work}/ldd.txt" +if grep -E "libtorch|libc10|libpython|libtriton|not found" "${work}/ldd.txt"; then + fail "executor_runner must resolve every library and link no PyTorch, Python or Triton library" +fi +# The directories the runner's libraries resolve to, plus the inherited ones +# (e.g. where the container mounts the CUDA driver) that are not part of a +# Python installation. The system's default directories are left out: the +# loader searches them last anyway, and listed first, their older libstdc++ +# would shadow the newer one the program's shared library needs. +library_path="$( + { + awk '$3 ~ /^\// {print $3}' "${work}/ldd.txt" | xargs -n1 dirname + tr ':' '\n' <<<"${LD_LIBRARY_PATH:-}" | grep -vE 'site-packages|/python[0-9.]*/' || true + } | grep -vE '^$|^/(usr/)?lib(64)?(/x86_64-linux-gnu)?/?$' | awk '!seen[$0]++' | paste -sd: - +)" +# strace may itself come from a conda environment. +strace_library_path="$(dirname "$(dirname "${strace_bin}")")/lib" + +forbidden='site-packages|/python[0-9.]*/|libpython|libtorch|libc10|triton|inductor|nvcc|ptxas|nvrtc|/bin/(cc|c\+\+|gcc|g\+\+|clang|ld)$' + +for run in 1 2; do + trace="${work}/strace.${run}" + ( + cd "${work}/cwd" + env -i \ + HOME="${work}/home" \ + TMPDIR="${work}/tmp" \ + XDG_CACHE_HOME="${work}/cache/xdg" \ + TRITON_CACHE_DIR="${work}/cache/triton" \ + TORCHINDUCTOR_CACHE_DIR="${work}/cache/inductor" \ + CUDA_CACHE_PATH="${work}/cache/cuda" \ + CUDA_CACHE_DISABLE=1 \ + PATH="${work}/bin" \ + LD_LIBRARY_PATH="${library_path}:${strace_library_path}" \ + "${strace_bin}" -f -qq -o "${trace}" \ + -e trace=execve,execveat,open,openat,creat,mkdir,mkdirat,rename,renameat,renameat2,link,linkat,symlink,symlinkat \ + "${runner}" \ + --model_path "${model_dir}/${model}.pte" \ + --data_path "${model_dir}/aoti_cuda_blob.ptd" \ + --inputs "${inputs}" \ + --output_file "${work}/out/run${run}" \ + --print_output none + ) + + execs="$(grep -E 'execve(at)?\(' "${trace}" | grep -v "\"${runner}\"" || true)" + [ -z "${execs}" ] || fail "run ${run} started another program: ${execs}" + opened="$(grep -E '(open|openat|creat)\(' "${trace}" | grep -vE '= -1 ' \ + | grep -oE '"[^"]*"' | tr -d '"' | sort -u)" + echo "run ${run} opened:" + echo "${opened}" + bad="$(grep -E "${forbidden}" <<<"${opened}" || true)" + [ -z "${bad}" ] || fail "run ${run} opened: ${bad}" + written="$(grep -E '(open|openat)\(.*O_(WRONLY|RDWR|CREAT)|creat\(|mkdir|rename|link\(' "${trace}" \ + | grep -vE '= -1 ' | grep -oE '"[^"]*"' | tr -d '"' \ + | grep -vE "^(${work}/(tmp|out)/|/dev/|/proc/self/)" || true)" + [ -z "${written}" ] || fail "run ${run} wrote outside its output and temporary directories: ${written}" + left="$(find "${work}/home" "${work}/tmp" "${work}/cwd" "${work}"/cache/* -mindepth 1)" + [ -z "${left}" ] || fail "run ${run} left files behind: ${left}" +done + +cmp "${work}/out/run1-0.bin" "${work}/out/run2-0.bin" +python - "${work}/out/run1-0.bin" "${model_dir}/${model}_expected_0.bin" "${dtype}" <<'PY' +import sys + +import numpy as np + +output, expected = (np.fromfile(path, dtype=sys.argv[3]) for path in sys.argv[1:3]) +assert output.shape == expected.shape, (output.shape, expected.shape) +np.testing.assert_allclose(output, expected, rtol=1e-4, atol=1e-4) +print(f"output matches the eager reference ({output.size} values)") +PY +echo "Clean runtime check passed for ${model}" diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 4675e757c6f..a392c10b05d 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -237,7 +237,7 @@ jobs: -Bcmake-out . cmake --build cmake-out -j4 - for model in linear add add_mul resnet18 conv1d sdpa mv2 mv3; do + for model in linear add add_mul resnet18 conv1d depthwise_conv sdpa mv2 mv3; do python -m examples.cuda.scripts.export --model_name="$model" --output_dir "./" ./cmake-out/executor_runner --model_path "./${model}.pte" --data_path "./aoti_cuda_blob.ptd" done @@ -293,7 +293,9 @@ jobs: # Export the explicit PTX fallback separately. A10G supports 99 KiB of # opt-in shared memory per block, so constrain the A100 export to it. - for model in linear sdpa; do + # depthwise_conv also carries its example input and eager output, for + # the clean-environment runs of the consuming jobs. + for model in linear sdpa depthwise_conv; do model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" mkdir -p "${model_dir}" ET_CUDA_TARGET_SMEM_BYTES=101376 \ @@ -302,8 +304,12 @@ jobs: --model_name="${model}" \ --output_dir="${model_dir}" \ --cuda_include_ptx=ON \ - --seed=0 + --seed=0 \ + --save_example_io done + # The depthwise convolution must run as the backend's Triton kernel. + grep -aqE "_(flat|row8)_kernel" \ + "${RUNNER_ARTIFACT_DIR}/fallback-sm80/depthwise_conv/depthwise_conv.pte" test-cuda-target-smem-cross-arch: name: merge-and-test-cuda-multi-arch-a10g @@ -381,13 +387,18 @@ jobs: -Bcmake-out . cmake --build cmake-out --target executor_runner -j4 - for model in linear sdpa; do + for model in linear sdpa depthwise_conv; do model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" ./cmake-out/executor_runner \ --model_path "${model_dir}/${model}.pte" \ --data_path "${model_dir}/aoti_cuda_blob.ptd" done + # The A100 export runs here from its PTX, twice, with only the runner + # and its libraries in an empty environment. + bash .ci/scripts/test_cuda_clean_runtime.sh ./cmake-out/executor_runner \ + "${RUNNER_ARTIFACT_DIR}/fallback-sm80/depthwise_conv" depthwise_conv + # The local native PTE and merged PTE must run on sm86. ./cmake-out/executor_runner \ --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ @@ -454,6 +465,11 @@ jobs: cmp "${sm80_ptd}" "${fallback_ptd}" cmp "${sm80_ptd}" "${merged_ptd}" + # The A100 export runs natively on this fresh runner, twice, with only + # the runner and its libraries in an empty environment. + bash .ci/scripts/test_cuda_clean_runtime.sh ./cmake-out/executor_runner \ + "${RUNNER_ARTIFACT_DIR}/fallback-sm80/depthwise_conv" depthwise_conv + # The local native PTE and merged PTE must run on sm80. ./cmake-out/executor_runner \ --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 80c705bcc0d..1da2346bc15 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -173,6 +173,7 @@ fbcode_target( "triton/kernels/conv/conv.py", "triton/kernels/conv/core/__init__.py", "triton/kernels/conv/core/common.py", + "triton/kernels/conv/core/direct.py", "triton/kernels/conv/routing.py", "triton/kernels/fused_moe.py", "triton/kernels/int4_quantized_gemm.py", diff --git a/backends/cuda/benchmarks/conv/benchmark_depthwise.py b/backends/cuda/benchmarks/conv/benchmark_depthwise.py new file mode 100644 index 00000000000..c92dfbfa03e --- /dev/null +++ b/backends/cuda/benchmarks/conv/benchmark_depthwise.py @@ -0,0 +1,528 @@ +#!/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 triton::convolution on depthwise and few-channel grouped shapes +against Inductor's Triton conv template (what the CUDA backend lowers a +convolution to otherwise) and cuDNN; grouped shapes also against a sum over +kernel taps, a decomposition of grouped convolutions into elementwise products +and sums that Inductor compiles. + +Each call is timed inside a CUDA graph, preceded by a write larger than L2 +(cold L2), with timing events around the call only; the graph is replayed +several times and the median and the 10th-90th percentile spread of the +per-call times are reported. ``--warm`` times without the flush. + +triton::convolution runs with the schedule and config the backend's compile +would pick: the schedules are timed once under ``autotune_launch_params`` and +Triton's autotuner picks the config. Inductor compiles the convolution with +the backend's options (``max_autotune``, Triton-only conv and GEMM backends). + +A shape is "significantly faster" when the median speedup over Inductor is at +least ``--threshold`` (default 1.25x) and our 90th percentile is below +Inductor's 10th. GB/s counts the compulsory traffic: one read of the input, +weight and bias, one write of the output. + +Suites: ``representative`` (model shapes, one dtype and layout per run), +``envelope`` (each axis of what the template supports varied around base +shapes, every dtype and layout), ``extent`` (the largest shapes and the +remaining weight layouts) and ``grouped`` (1 to 8 input channels per group, +channel multipliers 1 to 4). The router's ``direct.worth_routing`` takes the +range these suites cover. For grouped shapes a case is significantly faster +when it is against both the template and the tap sum. + +Usage: + python benchmark_depthwise.py + python benchmark_depthwise.py --dtype float16 --layout channels_last + python benchmark_depthwise.py --suite envelope --shard 0/8 --json results.jsonl + python benchmark_depthwise.py --suite grouped --json grouped.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, +) + +# (label, (N, C, H, W), kernel, stride, padding) +SHAPES = [ + ("3x3 x3", (3, 256, 56, 56), 3, 1, 1), + ("3x3 x3", (3, 512, 28, 28), 3, 1, 1), + ("3x3 x3", (3, 1024, 14, 14), 3, 1, 1), + ("3x3 x3", (3, 2048, 7, 7), 3, 1, 1), + ("large 7x7", (1, 128, 128, 160), 7, 1, 3), + ("7x7", (2, 160, 32, 40), 7, 1, 3), + ("9x9", (2, 304, 16, 20), 9, 1, 4), + ("5x5", (2, 96, 64, 80), 5, 1, 2), + ("3x3", (2, 48, 128, 160), 3, 1, 1), + ("3x3", (2, 48, 64, 80), 3, 1, 1), + ("3x3", (2, 54, 32, 40), 3, 1, 1), + ("3x3", (2, 76, 16, 20), 3, 1, 1), + ("7x7 x8", (8, 96, 56, 56), 7, 1, 3), + ("7x7 x8", (8, 192, 28, 28), 7, 1, 3), + ("7x7 x8", (8, 384, 14, 14), 7, 1, 3), + ("7x7 x8", (8, 768, 7, 7), 7, 1, 3), + ("3x3 x8", (8, 32, 112, 112), 3, 1, 1), + ("3x3 s2 x8", (8, 96, 112, 112), 3, 2, 1), + ("3x3 x8", (8, 144, 56, 56), 3, 1, 1), + ("3x3 s2 x8", (8, 144, 56, 56), 3, 2, 1), + ("3x3 x8", (8, 192, 28, 28), 3, 1, 1), + ("3x3 x8", (8, 384, 14, 14), 3, 1, 1), + ("3x3 x8", (8, 576, 14, 14), 3, 1, 1), + ("3x3 x8", (8, 960, 7, 7), 3, 1, 1), + ("batch 1 3x3", (1, 32, 8, 8), 3, 1, 1), +] + + +@dataclass(frozen=True) +class Case: + label: str + shape: tuple[int, int, int, int] + kernel: tuple[int, int] + stride: tuple[int, int] + before: tuple[int, int] + after: tuple[int, int] + dtype: str = "float32" + layout: str = "nchw" + # Input channels per group and output channels per group; (1, 1) is a + # depthwise convolution. + group_channels: int = 1 + multiplier: int = 1 + + @property + def groups(self): + return self.shape[1] // self.group_channels + + +def representative_cases(dtype, layout): + return [ + Case(label, shape, (k, k), (s, s), (p, p), (p, p), dtype, layout) + for label, shape, k, s, p in SHAPES + ] + + +def _same(kernel): + before = tuple((k - 1) // 2 for k in kernel) + return before, tuple(k - 1 - b for k, b in zip(kernel, before)) + + +def envelope_cases(): + """Shapes spanning what ``direct.supports()`` accepts: each axis (kernel, + stride, padding, channels, spatial size, batch) is varied around a few + base shapes, in every dtype with NCHW and channels-last inputs, plus + channels-last weights in fp32.""" + geometries = [] + for base in ((1, 16, 16, 20), (2, 64, 56, 56), (8, 256, 28, 28)): + for kernel in ( + (1, 1), + (3, 3), + (5, 5), + (7, 7), + (9, 9), + (1, 9), + (9, 1), + (3, 5), + (7, 3), + (4, 4), + (2, 6), + ): + geometries.append(("kernel", base, kernel, (1, 1), *_same(kernel))) + for stride in ((2, 2), (3, 3), (1, 2), (2, 1), (3, 1)): + geometries.append(("stride", base, (3, 3), stride, (1, 1), (1, 1))) + for before, after in ( + ((0, 0), (0, 0)), + ((4, 4), (4, 4)), + ((0, 0), (4, 4)), + ((4, 4), (0, 0)), + ((1, 3), (3, 0)), + ): + geometries.append(("padding", base, (5, 5), (1, 1), before, after)) + for channels in (1, 2, 3, 8, 32, 96, 512, 2048): + geometries.append( + ("channels", (2, channels, 28, 28), (3, 3), (1, 1), (1, 1), (1, 1)) + ) + for spatial in ( + (1, 1), + (2, 2), + (7, 7), + (1, 1024), + (1024, 1), + (256, 256), + (512, 512), + ): + geometries.append( + ("spatial", (1, 32, *spatial), (3, 3), (1, 1), (1, 1), (1, 1)) + ) + for batch in (4, 32, 128): + geometries.append( + ("batch", (batch, 64, 28, 28), (3, 3), (1, 1), (1, 1), (1, 1)) + ) + cases = [] + for label, shape, kernel, stride, before, after in geometries: + for dtype in ("float32", "float16", "bfloat16"): + for layout in ("nchw", "channels_last"): + cases.append( + Case(label, shape, kernel, stride, before, after, dtype, layout) + ) + cases.append( + Case(label, shape, kernel, stride, before, after, "float32", "weight_cl") + ) + return cases + + +def extent_cases(): + """The largest shapes and the remaining weight layouts the router takes: + 4096 channels, 2048x2048 images, batch 256, ~50M-element inputs, and + channels-last weights with NCHW or channels-last inputs in every dtype.""" + cases = [] + for shape, kernel in ( + ((1, 4096, 32, 32), (3, 3)), + ((2, 4096, 14, 14), (7, 7)), + ((1, 16, 1024, 1024), (3, 3)), + ((1, 4, 2048, 2048), (5, 5)), + ((256, 32, 28, 28), (3, 3)), + ((8, 96, 224, 224), (3, 3)), + ((32, 128, 112, 112), (3, 3)), + ((4, 512, 96, 96), (9, 9)), + ): + for dtype in ("float32", "float16"): + for layout in ("nchw", "channels_last"): + cases.append( + Case("extent", shape, kernel, (1, 1), *_same(kernel), dtype, layout) + ) + for kernel in ((3, 3), (7, 7), (1, 9), (4, 4)): + for dtype in ("float32", "float16", "bfloat16"): + for layout in ("weight_cl", "both_cl"): + if (dtype, layout) != ("float32", "weight_cl"): + cases.append( + Case( + "weights", + (2, 64, 56, 56), + kernel, + (1, 1), + *_same(kernel), + dtype, + layout, + ) + ) + return cases + + +def grouped_cases(): + """Grouped convolutions with 1 to 8 input channels per group and channel + multipliers 1 to 4 at one size: every combination with 3x3, 5x5 and 7x7 + kernels and strides 1 and 2 in fp32 and bf16, odd group widths, and + channels-last inputs; then four combinations from the smallest to the + largest sizes.""" + shape = (8, 192, 28, 28) + combinations = [(g, m) for g in (1, 2, 4, 8) for m in (1, 2, 4) if (g, m) != (1, 1)] + geometries = [ + (g, m, k, s, "nchw") for g, m in combinations for k in (3, 5, 7) for s in (1, 2) + ] + geometries += [(g, m, 3, s, "nchw") for g in (3, 6) for m in (1, 2) for s in (1, 2)] + geometries += [(g, m, 3, 1, "channels_last") for g, m in combinations] + sized = [(shape, *geometry) for geometry in geometries] + for size in ( + (1, 24, 8, 8), + (32, 192, 28, 28), + (1, 96, 224, 224), + (2, 1536, 14, 14), + (8, 384, 112, 112), + ): + sized += [ + (size, g, m, k, 1, "nchw") + for g, m in ((1, 2), (2, 1), (4, 2), (8, 2)) + for k in (3, 7) + ] + return [ + Case( + f"g{g} m{m}", + size, + (k, k), + (s, s), + (k // 2, k // 2), + (k // 2, k // 2), + dtype, + layout, + g, + m, + ) + for size, g, m, k, s, layout in sized + for dtype in ("float32", "bfloat16") + ] + + +_MAX_TAP_SUM_GROUP_CHANNELS = 8 + + +def tap_sum_conv2d(input, weight, bias, stride, padding, groups): + """A grouped convolution with few channels per group as a sum over kernel + taps of elementwise products, accumulated in fp32 (symmetric padding, no + dilation).""" + if groups == 1 or weight.shape[1] > _MAX_TAP_SUM_GROUP_CHANNELS: + return F.conv2d(input, weight, bias, stride, padding, 1, groups) + dtype = input.dtype + x = input.float() + batch = x.shape[0] + out_channels, group_channels, kernel_h, kernel_w = weight.shape + out_h, out_w = ( + (size + 2 * pad - kernel) // step + 1 + for size, pad, kernel, step in zip( + x.shape[2:], padding, (kernel_h, kernel_w), stride + ) + ) + if any(padding): + x = F.pad(x, [padding[1], padding[1], padding[0], padding[0]]) + x = x.reshape(batch, groups, 1, group_channels, x.shape[2], x.shape[3]) + w = weight.float().reshape( + 1, groups, out_channels // groups, group_channels, kernel_h, kernel_w + ) + total = None + for i in range(kernel_h): + for j in range(kernel_w): + window = x[ + ..., + i : i + stride[0] * (out_h - 1) + 1 : stride[0], + j : j + stride[1] * (out_w - 1) + 1 : stride[1], + ] + term = (window * w[..., i, j, None, None]).sum(3) + total = term if total is None else total + term + out = total.reshape(batch, out_channels, out_h, out_w) + if bias is not None: + out = out + bias.float().reshape(1, out_channels, 1, 1) + return out.to(dtype) + + +_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 _pad(x, case): + return F.pad(x, [case.before[1], case.after[1], case.before[0], case.after[0]]) + + +def triton_runner(x, w, b, case): + args = ( + list(case.stride), + list(case.before), + list(case.after), + [1, 1], + False, + [0, 0], + case.groups, + ) + 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 "dw_flat" + + 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.""" + if case.before == case.after: + return F.conv2d(x, w, b, case.stride, case.before, 1, case.groups) + return F.conv2d(_pad(x, case), w, b, case.stride, 0, 1, case.groups) + + +def tap_sum(x, w, b, case): + return tap_sum_conv2d(x, w, b, case.stride, case.before, case.groups) + + +# The Inductor options of the backend's AOTI compile that affect kernels. +INDUCTOR_OPTIONS = { + "loop_ordering_after_fusion": False, + "emulate_precision_casts": True, + "max_autotune": True, + "max_autotune_conv_backends": "TRITON", + "max_autotune_gemm_backends": "TRITON", +} + + +def inductor_runner(x, w, b, case, conv=aten_conv): + torch._dynamo.reset() + with torch._inductor.config.patch(INDUCTOR_OPTIONS): + compiled = torch.compile( + lambda x, w, b: conv(x, w, b, case), dynamic=False, fullgraph=True + ) + compiled(x, w, b) + return lambda: compiled(x, w, b) + + +def run_case(case, flush, threshold): + torch.manual_seed(0) + dtype = getattr(torch, case.dtype) + out_channels = case.groups * case.multiplier + x = torch.randn(case.shape, device="cuda", dtype=dtype) + w = torch.randn( + out_channels, case.group_channels, *case.kernel, device="cuda", dtype=dtype + ) + w = w / (case.group_channels * case.kernel[0] * case.kernel[1]) ** 0.5 + b = torch.randn(out_channels, device="cuda", dtype=dtype) + if case.layout in ("channels_last", "both_cl"): + x = x.to(memory_format=torch.channels_last) + if case.layout in ("weight_cl", "both_cl"): + w = w.to(memory_format=torch.channels_last) + ours, schedule = triton_runner(x, w, b, case) + runners = [ + ("ours", ours), + ("inductor", inductor_runner(x, w, b, case)), + ("cudnn", lambda: aten_conv(x, w, b, case)), + ] + grouped = case.groups != case.shape[1] or case.multiplier != 1 + if grouped: + runners.append(("tap_sum", inductor_runner(x, w, b, case, tap_sum))) + reference = aten_conv(x.double(), w.double(), b.double(), case) + error = (ours().double() - reference).abs().max().item() + results = {name: summarize(time_us(fn, flush)) for name, fn in runners} + nbytes = sum(t.numel() * t.element_size() for t in (x, w, b, ours())) + speedup = results["inductor"][0] / results["ours"][0] + significant = all( + results[name][0] / results["ours"][0] >= threshold + and results["ours"][2] < results[name][1] + for name in ("inductor", "tap_sum") + if name in results + ) + extra = {} + if grouped: + extra["speedup_vs_tap_sum"] = results["tap_sum"][0] / results["ours"][0] + return { + **asdict(case), + "schedule": schedule, + "max_abs_error": error, + "bytes": nbytes, + "speedup_vs_inductor": speedup, + "speedup_vs_cudnn": results["cudnn"][0] / results["ours"][0], + **extra, + "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="representative", + choices=["representative", "envelope", "extent", "grouped"], + ) + parser.add_argument( + "--dtype", default="float32", choices=["float32", "float16", "bfloat16"] + ) + 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() + + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + cases = { + "representative": lambda: representative_cases(args.dtype, args.layout), + "envelope": envelope_cases, + "extent": extent_cases, + "grouped": grouped_cases, + }[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" + ) + print( + f"{'case':10s} {'dtype':9s} {'layout':13s} {'input':20s} {'k':6s} {'s':6s} " + f"{'pad (before/after)':18s} {'ours':>22s} {'inductor':>22s} {'cudnn':>22s} " + f"{'vs ind':>8s} {'vs cudnn':>8s} {'GB/s':>6s} schedule significant " + "[tap sum, vs tap sum]" + ) + for case in cases[index::count]: + r = run_case(case, flush, args.threshold) + times = " ".join( + f"{t['median']:7.1f} [{t['p10']:5.1f},{t['p90']:6.1f}]" + for t in (r["ours_us"], r["inductor_us"], r["cudnn_us"]) + ) + geometry = ( + f"{case.kernel[0]}x{case.kernel[1]:<4d} {case.stride[0]}x{case.stride[1]:<4d} " + f"{str(case.before) + '/' + str(case.after):18s}" + ) + print( + f"{case.label:10s} {case.dtype:9s} {case.layout:13s} {str(case.shape):20s} " + f"{geometry} {times} {r['speedup_vs_inductor']:7.2f}x " + f"{r['speedup_vs_cudnn']:7.2f}x {r['bytes'] / r['ours_us']['median'] / 1e3:6.0f} " + f"{r['schedule']:9s} {'yes' if r['significantly_faster'] else 'no ':11s}" + + ( + f" {r['tap_sum_us']['median']:7.1f} {r['speedup_vs_tap_sum']:7.2f}x" + if "tap_sum_us" in r + else "" + ), + 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/targets.bzl b/backends/cuda/tests/targets.bzl index 3829ded98ed..163b610bae3 100644 --- a/backends/cuda/tests/targets.bzl +++ b/backends/cuda/tests/targets.bzl @@ -97,6 +97,30 @@ def define_common_targets(is_fbcode = False): ), ) + python_unittest_remote_gpu( + name = "test_conv_depthwise", + srcs = [ + "test_conv_depthwise.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_depthwise.py b/backends/cuda/tests/test_conv_depthwise.py new file mode 100644 index 00000000000..c2be0b968f9 --- /dev/null +++ b/backends/cuda/tests/test_conv_depthwise.py @@ -0,0 +1,523 @@ +# 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 2-D depthwise convolutions and grouped convolutions +with few channels per group: every schedule of the direct-stencil template +against an FP64 reference, and models lowered through the CUDA backend.""" + +import glob +import os +import tempfile +import unittest +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} + + +def _reference(x, w, b, stride, before, after): + padded = F.pad(x.double(), [before[1], after[1], before[0], after[0]]) + bias = None if b is None else b.double() + return F.conv2d(padded, w.double(), bias, stride, 0, 1, x.shape[1] // w.shape[1]) + + +class DepthwiseScheduleTest(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) + + def convolution(self, x, w, b, stride, before, after): + from executorch.backends.cuda.triton.kernels.conv.conv import convolution + + return convolution( + x, + w, + b, + list(stride), + list(before), + list(after), + [1, 1], + False, + [0, 0], + x.shape[1] // w.shape[1], + ) + + def check(self, x, w, b, stride, before, after, every_config=False): + """Runs every schedule that supports the call, with its first config or + every config, and checks the result and the layout.""" + from executorch.backends.cuda.triton.kernels.conv.conv import ( + force_conv_candidate, + ) + from executorch.backends.cuda.triton.kernels.conv.core import direct + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + out_shape_and_stride, + ) + + params = ConvParams( + tuple(stride), + tuple(before), + tuple(after), + (1, 1), + False, + (0, 0), + x.shape[1] // w.shape[1], + ) + call = ConvCall(x, w, b, params) + self.assertIsNone(direct.supports(call)) + expected = _reference(x, w, b, stride, before, after).to(x.dtype) + layout = out_shape_and_stride(call) + ran = [] + for schedule, all_configs in ( + ("dw_flat", direct.FLAT_CONFIGS), + ("dw_row8", direct.ROW8_CONFIGS), + ("dw_mult", direct.MULT_CONFIGS), + ("grp_plane", direct.PLANE_CONFIGS), + ("grp_reduce", direct.REDUCE_CONFIGS), + ): + if direct.schedule_supports(schedule, call) is not None: + with force_conv_candidate(schedule), self.assertRaises(ValueError): + self.convolution(x, w, b, stride, before, after) + continue + for config in all_configs if every_config else all_configs[:1]: + with force_conv_candidate(schedule, config): + out = self.convolution(x, w, b, stride, before, after) + self.assertEqual((tuple(out.shape), out.stride()), layout) + tolerance = _TOLERANCE[x.dtype] + torch.testing.assert_close( + out.double(), + expected.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}", + ) + ran.append(schedule) + self.assertIn("dw_flat", ran) + return ran + + def tensors( + self, + n, + c, + h, + w, + kh, + kw, + dtype=torch.float32, + bias=True, + group_channels=1, + multiplier=1, + ): + """Input, weight and bias of a convolution with ``group_channels`` input + and ``multiplier`` output channels per group.""" + out_channels = c // group_channels * multiplier + x = torch.randn(n, c, h, w, device="cuda", dtype=dtype) + weight = ( + torch.randn( + out_channels, group_channels, kh, kw, device="cuda", dtype=dtype + ) + / (group_channels * kh * kw) ** 0.5 + ) + b = torch.randn(out_channels, device="cuda", dtype=dtype) if bias else None + return x, weight, b + + def test_every_config_of_every_schedule(self): + for layout in ("nchw", "channels_last"): + x, w, b = self.tensors(2, 40, 13, 19, 3, 3) + if layout == "channels_last": + x = x.to(memory_format=torch.channels_last) + with self.subTest(layout=layout): + ran = self.check(x, w, b, (1, 1), (1, 1), (1, 1), every_config=True) + self.assertEqual("dw_row8" in ran, layout == "channels_last") + + def test_geometry_dtypes_and_layouts(self): + geometries = ( + # (n, c, h, w), (kh, kw), stride, before, after + ((2, 48, 17, 23), (3, 3), (1, 1), (1, 1), (1, 1)), + ((1, 7, 9, 31), (5, 3), (2, 1), (2, 1), (2, 1)), + ((3, 33, 11, 10), (7, 7), (1, 2), (3, 3), (3, 3)), + ((1, 16, 8, 9), (9, 9), (1, 1), (4, 4), (4, 4)), + ((2, 5, 6, 7), (1, 9), (3, 1), (0, 4), (0, 4)), + ((1, 24, 10, 12), (4, 4), (1, 1), (1, 1), (2, 2)), + ((1, 12, 9, 9), (3, 3), (2, 2), (0, 0), (0, 0)), + ((1, 3, 1, 1), (1, 1), (1, 1), (0, 0), (0, 0)), + ((1, 64, 2, 2), (3, 3), (2, 2), (1, 0), (0, 1)), + ) + cases = [ + (geometry, dtype, layout, i % 2 == 0) + for i, geometry in enumerate(geometries) + for dtype in (torch.float32, torch.float16, torch.bfloat16) + for layout in ("nchw", "channels_last") + ] + cases += [ + (geometry, torch.float32, "weight_cl", True) for geometry in geometries + ] + for (shape, kernel, stride, before, after), dtype, layout, bias in cases: + x, w, b = self.tensors(*shape, *kernel, dtype=dtype, bias=bias) + if layout == "channels_last": + x = x.to(memory_format=torch.channels_last) + if layout == "weight_cl": + w = w.to(memory_format=torch.channels_last) + with self.subTest( + shape=shape, kernel=kernel, stride=stride, dtype=dtype, layout=layout + ): + self.check(x, w, b, stride, before, after) + + def test_grouped_geometry_dtypes_and_layouts(self): + geometries = ( + # (n, c, h, w), (kh, kw), stride, before, after, (group channels, + # multiplier) + ((2, 16, 13, 17), (3, 3), (1, 1), (1, 1), (1, 1), (1, 2)), + ((1, 12, 11, 9), (5, 5), (2, 2), (2, 2), (2, 2), (1, 4)), + ((2, 16, 10, 12), (3, 3), (1, 1), (1, 1), (1, 1), (2, 1)), + ((1, 24, 15, 13), (7, 7), (1, 1), (3, 3), (3, 3), (3, 2)), + ((1, 32, 9, 11), (3, 3), (2, 2), (1, 1), (1, 1), (4, 1)), + ((2, 32, 12, 10), (5, 5), (1, 2), (2, 2), (2, 2), (8, 1)), + ((1, 16, 14, 14), (7, 7), (2, 2), (3, 3), (3, 3), (8, 2)), + ((1, 16, 9, 8), (4, 4), (1, 1), (1, 1), (2, 2), (4, 4)), + ) + cases = [ + (geometry, dtype, layout, i % 2 == 0) + for i, geometry in enumerate(geometries) + for dtype in (torch.float32, torch.float16, torch.bfloat16) + for layout in ("nchw", "channels_last") + ] + cases += [ + (geometry, torch.float32, "weight_cl", True) for geometry in geometries + ] + for geometry, dtype, layout, bias in cases: + shape, kernel, stride, before, after, (group_channels, multiplier) = ( + geometry + ) + x, w, b = self.tensors( + *shape, + *kernel, + dtype=dtype, + bias=bias, + group_channels=group_channels, + multiplier=multiplier, + ) + if layout == "channels_last": + x = x.to(memory_format=torch.channels_last) + if layout == "weight_cl": + w = w.to(memory_format=torch.channels_last) + with self.subTest( + shape=shape, + kernel=kernel, + group_channels=group_channels, + multiplier=multiplier, + dtype=dtype, + layout=layout, + ): + self.check(x, w, b, stride, before, after) + + def test_every_config_of_every_schedule_grouped(self): + for (group_channels, multiplier), layout in ( + ((3, 2), "nchw"), + ((3, 2), "channels_last"), + ((1, 3), "nchw"), + ((1, 3), "channels_last"), + ): + x, w, b = self.tensors( + 2, + 24, + 11, + 13, + 3, + 3, + group_channels=group_channels, + multiplier=multiplier, + ) + if layout == "channels_last": + x = x.to(memory_format=torch.channels_last) + with self.subTest(group_channels=group_channels, layout=layout): + ran = self.check(x, w, b, (1, 1), (1, 1), (1, 1), every_config=True) + self.assertEqual("dw_row8" in ran, layout == "channels_last") + self.assertEqual("dw_mult" in ran, group_channels == 1) + self.assertIn("grp_plane", ran) + self.assertIn("grp_reduce", ran) + + def test_high_channel_multipliers(self): + from executorch.backends.cuda.autotune.launch_params import ( + autotune_launch_params, + clear_launch_param_cache, + ) + from executorch.backends.cuda.triton.kernels.conv.core import direct + + x, w, b = self.tensors(1, 3, 9, 11, 3, 3, multiplier=direct.MAX_MULT_MULTIPLIER) + ran = self.check(x, w, b, (1, 1), (1, 1), (1, 1), every_config=True) + self.assertIn("dw_mult", ran) + # 4096 output channels from two input channels: dw_mult declines the + # multiplier and the timed choice runs the other schedules. + x, w, b = self.tensors(1, 2, 12, 12, 3, 3, multiplier=2048) + self.assertNotIn("dw_mult", self.check(x, w, b, (1, 1), (1, 1), (1, 1))) + clear_launch_param_cache() + with autotune_launch_params() as stats: + out = self.convolution(x, w, b, (1, 1), (1, 1), (1, 1)) + self.assertEqual(stats.measured, 1) + self.assertEqual( + set(stats.picks[0].times_us), {"dw_flat", "grp_plane", "grp_reduce"} + ) + torch.testing.assert_close( + out.double(), + _reference(x, w, b, (1, 1), (1, 1), (1, 1)), + atol=1e-4, + rtol=1e-4, + ) + + def test_unsupported_groupings(self): + from executorch.backends.cuda.triton.kernels.conv.core import direct + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + ) + + x = torch.randn(1, 18, 8, 8, device="cuda") + for groups, out_channels, reason in ( + (1, 4, "not a depthwise convolution"), + (2, 4, "not a depthwise convolution"), + ): + w = torch.randn(out_channels, 18 // groups, 3, 3, device="cuda") + call = ConvCall( + x, + w, + None, + ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), groups), + ) + with self.subTest(groups=groups): + self.assertIn(reason, direct.supports(call)) + + 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, 3, 3) + for x in ( + base[:, 2:18, 1:-1, 2:], + base.to(memory_format=torch.channels_last)[:, 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, 8, 6, 6, 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_every_gpu_runs_the_base_schedule(self): + from executorch.backends.cuda.triton.kernels.conv.core import direct + from executorch.backends.cuda.triton.kernels.conv.core.common import ( + ConvCall, + ConvParams, + ) + + self.assertEqual(direct.SCHEDULES[0], direct.Schedule("dw_flat", 0)) + x, w, b = self.tensors(1, 8, 6, 6, 3, 3) + call = ConvCall( + x, w, b, ConvParams((1, 1), (1, 1), (1, 1), (1, 1), False, (0, 0), 8) + ) + for capability in ((6, 0), (7, 5), (8, 0), (9, 0)): + with mock.patch.object( + torch.cuda, "get_device_capability", return_value=capability + ): + self.assertIsNone(direct.schedule_supports("dw_flat", call)) + + def test_timed_schedule_choice(self): + from executorch.backends.cuda.autotune.launch_params import ( + autotune_launch_params, + clear_launch_param_cache, + ) + + clear_launch_param_cache() + x, w, b = self.tensors(2, 64, 28, 28, 3, 3) + x = x.to(memory_format=torch.channels_last) + 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), {"dw_flat", "dw_row8"}) + torch.testing.assert_close( + out.double(), + _reference(x, w, b, (1, 1), (1, 1), (1, 1)), + atol=1e-4, + rtol=1e-4, + ) + + def test_64_bit_indexing(self): + from executorch.backends.cuda.triton.kernels.conv.conv import ( + force_conv_candidate, + ) + + # 2^31 + 2^17 elements: offsets past the int32 range. + shape = (1, 2, 32769, 32768) + if torch.cuda.mem_get_info()[0] < 12 * 2**30: + self.skipTest("needs 12 GiB of free device memory") + w = torch.randn(2, 1, 3, 3, device="cuda", dtype=torch.float16) + for schedule, memory_format in ( + ("dw_flat", torch.contiguous_format), + ("dw_row8", torch.channels_last), + ): + x = torch.randn(shape, device="cuda", dtype=torch.float16) + x = x.contiguous(memory_format=memory_format) + with self.subTest(schedule=schedule), force_conv_candidate(schedule): + 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=2e-3, rtol=2e-3 + ) + del x, out + + +class _DepthwiseBlock(nn.Module): + def __init__(self): + super().__init__() + self.dw3 = nn.Conv2d(32, 32, 3, padding=1, groups=32) + self.dw7 = nn.Conv2d(32, 32, 7, padding="same", groups=32, bias=False) + self.pw = nn.Conv2d(32, 16, 1) + + def forward(self, x): + return self.pw(F.gelu(self.dw7(self.dw3(x)))) + + +class _GroupedBlock(nn.Module): + """Grouped convolutions with 2 and 4 input channels per group and a + depthwise convolution with channel multiplier 2.""" + + def __init__(self): + super().__init__() + self.g2 = nn.Conv2d(32, 32, 3, padding=1, groups=16) + self.dw_m2 = nn.Conv2d(32, 64, 5, padding=2, groups=32, bias=False) + self.g4 = nn.Conv2d(64, 16, 7, stride=2, padding=3, groups=16) + + def forward(self, x): + return self.g4(F.gelu(self.dw_m2(self.g2(x)))) + + +class DepthwiseLoweringTest(unittest.TestCase): + def setUp(self): + if not torch.cuda.is_available() or torch.version.hip is not None: + self.skipTest("CUDA required") + + def test_depthwise_block_lowers_and_runs(self): + 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 + + torch.manual_seed(0) + model = _DepthwiseBlock().cuda().eval() + inputs = (torch.randn(2, 32, 20, 24, device="cuda"),) + 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, direct template)"), 2) + self.assertIn("Skipping conv replacement", output) + self.assertIn("Replaced 2 nodes with Triton kernels", output) + + from executorch.export.utils import is_fbcode + + if is_fbcode(): + # Running the program is skipped in fbcode. + return + 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) + (output,) = method.load_method("forward").execute([inputs[0].cpu()]) + with torch.no_grad(): + expected = model(*inputs) + torch.testing.assert_close(output, expected.cpu(), atol=1e-4, rtol=1e-4) + + def test_grouped_block_lowers_and_runs_with_both_apis(self): + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + from executorch.exir import ( + EdgeCompileConfig, + to_edge, + to_edge_transform_and_lower, + ) + from executorch.export.utils import is_fbcode + + torch.manual_seed(0) + model = _GroupedBlock().cuda().eval() + inputs = (torch.randn(2, 32, 20, 24, device="cuda"),) + with torch.no_grad(): + program = torch.export.export(model, inputs, strict=True) + expected = model(*inputs) + config = EdgeCompileConfig(_check_ir_validity=False) + for api in ("to_edge_transform_and_lower", "to_edge_to_backend"): + with self.subTest(api=api): + partitioner = CudaPartitioner( + [CudaBackend.generate_method_name_compile_spec("forward")] + ) + with self.assertLogs(_PASS_LOGGER, "INFO") as logs: + if api == "to_edge_transform_and_lower": + edge = to_edge_transform_and_lower( + program, partitioner=[partitioner], compile_config=config + ) + else: + edge = to_edge(program, compile_config=config).to_backend( + partitioner + ) + output = "\n".join(logs.output) + self.assertEqual(output.count("(optional, direct template)"), 3) + self.assertIn("Replaced 3 nodes with Triton kernels", output) + if is_fbcode(): + # Running the program is skipped in fbcode. + continue + 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) + (out,) = method.load_method("forward").execute([inputs[0].cpu()]) + torch.testing.assert_close(out, expected.cpu(), atol=1e-4, rtol=1e-4) diff --git a/backends/cuda/tests/test_conv_routing.py b/backends/cuda/tests/test_conv_routing.py index b594c77e979..9541f34e644 100644 --- a/backends/cuda/tests/test_conv_routing.py +++ b/backends/cuda/tests/test_conv_routing.py @@ -95,7 +95,7 @@ def __init__(self): kwargs.pop("out_channels", channels), kernel, padding=padding, - groups=channels, + groups=kwargs.pop("groups", channels), dtype=dtype, **kwargs, ) @@ -405,3 +405,190 @@ def fail(*args, **kwargs): raise RuntimeError("injected fault") return ConvTemplate("failing", fail, ("any",), fail, fail, fail, fail) + + +class DirectTemplateRoutingTest(_PassTest): + """Which convolutions the direct-stencil template takes.""" + + def test_depthwise_is_routed(self): + for dtype in (torch.float32, torch.float16, torch.bfloat16): + with self.subTest(dtype=dtype): + self.assertRouted( + _depthwise(dtype=dtype), + lambda dtype=dtype: torch.randn(2, 8, 12, 10, dtype=dtype), + template="direct", + ) + self.assertRouted( + _depthwise(kernel=4, padding="same"), + lambda: torch.randn(1, 8, 9, 9), + template="direct", + canonical_args=([1, 1], [1, 1], [2, 2], [1, 1], False, [0, 0], 8), + ) + self.assertRouted( + _depthwise(kernel=7, padding=3, stride=2), + lambda: torch.randn(1, 8, 15, 9).to(memory_format=torch.channels_last), + template="direct", + device="cpu", + ) + + def test_grouped_is_routed(self): + # (channels, out channels, groups, kernel, stride) + for channels, out_channels, groups, kernel, stride in ( + (8, 16, 8, 3, 1), + (8, 8, 4, 5, 2), + (16, 4, 2, 7, 1), + # 4 input and output channels per group with 1x1 and 3x3 kernels. + (16, 16, 4, 1, 1), + (16, 16, 4, 3, 1), + ): + with self.subTest(channels=channels, groups=groups): + self.assertRouted( + _depthwise( + channels=channels, + out_channels=out_channels, + groups=groups, + kernel=kernel, + padding=kernel // 2, + stride=stride, + ), + lambda channels=channels: torch.randn(2, channels, 12, 10), + template="direct", + ) + + def test_other_convolutions_are_kept(self): + x = lambda: torch.randn(1, 8, 9, 9) # noqa: E731 + self.assertKept( + _depthwise(channels=18, groups=2), + lambda: torch.randn(1, 18, 9, 9), + reason="at most 8 input channels per group", + ) + for channels, out_channels, groups, kernel, stride, reason in ( + (16, 16, 4, 5, 1, "where Inductor's template is as fast"), + (8, 16, 4, 9, 1, "grouped kernel [9, 9] is not benchmarked"), + (8, 16, 4, 3, 3, "grouped stride [3, 3] is not benchmarked"), + ): + with self.subTest(groups=groups, out_channels=out_channels, kernel=kernel): + self.assertKept( + _depthwise( + channels=channels, + out_channels=out_channels, + groups=groups, + kernel=kernel, + padding=kernel // 2, + stride=stride, + ), + lambda channels=channels: torch.randn(1, channels, 12, 12), + reason=reason, + ) + for dtype, reason in ( + (torch.bfloat16, "reduction of 196 per group"), + (torch.float16, "grouped float16 is not benchmarked"), + ): + with self.subTest(dtype=dtype): + self.assertKept( + _depthwise( + channels=16, + out_channels=8, + groups=4, + kernel=7, + padding=3, + dtype=dtype, + ), + lambda dtype=dtype: torch.randn(1, 16, 12, 12, dtype=dtype), + reason=reason, + ) + self.assertKept(_depthwise(dilation=2, padding=2), x, reason="dilation [2, 2]") + self.assertKept(_depthwise(kernel=11, padding=5), x, reason="larger than 9") + self.assertKept( + _depthwise(dtype=torch.float64), + lambda: torch.randn(1, 8, 9, 9, dtype=torch.float64), + reason="dtype torch.float64", + ) + + class Dense(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(8, 4, 3) + + def forward(self, x): + return self.conv(x) + + self.assertKept(Dense, x, reason="not a depthwise convolution") + + class Depthwise3d(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv3d(8, 8, 3, groups=8) + + def forward(self, x): + return self.conv(x) + + self.assertKept( + Depthwise3d, lambda: torch.randn(1, 8, 5, 5, 5), reason="3-D convolution" + ) + + class Transposed(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.ConvTranspose2d(8, 8, 3, groups=8) + + def forward(self, x): + return self.conv(x) + + self.assertKept(Transposed, x, reason="transposed convolution") + + def test_benchmarked_range(self): + from executorch.backends.cuda.triton.kernels.conv.core import direct + + # (channels, input size, kernel, stride, padding): the edge of the + # benchmarked range is routed, one step past it is not. + edges = ( + ( + (8, (1, 8, 20, 20), 3, 3, 1), + (8, (1, 8, 20, 20), 3, 4, 1), + "stride [4, 4] is not benchmarked", + ), + ( + (8, (1, 8, 20, 20), 5, 1, 4), + (8, (1, 8, 20, 20), 5, 1, 5), + "beyond kernel - 1 is not benchmarked", + ), + ( + ( + direct.MAX_ROUTED_CHANNELS, + (1, direct.MAX_ROUTED_CHANNELS, 4, 4), + 3, + 1, + 1, + ), + (4100, (1, 4100, 4, 4), 3, 1, 1), + "4100 channels are not benchmarked", + ), + ( + (2, (1, 2, 8, direct.MAX_ROUTED_SPATIAL_SIZE), 3, 1, 1), + (2, (1, 2, 8, direct.MAX_ROUTED_SPATIAL_SIZE + 1), 3, 1, 1), + "spatial size [8, 2049] is not benchmarked", + ), + ( + (128, (32, 128, 112, 112), 3, 1, 1), + (128, (33, 128, 112, 112), 3, 1, 1), + "element input is not benchmarked", + ), + ) + for inside, outside, reason in edges: + for case, routed in ((inside, True), (outside, False)): + channels, shape, kernel, stride, padding = case + model = _depthwise( + channels=channels, kernel=kernel, padding=padding, stride=stride + ) + with self.subTest(case=case): + if routed: + self.assertRouted( + model, + lambda shape=shape: torch.randn(shape), + template="direct", + ) + else: + self.assertKept( + model, lambda shape=shape: torch.randn(shape), reason=reason + ) diff --git a/backends/cuda/triton/kernels/conv/conv.py b/backends/cuda/triton/kernels/conv/conv.py index 7031e79735b..aabbc138679 100644 --- a/backends/cuda/triton/kernels/conv/conv.py +++ b/backends/cuda/triton/kernels/conv/conv.py @@ -20,6 +20,7 @@ import torch import triton +from executorch.backends.cuda.triton.kernels.conv.core import direct from executorch.backends.cuda.triton.kernels.conv.core.common import ( check_params, ConvCall, @@ -42,7 +43,17 @@ class ConvTemplate: worth_routing: Callable[[ConvCall], Optional[str]] -TEMPLATES: tuple[ConvTemplate, ...] = () +TEMPLATES = ( + ConvTemplate( + "direct", + direct.supports, + tuple(s.name for s in direct.SCHEDULES), + direct.schedule_supports, + direct.launch, + direct.run_schedule, + direct.worth_routing, + ), +) _FORCED: contextvars.ContextVar[Optional[tuple[str, Optional[triton.Config]]]] = ( contextvars.ContextVar("forced_conv_schedule", default=None) diff --git a/backends/cuda/triton/kernels/conv/core/direct.py b/backends/cuda/triton/kernels/conv/core/direct.py new file mode 100644 index 00000000000..69397373a56 --- /dev/null +++ b/backends/cuda/triton/kernels/conv/core/direct.py @@ -0,0 +1,1109 @@ +# 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. + +"""Direct-stencil template: 2-D depthwise convolutions and grouped convolutions +with at most ``MAX_GROUP_CHANNELS`` input channels per group (any channel +multiplier). + +Each output reads its ``KH x KW`` window of its group's input channels straight +from the input; the stencil is unrolled over constexpr taps (one kernel row at a +time for more than one channel per group, which keeps compile time flat in the +kernel height) and accumulated in FP32, and the bias is added in FP32 before the +single store. Shapes, strides and geometry are constexpr (static shapes), tiles +are autotuned. Five schedules write the same output layout: + +- ``dw_flat``: one flat tile of outputs in the output's memory order. Legal for + every supported call and correct on every NVIDIA GPU; it is the base schedule. +- ``dw_row8``: eight adjacent outputs of one row times a block of channels, + sliding the input window through registers so each input vector is loaded + once per kernel row (with unit W stride), and each weight vector once for + the eight. Needs channel-contiguous input. +- ``dw_mult``: depthwise convolutions with a channel multiplier of at most + ``MAX_MULT_MULTIPLIER``; a block of positions of one input channel computes + all of its output channels, so each input value is loaded once for them. +- ``grp_plane``: grouped convolutions and channel multipliers; a block of + positions of one output channel, so the weights are scalars shared by the + block, loaded once per warp. +- ``grp_reduce``: grouped convolutions and channel multipliers with few + outputs; a tile of outputs times a block of their taps summed across the + tile, so the reduction runs in parallel. + +The router sends a call here when ``supports()`` accepts it and it lies in the +range ``benchmarks/conv/benchmark_depthwise.py`` measures (``worth_routing``), +where the measured calls are faster than Inductor's Triton conv template; other +supported calls keep the regular lowering. +""" + +from __future__ import annotations + +from typing import NamedTuple, 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 ( + ConvCall, + ConvParams, + max_offset, + SUPPORTED_DTYPES, +) +from executorch.backends.cuda.triton.routing import device_supports +from torch.library import wrap_triton + +MAX_KERNEL_SIZE = 9 +MAX_GROUP_CHANNELS = 8 +# dw_mult holds a (multiplier, BLOCK) FP32 accumulator per program: up to 16 +# output channels it stays within 128 registers per thread in every config. +# Larger multipliers take the other schedules. +MAX_MULT_MULTIPLIER = 16 +_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), + Schedule("dw_mult", 0), + Schedule("grp_plane", 0), + Schedule("grp_reduce", 0), +) + + +def _strides_ordered(t: torch.Tensor, order: tuple[int, ...]) -> bool: + strides = [t.stride(d) for d in order if t.shape[d] > 1] + return all(s > 0 for s in strides) and all( + a > b for a, b in zip(strides, strides[1:]) + ) + + +def _geometry_reason(call: ConvCall) -> Optional[str]: + x, w, p = call.input, call.weight, call.params + if p.rank != 2: + return f"direct stencil: {p.rank}-D convolution, needs 2-D" + if p.transposed: + return "direct stencil: transposed convolution" + if any(d != 1 for d in p.dilation): + return f"direct stencil: dilation {list(p.dilation)}" + if any(o != 0 for o in p.output_padding): + return f"direct stencil: output_padding {list(p.output_padding)}" + channels = x.shape[1] + depthwise = p.groups == channels == w.shape[0] + if not (depthwise or (p.groups > 1 and w.shape[1] <= MAX_GROUP_CHANNELS)): + return ( + "direct stencil: not a depthwise convolution or a grouped convolution " + f"with at most {MAX_GROUP_CHANNELS} input channels per group " + f"(groups {p.groups}, input {list(x.shape)}, weight {list(w.shape)})" + ) + if max(w.shape[2:]) > MAX_KERNEL_SIZE: + return ( + f"direct stencil: kernel {list(w.shape[2:])} larger than {MAX_KERNEL_SIZE}" + ) + out_hw = [ + (i + pb + pa - k) // s + 1 + for i, k, s, pb, pa in zip( + x.shape[2:], w.shape[2:], p.stride, p.padding_before, p.padding_after + ) + ] + if min(*x.shape, *out_hw) <= 0: + return f"direct stencil: 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``: positive strides, no negative padding).""" + 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"direct stencil: dtype {x.dtype}" + if w.dtype != x.dtype or (b is not None and b.dtype != x.dtype): + return "direct stencil: input, weight and bias dtypes differ" + if not (_strides_ordered(x, (0, 1, 2, 3)) or _strides_ordered(x, (0, 2, 3, 1))): + return f"direct stencil: input strides {list(x.stride())} are neither NCHW nor channels-last" + if b is not None and (tuple(b.shape) != (w.shape[0],) or not b.is_contiguous()): + return "direct stencil: bias is not a contiguous [C_out] tensor" + return None + + +# The benchmarked range (the benchmark's envelope and extent suites cover every +# dtype and layout supports() accepts within it). +MAX_ROUTED_STRIDE = 3 +MAX_ROUTED_CHANNELS = 4096 +MAX_ROUTED_SPATIAL_SIZE = 2048 +MAX_ROUTED_INPUT_ELEMENTS = 32 * 128 * 112 * 112 +# Grouped convolutions (the grouped suite): kernels up to 7x7 and strides up to +# 2. Inductor's template runs each group as a matrix product (reduction over the +# group's input channels and taps, one column per output channel); it is as +# fast as the stencil with 4 or more input and output channels per group and a +# kernel larger than 3x3, and in fp16/bf16, where it uses tensor cores, with two +# or more output channels per group and a reduction of 4 x 7 x 7 or more. The +# suite runs fp32 and bf16. +MAX_ROUTED_GROUPED_KERNEL = 7 +MAX_ROUTED_GROUPED_STRIDE = 2 +MIN_TEMPLATE_GROUP_CHANNELS = 4 +MIN_TEMPLATE_HALF_REDUCTION = 4 * 7 * 7 + + +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 max(p.stride) > MAX_ROUTED_STRIDE: + return f"direct stencil: stride {list(p.stride)} is not benchmarked" + if any( + max(b, a) > k - 1 + for b, a, k in zip(p.padding_before, p.padding_after, w.shape[2:]) + ): + return ( + f"direct stencil: padding {list(p.padding_before)}/" + f"{list(p.padding_after)} beyond kernel - 1 is not benchmarked" + ) + channels = max(x.shape[1], w.shape[0]) + if channels > MAX_ROUTED_CHANNELS: + return f"direct stencil: {channels} channels are not benchmarked" + if max(x.shape[2:]) > MAX_ROUTED_SPATIAL_SIZE: + return f"direct stencil: spatial size {list(x.shape[2:])} is not benchmarked" + if x.numel() > MAX_ROUTED_INPUT_ELEMENTS: + return f"direct stencil: a {x.numel()}-element input is not benchmarked" + if w.shape[1] == 1 and w.shape[0] == p.groups: + return None + group_channels, multiplier = w.shape[1], w.shape[0] // p.groups + if ( + min(group_channels, multiplier) >= MIN_TEMPLATE_GROUP_CHANNELS + and max(w.shape[2:]) > 3 + ): + return ( + f"direct stencil: {group_channels} input and {multiplier} output " + "channels per group, where Inductor's template is as fast" + ) + if x.dtype == torch.float16: + return "direct stencil: grouped float16 is not benchmarked" + reduction = group_channels * w.shape[2] * w.shape[3] + if ( + x.dtype != torch.float32 + and multiplier >= 2 + and reduction >= MIN_TEMPLATE_HALF_REDUCTION + ): + return ( + f"direct stencil: {x.dtype} with {multiplier} output channels and a " + f"reduction of {reduction} per group, where Inductor's template is as fast" + ) + if max(w.shape[2:]) > MAX_ROUTED_GROUPED_KERNEL: + return f"direct stencil: grouped kernel {list(w.shape[2:])} is not benchmarked" + if max(p.stride) > MAX_ROUTED_GROUPED_STRIDE: + return f"direct stencil: grouped stride {list(p.stride)} is not benchmarked" + return None + + +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 == "dw_mult": + w = call.weight + if w.shape[1] != 1 or w.shape[0] == call.params.groups: + return "dw_mult needs a depthwise convolution with a channel multiplier" + multiplier = w.shape[0] // call.params.groups + if multiplier > MAX_MULT_MULTIPLIER: + return ( + f"dw_mult: channel multiplier {multiplier} above {MAX_MULT_MULTIPLIER}" + ) + if name in ("grp_plane", "grp_reduce"): + w = call.weight + if w.shape[1] == 1 and w.shape[0] == call.params.groups: + return f"{name} needs a grouped convolution or a channel multiplier" + if name == "dw_row8": + if call.input.stride(1) != 1: + return "dw_row8 needs channel-contiguous input" + return None + + +@triton.jit +def _group(c, M: tl.constexpr): + if M == 1: + return c + return c // M + + +@triton.jit +def _flat_row( + acc, + x_base, + w_base, + mask, + oh, + ow, + kh, + CG: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, +): + """Adds kernel row ``kh`` of every input channel of the group to ``acc``.""" + ih = oh * STRIDE_H + kh - PAD_H + row_mask = mask & (ih >= 0) & (ih < H) + for kw in tl.static_range(KW): + iw = ow * STRIDE_W + kw - PAD_W + tap_mask = row_mask & (iw >= 0) & (iw < W) + for ci in tl.static_range(CG): + xv = tl.load( + x_base + ci * SX_C + ih * SX_H + iw * SX_W, mask=tap_mask, other=0.0 + ) + wv = tl.load( + w_base + ci * SW_CI + kh * SW_H + kw * SW_W, mask=mask, other=0.0 + ) + acc += xv.to(tl.float32) * wv.to(tl.float32) + return acc + + +@triton.jit +def _flat_row_same_size( + acc, + x_pos, + w_base, + mask, + oh, + ow, + kh, + CG: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + KW: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, +): + """``_flat_row`` for unit strides, an output the size of the input and + contiguous input rows: a tap reads the output's own flat spatial offset + plus a constant, so adjacent outputs read adjacent inputs and their loads + can be vectorized.""" + # The bounds compare the output index itself with constants, which lets + # Triton prove a mask uniform over adjacent outputs and vectorize the load. + row_mask = mask & (oh >= PAD_H - kh) & (oh < H + PAD_H - kh) + row = x_pos + (kh - PAD_H) * W + for kw in tl.static_range(KW): + tap_mask = row_mask & (ow >= PAD_W - kw) & (ow < W + PAD_W - kw) + for ci in tl.static_range(CG): + xv = tl.load(row + ci * SX_C + (kw - PAD_W), mask=tap_mask, other=0.0) + wv = tl.load( + w_base + ci * SW_CI + kh * SW_H + kw * SW_W, mask=mask, other=0.0 + ) + acc += xv.to(tl.float32) * wv.to(tl.float32) + return acc + + +@triton.jit +def _flat_kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + CG: tl.constexpr, + M: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNELS_LAST: tl.constexpr, + OUT_DENSE: tl.constexpr, + INDEX_64: tl.constexpr, + BLOCK: tl.constexpr, +): + TOTAL: tl.constexpr = N * C * OH * OW + pid = tl.program_id(0) + if INDEX_64: + pid = pid.to(tl.int64) + idx = pid * BLOCK + tl.arange(0, BLOCK) + if TOTAL % BLOCK == 0: + mask = tl.full((BLOCK,), True, tl.int1) + else: + mask = idx < TOTAL + if CHANNELS_LAST: + c = idx % C + rest = idx // C + ow = rest % OW + rest = rest // OW + oh = rest % OH + n = rest // OH + else: + ow = idx % OW + rest = idx // OW + oh = rest % OH + rest = rest // OH + c = rest % C + n = rest // C + x_base = x_ptr + n * SX_N + _group(c, M) * (CG * SX_C) + w_base = w_ptr + c * SW_C + acc = tl.zeros((BLOCK,), dtype=tl.float32) + # Unit strides, an output the size of the input and contiguous input rows: + # a tap reads the output's own flat offset plus a constant. + SAME_SIZE: tl.constexpr = ( + not CHANNELS_LAST + and STRIDE_H == 1 + and STRIDE_W == 1 + and OH == H + and OW == W + and SX_W == 1 + and SX_H == W + ) + # Depthwise calls unroll every tap; grouped calls unroll one kernel row, + # which keeps their compile time flat in the kernel height. Depthwise + # kernels narrower than 5 have no tap the flat offset lets Triton vectorize + # beyond the center one it vectorizes anyway. + if SAME_SIZE and (CG > 1 or KW >= 5): + x_pos = x_base + idx % (OH * OW) + if CG == 1: + for kh in tl.static_range(KH): + acc = _flat_row_same_size( + acc, + x_pos, + w_base, + mask, + oh, + ow, + kh, + CG, + H, + W, + KW, + PAD_H, + PAD_W, + SX_C, + SW_CI, + SW_H, + SW_W, + ) + else: + for kh in range(KH): + acc = _flat_row_same_size( + acc, + x_pos, + w_base, + mask, + oh, + ow, + kh, + CG, + H, + W, + KW, + PAD_H, + PAD_W, + SX_C, + SW_CI, + SW_H, + SW_W, + ) + elif CG == 1: + for kh in tl.static_range(KH): + acc = _flat_row( + acc, + x_base, + w_base, + mask, + oh, + ow, + kh, + CG, + H, + W, + KW, + STRIDE_H, + STRIDE_W, + PAD_H, + PAD_W, + SX_C, + SX_H, + SX_W, + SW_CI, + SW_H, + SW_W, + ) + else: + for kh in range(KH): + acc = _flat_row( + acc, + x_base, + w_base, + mask, + oh, + ow, + kh, + CG, + H, + W, + KW, + STRIDE_H, + STRIDE_W, + PAD_H, + PAD_W, + SX_C, + SX_H, + SX_W, + SW_CI, + SW_H, + SW_W, + ) + if HAS_BIAS: + acc += tl.load(b_ptr + c, mask=mask, other=0.0).to(tl.float32) + if OUT_DENSE: + out_ptrs = out_ptr + idx + else: + out_ptrs = out_ptr + n * SO_N + c * SO_C + oh * SO_H + ow * SO_W + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _load_row(ptrs, mask, iw, W: tl.constexpr, SX_W: tl.constexpr): + v = tl.load(ptrs + iw * SX_W, mask=mask & (iw >= 0) & (iw < W), other=0.0) + return v.to(tl.float32) + + +@triton.jit +def _store_row(out_ptrs, acc, mask, ow, OW: tl.constexpr, SO_W: tl.constexpr): + if OW % 8 != 0: + mask = mask & (ow < OW) + tl.store(out_ptrs + ow * SO_W, acc.to(out_ptrs.dtype.element_ty), mask=mask) + + +@triton.jit +def _row8_input_row( + acc0, + acc1, + acc2, + acc3, + acc4, + acc5, + acc6, + acc7, + ptrs, + row_mask, + w_row, + c_mask, + iw0, + W: tl.constexpr, + KW: tl.constexpr, + STRIDE_W: tl.constexpr, + SX_W: tl.constexpr, + SW_W: tl.constexpr, +): + """Adds one input row times one kernel row to the eight outputs, sliding + the input window through registers (with unit W stride; otherwise each + output loads its own window, and the weights still serve all eight).""" + if STRIDE_W != 1: + for kw in tl.static_range(KW): + wv = tl.load(w_row + kw * SW_W, mask=c_mask, other=0.0).to(tl.float32) + acc0 += _load_row(ptrs, row_mask, iw0 + kw, W, SX_W) * wv + acc1 += _load_row(ptrs, row_mask, iw0 + STRIDE_W + kw, W, SX_W) * wv + acc2 += _load_row(ptrs, row_mask, iw0 + 2 * STRIDE_W + kw, W, SX_W) * wv + acc3 += _load_row(ptrs, row_mask, iw0 + 3 * STRIDE_W + kw, W, SX_W) * wv + acc4 += _load_row(ptrs, row_mask, iw0 + 4 * STRIDE_W + kw, W, SX_W) * wv + acc5 += _load_row(ptrs, row_mask, iw0 + 5 * STRIDE_W + kw, W, SX_W) * wv + acc6 += _load_row(ptrs, row_mask, iw0 + 6 * STRIDE_W + kw, W, SX_W) * wv + acc7 += _load_row(ptrs, row_mask, iw0 + 7 * STRIDE_W + kw, W, SX_W) * wv + return acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7 + x0 = _load_row(ptrs, row_mask, iw0, W, SX_W) + x1 = _load_row(ptrs, row_mask, iw0 + 1, W, SX_W) + x2 = _load_row(ptrs, row_mask, iw0 + 2, W, SX_W) + x3 = _load_row(ptrs, row_mask, iw0 + 3, W, SX_W) + x4 = _load_row(ptrs, row_mask, iw0 + 4, W, SX_W) + x5 = _load_row(ptrs, row_mask, iw0 + 5, W, SX_W) + x6 = _load_row(ptrs, row_mask, iw0 + 6, W, SX_W) + x7 = _load_row(ptrs, row_mask, iw0 + 7, W, SX_W) + for kw in tl.static_range(KW): + wv = tl.load(w_row + kw * SW_W, mask=c_mask, other=0.0) + wv = wv.to(tl.float32) + acc0 += x0 * wv + acc1 += x1 * wv + acc2 += x2 * wv + acc3 += x3 * wv + acc4 += x4 * wv + acc5 += x5 * wv + acc6 += x6 * wv + acc7 += x7 * wv + if kw + 1 < KW: + x0, x1, x2, x3, x4, x5, x6 = x1, x2, x3, x4, x5, x6, x7 + x7 = _load_row(ptrs, row_mask, iw0 + kw + 8, W, SX_W) + return acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7 + + +@triton.jit +def _row8_kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + CG: tl.constexpr, + M: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNELS_LAST: tl.constexpr, + OUT_DENSE: tl.constexpr, + INDEX_64: tl.constexpr, + BLOCK_C: tl.constexpr, +): + row = tl.program_id(0) + ow0 = tl.program_id(1) * 8 + c = tl.program_id(2) * BLOCK_C + tl.arange(0, BLOCK_C) + if INDEX_64: + row = row.to(tl.int64) + ow0 = ow0.to(tl.int64) + c = c.to(tl.int64) + n = row // OH + oh = row % OH + if C % BLOCK_C == 0: + c_mask = tl.full((BLOCK_C,), True, tl.int1) + else: + c_mask = c < C + x_base = x_ptr + n * SX_N + _group(c, M) * (CG * SX_C) + w_base = w_ptr + c * SW_C + iw0 = ow0 * STRIDE_W - PAD_W + acc0 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc1 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc2 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc3 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc4 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc5 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc6 = tl.zeros((BLOCK_C,), dtype=tl.float32) + acc7 = tl.zeros((BLOCK_C,), dtype=tl.float32) + # Depthwise calls unroll every tap; grouped calls unroll one kernel row, + # which keeps their compile time flat in the kernel height. + if CG == 1: + for kh in tl.static_range(KH): + ih = oh * STRIDE_H + kh - PAD_H + row_mask = c_mask & (ih >= 0) & (ih < H) + acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7 = _row8_input_row( + acc0, + acc1, + acc2, + acc3, + acc4, + acc5, + acc6, + acc7, + x_base + ih * SX_H, + row_mask, + w_base + kh * SW_H, + c_mask, + iw0, + W, + KW, + STRIDE_W, + SX_W, + SW_W, + ) + else: + for kh in range(KH): + ih = oh * STRIDE_H + kh - PAD_H + row_mask = c_mask & (ih >= 0) & (ih < H) + for ci in tl.static_range(CG): + acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7 = _row8_input_row( + acc0, + acc1, + acc2, + acc3, + acc4, + acc5, + acc6, + acc7, + x_base + ci * SX_C + ih * SX_H, + row_mask, + w_base + ci * SW_CI + kh * SW_H, + c_mask, + iw0, + W, + KW, + STRIDE_W, + SX_W, + SW_W, + ) + if HAS_BIAS: + bias = tl.load(b_ptr + c, mask=c_mask, other=0.0).to(tl.float32) + acc0 += bias + acc1 += bias + acc2 += bias + acc3 += bias + acc4 += bias + acc5 += bias + acc6 += bias + acc7 += bias + out_ptrs = out_ptr + n * SO_N + c * SO_C + oh * SO_H + _store_row(out_ptrs, acc0, c_mask, ow0, OW, SO_W) + _store_row(out_ptrs, acc1, c_mask, ow0 + 1, OW, SO_W) + _store_row(out_ptrs, acc2, c_mask, ow0 + 2, OW, SO_W) + _store_row(out_ptrs, acc3, c_mask, ow0 + 3, OW, SO_W) + _store_row(out_ptrs, acc4, c_mask, ow0 + 4, OW, SO_W) + _store_row(out_ptrs, acc5, c_mask, ow0 + 5, OW, SO_W) + _store_row(out_ptrs, acc6, c_mask, ow0 + 6, OW, SO_W) + _store_row(out_ptrs, acc7, c_mask, ow0 + 7, OW, SO_W) + + +@triton.jit +def _multiplier_kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + CG: tl.constexpr, + M: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNELS_LAST: tl.constexpr, + OUT_DENSE: tl.constexpr, + INDEX_64: tl.constexpr, + MP: tl.constexpr, + BLOCK: tl.constexpr, +): + """A depthwise convolution with channel multiplier ``M``: a block of output + positions of one image and input channel, all ``M`` output channels of the + channel per position. Each input value is loaded once for the ``M`` + outputs, and the weights are the same across the block.""" + G: tl.constexpr = C // M + pid = tl.program_id(0) + ng = tl.program_id(1) + if INDEX_64: + pid = pid.to(tl.int64) + ng = ng.to(tl.int64) + n = ng // G + g = ng % G + pos = pid * BLOCK + tl.arange(0, BLOCK) + if (OH * OW) % BLOCK == 0: + mask = tl.full((BLOCK,), True, tl.int1) + else: + mask = pos < OH * OW + oh = pos // OW + ow = pos % OW + m = tl.arange(0, MP) + m_mask = m < M + c = g * M + m + # Output channels on the outer dimension: each thread holds every channel + # of its positions, so an input value feeds them without being reloaded. + acc = tl.zeros((MP, BLOCK), dtype=tl.float32) + x_base = x_ptr + n * SX_N + g * SX_C + w_base = w_ptr + c * SW_C + for kh in tl.static_range(KH): + ih = oh * STRIDE_H + kh - PAD_H + row_mask = mask & (ih >= 0) & (ih < H) + for kw in tl.static_range(KW): + iw = ow * STRIDE_W + kw - PAD_W + xv = tl.load( + x_base + ih * SX_H + iw * SX_W, + mask=row_mask & (iw >= 0) & (iw < W), + other=0.0, + ) + wv = tl.load(w_base + kh * SW_H + kw * SW_W, mask=m_mask, other=0.0) + acc += wv.to(tl.float32)[:, None] * xv.to(tl.float32)[None, :] + if HAS_BIAS: + acc += tl.load(b_ptr + c, mask=m_mask, other=0.0).to(tl.float32)[:, None] + out_ptrs = ( + out_ptr + n * SO_N + c[:, None] * SO_C + oh[None, :] * SO_H + ow[None, :] * SO_W + ) + tl.store( + out_ptrs, + acc.to(out_ptr.dtype.element_ty), + mask=m_mask[:, None] & mask[None, :], + ) + + +@triton.jit +def _reduce_kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + CG: tl.constexpr, + M: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNELS_LAST: tl.constexpr, + OUT_DENSE: tl.constexpr, + INDEX_64: tl.constexpr, + BLOCK: tl.constexpr, + BLOCK_R: tl.constexpr, +): + """A tile of outputs times a block of their taps (group channel, kernel + row, kernel column), summed across the tile: the reduction runs in + parallel, for calls with few outputs and many taps per output.""" + TOTAL: tl.constexpr = N * C * OH * OW + TAPS: tl.constexpr = CG * KH * KW + pid = tl.program_id(0) + if INDEX_64: + pid = pid.to(tl.int64) + idx = pid * BLOCK + tl.arange(0, BLOCK) + mask = idx < TOTAL + if CHANNELS_LAST: + c = idx % C + rest = idx // C + ow = rest % OW + rest = rest // OW + oh = rest % OH + n = rest // OH + else: + ow = idx % OW + rest = idx // OW + oh = rest % OH + rest = rest // OH + c = rest % C + n = rest // C + x_base = x_ptr + n * SX_N + _group(c, M) * (CG * SX_C) + w_base = w_ptr + c * SW_C + acc = tl.zeros((BLOCK,), dtype=tl.float32) + for r0 in range(0, TAPS, BLOCK_R): + r = r0 + tl.arange(0, BLOCK_R) + r_mask = r < TAPS + ci = r // (KH * KW) + kh = (r // KW) % KH + kw = r % KW + ih = oh[:, None] * STRIDE_H + kh[None, :] - PAD_H + iw = ow[:, None] * STRIDE_W + kw[None, :] - PAD_W + tap_mask = ( + mask[:, None] + & r_mask[None, :] + & (ih >= 0) + & (ih < H) + & (iw >= 0) + & (iw < W) + ) + xv = tl.load( + x_base[:, None] + ci[None, :] * SX_C + ih * SX_H + iw * SX_W, + mask=tap_mask, + other=0.0, + ) + wv = tl.load( + w_base[:, None] + (ci * SW_CI + kh * SW_H + kw * SW_W)[None, :], + mask=mask[:, None] & r_mask[None, :], + other=0.0, + ) + acc += tl.sum(xv.to(tl.float32) * wv.to(tl.float32), axis=1) + if HAS_BIAS: + acc += tl.load(b_ptr + c, mask=mask, other=0.0).to(tl.float32) + if OUT_DENSE: + out_ptrs = out_ptr + idx + else: + out_ptrs = out_ptr + n * SO_N + c * SO_C + oh * SO_H + ow * SO_W + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _plane_kernel( + x_ptr, + w_ptr, + b_ptr, + out_ptr, + N: tl.constexpr, + C: tl.constexpr, + CG: tl.constexpr, + M: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + OH: tl.constexpr, + OW: tl.constexpr, + KH: tl.constexpr, + KW: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_W: tl.constexpr, + PAD_H: tl.constexpr, + PAD_W: tl.constexpr, + SX_N: tl.constexpr, + SX_C: tl.constexpr, + SX_H: tl.constexpr, + SX_W: tl.constexpr, + SW_C: tl.constexpr, + SW_CI: tl.constexpr, + SW_H: tl.constexpr, + SW_W: tl.constexpr, + SO_N: tl.constexpr, + SO_C: tl.constexpr, + SO_H: tl.constexpr, + SO_W: tl.constexpr, + HAS_BIAS: tl.constexpr, + CHANNELS_LAST: tl.constexpr, + OUT_DENSE: tl.constexpr, + INDEX_64: tl.constexpr, + BLOCK: tl.constexpr, +): + """A block of positions of one output channel of one image: the weights + are scalars shared by the block, loaded once per warp.""" + pid = tl.program_id(0) + nc = tl.program_id(1) + if INDEX_64: + pid = pid.to(tl.int64) + nc = nc.to(tl.int64) + n = nc // C + c = nc % C + pos = pid * BLOCK + tl.arange(0, BLOCK) + mask = pos < OH * OW + oh = pos // OW + ow = pos % OW + x_base = x_ptr + n * SX_N + (c // M) * (CG * SX_C) + w_base = w_ptr + c * SW_C + acc = tl.zeros((BLOCK,), dtype=tl.float32) + for kh in range(KH): + ih = oh * STRIDE_H + kh - PAD_H + row_mask = mask & (ih >= 0) & (ih < H) + for kw in tl.static_range(KW): + iw = ow * STRIDE_W + kw - PAD_W + tap_mask = row_mask & (iw >= 0) & (iw < W) + for ci in tl.static_range(CG): + xv = tl.load( + x_base + ci * SX_C + ih * SX_H + iw * SX_W, mask=tap_mask, other=0.0 + ) + wv = tl.load(w_base + ci * SW_CI + kh * SW_H + kw * SW_W) + acc += xv.to(tl.float32) * wv.to(tl.float32) + if HAS_BIAS: + acc += tl.load(b_ptr + c).to(tl.float32) + out_ptrs = out_ptr + n * SO_N + c * SO_C + oh * SO_H + ow * SO_W + tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=mask) + + +_TUNING_KEY = [ + "N", + "C", + "CG", + "M", + "H", + "W", + "KH", + "KW", + "STRIDE_H", + "STRIDE_W", + "CHANNELS_LAST", +] +FLAT_CONFIGS = [ + triton.Config({"BLOCK": block}, num_warps=warps) + for block, warps in ((256, 4), (512, 4), (512, 8), (1024, 8)) +] +ROW8_CONFIGS = [ + triton.Config({"BLOCK_C": block}, num_warps=warps) + for block, warps in ((16, 1), (16, 2), (32, 2), (32, 4), (64, 4)) +] +MULT_CONFIGS = [ + triton.Config({"BLOCK": block}, num_warps=warps) + for block, warps in ((256, 2), (512, 4), (1024, 4), (1024, 8), (2048, 8)) +] +PLANE_CONFIGS = [ + triton.Config({"BLOCK": block}, num_warps=warps) + for block, warps in ((64, 2), (128, 4), (256, 4), (512, 4), (1024, 8)) +] +REDUCE_CONFIGS = [ + triton.Config({"BLOCK": block, "BLOCK_R": r}, num_warps=warps) + for block, r, warps in ( + (16, 64, 4), + (32, 32, 4), + (8, 128, 4), + (64, 16, 4), + (32, 64, 8), + ) +] +_KERNELS = { + "dw_flat": triton.autotune(configs=FLAT_CONFIGS, key=_TUNING_KEY)(_flat_kernel), + "dw_row8": triton.autotune(configs=ROW8_CONFIGS, key=_TUNING_KEY)(_row8_kernel), + "dw_mult": triton.autotune(configs=MULT_CONFIGS, key=_TUNING_KEY)( + _multiplier_kernel + ), + "grp_plane": triton.autotune(configs=PLANE_CONFIGS, key=_TUNING_KEY)(_plane_kernel), + "grp_reduce": triton.autotune(configs=REDUCE_CONFIGS, key=_TUNING_KEY)( + _reduce_kernel + ), +} + + +def _grid(name: str, n: int, c: int, oh: int, ow: int, groups: int): + if name in ("dw_flat", "grp_reduce"): + return lambda meta: (triton.cdiv(n * c * oh * ow, meta["BLOCK"]),) + if name == "grp_plane": + return lambda meta: (triton.cdiv(oh * ow, meta["BLOCK"]), n * c) + if name == "dw_mult": + return lambda meta: (triton.cdiv(oh * ow, meta["BLOCK"]), n * groups) + return lambda meta: (n * oh, triton.cdiv(ow, 8), triton.cdiv(c, meta["BLOCK_C"])) + + +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.""" + n, _, h, w = x.shape + _, c, oh, ow = out.shape + channels_last = out.stride(1) == 1 and c > 1 + dense = ( + (c * ow * oh, 1, c * ow, c) if channels_last else (c * oh * ow, oh * ow, ow, 1) + ) + kernel = _KERNELS[name] + options = {} + if config is not None: + kernel = kernel.fn + options = {**config.kwargs, "num_warps": config.num_warps} + if name == "dw_mult": + # The group's output channels, padded to a power of two. + options["MP"] = triton.next_power_of_2(c // params.groups) + wrap_triton(kernel)[_grid(name, n, c, oh, ow, params.groups)]( + x, + weight, + x if bias is None else bias, + out, + n, + c, + weight.shape[1], + c // params.groups, + h, + w, + oh, + ow, + weight.shape[2], + weight.shape[3], + *params.stride, + *params.padding_before, + *x.stride(), + *weight.stride(), + *out.stride(), + HAS_BIAS=bias is not None, + CHANNELS_LAST=channels_last, + OUT_DENSE=all( + size == 1 or s == d for size, s, d in zip(out.shape, out.stride(), dense) + ), + INDEX_64=max(max_offset(x), max_offset(out), max_offset(weight)) > _INT32_MAX + or n * c * oh * ow > _INT32_MAX, + **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) diff --git a/examples/cuda/scripts/export.py b/examples/cuda/scripts/export.py index 8eb9c651b24..00a15eace79 100644 --- a/examples/cuda/scripts/export.py +++ b/examples/cuda/scripts/export.py @@ -62,6 +62,13 @@ def parse_args() -> argparse.Namespace: type=int, help="Seed model initialization and example input generation", ) + parser.add_argument( + "--save_example_io", + action="store_true", + help="Also write the example inputs and the eager outputs as raw files " + "(_input_.bin, _expected_.bin) for executor_runner " + "--inputs and --output_file", + ) args = parser.parse_args() return args @@ -75,6 +82,25 @@ def save_processed_bytes(processed_bytes, base_name: str): return +def save_example_io(model, example_args, base_name: str, output_dir: pathlib.Path): + with torch.no_grad(): + outputs = model(*example_args) + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + output_dir.mkdir(parents=True, exist_ok=True) + for kind, tensors in (("input", example_args), ("expected", outputs)): + if not all(isinstance(t, torch.Tensor) for t in tensors): + raise TypeError( + f"--save_example_io writes tensors only; {base_name} has {kind}s of " + f"types {[type(t).__name__ for t in tensors]}" + ) + for i, tensor in enumerate(tensors): + path = output_dir / f"{base_name}_{kind}_{i}.bin" + path.write_bytes( + tensor.contiguous().view(-1).view(torch.uint8).numpy().tobytes() + ) + + def main(): args = parse_args() @@ -94,6 +120,8 @@ def main(): dynamic_shapes, ) = EagerModelFactory.create_model(*MODEL_NAME_TO_MODEL[args.model_name]) model = model.eval() + if args.save_example_io: + save_example_io(model, example_args, args.model_name, args.output_dir) exported_programs = torch.export.export( model, args=example_args, diff --git a/examples/models/__init__.py b/examples/models/__init__.py index dbbe142337c..fb3a3dc69e2 100644 --- a/examples/models/__init__.py +++ b/examples/models/__init__.py @@ -46,6 +46,7 @@ class Model(str, Enum): ResNet8 = "resnet8" StreamingWakeWord = "streaming_wakeword" Sdpa = "sdpa" + DepthwiseConv = "depthwise_conv" Qwen3 = "qwen3" SmolVLM = "smolvlm" YOLO26 = "yolo26" @@ -114,6 +115,7 @@ def __str__(self) -> str: "StreamingWakeWordModel", ), str(Model.Sdpa): ("toy_model", "SdpaModule"), + str(Model.DepthwiseConv): ("toy_model", "DepthwiseConvModule"), str(Model.Qwen3): ("qwen3", "Qwen3Model"), str(Model.SmolVLM): ("smolvlm", "SmolVLMModel"), str(Model.YOLO26): ("yolo26", "YOLO26Model"), diff --git a/examples/models/toy_model/__init__.py b/examples/models/toy_model/__init__.py index 87456e3fd4c..09fd387348a 100644 --- a/examples/models/toy_model/__init__.py +++ b/examples/models/toy_model/__init__.py @@ -8,6 +8,7 @@ AddModule, AddMulModule, Conv1dModule, + DepthwiseConvModule, LinearModule, MulModule, SdpaModule, @@ -18,6 +19,7 @@ AddModule, AddMulModule, Conv1dModule, + DepthwiseConvModule, LinearModule, MulModule, SdpaModule, diff --git a/examples/models/toy_model/model.py b/examples/models/toy_model/model.py index a31149c29af..22c1797ff90 100644 --- a/examples/models/toy_model/model.py +++ b/examples/models/toy_model/model.py @@ -107,6 +107,24 @@ def get_example_inputs(self): return (torch.randn(1, 3, 10),) +class DepthwiseConvModule(torch.nn.Module, EagerModelBase): + """A depthwise 3x3 convolution with a residual, as in the convolutional + position encodings of hybrid vision transformers.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(64, 64, kernel_size=3, padding=1, groups=64) + + def forward(self, x): + return x + self.conv(x) + + def get_eager_model(self) -> torch.nn.Module: + return self + + def get_example_inputs(self): + return (torch.randn(3, 64, 28, 28),) + + class SdpaModule(torch.nn.Module, EagerModelBase): def __init__(self): super().__init__()