diff --git a/onnxruntime/core/providers/webgpu/tensor/pad.cc b/onnxruntime/core/providers/webgpu/tensor/pad.cc new file mode 100644 index 0000000000000..9ee13aada67fe --- /dev/null +++ b/onnxruntime/core/providers/webgpu/tensor/pad.cc @@ -0,0 +1,261 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include + +#include "core/util/math.h" +#include "core/providers/webgpu/tensor/pad.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" + +namespace onnxruntime { +namespace webgpu { + +Status PadProgram::GenerateShaderCode(ShaderHelper& shader) const { + if (!dim_value_zero_) { + shader.AddInput("data", ShaderUsage::UseUniform | ShaderUsage::UseShapeAndStride); + } + const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseShapeAndStride | ShaderUsage::UseValueTypeAlias); + + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size"); + std::string constant_value_str = std::string("let constant_value = ") + + (is_float16_ ? "bitcast>(uniforms.constant_value)[0];\n" : "bitcast(uniforms.constant_value);\n"); + if (dim_value_zero_) { + // Only Constant mode needs fill output if the one dim value or mores dims' values of input are zero. + shader.MainFunctionBody() << constant_value_str + << "output[global_idx] = constant_value;\n"; + return Status::OK(); + } + + shader.MainFunctionBody() << " let output_indices = " << output.OffsetToIndices("global_idx") << ";\n" + << " var input_index = u32(0);\n" + << " var use_pad_value = false;\n" + << " var in_coord = i32(0);\n"; + + const int rank = output.Rank(); + std::string output_indices_str = "i32(" + GetElementAt("output_indices", "dim", rank) + ")"; + std::string lower_pads_str = GetElementAt("uniforms.lower_pads", "dim", rank); + std::string data_shape_str = "i32(" + GetElementAt("uniforms.data_shape", "dim", rank) + ")"; + std::string data_stride_str = rank == 1 ? "" : " * " + GetElementAt("uniforms.data_stride", "dim", rank - 1); + std::string begin_axis_statement = "in_coord = "; + std::string end_axis_statement = "in_coord = "; + std::string in_axis_statement = "in_coord = " + output_indices_str + " - " + lower_pads_str + ";\n"; + switch (mode_) { + case Mode::Constant: + begin_axis_statement = "use_pad_value = true;\n"; + end_axis_statement = "use_pad_value = true;\n"; + break; + case Mode::Edge: + begin_axis_statement += "0;\n"; + end_axis_statement += data_shape_str + " - 1;\n"; + break; + case Mode::Reflect: + begin_axis_statement += lower_pads_str + " - " + output_indices_str + ";\n"; + end_axis_statement += data_shape_str + " - 2 - (" + output_indices_str + + " - (" + lower_pads_str + " + " + data_shape_str + "));\n"; + break; + case Mode::Wrap: + begin_axis_statement += data_shape_str + " + " + output_indices_str + " - " + lower_pads_str + ";\n"; + end_axis_statement += output_indices_str + " - " + lower_pads_str + " - " + data_shape_str + ";\n"; + break; + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported mode type: ", static_cast(mode_)); + } + + shader.MainFunctionBody() << " for (var dim = 0; dim < " << rank << " && !use_pad_value; dim++) {\n" + << " if (" << output_indices_str << " < " << lower_pads_str << ") {\n" + << " " << begin_axis_statement << " }\n" + << " else if (" << output_indices_str << " >= " << lower_pads_str << " + " << data_shape_str << ") {\n" + << " " << end_axis_statement << " }\n" + << " else {\n" + << " " << in_axis_statement << " }\n" + << " input_index += select(u32(in_coord)" << data_stride_str << ", u32(in_coord), dim == " << rank - 1 << ");\n" + << " }\n" + << " " << constant_value_str + << " " << output.SetByOffset("global_idx", "select(data[input_index], constant_value, use_pad_value)"); + + return Status::OK(); +} + +Status Pad::ComputeInternal(ComputeContext& context) const { + const Tensor* input_tensor = context.Input(0); + auto const& input_shape = input_tensor->Shape(); + size_t dimension_count = input_shape.NumDimensions(); + + const PadsVector* p_pads = &pads_; + const PadsVector* p_slices = &slices_; + + PadsVector pads; + PadsVector slices; + // kOnnxDomain Pad opset >= 11 (Or) kMsDomain opset == 1 + if (is_dynamic_) { + size_t data_rank = input_tensor->Shape().NumDimensions(); + + const Tensor* pads_tensor = context.Input(1); + auto pads_tensor_dims = pads_tensor->Shape().GetDims(); + ORT_ENFORCE(pads_tensor_dims.size() == 1 || (pads_tensor_dims.size() == 2 && pads_tensor_dims[0] == 1), + "Pads tensor should be a 1D tensor of shape [2 * num_axes] " + "or a 2D tensor of shape [1, 2 * num_axes]"); + + const auto pads_data = pads_tensor->DataAsSpan(); + + // Compute Pads by applying axes if specified otherwise copy the supplied pads. + PadBase::ComputePads(context.KernelContext(), data_rank, pads_data, pads); + + // Separate out any negative pads into the slices array + PadBase::SeparateNegativeToSlices(pads, slices); + + p_pads = &pads; + p_slices = &slices; + } + + auto output_dims(input_shape.AsShapeVector()); + ORT_ENFORCE(dimension_count * 2 == p_pads->size(), "'pads' attribute has wrong number of values"); + + // Calculate output dimensions, and handle any negative padding + std::vector lower_pads(dimension_count); + for (size_t i = 0; i < dimension_count; i++) { + int64_t lower_pad = (*p_pads)[i] + (*p_slices)[i]; + int64_t upper_pad = (*p_pads)[i + dimension_count] + (*p_slices)[i + dimension_count]; + lower_pads[i] = static_cast(lower_pad); + output_dims[i] += lower_pad + upper_pad; + } + TensorShape output_shape(output_dims); + + // special case when there is a dim value of 0 in the shape. behavior depends on mode + bool dim_value_zero = input_shape.Size() == 0; + if (dim_value_zero) { + ORT_RETURN_IF_ERROR(PadBase::HandleDimValueZero(mode_, input_shape, output_shape)); + } + + auto* output_tensor = context.Output(0, output_shape); + uint32_t output_size = gsl::narrow(output_shape.Size()); + if (output_size == 0) { + // Do not need to fill output, return + return Status::OK(); + } + + // Read constant value and bitcast to uint32. + uint32_t value_uint32 = 0; + const auto data_type = input_tensor->GetElementType(); + bool is_float16 = data_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16; + const Tensor* value_tensor = context.Input(2); + if (!is_dynamic_) { + if (is_float16) { + uint16_t value = math::floatToHalf(value_); + std::memcpy(&value_uint32, &value, sizeof(value)); + } else { + value_uint32 = *reinterpret_cast(&value_); + } + } else if (value_tensor) { + ORT_ENFORCE(value_tensor->DataType() == input_tensor->DataType() && value_tensor->Shape().Size() == 1, + "Value tensor should be a 1D tensor of size 1 with the same type as that of the input tensor"); + switch (data_type) { + case ONNX_NAMESPACE::TensorProto_DataType_INT32: { + int32_t value = value_tensor->Data()[0]; + value_uint32 = *reinterpret_cast(&value); + } break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: { + float value = value_tensor->Data()[0]; + value_uint32 = *reinterpret_cast(&value); + } break; + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: { + uint16_t value = value_tensor->Data()[0].val; + std::memcpy(&value_uint32, &value, sizeof(value)); + } break; + case ONNX_NAMESPACE::TensorProto_DataType_UINT32: { + value_uint32 = value_tensor->Data()[0]; + } break; + default: + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Unsupported input type: ", static_cast(data_type)); + } + } + + PadProgram program{mode_, dim_value_zero, is_float16}; + if (!dim_value_zero) { + program.AddInput({input_tensor, ProgramTensorMetadataDependency::TypeAndRank}); + } + program.AddOutput({output_tensor, ProgramTensorMetadataDependency::Rank}) + .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .CacheHint(std::to_string(static_cast(mode_)), dim_value_zero) + .AddUniformVariables({{gsl::span(lower_pads.data(), lower_pads.size())}, {output_size}, {value_uint32}}); + + return context.RunProgram(program); +} + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 2, 10, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 11, 12, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 13, 17, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 18, 18, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 19, 20, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Pad, + kOnnxDomain, + 21, 22, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); +ONNX_OPERATOR_KERNEL_EX( + Pad, + kOnnxDomain, + 23, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .InputMemoryType(OrtMemTypeCPUInput, 1) + .InputMemoryType(OrtMemTypeCPUInput, 2) + .InputMemoryType(OrtMemTypeCPUInput, 3) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Pad); + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/pad.h b/onnxruntime/core/providers/webgpu/tensor/pad.h new file mode 100644 index 0000000000000..58049ddb0e5ce --- /dev/null +++ b/onnxruntime/core/providers/webgpu/tensor/pad.h @@ -0,0 +1,40 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/program.h" +#include "core/providers/webgpu/webgpu_kernel.h" +#include "core/providers/cpu/tensor/padbase.h" + +namespace onnxruntime { +namespace webgpu { + +class PadProgram final : public Program { + public: + PadProgram(const Mode mode, bool dim_value_zero, bool is_float16) : Program{"Pad"}, + mode_{mode}, + dim_value_zero_{dim_value_zero}, + is_float16_{is_float16} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"lower_pads", ProgramUniformVariableDataType::Int32}, + {"output_size", ProgramUniformVariableDataType::Uint32}, + {"constant_value", ProgramUniformVariableDataType::Uint32}); + + private: + Mode mode_; + bool dim_value_zero_; + bool is_float16_; +}; + +class Pad final : public PadBase, public WebGpuKernel { + public: + Pad(const OpKernelInfo& info) : PadBase(info), WebGpuKernel(info) {} + + Status ComputeInternal(ComputeContext& context) const override; +}; + +} // 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 d44cf4674d8a3..981c467ad725a 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -363,7 +363,9 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxD class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Pad); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, 17, Pad); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 18, 18, Pad); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, Pad); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 19, 20, Pad); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 21, 22, Pad); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 23, Pad); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, If); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, If); @@ -685,11 +687,13 @@ std::unique_ptr RegisterKernels() { // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, // BuildKernelCreateInfo, // BuildKernelCreateInfo,