From ee81e5cd553e14984a5fff4d5a883a98ad8a0a2b Mon Sep 17 00:00:00 2001 From: mingmingtasd Date: Wed, 22 Jul 2026 14:09:25 +0800 Subject: [PATCH] [WebGPU EP] Support Cast to and from uint8 Add uint8 support to the WebGPU Cast in both directions, so uint8 Cast nodes stop falling back to the CPU EP (which aborts session init under the compile-only flow). uint8 was in neither the input (T1) nor output (T2) type constraint, so both casts *to* uint8 (e.g. a terminal bool->uint8 at a graph output) and casts *from* uint8 (back to int32/float/bool) fell back. WebGPU stores uint8 packed as Uint8x4 (4 per u32, lane 0 -> low byte), the same layout already used for bool. * to uint8 (cast.cc, shader_variable.cc): add uint8 to the output (T2) type constraint, a to-uint8 shader expression, and a Uint8x4 packing path in SetByOffset. * from uint8 (cast.cc, shader_variable.cc): add uint8 to the input (T1) type constraint, a Uint8x4 unpacking path in GetByOffset (reads the 4 bytes into a vec4, inverse of the packing), and include uint8 in is_from_unsigned so uint8->int64 zero-extends. Tests (explicit WebGPU EP), cast_op_test.cc: BoolToUint8_WebGpu (to uint8; non-multiple-of-4 shape for the partial packed word); Uint8ToInt32_WebGpu (from uint8; values 0/1/128/255 spanning the full byte range to prove each byte unpacks); Uint8ToFloat_WebGpu (from uint8, float path); Uint8ToBool_WebGpu (from uint8; the distinct Boolx4 output-packing path). --- .../core/providers/webgpu/shader_variable.cc | 11 +++ .../core/providers/webgpu/tensor/cast.cc | 36 +++++++-- .../core/providers/webgpu/tensor/cast.h | 10 ++- .../test/providers/cpu/tensor/cast_op_test.cc | 81 +++++++++++++++++++ 4 files changed, 130 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index 6940d5a62d80f..5e3b7953539b2 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -361,6 +361,10 @@ std::string ShaderVariableHelper::GetByOffsetImpl(std::string_view offset) const << name_ << "[" << offset << "] & 0xFF000000u))"; break; default: + // Uint8x4 falls through here intentionally: GetByOffset returns the raw packed u32 storage + // word, matching the convention other kernels rely on for byte-packed uint8 tensors (they + // unpack sub-byte fields themselves, e.g. via unpack4xU8). Callers that want the 4 unpacked + // byte values apply unpack4xU8 at the use site. ss << name_ << "[" << offset << "]"; } @@ -393,6 +397,13 @@ std::string ShaderVariableHelper::SetByOffsetImpl(std::string_view offset, std:: case onnxruntime::webgpu::ProgramVariableDataType::Boolx4: ss << name_ << "[" << offset << "]=dot(vec4(0x1, 0x100, 0x10000, 0x1000000), vec4(" << value << "));"; break; + case onnxruntime::webgpu::ProgramVariableDataType::Uint8x4: + // Pack 4 uint8 elements (supplied as a vec4, one value per lane) into a single u32 + // storage word, lane 0 -> low byte. Same byte layout as Boolx4, but mask each lane to a + // byte so values > 1 (real uint8, not just 0/1) pack correctly. + ss << name_ << "[" << offset << "]=dot(vec4(0x1u, 0x100u, 0x10000u, 0x1000000u), (vec4(" + << value << ") & vec4(0xFFu)));"; + break; default: ss << name_ << "[" << offset << "]=" << value << ";"; } diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index a494ac788edab..a2cfbea5e6069 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -23,7 +23,9 @@ Status Cast::ComputeInternal(ComputeContext& context) const { bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType(); bool is_from_float = input_tensor->DataType() == DataTypeImpl::GetType() || input_tensor->DataType() == DataTypeImpl::GetType(); + bool is_from_uint8 = input_tensor->DataType() == DataTypeImpl::GetType(); bool is_from_unsigned = input_tensor->DataType() == DataTypeImpl::GetType() || + is_from_uint8 || input_tensor->DataType() == DataTypeImpl::GetType(); const int in_components = is_from_int64 ? 1 : 4; const int out_components = to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64 ? 1 : 4; @@ -31,7 +33,7 @@ Status Cast::ComputeInternal(ComputeContext& context) const { uint32_t in_vec_size = onnxruntime::narrow(in_components == 1 ? size : vec_size); uint32_t out_vec_size = onnxruntime::narrow(out_components == 1 ? size : vec_size); - CastProgram program{to_, is_from_int64, is_from_float, is_from_unsigned}; + CastProgram program{to_, is_from_int64, is_from_float, is_from_unsigned, is_from_uint8}; program .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {in_vec_size}, in_components}) .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {out_vec_size}, out_components}) @@ -61,6 +63,11 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { case ONNX_NAMESPACE::TensorProto_DataType_UINT32: expression = "vec4(a)"; break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT8: + // Output uint8 is stored packed (4 bytes per u32); the SetByOffset for Uint8x4 does the + // packing, so here we just widen to a vec4 of per-lane values (e.g. bool 0/1 -> uint8). + expression = "vec4(a)"; + break; case ONNX_NAMESPACE::TensorProto_DataType_BOOL: expression = "vec4(a)"; break; @@ -151,7 +158,11 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { // cast to int64 (non-int64 inputs only) std::array values; constexpr std::array kLanes{'x', 'y', 'z', 'w'}; - sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n" + // A uint8 input is stored packed (4 bytes per u32); GetByOffset returns the raw word, so + // unpack4xU8 recovers the 4 per-lane byte values. Other inputs are already vec4-shaped. + const std::string load_a = is_from_uint8_ ? "unpack4xU8(" + input.GetByOffset("global_idx") + ")" + : input.GetByOffset("global_idx"); + sh.MainFunctionBody() << " let a = " << load_a << ";\n" << " let base = global_idx * 4u;\n"; for (size_t i = 0; i < 4; ++i) { if (is_from_float_) { @@ -175,8 +186,13 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { << " }\n"; } } else { - // generic cast (no int64 involved). - sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; + // generic cast (no int64 involved). A uint8 input is stored packed (4 bytes per u32); + // GetByOffset returns the raw word, so unpack4xU8 recovers the 4 per-lane byte values. + if (is_from_uint8_) { + sh.MainFunctionBody() << " let a = unpack4xU8(" << input.GetByOffset("global_idx") << ");\n"; + } else { + sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; + } sh.MainFunctionBody() << output.SetByOffset("global_idx", expression); } @@ -186,8 +202,16 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { template KernelCreateInfo CreateCastKernelInfo(bool enable_int64) { // Casting to int64 is always supported. Casting *from* int64 (int64 in T1/input) stays guarded by enable_int64. - const auto& t1_constraints = GetOpTypeConstraints(/*enable_int64=*/enable_int64, /*enable_bool=*/true); - const auto& t2_constraints = GetOpTypeConstraints(/*enable_int64=*/true, /*enable_bool=*/true); + std::vector t1_constraints = GetOpTypeConstraints(/*enable_int64=*/enable_int64, /*enable_bool=*/true); + // T1 (input): plus uint8, so casts *from* a uint8 tensor (to int32/float/bool/etc.) run on the + // WebGPU EP instead of falling back. Reading uint8 input unpacks the packed Uint8x4 storage via + // GetByOffset. + t1_constraints.push_back(DataTypeImpl::GetTensorType()); + // T2 (output): the int64+bool set, plus uint8 so the WebGPU EP can produce uint8 output directly + // (e.g. a terminal bool->uint8 cast at a graph output) instead of falling back. Casting *to* uint8 + // packs via the Uint8x4 SetByOffset. + std::vector t2_constraints = GetOpTypeConstraints(/*enable_int64=*/true, /*enable_bool=*/true); + t2_constraints.push_back(DataTypeImpl::GetTensorType()); KernelCreatePtrFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { out = std::make_unique(info); diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index e81d5e594a310..6ecbefe4a7ac7 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -12,8 +12,13 @@ namespace webgpu { class CastProgram final : public Program { public: - CastProgram(int32_t to, bool is_from_int64, bool is_from_float, bool is_from_unsigned) - : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64}, is_from_float_{is_from_float}, is_from_unsigned_{is_from_unsigned} {} + CastProgram(int32_t to, bool is_from_int64, bool is_from_float, bool is_from_unsigned, bool is_from_uint8) + : Program{"Cast"}, + to_{to}, + is_from_int64_{is_from_int64}, + is_from_float_{is_from_float}, + is_from_unsigned_{is_from_unsigned}, + is_from_uint8_{is_from_uint8} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -25,6 +30,7 @@ class CastProgram final : public Program { bool is_from_int64_; bool is_from_float_; bool is_from_unsigned_; + bool is_from_uint8_; }; class Cast final : public WebGpuKernel { diff --git a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc index 956cc09c3efd9..3727ae28891a8 100644 --- a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc @@ -18,6 +18,7 @@ #include "test/common/cuda_op_test_utils.h" #include "test/providers/provider_test_utils.h" +#include "test/util/include/default_providers.h" namespace onnxruntime { namespace test { @@ -193,6 +194,86 @@ TEST(CastOpTest, NonStringTypes) { CastNonStringTester{}); } +// bool -> uint8 Cast, run explicitly on the WebGPU EP (e.g. a terminal bool->uint8 cast at a graph +// output). WebGPU stores uint8 packed (4 per u32); the Cast writes it via the Uint8x4 SetByOffset +// packing. The 6-element shape deliberately isn't a multiple of 4, exercising the final partial +// packed word. +TEST(CastOpTest, BoolToUint8_WebGpu) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (webgpu_ep == nullptr) { + GTEST_SKIP() << "WebGPU EP is not available in this build."; + } + + const std::vector dims{2, 3}; + OpTester test("Cast", 13); + test.AddAttribute("to", utils::ToTensorProtoElementType()); + test.AddInput("input", dims, {true, false, true, false, true, true}); + test.AddOutput("output", dims, {1, 0, 1, 0, 1, 1}); + + std::vector> execution_providers; + execution_providers.push_back(std::move(webgpu_ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// uint8 -> int32 Cast, run explicitly on the WebGPU EP. WebGPU stores uint8 packed (4 per u32); +// reading it unpacks via the Uint8x4 GetByOffset. Values span the full byte range (0, 1, 128, 255) +// to prove each byte is unpacked correctly, and the 6-element shape isn't a multiple of 4 so the +// final partial packed word is exercised. +TEST(CastOpTest, Uint8ToInt32_WebGpu) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (webgpu_ep == nullptr) { + GTEST_SKIP() << "WebGPU EP is not available in this build."; + } + + const std::vector dims{2, 3}; + OpTester test("Cast", 13); + test.AddAttribute("to", utils::ToTensorProtoElementType()); + test.AddInput("input", dims, {0, 1, 128, 255, 42, 7}); + test.AddOutput("output", dims, {0, 1, 128, 255, 42, 7}); + + std::vector> execution_providers; + execution_providers.push_back(std::move(webgpu_ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// uint8 -> float Cast on the WebGPU EP, exercising the uint8-input read on the float conversion path. +TEST(CastOpTest, Uint8ToFloat_WebGpu) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (webgpu_ep == nullptr) { + GTEST_SKIP() << "WebGPU EP is not available in this build."; + } + + const std::vector dims{2, 3}; + OpTester test("Cast", 13); + test.AddAttribute("to", utils::ToTensorProtoElementType()); + test.AddInput("input", dims, {0, 1, 128, 255, 42, 7}); + test.AddOutput("output", dims, {0.0f, 1.0f, 128.0f, 255.0f, 42.0f, 7.0f}); + + std::vector> execution_providers; + execution_providers.push_back(std::move(webgpu_ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// uint8 -> bool Cast on the WebGPU EP. Exercises the distinct Boolx4 output-packing path (vs. the +// direct write used for int32/float above) and the standard cast-to-bool semantics (any nonzero +// byte -> true) -- the natural inverse of the bool->uint8 case. +TEST(CastOpTest, Uint8ToBool_WebGpu) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (webgpu_ep == nullptr) { + GTEST_SKIP() << "WebGPU EP is not available in this build."; + } + + const std::vector dims{2, 3}; + OpTester test("Cast", 13); + test.AddAttribute("to", utils::ToTensorProtoElementType()); + test.AddInput("input", dims, {0, 1, 128, 255, 0, 7}); + test.AddOutput("output", dims, {false, true, true, true, false, true}); + + std::vector> execution_providers; + execution_providers.push_back(std::move(webgpu_ep)); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + TEST(CastOpTest, FromString) { const std::vector shape{2, 2, 2}; const std::vector string_data = {"-inf", "+INF", "0.9767611", "0.28280696",