From 91f635ce5d17ae29835f5221e543db79241d3b1f Mon Sep 17 00:00:00 2001 From: vraspar Date: Mon, 3 Feb 2025 14:25:01 -0800 Subject: [PATCH] Refactor Softmax implementation for improved readability and consistency in variable naming --- .../core/providers/webgpu/math/softmax.cc | 48 ++++++++++--------- .../core/providers/webgpu/math/softmax.h | 6 ++- .../core/providers/webgpu/tensor/transpose.cc | 40 +--------------- .../core/providers/webgpu/tensor/transpose.h | 2 +- .../test/providers/cpu/math/softmax_test.cc | 2 +- 5 files changed, 33 insertions(+), 65 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc index 1760acae95..4685d8eef0 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -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 rowMaxShared : x_value_t;\n" - << "var rowSumShared : x_value_t;\n" - << "var threadShared : array;\n"; + << "var row_max_shared : x_value_t;\n" + << "var row_sum_shared : x_value_t;\n" + << "var thread_shared : array;\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(components)}}) diff --git a/onnxruntime/core/providers/webgpu/math/softmax.h b/onnxruntime/core/providers/webgpu/math/softmax.h index 5eb6bd0ccd..cc97611dcb 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.h +++ b/onnxruntime/core/providers/webgpu/math/softmax.h @@ -37,7 +37,8 @@ class Softmax final : public WebGpuKernel { class SoftmaxProgram final : public Program { 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 { WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"packedCols", ProgramUniformVariableDataType::Int32}); private: - uint32_t WG; + uint32_t wg_; + bool is_fp32_; }; } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.cc b/onnxruntime/core/providers/webgpu/tensor/transpose.cc index c06a742239..97d7332507 100644 --- a/onnxruntime/core/providers/webgpu/tensor/transpose.cc +++ b/onnxruntime/core/providers/webgpu/tensor/transpose.cc @@ -97,7 +97,7 @@ Status TransposeProgram::GenerateShaderCode(ShaderHelper& shader) const { return Status::OK(); } -Status Transpose::DoTranspose(onnxruntime::webgpu::ComputeContext& context, const gsl::span& permutations, const Tensor& input, Tensor& output) { +Status Transpose::DoTranspose(onnxruntime::webgpu::ComputeContext& context, gsl::span 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(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 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; - 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(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((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); + return DoTranspose(context, *p_perm, *input_tensor, *output_tensor); } } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/tensor/transpose.h b/onnxruntime/core/providers/webgpu/tensor/transpose.h index 81706dde33..b62a419fa1 100644 --- a/onnxruntime/core/providers/webgpu/tensor/transpose.h +++ b/onnxruntime/core/providers/webgpu/tensor/transpose.h @@ -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& permutations, const Tensor& input, Tensor& output); + static Status DoTranspose(onnxruntime::webgpu::ComputeContext& context, gsl::span permutations, const Tensor& input, Tensor& output); constexpr static uint32_t TILE_SIZE = 16; }; diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 3808d62a10..1c6375ebdb 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -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 ); }