Skip to content

Commit 04f9002

Browse files
msluszniakSS-JIA
authored andcommitted
[ET-VK] Fix squeeze_copy of the outermost dim under dynamic shapes
add_squeeze_copy_dims_node() skips dim 0 and falls back to add_clone_node(). resize_clone_node() only propagates sizes when input and output have the same dim count, which a squeeze never does, so the output keeps the extents it was built with. With static shapes that is invisible. With dynamic shapes the output holds its upper-bound extents while consumers read it at the real size, so the copy lands in the wrong places and roughly half the output comes back zeroed -- silently, with no error. Route dim 0 through the permute path like every other squeeze dim; resize_permute_node() already has an explicit branch for the rank-reducing case. Repro: any model that ends up with torch.cat(list(x), -1) over a rank-4 tensor with a dynamic dim. The unbind lowers to slice_copy plus squeeze_copy.dims, and the second slice comes back zeroed for every extent below the bound. Reduced to a 15-line case: y[1:2] is correct while y[1:2].squeeze(0) returns exactly half zeros (cosine 0.704 = sqrt of 0.5 against the reference), correct only at the bound. Verified on a Galaxy S26 Ultra (Adreno 840): the reduced case goes from 0.704 to 1.000000 at extents 200, 500 and 1000, and a TTS model whose classifier-free-guidance batch is built this way goes from cosine 0.36 to 0.99993 against its CPU reference.
1 parent eadb4c3 commit 04f9002

1 file changed

Lines changed: 12 additions & 6 deletions

File tree

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,21 +23,27 @@ void add_squeeze_copy_dims_node(
2323
const ValueRef out) {
2424
const int64_t in_dim = graph.dim_of(in);
2525
const std::vector<int64_t> in_sizes = graph.sizes_of(in);
26-
const std::vector<int64_t> out_sizes = graph.sizes_of(in);
2726

2827
const std::vector<int64_t> dims = graph.extract_int_or_symint_list(dims_ref);
2928
std::vector<int64_t> squeeze_dims;
30-
// Filter out edge cases that we don't need squeeze:
31-
// 1. The size of squeeze dim is larger than 1.
32-
// 2. Squeeze outter most dim
33-
// For these cases, just pass input to output via clone.
29+
// Filter out the edge case that we don't need to squeeze: the size of the
30+
// squeeze dim is larger than 1. For that case, just pass input to output via
31+
// clone.
32+
//
33+
// Note that the outermost dim must NOT be excluded here. Routing it to
34+
// add_clone_node() leaves the output unresized at runtime, because
35+
// resize_clone_node() only propagates sizes when input and output have the
36+
// same dim count -- which is never true for a squeeze. Under dynamic shapes
37+
// the output then keeps its upper-bound extents while consumers read it at
38+
// the real size, silently producing wrong values. add_permute_node()'s
39+
// resize function handles the rank-reducing case explicitly.
3440
for (int i = 0; i < dims.size(); ++i) {
3541
// adjust negative dims
3642
int64_t dim_val = dims.at(i);
3743
if (dim_val < 0) {
3844
dim_val += in_dim;
3945
}
40-
if (dims.at(i) != 0 && in_sizes.at(dim_val) == 1) {
46+
if (in_sizes.at(dim_val) == 1) {
4147
squeeze_dims.push_back(dim_val);
4248
}
4349
}

0 commit comments

Comments
 (0)