Relay->Relax translator (ResNet example) #75

YuchenJin · 2022-02-02T18:55:42Z

This PR implements a Relay to Relax translator, which allows us to import Relay workloads to Relax for benchmarking and development purposes:

Support converting the most commonly used operators to Relax through EmitTE in relay_translator.py . @yongwww will add more ops and implement a LSTM demo in the future.
A ResNet demo by translating the Relay ResNet model is in apps/relax_examples/resnet.py.

cc: @yongwww @hypercubestart @ZihengJiang

ZihengJiang · 2022-02-04T01:44:42Z

include/tvm/relax/expr.h

+class ConstantNode : public ExprNode {
+ public:
+  /*! \brief The data of the tensor */
+  runtime::NDArray data;


Can we use relay::ConstantNode?

The reason I added relax::ConstantNode was that when we meet a ConstantNode in relay in the translation, we need to fill its shape_ field so that it can be converted to a te tensor. The constructor of relax::ConstantNode fills the shape_ field here: https://github.com/tlc-pack/relax/blob/relax-resnet/src/relax/ir/expr.cc#L61.

Can we change the original constructor to make shape as an optional field?

Now I reuse the relay::ConstantNode and fill its shape and checked_type during construction.

ZihengJiang · 2022-02-04T01:50:09Z

python/tvm/relax/testing/relay_translator.py

+        if attrs is not None:
+            shape = attrs["shape"]
+            dtype = attrs["dtype"]
+            return nn.emit_te(topi.full, shape, dtype, 0.0)


I still feel it is better to make those Op conversion more automatic. We should be able to get the op's corresponding TE function and attributes according to the Op object.

Thanks for the suggestion! Yes, I totally agree it'd be great to make the Op conversions more generic and automatic.

I couldn't find a good API to get an op's TE function in the TVM codebase (op.get_attr("FTVMCompute") returns a PackedFunc), do you happen to have a pointer? And the other question: is Op and TE function a one-to-one mapping? For example, nn.conv2d maps to multiple TE functions/strategies depending on the layout.

Another inconvenient and cumbersome part when I implemented the converter was some topi functions take positional arguments (for example one_hot), but some of these positional arguments are stored in the attribute (which is a dict). So I had to unpack them from the attr and insert them like the shape and dtype here. I think we can write a signature inspector to automate this part in the future if we really want it.

I would suggest to take a look at how the current lowering flow works in Relay, since I think it also need to do this job: according to the op and target, get the schedule out and do the leftover lowering. I can also do some check this week.

But if the PR is needed urgently, we can merge it now

No worries, let me try something and we can (probably) merge it tomorrow. :)

automatic relay->topi mapping looks helpful for many cases, but there are still lots of models that are not able to be supported even we have the automatic op mapping. For example, almost all object detection/segmentation models like Mark-RCNN are composed with Relay ADT stuffs like TensorArray, which is not straightforward to be converted to relax. Probably we will need the converter from original DL framework to relax in the end (more discussion is needed)

Yes, we definitely need the converter builtin relax eventually. Otherwise, the expressiveness cannot beyond of the relay's scope. The translator here can be used as a temporary solution for now when we don't have the optimization passes yet.

I updated the translator to leverage FTVMCompute for quite a lot of Ops' translations, which saves a lot of code 😎. Please take a look. Thanks! :)

apps/relax_examples/resnet.py

python/tvm/relax/testing/relay_translator.py

yongwww

overall LGTM!

cc @ZihengJiang @hypercubestart could you please take another look when you have some cycles?

hypercubestart

LGTM! thanks for the hard work!

yongwww · 2022-02-11T04:28:58Z

apps/relax_examples/resnet.py

+    relax_mod = relay_translator.from_relay(relay_mod["main"])
+
+    # print the ResNet IRmodule got translated
+    print(R.parser.astext(relax_mod))


might not need this print? and how about moving this file to tests/python/relax/

I'd prefer to put it into apps/relax_examples because it's a demo to translate the Relay ResNet workload to Relax. And I put the translator in python/tvm/relax/testing/ because the relay->relax translator is a current workaround to allow us to construct real workloads quickly. We will need pytorch/tf/onnx importers once we have high-level Ops on Relax.

src/printer/relax_script_printer.cc

python/tvm/relax/testing/topi.py

include/tvm/relax/expr.h

include/tvm/relax/topi.h

