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

Refactor Softmax implementation for improved readability and consistency in variable naming

Browse source
This commit is contained in:
vraspar 2025-02-03 14:25:01 -08:00
parent 2d8b47de27
commit 91f635ce5d
5 changed files with 33 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
onnxruntime/core/providers/webgpu/math/softmax.cc
View file

@ -81,13 +81,13 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
shader.AddOutput("result", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias);
int components = input.NumComponents();
std::string threadMaxDecl = input.ElementType() == "f32" ? "var threadMax = x_value_t(-3.402823e+38f);\n" : "var threadMax = x_value_t(-65504.0h);\n";
std::string threadMaxDecl = is_fp32_ ? "var thread_max = x_value_t(-3.402823e+38f);\n" : "var thread_max = 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";
<< "var<workgroup> row_max_shared : x_value_t;\n"
<< "var<workgroup> row_sum_shared : x_value_t;\n"
<< "var<workgroup> thread_shared : array<x_value_t, " << wg_ << ">;\n";
// Define helper functions to get and set values
shader.AdditionalImplementation()
@ -104,7 +104,7 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
shader.MainFunctionBody()
<< " let gindex = i32(global_idx);\n"
<< " let lindex = i32(local_idx);\n"
<< " const wg = " << WG << ";\n"
<< " const wg = " << wg_ << ";\n"
<< " let row = gindex / wg;\n"
<< " let cols = uniforms.packedCols;\n"
<< " let row_stride : i32 = uniforms.packedCols;\n"
@ -113,51 +113,51 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
<< threadMaxDecl
<< " for (var col = lindex; col < cols; col += wg) {\n"
<< " let value = getValue(row, col, row_stride);\n"
<< " threadMax = max(threadMax, value);\n"
<< " thread_max = max(thread_max, value);\n"
<< " }\n"
<< " if (lindex < cols) {\n"
<< " threadShared[lindex] = threadMax;\n"
<< " thread_shared[lindex] = thread_max;\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"
<< " var reduce_size = min(cols, wg);\n"
<< " for (var curr_size = reduce_size >> 1; curr_size > 0; curr_size = reduce_size >> 1) {\n"
<< " reduce_size = curr_size + (reduce_size & 1);\n"
<< " if (lindex < curr_size) {\n"
<< " thread_shared[lindex] = max(thread_shared[lindex], thread_shared[lindex + reduce_size]);\n"
<< " }\n"
<< " workgroupBarrier();\n"
<< " }\n"
<< " if (lindex == 0) {\n"
<< " rowMaxShared = x_value_t(" << MaxVector("threadShared[0]", components) << ");\n"
<< " row_max_shared = x_value_t(" << MaxVector("thread_shared[0]", components) << ");\n"
<< " }\n"
<< " workgroupBarrier();\n"
// Find the row's sum of exponentials
<< " var threadSum = x_value_t(0.0);\n"
<< " var thread_sum = 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"
<< " let sub_exp = exp(getValue(row, col, row_stride) - row_max_shared);\n"
<< " thread_sum += sub_exp;\n"
<< " }\n"
<< " threadShared[lindex] = threadSum;\n"
<< " thread_shared[lindex] = thread_sum;\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"
<< " for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n"
<< " if (lindex < curr_size) {\n"
<< " thread_shared[lindex] = thread_shared[lindex] + thread_shared[lindex + curr_size];\n"
<< " }\n"
<< " workgroupBarrier();\n"
<< " }\n"
<< " if (lindex == 0) {\n"
<< " rowSumShared = x_value_t(" << SumVector("threadShared[0]", components) << ");\n"
<< " row_sum_shared = x_value_t(" << SumVector("thread_shared[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"
<< " let value = exp(getValue(row, col, row_stride) - row_max_shared) / row_sum_shared;\n"
<< " setValue(row, col, row_stride, value);\n"
<< " }\n";
@ -200,8 +200,10 @@ Status Softmax::ComputeInternal(ComputeContext& context) const {
const int64_t components = GetMaxComponents(cols);
const auto packedCols = cols / components;
uint32_t WG = rows == 1 ? 256 : 64;
// check input tensor element type is float
const bool is_fp32 = input_tensor->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
SoftmaxProgram program{WG};
SoftmaxProgram program{WG, is_fp32};
if (is_transpose_required) {
program
.AddInputs({{&transposed_input_tensor, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(components)}})

6
onnxruntime/core/providers/webgpu/math/softmax.h
View file

@ -37,7 +37,8 @@ class Softmax final : public WebGpuKernel {
class SoftmaxProgram final : public Program<SoftmaxProgram> {
public:
SoftmaxProgram(uint32_t wg) : Program{"Softmax"}, WG{wg} {
SoftmaxProgram(uint32_t wg, bool is_fp32)
: Program{"Softmax"}, wg_{wg}, is_fp32_{is_fp32} {
}
Status GenerateShaderCode(ShaderHelper& sh) const override;
@ -45,7 +46,8 @@ class SoftmaxProgram final : public Program<SoftmaxProgram> {
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"packedCols", ProgramUniformVariableDataType::Int32});
private:
uint32_t WG;
uint32_t wg_;
bool is_fp32_;
};
} // namespace webgpu

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

@ -97,7 +97,7 @@ Status TransposeProgram::GenerateShaderCode(ShaderHelper& shader) const {
return Status::OK();
}
Status Transpose::DoTranspose(onnxruntime::webgpu::ComputeContext& context, const gsl::span<const size_t>& permutations, const Tensor& input, Tensor& output) {
Status Transpose::DoTranspose(onnxruntime::webgpu::ComputeContext& context, gsl::span<const size_t> permutations, const Tensor& input, Tensor& output) {
const auto& input_shape = input.Shape();
const auto& input_dims = input_shape.GetDims();
int32_t rank = gsl::narrow_cast<int32_t>(input_shape.NumDimensions());
@ -162,43 +162,7 @@ Status Transpose::ComputeInternal(ComputeContext& context) const {
TensorShape output_shape(output_dims);
auto* output_tensor = context.Output(0, output_shape);
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;
TensorShape new_output_shape(output_dims);
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_tensor->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_tensor, ProgramTensorMetadataDependency::TypeAndRank, new_input_shape, 1}})
.AddOutputs({{output_tensor, 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);
return DoTranspose(context, *p_perm, *input_tensor, *output_tensor);
}
} // namespace webgpu

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

@ -16,7 +16,7 @@ class Transpose final : public WebGpuKernel, public TransposeBase {
Transpose(const OpKernelInfo& info) : WebGpuKernel{info}, TransposeBase{info} {
}
Status ComputeInternal(ComputeContext& context) const override;
static Status DoTranspose(onnxruntime::webgpu::ComputeContext& context, const gsl::span<const size_t>& permutations, const Tensor& input, Tensor& output);
static Status DoTranspose(onnxruntime::webgpu::ComputeContext& context, gsl::span<const size_t> permutations, const Tensor& input, Tensor& output);
constexpr static uint32_t TILE_SIZE = 16;
};

2
onnxruntime/test/providers/cpu/math/softmax_test.cc
View file

@ -377,7 +377,7 @@ TEST(SoftmaxOperator, DimWithZero) {
RunTest(x_vals, expected_vals, dimensions, /*opset*/ -1, /*axis*/ 0,
{kTensorrtExecutionProvider,
kNnapiExecutionProvider, // NNAPI softmax does not support empty input
kWebGpuExecutionProvider, // WebGPU does not dim 0
kWebGpuExecutionProvider, // WebGPU does not support dim 0
kQnnExecutionProvider} // QNN doesn't support dim 0
);
}