[Relay][VM] Add support for references.#6798
Conversation
| scope.ret(relay.RefRead(ref_create)) | ||
|
|
||
| f = relay.Function([], scope.get()) | ||
| print(f) |
There was a problem hiding this comment.
do we remove the print in the test?
| } | ||
| case Opcode::RefCreate: { | ||
| auto value = ReadRegister(instr.ref_create.initial_value); | ||
| auto ref = ADT(111, {value}); |
There was a problem hiding this comment.
How we do ensure this tag does not collide with other ADT's tag number?
There was a problem hiding this comment.
It doesn't actually matter, I just used ADT container since typing guarantees we already know ahead of time the value is a reference we don't need to check or inspect the tag.
There was a problem hiding this comment.
Maybe in this case Tuple({values}) is okay which defaults tag to 0?
There was a problem hiding this comment.
Agree, probably just use Tuple.
There was a problem hiding this comment.
Let me speculate a case that in a relay program we defined many ADTs (150 ADT types), and VM compiler generates tags for each ADT type, is it possible to collide with 111? I recall it includes hashes in the tag number, so probably unlikely. And Tuple seems perfect fix.
| Expr VisitExpr_(const LetNode* op) final { | ||
| Var v = op->var; | ||
| if (HasLet(v)) { | ||
| if (HasLet(v) || op->value.as<RefWriteNode>()) { |
There was a problem hiding this comment.
I don't believe this fully fixes DCE due to nesting ref_write deep inside lets that get DCEd
There was a problem hiding this comment.
This should stop it from removing RefWrites anywhere and by extension keep all ref operations alive as the reference is now live too.
There was a problem hiding this comment.
I will give you a test case for this since I made this exact change when debugging and it wasn't enough to fix the problem. (In fact I remember segfaulting...)
There was a problem hiding this comment.
BTW, out of the scope of this PR, do we need to do some alias analysis to handle reference for DCE?
There was a problem hiding this comment.
I believe so, for now we should probably require DCE to error when the incoming module contains references. @jroesch and I are planning on writing an alias analysis pass soon.
mbrookhart
left a comment
mbrookhart
left a comment
There was a problem hiding this comment.
I'd like to see some of the debug stuff more formalized.
Looks like we're failing a bunch of unit tests?
I'm not sure I understand the goal of this PR, is there an RFC for references in TVM more generally?
| std::string DebugPrint(NDArray array) { | ||
| const PackedFunc* fprint = Registry::Get("relay._ndarray_repr"); | ||
| ICHECK(fprint) << "unable to find printing function for constants"; | ||
| return (*fprint)(array); | ||
| } | ||
|
|
There was a problem hiding this comment.
Should we make this part of the ReprPrinter in ndarray.cc? I've been using something very similar in the VM, but I feel like we should make it more general.
|
|
||
| if __name__ == "__main__": | ||
| pytest.main([__file__]) | ||
| import sys |
| } | ||
| case Opcode::RefCreate: { | ||
| auto value = ReadRegister(instr.ref_create.initial_value); | ||
| auto ref = ADT(111, {value}); |
There was a problem hiding this comment.
Agree, probably just use Tuple.
|
I've addressed the ADT tag issue as suggested (using Tuple), and left a TODO comment in DCE. For the record, I tried adding Feature set checking in DCE to error on detecting RefWrite. However, the VM compilation process requires DCE in other places (I think most critically in re @mbrookhart , I'm not totally sure what you meant by the debug stuff, but references have been in Relay since higher-order AD was introduced. I'm not sure if there is/was an RFC for adding references, although I agree we should probably make one especially as we are working towards supporting stateful stuff going forward. cc @jroesch who might have more thoughts |
|
I think the failing unit tests are actually unsound and rely on DCE for refs, or perhaps they are sound but correctness is definitely not guaranteed by what we currently have. cc @MarisaKirisame, how should we proceed? We might have to block this PR until DCE gets fixed once and for all, or disable all the offending tests. |
|
Maybe we should disable the offending test rn and push on the DCE fixes? |
| case Opcode::DeviceCopy: | ||
| case Opcode::RefCreate: | ||
| case Opcode::RefRead: | ||
| case Opcode::RefWrite: |
There was a problem hiding this comment.
RefWrite doesn't have last register, right? Maybe set the last_register_ to -1?
|
What's the status on this PR now? @jroesch |
9 participants
Add initial support for references in the VM. cc @altanh @zhiics @icemelon9 @antinucleon @tmoreau89 @mbrookhart.