python/tvm/relax/testing/relay_translator.py

Robslhc · 2022-02-11T09:31:46Z

Hi @YuchenJin ,
I have a question about this translator. It seems that all relay ops are directly transformed into PrimFuncs. Then how should graph optimizations be performed? For example, we can see the TIR function of batch_norm, which is optimized originally by SimplifyInference Pass in Relay. Also, we can not see any fused ops in the translated IR. I guess it is not easy to perform such graph optimizations on this result, because we do not have the concept of "op". What's your idea about this?

YuchenJin · 2022-02-11T16:10:54Z

Hi @YuchenJin , I have a question about this translator. It seems that all relay ops are directly transformed into PrimFuncs. Then how should graph optimizations be performed? For example, we can see the TIR function of batch_norm, which is optimized originally by SimplifyInference Pass in Relay. Also, we can not see any fused ops in the translated IR. I guess it is not easy to perform such graph optimizations on this result, because we do not have the concept of "op". What's your idea about this?

Hi @Robslhc, glad to see you here, it's a great question! :)

In last week's Relax open dev meeting, @junrushao1994 had a good idea: we can try translating the after-fused Relay program (program after FuseOps pass) to Relax directly, so we don't have to write a Relax fusion pass for now to achieve comparable performance with Relay. We can enhance this translator later to add this feature.

include/tvm/relax/topi.h

src/relax/topi/topi.cc

yongwww · 2022-02-14T17:08:46Z

Thanks @YuchenJin for the hard work!!! thanks @ZihengJiang @tqchen @hypercubestart @Robslhc @junrushao1994 for the review/discussion!

ZihengJiang

Great work. Just a minor question

ZihengJiang · 2022-02-14T18:00:13Z

src/printer/relax_script_printer.cc

@@ -215,6 +216,12 @@ Doc RelaxScriptPrinter::VisitNode_(const relay::TupleGetItemNode* op) {
  return doc;
 }

+Doc RelaxScriptPrinter::VisitNode_(const relay::ConstantNode* op) {
+  Doc doc;
+  doc << "meta[relax.Constant][" << constant_counter_++ << "]";


Will this emit different counter number for the same constant?

good catch! i think it might emit different counter number. Actually this block is more like a todo mark, it will be implemented and tested in #76.

Sounds good :)

Thanks Ziheng for the good catch! You are indeed right. Yong will add the "real" metadata section support in his PR. :)

* Relay translator; build static bert forward/backward. * rebase * Add ops. * resnet demo * cleanup code. * Rebase. * Address comments. * leverage FTVMCompute for most op translation; reuse relay.Constant. * lint.

