diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc new file mode 100644 index 0000000000..796c56f67c --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -0,0 +1,241 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/common/inlined_containers.h" +#include "core/providers/webgpu/tensor/softmax.h" +#include "core/providers/webgpu/tensor/transpose.h" +#include "core/providers/cpu/tensor/utils.h" +#include "core/providers/webgpu/shader_variable.h" +#include "core/providers/webgpu/shader_helper.h" +#include "core/providers/webgpu/webgpu_supported_types.h" + +namespace onnxruntime { +namespace webgpu { + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Softmax, + kOnnxDomain, + 1, 10, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Softmax); + +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Softmax, + kOnnxDomain, + 11, 12, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Softmax); + +ONNX_OPERATOR_KERNEL_EX( + Softmax, + kOnnxDomain, + 13, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .TypeConstraint("T", WebGpuSupportedNumberTypes()), + Softmax); + +static std::string MaxVector(std::string name, int components) { + switch (components) { + case 1: + return name; + case 2: + return "max(" + name + ".x, " + name + ".y)"; + case 4: + return "max(max(" + name + ".x, " + name + ".y), max(" + name + ".z, " + name + ".w))"; + default: + ORT_THROW("Unsupported number of components: ", components); + } +} + +static std::string SumVector(std::string x, int components) { + switch (components) { + case 1: + return x; + case 2: + return "(" + x + ".x + " + x + ".y" + ")"; + case 4: + return "(" + x + ".x + " + x + ".y + " + x + ".w + " + x + ".z" + ")"; + default: + ORT_THROW("Unsupported number of components: ", components); + } +} + +static int GetMaxComponents(int64_t size) { + if (size % 4 == 0) { + return 4; + } else if (size % 2 == 0) { + return 2; + } + return 1; +} + +Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { + // Add input and output variables + const auto& input = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + const auto& output = shader.AddOutput("result", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); + int components = input.NumComponents(); + + std::string threadMaxDecl = input.StorageType() == "f32" ? + "val threadMax = x_value_t(-3.402823e+38f);\n" : + "val threadMax = x_value_t(-65504.0h));\n"; + + + // Define shared memory for row max and row sum + shader.AdditionalImplementation() + << "var rowMaxShared : x_value_t;\n" + << "var rowSumShared : x_value_t;\n" + << "var threadShared : array;\n"; + + // Define helper functions to get and set values + shader.AdditionalImplementation() + << "fn getValue(row: i32, col: i32, row_stride: i32) -> x_value_t {\n" + << " let index = row * row_stride + col;\n" + << " return x[index];\n" + << "}\n" + << "fn setValue(row: i32, col: i32, row_stride: i32, value: x_value_t) {\n" + << " let index = row * row_stride + col;\n" + << " result[index] = value;\n" + << "}\n"; + + // Main function body + shader.MainFunctionBody() + << " let gindex = i32(global_idx);\n" + << " let lindex = i32(local_idx);\n" + << " const wg = " << WG << ";\n" + << " let row = gindex / wg;\n" + << " let cols = uniforms.packedCols;\n" + << " let row_stride : i32 = uniforms.packedCols;\n" + + // Find the row's max value + << threadMaxDecl + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = getValue(row, col, row_stride);\n" + << " threadMax = max(threadMax, value);\n" + << " }\n" + << " if (lindex < cols) {\n" + << " threadShared[lindex] = threadMax;\n" + << " }\n" + << " workgroupBarrier();\n" + + // Reduce to find the max value + << " var reduceSize = min(cols, wg);\n" + << " for (var currSize = reduceSize >> 1; currSize > 0; currSize = reduceSize >> 1) {\n" + << " reduceSize = currSize + (reduceSize & 1);\n" + << " if (lindex < currSize) {\n" + << " threadShared[lindex] = max(threadShared[lindex], threadShared[lindex + reduceSize]);\n" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " rowMaxShared = x_value_t(" << MaxVector('threadShared[0]', components) << ");\n" + << " }\n" + << " workgroupBarrier();\n" + + // Find the row's sum of exponentials + << " var threadSum = x_value_t(0.0);\n" + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let subExp = exp(getValue(row, col, row_stride) - rowMaxShared);\n" + << " threadSum += subExp;\n" + << " }\n" + << " threadShared[lindex] = threadSum;\n" + << " workgroupBarrier();\n" + + // Reduce to find the sum of exponentials + << " for (var currSize = wg >> 1; currSize > 0; currSize = currSize >> 1) {\n" + << " if (lindex < currSize) {\n" + << " threadShared[lindex] = threadShared[lindex] + threadShared[lindex + currSize];\n" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " rowSumShared = x_value_t(" << SumVector("threadShared[0]", components) << ");\n" + << " }\n" + << " workgroupBarrier();\n" + + // Calculate the final value for each element in the row + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = exp(getValue(row, col, row_stride) - rowMaxShared) / rowSumShared;\n" + << " setValue(row, col, row_stride, value);\n" + << " }\n"; + + return Status::OK(); +} + +Status Softmax::ComputeInternal(ComputeContext& context) const { + const auto* input_tensor = context.Input(0); + const TensorShape& input_shape = input_tensor->Shape(); + size_t input_rank = input_shape.NumDimensions(); + + auto* output_tensor = context.Output(0, input_shape); + + // normalize axis + int64_t axis = axis < 0 ? axis_ + input_rank : axis_; + + bool is_transpose_required = axis < input_rank - 1; + TensorShape transposed_input_shape = input_shape; + Tensor transposed_input_tensor; + Tensor intermediate_output; + InlinedVector perm; + + if (is_transpose_required) { + AllocatorPtr alloc; + perm.reserve(input_rank); + for (size_t i = 0; i < input_rank; ++i) { + perm[i] = i; + } + perm[axis] = input_rank - 1; + perm[input_rank - 1] = axis; + + // allocate a temporary tensor to hold transposed input + Tensor temp_input(input_tensor->DataType(), TensorShape(transposed_input_shape), alloc); + + ORT_RETURN_IF_ERROR(Transpose::DoTranspose( perm, *input_tensor, temp_input)); + transposed_input_tensor = std::move(temp_input); + transposed_input_shape = transposed_input_tensor.Shape(); + + // Allocate memory for the intermediate output + Tensor temp_output(output_tensor->DataType(), TensorShape(transposed_input_shape), alloc); + intermediate_output = std::move(temp_output); + } else { + transposed_input_tensor = *input_tensor; + } + + + const size_t cols = transposed_input_shape[input_rank - 1]; + const size_t rows = input_shape.Size() / cols; + const size_t components = GetMaxComponents(cols); + const auto packedCols = cols / components; + + size_t WG = rows == 1 ? 256: 64; + + SoftmaxProgram program{WG}; + + + program + .CacheHint(std::to_string(components), std::to_string(WG)) + .AddInputs({*transposed_input_tensor, ProgramTensorMetadataDependency::TypeAndRank}}) + .AddOutputs({ is_transpose_required ? *intermediate_output : output_tensor}) + .SetWorkgroupSize(WG) + .SetDispatchGroupSize(rows) + .AddUniformVariables({ + {static_cast(packedCols)} + }); + + + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + + // If transpose was required, transpose the result back + if (is_transpose_required) { + Tensor transposed_output_tensor; + ORT_RETURN_IF_ERROR(Transpose::DoTranspose(perm, intermediate_output, *output_tensor)); + } + + return Status::OK(); +} +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/math/softmax.h b/onnxruntime/core/providers/webgpu/math/softmax.h new file mode 100644 index 0000000000..b8bc37a0c0 --- /dev/null +++ b/onnxruntime/core/providers/webgpu/math/softmax.h @@ -0,0 +1,52 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/cpu/math/softmax.h" +#include "core/providers/webgpu/webgpu_kernel.h" +#include "core/providers/webgpu/program.h" + +namespace onnxruntime { +namespace webgpu { + +class Softmax final : public WebGpuKernel { + public: + Softmax(const OpKernelInfo& info) : WebGpuKernel{info} { + int opset_ = info.node().SinceVersion(); + size_t axis; + Status status = info.GetAttr("axis", &axis); + + if (status.IsOK()) { + axis_ = axis; + } else { + if (opset_ < 13) { + axis_ = 1; // opset-12 and below, the default axis value is 1 + } else { + axis_ = -1; // opset-13, the default axis value is -1 + } + } + } + + Status ComputeInternal(ComputeContext& context) const override; + + private: + size_t axis_; +}; + +class SoftmaxProgram final : public Program { + public: + SoftmaxProgram(size_t axis, int wg) : Program{"Softmax"}, axis_{axis}, WG_{wg} { + } + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"packedCols", ProgramUniformVariableDataType::Int32}); + + private: + int WG; +}; + +} // namespace webgpu +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.cc b/onnxruntime/core/providers/webgpu/tensor/transpose.cc index c40ec43dd0..062500055e 100644 --- a/onnxruntime/core/providers/webgpu/tensor/transpose.cc +++ b/onnxruntime/core/providers/webgpu/tensor/transpose.cc @@ -97,6 +97,56 @@ Status TransposeProgram::GenerateShaderCode(ShaderHelper& shader) const { return Status::OK(); } +Status Transpose::DoTranspose(const gsl::span& permutations, const Tensor& input, Tensor& output) { + const auto& input_shape = input.Shape(); + int32_t rank = gsl::narrow_cast(input_shape.NumDimensions()); + + + TensorShapeVector output_dims(rank); + InlinedVector default_perm(rank); + const InlinedVector* p_perm = nullptr; + ORT_RETURN_IF_ERROR(ComputeOutputShape(input, output_dims, default_perm, p_perm)); + TensorShape output_shape(output_dims); + + InlinedVector new_shape{}; + InlinedVector new_perm{}; + SqueezeShape(input_shape.GetDims(), *p_perm, new_shape, new_perm); + const bool channels_last = new_perm == InlinedVector({2, 3, 1}); + const bool channels_first = new_perm == InlinedVector({3, 1, 2}); + const bool use_shared = (new_shape.size() == 2 && new_perm[0] > new_perm[1]) || channels_last || channels_first; + auto new_input_shape = input_shape; + + if (use_shared) { + new_input_shape = channels_last + ? TensorShape({new_shape[0], new_shape[1] * new_shape[2]}) + : channels_first + ? TensorShape({new_shape[0] * new_shape[1], new_shape[2]}) + : new_shape; + new_output_shape = TensorShape({new_input_shape[1], new_input_shape[0]}); + } + + uint32_t output_size = gsl::narrow_cast(input.Shape().Size()); + TransposeProgram program{*p_perm, use_shared}; + if (use_shared) { + program.SetWorkgroupSize(TILE_SIZE, TILE_SIZE, 1); + } + + program + .CacheHint(absl::StrJoin(*p_perm, "-")) + .AddInputs({{*input, ProgramTensorMetadataDependency::TypeAndRank, new_input_shape, 1}}) + .AddOutputs({{*output, ProgramTensorMetadataDependency::None, new_output_shape, 1}}) + .SetDispatchGroupSize(static_cast((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE), + static_cast(((new_output_shape[0] + TILE_SIZE - 1) / TILE_SIZE))) + .AddUniformVariables({ + {static_cast(output_size)}, + }); + + use_shared ? program.SetDispatchGroupSize(static_cast((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE), + static_cast(((new_output_shape[0] + TILE_SIZE - 1) / TILE_SIZE))) + : program.SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE); + return context.RunProgram(program); +} + Status Transpose::ComputeInternal(ComputeContext& context) const { const auto* input_tensor = context.Input(0); const TensorShape& input_shape = input_tensor->Shape(); diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.h b/onnxruntime/core/providers/webgpu/tensor/transpose.h index 7cf5c1fe08..3eb672d1c6 100644 --- a/onnxruntime/core/providers/webgpu/tensor/transpose.h +++ b/onnxruntime/core/providers/webgpu/tensor/transpose.h @@ -16,6 +16,8 @@ class Transpose final : public WebGpuKernel, public TransposeBase { Transpose(const OpKernelInfo& info) : WebGpuKernel{info}, TransposeBase{info} { } Status ComputeInternal(ComputeContext& context) const override; + static Status DoTranspose(const gsl::span& permutations, const Tensor& input, Tensor& output); + constexpr static uint32_t TILE_SIZE = 16; }; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index dec7e48786..04b5965177 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -625,9 +625,9 @@ std::unique_ptr RegisterKernels() { // BuildKernelCreateInfo, // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, - // BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo,