Helion squeeze and excitation net Example

This example demonstrates a Helion kernel implementation of squeeze and excitation net as those used in https://arxiv.org/abs/1709.01507.

from __future__ import annotations

import torch
from torch import Tensor

import helion
from helion._testing import DEVICE
from helion._testing import HALF_DTYPE
from helion._testing import run_example
import helion.language as hl

@helion.kernel(
    # static_shapes=True gives a performance boost for matmuls
    static_shapes=True,
)
def squeeze_and_excitation_net_fwd(
    x: Tensor, a: Tensor, b: Tensor
) -> tuple[Tensor, Tensor, Tensor]:
    """
    Performs torch.mul(x, torch.sigmoid(torch.relu((x @ a)) @ b))
    Args:
        x: 2D tensor of shape [m, n].
        a: 2D tensor of shape [n, k].
        b: 2D tensor of shape [k, n].
    Returns:
        out: Resulting matrix of shape [m, n].
        c = torch.relu(x @ a) of shape [m, k].
        d = torch.sigmoid(c @ b) of shape [m, n].
    """
    m, n = x.size()
    k = a.size(1)

    out = torch.empty([m, n], dtype=x.dtype, device=x.device)
    c = torch.empty([m, k], dtype=x.dtype, device=x.device)
    d = torch.empty([m, n], dtype=x.dtype, device=x.device)

    for tile_m in hl.tile(m):
        # Compute c = relu(x @ a) for this tile_m
        for tile_k in hl.tile(k):
            partial_xa = x[tile_m, :] @ a[:, tile_k]
            c[tile_m, tile_k] = torch.relu(partial_xa)

        # Compute d = sigmoid(c @ b) and out = x * d for this tile_m
        for tile_n in hl.tile(n):
            acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
            for tile_k in hl.tile(k):
                acc = torch.addmm(acc, c[tile_m, tile_k], b[tile_k, tile_n])
            # Match the input-dtype matmul result before applying sigmoid.
            d[tile_m, tile_n] = torch.sigmoid(acc.to(x.dtype))
            out[tile_m, tile_n] = x[tile_m, tile_n] * d[tile_m, tile_n]

    return out, c, d

@helion.kernel(static_shapes=True)
def squeeze_and_excitation_net_bwd_dc(
    grad_out: Tensor, x: Tensor, b: Tensor, c: Tensor, d: Tensor
) -> tuple[Tensor, Tensor]:
    """
    Compute the intermediates shared by every weight and input gradient.
    grad_cb = grad_out * x * d * (1 - d)   (gradient of c @ b, shape [m, n])
    grad_c = (grad_cb @ b.T) * (c > 0)     (gradient of x @ a, shape [m, k])

    The pointwise phase materializes grad_cb, a barrier makes it visible, and
    the tiled matmul phase reduces over n with an fp32 accumulator before the
    relu-backward mask. Both intermediates are stored in the input dtype so
    the dx, da and db kernels can each be a single tiled matmul.
    """
    m, n = x.size()
    k = c.size(1)

    grad_cb = torch.empty([m, n], dtype=x.dtype, device=x.device)
    grad_c = torch.empty([m, k], dtype=x.dtype, device=x.device)

    for point_m, point_n in hl.tile((m, n)):
        grad_cb[point_m, point_n] = (
            grad_out[point_m, point_n]
            * x[point_m, point_n]
            * d[point_m, point_n]
            * (1.0 - d[point_m, point_n])
        )
    hl.barrier()

    for row_c, col_c in hl.tile((m, k)):
        acc_c = hl.zeros([row_c, col_c], dtype=torch.float32)
        for reduce_n in hl.tile(n):
            acc_c = torch.addmm(acc_c, grad_cb[row_c, reduce_n], b[col_c, reduce_n].T)
        grad_c[row_c, col_c] = acc_c.to(x.dtype) * (c[row_c, col_c] > 0)

    return grad_cb, grad_c

@helion.kernel(static_shapes=True)
def squeeze_and_excitation_net_bwd_dx(
    grad_out: Tensor, a: Tensor, d: Tensor, grad_c: Tensor
) -> Tensor:
    """
    Compute grad_x for the squeeze and excitation network.
    grad_x = grad_out * d + grad_c @ a.T

    A single tiled matmul over [m, n] reducing over k, with the gating
    branch's gradient grad_out * d added in the epilogue. grad_c comes from
    squeeze_and_excitation_net_bwd_dc.
    """
    m, n = grad_out.size()
    k = a.size(1)

    grad_x = torch.empty([m, n], dtype=grad_out.dtype, device=grad_out.device)

    for row_x, col_x in hl.tile((m, n)):
        acc_x = hl.zeros([row_x, col_x], dtype=torch.float32)
        for reduce_k in hl.tile(k):
            acc_x = torch.addmm(acc_x, grad_c[row_x, reduce_k], a[col_x, reduce_k].T)
        grad_x[row_x, col_x] = grad_out[row_x, col_x] * d[row_x, col_x] + acc_x.to(
            grad_out.dtype
        )

    return grad_x

