Skip to content
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

[webgpu] support Pad operator #23141

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
+389 −10
Open

[webgpu] support Pad operator #23141

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fc309db
[webgpu] support Pad operator
xhcao Dec 18, 2024
dbe430f
Fix compiling error on Mac OS
xhcao Dec 20, 2024
491e597
Merge branch 'main' into pad
xhcao Dec 20, 2024
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
262 changes: 262 additions & 0 deletions onnxruntime/core/providers/webgpu/tensor/pad.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,262 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "core/providers/webgpu/tensor/pad.h"
#include "core/providers/webgpu/shader_helper.h"
#include "core/providers/webgpu/webgpu_common.h"
#include "core/providers/webgpu/webgpu_supported_types.h"

namespace onnxruntime {
namespace webgpu {

template <typename T>
Status PadProgram<T>::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);

shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.output_size");
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() << "output[global_idx] = uniforms.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";

std::string shapeDimStr = output.Rank() == 1 ? "" : "[dim]";
std::string strideDimStr = output.Rank() < 3 ? "" : "[dim]";
std::string begin_axis_statement, end_axis_statement;
std::string in_axis_statement = "in_coord = i32(output_indices" + shapeDimStr + ") - uniforms.lower_pads" +
shapeDimStr + ";\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 = "in_coord = 0;\n";
end_axis_statement = "in_coord = i32(uniforms.data_shape" + shapeDimStr + ") - 1;\n";
break;
case Mode::Reflect:
begin_axis_statement = "in_coord = uniforms.lower_pads" + shapeDimStr + " - i32(output_indices" +
shapeDimStr + ");\n";
end_axis_statement = "in_coord = i32(uniforms.data_shape" + shapeDimStr + ") - 2 - (i32(output_indices" +
shapeDimStr + ") - (uniforms.lower_pads" + shapeDimStr + " + i32(uniforms.data_shape" +
shapeDimStr + ")));\n";
break;
case Mode::Wrap:
begin_axis_statement = "in_coord = i32(uniforms.data_shape" + shapeDimStr + " + output_indices" +
shapeDimStr + ") - uniforms.lower_pads" + shapeDimStr + ";\n";
end_axis_statement = "in_coord = i32(output_indices" + shapeDimStr + ") - uniforms.lower_pads" +
shapeDimStr + " - i32(uniforms.data_shape" + shapeDimStr + ");\n";
break;
default:
break;
}

std::string input_index_statement = output.Rank() < 2 ? "" : " if (dim + 1 < " + std::to_string(output.Rank()) + ") {\n" + " input_index += uniforms.data_stride" + strideDimStr + " * u32(in_coord);\n" + " }\n";

Check warning on line 62 in onnxruntime/core/providers/webgpu/tensor/pad.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

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

Raw Output:
onnxruntime/core/providers/webgpu/tensor/pad.cc:62:  Add #include <string> for string  [build/include_what_you_use] [4]
shader.MainFunctionBody() << " for (var dim = 0; dim < " << output.Rank() << " && !use_pad_value; dim++) {\n"
<< " if (i32(output_indices" << shapeDimStr << ") < uniforms.lower_pads" << shapeDimStr << ") {\n"
<< " " << begin_axis_statement << " }\n"
<< " else if (i32(output_indices" << shapeDimStr << ") >= uniforms.lower_pads"
<< shapeDimStr << " + i32(uniforms.data_shape" << shapeDimStr << ")) {\n"
<< " " << end_axis_statement << " }\n"
<< " else {\n"
<< " " << in_axis_statement << " }\n"
<< input_index_statement
<< " }\n"
<< " input_index += u32(in_coord);\n"
<< " output[global_idx] = select(data[input_index], uniforms.constant_value, use_pad_value);\n";

return Status::OK();
}

template <typename T>
typename ToWebGpuType<T>::MappedType ToWebGpuValue(const T& value) {
return value;
}

template <>
typename ToWebGpuType<MLFloat16>::MappedType ToWebGpuValue<MLFloat16>(const MLFloat16& value) {
return *reinterpret_cast<const typename ToWebGpuType<MLFloat16>::MappedType*>(&value.val);
}

