Skip to content

[webgpu] And int64 to cast #25610

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fs-eire merged 10 commits into from
Oct 15, 2025
Merged

[webgpu] And int64 to cast #25610

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
153e187
[webgpu] And int64 to cast
qjia7 Jul 31, 2025
0bae27b
Register conditional Cast
qjia7 Sep 29, 2025
a403660
set default to false
qjia7 Sep 29, 2025
c38e04c
fix CI errors
qjia7 Sep 29, 2025
2c46922
fix CI errors
qjia7 Sep 29, 2025
dbd1430
fix the CI errors
qjia7 Sep 29, 2025
f7aeb78
move the implementation to cc
qjia7 Sep 30, 2025
622795f
address comments
qjia7 Oct 11, 2025
4bff64c
remove the wrong the annotation
qjia7 Oct 11, 2025
eb34d8c
address comments
qjia7 Oct 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion onnxruntime/core/providers/webgpu/shader_variable.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>(u32(" << value << "), select(0u, 0xFFFFFFFFu, " << value << " < 0));";
ss << name_ << "[" << offset << "]=vec2<u32>(u32(" << value << "), select(0u, 0xFFFFFFFFu, i32(" << value << ") < 0));";
break;
case onnxruntime::webgpu::ProgramVariableDataType::Uint64:
ss << name_ << "[" << offset << "]=vec2<u32>(u32(" << value << "), 0u);";
Expand Down
154 changes: 88 additions & 66 deletions onnxruntime/core/providers/webgpu/tensor/cast.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,88 +11,47 @@
namespace webgpu {

namespace {
const std::vector<MLDataType>& CastOpTypeConstraints() {
// currently support boolean, integer and float types that explicitly allowed in WGSL:
const std::vector<MLDataType>& 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<MLDataType> types{
static std::vector<MLDataType> base_types{
DataTypeImpl::GetTensorType<MLFloat16>(),
DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<uint32_t>(),
DataTypeImpl::GetTensorType<bool>()};
return types;

if (enable_graph_capture) {
static std::vector<MLDataType> types_with_int64 = []() {
auto types = base_types;
types.push_back(DataTypeImpl::GetTensorType<int64_t>());
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());
int64_t size = input_tensor->Shape().Size();
if (size == 0) {
return Status::OK();
}
bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType<int64_t>();
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_};
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<uint32_t>(vec_size)},
Expand Down Expand Up @@ -121,15 +80,78 @@
case ONNX_NAMESPACE::TensorProto_DataType_BOOL:
expression = "vec4<bool>(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<i32>(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 <int StartVersion, int EndVersion>
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<OpKernel>& out) -> Status {
out = std::make_unique<Cast>(info);

Check warning on line 117 in onnxruntime/core/providers/webgpu/tensor/cast.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Add #include <memory> for make_unique<>  [build/include_what_you_use] [4]

Raw Output:
onnxruntime/core/providers/webgpu/tensor/cast.cc:117:  Add #include <memory> for make_unique<>  [build/include_what_you_use] [4]
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
9 changes: 8 additions & 1 deletion onnxruntime/core/providers/webgpu/tensor/cast.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,21 +3,24 @@

#pragma once

#include "core/framework/kernel_registry.h"
#include "core/framework/op_kernel.h"
#include "core/providers/webgpu/webgpu_kernel.h"

namespace onnxruntime {
namespace webgpu {

class CastProgram final : public Program<CastProgram> {
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;

WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"vec_size", ProgramUniformVariableDataType::Uint32});

private:
int32_t to_;
bool is_from_int64_;
};

class Cast final : public WebGpuKernel {
Expand All @@ -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 <int StartVersion, int EndVersion = StartVersion>
KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture);

} // namespace webgpu
} // namespace onnxruntime
28 changes: 17 additions & 11 deletions onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {

Expand Down Expand Up @@ -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<KernelRegistry> RegisterKernels() {
std::unique_ptr<KernelRegistry> RegisterKernels(bool enable_graph_capture = false) {
auto kernel_registry = std::make_unique<onnxruntime::KernelRegistry>();

static const BuildKernelCreateInfoFn function_table[] = {
Expand Down Expand Up @@ -464,13 +465,6 @@ std::unique_ptr<KernelRegistry> 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<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 11, float, Clip)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 12, 12, float, Clip)>,
Expand Down Expand Up @@ -771,6 +765,14 @@ std::unique_ptr<KernelRegistry> 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());
Expand Down Expand Up @@ -869,9 +871,13 @@ std::vector<std::unique_ptr<ComputeCapability>> WebGpuExecutionProvider::GetCapa
}

std::shared_ptr<KernelRegistry> WebGpuExecutionProvider::GetKernelRegistry() const {
static std::shared_ptr<KernelRegistry> registry = webgpu::RegisterKernels();

return registry;
if (enable_graph_capture_) {
static std::shared_ptr<KernelRegistry> registry = webgpu::RegisterKernels(true);
return registry;
} else {
static std::shared_ptr<KernelRegistry> registry = webgpu::RegisterKernels(false);
return registry;
}
}
Comment thread
qjia7 marked this conversation as resolved.

std::unique_ptr<onnxruntime::IDataTransfer> WebGpuExecutionProvider::GetDataTransfer() const {
Expand Down
Loading