depthwise.py

"""LFM's depthwise causal convolution, the same numbers in every process: on ROCm as shifted sums, never MIOpen's
kernel, whose solver search runs again for every new length."""

from __future__ import annotations

import torch
import torch.nn.functional as F


def on_rocm(t: torch.Tensor) -> bool:
    return t.is_cuda and bool(torch.version.hip)


def depthwise(padded: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """Depthwise conv, `(B, C, T + k - 1)` to `(B, C, T)`. On ROCm, `shifted_sum`: MIOpen's depthwise kernel
    searches a solver for every new length, logs a warning each time, may settle on another one in another process,
    and takes 1.25 ms at 5k tokens against 0.07 ms fused (MI325X)."""
    if not on_rocm(padded):
        return F.conv1d(padded, weight, groups=padded.shape[1])
    return shifted_sum(padded, weight)


def shifted_sum(padded: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    """The depthwise conv as k shifted products summed in fp32: one fp32 copy of `padded`, the shifts views of it."""
    k, t = weight.shape[-1], padded.shape[-1] - weight.shape[-1] + 1
    p, w = padded.float(), weight[:, 0, :, None].float()
    return sum(p[..., i:i + t] * w[:, i] for i in range(k)).to(padded.dtype)


def pinned_conv(conv: torch.nn.Conv1d | torch.nn.Conv2d, x: torch.Tensor) -> torch.Tensor:
    """`conv(x)` without MIOpen on ROCm: a 1x1 conv as the matmul it is, a 2D conv with one input channel a group
    (`groups` = input channels) as `shifted_sum2d`. MIOpen's solver for the omni downsampler's 1x1 `lin_out` changed
    from one call to the next within a process (MI325X). Elsewhere, and for any other conv, `conv`'s own forward."""
    if on_rocm(x) and conv.padding_mode == "zeros" and set(conv.dilation) == {1}:
        if set(conv.kernel_size) == {1} and set(conv.stride) == {1} and conv.padding in ((0,), (0, 0)) and conv.groups == 1:
            return F.linear(x.movedim(1, -1), conv.weight.flatten(1), conv.bias).movedim(-1, 1)
        if x.ndim == 4 and conv.groups == conv.in_channels and not isinstance(conv.padding, str):
            return shifted_sum2d(x, conv.weight, conv.bias, conv.stride, conv.padding)
    return type(conv).forward(conv, x)


def shifted_sum2d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, stride: tuple[int, int],
                  padding: tuple[int, int]) -> torch.Tensor:
    """A 2D conv with one input channel a group, `weight` `(C * m, 1, kh, kw)`, as kh x kw strided products summed in
    fp32: each of the C input channels feeds its m output channels."""
    (kh, kw), (sh, sw), (ph, pw) = weight.shape[-2:], stride, padding
    c = x.shape[1]
    ho, wo = (x.shape[-2] + 2 * ph - kh) // sh + 1, (x.shape[-1] + 2 * pw - kw) // sw + 1
    p = F.pad(x.float(), (pw, pw, ph, ph))[:, :, None]
    w = weight.float().view(c, -1, kh, kw)[..., None, None]
    out = sum(p[..., i:i + sh * (ho - 1) + 1:sh, j:j + sw * (wo - 1) + 1:sw] * w[:, :, i, j]
              for i in range(kh) for j in range(kw)).flatten(1, 2)
    return (out if bias is None else out + bias.float()[:, None, None]).to(x.dtype)