Flow warp ops: et_vk.warp, et_vk.warp_blend (#24692)

Summary:

Custom ops for optical-flow frame interpolation networks:
- `et_vk.warp(input, flow)` samples `input` bilinearly at `out_px + flow`, with border padding. A flow warp written in PyTorch builds a normalized grid and calls `grid_sampler_2d`; this op skips the grid, which removes its dispatches and is more accurate in fp16, because the grid's normalized coordinates lose precision at large resolutions.
- `et_vk.warp_blend(img0, img1, flow, mask)` fuses the final step of such a network into one shader: a bilinear upsample of `flow` and `mask` to the image size (when they are smaller), two warps (flow channels 0-1 and 2-3), a sigmoid on the mask, the blend and the clamp.

Both have a reference implementation and a `Meta` kernel, so H and W can be dynamic at export. The partitioner only delegates shapes the shaders handle: flow `[N, 2, H, W]` for `warp`, flow `[N, 4, h, w]` and mask `[N, 1, h, w]` for `warp_blend`, with no broadcasting. The ops have no CPU kernel, so a rejected node fails at load. The two shaders share a bilinear border-sampling helper. Graphs that do not emit these ops are unchanged.

`etvk_ops_bench` gets CPU references for both ops and `operators/warp/cases/correctness.json`: fp32 and fp16, N=2, C in {3, 4, 5, 8}, a 1-pixel-wide input, flows up to +-100 px, and same-size, 2x, non-integer and downsampled flow for `warp_blend`. The `warp` cases are added to the operator lists of `test-vulkan-operators-linux` in `pull.yml` and of the GPU job in `vulkan.yml`.

Reviewed By: SS-JIA

Differential Revision: D129150747

Pull Request resolved: #24692

Author: youxie

.github/workflows/pull.yml

@@ -2084,7 +2084,8 @@ jobs:
         # Custom operator tests
         PYTHON_EXECUTABLE=python bash backends/vulkan/test/ops/benchmark_v2/build.sh
         for op in binary/add gemm/linear/qw/q8csw convolution/qw/q8csw gemm/linear/qw/q4gsw \
-                  quantization/choose_qparams_per_row quantize/q8ta clone/q8ta binary/q8ta; do
+                  quantization/choose_qparams_per_row quantize/q8ta clone/q8ta binary/q8ta \
+                  warp; do
           ./cmake-out/backends/vulkan/test/ops/benchmark_v2/etvk_ops_bench --validation_only \
             --config=backends/vulkan/test/ops/benchmark_v2/operators/${op}/cases/correctness.json
         done

.github/workflows/vulkan.yml

@@ -123,7 +123,8 @@ jobs:
         # selected above.
         PYTHON_EXECUTABLE=python bash backends/vulkan/test/ops/benchmark_v2/build.sh
         for op in binary/add gemm/linear/qw/q8csw convolution/qw/q8csw gemm/linear/qw/q4gsw \
-                  quantization/choose_qparams_per_row quantize/q8ta clone/q8ta binary/q8ta; do
+                  quantization/choose_qparams_per_row quantize/q8ta clone/q8ta binary/q8ta \
+                  warp; do
           ./cmake-out/backends/vulkan/test/ops/benchmark_v2/etvk_ops_bench --validation_only \
             --config=backends/vulkan/test/ops/benchmark_v2/operators/${op}/cases/correctness.json
         done

backends/vulkan/custom_ops_lib.py

@@ -1349,3 +1349,133 @@ def q4gsw_requant_meta(
 lib.impl(name, q4gsw_requant_impl, "CompositeExplicitAutograd")
 lib.impl(name, q4gsw_requant_meta, "Meta")
 q4gsw_requant_op = getattr(getattr(torch.ops, namespace), name)
+
+
+################
+## warp_blend ##
+################
+
+
+def _warp_blend_warp(img: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
+    # Bilinear warp by a pixel-unit flow field with border padding, written as
+    # a normalized grid + grid_sample (align_corners=True). The Vulkan shader
+    # samples at out_px + flow directly, which is the same coordinate.
+    _, _, height, width = img.shape
+    xs = torch.linspace(-1.0, 1.0, width, dtype=img.dtype, device=img.device)
+    ys = torch.linspace(-1.0, 1.0, height, dtype=img.dtype, device=img.device)
+    grid_y, grid_x = torch.meshgrid(ys, xs, indexing="ij")
+    base_grid = torch.stack((grid_x, grid_y), dim=-1).unsqueeze(0)
+    # A 1-pixel axis has no extent to scale by; border padding clamps every
+    # sample to it, so any scale gives the same result.
+    norm_scale = torch.tensor(
+        [2.0 / max(width - 1.0, 1.0), 2.0 / max(height - 1.0, 1.0)],
+        dtype=img.dtype,
+        device=img.device,
+    )
+    grid = base_grid + flow.permute(0, 2, 3, 1) * norm_scale
+    return torch.nn.functional.grid_sample(
+        img, grid, mode="bilinear", padding_mode="border", align_corners=True
+    )
+
+
+def warp_blend_impl(
+    img0: torch.Tensor,
+    img1: torch.Tensor,
+    flow: torch.Tensor,
+    mask: torch.Tensor,
+) -> torch.Tensor:
+    # img0/img1: [N, C, Hout, Wout]. flow: [N, 4, Hin, Win] displacements in
+    # flow pixels (ch 0,1 warp img0; ch 2,3 warp img1). mask: [N, 1, Hin, Win]
+    # logits. A smaller flow/mask is bilinearly upsampled to the output size and
+    # the displacement rescaled per axis; the shader does the same in place.
+    _, _, out_h, out_w = img0.shape
+    _, _, in_h, in_w = flow.shape
+    if in_h != out_h or in_w != out_w:
+        flow = torch.nn.functional.interpolate(
+            flow, size=(out_h, out_w), mode="bilinear", align_corners=False
+        )
+        sx = out_w / in_w
+        sy = out_h / in_h
+        flow = flow * torch.tensor(
+            [sx, sy, sx, sy], dtype=flow.dtype, device=flow.device
+        ).view(1, 4, 1, 1)
+        mask = torch.nn.functional.interpolate(
+            mask, size=(out_h, out_w), mode="bilinear", align_corners=False
+        )
+    warped0 = _warp_blend_warp(img0, flow[:, :2])
+    warped1 = _warp_blend_warp(img1, flow[:, 2:4])
+    m = torch.sigmoid(mask)
+    merged = warped0 * m + warped1 * (1.0 - m)
+    return torch.clamp(merged, 0.0, 1.0)
+
+
+def warp_blend_meta(
+    img0: torch.Tensor,
+    img1: torch.Tensor,
+    flow: torch.Tensor,
+    mask: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(img0)
+
+
+name = "warp_blend"
+lib.define(f"{name}(Tensor img0, Tensor img1, Tensor flow, Tensor mask) -> Tensor")
+lib.impl(name, warp_blend_impl, "CompositeExplicitAutograd")
+# Export only needs the output shape; tracing the reference would specialize
+# H and W through linspace and the grid scale.
+lib.impl(name, warp_blend_meta, "Meta")
+warp_blend_op = getattr(getattr(torch.ops, namespace), name)
+
+
+# When lowering to executorch, ops are converted from default to out variant.
+def warp_blend_out_impl(
+    img0: torch.Tensor,
+    img1: torch.Tensor,
+    flow: torch.Tensor,
+    mask: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    out = warp_blend_impl(img0, img1, flow, mask)
+    return out
+
+
+name = "warp_blend.out"
+lib.define(
+    f"{name}(Tensor img0, Tensor img1, Tensor flow, Tensor mask, *, Tensor(a!) out) -> Tensor(a!)"
+)
+lib.impl(name, warp_blend_out_impl, "CompositeExplicitAutograd")
+
+##########
+## warp ##
+##########
+
+
+def warp_impl(input: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
+    # input: [N, C, H, W]; flow: [N, 2, H, W] displacement in pixels.
+    return _warp_blend_warp(input, flow)
+
+
+def warp_meta(input: torch.Tensor, flow: torch.Tensor) -> torch.Tensor:
+    return torch.empty_like(input)
+
+
+name = "warp"
+lib.define(f"{name}(Tensor input, Tensor flow) -> Tensor")
+lib.impl(name, warp_impl, "CompositeExplicitAutograd")
+lib.impl(name, warp_meta, "Meta")
+warp_op = getattr(getattr(torch.ops, namespace), name)
+
+
+# When lowering to executorch, ops are converted from default to out variant.
+def warp_out_impl(
+    input: torch.Tensor,
+    flow: torch.Tensor,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    out = warp_impl(input, flow)
+    return out
+
+
+name = "warp.out"
+lib.define(f"{name}(Tensor input, Tensor flow, *, Tensor(a!) out) -> Tensor(a!)")
+lib.impl(name, warp_out_impl, "CompositeExplicitAutograd")

backends/vulkan/op_registry.py

@@ -1860,6 +1860,84 @@ def check_grid_sampler_2d_node(node: torch.fx.Node) -> bool:
     )
 
 
+# =============================================================================
+# Warp.cpp
+# =============================================================================
+
+
+def _node_vals(node: torch.fx.Node, n_args: int) -> Optional[List[torch.Tensor]]:
+    if len(node.args) < n_args:
+        return None
+    vals = []
+    for arg in list(node.args[:n_args]) + [node]:
+        if not isinstance(arg, torch.fx.Node) or arg.meta.get("val") is None:
+            return None
+        vals.append(arg.meta["val"])
+    return vals
+
+
+def _check_warp_shapes(node: torch.fx.Node) -> bool:
+    """Reject shapes the Vulkan warp shader can't handle, such as a flow that
+    broadcasts over batch or space. These ops have no CPU kernel, so a rejected
+    node fails at load time."""
+    vals = _node_vals(node, 2)
+    if vals is None:
+        return False
+    inp, flow, out = vals
+    return (
+        inp.dtype == flow.dtype == out.dtype
+        and inp.shape == out.shape
+        and flow.dim() == 4
+        and flow.shape[:2] == (out.shape[0], 2)
+        and flow.shape[-2:] == out.shape[-2:]
+    )
+
+
+def _check_warp_blend_shapes(node: torch.fx.Node) -> bool:
+    """Reject shapes the Vulkan warp_blend shader can't handle, such as a mask
+    with more than one channel or a flow or mask that broadcasts over batch.
+    These ops have no CPU kernel, so a rejected node fails at load time."""
+    vals = _node_vals(node, 4)
+    if vals is None:
+        return False
+    img0, img1, flow, mask, out = vals
+    return (
+        img0.dtype == img1.dtype == flow.dtype == mask.dtype == out.dtype
+        and img0.shape == out.shape
+        and img1.shape == out.shape
+        and flow.dim() == 4
+        and mask.dim() == 4
+        and flow.shape[:2] == (out.shape[0], 4)
+        and mask.shape[:2] == (out.shape[0], 1)
+        and flow.shape[-2:] == mask.shape[-2:]
+    )
+
+
+@update_features(exir_ops.edge.et_vk.warp.default)
+def register_warp():
+    return OpFeatures(
+        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        are_node_inputs_supported_fn=_check_warp_shapes,
+    )
+
+
+# =============================================================================
+# WarpBlend.cpp
+# =============================================================================
+
+
+@update_features(exir_ops.edge.et_vk.warp_blend.default)
+def register_warp_blend():
+    return OpFeatures(
+        inputs_storage=utils.CHANNELS_PACKED_TEXTURE,
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        are_node_inputs_supported_fn=_check_warp_blend_shapes,
+    )
+
+
 # =============================================================================
 # Repeat.cpp
 # =============================================================================

backends/vulkan/runtime/graph/ops/glsl/bilinear_border_sample.glslh

@@ -0,0 +1,33 @@
+/*
+ * 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.
+ */
+
+#ifndef BILINEAR_BORDER_SAMPLE_GLSLH
+#define BILINEAR_BORDER_SAMPLE_GLSLH
+
+/*
+ * Bilinear sample of a channels-packed texture at pixel coordinate c on texel
+ * slice z, with border padding: c is clamped to [0, max_xy].
+ */
+VEC4_T bilinear_border_sample(
+    sampler3D tex,
+    vec2 c,
+    const int z,
+    const ivec2 max_xy) {
+  c = clamp(c, vec2(0.0), vec2(max_xy));
+  const ivec2 lower = ivec2(floor(c));
+  const ivec2 upper = min(lower + ivec2(1), max_xy);
+  const vec2 w = c - vec2(lower);
+
+  const VEC4_T s00 = VEC4_T(texelFetch(tex, ivec3(lower.x, lower.y, z), 0));
+  const VEC4_T s10 = VEC4_T(texelFetch(tex, ivec3(upper.x, lower.y, z), 0));
+  const VEC4_T s01 = VEC4_T(texelFetch(tex, ivec3(lower.x, upper.y, z), 0));
+  const VEC4_T s11 = VEC4_T(texelFetch(tex, ivec3(upper.x, upper.y, z), 0));
+  return mix(mix(s00, s10, w.x), mix(s01, s11, w.x), w.y);
+}
+
+#endif // BILINEAR_BORDER_SAMPLE_GLSLH

backends/vulkan/runtime/graph/ops/glsl/warp.glsl

@@ -0,0 +1,57 @@
+/*
+ * 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.
+ */
+
+#version 450 core
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+#define PRECISION ${PRECISION}
+
+#define VEC4_T ${texel_load_type(DTYPE, STORAGE)}
+
+${define_active_storage_type(STORAGE)}
+
+layout(std430) buffer;
+
+#include "indexing.glslh"
+#include "bilinear_border_sample.glslh"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_flow", DTYPE, STORAGE)}
+
+${layout_declare_ubo(B, "TextureMetadata", "outp")}
+${layout_declare_ubo(B, "TextureMetadata", "inp")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+/*
+ * Bilinear warp by a pixel-unit flow field, with border padding:
+ *   out[n, c, y, x] = in[n, c, y + flow[n, 1, y, x], x + flow[n, 0, y, x]]
+ *
+ * in, out: channels-packed texture3d, [N, C, H, W]
+ * flow   : channels-packed texture3d, [N, 2, H, W] (one z-slice per batch)
+ */
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  if (out_of_bounds(pos, outp)) {
+    return;
+  }
+
+  const int num_z_per_n = outp.limits.z / outp.sizes.w;
+  const int n = pos.z / num_z_per_n;
+
+  const VEC4_T flow_texel = texelFetch(t_flow, ivec3(pos.x, pos.y, n), 0);
+  const vec2 c = vec2(pos.xy) + vec2(flow_texel.xy);
+
+  imageStore(
+      t_out,
+      pos,
+      bilinear_border_sample(t_in, c, pos.z, ivec2(inp.sizes.xy) - 1));
+}

backends/vulkan/runtime/graph/ops/glsl/warp.yaml

@@ -0,0 +1,16 @@
+# 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.
+
+warp:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: texture3d
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: warp