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
Add two more einsum diagonal onnx verify tests
music-dino committed Mar 7, 2024

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.
commit 047022eec5cbfd0325ffe5dd1d15b3dadb043fb9
17 changes: 17 additions & 0 deletions test/onnx/einsum_common_8_test.onnx
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
 einsum_common_8_test:†
-
x1
x2y"Einsum*
equation" ii,jj->ij einsum_common_8_testZ
x1


Z
x2


b
y


B
17 changes: 17 additions & 0 deletions test/onnx/einsum_diag_vector_multiply_test.onnx
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
  einsum_diag_vector_multiply_test:ˆ
+
x1
x2y"Einsum*
equation"ii,i->i  einsum_diag_vector_multiply_testZ
x1


Z
x2


b
y


B
28 changes: 28 additions & 0 deletions test/onnx/gen_onnx.py
Original file line number Diff line number Diff line change
@@ -2226,6 +2226,20 @@ def einsum_3d_diagonal_test():
return ([node], [x], [y])


@onnx_test()
def einsum_diag_vector_multiply_test():
x1 = helper.make_tensor_value_info('x1', TensorProto.FLOAT, [3, 3])
x2 = helper.make_tensor_value_info('x2', TensorProto.FLOAT, [3])
y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3])

node = onnx.helper.make_node('Einsum',
inputs=['x1', 'x2'],
outputs=['y'],
equation='ii,i->i')

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


@onnx_test()
def einsum_matrix_trace_test():
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 3])
@@ -2519,6 +2533,20 @@ def einsum_common_7_test():
return ([node], [x], [y])


@onnx_test()
def einsum_common_8_test():
x1 = helper.make_tensor_value_info('x1', TensorProto.FLOAT, [3, 3])
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'],
equation='ii,jj->ij')

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


@onnx_test()
def elu_test():
x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [3])
82 changes: 82 additions & 0 deletions test/onnx/verify/einsum_test.cpp
Original file line number Diff line number Diff line change
@@ -629,6 +629,39 @@ TEST_CASE(einsum_3d_diagonal_test)
EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
}

TEST_CASE(einsum_diag_vector_multiply_test)
{
migraphx::program p = migraphx::parse_onnx("einsum_diag_vector_multiply_test.onnx");
p.compile(migraphx::make_target("ref"));

migraphx::shape x1_shape{migraphx::shape::float_type, {3, 3}};
std::vector<float> x1_data = {0.8628764,
0.96045198,
0.14103307,
0.89249896,
0.97520951,
0.7015561,
0.06408759,
0.59921615,
0.76173894};

migraphx::shape x2_shape{migraphx::shape::float_type, {3}};
std::vector<float> x2_data = {0.79284103, 0.61505765, 0.70876231};

migraphx::parameter_map pm;
pm["x1"] = migraphx::argument{x1_shape, x1_data.data()};
pm["x2"] = migraphx::argument{x2_shape, x2_data.data()};

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

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

std::vector<float> gold = {0.68412382, 0.59981008, 0.53989185};
EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
}

TEST_CASE(einsum_matrix_trace_test)
{
migraphx::program p = migraphx::parse_onnx("einsum_matrix_trace_test.onnx");
@@ -1593,3 +1626,52 @@ TEST_CASE(einsum_common_7_test)
std::vector<float> gold = {2.90563922, 2.5946174, 2.82818581, 2.47204655, 2.28814157};
EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
}

TEST_CASE(einsum_common_8_test)
{
migraphx::program p = migraphx::parse_onnx("einsum_common_8_test.onnx");
p.compile(migraphx::make_target("ref"));

migraphx::shape x1_shape{migraphx::shape::float_type, {3, 3}};
std::vector<float> x1_data = {0.31281588,
0.34922652,
0.79181082,
0.55581571,
0.34963734,
0.39777707,
0.43040396,
0.19965846,
0.68818176};

migraphx::shape x2_shape{migraphx::shape::float_type, {3, 3}};
std::vector<float> x2_data = {0.94199384,
0.06564557,
0.36439139,
0.30556677,
0.25776106,
0.59531702,
0.21481152,
0.09608821,
0.41203512};

migraphx::parameter_map pm;
pm["x1"] = migraphx::argument{x1_shape, x1_data.data()};
pm["x2"] = migraphx::argument{x2_shape, x2_data.data()};

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

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

std::vector<float> gold = {0.29467063,
0.08063175,
0.12889113,
0.32935622,
0.09012289,
0.14406286,
0.64826297,
0.17738646,
0.28355505};
EXPECT(migraphx::verify::verify_rms_range(result_vector, gold));
}

