Skip to content

[microNPU] Refactor Relay to TIR hook#10599

Merged
manupak merged 4 commits into
apache:mainfrom
lhutton1:refactor-relay-to-tir-hook
Mar 21, 2022
Merged

[microNPU] Refactor Relay to TIR hook#10599
manupak merged 4 commits into
apache:mainfrom
lhutton1:refactor-relay-to-tir-hook

Conversation

@lhutton1

Copy link
Copy Markdown
Contributor

Refactors the Relay to TIR python hook for the NPU so that optimizations can be applied across the whole module and not just functions that will be offloaded to the NPU. A pass OutlineCompilerFunctions is introduced to outline NPU functions, which now happens before optimization passes are run (this previously happened after the prim_func had been created).

In addition, optimization passes that should only run on NPU functions are now limited to running on outlined functions for the NPU (by checking the "Compiler" attribute). To help avoid code duplication, a helpful decorator npu_pass has been created for python passes that should only run on NPU functions.

This refactor helps move a number of passes in the microNPU codegen to use an IRModule -> IRModule philosophy.

cc @manupa-arm @ekalda @NicolaLancellotti @dchauhan-arm

@github-actions github-actions Bot requested a review from manupak March 14, 2022 11:08

@manupak manupak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An early comment -- I ll do a full pass later.

Comment thread tests/python/contrib/test_ethosu/test_outline_compiler_functions.py
@github-actions github-actions Bot requested a review from manupak March 14, 2022 15:12
@lhutton1 lhutton1 force-pushed the refactor-relay-to-tir-hook branch from e95d9f2 to cf90859 Compare March 14, 2022 15:30

@manupak manupak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the refactor. Overall it looks good.

Just two comments regarding the decorator name and I think we can avoid needing to declare to RelayToTIR() passes.

Comment thread python/tvm/relay/backend/contrib/ethosu/util.py Outdated
Comment thread python/tvm/relay/backend/contrib/ethosu/codegen.py Outdated
Comment thread src/relay/backend/contrib/ethosu/codegen.cc
Refactors the Relay to TIR python hook for the NPU so that optimizations
can be applied across the whole module and not just functions that will
be offloaded to the NPU. A pass `OutlineCompilerFunctions` is introduced
to outline NPU functions, which now happens before optimization passes
are run (this previously happened after the prim_func had been created).

In addition, optimization passes that should only run on NPU functions
are now limited to running on outlined functions for the NPU (by
checking the "Compiler" attribute). To help avoid code duplication, a
helpful decorator `create_npu_function_pass` has been created for python
passes that should only run on NPU functions.

This refactor helps move a number of passes in the microNPU codegen to
use an IRModule -> IRModule philosophy.

Change-Id: Icdea9ba43da0157d5ee17529d2b23b761396d112
Change-Id: I3ca48738e096bb0f4dc362f0e9550317fc0d5afd
This commit renames `npu_pass` -> `create_npu_function_pass`.

It also renames the `RelayToTIR` pass created in Python to `LowerToTIR`,
along with moving it to compiler.py to make it clear that this pass is a
wrapper around the `_lower_to_tir` function. In addition, to make it
explicit that the `lower_to_tir` func->func pass should not be used
directly it has been renamed to `_lower_to_tir` - it is being maintained
since it is used in many tests.

Change-Id: I3a0a06801f029aeaa4a51c2d86d8703bb0d7afbb
@lhutton1 lhutton1 force-pushed the refactor-relay-to-tir-hook branch from cf90859 to 92dee34 Compare March 17, 2022 13:54
@github-actions github-actions Bot requested a review from manupak March 17, 2022 13:55

@manupak manupak left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@ekalda @NicolaLancellotti please take a look..

@ekalda ekalda left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice work, looks good to me! :)

Comment thread python/tvm/relay/backend/contrib/ethosu/codegen.py Outdated

@NicolaLancellotti NicolaLancellotti left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@github-actions github-actions Bot requested a review from manupak March 18, 2022 15:02
@manupak

manupak commented Mar 18, 2022

Copy link
Copy Markdown
Contributor

Nice! I ll merge when the CI is green

Change-Id: I44c64de15fa8680cc89ce0440ffa6c9e0ec62a50
@lhutton1 lhutton1 force-pushed the refactor-relay-to-tir-hook branch from a12e806 to 94fd3cd Compare March 21, 2022 09:28
@manupak manupak merged commit 290395a into apache:main Mar 21, 2022
@manupak

manupak commented Mar 21, 2022

Copy link
Copy Markdown
Contributor

Thanks @lhutton1 @NicolaLancellotti @ekalda

@lhutton1 lhutton1 deleted the refactor-relay-to-tir-hook branch March 21, 2022 17:28
pfk-beta pushed a commit to pfk-beta/tvm that referenced this pull request Apr 11, 2022
* [microNPU] Refactor Relay to TIR hook

Refactors the Relay to TIR python hook for the NPU so that optimizations
can be applied across the whole module and not just functions that will
be offloaded to the NPU. A pass `OutlineCompilerFunctions` is introduced
to outline NPU functions, which now happens before optimization passes
are run (this previously happened after the prim_func had been created).

In addition, optimization passes that should only run on NPU functions
are now limited to running on outlined functions for the NPU (by
checking the "Compiler" attribute). To help avoid code duplication, a
helpful decorator `create_npu_function_pass` has been created for python
passes that should only run on NPU functions.

This refactor helps move a number of passes in the microNPU codegen to
use an IRModule -> IRModule philosophy.

Change-Id: Icdea9ba43da0157d5ee17529d2b23b761396d112

* add mixed compilers to test

Change-Id: I3ca48738e096bb0f4dc362f0e9550317fc0d5afd

* Address comments including renaming both npu_pass and RelayToTIR

This commit renames `npu_pass` -> `create_npu_function_pass`.

It also renames the `RelayToTIR` pass created in Python to `LowerToTIR`,
along with moving it to compiler.py to make it clear that this pass is a
wrapper around the `_lower_to_tir` function. In addition, to make it
explicit that the `lower_to_tir` func->func pass should not be used
directly it has been renamed to `_lower_to_tir` - it is being maintained
since it is used in many tests.

Change-Id: I3a0a06801f029aeaa4a51c2d86d8703bb0d7afbb

* address nit and small fix to example

Change-Id: I44c64de15fa8680cc89ce0440ffa6c9e0ec62a50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants