diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index c197e227e2a8c..ac8b3d07e6e80 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -308,7 +308,7 @@ std::string ShaderVariableHelper::SetByOffsetImpl(std::string_view offset, std:: ORT_THROW("Invalid type"); break; case onnxruntime::webgpu::ProgramVariableDataType::Int64: - ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, " << value << " < 0));"; + ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, i32(" << value << ") < 0));"; break; case onnxruntime::webgpu::ProgramVariableDataType::Uint64: ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), 0u);"; diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 313a96ba25509..daf4aa323c12e 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -11,75 +11,29 @@ namespace onnxruntime { namespace webgpu { namespace { -const std::vector& CastOpTypeConstraints() { - // currently support boolean, integer and float types that explicitly allowed in WGSL: +const std::vector& CastOpTypeConstraints(bool enable_graph_capture) { + // Base types that are always supported - boolean, integer and float types that explicitly allowed in WGSL: // https://gpuweb.github.io/gpuweb/wgsl/#plain-types-section - // - static std::vector types{ + static std::vector base_types{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}; - return types; + + if (enable_graph_capture) { + static std::vector types_with_int64 = []() { + auto types = base_types; + types.push_back(DataTypeImpl::GetTensorType()); + return types; + }(); + return types_with_int64; + } else { + return base_types; + } } } // namespace -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 6, 8, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 9, 12, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 13, 18, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 19, 20, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 21, 22, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_KERNEL_EX( - Cast, - kOnnxDomain, - 23, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); - Status Cast::ComputeInternal(ComputeContext& context) const { const auto* input_tensor = context.Input(0); auto* output_tensor = context.Output(0, input_tensor->Shape()); @@ -87,12 +41,17 @@ Status Cast::ComputeInternal(ComputeContext& context) const { if (size == 0) { return Status::OK(); } + bool is_from_int64 = 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; uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); + 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_}; + CastProgram program{to_, is_from_int64}; program - .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {vec_size}, 4}) + .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {in_vec_size}, in_components}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {out_vec_size}, out_components}) .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(vec_size)}, @@ -121,15 +80,78 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { case ONNX_NAMESPACE::TensorProto_DataType_BOOL: expression = "vec4(a)"; break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + expression = "int32(a)"; + break; default: ORT_NOT_IMPLEMENTED("Cast to type ", to_, " is not supported."); } - sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") - << " let a = " << input.GetByOffset("global_idx") << ";\n " - << output.SetByOffset("global_idx", expression); + + sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); + if (is_from_int64_) { + sh.MainFunctionBody() << " let a0 = " << input.GetByOffset("global_idx * 4") << ";\n" + << " let a1 = " << input.GetByOffset("global_idx * 4 + 1") << ";\n" + << " let a2 = " << input.GetByOffset("global_idx * 4 + 2") << ";\n" + << " let a3 = " << input.GetByOffset("global_idx * 4 + 3") << ";\n" + << " let a = vec4(a0, a1, a2, a3);\n"; + } else { + sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; + } + if (to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64) { + sh.MainFunctionBody() << output.SetByOffset("global_idx * 4", "a.x") << "\n" + << output.SetByOffset("global_idx * 4 + 1", "a.y") << "\n" + << output.SetByOffset("global_idx * 4 + 2", "a.z") << "\n" + << output.SetByOffset("global_idx * 4 + 3", "a.w") << "\n"; + } else { + sh.MainFunctionBody() << output.SetByOffset("global_idx", expression); + } return Status::OK(); } +template +KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { + const auto& type_constraints = CastOpTypeConstraints(enable_graph_capture); + + KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { + out = std::make_unique(info); + return Status::OK(); + }; + + if constexpr (StartVersion == EndVersion) { + // Non-versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } else { + // Versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion, EndVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } +} + +// Explicit template instantiations +template KernelCreateInfo CreateCastKernelInfo<6, 8>(bool); +template KernelCreateInfo CreateCastKernelInfo<9, 12>(bool); +template KernelCreateInfo CreateCastKernelInfo<13, 18>(bool); +template KernelCreateInfo CreateCastKernelInfo<19, 20>(bool); +template KernelCreateInfo CreateCastKernelInfo<21, 22>(bool); +template KernelCreateInfo CreateCastKernelInfo<23>(bool); + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 925cd200f0aba..7dfb50e3241c8 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -3,6 +3,8 @@ #pragma once +#include "core/framework/kernel_registry.h" +#include "core/framework/op_kernel.h" #include "core/providers/webgpu/webgpu_kernel.h" namespace onnxruntime { @@ -10,7 +12,7 @@ namespace webgpu { class CastProgram final : public Program { public: - CastProgram(int32_t to) : Program{"Cast"}, to_{to} {} + CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -18,6 +20,7 @@ class CastProgram final : public Program { private: int32_t to_; + bool is_from_int64_; }; class Cast final : public WebGpuKernel { @@ -37,5 +40,9 @@ class Cast final : public WebGpuKernel { int32_t to_; }; +// Create Cast kernel info with appropriate type constraints based on graph capture support +template +KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture); + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index bbb3fbdd221d3..0f7607ac1dbfe 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -28,6 +28,7 @@ #include "core/providers/webgpu/data_transfer.h" #include "core/providers/webgpu/external_data_loader.h" #include "core/providers/webgpu/webgpu_profiler.h" +#include "core/providers/webgpu/tensor/cast.h" namespace onnxruntime { @@ -417,7 +418,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxD class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 16, 17, ScatterND); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 18, ScatterND); -std::unique_ptr RegisterKernels() { +std::unique_ptr RegisterKernels(bool enable_graph_capture = false) { auto kernel_registry = std::make_unique(); static const BuildKernelCreateInfoFn function_table[] = { @@ -464,13 +465,6 @@ std::unique_ptr RegisterKernels() { KERNEL_CREATE_INFO(13, Tanh), KERNEL_CREATE_INFO(1, Not), - KERNEL_CREATE_INFO_VERSIONED(6, 8, Cast), - KERNEL_CREATE_INFO_VERSIONED(9, 12, Cast), - KERNEL_CREATE_INFO_VERSIONED(13, 18, Cast), - KERNEL_CREATE_INFO_VERSIONED(19, 20, Cast), - KERNEL_CREATE_INFO_VERSIONED(21, 22, Cast), - KERNEL_CREATE_INFO(23, Cast), - // // activations BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -771,6 +765,14 @@ std::unique_ptr RegisterKernels() { } } + // Register Cast kernels with conditional int64 support based on graph capture + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<6, 8>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<9, 12>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<13, 18>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<19, 20>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<21, 22>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<23>(enable_graph_capture))); + #ifndef DISABLE_CONTRIB_OPS Status status = ::onnxruntime::contrib::webgpu::RegisterWebGpuContribKernels(*kernel_registry); ORT_ENFORCE(status.IsOK(), "Failed to register WebGPU contrib kernels: " + status.ErrorMessage()); @@ -869,9 +871,13 @@ std::vector> WebGpuExecutionProvider::GetCapa } std::shared_ptr WebGpuExecutionProvider::GetKernelRegistry() const { - static std::shared_ptr registry = webgpu::RegisterKernels(); - - return registry; + if (enable_graph_capture_) { + static std::shared_ptr registry = webgpu::RegisterKernels(true); + return registry; + } else { + static std::shared_ptr registry = webgpu::RegisterKernels(false); + return registry; + } } std::unique_ptr WebGpuExecutionProvider::GetDataTransfer() const {