Unchanged files with check annotations Beta

const std::vector<instruction_ref>& args) const
{
if(not contains(info.attributes, "equation"))
MIGRAPHX_THROW("Equation attribute is required");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L51

Added line #L51 was not covered by tests
std::string equation = info.attributes.at("equation").s();
auto [terms, unique_labels, ellipses_ndim] = analyze_equation(equation, args);
for(int d = 0; d < mat[0].size(); ++d)
{
if(rows[0][d] > 0 and rows[1][d] == -1)
red.push_back(d);

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L193

Added line #L193 was not covered by tests
else if(rows[0][d] == -1 and rows[1][d] >= 0)
MIGRAPHX_THROW("Issue in equation");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L195

Added line #L195 was not covered by tests
}
op = apply_reduce_sum_op(info, op, red, rows[1]);
std::vector<std::tuple<int, int_vec>> diag) const
{
if(diag.size() != 1)
MIGRAPHX_THROW("Not implemented with more than one duplicated label");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L210

Added line #L210 was not covered by tests
auto axis = std::get<0>(diag[0]);
auto axes = std::get<1>(diag[0]);
auto min_axes = *(std::min_element(axes.begin(), axes.end()));
if(not all_of(batch_axes, [=](int ba) { return ba < min_axes; }))
MIGRAPHX_THROW("Currently batch axes have to be partitioned to the left");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L222

Added line #L222 was not covered by tests
auto op_shape = op->get_shape().lens();
if(not all_of(axes, [op_shape, axis](int a) { return op_shape[axis] == op_shape[a]; }))
MIGRAPHX_THROW("All duplicated indices have to be the same dimension");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L227

Added line #L227 was not covered by tests
size_t batch_size =
std::accumulate(batch_axes.begin(), batch_axes.end(), 1, [&](auto acc, auto dim) {
for(auto& r : rows[1])
{
if(r == t)
MIGRAPHX_THROW("Unexpected result");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L273

Added line #L273 was not covered by tests
if(r > t)
r -= 1;

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L276

Added line #L276 was not covered by tests
}
}
case ' ': break;
case '-':
if(explicit_form)
MIGRAPHX_THROW("Einsum equation has multiple '->' symbols");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L579

Added line #L579 was not covered by tests
if(i + 1 >= equation.size() or equation[i + 1] != '>')
MIGRAPHX_THROW("Invalid '->' in einsum equation");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L581

Added line #L581 was not covered by tests
++i;
explicit_form = true;
break;
case '.':
if(has_ellipsis)
MIGRAPHX_THROW("Ellipsis can only appear once per einsum equation term");

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L593

Added line #L593 was not covered by tests
if(i + 2 >= equation.size() or equation[i + 1] != '.' or equation[i + 2] != '.')
MIGRAPHX_THROW("Incomplete ellipsis in einsum equation " +

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L596

Added line #L596 was not covered by tests
std::string(equation));
i += 2;
break;
default:
if(std::isalpha(c) == 0)
MIGRAPHX_THROW(std::string("Invalid character '") + c +

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L605

Added line #L605 was not covered by tests
"' in einsum equation term");
term += c;
return not contains(label_count, l) and l != '*';
});
if(it != output_term.end())
MIGRAPHX_THROW("Output term contains label " + std::to_string(*it) +

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L640

Added line #L640 was not covered by tests
", which is not present in any of the input terms");
if(ellipses_ndim != 0 and not contains(output_term, "*"))
MIGRAPHX_THROW(

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L644

Added line #L644 was not covered by tests
"Output term does not contain ellipsis (...) even though an input term does");
}
const std::vector<instruction_ref>& args) const
{
if(input_terms.size() != args.size())
MIGRAPHX_THROW(

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L652

Added line #L652 was not covered by tests
"Number of terms in the input equation - " + std::to_string(input_terms.size()) +
" does not match the number of input tensors " + std::to_string(args.size()));
{
auto ellipses_dims = rank - term.size() + 1;
if(global_ellipses_dims > 0 and ellipses_dims != global_ellipses_dims)
MIGRAPHX_THROW("Every occurrence of ellipsis in the equation must "

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L670

Added line #L670 was not covered by tests
"represent the same number of dimensions");
global_ellipses_dims = ellipses_dims;
current_dim += ellipses_dims;
}
if(current_dim != rank)
MIGRAPHX_THROW("Number of labels in " + std::to_string(i + 1) + ". input_term (" +

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

Codecov / codecov/patch

src/onnx/parse_einsum.cpp#L680

Added line #L680 was not covered by tests
term + ") does not match the rank (" + std::to_string(rank) +
") of corresponding input");
}