* Squashed minor cases (tlc-pack#74) * squashed * minor cases * minor cases `T.void` refactor `bind_value` apply code review suggestions add `T.Let` add `node` in `bind_value` for error reporting * update annotation

* Relay translator; build static bert forward/backward. * rebase * Add ops. * resnet demo * cleanup code. * Rebase. * Address comments. * leverage FTVMCompute for most op translation; reuse relay.Constant. * lint.

YuchenJin added 5 commits February 2, 2022 09:47

Relay translator; build static bert forward/backward.

1151660

rebase

598782c

Add ops.

06f5f40

resnet demo

274fa28

cleanup code.

9554410

YuchenJin mentioned this pull request Feb 2, 2022

[WIP] Relay->Relax translator (dynamic Bert forward/backward example) #63

Closed

YuchenJin force-pushed the relax-resnet branch from dbe2e50 to c34048e Compare February 2, 2022 19:26

Rebase.

c34048e

ZihengJiang reviewed Feb 4, 2022

View reviewed changes

yongwww mentioned this pull request Feb 4, 2022

Add metadata section, support constant and metadata in parser & printer #76

Merged

yongwww reviewed Feb 4, 2022

View reviewed changes

apps/relax_examples/resnet.py Show resolved Hide resolved

hypercubestart reviewed Feb 4, 2022

View reviewed changes

python/tvm/relax/testing/relay_translator.py Outdated Show resolved Hide resolved

hypercubestart reviewed Feb 4, 2022

View reviewed changes

python/tvm/relax/testing/relay_translator.py Outdated Show resolved Hide resolved

python/tvm/relax/testing/relay_translator.py Outdated Show resolved Hide resolved

Address comments.

99a07c1

yongwww approved these changes Feb 9, 2022

View reviewed changes

hypercubestart approved these changes Feb 10, 2022

View reviewed changes

yongwww reviewed Feb 11, 2022

View reviewed changes

tqchen reviewed Feb 11, 2022

View reviewed changes

include/tvm/relax/topi.h Outdated Show resolved Hide resolved

tqchen reviewed Feb 11, 2022

View reviewed changes

src/relax/topi/topi.cc Outdated Show resolved Hide resolved

YuchenJin added 2 commits February 13, 2022 11:08

leverage FTVMCompute for most op translation; reuse relay.Constant.

ffd2e06

lint.

b95a7e5

tqchen approved these changes Feb 14, 2022

View reviewed changes

yongwww merged commit 27c1ea2 into relax Feb 14, 2022

YuchenJin deleted the relax-resnet branch February 14, 2022 17:47

ZihengJiang reviewed Feb 14, 2022

View reviewed changes

MasterJH5574 pushed a commit to MasterJH5574/tlc-relax that referenced this pull request Jan 1, 2023

[Op] cast, wrap_param (tlc-pack#75)

371658e

vinx13 pushed a commit to vinx13/relax that referenced this pull request Jan 9, 2023

[Op] cast, wrap_param (tlc-pack#75)

1bd2329

MasterJH5574 pushed a commit to MasterJH5574/tlc-relax that referenced this pull request Jan 13, 2023

[Op] cast, wrap_param (tlc-pack#75)

39dcffa

MasterJH5574 pushed a commit to MasterJH5574/tlc-relax that referenced this pull request Jan 16, 2023

[Op] cast, wrap_param (tlc-pack#75)

c576371

vinx13 pushed a commit to vinx13/relax that referenced this pull request Jan 31, 2023

[Op] cast, wrap_param (tlc-pack#75)

46d952f

vinx13 pushed a commit to vinx13/relax that referenced this pull request Jan 31, 2023

[Op] cast, wrap_param (tlc-pack#75)

e759db6

vinx13 pushed a commit to vinx13/relax that referenced this pull request Feb 8, 2023

[Op] cast, wrap_param (tlc-pack#75)

608f398

MasterJH5574 pushed a commit to MasterJH5574/tlc-relax that referenced this pull request Feb 12, 2023

[Op] cast, wrap_param (tlc-pack#75)

fa09a5b

vinx13 pushed a commit to vinx13/relax that referenced this pull request Feb 13, 2023

[Op] cast, wrap_param (tlc-pack#75)

9b5e90c

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Relay->Relax translator (ResNet example) #75

Relay->Relax translator (ResNet example) #75

YuchenJin commented Feb 2, 2022

ZihengJiang Feb 4, 2022

YuchenJin Feb 4, 2022

ZihengJiang Feb 10, 2022

YuchenJin Feb 13, 2022

ZihengJiang Feb 4, 2022

YuchenJin Feb 4, 2022

YuchenJin Feb 4, 2022

ZihengJiang Feb 10, 2022 •

edited

ZihengJiang Feb 11, 2022

YuchenJin Feb 11, 2022

yongwww Feb 11, 2022 •

edited

ZihengJiang Feb 11, 2022 •

edited

YuchenJin Feb 13, 2022

yongwww left a comment

hypercubestart left a comment

yongwww Feb 11, 2022

YuchenJin Feb 11, 2022

Robslhc commented Feb 11, 2022

YuchenJin commented Feb 11, 2022

yongwww commented Feb 14, 2022

ZihengJiang left a comment

ZihengJiang Feb 14, 2022

yongwww Feb 14, 2022

ZihengJiang Feb 14, 2022

YuchenJin Feb 14, 2022

Relay->Relax translator (ResNet example) #75

Relay->Relax translator (ResNet example) #75

Conversation

YuchenJin commented Feb 2, 2022

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

ZihengJiang Feb 10, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yongwww Feb 11, 2022 • edited

Choose a reason for hiding this comment

ZihengJiang Feb 11, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yongwww left a comment

Choose a reason for hiding this comment

hypercubestart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Robslhc commented Feb 11, 2022

YuchenJin commented Feb 11, 2022

yongwww commented Feb 14, 2022

ZihengJiang 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

ZihengJiang Feb 10, 2022 •

edited

yongwww Feb 11, 2022 •

edited

ZihengJiang Feb 11, 2022 •

edited