From ce549738df0c1dceeb3615bca2c15740829f696d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Mon, 14 Sep 2026 18:07:47 +0200 Subject: [PATCH] Vulkan: support zeros padding and align_corners=false in grid_sampler_2d The shader implemented one configuration, the one RIFE's WarpModule uses: bilinear, border padding, align_corners=true. Every other configuration was rejected by the partitioner and fell back to CPU, including torch's own default of zeros padding with align_corners=false, which is what the deformable attention in DETR-derived detectors emits. Make padding_mode and align_corners specialization constants so the four supported combinations share one shader variant and each pipeline still compiles with the branches folded away. border clamps the sample coordinate, which pins the interpolation weights at the edge. zeros must not clamp: the weights stay as computed and an out-of-range corner contributes a zero value instead, so the two paths differ for any grid that leaves [-1, 1]. On an RF-DETR keypoint model this takes all 8 grid_sampler_2d calls off the host and the graph from 7 delegate blobs to 2. (cherry picked from commit 3433a4f1779e04930043398727d21d176ac387b0) --- backends/vulkan/op_registry.py | 15 ++-- .../graph/ops/glsl/grid_sampler_2d.glsl | 80 ++++++++++++++----- .../runtime/graph/ops/impl/GridSampler2d.cpp | 22 ++--- backends/vulkan/test/test_vulkan_delegate.py | 42 ++++++++++ 4 files changed, 125 insertions(+), 34 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index add4d01a78e..9ce343f0c0c 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1627,9 +1627,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. @@ -1669,8 +1672,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=[ diff --git a/backends/vulkan/runtime/graph/ops/glsl/grid_sampler_2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/grid_sampler_2d.glsl index b697d66dfaf..687461164fe 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/grid_sampler_2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/grid_sampler_2d.glsl @@ -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 = diff --git a/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp b/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp index 0f5515233dc..d0c6fbf0dcd 100644 --- a/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/GridSampler2d.cpp @@ -51,13 +51,13 @@ void add_grid_sampler_2d_node( VK_CHECK_COND( graph.extract_scalar(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(padding_mode); VK_CHECK_COND( - graph.extract_scalar(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(padding_mode_val), + align_corners_val}, // Resize Args {}, // Resizing Logic diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 3e4b1c19433..3841eb61a04 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -1266,6 +1266,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):