[relay][VM] Miscompile with double-if

[abergeron]

This relay code:

def @main(%b: bool) -> ((int32,), (int32,)) {
  let %v0: (int32,) = (0 /* ty=int32 */,);
  let %v1: (int32,) = (1 /* ty=int32 */,);
  let %v2: (int32,) = if (%b) {
    %v0
  } else {
    %v1
  };
  let %v3: (int32,) = if (%b) {
    %v1
  } else {
    %v0
  };
  (%v2, %v3)
}

Compiles to this in VM bytcode:

main: 
0: load_const $1 Const[0];
1: alloc_data $2 tag(0) [$1];
2: load_const $3 Const[1];
3: alloc_data $4 tag(0) [$3];
4: load_consti $5 1;
5: if $0 $5 1 2;
6: goto 2;
7: move $2 $4;
8: load_consti $6 1;
9: if $0 $6 1 2;
10: goto 2;
11: move $4 $2;
12: alloc_data $7 tag(0) [$2,$4];
13: ret $7;

If you pass False to the original function this makes the VM run both 7: move $2 $4 and 11: move $4 $2 which sets both of them to (1,) thus returning the wrong result.

The debug interpreter doesn't have this problem.

This script reproduces the problem:

import tvm
from tvm import relay

mod = tvm.IRModule({})

b = relay.var('b')

v0 = relay.var('v0')
v1 = relay.var('v1')
v2 = relay.var('v2')
v3 = relay.var('v3')

out = relay.Tuple([v2, v3])
out = relay.Let(v3, relay.If(b, v1, v0), out)
out = relay.Let(v2, relay.If(b, v0, v1), out)
out = relay.Let(v1, relay.Tuple([relay.const(1)]), out)
out = relay.Let(v0, relay.Tuple([relay.const(0)]), out)

fn = relay.Function([b], out)

mod['main'] = fn

print(str(mod))

ctx = tvm.runtime.ndarray.context('llvm', 0)

debug = relay.create_executor(ctx=ctx, mod=mod, kind='debug')
res = debug.evaluate()(False)
res = ((int(res[0][0].asnumpy()),), (int(res[1][0].asnumpy()),))
print("debug:", res)

vm = relay.create_executor(ctx=ctx, mod=mod, kind='vm')
res = vm.evaluate()(False)
res = ((int(res[0][0].asnumpy()),), (int(res[1][0].asnumpy()),))
print("vm:", res)

[tqchen]

cc @icemelon9 @wweic @jroesch

[jroesch]

@abergeron I'll look into this one, it looks like a bug was introduced when compiling if statements, they seem to be sharing the same set of registers incorrectly.

[wweic]

@jroesch The problem is in when we compile if we don't allocate a new register to store the result. We ended up with overwriting registers we shouldn't write to. The fix is actually simple. we create a merge register to keep the result of if expression.