Skip to content

Fix CUDA ReduceSum erroring out on empty tensors with explicit axes #28353

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Justin Chu (justinchuby) wants to merge 16 commits into from
Closed

Fix CUDA ReduceSum erroring out on empty tensors with explicit axes #28353

Show file tree
Hide file tree
Changes from 7 commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
3e0e208
Fix CUDA ReduceSum crash on empty tensors with explicit axes
justinchuby May 4, 2026
6732857
Fix CUDA reduction ops for empty tensor inputs per ONNX spec
justinchuby May 4, 2026
2c2003f
Improve empty tensor reduction test coverage
justinchuby May 5, 2026
657664e
Address review: comprehensive empty tensor reduction fixes
justinchuby May 5, 2026
0b7616e
Fix CI: MLFloat16 type conversion and test corrections
justinchuby May 6, 2026
1cc286e
Fix clang-format: trailing whitespace alignment in macro
justinchuby May 6, 2026
b82b4f1
Remove bool ReduceMax/Min empty tests (ONNX schema rejects bool type)
justinchuby May 6, 2026
ab891fc
Address review: replace std::set with linear scan for axis lookup
justinchuby May 6, 2026
7e84448
Address review: consistent CPU/CUDA empty-set behavior and extended t…
justinchuby May 6, 2026
235ac37
Exclude unsupported EPs from empty tensor reduction tests
justinchuby May 7, 2026
5b8a5a4
Narrow empty tensor EP exclusions to only confirmed failures
justinchuby May 7, 2026
db8739b
Fix WebGPU reduction kernels for empty tensor inputs
justinchuby May 7, 2026
f00d512
Fix lint and test failures for empty tensor reductions
justinchuby May 7, 2026
bb43216
Fix WebGPU empty tensor identity values: use division-by-zero for inf…
justinchuby May 8, 2026
90c5b9f
Apply suggestions from code review
justinchuby May 8, 2026
5dbaaf9
Fix WebGPU shader: use bitcast for infinity instead of division by zero
justinchuby May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 30 additions & 13 deletions onnxruntime/core/providers/cpu/reduction/reduction_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
#include "core/common/narrow.h"
#include "core/common/span_utils.h"
#include "core/providers/common.h"

#include <set>

Check warning on line 11 in onnxruntime/core/providers/cpu/reduction/reduction_ops.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Found C++ system header after other header. Should be: reduction_ops.h, c system, c++ system, other.  [build/include_order] [4]

Raw Output:
onnxruntime/core/providers/cpu/reduction/reduction_ops.cc:11:  Found C++ system header after other header. Should be: reduction_ops.h, c system, c++ system, other.  [build/include_order] [4]
// TODO: fix the warnings
#if defined(_MSC_VER) && !defined(__clang__)
#pragma warning(disable : 26451)
Expand Down Expand Up @@ -879,25 +881,35 @@
return false;
}

// input is an empty set
// input is an empty set — resolve effective axes
std::vector<int64_t> input_axes;
if (ctx->InputCount() == 2) {
ORT_ENFORCE(axes.empty(), "Axes input and attribute should not both be present for reduction.");
// second input holds the axes.
const Tensor* axes_tensor = ctx->Input<Tensor>(1);
auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->Data<int64_t>();
input_axes.insert(input_axes.begin(), data, data + nDims);
if (axes_tensor != nullptr) {
ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor.");
auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->Data<int64_t>();
input_axes.insert(input_axes.begin(), data, data + nDims);
}
// axes_tensor == nullptr means no axes provided → reduce all dims
} else {
input_axes.resize(axes.size());
std::copy(axes.begin(), axes.end(), input_axes.begin());
}

gsl::span<const int64_t> shape_dims = input_shape.GetDims();
const int64_t input_shape_size = narrow<int64_t>(shape_dims.size());
// Normalize negative axes
const int64_t rank = narrow<int64_t>(input_shape.NumDimensions());
for (auto& axis : input_axes) {
axis = HandleNegativeAxis(axis, rank);
}

