[TUZ-6] Add a direct Onnx to Relax Importer #14

jwfromm · 2023-02-03T17:12:52Z

This PR is the culmination of work for the Onnx importer epic. It implements a converter similar in spirit to the Onnx -> Relay importer and even reuses many of the operator converters directly. Other operators instead are converted by directly emitting tensorir functions using topi.

What should we put here?

What are the right tests to run?

This PR depends on importing onnx, how should we update CI to support new dependencies?

python/tvm/relax/frontend/onnx_frontend.py

tests/python/relax/frontend/test_onnx_frontend.py

python/tvm/relax/frontend/onnx_frontend.py

areusch · 2023-02-03T18:00:27Z

python/tvm/relax/frontend/onnx_frontend.py

+        name = dim.dim_param
+        value = dim.dim_value
+        if value is None or value == 0:
+            value = tvm.tir.Var("d", "int64")


why name the var "d"?

I think we usually use 'd' for dynamic, although I'm not sure..

Yeah its just a stand-in dynamic shape. I've updated it to be dyn for more clarity.

python/tvm/relax/frontend/onnx_frontend.py

slyubomirsky · 2023-02-03T18:13:05Z

My overall high-level reactions (I will leave finer-grained comments in the code) from the perspective of documentation:

Gut reaction for what docs we would expect

A high-level comment outlining the overall algorithm for the conversion, explaining the relevant portions of an ONNX model and what each component means (at the very least, have links to ONNX docs)
We probably don't need a fully fledged tutorial explaining how to add new operators because the implementation of the importer might change over time, but a list of steps outlined in a comment near the implementation might be nice
A possible tutorial might be showing how to use the front-end interface for the importer and explaining, broadly, what happens and how the ONNX model is transformed into a Relax program

Specific reactions

The operator implementations are generally straightforward and so is the convert map, definitely good in terms of being "self-documenting" code
The graph conversion is an area where things can become very complicated and depend on details of ONNX
Name sanitization is a very potentially tricky issue that should be outlined in detail. This has been an enormous source of bugs and confusion in practice with Relay, so I think it should be explained as clearly as possible and made as idiot-proof as possible
The tests for individual operators seem to be lacking, as far as the goal of reaching branch coverage.

How the importer should be tested

The importer is really a compiler, so tests should be checking differential equality: Make sure the importer version gets the same result as the original
Fuzz random inputs to individual operators as a start, aim for branch coverage
Also have some integration tests making sure that the operators can be built up as expected
Less useful form of testing: Making sure that some specific AST is produced given an input program. That is just coupling too closely with the implementation. Better to ensure the equivalence of the results

src/relax/op/tensor/manipulate.cc

tqchen · 2023-02-03T22:56:45Z

python/tvm/relax/op/manipulate.py

+
+
+@tvm.register_func("relax.run.broadcast_to")
+def torch_broadcast_to(data: tvm.nd.array, shape: tvm.nd.array) -> tvm.nd.array:


likely this is no longer needed if you can lower bcast to via emit_te

True, i dont think that lowering exists yet but as soon as it does we'll remove this function.

slyubomirsky

