Implement aten.grid_sampler_2d.default op (#19982)

Differential Revision: D106866109

Pull Request resolved: #19982

Author: youxie

backends/vulkan/op_registry.py

@@ -1551,6 +1551,69 @@ def register_grid_priors():
     )
 
 
+# =============================================================================
+# GridSampler2d.cpp
+# =============================================================================
+
+
+@update_features(exir_ops.edge.aten.grid_sampler_2d.default)
+def register_grid_sampler_2d():
+    # The Vulkan implementation only supports the configuration used by RIFE's
+    # WarpModule: bilinear interpolation (0), border padding (1),
+    # align_corners=True. The C++ side has VK_CHECK_COND asserts for these,
+    # but those abort the whole inference at graph build — for any other model
+    # that contains a differently-configured grid_sampler_2d we want graceful
+    # CPU fallback, so we gate delegation here.
+    #
+    # Edge IR can hand us these scalar args as plain Python literals, SymInt /
+    # SymBool wrappers, or get_attr-style fx.Node references, so we unwrap
+    # each one defensively (mirrors the `isinstance(groups, int)` guard in
+    # check_conv_node / pick_conv_storage above). If we can't confidently pull
+    # a literal out of any arg, return False so the node stays on CPU instead
+    # of hitting a runtime VK_CHECK_COND.
+    def _unwrap_literal(arg: object) -> object:
+        # Plain Python literal (covers bool, since bool is a subclass of int).
+        if isinstance(arg, (bool, int, float)):
+            return arg
+        # get_attr / constant fx.Node — read the materialized value from meta.
+        if isinstance(arg, torch.fx.Node):
+            val = arg.meta.get("val", None)
+            if isinstance(val, (bool, int, float)):
+                return val
+            return None
+        # Symbolic int/bool (or anything else int-convertible) — try once.
+        try:
+            return int(arg)  # pyre-ignore[6]
+        except (TypeError, ValueError):
+            return None
+
+    def check_grid_sampler_2d_node(node: torch.fx.Node) -> bool:
+        # Schema: aten::grid_sampler_2d(input, grid, interpolation_mode,
+        #                               padding_mode, align_corners)
+        if len(node.args) < 5:
+            return False
+
+        interp = _unwrap_literal(node.args[2])
+        padding = _unwrap_literal(node.args[3])
+        align_corners = _unwrap_literal(node.args[4])
+
+        if interp is None or padding is None or align_corners is None:
+            return False
+
+        # mode: 0 = bilinear; padding: 1 = border; align_corners must be True.
+        return interp == 0 and padding == 1 and bool(align_corners) is True
+
+    return OpFeatures(
+        inputs_storage=[
+            utils.CHANNELS_PACKED_TEXTURE,  # input  : [N, C, Hin, Win]
+            utils.CONTIGUOUS_BUFFER,  # grid   : [N, Hout, Wout, 2]
+        ],
+        inputs_dtypes=utils.FP_T,
+        supports_resize=True,
+        are_node_inputs_supported_fn=check_grid_sampler_2d_node,
+    )
+
+
 # =============================================================================
 # Repeat.cpp
 # =============================================================================

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

