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):