pytorch · msluszniak · Sep 14, 2026
@@ -1566,9 +1566,12 @@ def register_grid_priors():
 
 @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,
+    # The Vulkan implementation supports bilinear interpolation (0) with
+    # either zeros (0) or border (1) padding and either align_corners. That
+    # covers RIFE's WarpModule (border, align_corners=True) and the deformable
+    # attention in DETR derivatives (zeros, align_corners=False). Reflection
+    # padding and nearest/bicubic interpolation are not implemented.
+    # 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.
@@ -1608,8 +1611,10 @@ def check_grid_sampler_2d_node(node: torch.fx.Node) -> bool:
         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
+        # mode: 0 = bilinear. padding: 0 = zeros, 1 = border (2 = reflection
+        # needs a coordinate fold the shader does not implement).
+        # align_corners is free: both settings are specialization constants.
+        return interp == 0 and padding in (0, 1)
 
     return OpFeatures(
         inputs_storage=[

@@ -40,11 +40,20 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 // 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")}
+// padding_mode: 0 = zeros, 1 = border. align_corners: 0 = false, 1 = true.
+// Both are specialization constants so the four supported configurations
+// share one shader variant and each pipeline still compiles to straight-line
+// code with the branches folded away.
+${layout_declare_spec_const(C, "int", "padding_mode", "1")}
+${layout_declare_spec_const(C, "int", "align_corners", "1")}
 
 /*
- * 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.
+ * Vulkan implementation of `aten.grid_sampler_2d.default` for
+ *   mode=bilinear, padding_mode in {zeros, border}, align_corners in {0, 1}.
+ *
+ * RIFE's `WarpModule` uses (border, align_corners=true); the deformable
+ * attention in RF-DETR and other DETR derivatives uses (zeros,
+ * align_corners=false), which is the torch default.
  *
  * Layout assumptions (validated in add_grid_sampler_2d_node):
  *   - input  : channels-packed texture3d, shape [N, C, Hin, Win]
@@ -60,6 +69,23 @@ ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
  *   outp.sizes.x = Wout, outp.sizes.y = Hout, outp.sizes.w = N.
  *   outp.limits.z = N * ceil(C/4) (texel slices along z).
  */
+/*
+ * Read one input texel at integer pixel (c.x, c.y) of channel slice `z`.
+ *
+ * padding_mode=zeros returns 0 for a corner outside the input, which is what
+ * makes the bilinear weights of a partially out-of-range sample sum to less
+ * than 1 -- exactly the aten semantics. padding_mode=border clamps instead,
+ * which is a no-op when the caller has already clamped the coordinate.
+ */
+VEC4_T sample_or_zero(const ivec2 c, const ivec2 max_in_xy, const int z) {
+  if (padding_mode == 0 &&
+      (c.x < 0 || c.y < 0 || c.x > max_in_xy.x || c.y > max_in_xy.y)) {
+    return VEC4_T(0);
+  }
+  const ivec2 cc = clamp(c, ivec2(0), max_in_xy);
+  return texelFetch(t_in, ivec3(cc.x, cc.y, z), 0);
+}
+
 void main() {
   const ivec3 pos = ivec3(gl_GlobalInvocationID);
 
@@ -86,29 +112,43 @@ void main() {
   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);
+  // Unnormalize. Input W/H come from inp.sizes (WHCN), not inp.limits
+  // (texel space).
+  //   align_corners=true : coord = (g + 1) * 0.5 * (size - 1)
+  //   align_corners=false: coord = ((g + 1) * size - 1) * 0.5
+  // The second form places the normalized range over pixel *edges* rather
+  // than pixel centers, so it can legitimately land outside [0, size-1].
+  const ivec2 in_size = ivec2(inp.sizes.xy);
+  const ivec2 max_in_xy = in_size - 1;
+  vec2 g_pixel;
+  if (align_corners == 1) {
+    g_pixel = (vec2(gx_norm, gy_norm) + 1.0) * 0.5 * vec2(max_in_xy);
+  } else {
+    g_pixel = ((vec2(gx_norm, gy_norm) + 1.0) * vec2(in_size) - 1.0) * 0.5;
+  }
 
-  // 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));
+  // padding_mode=border clamps the sample coordinate itself, which also pins
+  // the interpolation weights at the edge. padding_mode=zeros must NOT clamp:
+  // the weights stay as computed and each out-of-range corner contributes a
+  // zero value instead, so clamping here would change the result.
+  if (padding_mode == 1) {
+    g_pixel = clamp(g_pixel, vec2(0.0), vec2(max_in_xy));
+  }
 
-  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);
+  const ivec2 lower = ivec2(floor(g_pixel));
+  const ivec2 upper = lower + ivec2(1);
+  const vec2 w = g_pixel - 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);
+  //
+  // For border the coordinate is already clamped, so clamping the corner
+  // index is exact. For zeros an out-of-range corner reads as 0.
+  VEC4_T s00 = sample_or_zero(ivec2(lower.x, lower.y), max_in_xy, pos.z);
+  VEC4_T s10 = sample_or_zero(ivec2(upper.x, lower.y), max_in_xy, pos.z);
+  VEC4_T s01 = sample_or_zero(ivec2(lower.x, upper.y), max_in_xy, pos.z);
+  VEC4_T s11 = sample_or_zero(ivec2(upper.x, upper.y), max_in_xy, pos.z);
 
   // Bilinear interpolation. Weights are scalars; mix() acts on all 4 channels.
   VEC4_T out_tex =

@@ -51,13 +51,13 @@ void add_grid_sampler_2d_node(
   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
+  // padding_mode: 0 = zeros, 1 = border, 2 = reflection. Reflection would need
+  // a different coordinate fold and is not implemented.
+  const int64_t padding_mode_val = graph.extract_scalar<int64_t>(padding_mode);
   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");
+      padding_mode_val == 0 || padding_mode_val == 1,
+      "Vulkan grid_sampler_2d only supports zeros and border padding");
+  const int32_t align_corners_val = graph.get_bool(align_corners) ? 1 : 0;
 
   // Defense-in-depth layout validation. The partitioner enforces these
   // layouts via `inputs_storage` in op_registry.py::register_grid_sampler_2d,
@@ -103,9 +103,13 @@ void add_grid_sampler_2d_node(
       {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)},
+      // Specialization Constants — the output tensor's hashed layout lets the
+      // shader specialize on packed_dim at pipeline creation time, and the
+      // padding/align_corners pair folds the configuration branches away so
+      // the four supported combinations share one shader variant.
+      {graph.hashed_layout_of(out),
+       static_cast<int32_t>(padding_mode_val),
+       align_corners_val},
       // Resize Args
       {},
       // Resizing Logic

@@ -1192,6 +1192,48 @@ def forward(self, x):
             sample_inputs,
         )
 
