Skip to content

[Relay] Refactor Interpreter to treat lowering as IRModule->IRModule rewrite. - #8597

Merged
jroesch merged 2 commits into
apache:mainfrom
mbs-octoml:mbs-tecompiler-refactoring
Aug 17, 2021
Merged

[Relay] Refactor Interpreter to treat lowering as IRModule->IRModule rewrite.#8597
jroesch merged 2 commits into
apache:mainfrom
mbs-octoml:mbs-tecompiler-refactoring

Conversation

@mbs-octoml

@mbs-octoml mbs-octoml commented Jul 30, 2021

Copy link
Copy Markdown
Contributor

This continues the work outlined in
TECompiler: Rewrite existing compile engine to match updated compiler flow to centralize the TE->TIR lowering logic into a single IRModule->IRModule transformation. This gets about halfway there for the Interpreter:

  • Remove direct access to TECompiler from interpreter, and instead call tec::LowerTEExpr when 'preparing' a module and expression for evaluation.
  • Make clear there's no phase distinction between create_interpreter and evaluate on the Python side -- both must be prepared together as a single IRModule.
  • But in return make sure the result of evaluate on the Python side is a packed func ready to directly apply 'simple' arguments to an already interpreted closure.
  • The interpreter builds and caches primitive TIR functions (and their corresponding dynamic shape functions) as packed funcs as they are encountered.
  • Clean uses of interpreter for constant folding on the C++ side.

Future work:

  • Fold LoweredModule into IRModule so tec::LowerTEExpr is just another pass.
  • Get rid of the implicit caching of lowered functions in TECompiler.
  • Make calling convention from Relay to TIR explicit, and remove all the function attribute hackery currently needed so the interpreter can correctly invoke lowered functions as it encounters them.
  • Make TECompiler private. Though could do this now it will make migrating the VM and AOT uses of CompilerEngine harder.

@mbs-octoml mbs-octoml changed the title Rough cut and Rough cut at removing direct dependence of interpreter on TECompiler. Jul 30, 2021
@mbs-octoml

mbs-octoml commented Jul 30, 2021

Copy link
Copy Markdown
Contributor Author

@jroesch and @electriclilies for review

@mbs-octoml
mbs-octoml marked this pull request as draft July 30, 2021 00:14
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch from 9b894d7 to feb036d Compare July 30, 2021 18:34
@mbs-octoml mbs-octoml changed the title Rough cut at removing direct dependence of interpreter on TECompiler. [Relay] Rough cut at removing direct dependence of interpreter on TECompiler. Jul 30, 2021
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch 2 times, most recently from 34e00fb to 01c73c5 Compare August 3, 2021 15:51
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch 3 times, most recently from bc233cd to 020af1a Compare August 5, 2021 23:28
Comment thread tests/python/relay/test_backend_interpreter.py Outdated
Comment thread tests/python/relay/test_backend_interpreter.py Outdated
Comment thread tests/python/relay/test_backend_interpreter.py Outdated
Comment thread tests/python/relay/test_backend_interpreter.py Outdated
Comment thread include/tvm/ir/module.h
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch from a602e4b to 1f741bd Compare August 6, 2021 21:09
@mbs-octoml

Copy link
Copy Markdown
Contributor Author

@jroesch PTAL:

  • BUG: I was compiling dynamic shape functions for the target even though they must always stay on the cpu. Fixed, but it means a bit more bookkeeping in Interpreter.
  • Cleaned up flattening of tuples, both for TupleType/TensorType and runtime::ADT/runtime::NDArray.
    Thanks!

Comment thread src/relay/backend/interpreter.cc
Comment thread src/relay/backend/interpreter.cc
Comment thread src/relay/backend/te_compiler.cc
Comment thread src/relay/backend/te_compiler_cache.h

@jroesch jroesch left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

LGTM, just some stuff we should do as follow up.

@electriclilies electriclilies 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 all your hard work Mark! The PR looks good to me so far, I really appreciate the high-quality comments and documentation. I'm still working my way through because it is quite long :)

@electriclilies electriclilies 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! Thanks again :)

@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch from 1f741bd to 99951d8 Compare August 8, 2021 15:25
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch 3 times, most recently from 8a83600 to ccd62ef Compare August 11, 2021 22:01
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch 2 times, most recently from d8fca6c to f1fb28f Compare August 16, 2021 22:14
  https://discuss.tvm.apache.org/t/rfc-relay-tecompiler-rewrite-existing-compile-engine-to-match-updated-compiler-flow/9233
This gets about halfway there for the Interpreter:

* Remove direct access to TECompiler from interpreter, and instead call
  tec::LowerTEExpr when 'preparing' a module and expression for evaluation.
* Make clear there's no phase distinction between create_interpreter and
  evaluate on the Python side -- both must be prepared together as a single IRModule.
* But in return make sure the result of evaluate on the Python side is a packed func
  ready to directly apply 'simple' arguments to an already interpreted closure.
* The interpreter builds and caches primitive TIR functions (and their corresponding
  dynamic shape functions) as packed funcs as they are encountered.
* Cleanup uses of interpreter for constant folding on the C++ side.

Future work:
* Fold LoweredModule into IRModule so tec::LowerTEExpr is just another pass.
* Get rid of the implicit caching of lowered functions in TECompiler.
* Make calling convention from Relay to TIR explicit, and remove all the function
  attribute hackery currently needed so the interpreter can correctly invoke lowered
  functions as it encounters them.
* Make TECompiler private. Though could do this now it will make migrating the VM and
  AOT uses of CompilerEngine harder.

Force a gc between sphinx-gallery items to reclaim GPU memory. (apache#8722)

GPU memory is only released once the PackedFunc for evaling the model is gced
by Python. In CI we're noticing intermittent 'CUDA: Out of memory' failures
while processing the tutorials, and tracing showed there was no gc happening
between items. Not confident this will solve the problem but worth a try.
@mbs-octoml
mbs-octoml force-pushed the mbs-tecompiler-refactoring branch from f1fb28f to a249ea8 Compare August 17, 2021 16:20
@jroesch
jroesch merged commit c8a6089 into apache:main Aug 17, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants