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.

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

Add support for the Einsum operator #2772

Merged
merged 42 commits into from
May 29, 2024
Merged
Changes from 1 commit
Commits
Show all changes
42 commits
Select commit Hold shift + click to select a range
e98cf49
Add support for the Einsum operator [WIP]
mirza-halilcevic Feb 14, 2024
8472cce
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Feb 28, 2024
b44d80b
Add support for ... broadcasting and limited diagonal support
music-dino Mar 6, 2024
d7c7793
Implement more onnx_verify tests
music-dino Mar 6, 2024
d31932d
Merge remote-tracking branch 'upstream/develop' into einsum_support
music-dino Mar 6, 2024
bcf5015
Add missing gen_onnx.py tests for Einsum.
mirza-halilcevic Mar 6, 2024
fd38cf0
Add missing .onnx verify test files
music-dino Mar 7, 2024
047022e
Add two more einsum diagonal onnx verify tests
music-dino Mar 7, 2024
e2c6279
Add partial documentation
music-dino Mar 7, 2024
31c8aa6
Additional refactoring and documentation.
mirza-halilcevic Mar 7, 2024
dd1306a
Additional refactoring and renaming.
mirza-halilcevic Mar 11, 2024
26637f2
Re-enable onnx backend tests.
mirza-halilcevic Mar 11, 2024
a15ee68
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Mar 11, 2024
4ff489e
Refactor gather_diagonal.
mirza-halilcevic Mar 11, 2024
c1448dd
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Mar 11, 2024
596f1b9
Addressed several PR review comments.
mirza-halilcevic Mar 20, 2024
8a7e55d
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Mar 20, 2024
8537b14
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Mar 25, 2024
66de4ba
Address PR review comments
mirza-halilcevic Mar 27, 2024
34c3426
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic Mar 27, 2024
6b20225
Fix licensing, tidy and format checks.
mirza-halilcevic Mar 27, 2024
01d1cef
Add negative tests.
mirza-halilcevic Mar 28, 2024
3dceee1
Fix formatting and address review comments.
mirza-halilcevic Apr 3, 2024
8df341d
Merge branch 'develop' into einsum_support
causten Apr 19, 2024
366f6dc
Address code review comments.
mirza-halilcevic Apr 29, 2024
53d78db
Add comments explaining unsqueeze_transpose.
mirza-halilcevic Apr 30, 2024
40c9e28
Fix formatting.
mirza-halilcevic May 14, 2024
3db8573
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic May 14, 2024
3185360
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic May 14, 2024
45e4ec5
Refactor unsqueeze_transpose and squeeze_transpose.
mirza-halilcevic May 14, 2024
2c3144d
Rename unsqueeze_transpose to transpose_unsqueeze and transpose_squeeze
mirza-halilcevic May 14, 2024
c35c66c
Merge branch 'develop' into einsum_support
causten May 15, 2024
36279ff
Replace ssize_t with int64_t because windows build is failing.
mirza-halilcevic May 16, 2024
350b27c
Merge remote-tracking branch 'upstream/einsum_support' into einsum_su…
mirza-halilcevic May 16, 2024
02418fa
Merge branch 'develop' into einsum_support
causten May 17, 2024
6813ef2
Merge branch 'develop' into einsum_support
causten May 17, 2024
bc7d5cb
Merge remote-tracking branch 'upstream/develop' into einsum_support
mirza-halilcevic May 23, 2024
2542894
Fix onnx test loading.
mirza-halilcevic May 29, 2024
6034347
Merge branch 'develop' into einsum_support
mirza-halilcevic May 29, 2024
2641933
Fix clang-tidy error.
mirza-halilcevic May 29, 2024
980dffc
Merge remote-tracking branch 'upstream/einsum_support' into einsum_su…
mirza-halilcevic May 29, 2024
7140ae2
Fix string_view iterator issue with windows build.
mirza-halilcevic May 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
Prev Previous commit
Next Next commit
Fix formatting and address review comments.
mirza-halilcevic committed Apr 3, 2024

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.
commit 3dceee176d4beb3e5877f85a9772ef6ffca2ec93
16 changes: 8 additions & 8 deletions test/onnx/einsum_matrix_vector_multiplication_test.onnx
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
 (einsum_matrix_vector_multiplication_test:
+
 (einsum_matrix_vector_multiplication_test:
)
x
vy"Einsum*
equation" ij,kj->ik (einsum_matrix_vector_multiplication_testZ
vy"Einsum*
equation"ij,j->i (einsum_matrix_vector_multiplication_testZ
x


Z
v


Z
v


b
y

8 changes: 3 additions & 5 deletions test/onnx/gen_onnx.py
Original file line number Diff line number Diff line change
@@ -2064,13 +2064,13 @@ def einsum_row_sum_test():
@onnx_test()
def einsum_matrix_vector_multiplication_test():
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 3])
v = helper.make_tensor_value_info('v', TensorProto.FLOAT, [1, 3])
v = helper.make_tensor_value_info('v', TensorProto.FLOAT, [3])
y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 1])