// Build reduced output shape
std::set<int64_t> reduced_axes(input_axes.begin(), input_axes.end());
Comment thread
justinchuby marked this conversation as resolved.
Outdated
TensorShapeVector output_shape_vector;
for (int64_t i = 0; i < input_shape_size; ++i) {
if (input_axes.empty() || std::find(input_axes.begin(), input_axes.end(), i) != input_axes.end()) {
for (int64_t i = 0; i < rank; ++i) {
bool is_reduced = reduced_axes.empty() || reduced_axes.count(i) > 0;
if (is_reduced) {
if (keepdims) {
output_shape_vector.push_back(1);
}
Expand Down Expand Up @@ -968,17 +980,22 @@
void CommonReduce1Loop(OpKernelContext* ctx,
const gsl::span<const int64_t>& axes_, int64_t keepdims_,
bool noop_with_empty_axes) {
if (check_and_reduce_empty_set_input<AGG>(ctx, axes_, keepdims_ != 0)) {
return;
}

// Resolve effective axes first (from input tensor or attribute).
TensorShapeVector tmp_axes;
auto effective_axes = GetEffectiveAxes(ctx, axes_, tmp_axes);

// noop_with_empty_axes takes precedence: if no axes, copy input as-is
// (applying element-wise transforms if any). This applies even to empty
// tensors — a {1,0} input should stay {1,0}, not be reduced to scalar.
if (effective_axes.empty() && noop_with_empty_axes) {
ApplyNoopEmptyAxesElementwise<AGG>(ctx);
return;
}

if (check_and_reduce_empty_set_input<AGG>(ctx, axes_, keepdims_ != 0)) {
return;
}

FastReduceKind fast_kind;
TensorShapeVector fast_shape;
TensorShapeVector output_shape;
Expand Down
12 changes: 8 additions & 4 deletions onnxruntime/core/providers/cpu/reduction/reduction_ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -386,8 +386,10 @@ class ReduceAggregatorMax : public ReduceAggregator<T> {
inline void update(const T& v) { this->accumulator_ = v > this->accumulator_ ? v : this->accumulator_; }

static void fill_for_empty_set(Tensor& output) {
if constexpr (std::is_same_v<bool, T>) { /* bool specific impl */
ORT_NOT_IMPLEMENTED();
if constexpr (std::is_same_v<bool, T>) {
// ONNX spec: ReduceMax on empty bool set → false (boolean zero/identity)
auto* data = output.MutableData<bool>();
std::fill(data, data + output.Shape().Size(), false);
} else {
EigenMap<T>(output).array() = -std::numeric_limits<T>::infinity();
}
Expand Down Expand Up @@ -596,8 +598,10 @@ class ReduceAggregatorMin : public ReduceAggregator<T, T> {
inline void update(const T& v) { this->accumulator_ = v < this->accumulator_ ? v : this->accumulator_; }

static void fill_for_empty_set(Tensor& output) {
if constexpr (std::is_same_v<bool, T>) { /* bool specific impl */
ORT_NOT_IMPLEMENTED();
if constexpr (std::is_same_v<bool, T>) {
// ONNX spec: ReduceMin on empty bool set → true (boolean max/identity)
auto* data = output.MutableData<bool>();
std::fill(data, data + output.Shape().Size(), true);
} else {
EigenMap<T>(output).array() = std::numeric_limits<T>::infinity();
}
Expand Down
76 changes: 64 additions & 12 deletions onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -298,22 +298,17 @@
prepare_reduce_metadata.output_dims = input_shape.AsShapeVector();
for (auto axis : axes) {
axis = HandleNegativeAxis(axis, rank);
ORT_ENFORCE(input_dims[axis] != 0,
Comment thread
justinchuby marked this conversation as resolved.
"Can't reduce on dim with value of 0 if 'keepdims' is false. "
"Invalid output shape would be produced. input_shape:",
input_shape);
prepare_reduce_metadata.output_dims[axis] = 1;
reduced[axis] = true;
}
} else {
// no axes provided (i.e.) default axes => reduce on all dims
// Each reduced dim becomes 1 (even if the original dim was 0 — the
// reduction collapses the axis regardless of its size).
prepare_reduce_metadata.output_dims.reserve(input_dims.size());
for (auto dim : input_dims) {
ORT_ENFORCE(keepdims || dim != 0,
"Can't reduce on dim with value of 0 if 'keepdims' is false. "
"Invalid output shape would be produced. input_shape:",
input_shape);
prepare_reduce_metadata.output_dims.push_back(dim == 0 ? 0 : 1);
for (size_t i = 0; i < input_dims.size(); ++i) {
prepare_reduce_metadata.output_dims.push_back(1);
reduced[i] = true;
}
}

Expand Down Expand Up @@ -377,7 +372,37 @@
auto& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn;
// special case when there is a dim value of 0 in the shape.
if (input_count == 0) {
assert(output.Shape().Size() == 0);
// Empty input reduction: output may still be non-empty when only some
// axes are reduced. Per ONNX spec, fill with the reduction identity.
if (output_count > 0) {
// For types that don't support std::numeric_limits natively (MLFloat16,
// BFloat16), use float intermediary and convert via CudaT.
if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_AVG) {

@yuslepukhin Dmitri Smirnov (yuslepukhin) May 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

CPU and CUDA now disagree on empty-set ReduceMean, and the new generic tests encode only the CUDA behavior. The CPU empty-input fast path still routes through reduction_ops.cc:924, ReduceAggregatorMean still inherits ReduceAggregatorSum at reduction_ops.h:307, and the inherited empty-set fill is still 0 at reduction_ops.h:224.

CUDA was changed to emit quiet_NaN() for empty ReduceMean at reduction_ops.cc:380, reduction_ops.cc:381, reduction_ops.cc:816, and reduction_ops.cc:817, and the new test expects NaN at reduction_ops_test.cc:6512. ONNX says empty-set ReduceMean is undefined, so choosing NaN is defensible, but this PR has not made ORT internally consistent. As written, the “tests for all EPs” claim is not satisfactorily addressed since it should fail on CPU.

This needs to be consistent and the test should cover it. Please, consider using zeros from historical perspective.

@justinchuby Justin Chu (justinchuby) May 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed in 7e84448. CPU and CUDA are now consistent — both fill ReduceMean empty-set with 0:

  • CPU: ReduceAggregatorMean inherits ReduceAggregatorSum::fill_for_empty_set (fills 0) at reduction_ops.h:307/224
  • CUDA: CUDNN_REDUCE_TENSOR_AVG branch uses cudaMemsetAsync(..., 0, ...) at reduction_ops.cc:382
  • Macro path: else branch at line 821 uses cudaMemsetAsync(..., 0, ...) with comment 'Sum, SumSquare, Mean, L1, L2, Amax: identity is 0'

Test ReduceMean_EmptyTensor_ExplicitAxis expects 0.0f and runs on all EPs. Added ReduceMean_EmptyTensor_DefaultAxes for reduce-all coverage. CUDA comment updated to explain the consistency rationale.

// ReduceMean on empty set is undefined (0/0). Fill with 0.
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0,
output.SizeInBytes(), stream));
} else if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MUL) {
// ReduceProd identity is 1.
CudaT one_val = ToCudaType<T>::FromFloat(1.0f);
std::vector<CudaT> ones(output_count, one_val);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output.MutableDataRaw(), ones.data(),
output.SizeInBytes(), cudaMemcpyHostToDevice, stream));
} else if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) {
CudaT inf_val = ToCudaType<T>::FromFloat(std::numeric_limits<float>::infinity());
std::vector<CudaT> vals(output_count, inf_val);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output.MutableDataRaw(), vals.data(),
output.SizeInBytes(), cudaMemcpyHostToDevice, stream));
} else if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) {
CudaT neg_inf_val = ToCudaType<T>::FromFloat(-std::numeric_limits<float>::infinity());
std::vector<CudaT> vals(output_count, neg_inf_val);
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output.MutableDataRaw(), vals.data(),
output.SizeInBytes(), cudaMemcpyHostToDevice, stream));
} else {
// Sum, SumSquare, L1, L2: identity is 0.
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(output.MutableDataRaw(), 0,
output.SizeInBytes(), stream));
}
}
return Status::OK();
}

