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); }); //===------------------------------------------------------------------===//