[KernelGen][thead] Add upsample_nearest_exact2d_backward vendor specialization (#207)

Co-authored-by: Dongxu-H <dxhan@baai.ac.cn>

Author: chx7514

src/flag_gems/runtime/backend/_thead/ops/__init__.py

@@ -31,11 +31,13 @@
     special_shifted_chebyshev_polynomial_w,
 )
 from .unsafe_masked_index_put_accumulate import _unsafe_masked_index_put_accumulate
+from .upsample_nearest_exact2d_backward import _upsample_nearest_exact2d_backward
 
 __all__ = [
     "_conv_depthwise2d",
     "_thnn_fused_lstm_cell_backward_impl",
     "_unsafe_masked_index_put_accumulate",
+    "_upsample_nearest_exact2d_backward",
     "adaptive_max_pool3d_backward",
     "broadcast_tensors",
     "broadcast_to",

src/flag_gems/runtime/backend/_thead/ops/upsample_nearest_exact2d_backward.py

@@ -0,0 +1,185 @@
+# Copyright 2026, The FlagOS Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Generated by KernelGen: https://github.com/flagos-ai/KernelGen
+import logging
+import math
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from flag_gems.runtime import device, torch_device_fn
+
+logger = logging.getLogger(__name__)
+
+
+# Gather-style backward kernel.
+#
+# The generic implementation scatters every output position into the input
+# gradient with ``tl.atomic_add`` accumulated in a float32 staging buffer,
+# then casts back to the input dtype in a separate kernel. On the PPU that
+# means three launches (zeros + scatter + cast) plus atomic contention, which
+# dominates latency on the small feature maps this op sees.
+#
+# Here we invert the mapping: each program owns one *input* gradient pixel and
+# sums the contiguous block of output positions that map to it. For the
+# nearest(-exact) forward index ``ih = floor(oh * reciprocal_scale)`` the set
+# of ``oh`` mapping to a given ``ih`` is the half-open range
+# ``[ceil(ih / reciprocal_scale), ceil((ih + 1) / reciprocal_scale))``. These
+# ranges tile ``[0, OH)`` exactly, so the gathered sum is over the identical
+# set of gradient values as the scatter — numerics are unchanged. Because no
+# lane writes the same address twice, atomics are unnecessary and the sum is
+# accumulated in fp32 registers and cast once on store, removing the float32
+# buffer and the trailing cast kernel entirely.
+@triton.jit
+def _upsample_nearest_exact2d_backward_gather_kernel(
+    grad_output_ptr,
+    grad_input_ptr,
+    M,  # total input gradient elements = N * C * IH * IW
+    OH,
+    OW,
+    IH,
+    IW,
+    scale_h,  # scales_h (outputs per input row) when USE_SCALES else unused
+    scale_w,  # scales_w when USE_SCALES else unused
+    USE_SCALES: tl.constexpr,
+    MAX_H: tl.constexpr,  # upper bound on output rows mapping to one input row
+    MAX_W: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < M
+
+    # Decompose the flat contiguous input offset into (nc, ih, iw).
+    iw = offs % IW
+    t = offs // IW
+    ih = t % IH
+    nc = t // IH  # flattened (n, c) plane index
+
+    base_out = nc * OH * OW
+
+    # Half-open output ranges [oh0, oh1) x [ow0, ow1) that map onto this input
+    # pixel, i.e. the inverse of the forward index floor(oh / scale).
+    if USE_SCALES:
+        oh0 = tl.ceil(ih.to(tl.float32) * scale_h).to(tl.int32)
+        oh1 = tl.ceil((ih + 1).to(tl.float32) * scale_h).to(tl.int32)
+        ow0 = tl.ceil(iw.to(tl.float32) * scale_w).to(tl.int32)
+        ow1 = tl.ceil((iw + 1).to(tl.float32) * scale_w).to(tl.int32)
+    else:
+        # Exact integer ceil(ih * OH / IH) = (ih * OH + IH - 1) // IH.
+        oh0 = (ih * OH + IH - 1) // IH
+        oh1 = ((ih + 1) * OH + IH - 1) // IH
+        ow0 = (iw * OW + IW - 1) // IW
+        ow1 = ((iw + 1) * OW + IW - 1) // IW
+
+    # The forward index is clamped to IH-1 / IW-1, so the last input row/col
+    # absorbs every trailing output position. Extend its range to the edge to
+    # keep the ranges a perfect partition of [0, OH) x [0, OW).
+    oh1 = tl.where(ih == IH - 1, OH, oh1)
+    ow1 = tl.where(iw == IW - 1, OW, ow1)
+
+    acc = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
+    for dh in tl.static_range(MAX_H):
+        oh = oh0 + dh
+        valid_h = mask & (oh < oh1)
+        row_off = base_out + oh * OW
+        for dw in tl.static_range(MAX_W):
+            ow = ow0 + dw
+            valid = valid_h & (ow < ow1)
+            g = tl.load(grad_output_ptr + row_off + ow, mask=valid, other=0.0)
+            acc += g.to(tl.float32)
+
+    tl.store(grad_input_ptr + offs, acc, mask=mask)
+
+
+def _upsample_nearest_exact2d_backward(
+    grad_output: torch.Tensor,
+    output_size: Tuple[int, int],
+    input_size: Tuple[int, int, int, int],
+    scales_h: Optional[float] = None,
+    scales_w: Optional[float] = None,
+) -> torch.Tensor:
+    logger.debug("GEMS_THEAD _UPSAMPLE_NEAREST_EXACT2D_BACKWARD")
+
+    assert grad_output.device.type == device.name
+    assert grad_output.ndim == 4, "The ndim of grad_output must be 4"
+    assert len(output_size) == 2, "The len of output_size must be 2"
+    assert len(input_size) == 4, "The len of input_size must be 4"
+
+    OH, OW = output_size
+    N, C, IH, IW = input_size
+
+    assert grad_output.shape == (N, C, OH, OW), (
+        f"grad_output shape {grad_output.shape} does not match "
+        f"expected shape (N={N}, C={C}, OH={OH}, OW={OW})"
+    )
+
+    grad_input = torch.empty(
+        (N, C, IH, IW), device=grad_output.device, dtype=grad_output.dtype
+    )
+
+    if grad_input.numel() == 0:
+        return grad_input
+
+    if not grad_output.is_contiguous():
+        grad_output = grad_output.contiguous()
+
+    # Forward source index uses reciprocal_scale = IH/OH (or 1/scales_h). When
+    # scales are given the gather range uses that scale directly; otherwise we
+    # use exact integer arithmetic in-kernel from OH/IH.
+    use_scales = scales_h is not None or scales_w is not None
+    scale_h = float(scales_h) if scales_h is not None else OH / IH
+    scale_w = float(scales_w) if scales_w is not None else OW / IW
+
+    # Tight upper bound on how many output rows/cols map to a single input
+    # pixel. Non-final ranges have width <= ceil(OH / IH); the clamped final
+    # pixel spans to OH so account for it explicitly.
+    def _max_span(o, i):
+        if i <= 0 or o <= 0:
+            return 1
+        base = (o + i - 1) // i
+        last0 = math.ceil((i - 1) * o / i)
+        return max(base, o - last0, 1)
+
+    MAX_H = _max_span(OH, IH)
+    MAX_W = _max_span(OW, IW)
+
+    M = N * C * IH * IW
+    # BLOCK=256: tuned for thead PPU — balances parallelism (grid = ceil(M/256))
+    # against per-block register pressure for the backward gather loop.
+    BLOCK = 256
+    grid = (triton.cdiv(M, BLOCK),)
+
+    with torch_device_fn.device(grad_output.device):
+        _upsample_nearest_exact2d_backward_gather_kernel[grid](
+            grad_output,
+            grad_input,
+            M,
+            OH,
+            OW,
+            IH,
+            IW,
+            scale_h,
+            scale_w,
+            use_scales,
+            MAX_H,
+            MAX_W,
+            BLOCK,
+            num_warps=4,
+        )
+
+    return grad_input