Expand Down Expand Up @@ -770,7 +795,34 @@
auto& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; \
\
if (input_count == 0) { \
assert(Y->Shape().Size() == 0); \
/* Empty input reduction: fill output with the reduction identity. */ \
/* ONNX spec: Sum→0, Prod→1, Min→+inf, Max→-inf, Mean→0. */ \
if (Y->Shape().Size() > 0) { \
typedef typename ToCudaType<T>::MappedType CudaT_local; \
if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MUL) { \
/* Identity is 1 for product */ \
CudaT_local one_val = ToCudaType<T>::FromFloat(1.0f); \
std::vector<CudaT_local> ones(Y->Shape().Size(), one_val); \
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->MutableDataRaw(), ones.data(), \
Y->SizeInBytes(), cudaMemcpyHostToDevice, Stream(ctx))); \
} else if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MIN) { \
/* ONNX spec: "yields plus infinity (if supported) or max value" */ \
CudaT_local inf_val = ToCudaType<T>::FromFloat(std::numeric_limits<float>::infinity()); \
std::vector<CudaT_local> vals(Y->Shape().Size(), inf_val); \
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->MutableDataRaw(), vals.data(), \
Y->SizeInBytes(), cudaMemcpyHostToDevice, Stream(ctx))); \
} else if (cudnn_reduce_op == CUDNN_REDUCE_TENSOR_MAX) { \
/* ONNX spec: "yields minus infinity (if supported) or minimum value" */ \
CudaT_local neg_inf_val = ToCudaType<T>::FromFloat(-std::numeric_limits<float>::infinity()); \
std::vector<CudaT_local> vals(Y->Shape().Size(), neg_inf_val); \

Check warning on line 817 in onnxruntime/core/providers/cuda/reduction/reduction_ops.cc

View workflow job for this annotation

GitHub Actions / Optional Lint C++ 

[cpplint] reported by reviewdog 🐶
 Add #include <vector> for vector<>  [build/include_what_you_use] [4]

Raw Output:
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc:817:  Add #include <vector> for vector<>  [build/include_what_you_use] [4]
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->MutableDataRaw(), vals.data(), \
Y->SizeInBytes(), cudaMemcpyHostToDevice, Stream(ctx))); \
} else { \
/* Sum, SumSquare, Mean, L1, L2, Amax: identity is 0 */ \
CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, \
Y->SizeInBytes(), Stream(ctx))); \
} \
} \
return Status::OK(); \
} \
\
Expand Down
Loading
Loading