Change Call with TIRCallAttrs to call_lowered op by electriclilies · Pull Request #9312 · apache/tvm

electriclilies · 2021-10-18T23:54:26Z

In this PR, I change Call with TIRCallAttrs to a relay op called call_lowered. (I also rename CallTIRAttrs to CallLoweredAttrs). This means that after lowering, calls are of the form Call("call_lowered", [fn: GlobalVar, args: Tuple] Attrs(CallLoweredAttrs)).

One benefit to this approach is it is easy to identify calls to lowered functions-- instead of checking the type of the call's attributes, we can just check what the op is. This is especially helpful for passes that run both before an after lowering, like device planning and memory allocation.

I left the calling convention for extern functions as-is. A codegened extern call looks like this: Call(fn: GlobalVar, [arg1: Expr, ... argn: Expr], Attrs()) with null attributes. In the future we will probably need to change the extern calling convention to have attributes that store information like the target, but that is outside the scope of this work.

electriclilies · 2021-11-03T20:49:04Z

@mbs-octoml PTAL, I'm pretty sure I've resolved all the issues, and this is ready for review.
One thing I'm unsure of is that I don't think I need the explicit label for extern functions anymore-- it would be easy to take out but I could see it being useful in the future. LMK what you think about that.

mbs-octoml · 2021-11-04T18:47:58Z

In the overview comment prob should show that args must be a tuple (even for the zero and one arity case)

mbs-octoml

Thanks Lily. Mostly nits and a few of my personal style preferences (feel free to ignore those), so don't panic! Otherwise:

think we can tune up the Extract helper.
a bit worries about hiding the lowered vs original forms of the operators (eg device_copy), see comment inline.
feeling skittish about folding extern calls into the call_lowered convention just yet and suggest we back out of that.
Happy to look again on a PTAL.

mbs-octoml · 2021-11-04T18:30:14Z

+ */
+
+/*!
+ * \file tvm/relay/attrs/annotation.h


nit: update comment

mbs-octoml · 2021-11-04T18:43:52Z

+namespace tvm {
+namespace relay {
+
+Expr CallLowered(Expr func, Expr inputs, Attrs attrs, Array<Type> type_args, Span span);


nit: comments thanks!

mbs-octoml · 2021-11-04T18:45:20Z

+
+Expr CallLowered(Expr func, Expr inputs, Attrs attrs, Array<Type> type_args, Span span);
+const Op& CallLoweredOp();
+std::pair<GlobalVar, Array<Expr>> ExtractFunctionAndArgs(const CallNode* call_node);


nit: it's a lot of boilerplate i know, but a struct for the return result is probably best long term.
super duper nit: for on_device and device_copy i called these Get.... Do I get to claim precedence?

mbs-octoml · 2021-11-04T18:53:34Z

+ const TypeReporter& reporter) {
+ // Types = [func, args, ret_type]
+ ICHECK_EQ(types.size(), 3u);
+ auto func_type = types[0].as<FuncTypeNode>();


return false for these checks since we know nothing about how the expression was constructed.
even for the arity check i think?

and then can just the Assign for the TupleType do the check for types[1]

mbs-octoml · 2021-11-04T19:05:45Z

+
+const Op& CallLoweredOp() { return Op::Get("call_lowered"); }
+
+Expr CallLowered(Expr func, Expr inputs, Attrs attrs, Array<Type> type_args, Span span) {


nit: for symmetry with the Extract, inputs can be Array args.
You can std::move into the Call ctor (kinda a mico optimization, but good habit to get into to for when the perf really does matter).

mbs-octoml · 2021-11-04T20:33:45Z

- return {ext_func->prim_fn_var, Attrs()};
+ auto call_lowered_attrs = make_object<CallLoweredAttrs>();
+ // Mark the function as a extern function so that AOT knows what to do with
+ call_lowered_attrs->metadata.Set("extern_func", Integer(1));


I see the check in the graph_executor_codegen.cc but not aot_executor_codegen.cc. I wonder if we should leave the extern calls as regular old boring calls with no attributes and tackle them separately. Perhaps this method can return the call instead of pair so that the different call representations can be handled.

mbs-octoml · 2021-11-04T20:36:52Z

 } else {
 return {call_node->args[0], src_dev_type, dst_dev_type};
 }
+ } else if (call_node->op == CallLoweredOp()) {


I'm a bit nervous about this one since there's a lot of usage patterns such as
props = GetProps
if (prop.body.defined()) { .... MakeDeviceCopy(...) ... }
which would silently demote call_lowered copies to regular copies.
I think safer would be to have a GetLoweredDeviceCopyProps so it's always clear which form is being matched against.

Talked with mark and we determined this is fine, he was just being a bit skittish

I think that this chunk of code is not necessary since I also deal with CallLoweredOp in GetPrimitiveDeviceCopyProps, I'm not sure why I put it in. Removing it

mbs-octoml · 2021-11-04T20:37:56Z

 * \brief Dialect operators for Relay VM.
 */