node = onnx.helper.make_node('Einsum',
inputs=['x', 'v'],
outputs=['y'],
equation='ij,kj->ik')
equation='ij,j->i')

return ([node], [x, v], [y])

@@ -2554,9 +2554,7 @@ def einsum_missing_equation_negative_test():
x2 = helper.make_tensor_value_info('x2', TensorProto.FLOAT, [3, 3])
y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 3])

node = onnx.helper.make_node('Einsum',
inputs=['x1', 'x2'],
outputs=['y'])
node = onnx.helper.make_node('Einsum', inputs=['x1', 'x2'], outputs=['y'])

return ([node], [x1, x2], [y])

109 changes: 88 additions & 21 deletions test/onnx/parse/einsum_negative_tests.cpp
Original file line number Diff line number Diff line change
@@ -24,24 +24,91 @@

#include <onnx_test.hpp>

#define NEGATIVE_TEST(NAME) \
TEST_CASE(NAME) \
{ \
EXPECT(test::throws([&] { migraphx::parse_onnx(#NAME ".onnx"); })); \
}

NEGATIVE_TEST(einsum_missing_equation_negative_test)
NEGATIVE_TEST(einsum_multiple_arrows_negative_test)
NEGATIVE_TEST(einsum_empty_term_before_arrow_negative_test)
NEGATIVE_TEST(einsum_multiple_ellipses_negative_test)
NEGATIVE_TEST(einsum_comma_in_output_negative_test)
NEGATIVE_TEST(einsum_empty_term_before_comma_negative_test)
NEGATIVE_TEST(einsum_last_input_missing_negative_test)
NEGATIVE_TEST(einsum_term_input_mismatch_negative_test)
NEGATIVE_TEST(einsum_ellipsis_mismatch_negative_test)
NEGATIVE_TEST(einsum_rank_mismatch_negative_test)
NEGATIVE_TEST(einsum_output_surplus_label_negative_test)
NEGATIVE_TEST(einsum_output_missing_ellipsis_negative_test)
NEGATIVE_TEST(einsum_multiple_diagonals_negative_test)
NEGATIVE_TEST(einsum_diagonal_dim_mismatch_negative_test)
NEGATIVE_TEST(einsum_right_batch_diagonal_negative_test)
TEST_CASE(einsum_missing_equation_negative_test)
{
EXPECT(
test::throws([&] { migraphx::parse_onnx("einsum_missing_equation_negative_test.onnx"); }));
}

TEST_CASE(einsum_multiple_arrows_negative_test)
{
EXPECT(
test::throws([&] { migraphx::parse_onnx("einsum_multiple_arrows_negative_test.onnx"); }));
}

TEST_CASE(einsum_empty_term_before_arrow_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_empty_term_before_arrow_negative_test.onnx"); }));
}

TEST_CASE(einsum_multiple_ellipses_negative_test)
{
EXPECT(
test::throws([&] { migraphx::parse_onnx("einsum_multiple_ellipses_negative_test.onnx"); }));
}

TEST_CASE(einsum_comma_in_output_negative_test)
{
EXPECT(
test::throws([&] { migraphx::parse_onnx("einsum_comma_in_output_negative_test.onnx"); }));
}

TEST_CASE(einsum_empty_term_before_comma_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_empty_term_before_comma_negative_test.onnx"); }));
}

TEST_CASE(einsum_last_input_missing_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_last_input_missing_negative_test.onnx"); }));
}

TEST_CASE(einsum_term_input_mismatch_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_term_input_mismatch_negative_test.onnx"); }));
}

TEST_CASE(einsum_ellipsis_mismatch_negative_test)
{
EXPECT(
test::throws([&] { migraphx::parse_onnx("einsum_ellipsis_mismatch_negative_test.onnx"); }));
}

TEST_CASE(einsum_rank_mismatch_negative_test)
{
EXPECT(test::throws([&] { migraphx::parse_onnx("einsum_rank_mismatch_negative_test.onnx"); }));
}

TEST_CASE(einsum_output_surplus_label_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_output_surplus_label_negative_test.onnx"); }));
}

TEST_CASE(einsum_output_missing_ellipsis_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_output_missing_ellipsis_negative_test.onnx"); }));
}

TEST_CASE(einsum_multiple_diagonals_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_multiple_diagonals_negative_test.onnx"); }));
}

TEST_CASE(einsum_diagonal_dim_mismatch_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_diagonal_dim_mismatch_negative_test.onnx"); }));
}

TEST_CASE(einsum_right_batch_diagonal_negative_test)
{
EXPECT(test::throws(
[&] { migraphx::parse_onnx("einsum_right_batch_diagonal_negative_test.onnx"); }));
}
4 changes: 2 additions & 2 deletions test/onnx/verify/einsum_tests.cpp
Original file line number Diff line number Diff line change
@@ -130,15 +130,15 @@ TEST_CASE(einsum_matrix_vector_multiplication_test)
std::vector<float> x_data = {
0.4834133, 0.14106742, 0.50055824, 0.91764271, 0.95528452, 0.98199955};

migraphx::shape v_shape{migraphx::shape::float_type, {1, 3}};
migraphx::shape v_shape{migraphx::shape::float_type, {3}};
std::vector<float> v_data = {0.73961958, 0.53071864, 0.34152803};

migraphx::parameter_map pm;
pm["x"] = migraphx::argument{x_shape, x_data.data()};
pm["v"] = migraphx::argument{v_shape, v_data.data()};

auto result = p.eval(pm).back();
EXPECT(result.get_shape() == make_shape({2, 1}));
EXPECT(result.get_shape() == make_shape({2}));

std::vector<float> result_vector;
result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });

Unchanged files with check annotations Beta

if(error)
{
MIGRAPHX_THROW("TOKENIZE: no token found!");

Check warning on line 88 in src/include/migraphx/lexing.hpp

Codecov / codecov/patch

src/include/migraphx/lexing.hpp#L88

Added line #L88 was not covered by tests
}
}
has_ellipsis = true;
term += "*";
}
else if(token == ",")

Check warning on line 173 in src/onnx/parse_einsum.cpp

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L173

Added line #L173 was not covered by tests
{
if(ret.explicit_form)
MIGRAPHX_THROW("Einsum equation can't have a ',' symbol in the output");
static_cast<ssize_t>(calc_dim(sum_axes, op1_lens))};
std::vector<ssize_t> dims2{static_cast<ssize_t>(calc_dim(batch_axes, op2_lens)),
-1,
static_cast<ssize_t>(calc_dim(sum_axes, op2_lens))};

Check warning on line 516 in src/onnx/parse_einsum.cpp

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L516

Added line #L516 was not covered by tests
op1 = info.add_instruction(make_op("reshape", {{"dims", dims1}}), op1);
op2 = info.add_instruction(make_op("reshape", {{"dims", dims2}}), op2);
for(int d = 0; d < map_mat[0].size(); ++d)
{
if(cur_pair[0][d] > 0 and cur_pair[1][d] == -1)
red.push_back(d);

Check warning on line 556 in src/onnx/parse_einsum.cpp

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L556

Added line #L556 was not covered by tests
}
op = apply_reduce_sum_op(info, op, red, cur_pair[1]);