Skip to content
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

Jump to bottom

Propagate NaNs in the CPU min and max operators #21492

Merged
merged 6 commits into from
Jul 29, 2024
Merged

Propagate NaNs in the CPU min and max operators #21492

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
7b89d01
Add tests to reproduce NaN propagation failure for min and max
adamreeve Jul 24, 2024
359b3cc
Use PropagateNaN behaviour for Eigen min and max methods
adamreeve Jul 24, 2024
77d403a
Also fix min and max with float16
adamreeve Jul 24, 2024
f4a6d48
Fix formatting
adamreeve Jul 28, 2024
51ed042
Test broadcasting scalar NaN
adamreeve Jul 28, 2024
4abbd41
Revert MLFloat16 changes (except checker fix and test name typos)
adamreeve Jul 29, 2024
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
34 changes: 20 additions & 14 deletions onnxruntime/core/providers/cpu/math/element_wise_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -705,7 +705,7 @@ Status Min_6<float>::Compute(OpKernelContext* ctx) const {
for (int index = 1; index < inputCount; index++) {
auto& data_n = *ctx->Input<Tensor>(index);
ORT_ENFORCE(data_n.Shape() == shape, "All inputs must have the same shape");
min = min.array().min(EigenMap<float>(data_n).array());
min = min.array().template min<Eigen::PropagateNaN>(EigenMap<float>(data_n).array());
}

return Status::OK();
Expand All @@ -721,15 +721,16 @@ struct Min_8::ComputeImpl {
ProcessBroadcastSpanFuncs funcs{
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput1<T>().array().min(per_iter_bh.ScalarInput0<T>());
per_iter_bh.EigenInput1<T>().array().template min<Eigen::PropagateNaN>(per_iter_bh.ScalarInput0<T>());
},
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput0<T>().array().min(per_iter_bh.ScalarInput1<T>());
per_iter_bh.EigenInput0<T>().array().template min<Eigen::PropagateNaN>(per_iter_bh.ScalarInput1<T>());
},
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput0<T>().array().min(per_iter_bh.EigenInput1<T>().array());
per_iter_bh.EigenInput0<T>().array().template min<Eigen::PropagateNaN>(
per_iter_bh.EigenInput1<T>().array());
}};

int input_count = inst.Node().InputArgCount().front();
Expand All @@ -756,9 +757,11 @@ static Status MinMaxMLFloat16(const OpKernel& inst, OpKernelContext* context) {
EigenVectorArrayMap<Eigen::half> output_vec_map(output, num_elements);

if (is_min) {
output_vec_map = input_1_vec_map.min(static_cast<Eigen::half>(per_iter_bh.ScalarInput0<MLFloat16>()));
output_vec_map = input_1_vec_map.template min<Eigen::PropagateNaN>(
static_cast<Eigen::half>(per_iter_bh.ScalarInput0<MLFloat16>()));
} else {
output_vec_map = input_1_vec_map.max(static_cast<Eigen::half>(per_iter_bh.ScalarInput0<MLFloat16>()));
output_vec_map = input_1_vec_map.template max<Eigen::PropagateNaN>(
static_cast<Eigen::half>(per_iter_bh.ScalarInput0<MLFloat16>()));
}
},
[](BroadcastHelper& per_iter_bh) {
Expand All @@ -771,9 +774,11 @@ static Status MinMaxMLFloat16(const OpKernel& inst, OpKernelContext* context) {
EigenVectorArrayMap<Eigen::half> output_vec_map(output, num_elements);

if (is_min) {
output_vec_map = input_0_vec_map.min(static_cast<Eigen::half>(per_iter_bh.ScalarInput1<MLFloat16>()));
output_vec_map = input_0_vec_map.template min<Eigen::PropagateNaN>(
static_cast<Eigen::half>(per_iter_bh.ScalarInput1<MLFloat16>()));
} else {
output_vec_map = input_0_vec_map.max(static_cast<Eigen::half>(per_iter_bh.ScalarInput1<MLFloat16>()));
output_vec_map = input_0_vec_map.template max<Eigen::PropagateNaN>(
static_cast<Eigen::half>(per_iter_bh.ScalarInput1<MLFloat16>()));
}
},
[](BroadcastHelper& per_iter_bh) {
Expand All @@ -789,9 +794,9 @@ static Status MinMaxMLFloat16(const OpKernel& inst, OpKernelContext* context) {
EigenVectorArrayMap<Eigen::half> output_vec_map(output, num_elements);

if (is_min) {
output_vec_map = input_0_vec_map.min(input_1_vec_map);
output_vec_map = input_0_vec_map.template min<Eigen::PropagateNaN>(input_1_vec_map);
} else {
output_vec_map = input_0_vec_map.max(input_1_vec_map);
output_vec_map = input_0_vec_map.template max<Eigen::PropagateNaN>(input_1_vec_map);
}
}};

Expand Down Expand Up @@ -827,7 +832,7 @@ Status Max_6<float>::Compute(OpKernelContext* ctx) const {
for (int index = 1; index < inputCount; index++) {
auto& data_n = *ctx->Input<Tensor>(index);
ORT_ENFORCE(data_n.Shape() == shape, "All inputs must have the same shape");
max = max.array().max(EigenMap<float>(data_n).array());
max = max.array().template max<Eigen::PropagateNaN>(EigenMap<float>(data_n).array());
}

return Status::OK();
Expand All @@ -843,15 +848,16 @@ struct Max_8::ComputeImpl {
ProcessBroadcastSpanFuncs funcs{
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput1<T>().array().max(per_iter_bh.ScalarInput0<T>());
per_iter_bh.EigenInput1<T>().array().template max<Eigen::PropagateNaN>(per_iter_bh.ScalarInput0<T>());
},
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput0<T>().array().max(per_iter_bh.ScalarInput1<T>());
per_iter_bh.EigenInput0<T>().array().template max<Eigen::PropagateNaN>(per_iter_bh.ScalarInput1<T>());
},
[](BroadcastHelper& per_iter_bh) {
per_iter_bh.OutputEigen<T>() =
per_iter_bh.EigenInput0<T>().array().max(per_iter_bh.EigenInput1<T>().array());
per_iter_bh.EigenInput0<T>().array().template max<Eigen::PropagateNaN>(
per_iter_bh.EigenInput1<T>().array());
}};

int input_count = inst.Node().InputArgCount().front();
Expand Down
2 changes: 1 addition & 1 deletion onnxruntime/test/providers/checkers.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,7 +411,7 @@ struct TensorCheck<MLFloat16> {

for (int64_t i = 0; i < size; ++i) {
if (std::isnan(f_expected[i])) {
EXPECT_TRUE(std::isnan(f_expected[i])) << "Expected NaN. i:" << i;
EXPECT_TRUE(std::isnan(f_actual[i])) << "Expected NaN. i:" << i;
} else if (std::isinf(f_expected[i])) { // Test infinity for equality
EXPECT_EQ(f_expected[i], f_actual[i]) << "Expected infinity. i:" << i;
} else {
Expand Down
Loading
Loading