-#include "vm.h"


nit: might as well keep it -- standard idiom even though not strictly needed.

mbs-octoml · 2021-11-04T20:39:07Z

- if (tir_call_attrs->metadata.count("source_device") != 1 ||
- tir_call_attrs->metadata.count("dst_device") != 1) {
- return {};
+ if (call_node->op == CallLoweredOp()) {


I think this would become the GetLoweredDeviceCopyProps mentioned above, so no need for the helper here.

we decided we don't need that actually so this is OK as is.

mbs-octoml · 2021-11-04T20:40:40Z

 const Op& alloc_tensor_op = Op::Get("memory.alloc_tensor");
 const Op& shape_of_op = Op::Get("vm.shape_of");
 const Op& invoke_tvm_op = Op::Get("vm.invoke_tvm_op");
+ const Op& call_lowered = Op::Get("call_lowered");


nit: your CallLoweredOp() helper is good enough.
some day all these will use the same idiom.

electriclilies · 2021-11-09T19:45:19Z

@mbs-octoml I think this is ready for another review!

Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op

mbs-octoml

Two minor suggestions which are easily done in follow up. Thanks!

mbs-octoml · 2021-11-10T18:57:49Z

 target = GetTargetFromInteger(device_type, targets_);
 }
-
+ Array<Expr> visited_args;


not a correctness issue, but we are revisiting the args (bad rebase).

mbs-octoml · 2021-11-10T19:01:28Z

+ /* Get device props for a TIR function */
+ CallLoweredProps call_lowered_props = GetCallLoweredProps(call_node);
+
+ if (call_lowered_props.attrs.metadata.count("source_device") == 1 &&


confirming looks like we can avoid this here
(eventually GetPrimitiveDeviceCopyProps from device_domains.cc will need to get moved alongside this one (and renamed) but no need here)

mbrookhart · 2021-11-10T22:20:10Z

Thanks @electriclilies @mbs-octoml

electriclilies · 2021-11-10T22:50:07Z

@mbrookhart @mbs-octoml thanks!

electriclilies · 2021-11-10T23:44:36Z

Followed up in #9491

* main: (119 commits) [Topi][Op][PyTorch][Vitas] Fix inconsistent kernel layout conventions for conv2d_transpose (apache#9336) Fix repository URL in ubuntu_install_rocm.sh (apache#9425) Add LLVM-13 installation to Docker setup (apache#9498) [Relay] Use target_host determined at Relay level instead of recalculating it (apache#9499) Arm(R) Ethos(TM)-U NPU BinaryElementwise operators support (apache#9442) [COMMUNITY] Junru's and Wuwei's PGP key for ASF release (apache#9488) Add default for split op (apache#9489) [HOTFIX][TARGET] Change LOG in compilation config to DLOG (apache#9486) Fixed some warnings about lambda's closures that are bigger than necessary (apache#9481) [Support] Add libinfo into the runtime build (apache#9310) Change Call with TIRCallAttrs to call_lowered op (apache#9312) [ETHOSN] Streamline Ethos(TM)-N cross-compile rpc usage (apache#9477) [CMSIS-NN] Assert correct amount of CMSIS-NN artifacts in MLF (apache#9480) [MicroTVM][PyTest] Explicitly skip MicroTVM unittests. (apache#9335) [microNPU] Replace ICHECK with diagnostic context in type inference (apache#9470) Better host handling in CompilationConfig & debug printing (apache#9460) [AOT][Tests] Use pre-built libraries in Reference System tests (apache#9271) [TIR] Add type hint for TIR (apache#9432) [TVMC] Add test for quantized pytorch model (apache#9467) [CMSIS-NN] Convert CMSIS-NN to use Target Hooks (apache#9397) ...

* followups from #9312 * remove unneeded moves

* Introduce call_lowered op Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op * lint * Fix AOT to deal with externs * add const auto& * Fix aot crt test

* followups from apache#9312 * remove unneeded moves

* Introduce call_lowered op Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op * lint * Fix AOT to deal with externs * add const auto& * Fix aot crt test

* followups from apache#9312 * remove unneeded moves

* Introduce call_lowered op Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op * lint * Fix AOT to deal with externs * add const auto& * Fix aot crt test

* followups from apache#9312 * remove unneeded moves

* Introduce call_lowered op Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op * lint * Fix AOT to deal with externs * add const auto& * Fix aot crt test

* followups from apache#9312 * remove unneeded moves

* Introduce call_lowered op Add op vm.call_tir Change from checking if CallNode has CallTIRAttrs to checking if the Op is vm.call_tir Change device_domains to use vm.call_tir op more explicitly Fixed issue in type checker, now have seg fault :( Fix typo -- most of VM tests pass now Interpreter now deals with call_tir properly Fix typo in te_compiler Use InvokeTVMOp and CallTIR Add some checks to graph_plan_memory.cc Make GetToken skip function types C++ TESTS PASS WOOHOO Remove prints formatting vm.call_tir -> call_tir and more comment removals call_tir -> call_lowered fix lint clang format Remove compute from non computational vm ops missed some semicolons in prev commit Fix warning Move call_lowered to relay/op/call/call.cc and rename util func Add helper fn that returns lowered_call op fix import order clang format Add constraint to call_lowered type rel clean up empty token vector comment Move CallTIRAttrs to include/tvm/relay/attrs/call.h Rename TIRCallAttrs as CallLoweredAttrs lint Add helper for extracting func and args from call_lowered Change graph_executor_codegen to use helper function Update interpreter to use helper Fix device_domains.cc -- could still use cleanup, also I am not sure why there are still direct calls to primfns in DomainforCallee Clean up DeviceCopyProps and lint lint return CallLoweredAttrs with the extern func comment note in comment Progress & notes. Realized that I am not handling externs correctly not sure why this ever worked before? Clean up CreateFuncCall signature, notes comments Fix extern function handling extern_function -> extern_func fix DeviceAwareVisitExpr_ -- now it handles both lowered and normal calls yay passes AOT tests! formatting and comment removal cleanup Introduce call_lowered op * lint * Fix AOT to deal with externs * add const auto& * Fix aot crt test

* followups from apache#9312 * remove unneeded moves

mbs-octoml reviewed Oct 19, 2021
View reviewed changes

Comment threadsrc/relay/op/vm/vm.cc Outdated
Comment threadsrc/relay/backend/graph_plan_memory.cc Outdated
Comment threadsrc/relay/op/vm/vm.cc Outdated

electriclilies force-pushed the call_tir_cleaned_up branch 2 times, most recently from 9435c6d to 5ebbe78Compare October 26, 2021 23:48

electriclilies marked this pull request as ready for review October 26, 2021 23:54

electriclilies requested review from MarisaKirisame, ZihengJiang, anijain2305, jroesch, junrushao, jwfromm, mbrookhart, slyubomirsky, vinx13, wweic, yzhliu and zhiics as code owners October 26, 2021 23:55

electriclilies changed the title ~~[DRAFT] Change Call with TIRCallAttrs to call_lowered op~~Change Call with TIRCallAttrs to call_lowered opOct 27, 2021

electriclilies requested review from comaniac, icemelon, manupak, mbaret and trevor-m as code owners November 3, 2021 00:04

electriclilies force-pushed the call_tir_cleaned_up branch 2 times, most recently from e8c644a to 831879cCompare November 3, 2021 20:47

mbs-octoml suggested changes Nov 4, 2021
View reviewed changes

electriclilies force-pushed the call_tir_cleaned_up branch 3 times, most recently from 5861dc6 to d5a33f9Compare November 9, 2021 19:42

electriclilies added 5 commits November 9, 2021 21:46

lint
8960dcd

Fix AOT to deal with externs
e0e50fa

add const auto&
2dd601b

Fix aot crt test
207c487

electriclilies force-pushed the call_tir_cleaned_up branch from fddc231 to 207c487Compare November 10, 2021 05:48

mbs-octoml approved these changes Nov 10, 2021
View reviewed changes

mbrookhart approved these changes Nov 10, 2021
View reviewed changes

mbrookhart merged commit 0812c07 into apache:mainNov 10, 2021

electriclilies deleted the call_tir_cleaned_up branch November 10, 2021 22:38

electriclilies added a commit to electriclilies/tvm that referenced this pull request Nov 10, 2021
followups fromapache#9312
9efc217

electriclilies mentioned this pull request Nov 10, 2021
Followup from #9312 (Introduce call_lowered op) #9491
Merged

electriclilies added a commit to electriclilies/tvm that referenced this pull request Nov 12, 2021
followups fromapache#9312
8e2e7fd

masahi pushed a commit that referenced this pull request Nov 13, 2021
Followup from#9312(Introduce call_lowered op) (#9491)
e4904c6
* followups from #9312 * remove unneeded moves

mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
Followup fromapache#9312(Introduce call_lowered op) (apache#9491)
087870a
* followups from apache#9312 * remove unneeded moves

mehrdadh pushed a commit to mehrdadh/tvm that referenced this pull request Dec 1, 2021
Followup fromapache#9312(Introduce call_lowered op) (apache#9491)
31679b2
* followups from apache#9312 * remove unneeded moves

electriclilies mentioned this pull request Dec 5, 2021
[DISCUSS] call_dps renaming tlc-pack/relax#56
Closed

ylc pushed a commit to ylc/tvm that referenced this pull request Jan 7, 2022
Followup fromapache#9312(Introduce call_lowered op) (apache#9491)
240a220
* followups from apache#9312 * remove unneeded moves

yangulei pushed a commit to yangulei/tvm that referenced this pull request Jan 11, 2022
Followup fromapache#9312(Introduce call_lowered op) (apache#9491)
dfb8045
* followups from apache#9312 * remove unneeded moves

ylc pushed a commit to ylc/tvm that referenced this pull request Jan 13, 2022
Followup fromapache#9312(Introduce call_lowered op) (apache#9491)
0a6ffac
* followups from apache#9312 * remove unneeded moves


		const Op& CallLoweredOp() { return Op::Get("call_lowered"); }

		Expr CallLowered(Expr func, Expr inputs, Attrs attrs, Array<Type> type_args, Span span) {

Uh oh!

Conversation

electriclilies commented Oct 18, 2021• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

electriclilies commented Nov 3, 2021

Uh oh!

mbs-octoml commented Nov 4, 2021

Uh oh!

mbs-octoml left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

electriclilies commented Nov 9, 2021

Uh oh!

mbs-octoml left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mbrookhart commented Nov 10, 2021

Uh oh!

electriclilies commented Nov 10, 2021

Uh oh!

electriclilies commented Nov 10, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

electriclilies commented Oct 18, 2021•
edited
Loading