[Relay] Prepare for merging context_analysis.cc and device_annotation.cc #9077

mbs-octoml · 2021-09-22T19:51:54Z

In preparation for #9038:

Improve construction and deconstruction of "on_device" and "device_copy" calls since they will be center stage.
Move "device_copy" support out of memory.h into own module to mirror "on_device".
Clearing out some DLOG -> VLOG changes I found helped me debug.
Clearing out some whitespace-only changes I accumulated.
Cleanup Device aliases and make sure we have some global common non-standard DLDeviceTypes.

- Improve construction and deconstruction of "on_device" and "device_copy" calls since they will be center stage. - Move "device_copy" support out of memory.h into own module to mirror "on_device". - Clearing out some DLOG -> VLOG changes I found helped me debug. - Clearing out some whitespace-only changes I accumulated.

Mousius

Looks good @mbs-octoml, thanks for splitting this out from the other PR - overall I think I understand where this is going, but I've left a few questions and suggestions to clarify 😸

Mousius · 2021-09-23T13:31:22Z

include/tvm/relay/attrs/annotation.h

+
+  /*! \brief Device type on which each of the function's arguments already resides. */
+  Array<Integer> param_device_types;
+  // TODO(mbs): Replace device types with TargetDevice.


Do we need these TODOs, I'm assuming you're also tracking this work elsewhere?

I'd like to be able to freely write "TODO(issues/1234)" so everything is cross linked. That would help anyone picking up the issue get a head start on where to look in the code, as well as just remind us all where things are likely to shift. That got a bit tangled up with questions about how we should reconcile GitHub and our internal OctoML Jira issues, and the granularity of GitHub issues, but I think we've decided to just switch to GitHub. I'm guessing you have the same dilemma at Arm? Anyway these "TODO(mbs)" is just me leaving enough cookie crumbs so I can go back and do it properly.

I think this might be the start of quite a long discussion, I've seen quite a lot of abandoned TODO comments so the best approach (in my opinion) would be either GitHub issues to track all the work or just tracking internally. In this case there may be contention as to whether we use TargetDevice or something else, so if we leave this for a prolonged period people coming to this code may be confused. One of my favourite examples of this is bugs which get fixed in an entirely different area leaving TODO(Mousius) - Fix bug which we will never get rid of.

Having said all of that, I think these will get cleaned up pretty quickly and it's probably just my thinking - so let's not block anything on it 😸

Mousius · 2021-09-23T13:33:15Z

src/relay/op/memory/device_copy.h