template <typename T>
Status Pad<T>::ComputeInternal(ComputeContext& context) const {
typedef typename ToWebGpuType<T>::MappedType WebGpuT;
const Tensor* input_tensor = context.Input<Tensor>(0);
auto const& input_shape = input_tensor->Shape();
int32_t dimension_count = static_cast<int32_t>(input_shape.NumDimensions());

const PadsVector* p_pads = &pads_;
const PadsVector* p_slices = &slices_;
WebGpuT value = ToWebGpuType<T>::FromFloat(value_);

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<Tensor>(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<int64_t>();

// 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);

T raw_value{};
const Tensor* value_tensor = context.Input<Tensor>(2);
if (nullptr != value_tensor) {
ORT_ENFORCE(utils::IsPrimitiveDataType<T>(value_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");
raw_value = value_tensor->Data<T>()[0];
value = ToWebGpuValue<T>(raw_value);
}
p_pads = &pads;
p_slices = &slices;
}

auto output_dims(input_shape.AsShapeVector());
ORT_ENFORCE(static_cast<size_t>(dimension_count) * 2 == p_pads->size(), "'pads' attribute has wrong number of values");

// Calculate output dimensions, and handle any negative padding
std::vector<int32_t> lower_pads(dimension_count);

Check warning on line 137 in onnxruntime/core/providers/webgpu/tensor/pad.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

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

Raw Output:
onnxruntime/core/providers/webgpu/tensor/pad.cc:137:  Add #include <vector> for vector<>  [build/include_what_you_use] [4]
for (auto i = 0; i < dimension_count; i++) {
int64_t lower_pad = (*p_pads)[i] + (*p_slices)[i];
int64_t upper_pad = (*p_pads)[static_cast<int64_t>(i) + dimension_count] + (*p_slices)[static_cast<int64_t>(i) + dimension_count];
lower_pads[i] = static_cast<int32_t>(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<uint32_t>(output_shape.Size());
if (output_size == 0) {
// Do not need to fill output, return
return Status::OK();
}

PadProgram<T> program{mode_, dim_value_zero};
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<int>(mode_)), dim_value_zero)
.AddUniformVariables({{gsl::span<const int32_t>(lower_pads.data(), lower_pads.size())}, {output_size}, {value}});

return context.RunProgram(program);
}

#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
2, 10, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
11, 12, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
13, 17, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
18, 18, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.InputMemoryType(OrtMemTypeCPUInput, 3) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
19, 20, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.InputMemoryType(OrtMemTypeCPUInput, 3) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
21, 22, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.InputMemoryType(OrtMemTypeCPUInput, 3) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Pad<T>); \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
Pad, \
kOnnxDomain, \
23, \
T, \
kWebGpuExecutionProvider, \
(*KernelDefBuilder::Create()) \
.InputMemoryType(OrtMemTypeCPUInput, 1) \
.InputMemoryType(OrtMemTypeCPUInput, 2) \
.InputMemoryType(OrtMemTypeCPUInput, 3) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems you needn't have bothered with all the specialized stuff if you had used WebGpuSupportedNumberTypes() like this.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pad is a template class, it should transfer template type when registering. I am not sure whether WebGpuSupportedNumberTypes() works correctly. I had referred CUDA EP.

Pad<T>);

#define SPECIALIZED_COMPUTE(T) \
REGISTER_KERNEL_TYPED(T) \
template Status Pad<T>::ComputeInternal(ComputeContext& context) const;

SPECIALIZED_COMPUTE(float)
SPECIALIZED_COMPUTE(MLFloat16)
SPECIALIZED_COMPUTE(uint32_t)
SPECIALIZED_COMPUTE(int32_t)

} // namespace webgpu
} // namespace onnxruntime
39 changes: 39 additions & 0 deletions onnxruntime/core/providers/webgpu/tensor/pad.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
// 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 {

template <typename T>
class PadProgram final : public Program<PadProgram<T> > {
public:
PadProgram(const Mode mode, bool dim_value_zero) : Program<PadProgram<T> >{"Pad"}, mode_{mode}, dim_value_zero_{dim_value_zero} {}

Status GenerateShaderCode(ShaderHelper& sh) const override;

WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"lower_pads", ProgramUniformVariableDataType::Int32},
{"output_size", ProgramUniformVariableDataType::Uint32},
{"constant_value",
std::is_same_v<T, float> ? ProgramUniformVariableDataType::Float32 : (std::is_same_v<T, int32_t> ? ProgramUniformVariableDataType::Int32 : (std::is_same_v<T, uint32_t> ? ProgramUniformVariableDataType::Uint32 : ProgramUniformVariableDataType::Float16))});

private:
Mode mode_;
bool dim_value_zero_;
};

template <typename T>
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
32 changes: 32 additions & 0 deletions onnxruntime/core/providers/webgpu/webgpu_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once

#include "core/framework/float16.h"
#include "core/util/math.h"

namespace onnxruntime {
namespace webgpu {

template <typename T>
class ToWebGpuType {
public:
typedef T MappedType;
static MappedType FromFloat(float f) {
return static_cast<T>(f);
}
};

template <>
class ToWebGpuType<MLFloat16> {
public:
typedef MLFloat16 MappedType;
static MappedType FromFloat(float f) {
uint16_t h = math::floatToHalf(f);
return *reinterpret_cast<MappedType*>(&h);
}
};

} // namespace webgpu
} // namespace onnxruntime
Loading
Loading