Skip to content

Commit 412233f

Browse files
borg323borg323
authored
coreml: Fix Gather fp16 support (microsoft#26442)
### Description <!-- Describe your changes. --> The coreml gather MLProgram operator supports fp16, but the check was missing in `GatherOpBuilder::HasSupportedInputsImpl()` ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> We use `Gather` in LeelaChessZero. --------- Co-authored-by: borg323 <borg323@users.noreply.github.com>
1 parent 1d149f5 commit 412233f

2 files changed

Lines changed: 64 additions & 56 deletions

File tree

‎onnxruntime/core/providers/coreml/builders/impl/gather_op_builder.cc‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,13 +55,15 @@ Status GatherOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
5555
return Status::OK();
5656
}
5757

58-
bool GatherOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& /*input_params*/,
58+
bool GatherOpBuilder::HasSupportedInputsImpl(const Node& node, const OpBuilderInputParams& input_params,
5959
const logging::Logger& logger) const {
6060
int32_t input_type;
6161
if (!GetType(*node.InputDefs()[0], input_type, logger))
6262
return false;
6363

6464
if (input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT &&
65+
(input_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ||
66+
!input_params.create_mlprogram || input_params.coreml_version < 6) &&
6567
input_type != ONNX_NAMESPACE::TensorProto_DataType_INT64) {
6668
LOGS(logger, VERBOSE) << "[" << node.OpType()
6769
<< "] Input type: [" << input_type

‎onnxruntime/test/providers/cpu/tensor/gather_op_test.cc‎

Lines changed: 61 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -5,54 +5,62 @@
55
#include "gtest/gtest.h"
66
#include "test/providers/provider_test_utils.h"
77
#include "test/util/include/default_providers.h"
8+
#include "test/common/tensor_op_test_utils.h"
89

910
namespace onnxruntime {
1011
namespace test {
1112

1213
// Some of the tests can't run on TensorrtExecutionProvider because of unsupported data types.
1314
// Those tests will fallback to other EPs
1415

15-
TEST(GatherOpTest, Gather_axis0) {
16+
template <typename T>
17+
class GatherOpTest : public ::testing::Test {
18+
};
19+
20+
using GatherOpTestTypes = ::testing::Types<float, MLFloat16>;
21+
TYPED_TEST_SUITE(GatherOpTest, GatherOpTestTypes);
22+
23+
TYPED_TEST(GatherOpTest, Gather_axis0) {
1624
// To test for NNAPI EP, we need the indices to be initializers
1725
auto run_test = [](bool indices_is_initializer) {
1826
OpTester test("Gather");
1927
test.AddAttribute<int64_t>("axis", 0LL);
20-
test.AddInput<float>("data", {2, 3, 4},
21-
{0.0f, 0.1f, 0.2f, 0.3f,
22-
1.0f, 1.1f, 1.2f, 1.3f,
23-
2.0f, 2.1f, 2.2f, 2.3f,
24-
10.0f, 10.1f, 10.2f, 10.3f,
25-
11.0f, 11.1f, 11.2f, 11.3f,
26-
12.0f, 12.1f, 12.2f, 12.3f});
28+
test.AddInput<TypeParam>("data", {2, 3, 4},
29+
GetTypedArray<TypeParam>({0.0f, 0.1f, 0.2f, 0.3f,
30+
1.0f, 1.1f, 1.2f, 1.3f,
31+
2.0f, 2.1f, 2.2f, 2.3f,
32+
10.0f, 10.1f, 10.2f, 10.3f,
33+
11.0f, 11.1f, 11.2f, 11.3f,
34+
12.0f, 12.1f, 12.2f, 12.3f}));
2735
test.AddInput<int64_t>("indices", {1}, {1LL}, indices_is_initializer);
28-
test.AddOutput<float>("output", {1, 3, 4},
29-
{10.0f, 10.1f, 10.2f, 10.3f,
30-
11.0f, 11.1f, 11.2f, 11.3f,
31-
12.0f, 12.1f, 12.2f, 12.3f});
36+
test.AddOutput<TypeParam>("output", {1, 3, 4},
37+
GetTypedArray<TypeParam>({10.0f, 10.1f, 10.2f, 10.3f,
38+
11.0f, 11.1f, 11.2f, 11.3f,
39+
12.0f, 12.1f, 12.2f, 12.3f}));
3240
test.Run();
3341
};
3442

3543
run_test(false);
3644
run_test(true);
3745
}
3846

39-
TEST(GatherOpTest, Gather_negative_axis) {
47+
TYPED_TEST(GatherOpTest, Gather_negative_axis) {
4048
// To test for NNAPI EP, we need the indices to be initializers
4149
auto run_test = [](bool indices_is_initializer) {
4250
OpTester test("Gather");
4351
test.AddAttribute<int64_t>("axis", -3LL);
44-
test.AddInput<float>("data", {2, 3, 4},
45-
{0.0f, 0.1f, 0.2f, 0.3f,
46-
1.0f, 1.1f, 1.2f, 1.3f,
47-
2.0f, 2.1f, 2.2f, 2.3f,
48-
10.0f, 10.1f, 10.2f, 10.3f,
49-
11.0f, 11.1f, 11.2f, 11.3f,
50-
12.0f, 12.1f, 12.2f, 12.3f});
52+
test.AddInput<TypeParam>("data", {2, 3, 4},
53+
GetTypedArray<TypeParam>({0.0f, 0.1f, 0.2f, 0.3f,
54+
1.0f, 1.1f, 1.2f, 1.3f,
55+
2.0f, 2.1f, 2.2f, 2.3f,
56+
10.0f, 10.1f, 10.2f, 10.3f,
57+
11.0f, 11.1f, 11.2f, 11.3f,
58+
12.0f, 12.1f, 12.2f, 12.3f}));
5159
test.AddInput<int64_t>("indices", {1}, {1LL}, indices_is_initializer);
52-
test.AddOutput<float>("output", {1, 3, 4},
53-
{10.0f, 10.1f, 10.2f, 10.3f,
54-
11.0f, 11.1f, 11.2f, 11.3f,
55-
12.0f, 12.1f, 12.2f, 12.3f});
60+
test.AddOutput<TypeParam>("output", {1, 3, 4},
61+
GetTypedArray<TypeParam>({10.0f, 10.1f, 10.2f, 10.3f,
62+
11.0f, 11.1f, 11.2f, 11.3f,
63+
12.0f, 12.1f, 12.2f, 12.3f}));
5664
test.Run();
5765
};
5866

@@ -206,65 +214,63 @@ TEST(GatherOpTest, Gather_axis0_indices2d) {
206214
run_test(true);
207215
}
208216

209-
TEST(GatherOpTest, Gather_axis1_indices2d) {
217+
TYPED_TEST(GatherOpTest, Gather_axis1_indices2d) {
210218
// To test for NNAPI EP, we need the indices to be initializers
211219
auto run_test = [](bool indices_is_initializer) {
212220
OpTester test("Gather");
213221
test.AddAttribute<int64_t>("axis", 1LL);
214-
test.AddInput<float>("data", {3, 3},
215-
{0.0f, 0.1f, 0.2f,
216-
1.0f, 1.1f, 1.2f,
217-
2.0f, 2.1f, 2.2f});
222+
test.AddInput<TypeParam>("data", {3, 3},
223+
GetTypedArray<TypeParam>({0.0f, 0.1f, 0.2f,
224+
1.0f, 1.1f, 1.2f,
225+
2.0f, 2.1f, 2.2f}));
218226
test.AddInput<int64_t>("indices", {2LL, 2LL},
219227
{1LL, 0LL,
220228
2LL, 1LL},
221229
indices_is_initializer);
222-
test.AddOutput<float>("output", {3, 2, 2},
223-
{0.1f, 0.0f, 0.2f, 0.1f,
224-
1.1f, 1.0f, 1.2f, 1.1f,
225-
2.1f, 2.0f, 2.2f, 2.1f});
230+
test.AddOutput<TypeParam>("output", {3, 2, 2},
231+
GetTypedArray<TypeParam>({0.1f, 0.0f, 0.2f, 0.1f,
232+
1.1f, 1.0f, 1.2f, 1.1f,
233+
2.1f, 2.0f, 2.2f, 2.1f}));
226234
test.Run();
227235
};
228236

229237
run_test(false);
230238
run_test(true);
231239
}
232240

233-
TEST(GatherOpTest, Gather_axis0_indicesInt32) {
241+
TYPED_TEST(GatherOpTest, Gather_axis0_indicesInt32) {
234242
// NNAPI EP only supports float input data for now,
235243
// the following two test cases cover int32_t indices with float input other than int64_t type for Nnapi
236244
OpTester test("Gather");
237245
test.AddAttribute<int64_t>("axis", 0LL);
238-
test.AddInput<float>("data", {2, 3, 4},
239-
{0.0f, 0.1f, 0.2f, 0.3f,
240-
1.0f, 1.1f, 1.2f, 1.3f,
241-
2.0f, 2.1f, 2.2f, 2.3f,
242-
10.0f, 10.1f, 10.2f, 10.3f,
243-
11.0f, 11.1f, 11.2f, 11.3f,
244-
12.0f, 12.1f, 12.2f, 12.3f});
246+
test.AddInput<TypeParam>("data", {2, 3, 4},
247+
GetTypedArray<TypeParam>({0.0f, 0.1f, 0.2f, 0.3f,
248+
1.0f, 1.1f, 1.2f, 1.3f,
249+
2.0f, 2.1f, 2.2f, 2.3f,
250+
10.0f, 10.1f, 10.2f, 10.3f,
251+
11.0f, 11.1f, 11.2f, 11.3f,
252+
12.0f, 12.1f, 12.2f, 12.3f}));
245253
test.AddInput<int32_t>("indices", {1}, {1});
246-
test.AddOutput<float>("output", {1, 3, 4},
247-
{10.0f, 10.1f, 10.2f, 10.3f,
248-
11.0f, 11.1f, 11.2f, 11.3f,
249-
12.0f, 12.1f, 12.2f, 12.3f});
250-
254+
test.AddOutput<TypeParam>("output", {1, 3, 4},
255+
GetTypedArray<TypeParam>({10.0f, 10.1f, 10.2f, 10.3f,
256+
11.0f, 11.1f, 11.2f, 11.3f,
257+
12.0f, 12.1f, 12.2f, 12.3f}));
251258
test.Run();
252259
}
253260

254-
TEST(GatherOpTest, Gather_axis0_indices2dInt32) {
261+
TYPED_TEST(GatherOpTest, Gather_axis0_indices2dInt32) {
255262
OpTester test("Gather");
256263
test.AddAttribute<int64_t>("axis", 0LL);
257-
test.AddInput<float>("data", {3, 3},
258-
{0.0f, 0.1f, 0.2f,
259-
1.0f, 1.1f, 1.2f,
260-
2.0f, 2.1f, 2.2f});
264+
test.AddInput<TypeParam>("data", {3, 3},
265+
GetTypedArray<TypeParam>({0.0f, 0.1f, 0.2f,
266+
1.0f, 1.1f, 1.2f,
267+
2.0f, 2.1f, 2.2f}));
261268
test.AddInput<int32_t>("indices", {2, 2},
262269
{1, 0,
263270
2, 1});
264-
test.AddOutput<float>("output", {2, 2, 3},
265-
{1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f,
266-
2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f});
267-
271+
test.AddOutput<TypeParam>("output", {2, 2, 3},
272+
GetTypedArray<TypeParam>({1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f,
273+
2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f}));
268274
test.Run();
269275
}
270276

0 commit comments

Comments
 (0)