Fix Min/Max CPU kernels for float16 type (#6205)

This commit is contained in:
Hariharan Seshadri 2021-01-08 13:02:52 +05:30 committed by GitHub
parent a72fcbd5fc
commit 7fc827a8a1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 161 additions and 6 deletions

View file

@ -569,10 +569,80 @@ struct Min_8::ComputeImpl {
}
};
template <bool is_min>
static Status MinMaxMLFloat16(const OpKernel& inst, OpKernelContext* context) {
const auto typed_allocator = [](const TensorAllocator& tensor_allocator, const TensorShape& shape) {
return tensor_allocator.Allocate<MLFloat16>(shape);
};
ProcessBroadcastSpanFuncs funcs{
[](BroadcastHelper& per_iter_bh) {
auto num_elements = per_iter_bh.NumOutputElements();
const auto* input_1 = reinterpret_cast<const Eigen::half*>(per_iter_bh.EigenInput1<MLFloat16>().data());
ConstEigenVectorArrayMap<Eigen::half> input_1_vec_map(input_1, num_elements);
auto* output = reinterpret_cast<Eigen::half*>(per_iter_bh.OutputEigen<MLFloat16>().data());
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>()));
} else {
output_vec_map = input_1_vec_map.max(static_cast<Eigen::half>(per_iter_bh.ScalarInput0<MLFloat16>()));
}
},
[](BroadcastHelper& per_iter_bh) {
auto num_elements = per_iter_bh.NumOutputElements();
const auto* input_0 = reinterpret_cast<const Eigen::half*>(per_iter_bh.EigenInput0<MLFloat16>().data());
ConstEigenVectorArrayMap<Eigen::half> input_0_vec_map(input_0, num_elements);
auto* output = reinterpret_cast<Eigen::half*>(per_iter_bh.OutputEigen<MLFloat16>().data());
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>()));
} else {
output_vec_map = input_0_vec_map.max(static_cast<Eigen::half>(per_iter_bh.ScalarInput1<MLFloat16>()));
}
},
[](BroadcastHelper& per_iter_bh) {
auto num_elements = per_iter_bh.NumOutputElements();
const auto* input_0 = reinterpret_cast<const Eigen::half*>(per_iter_bh.EigenInput0<MLFloat16>().data());
ConstEigenVectorArrayMap<Eigen::half> input_0_vec_map(input_0, num_elements);
const auto* input_1 = reinterpret_cast<const Eigen::half*>(per_iter_bh.EigenInput1<MLFloat16>().data());
ConstEigenVectorArrayMap<Eigen::half> input_1_vec_map(input_1, num_elements);
auto* output = reinterpret_cast<Eigen::half*>(per_iter_bh.OutputEigen<MLFloat16>().data());
EigenVectorArrayMap<Eigen::half> output_vec_map(output, num_elements);
if (is_min) {
output_vec_map = input_0_vec_map.min(input_1_vec_map);
} else {
output_vec_map = input_0_vec_map.max(input_1_vec_map);
}
}};
int input_count = inst.Node().InputArgCount().front();
UntypedBroadcastVariadic(input_count, *context, typed_allocator, funcs);
return Status::OK();
}
Status Min_8::Compute(OpKernelContext* context) const {
utils::MLTypeCallDispatcherRet<Status, ComputeImpl, float, double, MLFloat16, int32_t, uint32_t, int64_t, uint64_t>
t_disp(context->Input<Tensor>(0)->GetElementType());
return t_disp.Invoke(*this, context);
auto dt_type = context->Input<Tensor>(0)->GetElementType();
switch (dt_type) {
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
return MinMaxMLFloat16<true>(*this, context);
break;
default:
utils::MLTypeCallDispatcherRet<Status, ComputeImpl, float, double, int32_t, uint32_t, int64_t, uint64_t>
t_disp(dt_type);
return t_disp.Invoke(*this, context);
}
}
template <>
@ -622,9 +692,17 @@ struct Max_8::ComputeImpl {
};
Status Max_8::Compute(OpKernelContext* context) const {
utils::MLTypeCallDispatcherRet<Status, ComputeImpl, float, double, MLFloat16, int32_t, uint32_t, int64_t, uint64_t>
t_disp(context->Input<Tensor>(0)->GetElementType());
return t_disp.Invoke(*this, context);
auto dt_type = context->Input<Tensor>(0)->GetElementType();
switch (dt_type) {
case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
return MinMaxMLFloat16<false>(*this, context);
break;
default:
utils::MLTypeCallDispatcherRet<Status, ComputeImpl, float, double, int32_t, uint32_t, int64_t, uint64_t>
t_disp(dt_type);
return t_disp.Invoke(*this, context);
}
}
Status Not::Compute(OpKernelContext* context) const {

View file

@ -1241,6 +1241,44 @@ TEST(MathOpTest, Min_12_UInt64) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Min_12_MLFLoat16) {
OpTester test("Min", 12);
test.AddInput<MLFloat16>("data_0", {1, 3},
MakeMLFloat16({1.f, 1.f, 1.f}));
test.AddInput<MLFloat16>("data_1", {1, 3},
MakeMLFloat16({2.f, -1.f, -2.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({3.f, 2.f, -3.f}));
test.AddOutput<MLFloat16>("min", {1, 3},
MakeMLFloat16({1.f, -1.f, -3.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Min_12_MLFLoat16_Scalar0) {
OpTester test("Min", 12);
test.AddInput<MLFloat16>("data_0", {},
MakeMLFloat16({-10.f}));
test.AddInput<MLFloat16>("data_1", {1, 3},
MakeMLFloat16({2.f, -1.f, -2.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({3.f, 2.f, -3.f}));
test.AddOutput<MLFloat16>("min", {1, 3},
MakeMLFloat16({-10.f, -10.f, -10.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Min_12_MLFLoat16_Scalar1) {
OpTester test("Min", 12);
test.AddInput<MLFloat16>("data_0", {1, 3},
MakeMLFloat16({2.f, 3.f, 4.f}));
test.AddInput<MLFloat16>("data_1", {},
MakeMLFloat16({-10.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({3.f, 2.f, -3.f}));
test.AddOutput<MLFloat16>("min", {1, 3},
MakeMLFloat16({-10.f, -10.f, -10.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Max_6) {
OpTester test("Max", 6);
std::vector<int64_t> dims{3, 3};
@ -1418,6 +1456,45 @@ TEST(MathOpTest, Max_12_UInt64) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Max_12_MLFLoat16) {
OpTester test("Max", 12);
test.AddInput<MLFloat16>("data_0", {1, 3},
MakeMLFloat16({-1.f, -1.f, -1.f}));
test.AddInput<MLFloat16>("data_1", {1, 3},
MakeMLFloat16({-2.f, -1.f, -2.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({-3.f, -2.f, -3.f}));
test.AddOutput<MLFloat16>("max", {1, 3},
MakeMLFloat16({-1.f, -1.f, -1.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Max_12_MLFLoat16_Scalar0) {
OpTester test("Max", 12);
test.AddInput<MLFloat16>("data_0", {},
MakeMLFloat16({-1.f}));
test.AddInput<MLFloat16>("data_1", {1, 3},
MakeMLFloat16({-11.f, -12.f, -22.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({-10.f, -11.f, -13.f}));
test.AddOutput<MLFloat16>("max", {1, 3},
MakeMLFloat16({-1.f, -1.f, -1.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Max_12_MLFLoat16_Scalar1) {
OpTester test("Max", 12);
test.AddInput<MLFloat16>("data_0", {1, 3},
MakeMLFloat16({-1.f, -2.f, -3.f}));
test.AddInput<MLFloat16>("data_1", {},
MakeMLFloat16({2.f}));
test.AddInput<MLFloat16>("data_2", {1, 3},
MakeMLFloat16({-2.f, -3.f, -4.f}));
test.AddOutput<MLFloat16>("max", {1, 3},
MakeMLFloat16({2.f, 2.f, 2.f}));
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); //TensorRT: Input batch size is inconsistent
}
TEST(MathOpTest, Not) {
OpTester test("Not");
std::vector<int64_t> dims{2};