+    def test_vulkan_backend_grid_sampler_2d(self):
+        class GridSampler2d(torch.nn.Module):
+            def __init__(self, padding_mode, align_corners):
+                super().__init__()
+                self.padding_mode = padding_mode
+                self.align_corners = align_corners
+
+            def forward(self, x, grid):
+                return torch.nn.functional.grid_sample(
+                    x,
+                    grid,
+                    mode="bilinear",
+                    padding_mode=self.padding_mode,
+                    align_corners=self.align_corners,
+                )
+
+        # Deliberately push the grid past [-1, 1] on every side so the zeros
+        # and border paths actually diverge; an in-range grid is identical
+        # under both and would pass even with the padding branch broken.
+        grid = torch.stack(
+            torch.meshgrid(
+                torch.linspace(-1.6, 1.6, 7),
+                torch.linspace(-1.6, 1.6, 5),
+                indexing="ij",
+            )[::-1],
+            dim=-1,
+        ).unsqueeze(0)
+        sample_inputs = (
+            torch.rand(size=(1, 4, 6, 8), dtype=torch.float32),
+            grid.contiguous(),
+        )
+
+        for padding_mode in ("zeros", "border"):
+            for align_corners in (True, False):
+                with self.subTest(
+                    padding_mode=padding_mode, align_corners=align_corners
+                ):
+                    self.lower_module_and_test_output(
+                        GridSampler2d(padding_mode, align_corners),
+                        sample_inputs,
+                    )
+
     def test_vulkan_backend_minimum(self):
         class MinimumModule(torch.nn.Module):
             def __init__(self):