@@ -40,11 +40,20 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
4040// fp components of a texel share the same (N, Hout, Wout) and differ only in
4141// channel. This lets one bilinear interpolation produce all 4 output channels.
4242${layout_declare_spec_const(C, "int ", "out_layout", "CONTIG_LAYOUT_INT")}
43+ // padding_mode: 0 = zeros, 1 = border. align_corners: 0 = false, 1 = true.
44+ // Both are specialization constants so the four supported configurations
45+ // share one shader variant and each pipeline still compiles to straight-line
46+ // code with the branches folded away.
47+ ${layout_declare_spec_const(C, "int ", "padding_mode", "1 ")}
48+ ${layout_declare_spec_const(C, "int ", "align_corners", "1 ")}
4349
4450/*
45- * Vulkan implementation of `aten.grid_sampler_2d.default` for the
46- * specific configuration used by RIFE's `WarpModule`:
47- * mode=bilinear, padding_mode=border, align_corners=true.
51+ * Vulkan implementation of `aten.grid_sampler_2d.default` for
52+ * mode=bilinear, padding_mode in {zeros, border}, align_corners in {0, 1}.
53+ *
54+ * RIFE's `WarpModule` uses (border, align_corners=true); the deformable
55+ * attention in RF-DETR and other DETR derivatives uses (zeros,
56+ * align_corners=false), which is the torch default.
4857 *
4958 * Layout assumptions (validated in add_grid_sampler_2d_node):
5059 * - input : channels-packed texture3d, shape [N, C, Hin, Win]
@@ -60,6 +69,23 @@ ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")}
6069 * outp.sizes.x = Wout, outp.sizes.y = Hout, outp.sizes.w = N.
6170 * outp.limits.z = N * ceil(C/4) (texel slices along z).
6271 */
72+ /*
73+ * Read one input texel at integer pixel (c.x, c.y) of channel slice `z`.
74+ *
75+ * padding_mode=zeros returns 0 for a corner outside the input, which is what
76+ * makes the bilinear weights of a partially out-of-range sample sum to less
77+ * than 1 -- exactly the aten semantics. padding_mode=border clamps instead,
78+ * which is a no-op when the caller has already clamped the coordinate.
79+ */
80+ VEC4_T sample_or_zero(const ivec2 c, const ivec2 max_in_xy, const int z) {
81+ if (padding_mode == 0 &&
82+ (c.x < 0 || c.y < 0 || c.x > max_in_xy.x || c.y > max_in_xy.y)) {
83+ return VEC4_T(0 );
84+ }
85+ const ivec2 cc = clamp (c, ivec2 (0 ), max_in_xy);
86+ return texelFetch(t_in, ivec3 (cc.x, cc.y, z), 0 );
87+ }
88+
6389void main() {
6490 const ivec3 pos = ivec3 (gl_GlobalInvocationID);
6591
@@ -86,29 +112,43 @@ void main() {
86112 const float gx_norm = float (t_grid[grid_base + 0 ]);
87113 const float gy_norm = float (t_grid[grid_base + 1 ]);
88114
89- // Unnormalize for align_corners=true:
90- // coord_pixel = (coord_norm + 1) * 0.5 * (size - 1)
91- // Input W/H come from inp.sizes (WHCN), not inp.limits (texel space).
92- const ivec2 max_in_xy = ivec2 (inp.sizes.xy) - 1 ;
93- const float gx_pixel = (gx_norm + 1.0 ) * 0.5 * float (max_in_xy.x);
94- const float gy_pixel = (gy_norm + 1.0 ) * 0.5 * float (max_in_xy.y);
115+ // Unnormalize. Input W/H come from inp.sizes (WHCN), not inp.limits
116+ // (texel space).
117+ // align_corners=true : coord = (g + 1) * 0.5 * (size - 1)
118+ // align_corners=false: coord = ((g + 1) * size - 1) * 0.5
119+ // The second form places the normalized range over pixel *edges* rather
120+ // than pixel centers, so it can legitimately land outside [0, size-1].
121+ const ivec2 in_size = ivec2 (inp.sizes.xy);
122+ const ivec2 max_in_xy = in_size - 1 ;
123+ vec2 g_pixel;
124+ if (align_corners == 1 ) {
125+ g_pixel = (vec2 (gx_norm, gy_norm) + 1.0 ) * 0.5 * vec2 (max_in_xy);
126+ } else {
127+ g_pixel = ((vec2 (gx_norm, gy_norm) + 1.0 ) * vec2 (in_size) - 1.0 ) * 0.5 ;
128+ }
95129
96- // padding_mode=border: clamp coordinates to [0, size-1].
97- const float gx = clamp (gx_pixel, 0.0 , float (max_in_xy.x));
98- const float gy = clamp (gy_pixel, 0.0 , float (max_in_xy.y));
130+ // padding_mode=border clamps the sample coordinate itself, which also pins
131+ // the interpolation weights at the edge. padding_mode=zeros must NOT clamp:
132+ // the weights stay as computed and each out-of-range corner contributes a
133+ // zero value instead, so clamping here would change the result.
134+ if (padding_mode == 1 ) {
135+ g_pixel = clamp (g_pixel, vec2 (0.0 ), vec2 (max_in_xy));
136+ }
99137
100- const ivec2 lower = ivec2 (floor (vec2 (gx, gy)));
101- // Clamp ceil to valid range for samples on the border.
102- const ivec2 upper = clamp (lower + ivec2 (1 ), ivec2 (0 ), max_in_xy);
103- const vec2 w = vec2 (gx, gy) - vec2 (lower);
138+ const ivec2 lower = ivec2 (floor (g_pixel));
139+ const ivec2 upper = lower + ivec2 (1 );
140+ const vec2 w = g_pixel - vec2 (lower);
104141
105142 // Fetch the four nearest texels (each carries 4 channels). Because input
106143 // is channels-packed, pos.z indexes the same channel slice in input as in
107144 // output, so we can reuse pos.z directly without remapping.
108- VEC4_T s00 = texelFetch(t_in, ivec3 (lower.x, lower.y, pos.z), 0 );
109- VEC4_T s10 = texelFetch(t_in, ivec3 (upper.x, lower.y, pos.z), 0 );
110- VEC4_T s01 = texelFetch(t_in, ivec3 (lower.x, upper.y, pos.z), 0 );
111- VEC4_T s11 = texelFetch(t_in, ivec3 (upper.x, upper.y, pos.z), 0 );
145+ //
146+ // For border the coordinate is already clamped, so clamping the corner
147+ // index is exact. For zeros an out-of-range corner reads as 0.
148+ VEC4_T s00 = sample_or_zero(ivec2 (lower.x, lower.y), max_in_xy, pos.z);
149+ VEC4_T s10 = sample_or_zero(ivec2 (upper.x, lower.y), max_in_xy, pos.z);
150+ VEC4_T s01 = sample_or_zero(ivec2 (lower.x, upper.y), max_in_xy, pos.z);
151+ VEC4_T s11 = sample_or_zero(ivec2 (upper.x, upper.y), max_in_xy, pos.z);
112152
113153 // Bilinear interpolation. Weights are scalars; mix() acts on all 4 channels.
114154 VEC4_T out_tex =
0 commit comments