After taking a detailed look at the implementation, I've highlighted some areas in which the tests can be improved (many test cases do check all the branches, which is good). Overall, I stand by my earlier comments that there should be some outline of the overall structure of the converter. There should also be some explanation of how the existing Relay converter's methods are used (is that intended to remain permanently? Not sure it's good to depend on the Relay converter at all if we can avoid it, just another massive complication)

slyubomirsky · 2023-02-06T19:39:58Z

src/relax/op/tensor/manipulate.cc

@@ -88,21 +88,22 @@ StructInfo InferStructInfoBroadcastTo(const Call& call, const BlockBuilder& ctx)
    const auto* old_len_int = old_len.as<IntImmNode>();
    if (old_len_int != nullptr && old_len_int->value == 1) {
      continue;
-    } else if (!analyzer->CanProveEqual(old_len, tgt_len)) {


Could you explain why this was commented out?

Ah yeah this portion of code is specifically requiring shapes to be known and checking if those static shapes are compatible for broadcasting. My opinion is that we should not enforce that since dynamic broadcast_to shows up pretty frequently. This should be fully removed rather than commented out though.

I see. The intent of the original code was to require using a MatchCast to check the shape first. However, I agree with you that the operator implementation can do the checking, so that's unnecessary work for the user

python/tvm/relax/frontend/onnx_frontend.py

tests/python/relax/frontend/test_onnx_frontend.py

slyubomirsky · 2023-02-06T21:36:30Z

tests/python/relax/frontend/test_onnx_frontend.py

+    check_correctness(model)
+
+
+def test_const():


Perhaps it would be good to have tests for multiple types or for the edge case discussed in the implementation (not supporting strings)

tests/python/relax/frontend/test_onnx_frontend.py

slyubomirsky · 2023-02-06T21:54:26Z

python/tvm/relax/frontend/onnx_frontend.py

+        converter, which should be `_impl_vx`. Number x is the biggest
+            number smaller than or equal to opset belongs to all support versions.
+        """
+        versions = [int(d.replace("_impl_v", "")) for d in dir(cls) if "_impl_v" in d]


Might there be a simpler way to accomplish the dispatching by opset? E.g., as a property in the converter classes. Encoding it in the method names seems a little unusual. Is split the only operator that has multiple versions in this manner?

In principle, additionally, we should test the error case for rejecting an op due to version (we could test this function separately)

At the moment -- yes, split is the only operator that has multiple versions. Probably in the future, this will not hold true.

I am not sure about using a different kind of encoding -- what is nice with this implementation is that each converter does not have to contain a different mapping opset version -> function implementation

tests/python/relax/frontend/test_onnx_frontend.py

jwfromm · 2023-02-13T16:53:05Z

@driazati, it looks like the build is failing due to the image not having onnx. How should we update that as a dependency?

slyubomirsky

I think my principal concerns have been addressed and I thank those who made the changes. We can improve the documentation as we go along.

* [Relax][Onnx] Implement Div, Sigmoid, Softmax, Transpose and Unsqueeze ops * skip test_reshape * [Relax][ONNX] Implement BiasGelu and Gelu ops * [Relax][ONNX] Implement Where op

…hape / Not / Tanh (#3) * Rebase w/ Equal, Not, Tanh, Sqrt, Relu, Clip, Conv, Pow, Erf. * Fix cumsum but still needs work.

* Add squeeze. * Add Constant. * Add sub.

…tend (#8) * [WIP] Support using Relay ops in the Relax ONNX frontend Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * [WIP] small fixes Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * [WIP] Support dynamic matmul and reshape Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * Address PR comments --------- Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]>

* [WIP] add more ops. Some fail at the moment * skip some tests * Remove duplicate tests for squeeze

* [Relax][ONNX] Add Split op * Remove tmp

jwfromm · 2023-02-17T04:28:15Z

I'm going to merge this to help unblock other efforts. There are still a few CI related issues we should try to fix. Specifically, it seems like CI failed despite all tests passing because some of those tests triggered intentional errors.

* Initial importer and testing scaffolding. * Implement matmul operator and tests. * Add a bunch of new operators. * Add new ops * [Relax][Onnx] Implement Div, Sigmoid, Softmax, Transpose and Unsqueeze ops * skip test_reshape * [Relax][ONNX] Implement BiasGelu and Gelu ops * [Relax][ONNX] Implement Where op * [Relax][ONNX] Add Multiple ONNX Frontend Support for Clip / Equal / Shape / Not / Tanh (#3) * Rebase w/ Equal, Not, Tanh, Sqrt, Relu, Clip, Conv, Pow, Erf. * Fix cumsum but still needs work. * Fix initializer for CumSum. (#9) * Add Constant, Squeeze & Sub (#10) * Add squeeze. * Add Constant. * Add sub. * Support reusing Relay ONNX operator convertors in the Relax ONNX frontend (#8) * [WIP] Support using Relay ops in the Relax ONNX frontend Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * [WIP] small fixes Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * [WIP] Support dynamic matmul and reshape Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * Address PR comments --------- Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> * Add more ops (including all Reduce ops) using the relay frontend (#11) * [WIP] add more ops. Some fail at the moment * skip some tests * Remove duplicate tests for squeeze * Add Split op in the Relax ONNX frontend (#12) * [Relax][ONNX] Add Split op * Remove tmp * Fix layer normalizations and Shape operator. * Replace main loop with tvm testing. * Simplify Slice for opset 13. * [Relax][ONNX] Implement pad op * Incorporate pad op, add static constantofshape op. * Changes to shape to temporarily enable constantofshape in our models. * Add initial tensor_to_shape implementation. * Implemented dynamic broadcast_to to support expand and constantofshape. * Changes sufficient for vortex end to end run. * Formatting. * Format tests. * Re-add broadcast_to shape checking. * Fix formatting. * Remove overly strict manipulate check. * Fix typing * [Relax][Onnx] Implement Tile operator * Switch to native relax attention importer. * Address some of the PR comments * Check for the imported model IR version * switch from torch to numpy due to some incompatibility * Fix make format. * Clean up typing issues. * Clarify variable name. * Remove unneeded comprehension. * Remove circular dependency. * Add name sanitization for inputs * Disable reshape rewrite pass until fixed. * Fix long comment * Update cpu image. --------- Co-authored-by: Florin Blanaru <[email protected]> Co-authored-by: Xiyou Zhou <[email protected]> Co-authored-by: Matthew Barrett <[email protected]> Co-authored-by: Michalis Papadimitriou <[email protected]> Co-authored-by: Florin Blanaru <[email protected]> Co-authored-by: sung <[email protected]>

driazati reviewed Feb 3, 2023

View reviewed changes

yelite reviewed Feb 3, 2023

View reviewed changes

python/tvm/relax/frontend/onnx_frontend.py Show resolved Hide resolved

python/tvm/relax/frontend/onnx_frontend.py Show resolved Hide resolved

python/tvm/relax/frontend/onnx_frontend.py Show resolved Hide resolved

areusch reviewed Feb 3, 2023

View reviewed changes

jwfromm force-pushed the TUZ-6 branch from ac299df to 2221bcc Compare February 3, 2023 20:58

tqchen reviewed Feb 3, 2023

View reviewed changes

src/relax/op/tensor/manipulate.cc Outdated Show resolved Hide resolved

tqchen reviewed Feb 3, 2023

View reviewed changes

slyubomirsky suggested changes Feb 6, 2023

View reviewed changes

slyubomirsky approved these changes Feb 13, 2023

View reviewed changes

jwfromm force-pushed the TUZ-6 branch from 21c6372 to 621c52b Compare February 14, 2023 19:28

jwfromm pushed a commit that referenced this pull request Feb 14, 2023

Reorganize source code. (#14)

3087734

jwfromm force-pushed the relax branch from f6db158 to 1381a52 Compare February 14, 2023 22:31

jwfromm force-pushed the TUZ-6 branch 2 times, most recently from f36f4ea to 1e14a30 Compare February 15, 2023 23:46

Josh Fromm and others added 15 commits February 16, 2023 13:13

Initial importer and testing scaffolding.

16bc26d

Implement matmul operator and tests.

ed7ea1f

Add a bunch of new operators.

e9cafd7

Add new ops

54a495a

* [Relax][Onnx] Implement Div, Sigmoid, Softmax, Transpose and Unsqueeze ops * skip test_reshape * [Relax][ONNX] Implement BiasGelu and Gelu ops * [Relax][ONNX] Implement Where op

[Relax][ONNX] Add Multiple ONNX Frontend Support for Clip / Equal / S…

ba6ccf2

…hape / Not / Tanh (#3) * Rebase w/ Equal, Not, Tanh, Sqrt, Relu, Clip, Conv, Pow, Erf. * Fix cumsum but still needs work.

Fix initializer for CumSum. (#9)

99ed8d9

Add Constant, Squeeze & Sub (#10)

c492c54

* Add squeeze. * Add Constant. * Add sub.

Add more ops (including all Reduce ops) using the relay frontend (#11)

c59cfb1

* [WIP] add more ops. Some fail at the moment * skip some tests * Remove duplicate tests for squeeze

Add Split op in the Relax ONNX frontend (#12)

7bd93bb

* [Relax][ONNX] Add Split op * Remove tmp

Fix layer normalizations and Shape operator.

10dd45b

Replace main loop with tvm testing.

909473d

Simplify Slice for opset 13.

4205426

[Relax][ONNX] Implement pad op

cec2760

Incorporate pad op, add static constantofshape op.

a6f756c

Josh Fromm and others added 19 commits February 16, 2023 13:13

Changes sufficient for vortex end to end run.

3fae07b

Formatting.

370ed01

Format tests.

359615d

Re-add broadcast_to shape checking.

06f1114

Fix formatting.

acf84b9

Remove overly strict manipulate check.

1d594e4

Fix typing

c3dfe1d

[Relax][Onnx] Implement Tile operator

03d840e

Switch to native relax attention importer.

eca82df

Address some of the PR comments

bf9695c

Check for the imported model IR version

96331c4

switch from torch to numpy due to some incompatibility

853873e

Fix make format.

b0d7cd4

Clean up typing issues.

609718b

Clarify variable name.

cb0156e

Remove unneeded comprehension.

79de976

Remove circular dependency.

bf792ee

Add name sanitization for inputs

84ccc8e

Disable reshape rewrite pass until fixed.

20e3b92

jwfromm force-pushed the TUZ-6 branch from 1e14a30 to 20e3b92 Compare February 16, 2023 21:13

Fix long comment

8ce334c

zxybazh mentioned this pull request Feb 17, 2023

[TUZ-162] ONNX Frontend Supported By Relax Operators #22

Closed

4 tasks

Update cpu image.

9f0e3b1

jwfromm merged commit 62d5ad1 into relax Feb 17, 2023

jwfromm deleted the TUZ-6 branch March 7, 2023 21:30

vinx13 pushed a commit to vinx13/relax-octo that referenced this pull request Mar 29, 2023

Reorganize source code. (octoml#14)

3fed13c

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TUZ-6] Add a direct Onnx to Relax Importer #14

[TUZ-6] Add a direct Onnx to Relax Importer #14

jwfromm commented Feb 3, 2023 •

edited

Loading

areusch Feb 3, 2023

gigiblender Feb 9, 2023

jwfromm Feb 10, 2023

slyubomirsky commented Feb 3, 2023 •

edited

Loading

tqchen Feb 3, 2023

jwfromm Feb 3, 2023

slyubomirsky left a comment

slyubomirsky Feb 6, 2023

jwfromm Feb 10, 2023

slyubomirsky Feb 10, 2023

slyubomirsky Feb 6, 2023

slyubomirsky Feb 6, 2023

slyubomirsky Feb 6, 2023

gigiblender Feb 9, 2023

jwfromm commented Feb 13, 2023 •

edited

Loading

slyubomirsky left a comment

jwfromm commented Feb 17, 2023



		@tvm.register_func("relax.run.broadcast_to")
		def torch_broadcast_to(data: tvm.nd.array, shape: tvm.nd.array) -> tvm.nd.array:

[TUZ-6] Add a direct Onnx to Relax Importer #14

[TUZ-6] Add a direct Onnx to Relax Importer #14

Conversation

jwfromm commented Feb 3, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slyubomirsky commented Feb 3, 2023 • edited Loading

Gut reaction for what docs we would expect

Specific reactions

How the importer should be tested

Choose a reason for hiding this comment

Choose a reason for hiding this comment

slyubomirsky left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jwfromm commented Feb 13, 2023 • edited Loading

slyubomirsky left a comment

Choose a reason for hiding this comment

jwfromm commented Feb 17, 2023

jwfromm commented Feb 3, 2023 •

edited

Loading

slyubomirsky commented Feb 3, 2023 •

edited

Loading

jwfromm commented Feb 13, 2023 •

edited

Loading