[WebGPU EP] Support Cast to and from uint8

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<u32>(0x1, 0x100, 0x10000, 0x1000000), vec4<u32>(" << value << "));";
       break;
+    case onnxruntime::webgpu::ProgramVariableDataType::Uint8x4:
+      // Pack 4 uint8 elements (supplied as a vec4<u32>, 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<u32>(0x1u, 0x100u, 0x10000u, 0x1000000u), (vec4<u32>("
+         << value << ") & vec4<u32>(0xFFu)));";
+      break;
     default:
       ss << name_ << "[" << offset << "]=" << value << ";";
   }

onnxruntime/core/providers/webgpu/tensor/cast.cc

@@ -23,15 +23,17 @@ Status Cast::ComputeInternal(ComputeContext& context) const {
   bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType<int64_t>();
   bool is_from_float = input_tensor->DataType() == DataTypeImpl::GetType<float>() ||
                        input_tensor->DataType() == DataTypeImpl::GetType<MLFloat16>();
+  bool is_from_uint8 = input_tensor->DataType() == DataTypeImpl::GetType<uint8_t>();
   bool is_from_unsigned = input_tensor->DataType() == DataTypeImpl::GetType<uint32_t>() ||
+                          is_from_uint8 ||
                           input_tensor->DataType() == DataTypeImpl::GetType<bool>();
   const int in_components = is_from_int64 ? 1 : 4;
   const int out_components = to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64 ? 1 : 4;
   uint32_t vec_size = onnxruntime::narrow<uint32_t>((size + 3) / 4);
   uint32_t in_vec_size = onnxruntime::narrow<uint32_t>(in_components == 1 ? size : vec_size);
   uint32_t out_vec_size = onnxruntime::narrow<uint32_t>(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<u32>(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<u32> of per-lane values (e.g. bool 0/1 -> uint8).
+      expression = "vec4<u32>(a)";
+      break;
     case ONNX_NAMESPACE::TensorProto_DataType_BOOL:
       expression = "vec4<bool>(a)";
       break;
@@ -151,7 +158,11 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const {
     // cast to int64 (non-int64 inputs only)
     std::array<std::string, 4> values;
     constexpr std::array<char, 4> 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 <int StartVersion, int EndVersion>
 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<MLDataType> 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<uint8_t>());
+  // 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<MLDataType> t2_constraints = GetOpTypeConstraints(/*enable_int64=*/true, /*enable_bool=*/true);
+  t2_constraints.push_back(DataTypeImpl::GetTensorType<uint8_t>());
 
   KernelCreatePtrFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr<OpKernel>& out) -> Status {
     out = std::make_unique<Cast>(info);

onnxruntime/core/providers/webgpu/tensor/cast.h

@@ -12,8 +12,13 @@ namespace webgpu {
 
 class CastProgram final : public Program<CastProgram> {
  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<CastProgram> {
   bool is_from_int64_;
   bool is_from_float_;
   bool is_from_unsigned_;
+  bool is_from_uint8_;
 };
 
 class Cast final : public WebGpuKernel {

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<int64_t> dims{2, 3};
+  OpTester test("Cast", 13);
+  test.AddAttribute<int64_t>("to", utils::ToTensorProtoElementType<uint8_t>());
+  test.AddInput<bool>("input", dims, {true, false, true, false, true, true});
+  test.AddOutput<uint8_t>("output", dims, {1, 0, 1, 0, 1, 1});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<int64_t> dims{2, 3};
+  OpTester test("Cast", 13);
+  test.AddAttribute<int64_t>("to", utils::ToTensorProtoElementType<int32_t>());
+  test.AddInput<uint8_t>("input", dims, {0, 1, 128, 255, 42, 7});
+  test.AddOutput<int32_t>("output", dims, {0, 1, 128, 255, 42, 7});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<int64_t> dims{2, 3};
+  OpTester test("Cast", 13);
+  test.AddAttribute<int64_t>("to", utils::ToTensorProtoElementType<float>());
+  test.AddInput<uint8_t>("input", dims, {0, 1, 128, 255, 42, 7});
+  test.AddOutput<float>("output", dims, {0.0f, 1.0f, 128.0f, 255.0f, 42.0f, 7.0f});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<int64_t> dims{2, 3};
+  OpTester test("Cast", 13);
+  test.AddAttribute<int64_t>("to", utils::ToTensorProtoElementType<bool>());
+  test.AddInput<uint8_t>("input", dims, {0, 1, 128, 255, 0, 7});
+  test.AddOutput<bool>("output", dims, {false, true, true, true, false, true});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(std::move(webgpu_ep));
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
 TEST(CastOpTest, FromString) {
   const std::vector<int64_t> shape{2, 2, 2};
   const std::vector<std::string> string_data = {"-inf", "+INF", "0.9767611", "0.28280696",