@@ -0,0 +1,118 @@
+/*
+ * 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_required_extensions("buffer", DTYPE)}
+
+#define PRECISION ${PRECISION}
+
+#define VEC4_T ${texel_load_type(DTYPE, STORAGE)}
+#define T ${texel_load_component_type(DTYPE, "buffer")}
+
+${define_active_storage_type(STORAGE)}
+
+layout(std430) buffer;
+
+#include "indexing.glslh"
+
+${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
+// `t_grid` is always bound as a contiguous (width-packed) buffer of fp scalars
+// with logical shape [N, Hout, Wout, 2]. See add_grid_sampler_2d_node which
+// asserts this with `is_contiguous_buffer_tensor`.
+${layout_declare_tensor(B, "r", "t_grid", DTYPE, "buffer")}
+
+${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;
+
+// `out_layout` is passed for forward compatibility and is currently asserted
+// to be the standard channels-packed layout by `add_grid_sampler_2d_node`.
+// All texel math below assumes packed_dim = C (channels-packed), so the four
+// fp components of a texel share the same (N, Hout, Wout) and differ only in
+// channel. This lets one bilinear interpolation produce all 4 output channels.
+${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
+
+/*
+ * Vulkan implementation of `aten.grid_sampler_2d.default` for the
+ * specific configuration used by RIFE's `WarpModule`:
+ *   mode=bilinear, padding_mode=border, align_corners=true.
+ *
+ * Layout assumptions (validated in add_grid_sampler_2d_node):
+ *   - input  : channels-packed texture3d, shape [N, C, Hin, Win]
+ *   - grid   : contiguous (width-packed) buffer SSBO of fp scalars,
+ *              shape [N, Hout, Wout, 2] in normalized coords [-1, 1]
+ *   - output : channels-packed texture3d, shape [N, C, Hout, Wout]
+ *
+ * For channels-packed texture3d, the texel z extent is N * ceil(C/4),
+ * laid out as z = n * num_z_per_n + c_slice. Both input and output share
+ * the same N and C, so input z == output z.
+ *
+ * TextureMetadata layout (vtensor.md): sizes is WHCN order, so
+ *   outp.sizes.x = Wout, outp.sizes.y = Hout, outp.sizes.w = N.
+ *   outp.limits.z = N * ceil(C/4) (texel slices along z).
+ */
+void main() {
+  const ivec3 pos = ivec3(gl_GlobalInvocationID);
+
+  if (out_of_bounds(pos, outp)) {
+    return;
+  }
+
+  // Derive batch index from texel z. Each batch occupies `num_z_per_n`
+  // consecutive z-slices (one per 4-channel slice). Integer division by
+  // num_z_per_n picks out the batch.
+  const int N = outp.sizes.w;
+  const int num_z_per_n = outp.limits.z / N;
+  const int n = pos.z / num_z_per_n;
+
+  // Look up the (gx, gy) for this output pixel from the grid SSBO.
+  // The grid is a contiguous buffer of [N, Hout, Wout, 2], so the linear
+  // index for (n, h, w, comp) is ((n*Hout + h)*Wout + w)*2 + comp. This
+  // relies on `inputs_storage` in op_registry.py pinning grid to
+  // CONTIGUOUS_BUFFER and the C++ dispatcher re-checking with
+  // `is_contiguous_buffer_tensor` — see GridSampler2d.cpp.
+  const int Wout = outp.sizes.x;
+  const int Hout = outp.sizes.y;
+  const int grid_base = ((n * Hout + pos.y) * Wout + pos.x) * 2;
+  const float gx_norm = float(t_grid[grid_base + 0]);
+  const float gy_norm = float(t_grid[grid_base + 1]);
+
+  // Unnormalize for align_corners=true:
+  //   coord_pixel = (coord_norm + 1) * 0.5 * (size - 1)
+  // Input W/H come from inp.sizes (WHCN), not inp.limits (texel space).
+  const ivec2 max_in_xy = ivec2(inp.sizes.xy) - 1;
+  const float gx_pixel = (gx_norm + 1.0) * 0.5 * float(max_in_xy.x);
+  const float gy_pixel = (gy_norm + 1.0) * 0.5 * float(max_in_xy.y);
+
+  // padding_mode=border: clamp coordinates to [0, size-1].
+  const float gx = clamp(gx_pixel, 0.0, float(max_in_xy.x));
+  const float gy = clamp(gy_pixel, 0.0, float(max_in_xy.y));
+
+  const ivec2 lower = ivec2(floor(vec2(gx, gy)));
+  // Clamp ceil to valid range for samples on the border.
+  const ivec2 upper = clamp(lower + ivec2(1), ivec2(0), max_in_xy);
+  const vec2 w = vec2(gx, gy) - vec2(lower);
+
+  // Fetch the four nearest texels (each carries 4 channels). Because input
+  // is channels-packed, pos.z indexes the same channel slice in input as in
+  // output, so we can reuse pos.z directly without remapping.
+  VEC4_T s00 = texelFetch(t_in, ivec3(lower.x, lower.y, pos.z), 0);
+  VEC4_T s10 = texelFetch(t_in, ivec3(upper.x, lower.y, pos.z), 0);
+  VEC4_T s01 = texelFetch(t_in, ivec3(lower.x, upper.y, pos.z), 0);
+  VEC4_T s11 = texelFetch(t_in, ivec3(upper.x, upper.y, pos.z), 0);
+
+  // Bilinear interpolation. Weights are scalars; mix() acts on all 4 channels.
+  VEC4_T out_tex =
+      mix(mix(s00, s10, w.x), mix(s01, s11, w.x), w.y);
+
+  imageStore(t_out, pos, out_tex);
+}

backends/vulkan/runtime/graph/ops/glsl/grid_sampler_2d.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.
+
+grid_sampler_2d:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: texture3d
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: half
+      - VALUE: float
+  shader_variants:
+    - NAME: grid_sampler_2d

backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp

