Making vm.yield take a resume branch.
This ensures we yield at a safe point and can join on resume from multiple
yield points. If we ever get around to making register allocation we can
also treat the terminator as a good point for clearing ref registers to
eagerly drop resources.
diff --git a/iree/compiler/Dialect/VM/IR/VMOps.cpp b/iree/compiler/Dialect/VM/IR/VMOps.cpp
index e4008de..ee5e558 100644
--- a/iree/compiler/Dialect/VM/IR/VMOps.cpp
+++ b/iree/compiler/Dialect/VM/IR/VMOps.cpp
@@ -1151,6 +1151,22 @@
 // Async/fiber ops
 //===----------------------------------------------------------------------===//
 
+Block *YieldOp::getDest() { return getOperation()->getSuccessor(0); }
+
+void YieldOp::setDest(Block *block) {
+  return getOperation()->setSuccessor(block, 0);
+}
+
+void YieldOp::eraseOperand(unsigned index) {
+  getOperation()->eraseOperand(index);
+}
+
+Optional<MutableOperandRange> YieldOp::getMutableSuccessorOperands(
+    unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return destOperandsMutable();
+}
+
 //===----------------------------------------------------------------------===//
 // Debugging
 //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Dialect/VM/IR/VMOps.td b/iree/compiler/Dialect/VM/IR/VMOps.td
index eecb75d..5ce09bb 100644
--- a/iree/compiler/Dialect/VM/IR/VMOps.td
+++ b/iree/compiler/Dialect/VM/IR/VMOps.td
@@ -3739,21 +3739,57 @@
 //   await_any
 
 def VM_YieldOp : VM_Op<"yield", [
+    DeclareOpInterfaceMethods<BranchOpInterface>,
     DeclareOpInterfaceMethods<VM_SerializableOpInterface>,
     HasParent<"IREE::VM::FuncOp">,
+    Terminator,
     YieldPoint,
   ]> {
   let summary = [{unconditional fiber yield operation}];
   let description = [{
     Yields the fiber for some (likely short) amount of time. This can be used to
-    perform cooperative scheduling and ensure fair (enough) execution.
+    perform cooperative scheduling and ensure fair (enough) execution. Execution
+    resumes at the specified target branch.
+
+    ```
+    ^bb0:
+      vm.yield ^on_resume
+    ^on_resume:
+      ...
+   ```
   }];
 
-  let assemblyFormat = "attr-dict";
+  let arguments = (ins
+    Variadic<VM_AnyType>:$destOperands
+  );
+
+  let successors = (successor
+    AnySuccessor:$dest
+  );
+
+  let assemblyFormat = [{
+    $dest (`(` $destOperands^ `:` type($destOperands) `)`)? attr-dict
+  }];
 
   let encoding = [
     VM_EncOpcode<VM_OPC_Yield>,
+    VM_EncBranch<"dest", "getOperands", 0>,
   ];
+
+  let builders = [
+    OpBuilder<(ins "Block *":$dest, CArg<"ValueRange", "{}">:$destOperands), [{
+      $_state.addSuccessors(dest);
+      $_state.addOperands(destOperands);
+    }]>,
+  ];
+
+  let extraClassDeclaration = [{
+    Block *getDest();
+    void setDest(Block *block);
+
+    /// Erase the operand at 'index' from the operand list.
+    void eraseOperand(unsigned index);
+  }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -3851,11 +3887,11 @@
   ];
 
   let builders = [
-    OpBuilder<(ins "Block *":$dest, CArg<"ValueRange", "{}">:$destOperands),
-    [{
+    OpBuilder<(ins "Block *":$dest, CArg<"ValueRange", "{}">:$destOperands), [{
       $_state.addSuccessors(dest);
       $_state.addOperands(destOperands);
-  }]>];
+    }]>,
+  ];
 
   let extraClassDeclaration = [{
     Block *getDest();
diff --git a/iree/compiler/Dialect/VM/IR/test/control_flow_ops.mlir b/iree/compiler/Dialect/VM/IR/test/control_flow_ops.mlir
index beff9a6..4281968 100644
--- a/iree/compiler/Dialect/VM/IR/test/control_flow_ops.mlir
+++ b/iree/compiler/Dialect/VM/IR/test/control_flow_ops.mlir
@@ -217,8 +217,9 @@
 // CHECK-LABEL: @yield
 vm.module @my_module {
   vm.func @yield() {
-    // CHECK: vm.yield
-    vm.yield
+    // CHECK: vm.yield ^bb1
+    vm.yield ^bb1
+  ^bb1:
     vm.return
   }
 }
diff --git a/iree/vm/bytecode_disasm.c b/iree/vm/bytecode_disasm.c
index 2dbeed0..08400a0 100644
--- a/iree/vm/bytecode_disasm.c
+++ b/iree/vm/bytecode_disasm.c
@@ -1363,8 +1363,13 @@
     //===------------------------------------------------------------------===//
 
     DISASM_OP(CORE, Yield) {
+      int32_t block_pc = VM_DecBranchTarget("dest");
+      const iree_vm_register_remap_list_t* remap_list =
+          VM_ParseBranchOperands("operands");
       IREE_RETURN_IF_ERROR(
-          iree_string_builder_append_cstring(b, "vm.yield (TBD)"));
+          iree_string_builder_append_format(b, "vm.yield ^%08X(", block_pc));
+      EMIT_REMAP_LIST(remap_list);
+      IREE_RETURN_IF_ERROR(iree_string_builder_append_cstring(b, ")"));
       break;
     }
 
diff --git a/iree/vm/bytecode_dispatch.c b/iree/vm/bytecode_dispatch.c
index 23c7647..73e3a77 100644
--- a/iree/vm/bytecode_dispatch.c
+++ b/iree/vm/bytecode_dispatch.c
@@ -1486,8 +1486,18 @@
     //===------------------------------------------------------------------===//
 
     DISPATCH_OP(CORE, Yield, {
-      // TODO(benvanik): yield with execution results.
-      return iree_ok_status();
+      // Perform branch before yielding; in this way we will resume at the
+      // target without needing to retain any information about the yield.
+      int32_t block_pc = VM_DecBranchTarget("dest");
+      const iree_vm_register_remap_list_t* remap_list =
+          VM_DecBranchOperands("operands");
+      iree_vm_bytecode_dispatch_remap_branch_registers(regs, remap_list);
+      pc = block_pc;
+
+      // Return magic status code indicating a yield.
+      // This isn't an error, though callers not supporting coroutines will
+      // treat it as one and propagate it up.
+      return iree_status_from_code(IREE_STATUS_DEFERRED);
     });
 
     //===------------------------------------------------------------------===//