| # kernel_image_resize — how every op works (study notes) |
|
|
| This explains the whole package end to end: the resampling math, the data layout, and |
| every kernel op, plus the benchmark findings. It is meant for pen & paper — there is a |
| fully worked numeric example you can reproduce by hand in the "Worked example" section. |
|
|
| The package does one thing: **resize + rescale + normalize**, the op sequence a |
| `transformers` fast image processor runs (`TorchvisionBackend`: resize, then |
| `(x*rescale - mean)/std`), as a single GPU pipeline. Input: raw CHW `uint8` images |
| (any size, ragged). Output: `(N, C, out_h, out_w)` normalized `float32`. |
|
|
| --- |
|
|
| ## 1. The one idea behind resizing |
|
|
| Resizing does not "pick" pixels; each **output** pixel is a **weighted average of a small |
| window of input pixels**. Two things define that average: |
|
|
| 1. **Where** in the input an output pixel lands (its center). |
| 2. **Which** input pixels are in its window, and with **what weights**. |
|
|
| The weight of an input pixel falls off with distance from the center, following a filter |
| curve: |
|
|
| - **bilinear** → a triangle, width 1 on each side (so 2 input pixels per axis normally). |
| - **bicubic** → a cubic bump, width 2 on each side (so 4 input pixels per axis normally). |
|
|
| When you **shrink** an image, you must also **blur first** or you get aliasing. That is |
| what `antialias=True` does: it widens the window so each output pixel averages more input |
| pixels (a low-pass filter before throwing pixels away). Widening is proportional to the |
| shrink factor, so shrinking 3× turns a 4-tap bicubic into ~13 taps. |
|
|
| --- |
|
|
| ## 2. The resampling-weight formula (the heart of everything) |
|
|
| All kernels use the same formula, which matches PyTorch's aten `UpSampleKernel` |
| (`align_corners=False`, "half-pixel" convention). For one axis: |
|
|
| ``` |
| scale = in_size / out_size # > 1 means shrinking |
| interp_half = 1 (bilinear) or 2 (bicubic) # half-width of the filter |
| cubic_a = -0.75 (no antialias) or -0.5 (antialias) # the cubic curve's shape constant |
| |
| # antialias only widens the window, and only when shrinking: |
| if antialias and scale > 1: |
| eff = scale # window widens by the shrink factor |
| else: |
| eff = 1 # plain 2-tap / 4-tap interpolation |
| support = interp_half * eff # half-window width, in INPUT pixels |
| inv = 1 / eff # squashes the filter curve to match the widened window |
| |
| # for output index i: |
| center = scale * (i + 0.5) # input coordinate this output maps to |
| first_tap = floor(center - support + 0.5) # leftmost input pixel in the window |
| |
| # for each tap t = 0, 1, 2, ... (up to MAX_TAPS): |
| tap_pos = first_tap + t # an input pixel index |
| arg = (tap_pos - center + 0.5) * inv # distance from center, squashed |
| weight = filter(arg) # triangle or cubic, see below |
| ``` |
|
|
| The filter (`_resample_weight` in `_fused.py`), with `x = |arg|`: |
|
|
| ``` |
| bilinear: max(1 - x, 0) # triangle, zero past 1 |
| |
| bicubic: x <= 1 : (a+2)x^3 - (a+3)x^2 + 1 |
| 1<x<2 : a x^3 - 5a x^2 + 8a x - 4a |
| else : 0 |
| ``` |
|
|
| Two edge rules (both kernels do this identically): |
|
|
| - **non-antialias**: clamp the tap index into `[0, in_size-1]` → replicates the border |
| pixel. The filter weights of a standard 2/4-tap interpolation already sum to 1. |
| - **antialias**: instead set the weight to **0** for taps that fall off the image |
| (`tap_pos < 0` or `>= in_size`), then **renormalize** by dividing by the sum of the |
| realized weights. This keeps the average correct at the edges. |
|
|
| That renormalization is why every kernel computes a `weight_sum` and divides by it. For |
| the non-antialias case `weight_sum == 1`, so the division is a harmless no-op. |
|
|
| --- |
|
|
| ## 3. Worked example (do this by hand) |
|
|
| **bilinear, no antialias, one axis, in_size=4, out_size=2.** |
|
|
| ``` |
| scale = 4/2 = 2, interp_half = 1, eff = 1, support = 1, inv = 1 |
| ``` |
|
|
| Output pixel `i = 0`: |
| ``` |
| center = 2 * (0 + 0.5) = 1.0 |
| first_tap = floor(1.0 - 1 + 0.5) = floor(0.5) = 0 |
| t=0: tap_pos=0, arg=(0-1.0+0.5)= -0.5 -> weight = 1-0.5 = 0.5 |
| t=1: tap_pos=1, arg=(1-1.0+0.5)= 0.5 -> weight = 1-0.5 = 0.5 |
| t=2: tap_pos=2, arg=(2-1.0+0.5)= 1.5 -> weight = max(1-1.5,0) = 0 |
| weight_sum = 1.0 |
| output[0] = (0.5*in[0] + 0.5*in[1]) / 1.0 # halfway between in[0] and in[1] |
| ``` |
|
|
| Output pixel `i = 1`: |
| ``` |
| center = 2 * 1.5 = 3.0 |
| first_tap = floor(3.0 - 0.5) = 2 |
| t=0: tap_pos=2, arg=-0.5 -> 0.5 |
| t=1: tap_pos=3, arg= 0.5 -> 0.5 |
| t=2: tap_pos=4, arg= 1.5 -> 0 (index 4 would clamp to 3, but weight is 0 anyway) |
| output[1] = 0.5*in[2] + 0.5*in[3] |
| ``` |
|
|
| This 1-D operation is exactly one pass of the separable kernel. The 2-D result is the same |
| formula applied on both axes (rows and columns). |
|
|
| --- |
|
|
| ## 4. Data layout (host side, `_pack.py`) |
| |
| Ragged images (different H×W) cannot be stacked into one tensor, so they are flattened and |
| concatenated into one buffer, with side tables describing each image. |
| |
| `pack_images(images, dtype)` → |
| ``` |
| input_pixels : 1-D buffer, all images flattened (C,H,W row-major) and concatenated |
| offsets[n] : element index where image n starts |
| heights[n], widths[n] : that image's H and W |
| channels : C (shared by all images) |
| ``` |
| Address of input pixel `(channel, row, col)` of image `n`: |
| ``` |
| input_pixels[ offsets[n] + channel*(H*W) + row*W + col ] |
| ``` |
| The separable path packs as `uint8` (1 byte/pixel, half the memory traffic of float). |
|
|
| `fold_mean_std(mean, std, rescale)` → folds the rescale factor into the normalization |
| constants so the kernel does a single `(x - m)/s`: |
| ``` |
| m = mean / rescale s = std / rescale |
| (x - m)/s == (x*rescale - mean)/std # identical to the processor's fused normalize |
| ``` |
|
|
| `max_taps(images, out_size, axis, interp, antialias)` → the **widest** window in the batch |
| = `ceil(support) * 2 + 1`. A Triton loop bound must be a compile-time constant, so every |
| program loops this fixed count; taps beyond a given pixel's real window get ~0 weight. |
|
|
| `as_image_list` → accepts a stacked `(N,C,H,W)` tensor or a list, always returns a list. |
|
|
| --- |
|
|
| ## 5. Fused kernel (`_fused.py`, `backend="fused"`) |
| |
| One launch. **One program = one image + a BLOCK of its output pixels.** Each output pixel |
| reads the **full 2-D window** directly: `MAX_TAPS_H × MAX_TAPS_W` input pixels. |
| |
| ``` |
| grid = (num_images, ceil(out_h*out_w / BLOCK)) |
|
|
| per lane (one output pixel): |
| oy, ox = (flat_index // out_w, flat_index % out_w) |
| center_y, center_x, first_tap_y, first_tap_x # section 2, both axes |
|
|
| # weight_sum factorizes across axes (separable math, even though the LOADS are 2-D): |
| sum_wy = Σ_ty filter_y ; sum_wx = Σ_tx filter_x ; denom = sum_wy * sum_wx |
| |
| for channel: |
| acc = 0 |
| for ty in 0..MAX_TAPS_H: # <-- the 2-D window: TAPS_H * TAPS_W loads |
| for tx in 0..MAX_TAPS_W: |
| weight = filter_y(ty) * filter_x(tx) |
| pixel = input_pixels[channel, clamp(tap_y), clamp(tap_x)] |
| acc += weight * pixel |
| acc = acc / denom |
| out[image, channel, oy, ox] = (acc - mean[channel]) / std[channel] |
| ``` |
| |
| Cost per output pixel: `TAPS_H * TAPS_W` loads (e.g. 13×13 = **169**). Correct and simple, |
| but the 2-D load count is what makes it slow — hence the separable version. |
|
|
| --- |
|
|
| ## 6. Separable kernel (`_separable.py`, `backend="separable"`, the default) |
| |
| Same math, but the 2-D window is done as **two 1-D passes**, with a float intermediate |
| buffer in between. Loads per output pixel: `TAPS_W + TAPS_H` (e.g. 13+13 = **26**). |
| |
| ``` |
| input_pixels (uint8, C×H×W) --pass1--> intermediate (float, C×H×out_w) --pass2--> output (float, C×out_h×out_w) |
| resize W resize H + normalize |
| ``` |
| |
| The **intermediate** is the key object: same **height** as the input, but already the |
| **final width**. ("Tall and narrow.") It is also ragged in height, so it gets its own |
| offset table (built in `separable_resize_normalize`, same scheme as `pack_images`). |
|
|
| ### Pass 1 — `_horizontal_resize_kernel` (resize width only) |
| |
| ``` |
| grid = (num_images, ceil(H*out_w / BLOCK)) # work = every input row × every output col |
| per lane: |
| input_row = flat_index // out_w # row index, UNCHANGED by this pass |
| out_col = flat_index % out_w # output column being computed |
| |
| center_x, first_tap_x, col_weight_sum # section 2, COLUMN axis only |
| |
| for channel: |
| acc = 0 |
| for tap in 0..MAX_TAPS_COL: # <-- 1-D: only TAPS_W loads |
| weight = filter_x(tap) |
| pixel = input_pixels[channel, input_row, clamp(tap_col)] # uint8 -> float |
| acc += weight * pixel |
| acc = acc / col_weight_sum |
| intermediate[channel, input_row, out_col] = acc # NO normalize yet |
| ``` |
| |
| Reads original `uint8` bytes; writes `float32`. No normalization here. |
|
|
| ### Pass 2 — `_vertical_resize_normalize_kernel` (resize height, then normalize) |
|
|
| ``` |
| grid = (num_images, ceil(out_h*out_w / BLOCK)) # work = every output pixel |
| per lane: |
| out_row = flat_index // out_w |
| out_col = flat_index % out_w |
| |
| center_y, first_tap_y, row_weight_sum # section 2, ROW axis only |
| |
| for channel: |
| acc = 0 |
| for tap in 0..MAX_TAPS_ROW: # <-- 1-D: only TAPS_H loads |
| weight = filter_y(tap) |
| pixel = intermediate[channel, clamp(tap_row), out_col] # float |
| acc += weight * pixel |
| acc = acc / row_weight_sum |
| out[image, channel, out_row, out_col] = (acc - mean[channel]) / std[channel] # normalize here |
| ``` |
|
|
| Two launches (an implicit sync between them), so pass 2 always sees pass 1's finished |
| output. |
|
|
| ### Why separable wins |
|
|
| `TAPS_W + TAPS_H` loads instead of `TAPS_W * TAPS_H`. For a 13×13 window that is 26 vs 169. |
| This is exactly the algorithm PIL and torchvision use. The catch: an extra full-size float |
| intermediate buffer (more memory traffic), but the read-count reduction dominates. |
|
|
| Parity note: the intermediate here is **float32**; torchvision keeps a **fixed-point |
| uint8** intermediate. So the separable output is parity-*close* to torchvision, not |
| bit-identical — and the float version is actually the more accurate one. |
|
|
| --- |
|
|
| ## 7. Public API (`__init__.py`) |
|
|
| ``` |
| resize_normalize(images, size, image_mean, image_std, |
| rescale_factor=1/255, resample="bilinear", antialias=False, |
| backend="separable", block=256) |
| ``` |
| - `images`: stacked `(N,C,H,W)` tensor or a list of CHW tensors. |
| - `size`: int (square), `(H,W)`, or `{"height","width"}`. |
| - `resample`: `"bilinear"`/`"bicubic"`, or a PIL resample int (0/2→bilinear, 3→bicubic). |
| - `backend`: `"separable"` (default, fastest) or `"fused"` (2-D reference). |
| - `resize_normalize_ragged`: same kernels, list-only. |
|
|
| --- |
|
|
| ## 8. Benchmark findings (A100, CUDA_VISIBLE_DEVICES=1) |
|
|
| ### Standalone resize+normalize — SigLIP-so400m config, N=32 ragged 384–1024², out 384², bicubic+AA |
| ``` |
| torchvision eager loop : 2.91 ms (per-image float loop) |
| torchvision compiled : 5.70 ms (torch.compile dynamic, per-image; slower than eager) |
| torchvision compiled pkt: 2.55 ms (one graph over a padded stack; timing only) |
| fused triton (2D) : 11.49 ms (taps*taps; the slow reference) |
| separable triton (uint8): 1.29 ms (taps+taps) <-- fastest |
| real processor : 3.92 ms |
| ``` |
| **Separable is ~3× the real processor**, parity ≤1e-4 vs torchvision-float. The fused 2-D |
| loses for the algorithmic reason above (169 vs 26 loads). `torch.compile` does not help: |
| per-image it is *slower* (dispatch overhead over 32 ragged shapes); even as one packed |
| graph it only matches the eager loop, because inductor's interpolate is no faster than aten |
| resize. |
|
|
| ### End-to-end inference — Siglip2-base-patch16-224, **bf16** forward |
| ``` |
| preprocess forward(fixed input) preprocess+forward |
| processor 3.99 12.86 14.44 |
| separable 0.93 13.02 13.76 <-- ~5% faster e2e |
| fused 2.00 13.01 14.79 |
| compiled 6.14 12.89 14.00 |
| feature parity (separable/fused/compiled vs processor): 9.38e-2 = 1.2% of feature max |
| ``` |
| - `forward(fixed input)` is identical (~12.9 ms) for all → **no inference regression**; the |
| model does not care which preprocessor made the tensor. |
| - The 1.2% feature drift is the float-vs-uint8 resize difference, identical across all |
| float backends → not a bug. The float path is the more accurate one. |
| - End-to-end win is ~5% with a bf16 forward (was ~0.5% with fp32, where the forward was |
| ~80 ms). **The win scales with how preprocessing-bound you are.** |
|
|
| ### Data path from JPEG bytes — 552 KB/img |
| ``` |
| CPU decode + torchvision resize : 177.5 ms (status quo) |
| CPU decode + separable kernel : 176.4 ms (kernel saves ~1 ms; decode dominates) |
| GPU decode (nvJPEG) + kernel : 14.8 ms (fully on-GPU) |
| ``` |
| - ~175 ms of the 177 ms is **CPU JPEG decode + host→device copy**. Resize/normalize is ~1%. |
| - The 12× win (177→15) is **GPU decode (nvJPEG)**, i.e. `torchvision.io.decode_jpeg(device="cuda")` |
| — *not* the kernel. The kernel is the resize/normalize component of that GPU pipeline. |
|
|
| --- |
|
|
| ## 9. What is true / what to claim |
|
|
| - The kernel is **correct** (≤1e-4 vs torchvision-float, more accurate than the processor's |
| uint8 path) and feeds the model with **no inference regression**. |
| - It is **~3× the real processor at the resize/normalize stage** — a real, parity-clean win. |
| - It does **not** speed up preprocessing 12×. Decode dominates the data path; the GPU-decode |
| lever is nvJPEG, a torchvision feature, not this kernel. |
| - The kernel matters end-to-end only once you are **not decode-bound**: in a GPU-decode |
| pipeline it keeps resize/normalize minimal (~10% of that pipeline), and its standalone |
| preprocess win shows up when the forward is small (bf16, small model, large batch). |
| - Honest one-liner: *"GPU-native resize+normalize, 3× the fast processor at that stage, |
| drop-in for a GPU-decode pipeline."* |
|
|
