diff --git a/backends/webgpu/runtime/ops/select/Select.cpp b/backends/webgpu/runtime/ops/select/Select.cpp index f2e69823974..cb4575bc004 100644 --- a/backends/webgpu/runtime/ops/select/Select.cpp +++ b/backends/webgpu/runtime/ops/select/Select.cpp @@ -14,6 +14,7 @@ #include +#include #include #include #include @@ -29,12 +30,27 @@ struct SelectParams { uint32_t _pad[2]; }; -// dim/index are required Ints (SymInt throws); no Null default unlike slice. +// dim/index are static integer scalars (SymInt throws); Vulkan serialization +// can reuse an earlier Double value for an equal integer scalar. int64_t read_scalar(WebGPUGraph& graph, int id, const char* what) { - if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) { - return graph.get_int(id); + switch (graph.get_value_type(id)) { + case WebGPUGraph::ValueType::Int: + return graph.get_int(id); + case WebGPUGraph::ValueType::Double: { + const double d = graph.get_double(id); + constexpr double kInt64Limit = 0x1p63; + // Check every invalid case before casting: floating-to-integer conversion + // is undefined for out-of-range values, and fractional values must not be + // silently rounded or truncated. std::trunc(NaN) != NaN rejects NaN. + if (d < -kInt64Limit || d >= kInt64Limit || std::trunc(d) != d) { + throw std::runtime_error(std::string("select: non-integral ") + what); + } + return static_cast(d); + } + default: + throw std::runtime_error( + std::string("select: dynamic/unsupported ") + what); } - throw std::runtime_error(std::string("select: dynamic/unsupported ") + what); } // Build a TensorMeta from live dims, write it to buf, return numel. diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index 0a16b04a356..a21466b9a42 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -1824,6 +1824,190 @@ static bool test_slice_double_start() { return ok; } +// 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. +static void finish_select_scalar_graph( + ::flatbuffers::FlatBufferBuilder& fbb, + double dim, + double index, + uint32_t out_len, + bool symint_dim = false) { + 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, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &in_dims, + /*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())); + } + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::Double, + vk::CreateDouble(fbb, /*double_val=*/index).Union())); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &out_dims, + /*constant_id=*/-1, + /*mem_obj_id=*/1) + .Union())); + + std::vector args = {0, 1, 2, 3}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.select_copy.int", &args)); + std::vector input_ids = {0}; + std::vector output_ids = {3}; + auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); +} + +static bool test_select_double_scalar_case( + double dim, + double index, + const std::vector& expected) { + printf( + "\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())); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& e) { + printf("FAIL: graph build threw: %s\n", e.what()); + return false; + } + + std::vector in = {0.5f, 1.5f, 2.5f, 3.5f, 4.5f, 5.5f}; + std::vector inputs = { + {in.data(), in.size() * sizeof(float), false}}; + std::vector out(expected.size(), -1.0f); + std::vector> outputs = { + {out.data(), out.size() * sizeof(float)}}; + try { + graph.copy_inputs(inputs); + graph.execute(); + graph.copy_outputs(outputs); + } catch (const std::exception& e) { + printf("FAIL: select execute threw: %s\n", e.what()); + return false; + } + + if (out != expected) { + printf("FAIL: select Double-scalar gather mismatch\n"); + return false; + } + printf("PASS: select Double scalars\n"); + return true; +} + +static bool test_select_scalar_build_error( + double dim, + double index, + const char* expected_error, + bool symint_dim = false) { + ::flatbuffers::FlatBufferBuilder fbb; + finish_select_scalar_graph(fbb, dim, index, /*out_len=*/3u, symint_dim); + + WebGPUGraph graph; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& e) { + const std::string error = e.what(); + if (error.find(expected_error) != std::string::npos) { + printf("PASS: rejected with expected error: %s\n", e.what()); + return true; + } + printf( + "FAIL: rejection error %s did not contain %s\n", + e.what(), + expected_error); + return false; + } + printf("FAIL: expected graph build to reject select scalars\n"); + return false; +} + +static bool test_select_double_scalars() { + constexpr double kInt64Limit = 0x1p63; + const double below_int64_min = + std::nextafter(-kInt64Limit, -std::numeric_limits::infinity()); + + bool ok = true; + // Production artifact representation: both schema-int arguments alias 0.0. + ok = test_select_double_scalar_case( + /*dim=*/0.0, /*index=*/0.0, {0.5f, 1.5f, 2.5f}) && + ok; + // Negative Double indices retain normal select index normalization. + ok = test_select_double_scalar_case( + /*dim=*/1.0, /*index=*/-1.0, {2.5f, 5.5f}) && + ok; + + ok = test_select_scalar_build_error( + /*dim=*/0.5, + /*index=*/0.0, + "select: non-integral dim") && + ok; + ok = test_select_scalar_build_error( + /*dim=*/0.0, + /*index=*/0.5, + "select: non-integral index") && + ok; + ok = test_select_scalar_build_error( + std::numeric_limits::quiet_NaN(), + /*index=*/0.0, + "select: non-integral dim") && + ok; + ok = test_select_scalar_build_error( + std::numeric_limits::infinity(), + /*index=*/0.0, + "select: non-integral dim") && + ok; + ok = test_select_scalar_build_error( + kInt64Limit, /*index=*/0.0, "select: non-integral dim") && + ok; + ok = test_select_scalar_build_error( + below_int64_min, /*index=*/0.0, "select: non-integral dim") && + ok; + // -2^63 is representable: conversion succeeds, then the normal dim check + // rejects it as out of range for this rank-2 input. + ok = test_select_scalar_build_error( + -kInt64Limit, /*index=*/0.0, "select: dim out of range") && + ok; + ok = test_select_scalar_build_error( + /*dim=*/0.0, + /*index=*/0.0, + "select: dynamic/unsupported dim", + /*symint_dim=*/true) && + ok; + return ok; +} + // apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). struct RopeConfig { const char* name; @@ -2009,6 +2193,11 @@ TEST(WebGPUNative, SliceDoubleStart) { << "slice Double-start gather/reject checks failed"; } +TEST(WebGPUNative, SelectDoubleScalars) { + EXPECT_TRUE(test_select_double_scalars()) + << "select Double-scalar compatibility checks failed"; +} + TEST(WebGPUNative, Prepack2) { if (g_prepack2_model_path.empty() || g_prepack2_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK2_MODEL/GOLDEN not set";