Skip to content

Commit 20a9d8c

Browse files
msluszniakSS-JIA
authored andcommitted
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 3433a4f)
1 parent 86896d4 commit 20a9d8c

4 files changed

Lines changed: 125 additions & 34 deletions

File tree

‎backends/vulkan/op_registry.py‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1694,9 +1694,12 @@ def register_grid_priors():
16941694

16951695
@update_features(exir_ops.edge.aten.grid_sampler_2d.default)
16961696
def register_grid_sampler_2d():
1697-
# The Vulkan implementation only supports the configuration used by RIFE's
1698-
# WarpModule: bilinear interpolation (0), border padding (1),
1699-
# align_corners=True. The C++ side has VK_CHECK_COND asserts for these,
1697+
# The Vulkan implementation supports bilinear interpolation (0) with
1698+
# either zeros (0) or border (1) padding and either align_corners. That
1699+
# covers RIFE's WarpModule (border, align_corners=True) and the deformable
1700+
# attention in DETR derivatives (zeros, align_corners=False). Reflection
1701+
# padding and nearest/bicubic interpolation are not implemented.
1702+
# The C++ side has VK_CHECK_COND asserts for these,
17001703
# but those abort the whole inference at graph build — for any other model
17011704
# that contains a differently-configured grid_sampler_2d we want graceful
17021705
# CPU fallback, so we gate delegation here.
@@ -1736,8 +1739,10 @@ def check_grid_sampler_2d_node(node: torch.fx.Node) -> bool:
17361739
if interp is None or padding is None or align_corners is None:
17371740
return False
17381741

1739-
# mode: 0 = bilinear; padding: 1 = border; align_corners must be True.
1740-
return interp == 0 and padding == 1 and bool(align_corners) is True
1742+
# mode: 0 = bilinear. padding: 0 = zeros, 1 = border (2 = reflection
1743+
# needs a coordinate fold the shader does not implement).
1744+
# align_corners is free: both settings are specialization constants.
1745+
return interp == 0 and padding in (0, 1)
17411746

17421747
return OpFeatures(
17431748
inputs_storage=[

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

Lines changed: 60 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -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+
6389
void 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 =

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

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -51,13 +51,13 @@ void add_grid_sampler_2d_node(
5151
VK_CHECK_COND(
5252
graph.extract_scalar<int64_t>(interpolation_mode) == 0,
5353
"Vulkan grid_sampler_2d only supports bilinear interpolation");
54-
// padding_mode: 0 = zeros, 1 = border, 2 = reflection
54+
// padding_mode: 0 = zeros, 1 = border, 2 = reflection. Reflection would need
55+
// a different coordinate fold and is not implemented.
56+
const int64_t padding_mode_val = graph.extract_scalar<int64_t>(padding_mode);
5557
VK_CHECK_COND(
56-
graph.extract_scalar<int64_t>(padding_mode) == 1,
57-
"Vulkan grid_sampler_2d only supports border padding");
58-
VK_CHECK_COND(
59-
graph.get_bool(align_corners),
60-
"Vulkan grid_sampler_2d requires align_corners=true");
58+
padding_mode_val == 0 || padding_mode_val == 1,
59+
"Vulkan grid_sampler_2d only supports zeros and border padding");
60+
const int32_t align_corners_val = graph.get_bool(align_corners) ? 1 : 0;
6161

6262
// Defense-in-depth layout validation. The partitioner enforces these
6363
// layouts via `inputs_storage` in op_registry.py::register_grid_sampler_2d,
@@ -103,9 +103,13 @@ void add_grid_sampler_2d_node(
103103
{graph.meta_ubo(out), graph.meta_ubo(in)},
104104
// Push Constants
105105
{},
106-
// Specialization Constants — pass the output tensor's hashed layout so
107-
// the shader can specialize on packed_dim at pipeline creation time.
108-
{graph.hashed_layout_of(out)},
106+
// Specialization Constants — the output tensor's hashed layout lets the
107+
// shader specialize on packed_dim at pipeline creation time, and the
108+
// padding/align_corners pair folds the configuration branches away so
109+
// the four supported combinations share one shader variant.
110+
{graph.hashed_layout_of(out),
111+
static_cast<int32_t>(padding_mode_val),
112+
align_corners_val},
109113
// Resize Args
110114
{},
111115
// Resizing Logic

‎backends/vulkan/test/test_vulkan_delegate.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1280,6 +1280,48 @@ def forward(self, x):
12801280
sample_inputs,
12811281
)
12821282

1283+
def test_vulkan_backend_grid_sampler_2d(self):
1284+
class GridSampler2d(torch.nn.Module):
1285+
def __init__(self, padding_mode, align_corners):
1286+
super().__init__()
1287+
self.padding_mode = padding_mode
1288+
self.align_corners = align_corners
1289+
1290+
def forward(self, x, grid):
1291+
return torch.nn.functional.grid_sample(
1292+
x,
1293+
grid,
1294+
mode="bilinear",
1295+
padding_mode=self.padding_mode,
1296+
align_corners=self.align_corners,
1297+
)
1298+
1299+
# Deliberately push the grid past [-1, 1] on every side so the zeros
1300+
# and border paths actually diverge; an in-range grid is identical
1301+
# under both and would pass even with the padding branch broken.
1302+
grid = torch.stack(
1303+
torch.meshgrid(
1304+
torch.linspace(-1.6, 1.6, 7),
1305+
torch.linspace(-1.6, 1.6, 5),
1306+
indexing="ij",
1307+
)[::-1],
1308+
dim=-1,
1309+
).unsqueeze(0)
1310+
sample_inputs = (
1311+
torch.rand(size=(1, 4, 6, 8), dtype=torch.float32),
1312+
grid.contiguous(),
1313+
)
1314+
1315+
for padding_mode in ("zeros", "border"):
1316+
for align_corners in (True, False):
1317+
with self.subTest(
1318+
padding_mode=padding_mode, align_corners=align_corners
1319+
):
1320+
self.lower_module_and_test_output(
1321+
GridSampler2d(padding_mode, align_corners),
1322+
sample_inputs,
1323+
)
1324+
12831325
def test_vulkan_backend_minimum(self):
12841326
class MinimumModule(torch.nn.Module):
12851327
def __init__(self):

0 commit comments

Comments
 (0)