Merge pull request #8108 from google/benvanik-vm-yield
Preparing for VM coroutine yields via `vm.yield`.
diff --git a/iree/base/status.c b/iree/base/status.c
index 8fe228e..ceebea4 100644
--- a/iree/base/status.c
+++ b/iree/base/status.c
@@ -220,8 +220,6 @@
return "ALREADY_EXISTS";
case IREE_STATUS_PERMISSION_DENIED:
return "PERMISSION_DENIED";
- case IREE_STATUS_UNAUTHENTICATED:
- return "UNAUTHENTICATED";
case IREE_STATUS_RESOURCE_EXHAUSTED:
return "RESOURCE_EXHAUSTED";
case IREE_STATUS_FAILED_PRECONDITION:
@@ -238,6 +236,10 @@
return "UNAVAILABLE";
case IREE_STATUS_DATA_LOSS:
return "DATA_LOSS";
+ case IREE_STATUS_UNAUTHENTICATED:
+ return "UNAUTHENTICATED";
+ case IREE_STATUS_DEFERRED:
+ return "DEFERRED";
default:
return "";
}
diff --git a/iree/base/status.h b/iree/base/status.h
index a26352c..7192069 100644
--- a/iree/base/status.h
+++ b/iree/base/status.h
@@ -73,24 +73,85 @@
// Note that any code within IREE_STATUS_CODE_MASK is valid even if not
// enumerated here. Always check for unhandled errors/have default conditions.
typedef enum iree_status_code_e {
+ // Successful operation.
IREE_STATUS_OK = 0,
+
+ // Operation was cancelled by the caller.
IREE_STATUS_CANCELLED = 1,
+
+ // Unknown error, or error that could not be mapped to this enum.
IREE_STATUS_UNKNOWN = 2,
+
+ // The caller provided an invalid argument and that future calls with the same
+ // arguments will fail. If the failure is predicated on system state that may
+ // change prefer IREE_STATUS_OUT_OF_RANGE.
IREE_STATUS_INVALID_ARGUMENT = 3,
+
+ // A deadline was exceeded before the call could complete.
+ // This can be returned even if the operation would have completed
+ // successfully had the deadline not been met.
IREE_STATUS_DEADLINE_EXCEEDED = 4,
+
+ // A referenced resource could not be found or was unavailable to all
+ // requesters. IREE_STATUS_PERMISSION_DENIED should be used if only an
+ // individual requester is denied access.
IREE_STATUS_NOT_FOUND = 5,
+
+ // The resource the caller attempted to create already exists.
IREE_STATUS_ALREADY_EXISTS = 6,
+
+ // The caller does not have permission to execute the operation or have access
+ // to the requested resources.
IREE_STATUS_PERMISSION_DENIED = 7,
+
+ // Some resource type has been exhausted and the operation is unable to
+ // reserve what it requires, either by quota or underlying system exhaustion.
IREE_STATUS_RESOURCE_EXHAUSTED = 8,
+
+ // The operation was rejected because the system is not in a state required
+ // for the operation's execution.
+ //
+ // Use IREE_STATUS_UNAVAILABLE if the caller can retry the operation.
+ // Use IREE_STATUS_ABORTED if the caller should restart their transaction
+ // (the entire sequence of operations is invalid).
+ // Use IREE_STATUS_FAILED_PRECONDITION if the caller should not retry until
+ // the system state has been explicitly fixed.
IREE_STATUS_FAILED_PRECONDITION = 9,
+
+ // The operation was aborted by the system.
+ // If responding to a caller-requested cancellation use IREE_STATUS_CANCELLED.
IREE_STATUS_ABORTED = 10,
+
+ // The operation was attempted past the valid range (of a resource, etc).
+ // Indicates the operation can be retried if the system state is fixed.
IREE_STATUS_OUT_OF_RANGE = 11,
+
+ // Operation has not been implemented or is not supported.
IREE_STATUS_UNIMPLEMENTED = 12,
+
+ // An internal error has occurred and some invariants expected by an
+ // underlying system have been violated. This error code is reserved for
+ // serious errors.
IREE_STATUS_INTERNAL = 13,
+
+ // The system used to perform the operation is currently (and transiently)
+ // unavailable. Callers can retry with backoff.
IREE_STATUS_UNAVAILABLE = 14,
+
+ // An serious unrecoverable data loss or corruption has occurred.
+ // Indicates that an underlying system or resource has failed in such a way
+ // that all related operations may produce incorrect results.
IREE_STATUS_DATA_LOSS = 15,
+
+ // The requested operation does not have proper authentication.
+ // Callers can correct this and retry.
IREE_STATUS_UNAUTHENTICATED = 16,
+ // The operation has been deferred and must be resumed at a future point.
+ // Used by resumable operations as part of scheduling and execution systems.
+ // Callers that do not handle deferred execution can treat this as a failure.
+ IREE_STATUS_DEFERRED = 17,
+
IREE_STATUS_CODE_MASK = 0x1Fu,
} iree_status_code_t;
@@ -151,6 +212,8 @@
(iree_status_code(value) == IREE_STATUS_DATA_LOSS)
#define iree_status_is_unauthenticated(value) \
(iree_status_code(value) == IREE_STATUS_UNAUTHENTICATED)
+#define iree_status_is_deferred(value) \
+ (iree_status_code(value) == IREE_STATUS_DEFERRED)
#define IREE_STATUS_IMPL_CONCAT_INNER_(x, y) x##y
#define IREE_STATUS_IMPL_CONCAT_(x, y) IREE_STATUS_IMPL_CONCAT_INNER_(x, y)
diff --git a/iree/base/status_cc.h b/iree/base/status_cc.h
index 21fad56..4795dda 100644
--- a/iree/base/status_cc.h
+++ b/iree/base/status_cc.h
@@ -88,6 +88,7 @@
kUnavailable = IREE_STATUS_UNAVAILABLE,
kDataLoss = IREE_STATUS_DATA_LOSS,
kUnauthenticated = IREE_STATUS_UNAUTHENTICATED,
+ kDeferred = IREE_STATUS_DEFERRED,
};
static inline const char* StatusCodeToString(StatusCode code) {
diff --git a/iree/compiler/Dialect/Util/IR/UtilTypes.h b/iree/compiler/Dialect/Util/IR/UtilTypes.h
index 2d47d4f..6527a37 100644
--- a/iree/compiler/Dialect/Util/IR/UtilTypes.h
+++ b/iree/compiler/Dialect/Util/IR/UtilTypes.h
@@ -59,6 +59,7 @@
Unavailable = 14,
DataLoss = 15,
Unauthenticated = 16,
+ Deferred = 17,
DoNotUseReservedForFutureExpansionUseDefaultInSwitchInstead_ = 20
};
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);
});
//===------------------------------------------------------------------===//
