saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add Softmax kernel and transpose utility function for WebGPU execution provider

Browse source
This commit is contained in:
vraspar 2025-01-20 17:26:56 -08:00
parent e20b529a32
commit e7e373713e
5 changed files with 348 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

241
onnxruntime/core/providers/webgpu/math/softmax.cc Normal file
View file

@ -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<workgroup> rowMaxShared : x_value_t;\n"
<< "var<workgroup> rowSumShared : x_value_t;\n"
<< "var<workgroup> threadShared : array<x_value_t, " << WG << ">;\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<size_t> 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<int32_t>(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

52
onnxruntime/core/providers/webgpu/math/softmax.h Normal file
View file

@ -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<size_t>("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<SoftmaxProgram> {
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

50
onnxruntime/core/providers/webgpu/tensor/transpose.cc
View file

@ -97,6 +97,56 @@ Status TransposeProgram::GenerateShaderCode(ShaderHelper& shader) const {
return Status::OK();
}
Status Transpose::DoTranspose(const gsl::span<const size_t>& permutations, const Tensor& input, Tensor& output) {
const auto& input_shape = input.Shape();
int32_t rank = gsl::narrow_cast<int32_t>(input_shape.NumDimensions());
TensorShapeVector output_dims(rank);
InlinedVector<size_t> default_perm(rank);
const InlinedVector<size_t>* p_perm = nullptr;
ORT_RETURN_IF_ERROR(ComputeOutputShape(input, output_dims, default_perm, p_perm));
TensorShape output_shape(output_dims);
InlinedVector<int64_t> new_shape{};
InlinedVector<int64_t> new_perm{};
SqueezeShape(input_shape.GetDims(), *p_perm, new_shape, new_perm);
const bool channels_last = new_perm == InlinedVector<int64_t>({2, 3, 1});
const bool channels_first = new_perm == InlinedVector<int64_t>({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<int32_t>(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<uint32_t>((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE),
static_cast<uint32_t>(((new_output_shape[0] + TILE_SIZE - 1) / TILE_SIZE)))
.AddUniformVariables({
{static_cast<uint32_t>(output_size)},
});
use_shared ? program.SetDispatchGroupSize(static_cast<uint32_t>((new_output_shape[1] + TILE_SIZE - 1) / TILE_SIZE),
static_cast<uint32_t>(((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();

2
onnxruntime/core/providers/webgpu/tensor/transpose.h
View file

@ -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<const size_t>& permutations, const Tensor& input, Tensor& output);
constexpr static uint32_t TILE_SIZE = 16;
};

6
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc
View file

@ -625,9 +625,9 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, float, ArgMin)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, float, ArgMin)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, Softmax)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Softmax)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 3, Concat)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 4, 10, Concat)>,