@helion.kernel(static_shapes=True)
def squeeze_and_excitation_net_bwd_da(x: Tensor, grad_c: Tensor) -> Tensor:
    """
    Compute grad_a for the squeeze and excitation network.
    grad_a = x.T @ grad_c

    A single tiled matmul over [n, k] reducing over m. grad_c comes from
    squeeze_and_excitation_net_bwd_dc.
    """
    m, n = x.size()
    k = grad_c.size(1)

    grad_a = torch.empty([n, k], dtype=x.dtype, device=x.device)

    for row_a, col_a in hl.tile((n, k)):
        acc_a = hl.zeros([row_a, col_a], dtype=torch.float32)
        for reduce_m in hl.tile(m):
            acc_a = torch.addmm(acc_a, x[reduce_m, row_a].T, grad_c[reduce_m, col_a])
        grad_a[row_a, col_a] = acc_a

    return grad_a

@helion.kernel(static_shapes=True)
def squeeze_and_excitation_net_bwd_db(c: Tensor, grad_cb: Tensor) -> Tensor:
    """
    Compute grad_b for the squeeze and excitation network.
    grad_b = c.T @ grad_cb

    A single tiled matmul over [k, n] reducing over m. grad_cb comes from
    squeeze_and_excitation_net_bwd_dc.
    """
    m, n = grad_cb.size()
    k = c.size(1)

    grad_b = torch.empty([k, n], dtype=grad_cb.dtype, device=grad_cb.device)

    for tile_k, tile_n in hl.tile([k, n]):
        acc = hl.zeros([tile_k, tile_n], dtype=torch.float32)
        for tile_m in hl.tile(m):
            acc = torch.addmm(acc, c[tile_m, tile_k].T, grad_cb[tile_m, tile_n])
        grad_b[tile_k, tile_n] = acc

    return grad_b

Reference Implementation

def squeeze_and_excitation_net_pytorch(
    x: torch.Tensor, a: torch.Tensor, b: torch.Tensor
) -> torch.Tensor:
    """
    PyTorch reference implementation of squeeze_and_excitation_net.

    Args:
        x, a, b: Input tensors

    Returns:
        tensor of torch.mul(x, torch.sigmoid(torch.relu((x @ a)) @ b))
    """
    return torch.mul(x, torch.sigmoid(torch.relu(x @ a) @ b))

Autograd Function

class SqueezeAndExcitationNetFunction(torch.autograd.Function):
    @staticmethod
    def forward(  # type: ignore[override]
        ctx: object,
        x: torch.Tensor,
        a: torch.Tensor,
        b: torch.Tensor,
    ) -> torch.Tensor:
        """Forward pass for squeeze and excitation network."""
        out, c, d = squeeze_and_excitation_net_fwd(x, a, b)
        ctx.save_for_backward(x, a, b, c, d)  # type: ignore[attr-defined]
        return out

    @staticmethod
    def backward(  # type: ignore[override]
        ctx: object,
        grad_out: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Backward pass for squeeze and excitation network.

        One pointwise phase plus four matmuls: grad_cb and grad_c are
        computed once (a pointwise phase and the grad_c matmul, joined by a
        barrier), then each of grad_x, grad_a and grad_b is a single matmul
        over the shared intermediates.
        """
        x, a, b, c, d = ctx.saved_tensors  # type: ignore[attr-defined]

        grad_cb, grad_c = squeeze_and_excitation_net_bwd_dc(grad_out, x, b, c, d)
        grad_x = squeeze_and_excitation_net_bwd_dx(grad_out, a, d, grad_c)
        grad_a = squeeze_and_excitation_net_bwd_da(x, grad_c)
        grad_b = squeeze_and_excitation_net_bwd_db(c, grad_cb)
        return grad_x, grad_a, grad_b


def squeeze_and_excitation_net(
    x: torch.Tensor, a: torch.Tensor, b: torch.Tensor
) -> torch.Tensor:
    """
    Squeeze and excitation network with autograd support.

    Args:
        x: Input tensor [m, n]
        a: Weight matrix [n, k]
        b: Weight matrix [k, n]

    Returns:
        Output tensor [m, n]
    """
    return SqueezeAndExcitationNetFunction.apply(x, a, b)  # type: ignore[no-any-return]


def check(m: int, k: int, n: int) -> None:
    """
    Checks the correctness against PyTorch.
    Args:
        m (int): Number of rows in matrix x.
        n (int): Number of columns in matrix x.
        k (int): Number of columns in matrix a.
    """
    x = torch.randn([m, n], device=DEVICE, dtype=HALF_DTYPE, requires_grad=True)
    a = torch.randn([n, k], device=DEVICE, dtype=HALF_DTYPE, requires_grad=True)
    b = torch.randn([k, n], device=DEVICE, dtype=HALF_DTYPE, requires_grad=True)
    for bwd in [True, False]:
        run_example(
            squeeze_and_excitation_net,
            squeeze_and_excitation_net_pytorch,
            (x, a, b),
            bwd=bwd,
            # The backward kernels are exact given the forward's intermediates
            # (the shared grad_cb and grad_c are rounded to half precision
            # once, where the half-precision reference rounds them too), but
            # c = relu(x @ a) is rounded to half precision after an fp32
            # accumulation whose order differs from cuBLAS, so ~0.1% of its
            # elements differ by one ulp and the full-row reductions in the
            # backward turn each flip into a whole-row shift of grad_x. Judge
            # gradients by relative L2 instead: measured up to ~1.1e-2 at
            # k=1024, where the half-precision reference is itself ~7e-2 from
            # an fp32 ground truth.
            bwd_relative_l2=3e-2,
        )

def main() -> None:
    """
    Main function to run correctness checks.
    """
    check(1024, 1024, 1024)

if __name__ == "__main__":
    main()