Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[BACKPORT-0.6][Relay][VM] Fix compilation of If-Elses #5877

Merged
merged 1 commit into from Jun 21, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 5 additions & 3 deletions src/relay/backend/vm/compiler.cc
Expand Up @@ -375,7 +375,9 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
this->Emit(Instruction::If(test_register, target_register, 0, 0));
this->VisitExpr(if_node->true_branch);

size_t true_register = last_register_;
// It saves the result of If-Else expression.
auto merge_register = NewRegister();
Emit(Instruction::Move(last_register_, merge_register));
Emit(Instruction::Goto(0));

// Finally store how many instructions there are in the
Expand All @@ -387,7 +389,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
size_t false_register = last_register_;

// In else-branch, override the then-branch register
Emit(Instruction::Move(false_register, true_register));
Emit(Instruction::Move(false_register, merge_register));
// Compute the total number of instructions
// after generating false.
auto after_false = this->instructions_.size();
Expand All @@ -406,7 +408,7 @@ class VMFunctionCompiler : ExprFunctor<void(const Expr& expr)> {
// Patch the Goto.
this->instructions_[after_true - 1].pc_offset = (after_false - after_true) + 1;

this->last_register_ = true_register;
this->last_register_ = merge_register;
}

void EmitShapeFunc(Function func, Array<Expr> inputs, Array<Expr> outputs) {
Expand Down