+ */
+
+/*!
+ * \file relay/attrs/device_copy.h


Suggested change

* \file relay/attrs/device_copy.h

* \file relay/op/memory/device_copy.h

done.
begs the question why we have these?

(also for annotation.h)

The same reason we have type hints and types in our Python comments? 😸

Mousius · 2021-09-23T13:37:51Z

include/tvm/relay/attrs/annotation.h

+ * function parameters and result.
+ */
+struct FunctionOnDeviceAttrs : public tvm::AttrsNode<FunctionOnDeviceAttrs> {
+  constexpr static const char* kFunctionAttrsKey = "on_device";


Historically these would go into include/tvm/ir/function.h as the list of attributes on a function, is there a reason for them being here instead?

Just ignorance of the convention! Done, thx.

Hmm, except I'd like to keep this scoped to relay::, so I moved it into it's own include/tvm/attrs/function.h since ir/function.h is at the BaseFunc 'common to both Relay and TIR' scope. Though that's yet another duplicated basename which I don't like much.

I don't have an immediate solution, but this looks at least neater 😸

Mousius · 2021-09-23T13:39:30Z

include/tvm/runtime/ndarray.h

+// as a singleton target map indexed by the invalid DLDeviceType '0'.
+constexpr DLDeviceType kNullDeviceType = static_cast<DLDeviceType>(0);
+
+// An 'invalid' device type, does not correspond to any DLDeviceType enum.


Should these not be added as valid enum variants to dlpack.h ? There's a risk here that dlpack will change and it'll become incompatible?

I don't know. Note that these special values are endemic in the code (often stored in naked ints) so I'm at least not introducing a new convention. I think the argument for this approach is the enum values should always be valid at runtime and these invalid-but-distinguished values are just a compile time concept.

I this as a stepping stone to planning with hybrid targets/devices rather than just device or device types, in which case these distinguished values can go away.

Roger roger.

Mousius · 2021-09-23T13:40:06Z

python/tvm/relay/op/annotation/annotation.py

+    return _make.on_device(data, _device_to_int(device), is_fixed)
+
+
+# for testing only


If this is for testing, can it not be moved to the relevant test file?

Woops, I changed my mind on this and think it's a legit user-visible annotation, but forgot to update the comment. Done.

Mousius · 2021-09-23T13:41:22Z

python/tvm/relay/op/annotation/annotation.py

+        return device.device_type
+    if isinstance(device, str):
+        return _nd.device(device).device_type
+    raise ValueError(


I think this is just refactored from the below, but do we have a test for ValueError being raised?

Can't find any existing unit tests for these so added some.

Mousius · 2021-09-23T13:45:35Z

src/relay/op/annotation/annotation.cc

+  if (expr->IsInstance<OpNode>() || expr->IsInstance<GlobalVarNode>() ||
+      expr->IsInstance<VarNode>() || expr->IsInstance<ConstructorNode>()) {


May be worth documenting this condition in code:

Suggested change

if (expr->IsInstance<OpNode>() || expr->IsInstance<GlobalVarNode>() ||

expr->IsInstance<VarNode>() || expr->IsInstance<ConstructorNode>()) {

bool is_device_polymorphic = expr->IsInstance<OpNode>() || expr->IsInstance<ConstructorNode>();

bool has_implied_device = expr->IsInstance<VarNode>() || expr->IsInstance<GlobalVarNode>();

if (is_device_polymorphic || has_implied_device) {

Done, though I used comments and early returns rather than bool assignment since I personally find that clearer.

I'll avoid the clean code vs comments holy war here, it looks good now 😸

mbs-octoml · 2021-09-23T15:41:37Z

Note the CI failure is unrelated and fixed by #9076

jroesch · 2021-09-23T16:29:08Z

@mbs-octoml I already did a pass on the other PR, LGTM, I think let's quickly address (on this PR) @Mousius suggestions and keep it rolling for next PR.

Some stray py formatting changes snuck in since I just run black . at the root.

mbs-octoml · 2021-09-23T21:02:53Z

Thank you @Mousius & @jroesch.

Mousius

Overall LGTM, follow ups will address many of the concerns so we should move to them 😸

….cc (apache#9077) * [Relay] Prepare for merging context_analysis.cc and device_annotation.cc - Improve construction and deconstruction of "on_device" and "device_copy" calls since they will be center stage. - Move "device_copy" support out of memory.h into own module to mirror "on_device". - Clearing out some DLOG -> VLOG changes I found helped me debug. - Clearing out some whitespace-only changes I accumulated. * [checkpoint] Address Christopher's comments. Some stray py formatting changes snuck in since I just run black . at the root.

mbs-octoml requested review from anijain2305, areusch, comaniac, Hzfengsy, jroesch, junrushao, jwfromm, kazum, kparzysz-quic, liangfu, MarisaKirisame, masahi, mbrookhart, slyubomirsky, tmoreau89, tqchen, vinx13, wweic, yzhliu, zhiics and ZihengJiang as code owners September 22, 2021 19:51

mbs-octoml changed the title ~~[Relay] Prepare for merging context_analysis.cc and device_annotation.cc in #9038~~ [Relay] Prepare for merging context_analysis.cc and device_annotation.cc Sep 22, 2021

mbs-octoml mentioned this pull request Sep 22, 2021

[Relay] Merge analysis/context_analysis.cc and transforms/device_annotation.cc #9038

Merged

Mousius requested changes Sep 23, 2021

View reviewed changes

jroesch approved these changes Sep 23, 2021

View reviewed changes

[checkpoint] Address Christopher's comments.

65529cf

Some stray py formatting changes snuck in since I just run black . at the root.

mbs-octoml requested a review from merrymercy as a code owner September 23, 2021 21:02

Mousius approved these changes Sep 24, 2021

View reviewed changes

masahi merged commit 5871608 into apache:main Sep 24, 2021

mbs-octoml deleted the mbs-prepare-device-planning branch September 24, 2021 21:19

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relay] Prepare for merging context_analysis.cc and device_annotation.cc #9077

[Relay] Prepare for merging context_analysis.cc and device_annotation.cc #9077

mbs-octoml commented Sep 22, 2021 •

edited

Loading

Mousius left a comment

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 24, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 24, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 24, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 24, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 23, 2021

mbs-octoml Sep 23, 2021

Mousius Sep 24, 2021

mbs-octoml commented Sep 23, 2021

jroesch commented Sep 23, 2021 •

edited

Loading

mbs-octoml commented Sep 23, 2021

Mousius left a comment

	* \file relay/attrs/device_copy.h
	* \file relay/op/memory/device_copy.h

		return _make.on_device(data, _device_to_int(device), is_fixed)


		# for testing only

		if (expr->IsInstance<OpNode>() \|\| expr->IsInstance<GlobalVarNode>() \|\|
		expr->IsInstance<VarNode>() \|\| expr->IsInstance<ConstructorNode>()) {

-  if (expr->IsInstance<OpNode>() || expr->IsInstance<GlobalVarNode>() ||
-      expr->IsInstance<VarNode>() || expr->IsInstance<ConstructorNode>()) {
+  bool is_device_polymorphic = expr->IsInstance<OpNode>() || expr->IsInstance<ConstructorNode>();
+  bool has_implied_device = expr->IsInstance<VarNode>() || expr->IsInstance<GlobalVarNode>();
+  if (is_device_polymorphic || has_implied_device) {

[Relay] Prepare for merging context_analysis.cc and device_annotation.cc #9077

[Relay] Prepare for merging context_analysis.cc and device_annotation.cc #9077

Conversation

mbs-octoml commented Sep 22, 2021 • edited Loading

Mousius 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

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

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

mbs-octoml commented Sep 23, 2021

jroesch commented Sep 23, 2021 • edited Loading

mbs-octoml commented Sep 23, 2021

Mousius left a comment

Choose a reason for hiding this comment

mbs-octoml commented Sep 22, 2021 •

edited

Loading

jroesch commented Sep 23, 2021 •

edited

Loading