From ec1370648c43f565fd81c2ed7a9ab6df88b83be7 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Fri, 4 Sep 2026 10:36:16 -0700 Subject: [PATCH] select: accept a constant SymInt dim, reject driven ones Summary: `read_scalar` in the WebGPU `select` op treated every `ValueType::SymInt` as dynamic and threw `select: dynamic/unsupported dim`. That is stricter than it needs to be: the Vulkan serializer sometimes parks a plain integer constant in a SymInt slot, much as it can reuse an earlier `Double` value for an equal integer scalar. Such a graph is fully static, but failed to build. A SymInt is now rejected only when some op actually drives it at execute time. The check is deny-by-default -- the value counts as a build-time constant only if no producer registry claims it -- and consults all three SymInt producers: `et_vk.select_as_symint` (`symint_sources`), `sym_size.int` (`symint_dim_sources`) and SymInt arithmetic (`symint_computed`). Rejection has to stay loud for those, because `select` drops an axis: a dim that varies at runtime changes the output rank, which the graph cannot express, so silently accepting it would return a wrong-shaped result instead of an error. The arithmetic case needed a new registry. `register_sym_binary` keys its resize hook on the trigger operand rather than on the result, so neither existing registry can answer "does some op write this id?" for a value produced by `add`/`sub`/`mul`/`floordiv`. `WebGPUGraph::add_symint_computed` records the result instead, and `register_sym_binary` calls it whenever it registers a hook. That call is gated on operand liveness so `add(Int, Int) -> SymInt`, which registers no hook and genuinely cannot change, is not falsely rejected. Any future SymInt producer must be consulted in `is_runtime_driven_symint`; the comment there says so. Missing one puts `select` back to reading a build-time seed for a value that varies at execute. Authored with Claude Code. Differential Revision: D118248620 --- backends/webgpu/runtime/WebGPUGraph.h | 15 ++ backends/webgpu/runtime/ops/select/Select.cpp | 32 +++- .../ops/select_as_symint/SelectAsSymint.cpp | 15 +- backends/webgpu/test/test_webgpu_native.cpp | 179 ++++++++++++++++-- 4 files changed, 214 insertions(+), 27 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 23ce9df03ed..0dd9cba4640 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -251,6 +251,20 @@ class WebGPUGraph { void add_symint_dim_source(int symint_id, int tensor_id, int dim) { symint_dim_sources_.push_back({symint_id, tensor_id, dim}); } + const std::vector& symint_dim_sources() const { + return symint_dim_sources_; + } + + // Records a SymInt that a resize hook RECOMPUTES at execute (SymInt + // arithmetic). The SymIntSource registries above name where a SymInt is read + // from; resize_hooks_ is keyed by the hook's trigger operand, not its result, + // so neither can answer "does some op write this id?". + void add_symint_computed(int symint_id) { + symint_computed_.insert(symint_id); + } + const std::unordered_set& symint_computed() const { + return symint_computed_; + } bool tensor_has_dynamic_dims(int tensor_id) const { return dynamic_tensor_ids_.count(tensor_id) != 0; @@ -597,6 +611,7 @@ class WebGPUGraph { std::unordered_map symints_; std::vector symint_sources_; std::vector symint_dim_sources_; + std::unordered_set symint_computed_; std::unordered_set dynamic_tensor_ids_; // Resize hooks + the set of SymInts changed since the last propagate_resize. diff --git a/backends/webgpu/runtime/ops/select/Select.cpp b/backends/webgpu/runtime/ops/select/Select.cpp index cb4575bc004..832fd3e1b79 100644 --- a/backends/webgpu/runtime/ops/select/Select.cpp +++ b/backends/webgpu/runtime/ops/select/Select.cpp @@ -30,12 +30,40 @@ struct SelectParams { uint32_t _pad[2]; }; -// dim/index are static integer scalars (SymInt throws); Vulkan serialization -// can reuse an earlier Double value for an equal integer scalar. +// Deny by default: a SymInt holds a build-time constant only if NO producer +// registry claims it. The three SymInt producers are select_as_symint +// (symint_sources), sym_size.int (symint_dim_sources), and SymInt arithmetic +// (symint_computed). A NEW SYMINT PRODUCER MUST BE CONSULTED HERE -- if one is +// missed, select silently reads a build-time seed for a value that varies at +// execute and emits a wrong-rank result instead of failing loudly. +bool is_runtime_driven_symint(const WebGPUGraph& graph, int id) { + for (const auto& src : graph.symint_sources()) { + if (src.symint_id == id) { + return true; + } + } + for (const auto& src : graph.symint_dim_sources()) { + if (src.symint_id == id) { + return true; + } + } + return graph.symint_computed().count(id) != 0; +} + +// dim/index are static integer scalars; Vulkan serialization can reuse an +// earlier Double value for an equal integer scalar, or park one in a SymInt +// slot. Only a SymInt an op actually drives at runtime throws: select drops an +// axis, so a varying dim/index would change the output rank. int64_t read_scalar(WebGPUGraph& graph, int id, const char* what) { switch (graph.get_value_type(id)) { case WebGPUGraph::ValueType::Int: return graph.get_int(id); + case WebGPUGraph::ValueType::SymInt: + if (is_runtime_driven_symint(graph, id)) { + throw std::runtime_error( + std::string("select: dynamic/unsupported ") + what); + } + return graph.read_symint(id); case WebGPUGraph::ValueType::Double: { const double d = graph.get_double(id); constexpr double kInt64Limit = 0x1p63; diff --git a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp index be333c99d53..085f0f471cd 100644 --- a/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp +++ b/backends/webgpu/runtime/ops/select_as_symint/SelectAsSymint.cpp @@ -88,14 +88,21 @@ void register_sym_binary( g.set_symint(out, op(read_scalar(g, a), read_scalar(g, b))); }; recompute(graph); // seed the build-time value - if (graph.get_value_type(a) == WebGPUGraph::ValueType::SymInt) { - graph.add_resize_hook(a, recompute); - } + const bool a_live = + graph.get_value_type(a) == WebGPUGraph::ValueType::SymInt; // b != a: for a self-op (e.g. x + x) both operands share one id; register the // recompute hook once, not twice. - if (b != a && graph.get_value_type(b) == WebGPUGraph::ValueType::SymInt) { + const bool b_live = + b != a && graph.get_value_type(b) == WebGPUGraph::ValueType::SymInt; + if (a_live) { + graph.add_resize_hook(a, recompute); + } + if (b_live) { graph.add_resize_hook(b, recompute); } + if (a_live || b_live) { + graph.add_symint_computed(out); + } } void sym_add_impl(WebGPUGraph& graph, const std::vector& args) { diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index 6448568a66e..0ce29bc949e 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -2645,19 +2645,29 @@ static bool test_slice_double_start() { return ok; } +// How the select `dim` argument is serialized. The two dynamic kinds prepend +// the ops that make the SymInt runtime-driven: SymIntFromDim a sym_size.int, +// SymIntFromArithmetic a sym_size.int plus an `add` whose result feeds select. +enum class SelectDimKind { + Double, + Int, + SymIntConstant, + SymIntFromDim, + SymIntFromArithmetic, +}; + // Regression for serialized integer select arguments that alias an earlier // floating-point scalar in the Vulkan scalar cache. A production graph has -// select calls whose dim and index both reference Double 0.0. +// select calls whose dim and index both reference Double 0.0, and others whose +// dim lands in a SymInt slot holding a constant. static void finish_select_scalar_graph( ::flatbuffers::FlatBufferBuilder& fbb, double dim, double index, - uint32_t out_len, - bool symint_dim = false) { + const std::vector& out_dims, + SelectDimKind dim_kind = SelectDimKind::Double, + const std::vector& in_dims = {2u, 3u}) { namespace vk = vkgraph; - std::vector in_dims = {2u, 3u}; - std::vector out_dims = {out_len}; - std::vector<::flatbuffers::Offset> values; values.push_back(vk::CreateVkValue( fbb, @@ -2669,16 +2679,27 @@ static void finish_select_scalar_graph( /*constant_id=*/-1, /*mem_obj_id=*/0) .Union())); - if (symint_dim) { - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::SymInt, - vk::CreateSymInt(fbb, /*value=*/0).Union())); - } else { - values.push_back(vk::CreateVkValue( - fbb, - vk::GraphTypes::Double, - vk::CreateDouble(fbb, /*double_val=*/dim).Union())); + switch (dim_kind) { + case SelectDimKind::Double: + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/dim).Union())); + break; + case SelectDimKind::Int: + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Int, + vk::CreateInt(fbb, static_cast(dim)).Union())); + break; + case SelectDimKind::SymIntConstant: + case SelectDimKind::SymIntFromDim: + case SelectDimKind::SymIntFromArithmetic: + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::SymInt, + vk::CreateSymInt(fbb, static_cast(dim)).Union())); + break; } values.push_back(vk::CreateVkValue( fbb, @@ -2695,8 +2716,31 @@ static void finish_select_scalar_graph( /*mem_obj_id=*/1) .Union())); - std::vector args = {0, 1, 2, 3}; std::vector<::flatbuffers::Offset> chain; + const bool via_arithmetic = dim_kind == SelectDimKind::SymIntFromArithmetic; + // sym_size.int(in, dim=0) lands directly in the select dim (value 1), or in + // an intermediate SymInt (value 5) that `add` combines into value 1. + const std::vector sym_size_args = {0, 4, via_arithmetic ? 5 : 1}; + const std::vector sym_add_args = {5, 6, 1}; + if (dim_kind == SelectDimKind::SymIntFromDim || via_arithmetic) { + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + if (via_arithmetic) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::SymInt, + vk::CreateSymInt(fbb, /*value=*/0).Union())); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 1).Union())); + } + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "sym_size.int", &sym_size_args)); + if (via_arithmetic) { + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "add", &sym_add_args)); + } + } + std::vector args = {0, 1, 2, 3}; chain.push_back( vk::CreateOperatorCallDirect(fbb, 0, "aten.select_copy.int", &args)); std::vector input_ids = {0}; @@ -2714,7 +2758,7 @@ static bool test_select_double_scalar_case( "\n--- Test: select Double scalars (dim=%g, index=%g) ---\n", dim, index); ::flatbuffers::FlatBufferBuilder fbb; finish_select_scalar_graph( - fbb, dim, index, static_cast(expected.size())); + fbb, dim, index, {static_cast(expected.size())}); WebGPUGraph graph; try { @@ -2752,9 +2796,9 @@ static bool test_select_scalar_build_error( double dim, double index, const char* expected_error, - bool symint_dim = false) { + SelectDimKind dim_kind = SelectDimKind::Double) { ::flatbuffers::FlatBufferBuilder fbb; - finish_select_scalar_graph(fbb, dim, index, /*out_len=*/3u, symint_dim); + finish_select_scalar_graph(fbb, dim, index, /*out_dims=*/{3u}, dim_kind); WebGPUGraph graph; try { @@ -2775,6 +2819,85 @@ static bool test_select_scalar_build_error( return false; } +// A SymInt slot holding a serializer-parked constant is not dynamic: nothing +// drives it at execute, so select must accept it like a static Int. +static bool test_select_constant_symint_dim_builds() { + printf("\n--- Test: select constant SymInt dim builds ---\n"); + ::flatbuffers::FlatBufferBuilder fbb; + finish_select_scalar_graph( + fbb, + /*dim=*/0.0, + /*index=*/0.0, + /*out_dims=*/{3u}, + SelectDimKind::SymIntConstant); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + } catch (const std::exception& e) { + printf("FAIL: constant SymInt dim rejected: %s\n", e.what()); + return false; + } + printf("PASS: select constant SymInt dim builds\n"); + return true; +} + +// Runs select(dim=0, index=0) and writes the graph output into `out`. +static bool run_select_dim0( + SelectDimKind dim_kind, + const std::vector& in_dims, + const std::vector& out_dims, + const std::vector& in, + std::vector& out) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_select_scalar_graph( + fbb, /*dim=*/0.0, /*index=*/0.0, out_dims, dim_kind, in_dims); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + std::vector inputs = { + {in.data(), in.size() * sizeof(float), false}}; + std::vector outputs = { + {out.data(), out.size() * sizeof(float), true}}; + graph.copy_inputs(inputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); + } catch (const std::exception& e) { + printf("FAIL: select run threw: %s\n", e.what()); + return false; + } + return true; +} + +static bool test_select_symint_dim_matches_int_dim() { + printf("\n--- Test: select SymInt dim matches static Int dim ---\n"); + const std::vector in_dims = {1u, 2u, 3u}; + const std::vector out_dims = {2u, 3u}; + const std::vector in = {0.5f, 1.5f, 2.5f, 3.5f, 4.5f, 5.5f}; + + std::vector int_out(in.size(), -1.0f); + std::vector symint_out(in.size(), -2.0f); + if (!run_select_dim0(SelectDimKind::Int, in_dims, out_dims, in, int_out) || + !run_select_dim0( + SelectDimKind::SymIntConstant, in_dims, out_dims, in, symint_out)) { + return false; + } + + // Dropping a leading extent-1 axis is the identity on the buffer. + if (int_out != in) { + printf("FAIL: static-Int dim select did not reproduce the [1,N,M] input\n"); + return false; + } + if (symint_out != int_out) { + printf("FAIL: SymInt dim select disagrees with static Int dim\n"); + return false; + } + printf("PASS: select SymInt dim matches static Int dim\n"); + return true; +} + static bool test_select_double_scalars() { constexpr double kInt64Limit = 0x1p63; const double below_int64_min = @@ -2821,12 +2944,26 @@ static bool test_select_double_scalars() { ok = test_select_scalar_build_error( -kInt64Limit, /*index=*/0.0, "select: dim out of range") && ok; + // A SymInt fed by sym_size.int really is runtime-driven: select drops an + // axis, so a varying dim would change the output rank. Still rejected. + ok = test_select_scalar_build_error( + /*dim=*/0.0, + /*index=*/0.0, + "select: dynamic/unsupported dim", + SelectDimKind::SymIntFromDim) && + ok; + // Arithmetic on a runtime-driven SymInt is just as dynamic. This one reaches + // select through a resize hook rather than a source registry, and its + // build-time seed (dim=1) is a valid dim -- so a missed guard would build + // and silently emit the wrong rank rather than erroring. ok = test_select_scalar_build_error( /*dim=*/0.0, /*index=*/0.0, "select: dynamic/unsupported dim", - /*symint_dim=*/true) && + SelectDimKind::SymIntFromArithmetic) && ok; + ok = test_select_constant_symint_dim_builds() && ok; + ok = test_select_symint_dim_matches_int_dim() && ok; return ok; }