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",