@@ -0,0 +1,126 @@
+/*
+ * 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.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+void resize_grid_sampler_2d_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  (void)extra_args;
+  const ValueRef out = args.at(0).refs.at(0);
+  const ValueRef in = args.at(1).refs.at(0);
+  const ValueRef grid = args.at(1).refs.at(1);
+
+  const std::vector<int64_t> in_sizes = graph->sizes_of(in);
+  const std::vector<int64_t> grid_sizes = graph->sizes_of(grid);
+
+  // input  : [N, C, Hin, Win]
+  // grid   : [N, Hout, Wout, 2]
+  // output : [N, C, Hout, Wout]
+  std::vector<int64_t> out_sizes = {
+      in_sizes.at(0), in_sizes.at(1), grid_sizes.at(1), grid_sizes.at(2)};
+
+  graph->virtual_resize(out, out_sizes);
+}
+
+void add_grid_sampler_2d_node(
+    ComputeGraph& graph,
+    const ValueRef in,
+    const ValueRef grid,
+    const ValueRef interpolation_mode,
+    const ValueRef padding_mode,
+    const ValueRef align_corners,
+    const ValueRef out) {
+  // Runtime sanity checks. The Python partitioner is supposed to filter out
+  // unsupported configurations, but guard against bypass paths here too.
+  // mode: 0 = bilinear, 1 = nearest, 2 = bicubic
+  VK_CHECK_COND(
+      graph.extract_scalar<int64_t>(interpolation_mode) == 0,
+      "Vulkan grid_sampler_2d only supports bilinear interpolation");
+  // padding_mode: 0 = zeros, 1 = border, 2 = reflection
+  VK_CHECK_COND(
+      graph.extract_scalar<int64_t>(padding_mode) == 1,
+      "Vulkan grid_sampler_2d only supports border padding");
+  VK_CHECK_COND(
+      graph.get_bool(align_corners),
+      "Vulkan grid_sampler_2d requires align_corners=true");
+
+  // Defense-in-depth layout validation. The partitioner enforces these
+  // layouts via `inputs_storage` in op_registry.py::register_grid_sampler_2d,
+  // but the shader hard-codes channels-packed texture indexing for in/out and
+  // contiguous buffer indexing for grid, so a layout mismatch here would be a
+  // silent miscompute. Per the etvk-implement-operator skill ("Validate
+  // tensor layout assumptions"), assert these explicitly.
+  VK_CHECK_COND(
+      graph.is_standard_channels_packed_texture_tensor(in),
+      "Vulkan grid_sampler_2d requires input to be a channels-packed texture");
+  VK_CHECK_COND(
+      graph.is_standard_channels_packed_texture_tensor(out),
+      "Vulkan grid_sampler_2d requires output to be a channels-packed texture");
+  VK_CHECK_COND(
+      graph.is_contiguous_buffer_tensor(grid),
+      "Vulkan grid_sampler_2d requires grid to be a contiguous buffer");
+
+  // The shader binds t_in, t_out, and t_grid with a single DTYPE selected via
+  // `dtype_of(out)` below. The op registry allows `grid` to be fp16 or fp32
+  // independently of the input dtype, so without this guard a mixed-precision
+  // model (e.g., fp32 flow grid + fp16 activations) would bind the fp32 grid
+  // buffer as half and silently miscompute. Op tests use matching dtypes for
+  // all args, so they would not catch this.
+  VK_CHECK_COND(
+      graph.dtype_of(grid) == graph.dtype_of(out),
+      "Vulkan grid_sampler_2d requires grid and input to share dtype");
+
+  std::string kernel_name("grid_sampler_2d");
+  kernel_name.reserve(kShaderNameReserve);
+  add_dtype_suffix(kernel_name, graph.dtype_of(out));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      default_pick_global_wg_size,
+      default_pick_local_wg_size,
+      // Inputs and Outputs
+      {{out, vkapi::kWrite}, {{in, grid}, vkapi::kRead}},
+      // Shader params buffers. `meta_ubo` packs sizes, limits, axis_map, and
+      // packed_dim into the canonical TextureMetadata struct (see vtensor.md);
+      // the shader derives Wout/Hout/N/num_z_per_n from `outp.sizes` and
+      // `outp.limits`, so no extra params buffer is needed.
+      {graph.meta_ubo(out), graph.meta_ubo(in)},
+      // Push Constants
+      {},
+      // Specialization Constants — pass the output tensor's hashed layout so
+      // the shader can specialize on packed_dim at pipeline creation time.
+      {graph.hashed_layout_of(out)},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_grid_sampler_2d_node));
+}
+
+void grid_sampler_2d(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  // Argument order matches kernels/portable/cpu/op_grid_sampler_2d.cpp:
+  //   (input, grid, interpolation_mode, padding_mode, align_corners, out)
+  return add_grid_sampler_2d_node(
+      graph, args[0], args[1], args[2], args[3], args[4], args[5]);
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(aten.grid_sampler_2d.default, grid_sampler_2d);
+}
+
+} // namespace vkcompute