From 275b24ebdf63a587c3739652bf50c77c55bfd4e7 Mon Sep 17 00:00:00 2001 From: wp Date: Wed, 22 Apr 2026 15:47:57 +0800 Subject: [PATCH 1/3] webgpu: use naive reduction Use naive reduction when the output size of ReduceMean is far greater than reduce size. Shared reduction method may need to transpose input, which costs much time. --- onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc index db51675e81513..9fbb36bcad36f 100644 --- a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc @@ -372,7 +372,9 @@ Status ReduceKernel::ComputeInternal(ComputeContext& context) return Status::OK(); } - bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) || is_input_empty || input_tensor->Shape().NumDimensions() == 0; + bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) || + (name_ == "ReduceMean" && reduce_size <= 128 && output_size > 20000) || + is_input_empty || input_tensor->Shape().NumDimensions() == 0; if (use_naive_reduction) { ReduceNaiveProgram program(name_, reduce_op_type, keepdims_, noop_with_empty_axes_, input_axes, is_input_empty); From aadeda24130b98c12beebc6b88541676d2ddd8cb Mon Sep 17 00:00:00 2001 From: wp Date: Fri, 24 Apr 2026 16:15:21 +0800 Subject: [PATCH 2/3] Keep to use shared ReduceMean when axes are innermost --- .../webgpu/reduction/reduction_ops.cc | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc index 9fbb36bcad36f..5984e764fc7b6 100644 --- a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc @@ -372,8 +372,21 @@ Status ReduceKernel::ComputeInternal(ComputeContext& context) return Status::OK(); } + // Prefer the naive ReduceMean path when the shared path would pay extra overhead (transposing non-innermost reduce axes). + constexpr size_t kReduceMeanNaiveMaxReduceSize = 128; + constexpr size_t kReduceMeanNaiveMinOutputSize = 20000; + bool are_axes_innermost = true; + size_t axes_rank = input_axes.size(); + for (size_t i = 0; i < input_axes.size() && are_axes_innermost; ++i) { + if (input_axes[axes_rank - 1 - i] != rank - 1 - i) { + are_axes_innermost = false; + break; + } + } bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) || - (name_ == "ReduceMean" && reduce_size <= 128 && output_size > 20000) || + (name_ == "ReduceMean" && !are_axes_innermost && + reduce_size <= kReduceMeanNaiveMaxReduceSize && + output_size > kReduceMeanNaiveMinOutputSize) || is_input_empty || input_tensor->Shape().NumDimensions() == 0; if (use_naive_reduction) { @@ -397,14 +410,6 @@ Status ReduceKernel::ComputeInternal(ComputeContext& context) return context.RunProgram(program); } else { - bool are_axes_innermost = true; - size_t axes_rank = input_axes.size(); - for (size_t i = 0; i < input_axes.size() && are_axes_innermost; ++i) { - if (input_axes[axes_rank - 1 - i] != rank - 1 - i) { - are_axes_innermost = false; - break; - } - } Tensor input_transpose; if (!are_axes_innermost) { InlinedVector perm; From 2fcb28fafd033bfd3c4b23f7cd0347de1f819069 Mon Sep 17 00:00:00 2001 From: wp Date: Thu, 30 Apr 2026 16:39:49 +0800 Subject: [PATCH 3/3] Address jiajia's comment --- .../core/providers/webgpu/reduction/reduction_ops.cc | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc index 5984e764fc7b6..f88f3538a9171 100644 --- a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc @@ -372,9 +372,9 @@ Status ReduceKernel::ComputeInternal(ComputeContext& context) return Status::OK(); } - // Prefer the naive ReduceMean path when the shared path would pay extra overhead (transposing non-innermost reduce axes). - constexpr size_t kReduceMeanNaiveMaxReduceSize = 128; - constexpr size_t kReduceMeanNaiveMinOutputSize = 20000; + // Prefer the naive Reduce path when the shared path would pay extra overhead (transposing non-innermost reduce axes). + constexpr size_t kReduceNaiveMaxReduceSize = 128; + constexpr size_t kReduceNaiveMinOutputSize = 20000; bool are_axes_innermost = true; size_t axes_rank = input_axes.size(); for (size_t i = 0; i < input_axes.size() && are_axes_innermost; ++i) { @@ -384,9 +384,8 @@ Status ReduceKernel::ComputeInternal(ComputeContext& context) } } bool use_naive_reduction = name_ == "ArgMin" || name_ == "ArgMax" || (reduce_size < 32 && output_size > 1024) || - (name_ == "ReduceMean" && !are_axes_innermost && - reduce_size <= kReduceMeanNaiveMaxReduceSize && - output_size > kReduceMeanNaiveMinOutputSize) || + (!are_axes_innermost && reduce_size <= kReduceNaiveMaxReduceSize && + output_size > kReduceNaiveMinOutputSize) || is_input_empty || input_tensor->Shape().NumDimensions() == 0; if (use_naive_reduction) {