diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_ready.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_ready.h new file mode 100644 index 0000000000..1729c2a29a --- /dev/null +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_ready.h @@ -0,0 +1,1138 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#pragma once + +#include "scheduler_types.h" +#include "scheduler_memory.h" +#include "scheduler_graph.h" + +enum class SchedulerRouteResult : uint64_t { + READY = 0, + READY_TO_ENQUEUE = 1, + WAITING = 2, + COMPLETED = 3, + ERROR = 4, +}; + +enum class SchedulerPredicateResult : uint8_t { + PASS = 0, + FAIL = 1, + MALFORMED = 2, +}; + +struct SchedulerWakeStats { + uint64_t fanin_state_load_count{0}; + uint64_t wake_register_count{0}; + uint64_t wake_cas_retry_count{0}; + uint64_t wake_closed_retry_count{0}; + uint64_t wake_migrate_count{0}; + uint64_t wake_close_count{0}; +}; + +struct SchedulerReadyStats { + uint64_t enqueue_count{0}; + uint64_t batch_count{0}; + uint64_t pop_count{0}; + uint64_t steal_count{0}; + uint64_t cas_retry_count{0}; + uint64_t link_wait_count{0}; + uint64_t link_wait_max{0}; +}; + +struct SchedulerCompletionStats { + uint64_t enqueue_count{0}; + uint64_t resolve_count{0}; + uint64_t ready_to_kernel_cycles{0}; + uint64_t ready_to_kernel_max_cycles{0}; +}; + +struct SchedulerReadyBatch { + int64_t head{SCHEDULER_INBOX_EMPTY}; + int64_t tail{SCHEDULER_INBOX_EMPTY}; + uint64_t count{0}; +}; + +inline __aicore__ void scheduler_ready_batch_reset(SchedulerReadyBatch *batch) { + batch->head = SCHEDULER_INBOX_EMPTY; + batch->tail = SCHEDULER_INBOX_EMPTY; + batch->count = 0; +} + +inline __aicore__ uint64_t scheduler_ready_pending_pack(int64_t head, int64_t tail) { + const uint64_t packed_head = static_cast(static_cast(head)); + const uint64_t packed_tail = static_cast(static_cast(tail)); + return packed_head | (packed_tail << 32); +} + +inline __aicore__ int64_t scheduler_ready_pending_head(uint64_t packed) { + return static_cast(static_cast(static_cast(packed))); +} + +inline __aicore__ int64_t scheduler_ready_pending_tail(uint64_t packed) { + return static_cast(static_cast(static_cast(packed >> 32))); +} + +inline __aicore__ bool scheduler_ready_pending_endpoint_fits(int64_t task_id) { + return task_id >= SCHEDULER_INBOX_EMPTY && task_id <= INT32_MAX; +} + +inline __aicore__ void scheduler_ready_owner_init(__gm__ SchedulerReadyOwnerState *owner_state) { + if (owner_state == nullptr) return; + for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + scheduler_gm_store(owner_state->queues[type].pending_endpoints, SCHEDULER_READY_PENDING_EMPTY); + scheduler_gm_store(owner_state->queues[type].advertised, UINT64_C(0)); + } + scheduler_cache_barrier(); +} + +struct SchedulerReadyClaim { + int64_t task_id{SCHEDULER_TASK_ID_INVALID}; + uint64_t inbox_index{UINT64_MAX}; + SchedulerReadySource source{SchedulerReadySource::LOCAL}; + uint64_t claim_start_cycles{0}; + uint64_t claim_end_cycles{0}; +}; + +struct SchedulerFreeSlotClaim { + uint64_t worker_id{UINT64_MAX}; + uint32_t slot_index{UINT32_MAX}; + uint32_t generation{0}; +}; + +struct SchedulerCompletionServiceTiming { + uint64_t scan_cycles{0}; + uint64_t consume_cycles{0}; + uint64_t resolve_cycles{0}; + uint64_t ready_publish_cycles{0}; + uint64_t refill_cycles{0}; + uint64_t finalize_cycles{0}; +}; + +struct SchedulerDispatchFillTiming { + uint64_t prepare_cycles{0}; + uint64_t materialize_cycles{0}; + uint64_t publish_cycles{0}; +}; + +inline __aicore__ uint64_t scheduler_cycles() { +#if defined(__CCE_AICORE__) + return get_sys_cnt_aicore(); +#else + return 0; +#endif +} + +inline __aicore__ uint32_t scheduler_core_type_index(int32_t core_type) { + return core_type == static_cast(CoreType::AIC) ? 0U : 1U; +} + +inline __aicore__ uint32_t scheduler_metadata_core_type_index(uint8_t subtask_slot) { + return subtask_slot == 0 ? 0U : 1U; +} + +inline __aicore__ uint8_t scheduler_metadata_single_subtask_slot(uint8_t active_mask) { + if ((active_mask & 1U) != 0) return 0; + if ((active_mask & 2U) != 0) return 1; + return 2; +} + +inline __aicore__ SchedulerPredicateResult +scheduler_evaluate_task_predicate(const SchedulerGraphView &graph, int64_t task_id) { + __gm__ uint8_t *payload = scheduler_graph_payload(graph, task_id); + __gm__ SchedulerDispatchPredicate *predicate = + reinterpret_cast<__gm__ SchedulerDispatchPredicate *>(payload + SCHEDULER_GRAPH_PREDICATE_OFFSET); + scheduler_observe_cache_line(predicate); + if (predicate->op == 0) return SchedulerPredicateResult::PASS; + if (predicate->op > 6 || + (predicate->elem_size != 1 && predicate->elem_size != 2 && predicate->elem_size != 4 && + predicate->elem_size != 8) || + predicate->addr == 0 || (predicate->addr & (static_cast(predicate->elem_size) - 1)) != 0) + return SchedulerPredicateResult::MALFORMED; + + __gm__ void *operand = reinterpret_cast<__gm__ void *>(predicate->addr); + scheduler_observe_cache_line(operand); + int64_t value = 0; + switch (predicate->elem_size) { + case 1: + value = *reinterpret_cast<__gm__ int8_t *>(operand); + break; + case 2: + value = *reinterpret_cast<__gm__ int16_t *>(operand); + break; + case 4: + value = *reinterpret_cast<__gm__ int32_t *>(operand); + break; + case 8: + value = *reinterpret_cast<__gm__ int64_t *>(operand); + break; + default: + return SchedulerPredicateResult::MALFORMED; + } + switch (predicate->op) { + case 1: + return value == predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + case 2: + return value != predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + case 3: + return value > predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + case 4: + return value < predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + case 5: + return value >= predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + case 6: + return value <= predicate->target ? SchedulerPredicateResult::PASS : SchedulerPredicateResult::FAIL; + default: + return SchedulerPredicateResult::MALFORMED; + } +} + +inline __aicore__ bool +scheduler_lookup_callable_address(__gm__ uint64_t *callable_addresses, uint16_t kernel_id, uint64_t *callable_address) { + if (callable_addresses == nullptr || callable_address == nullptr || kernel_id >= SCHEDULER_CALLABLE_CAPACITY) + return false; + scheduler_observe_cache_line(&callable_addresses[kernel_id]); + *callable_address = callable_addresses[kernel_id]; + return *callable_address != 0; +} + +inline __aicore__ uint64_t +scheduler_completion_id(__gm__ const SchedulerWorkerContext *context, uint64_t local_completion_index) { + return local_completion_index * context->runtime_worker_count + context->worker_index; +} + +inline __aicore__ __gm__ SchedulerTaskControl *scheduler_task_control_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, int64_t task_id +) { + return scheduler_state_at( + scheduler_state_base, + context->task_controls_offset + static_cast(task_id) * sizeof(SchedulerTaskControl) + ); +} + +inline __aicore__ __gm__ SchedulerTaskMetadata *scheduler_task_metadata_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, int64_t task_id +) { + return scheduler_state_at( + scheduler_state_base, + context->task_metadata_offset + static_cast(task_id) * sizeof(SchedulerTaskMetadata) + ); +} + +inline __aicore__ __gm__ SchedulerCompletionInbox *scheduler_completion_inbox_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, uint64_t inbox_index +) { + return scheduler_state_at( + scheduler_state_base, context->completion_inboxes_offset + inbox_index * sizeof(SchedulerCompletionInbox) + ); +} + +inline __aicore__ __gm__ SchedulerReadyInbox *scheduler_ready_inbox_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, uint32_t core_type_index, + uint64_t inbox_index +) { + uint64_t linear = static_cast(core_type_index) * SCHEDULER_WORKER_CAPACITY + inbox_index; + return scheduler_state_at( + scheduler_state_base, context->ready_inboxes_offset + linear * sizeof(SchedulerReadyInbox) + ); +} + +inline __aicore__ __gm__ SchedulerReadyOwnerState * +scheduler_ready_owner_state_at(__gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context) { + return scheduler_state_at( + scheduler_state_base, + context->ready_owner_states_offset + context->inbox_index * sizeof(SchedulerReadyOwnerState) + ); +} + +inline __aicore__ __gm__ SchedulerReadyDirectory * +scheduler_ready_directory_at(__gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context) { + return scheduler_state_at(scheduler_state_base, context->ready_directory_offset); +} + +inline __aicore__ __gm__ SchedulerGangCoordinator * +scheduler_gang_coordinator_at(__gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context) { + return scheduler_state_at(scheduler_state_base, context->gang_coordinator_offset); +} + +inline __aicore__ void scheduler_publish_gang_ready( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, __gm__ SchedulerTaskControl *control, + uint8_t metadata_flags +) { + scheduler_gm_store(control->state, static_cast(SchedulerTaskState::READY)); + __gm__ SchedulerGangCoordinator *coordinator = scheduler_gang_coordinator_at(scheduler_state_base, context); + scheduler_gm_fetch_or(coordinator->ready_priority_bits, scheduler_task_priority_bit(metadata_flags)); +} + +inline __aicore__ __gm__ SchedulerWorkerContext *scheduler_worker_context_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, uint64_t worker_id +) { + return scheduler_state_at( + scheduler_state_base, context->worker_contexts_offset + worker_id * sizeof(SchedulerWorkerContext) + ); +} + +inline __aicore__ __gm__ SchedulerDispatchSlot *scheduler_dispatch_slot_at( + __gm__ void *scheduler_state_base, __gm__ const SchedulerWorkerContext *context, uint64_t worker_id, uint32_t slot +) { + return scheduler_state_at( + scheduler_state_base, + context->dispatch_slots_offset + + (worker_id * SCHEDULER_PENDING_SLOT_COUNT + static_cast(slot)) * sizeof(SchedulerDispatchSlot) + ); +} + +inline __aicore__ uint64_t scheduler_dispatch_publication(uint32_t generation, SchedulerDispatchSlotState state) { + return (static_cast(generation) << 8) | static_cast(state); +} + +inline __aicore__ uint32_t scheduler_dispatch_generation(uint64_t publication) { + return static_cast(publication >> 8); +} + +inline __aicore__ SchedulerDispatchSlotState scheduler_dispatch_state(uint64_t publication) { + return static_cast(publication & UINT64_C(0xff)); +} + +inline __aicore__ void scheduler_record_error( + __gm__ SchedulerRunControl *run_control, int64_t task_id, SchedulerGraphResult status, + const SchedulerGraphView *graph = nullptr, __gm__ const SchedulerWorkerContext *context = nullptr, + SchedulerErrorSite error_site = SchedulerErrorSite::UNKNOWN +) { + if (scheduler_gm_compare_exchange(run_control->error_claimed, UINT64_C(0), UINT64_C(1)) != 0) return; + scheduler_gm_store(run_control->error_task_id, static_cast(task_id)); + if (graph != nullptr) { + scheduler_gm_store(run_control->error_graph_task_count, graph->task_count); + scheduler_gm_store(run_control->error_descriptors_address, graph->descriptors_address); + scheduler_gm_store(run_control->error_payloads_address, graph->payloads_address); + scheduler_gm_store(run_control->error_task_window_mask, graph->task_window_mask); + } + if (context != nullptr) { + scheduler_gm_store(run_control->error_core_id, static_cast(context->physical_core_id)); + scheduler_gm_store(run_control->error_core_type, static_cast(context->core_type)); + } + scheduler_gm_store(run_control->error_site, static_cast(error_site)); + scheduler_gm_publish(run_control->scheduler_error, static_cast(status)); +} + +inline __aicore__ void scheduler_publish_waiter_metadata( + __gm__ SchedulerTaskControl *control, int64_t next_waiter, int32_t next_fanin_index, int32_t waiting_producer +) { + control->next_waiter = next_waiter; + control->next_fanin_index = next_fanin_index; + control->waiting_producer = waiting_producer; + scheduler_publish_cache_line(&control->next_waiter); +} + +inline __aicore__ int64_t scheduler_observe_next_waiter(__gm__ SchedulerTaskControl *control) { + scheduler_observe_cache_line(&control->next_waiter); + return control->next_waiter; +} + +inline __aicore__ SchedulerRouteResult scheduler_route_task( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, int64_t task_id, SchedulerWakeStats *stats, + bool validate_current_state = true +) { + if (task_id < 0 || static_cast(task_id) >= graph.task_count) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_TASK_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); + if (validate_current_state) { + int64_t state = scheduler_gm_query(control->state); + if (state == static_cast(SchedulerTaskState::DONE)) return SchedulerRouteResult::COMPLETED; + if (state == static_cast(SchedulerTaskState::READY)) return SchedulerRouteResult::READY; + if (state != static_cast(SchedulerTaskState::BLOCKED)) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::ROUTE_INVALID_STATE + ); + return SchedulerRouteResult::ERROR; + } + } + + __gm__ uint8_t *payload = scheduler_graph_payload(graph, task_id); + int32_t fanin_count = *reinterpret_cast<__gm__ int32_t *>(payload + SCHEDULER_GRAPH_FANIN_COUNT_OFFSET); + int32_t next_fanin = control->next_fanin_index; + if (fanin_count < 0 || next_fanin < 0 || next_fanin > fanin_count) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_FANIN_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + + while (next_fanin < fanin_count) { + int32_t producer = scheduler_graph_fanin_id(graph, task_id, next_fanin); + if (producer < 0 || producer >= task_id) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_FANIN_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + if (stats != nullptr) ++stats->fanin_state_load_count; + __gm__ SchedulerTaskControl *producer_control = + scheduler_task_control_at(scheduler_state_base, context, producer); + if (scheduler_gm_query(producer_control->state) == static_cast(SchedulerTaskState::DONE)) { + ++next_fanin; + continue; + } + while (true) { + int64_t observed = scheduler_gm_query(producer_control->wake_list_head); + if (observed == SCHEDULER_WAKE_LIST_CLOSED) { + if (stats != nullptr) ++stats->wake_closed_retry_count; + ++next_fanin; + break; + } + scheduler_publish_waiter_metadata(control, observed, next_fanin, producer); + int64_t actual = scheduler_gm_compare_exchange(producer_control->wake_list_head, observed, task_id); + if (actual == observed) { + if (stats != nullptr) ++stats->wake_register_count; + return SchedulerRouteResult::WAITING; + } + if (stats != nullptr) ++stats->wake_cas_retry_count; + if (actual == SCHEDULER_WAKE_LIST_CLOSED) { + if (stats != nullptr) ++stats->wake_closed_retry_count; + ++next_fanin; + break; + } + } + } + scheduler_publish_waiter_metadata( + control, SCHEDULER_TASK_ID_INVALID, fanin_count, static_cast(SCHEDULER_TASK_ID_INVALID) + ); + return SchedulerRouteResult::READY_TO_ENQUEUE; +} + +// No task can execute while the bootstrap barrier is closed. Executable +// producers therefore have open wake lists, and non-executable producers are +// the inline-completed tasks initialized by the host. The barrier makes it +// safe to publish the new head before publishing the waiter's link. +inline __aicore__ SchedulerRouteResult scheduler_bootstrap_route_task( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, int64_t task_id, SchedulerWakeStats *stats +) { + if (task_id < 0 || static_cast(task_id) >= graph.task_count) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_TASK_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + __gm__ uint8_t *payload = scheduler_graph_payload(graph, task_id); + int32_t fanin_count = *reinterpret_cast<__gm__ int32_t *>(payload + SCHEDULER_GRAPH_FANIN_COUNT_OFFSET); + if (fanin_count < 0) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_FANIN_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); + for (int32_t next_fanin = 0; next_fanin < fanin_count; ++next_fanin) { + int32_t producer = scheduler_graph_fanin_id(graph, task_id, next_fanin); + if (producer < 0 || producer >= task_id) { + scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_FANIN_ID, &graph, context); + return SchedulerRouteResult::ERROR; + } + __gm__ SchedulerTaskMetadata *producer_metadata = + scheduler_task_metadata_at(scheduler_state_base, context, producer); + scheduler_observe_cache_line(producer_metadata); + if (!scheduler_task_is_executable(producer_metadata->flags)) continue; + + __gm__ SchedulerTaskControl *producer_control = + scheduler_task_control_at(scheduler_state_base, context, producer); + int64_t previous = scheduler_gm_exchange(producer_control->wake_list_head, task_id); + if (previous < SCHEDULER_WAKE_LIST_OPEN) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::BOOTSTRAP_WAKE_INVALID_HEAD + ); + return SchedulerRouteResult::ERROR; + } + control->next_waiter = previous; + control->next_fanin_index = next_fanin; + control->waiting_producer = producer; + scheduler_writeback_cache_line(&control->next_waiter); + if (stats != nullptr) ++stats->wake_register_count; + return SchedulerRouteResult::WAITING; + } + + control->next_fanin_index = fanin_count; + control->waiting_producer = static_cast(SCHEDULER_TASK_ID_INVALID); + return SchedulerRouteResult::READY_TO_ENQUEUE; +} + +inline __aicore__ bool scheduler_bootstrap_ready_batch_append( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, int64_t task_id, + SchedulerReadyBatch *batch, SchedulerReadyStats *stats, bool trace_enabled = false +) { + if (batch == nullptr || task_id < 0 || static_cast(task_id) >= context->graph_task_count) return false; + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); + scheduler_gm_store(control->state, static_cast(SchedulerTaskState::READY)); + control->next_waiter = SCHEDULER_INBOX_EMPTY; + scheduler_writeback_cache_line(&control->next_waiter); + if (batch->head == SCHEDULER_INBOX_EMPTY) { + batch->head = task_id; + } else { + __gm__ SchedulerTaskControl *tail = scheduler_task_control_at(scheduler_state_base, context, batch->tail); + tail->next_waiter = task_id; + scheduler_writeback_cache_line(&tail->next_waiter); + } + batch->tail = task_id; + if (trace_enabled) { + __gm__ SchedulerTaskTrace *cells = + scheduler_state_at(scheduler_state_base, context->trace_cells_offset); + cells[task_id].ready_transition_cycles = scheduler_cycles(); + scheduler_writeback_cache_line(&cells[task_id].ready_transition_cycles); + } + ++batch->count; + if (stats != nullptr) ++stats->enqueue_count; + return true; +} + +inline __aicore__ bool scheduler_bootstrap_ready_batch_publish( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, uint32_t core_type_index, + uint64_t inbox_index, SchedulerReadyBatch *batch, SchedulerReadyStats *stats, uint64_t *ready_types +) { + if (batch == nullptr || batch->head == SCHEDULER_INBOX_EMPTY) return true; + if (core_type_index >= SCHEDULER_CORE_TYPE_COUNT || inbox_index >= SCHEDULER_WORKER_CAPACITY || batch->tail < 0 || + ready_types == nullptr) + return false; + __gm__ SchedulerReadyInbox *inbox = + scheduler_ready_inbox_at(scheduler_state_base, context, core_type_index, inbox_index); + scheduler_gm_store(inbox->head, batch->head); + *ready_types |= UINT64_C(1) << core_type_index; + if (stats != nullptr) ++stats->batch_count; + *batch = {}; + return true; +} + +inline __aicore__ void scheduler_bootstrap_ready_directory_publish( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, uint64_t resolver_count +) { + __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); + for (uint64_t inbox_index = 0; inbox_index < resolver_count; inbox_index += 8) + scheduler_invalidate_cache_line(&directory->bootstrap_ready_types[inbox_index]); + scheduler_cache_barrier(); + uint32_t shard_count = static_cast( + (resolver_count + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD - 1) / + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD + ); + for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + for (uint32_t shard = 0; shard < shard_count; ++shard) { + uint64_t bits = 0; + uint64_t shard_begin = static_cast(shard) * SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t shard_end = shard_begin + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + if (shard_end > resolver_count) shard_end = resolver_count; + for (uint64_t inbox_index = shard_begin; inbox_index < shard_end; ++inbox_index) { + uint64_t ready_types = directory->bootstrap_ready_types[inbox_index]; + if ((ready_types & (UINT64_C(1) << type)) != 0) bits |= UINT64_C(1) << (inbox_index - shard_begin); + } + directory->core_types[type][shard].bits = bits; + scheduler_publish_cache_line(&directory->core_types[type][shard]); + } + } +} + +inline __aicore__ bool scheduler_ready_batch_append( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, int64_t task_id, + SchedulerReadyBatch *batch, SchedulerReadyStats *stats, bool trace_enabled = false +) { + if (batch == nullptr || task_id < 0 || static_cast(task_id) >= context->graph_task_count) return false; + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); + scheduler_gm_store(control->state, static_cast(SchedulerTaskState::READY)); + control->next_waiter = SCHEDULER_INBOX_EMPTY; + scheduler_publish_cache_line(&control->next_waiter); + if (batch->head == SCHEDULER_INBOX_EMPTY) { + batch->head = task_id; + } else { + __gm__ SchedulerTaskControl *tail = scheduler_task_control_at(scheduler_state_base, context, batch->tail); + scheduler_observe_cache_line(&tail->next_waiter); + tail->next_waiter = task_id; + scheduler_publish_cache_line(&tail->next_waiter); + } + batch->tail = task_id; + if (trace_enabled) { + __gm__ SchedulerTaskTrace *cells = + scheduler_state_at(scheduler_state_base, context->trace_cells_offset); + __gm__ SchedulerTaskTrace *trace = &cells[task_id]; + scheduler_observe_cache_line(&trace->ready_transition_cycles); + trace->ready_transition_cycles = scheduler_cycles(); + scheduler_publish_cache_line(&trace->ready_transition_cycles); + } + ++batch->count; + if (stats != nullptr) ++stats->enqueue_count; + return true; +} + +inline __aicore__ void scheduler_ready_directory_set( + __gm__ SchedulerReadyDirectory *directory, uint32_t core_type_index, uint64_t inbox_index +) { + uint64_t shard = inbox_index / SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t bit = UINT64_C(1) << (inbox_index % SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD); + scheduler_gm_fetch_or(directory->core_types[core_type_index][shard].bits, bit); +} + +inline __aicore__ void scheduler_ready_directory_clear( + __gm__ SchedulerReadyDirectory *directory, uint32_t core_type_index, uint64_t inbox_index +) { + uint64_t shard = inbox_index / SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t bit = UINT64_C(1) << (inbox_index % SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD); + scheduler_gm_fetch_and(directory->core_types[core_type_index][shard].bits, ~bit); +} + +inline __aicore__ uint64_t scheduler_ready_owner_pending_load(__gm__ SchedulerReadyOwnerQueue *owner_queue) { + return scheduler_gm_query(owner_queue->pending_endpoints); +} + +inline __aicore__ bool +scheduler_ready_owner_pending_store(__gm__ SchedulerReadyOwnerQueue *owner_queue, int64_t head, int64_t tail) { + if (!scheduler_ready_pending_endpoint_fits(head) || !scheduler_ready_pending_endpoint_fits(tail)) return false; + if ((head == SCHEDULER_INBOX_EMPTY) != (tail == SCHEDULER_INBOX_EMPTY)) return false; + scheduler_gm_store(owner_queue->pending_endpoints, scheduler_ready_pending_pack(head, tail)); + scheduler_cache_barrier(); + return true; +} + +inline __aicore__ void scheduler_ready_owner_pending_reset(__gm__ SchedulerReadyOwnerQueue *owner_queue) { + scheduler_gm_store(owner_queue->pending_endpoints, SCHEDULER_READY_PENDING_EMPTY); + scheduler_cache_barrier(); +} + +inline __aicore__ bool scheduler_ready_owner_pending_append( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerReadyOwnerQueue *owner_queue, SchedulerReadyBatch *source +) { + // Only the owner mutates pending links; links reachable from a shared inbox head stay immutable. + if (owner_queue == nullptr || source == nullptr || source->head == SCHEDULER_INBOX_EMPTY) return true; + if (!scheduler_ready_pending_endpoint_fits(source->head) || source->head < 0 || + !scheduler_ready_pending_endpoint_fits(source->tail) || source->tail < 0) + return false; + const uint64_t pending = scheduler_ready_owner_pending_load(owner_queue); + const int64_t pending_head = scheduler_ready_pending_head(pending); + const int64_t pending_tail = scheduler_ready_pending_tail(pending); + if (pending_head == SCHEDULER_INBOX_EMPTY) { + if (pending_tail != SCHEDULER_INBOX_EMPTY || + !scheduler_ready_owner_pending_store(owner_queue, source->head, source->tail)) + return false; + scheduler_ready_batch_reset(source); + return true; + } + if (pending_head < 0 || pending_tail < 0) return false; + __gm__ SchedulerTaskControl *tail = scheduler_task_control_at(scheduler_state_base, context, pending_tail); + scheduler_observe_cache_line(&tail->next_waiter); + tail->next_waiter = source->head; + scheduler_publish_cache_line(&tail->next_waiter); + if (!scheduler_ready_owner_pending_store(owner_queue, pending_head, source->tail)) return false; + scheduler_ready_batch_reset(source); + return true; +} + +inline __aicore__ bool scheduler_ready_owner_maintain_type( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, uint32_t core_type_index, + __gm__ SchedulerReadyOwnerState *owner_state +) { + if (owner_state == nullptr || core_type_index >= SCHEDULER_CORE_TYPE_COUNT || + context->inbox_index >= SCHEDULER_WORKER_CAPACITY) + return false; + __gm__ SchedulerReadyOwnerQueue *owner_queue = &owner_state->queues[core_type_index]; + __gm__ SchedulerReadyInbox *inbox = + scheduler_ready_inbox_at(scheduler_state_base, context, core_type_index, context->inbox_index); + __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); + const int64_t head = scheduler_gm_query(inbox->head); + if (head < SCHEDULER_INBOX_EMPTY) return false; + if (head != SCHEDULER_INBOX_EMPTY) { + if (scheduler_gm_query(owner_queue->advertised) == 0) { + scheduler_ready_directory_set(directory, core_type_index, context->inbox_index); + scheduler_gm_store(owner_queue->advertised, UINT64_C(1)); + } + return true; + } + const uint64_t pending = scheduler_ready_owner_pending_load(owner_queue); + const int64_t pending_head = scheduler_ready_pending_head(pending); + const int64_t pending_tail = scheduler_ready_pending_tail(pending); + if (pending_head != SCHEDULER_INBOX_EMPTY) { + if (pending_head < 0 || pending_tail < 0) return false; + scheduler_cache_barrier(); + scheduler_gm_store(inbox->head, pending_head); + scheduler_cache_barrier(); + scheduler_ready_owner_pending_reset(owner_queue); + if (scheduler_gm_query(owner_queue->advertised) == 0) { + scheduler_ready_directory_set(directory, core_type_index, context->inbox_index); + scheduler_gm_store(owner_queue->advertised, UINT64_C(1)); + } + return true; + } else if (pending_tail != SCHEDULER_INBOX_EMPTY) { + return false; + } + if (scheduler_gm_query(owner_queue->advertised) != 0) { + scheduler_ready_directory_clear(directory, core_type_index, context->inbox_index); + scheduler_gm_store(owner_queue->advertised, UINT64_C(0)); + } + return true; +} + +inline __aicore__ bool scheduler_ready_owner_maintain( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerReadyOwnerState *owner_state +) { + for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + if (!scheduler_ready_owner_maintain_type(scheduler_state_base, context, type, owner_state)) return false; + } + return true; +} + +inline __aicore__ bool scheduler_ready_batch_push( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, uint32_t core_type_index, + uint64_t inbox_index, SchedulerReadyBatch *batch, SchedulerReadyStats *stats, + __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + if (batch == nullptr || batch->head == SCHEDULER_INBOX_EMPTY) return true; + if (core_type_index >= SCHEDULER_CORE_TYPE_COUNT || inbox_index >= SCHEDULER_WORKER_CAPACITY || batch->tail < 0) + return false; + __gm__ SchedulerReadyInbox *inbox = + scheduler_ready_inbox_at(scheduler_state_base, context, core_type_index, inbox_index); + if (owner_state == nullptr) { + if (scheduler_gm_query(inbox->head) != SCHEDULER_INBOX_EMPTY) return false; + scheduler_cache_barrier(); + scheduler_gm_store(inbox->head, batch->head); + scheduler_ready_directory_set( + scheduler_ready_directory_at(scheduler_state_base, context), core_type_index, inbox_index + ); + } else { + if (inbox_index != context->inbox_index) return false; + __gm__ SchedulerReadyOwnerQueue *owner_queue = &owner_state->queues[core_type_index]; + const int64_t head = scheduler_gm_query(inbox->head); + if (head < SCHEDULER_INBOX_EMPTY) return false; + const uint64_t pending = scheduler_ready_owner_pending_load(owner_queue); + const int64_t pending_head = scheduler_ready_pending_head(pending); + const int64_t pending_tail = scheduler_ready_pending_tail(pending); + if (head == SCHEDULER_INBOX_EMPTY && pending_head != SCHEDULER_INBOX_EMPTY) { + if (pending_head < 0 || pending_tail < 0) return false; + scheduler_cache_barrier(); + scheduler_gm_store(inbox->head, pending_head); + scheduler_cache_barrier(); + scheduler_ready_owner_pending_reset(owner_queue); + if (scheduler_gm_query(owner_queue->advertised) == 0) { + scheduler_ready_directory_set( + scheduler_ready_directory_at(scheduler_state_base, context), core_type_index, inbox_index + ); + scheduler_gm_store(owner_queue->advertised, UINT64_C(1)); + } + } else if (pending_head == SCHEDULER_INBOX_EMPTY && pending_tail != SCHEDULER_INBOX_EMPTY) { + return false; + } + const int64_t published_head = scheduler_gm_query(inbox->head); + if (published_head == SCHEDULER_INBOX_EMPTY) { + scheduler_cache_barrier(); + scheduler_gm_store(inbox->head, batch->head); + if (scheduler_gm_query(owner_queue->advertised) == 0) { + scheduler_ready_directory_set( + scheduler_ready_directory_at(scheduler_state_base, context), core_type_index, inbox_index + ); + scheduler_gm_store(owner_queue->advertised, UINT64_C(1)); + } + scheduler_ready_batch_reset(batch); + } else { + if (published_head < SCHEDULER_INBOX_EMPTY) return false; + if (scheduler_gm_query(owner_queue->advertised) == 0) { + scheduler_ready_directory_set( + scheduler_ready_directory_at(scheduler_state_base, context), core_type_index, inbox_index + ); + scheduler_gm_store(owner_queue->advertised, UINT64_C(1)); + } + if (!scheduler_ready_owner_pending_append(scheduler_state_base, context, owner_queue, batch)) return false; + } + } + if (stats != nullptr) ++stats->batch_count; + if (owner_state == nullptr) scheduler_ready_batch_reset(batch); + return true; +} + +inline __aicore__ bool scheduler_ready_pop_from_inbox( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, uint32_t core_type_index, uint64_t inbox_index, int64_t *task_id, + SchedulerReadyStats *stats +) { + if (task_id == nullptr) return false; + *task_id = SCHEDULER_TASK_ID_INVALID; + __gm__ SchedulerReadyInbox *inbox = + scheduler_ready_inbox_at(scheduler_state_base, context, core_type_index, inbox_index); + for (uint32_t attempt = 0; attempt < 64; ++attempt) { + int64_t head = scheduler_gm_query(inbox->head); + if (head == SCHEDULER_INBOX_EMPTY) return true; + if (head < 0 || static_cast(head) >= graph.task_count) { + scheduler_record_error( + run_control, head, SchedulerGraphResult::INVALID_TASK_ID, &graph, context, + SchedulerErrorSite::READY_POP_INVALID_HEAD + ); + return false; + } + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, head); + int64_t next = scheduler_observe_next_waiter(control); + uint64_t waits = 0; + bool head_changed = false; + while (next == SCHEDULER_INBOX_LINK_UNPUBLISHED && waits < UINT64_C(1048576)) { + if (scheduler_gm_query(inbox->head) != head) { + head_changed = true; + break; + } + ++waits; + next = scheduler_observe_next_waiter(control); + } + if (stats != nullptr) { + stats->link_wait_count += waits; + if (waits > stats->link_wait_max) stats->link_wait_max = waits; + } + if (head_changed) { + if (stats != nullptr) ++stats->cas_retry_count; + continue; + } + if (next < SCHEDULER_INBOX_EMPTY) { + scheduler_record_error( + run_control, head, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::READY_POP_INVALID_LINK + ); + return false; + } + int64_t actual = scheduler_gm_compare_exchange(inbox->head, head, next); + if (actual != head) { + if (stats != nullptr) ++stats->cas_retry_count; + continue; + } + if (stats != nullptr) ++stats->pop_count; + *task_id = head; + return true; + } + return true; +} + +inline __aicore__ uint64_t scheduler_load_ready_directory_shard( + __gm__ SchedulerReadyDirectory *directory, uint64_t resolver_count, uint32_t core_type_index, uint64_t inbox_index +) { + uint64_t shard = inbox_index / SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t shard_begin = shard * SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t shard_end = shard_begin + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + if (shard_end > resolver_count) shard_end = resolver_count; + uint64_t valid_bits = shard_end > shard_begin ? (UINT64_C(1) << (shard_end - shard_begin)) - 1 : 0; + return scheduler_gm_query(directory->core_types[core_type_index][shard].bits) & valid_bits; +} + +static __attribute__((noinline)) __aicore__ bool scheduler_steal_ready_from_shard( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, uint32_t core_type_index, uint64_t shard_begin, uint64_t shard_end, + uint64_t start, uint64_t *victim_cursor, uint64_t bits, SchedulerReadyStats *stats, SchedulerReadyClaim *claim, + bool trace_enabled +) { + int64_t task_id = SCHEDULER_TASK_ID_INVALID; + bits &= ~(UINT64_C(1) << (context->inbox_index - shard_begin)); + for (uint32_t pass = 0; pass < 2; ++pass) { + uint64_t range_begin = pass == 0 ? start : shard_begin; + uint64_t range_end = pass == 0 ? shard_end : start; + if (range_begin == range_end) continue; + uint32_t lower_bit = static_cast(range_begin - shard_begin); + uint32_t upper_bit = static_cast(range_end - shard_begin); + uint64_t candidates = bits & (((UINT64_C(1) << upper_bit) - 1) & ~((UINT64_C(1) << lower_bit) - 1)); + while (candidates != 0) { + uint32_t bit_index = static_cast(__builtin_ctzll(candidates)); + candidates &= candidates - 1; + uint64_t victim = shard_begin + bit_index; + *victim_cursor = victim + 1 == shard_end ? shard_begin : victim + 1; + if (!scheduler_ready_pop_from_inbox( + graph, scheduler_state_base, context, run_control, core_type_index, victim, &task_id, stats + )) + return false; + if (task_id >= 0) { + claim->task_id = task_id; + claim->inbox_index = victim; + claim->source = SchedulerReadySource::STOLEN; + claim->claim_end_cycles = trace_enabled ? scheduler_cycles() : 0; + if (stats != nullptr) ++stats->steal_count; + return true; + } + } + } + return true; +} + +inline __aicore__ bool scheduler_claim_ready_for_slot( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, uint64_t resolver_count, uint32_t core_type_index, uint64_t *victim_cursor, + SchedulerReadyStats *stats, SchedulerReadyClaim *claim, bool trace_enabled = false, + __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + if (victim_cursor == nullptr || claim == nullptr || resolver_count == 0) return false; + *claim = {}; + claim->claim_start_cycles = trace_enabled ? scheduler_cycles() : 0; + if (owner_state != nullptr && + !scheduler_ready_owner_maintain_type(scheduler_state_base, context, core_type_index, owner_state)) + return false; + int64_t task_id = SCHEDULER_TASK_ID_INVALID; + if (!scheduler_ready_pop_from_inbox( + graph, scheduler_state_base, context, run_control, core_type_index, context->inbox_index, &task_id, stats + )) + return false; + if (task_id >= 0) { + claim->task_id = task_id; + claim->inbox_index = context->inbox_index; + claim->claim_end_cycles = trace_enabled ? scheduler_cycles() : 0; + return true; + } + + __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); + uint64_t shard_begin = context->inbox_index / SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD * + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + uint64_t shard_end = shard_begin + SCHEDULER_READY_DIRECTORY_RESOLVERS_PER_SHARD; + if (shard_end > resolver_count) shard_end = resolver_count; + uint64_t start = *victim_cursor; + if (start < shard_begin || start >= shard_end) start = shard_begin; + uint64_t bits = + scheduler_load_ready_directory_shard(directory, resolver_count, core_type_index, context->inbox_index); + if (bits != 0 && !scheduler_steal_ready_from_shard( + graph, scheduler_state_base, context, run_control, core_type_index, shard_begin, shard_end, + start, victim_cursor, bits, stats, claim, trace_enabled + )) + return false; + if (claim->task_id >= 0) return true; + *victim_cursor = start + 1 == shard_end ? shard_begin : start + 1; + claim->claim_end_cycles = trace_enabled ? scheduler_cycles() : 0; + return true; +} + +inline __aicore__ bool scheduler_ready_directory_nonempty( + __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, uint64_t resolver_count, + uint32_t core_type_index +) { + __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); + return scheduler_load_ready_directory_shard(directory, resolver_count, core_type_index, context->inbox_index) != 0; +} + +inline __aicore__ void scheduler_initialize_free_slot(__gm__ SchedulerDispatchSlot *slot) { + uint32_t generation = slot->generation + 1; + if (generation == 0) generation = 1; + slot->task_id = SCHEDULER_TASK_ID_INVALID; + slot->generation = generation; + scheduler_publish_cache_line(slot); + scheduler_gm_publish( + slot->publication, scheduler_dispatch_publication(generation, SchedulerDispatchSlotState::FREE) + ); +} + +inline __aicore__ bool scheduler_fill_dispatch_slot( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, const SchedulerFreeSlotClaim &slot_claim, + const SchedulerReadyClaim &ready_claim, bool trace_enabled = false, SchedulerDispatchFillTiming *timing = nullptr +) { + if (ready_claim.task_id < 0 || static_cast(ready_claim.task_id) >= graph.task_count || + slot_claim.worker_id >= resolver->runtime_worker_count || slot_claim.slot_index >= SCHEDULER_PENDING_SLOT_COUNT) + return false; + const bool record_timeline = timing != nullptr; + uint64_t operation_start = record_timeline ? scheduler_cycles() : 0; + __gm__ SchedulerTaskMetadata *metadata_source = + scheduler_task_metadata_at(scheduler_state_base, resolver, ready_claim.task_id); + scheduler_observe_cache_line(metadata_source); + SchedulerTaskMetadata metadata{}; + metadata.kernel_ids[0] = metadata_source->kernel_ids[0]; + metadata.kernel_ids[1] = metadata_source->kernel_ids[1]; + metadata.kernel_ids[2] = metadata_source->kernel_ids[2]; + metadata.active_mask = metadata_source->active_mask; + metadata.flags = metadata_source->flags; + metadata.logical_block_num = metadata_source->logical_block_num; + metadata.total_required_subtasks = metadata_source->total_required_subtasks; + metadata.timing_slot = metadata_source->timing_slot; + const uint8_t subtask_slot = scheduler_metadata_single_subtask_slot(metadata.active_mask); + const uint16_t kernel_id = metadata.kernel_ids[subtask_slot]; + __gm__ SchedulerWorkerContext *target = + scheduler_worker_context_at(scheduler_state_base, resolver, slot_claim.worker_id); + scheduler_observe_cache_line(target); + if (!scheduler_task_is_executable(metadata.flags) || scheduler_task_is_gang(metadata.flags) || + scheduler_metadata_core_type_index(subtask_slot) != scheduler_core_type_index(target->core_type)) { + scheduler_record_error( + run_control, ready_claim.task_id, SchedulerGraphResult::UNSUPPORTED_SHAPE, &graph, resolver, + SchedulerErrorSite::DISPATCH_INVALID_SHAPE + ); + return false; + } + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, slot_claim.worker_id, slot_claim.slot_index); + uint32_t generation = slot_claim.generation + 1; + if (generation == 0) generation = 1; + __gm__ uint64_t *callable_addresses = + scheduler_state_at(scheduler_state_base, resolver->callable_addresses_offset); + const bool inline_task = scheduler_task_is_inline(metadata.flags); + uint64_t callable_address = UINT64_C(1); + if (!inline_task && !scheduler_lookup_callable_address(callable_addresses, kernel_id, &callable_address)) { + scheduler_record_error( + run_control, ready_claim.task_id, SchedulerGraphResult::INVALID_CALLABLE, &graph, resolver, + SchedulerErrorSite::DISPATCH_INVALID_CALLABLE + ); + return false; + } + + slot->task_id = ready_claim.task_id; + slot->ready_inbox_index = ready_claim.inbox_index; + slot->claim_start_cycles = ready_claim.claim_start_cycles; + slot->claim_end_cycles = ready_claim.claim_end_cycles; + slot->claim_worker_id = resolver->worker_index; + slot->kernel_id = kernel_id; + slot->subtask_slot = subtask_slot; + slot->has_fanin = scheduler_task_has_fanin(metadata.flags) ? 1 : 0; + slot->ready_source = static_cast(ready_claim.source); + slot->pending_slot = static_cast(slot_claim.slot_index); + slot->generation = generation; + slot->block_idx = 0; + slot->block_num = 1; + slot->cohort_generation = 0; + slot->cohort_index = UINT8_MAX; + slot->gang = 0; + scheduler_writeback_cache_line(slot); + + const uint64_t dispatch_payload_offset = + target->dispatch_payload_offset + static_cast(slot_claim.slot_index) * sizeof(DispatchPayload); + + uint64_t operation_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) timing->prepare_cycles += operation_end - operation_start; + + SchedulerTaskInfo task{ + ready_claim.task_id, + static_cast(kernel_id), + static_cast(subtask_slot), + subtask_slot == 0 ? CoreType::AIC : CoreType::AIV, + }; + __gm__ DispatchPayload *payload = + scheduler_state_at(scheduler_state_base, dispatch_payload_offset); + SchedulerGraphResult status = SchedulerGraphResult::OK; + if (inline_task) { + payload->function_bin_addr = 0; + payload->src_payload = 0; + } else { + status = scheduler_materialize_task_payload_resolved(graph, task, callable_address, payload); + if (status == SchedulerGraphResult::OK && scheduler_task_has_predicate(metadata.flags)) { + const SchedulerPredicateResult predicate = scheduler_evaluate_task_predicate(graph, ready_claim.task_id); + if (predicate == SchedulerPredicateResult::MALFORMED) { + scheduler_record_error( + run_control, ready_claim.task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::DISPATCH_INVALID_PREDICATE + ); + return false; + } + if (predicate == SchedulerPredicateResult::FAIL) payload->function_bin_addr = 0; + } + } + if (status != SchedulerGraphResult::OK) { + scheduler_record_error( + run_control, ready_claim.task_id, status, &graph, resolver, SchedulerErrorSite::DISPATCH_MATERIALIZE_FAILED + ); + return false; + } + uint64_t materialize_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) timing->materialize_cycles += materialize_end - operation_end; + scheduler_publish_dispatch_payload(payload); + __gm__ SchedulerTaskControl *control = + scheduler_task_control_at(scheduler_state_base, resolver, ready_claim.task_id); + if (trace_enabled) { + scheduler_observe_cache_line(&control->next_waiter); + control->ready_publish_cycles = scheduler_cycles(); + scheduler_publish_cache_line(&control->next_waiter); + } + uint64_t publish_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) timing->publish_cycles += publish_end - materialize_end; + scheduler_gm_publish( + slot->publication, scheduler_dispatch_publication(generation, SchedulerDispatchSlotState::READY) + ); + return true; +} + +inline __aicore__ bool scheduler_resolve_completion( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *context, + __gm__ SchedulerRunControl *run_control, int64_t task_id, SchedulerWakeStats *wake_stats, + SchedulerReadyStats *ready_stats, SchedulerCompletionStats *completion_stats, bool trace_enabled = false, + bool validate_done_state = true, uint64_t *ready_publish_cycles = nullptr, + __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); + if (validate_done_state && scheduler_gm_query(control->state) != static_cast(SchedulerTaskState::DONE)) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::COMPLETION_TASK_NOT_DONE + ); + return false; + } + uint64_t resolve_start = trace_enabled ? scheduler_cycles() : 0; + if (trace_enabled) { + scheduler_observe_cache_line(&control->next_waiter); + control->completion_resolve_start_cycles = resolve_start; + control->resolver_worker_id = context->worker_index; + } + int64_t waiter = scheduler_gm_exchange(control->wake_list_head, SCHEDULER_WAKE_LIST_CLOSED); + if (waiter == SCHEDULER_WAKE_LIST_CLOSED) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::COMPLETION_WAKE_ALREADY_CLOSED + ); + return false; + } + if (wake_stats != nullptr) ++wake_stats->wake_close_count; + SchedulerReadyBatch batches[SCHEDULER_CORE_TYPE_COUNT]{}; + while (waiter >= 0) { + if (static_cast(waiter) >= graph.task_count) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_TASK_ID, &graph, context, + SchedulerErrorSite::COMPLETION_INVALID_WAITER + ); + return false; + } + __gm__ SchedulerTaskControl *waiter_control = scheduler_task_control_at(scheduler_state_base, context, waiter); + int64_t next = scheduler_observe_next_waiter(waiter_control); + if (wake_stats != nullptr) ++wake_stats->wake_migrate_count; + SchedulerRouteResult route = + scheduler_route_task(graph, scheduler_state_base, context, run_control, waiter, wake_stats); + if (route == SchedulerRouteResult::ERROR) return false; + if (route == SchedulerRouteResult::READY_TO_ENQUEUE) { + __gm__ SchedulerTaskMetadata *metadata = scheduler_task_metadata_at(scheduler_state_base, context, waiter); + scheduler_observe_cache_line(metadata); + if (!scheduler_task_is_executable(metadata->flags)) { + scheduler_record_error( + run_control, waiter, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::COMPLETION_WAITER_NOT_EXECUTABLE + ); + return false; + } + if (scheduler_task_is_gang(metadata->flags)) { + scheduler_publish_gang_ready(scheduler_state_base, context, waiter_control, metadata->flags); + } else if (!scheduler_ready_batch_append( + scheduler_state_base, context, waiter, + &batches[scheduler_metadata_core_type_index( + scheduler_metadata_single_subtask_slot(metadata->active_mask) + )], + ready_stats, trace_enabled + )) { + scheduler_record_error( + run_control, waiter, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::COMPLETION_READY_APPEND_FAILED + ); + return false; + } + } + waiter = next; + } + uint64_t ready_publish_start = ready_publish_cycles == nullptr ? 0 : scheduler_cycles(); + for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + if (!scheduler_ready_batch_push( + scheduler_state_base, context, type, context->inbox_index, &batches[type], ready_stats, owner_state + )) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, + SchedulerErrorSite::COMPLETION_READY_PUBLISH_FAILED + ); + return false; + } + } + if (ready_publish_cycles != nullptr) *ready_publish_cycles += scheduler_cycles() - ready_publish_start; + if (trace_enabled) { + control->completion_resolve_end_cycles = scheduler_cycles(); + scheduler_publish_cache_line(&control->next_waiter); + } + if (completion_stats != nullptr) ++completion_stats->resolve_count; + return true; +} diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_types.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_types.h index e3307c0cbd..bec037451b 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_types.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_types.h @@ -895,6 +895,26 @@ struct alignas(128) SchedulerDispatchSlot { uint8_t publication_padding[56]; }; +// Stable device-side localization for the first scheduler failure. These values +// are diagnostic ABI: keep existing numbers stable when adding new sites. +enum class SchedulerErrorSite : uint64_t { + UNKNOWN = 0, + ROUTE_INVALID_STATE = 2, + READY_POP_INVALID_HEAD = 40, + READY_POP_INVALID_LINK = 41, + DISPATCH_INVALID_SHAPE = 44, + DISPATCH_INVALID_CALLABLE = 45, + DISPATCH_MATERIALIZE_FAILED = 46, + DISPATCH_INVALID_PREDICATE = 47, + BOOTSTRAP_WAKE_INVALID_HEAD = 60, + COMPLETION_TASK_NOT_DONE = 61, + COMPLETION_WAKE_ALREADY_CLOSED = 62, + COMPLETION_INVALID_WAITER = 63, + COMPLETION_WAITER_NOT_EXECUTABLE = 64, + COMPLETION_READY_APPEND_FAILED = 65, + COMPLETION_READY_PUBLISH_FAILED = 66, +}; + struct alignas(128) SchedulerRunControl { uint64_t config_reserved_prefix[2]; uint64_t active_worker_count; @@ -936,7 +956,8 @@ struct alignas(128) SchedulerRunControl { volatile uint64_t error_descriptors_address; volatile uint64_t error_payloads_address; volatile uint64_t error_task_window_mask; - uint64_t error_reserved[7]; + volatile uint64_t error_site; + uint64_t error_reserved[6]; }; struct alignas(128) AicpuCoreLifecycleTrace { @@ -1201,6 +1222,7 @@ static_assert(sizeof(SchedulerRunControl) == 384, "run control layout changed"); static_assert(alignof(SchedulerRunControl) == 128, "run control alignment changed"); static_assert(offsetof(SchedulerRunControl, executed_task_count) == 128, "lifecycle atomics need their own line"); static_assert(offsetof(SchedulerRunControl, error_claimed) == 256, "error state needs its own line"); +static_assert(offsetof(SchedulerRunControl, error_site) == 328, "error site ABI changed"); static_assert(sizeof(AicpuCoreLifecycleTrace) == 128, "AICPU lifecycle trace layout changed"); static_assert(sizeof(SchedulerTailTrace) == 256, "scheduler tail trace layout changed"); static_assert(sizeof(SchedulerWorkerContext) == 1024, "worker context layout changed"); diff --git a/src/common/host_build_graph/shared/orchestrator.cpp b/src/common/host_build_graph/shared/orchestrator.cpp index 6b12e117af..38c3cc0fab 100644 --- a/src/common/host_build_graph/shared/orchestrator.cpp +++ b/src/common/host_build_graph/shared/orchestrator.cpp @@ -1571,6 +1571,73 @@ static bool ensure_tensormap_capacity(OrchestratorState *orch, int32_t needed) { return false; } +static bool +resolve_dispatch_predicate(OrchestratorState *orch, const CoreTaskPredicate &predicate, DispatchPredicate *resolved) { + if (resolved == nullptr) return false; + *resolved = DispatchPredicate{}; + if (predicate.op == PredicateOp::NONE) return true; + + switch (predicate.op) { + case PredicateOp::EQ: + case PredicateOp::NE: + case PredicateOp::GT: + case PredicateOp::LT: + case PredicateOp::GE: + case PredicateOp::LE: + break; + case PredicateOp::NONE: + return true; + default: + orch->report_fatal(SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "dispatch predicate has an invalid operator"); + return false; + } + + const simpler::hbg::Tensor *operand = predicate.operand.tensor; + if (operand == nullptr || operand->buffer.addr == 0 || predicate.operand.ndims == 0 || + predicate.operand.ndims > operand->ndims || predicate.operand.ndims > MAX_TENSOR_DIMS) { + orch->report_fatal( + SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "dispatch predicate has an invalid operand tensor" + ); + return false; + } + + uint64_t element_offset = operand->start_offset; + for (uint32_t dim = 0; dim < predicate.operand.ndims; ++dim) { + if (predicate.operand.indices[dim] >= operand->shapes[dim] || + (predicate.operand.indices[dim] != 0 && + operand->strides[dim] > + (UINT64_MAX - element_offset) / static_cast(predicate.operand.indices[dim]))) { + orch->report_fatal( + SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "dispatch predicate index is outside the operand tensor" + ); + return false; + } + element_offset += + static_cast(predicate.operand.indices[dim]) * static_cast(operand->strides[dim]); + } + + const uint64_t element_size = get_element_size(operand->dtype); + if ((element_size != 1 && element_size != 2 && element_size != 4 && element_size != 8) || + operand->buffer.size < element_size || element_offset > (operand->buffer.size - element_size) / element_size) { + orch->report_fatal( + SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "dispatch predicate element is outside the operand buffer" + ); + return false; + } + const uint64_t byte_offset = element_offset * element_size; + if (operand->buffer.addr > UINT64_MAX - byte_offset || + ((operand->buffer.addr + byte_offset) & (element_size - 1)) != 0) { + orch->report_fatal(SIMPLER_ERROR_INVALID_ARGS, __FUNCTION__, "dispatch predicate operand address is invalid"); + return false; + } + + resolved->addr = operand->buffer.addr + byte_offset; + resolved->target = predicate.target; + resolved->elem_size = static_cast(element_size); + resolved->op = predicate.op; + return true; +} + // Shared body for submit_task / submit_dummy_task. Caller has already validated // args.has_error, decided active_mask (empty for dummy), and resolved the per-slot // kernel_ids (all INVALID_KERNEL_ID for dummy). Performs tensormap sync, fanin @@ -1583,6 +1650,8 @@ static TaskOutputTensors submit_task_common( CYCLE_COUNT_START(); ORCH_PHASE_START(); TaskOutputTensors result; + DispatchPredicate resolved_predicate{}; + if (!resolve_dispatch_predicate(orch, args.predicate(), &resolved_predicate)) return result; OutputLayout layout = calculate_output_layout(args); PreparedTask prepared; if (!prepare_task(orch, args, layout.total_output_size, active_mask, task_attrs, &prepared)) { @@ -1703,25 +1772,9 @@ static TaskOutputTensors submit_task_common( // must not touch either, or it would discard that. payload.init(args, result, prepared.alloc_result, layout); - // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address - // now so the scheduler can read it at the dispatch point with a single load, - // no Arg/simpler::hbg::Tensor access. Both branches write predicate.op explicitly because - // a payload slot is raw shared memory with no constructor; op == NONE means - // "always dispatch". - { - const CoreTaskPredicate &pred = args.predicate(); - if (pred.op != PredicateOp::NONE && pred.operand.tensor != nullptr && pred.operand.tensor->buffer.addr != 0) { - uint64_t elem_size = get_element_size(pred.operand.tensor->dtype); - uint64_t flat_offset = pred.operand.tensor->compute_flat_offset(pred.operand.indices, pred.operand.ndims); - payload.predicate.addr = pred.operand.tensor->buffer.addr + flat_offset * elem_size; - payload.predicate.target = pred.target; - payload.predicate.elem_size = static_cast(elem_size); - payload.predicate.op = pred.op; - } else { - payload.predicate.addr = 0; - payload.predicate.op = PredicateOp::NONE; - } - } + // Predicate validation runs before task allocation. Copy the resolved, bounded + // operand address into the device payload only after the rest of the payload exists. + payload.predicate = resolved_predicate; CYCLE_COUNT_LAP(g_orch_args_cycle); // === STEP 6: close the fanin region (device boot classifies) === diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 366b511c59..349362b780 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -939,6 +939,7 @@ target_sources(test_hbg_submit_poison PRIVATE ${CMAKE_SOURCE_DIR}/../../../src/common/platform/shared/aicpu/args_dump_aicpu.cpp ) add_a5_hbg_runtime_test(test_a5_hbg_scheduler_contracts a5/test_hbg_scheduler_contracts.cpp) +add_a5_hbg_runtime_test(test_a5_hbg_scheduler_ready a5/test_hbg_scheduler_ready.cpp) add_a2a3_hbg_runtime_test(test_hbg_graph_submit_failure common/test_hbg_graph_submit_failure.cpp) target_sources(test_hbg_graph_submit_failure PRIVATE ${HBG_ORCH_SHARED_SOURCES} diff --git a/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp b/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp index a945c26464..28aa3d7702 100644 --- a/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp +++ b/tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp @@ -26,6 +26,7 @@ #include +#include #include #include #include @@ -180,3 +181,41 @@ TEST_F(HbgSubmitPoisonTest, EveryDeviceReadFieldIsWrittenOverPoison) { EXPECT_EQ(cons_pl.fanin_count, 1); EXPECT_EQ(cons_pl.fanin_data()[0], static_cast(simpler::hbg::task_local_id(root.task_id()))); } + +TEST_F(HbgSubmitPoisonTest, InvalidDispatchPredicateIsRejectedBeforeTaskAllocation) { + orch.begin_scope(); + CoreTaskArgs args; + CoreTaskPredicate predicate; + predicate.op = PredicateOp::GT; + predicate.target = 0; + args.set_predicate(predicate); + MixedKernels kernels{}; + kernels.aiv0_kernel_id = 0; + + EXPECT_FALSE(orch.submit_task(kernels, args).task_id().is_valid()); + EXPECT_TRUE(orch.fatal); + EXPECT_EQ(orch.task_allocator.active_count(), 0); + EXPECT_EQ(sm_handle->header->orch_error_code.load(std::memory_order_acquire), SIMPLER_ERROR_INVALID_ARGS); +} + +TEST_F(HbgSubmitPoisonTest, OutOfRangeDispatchPredicateIndexIsRejected) { + orch.begin_scope(); + std::array values{}; + uint32_t shape[] = {static_cast(values.size())}; + simpler::hbg::Tensor operand = simpler::hbg::make_tensor_external(values.data(), shape, 1, DataType::INT32); + CoreTaskPredicate predicate; + predicate.operand.tensor = &operand; + predicate.operand.ndims = 1; + predicate.operand.indices[0] = static_cast(values.size()); + predicate.op = PredicateOp::GT; + CoreTaskArgs args; + args.add_input(operand); + args.set_predicate(predicate); + MixedKernels kernels{}; + kernels.aiv0_kernel_id = 0; + + EXPECT_FALSE(orch.submit_task(kernels, args).task_id().is_valid()); + EXPECT_TRUE(orch.fatal); + EXPECT_EQ(orch.task_allocator.active_count(), 0); + EXPECT_EQ(sm_handle->header->orch_error_code.load(std::memory_order_acquire), SIMPLER_ERROR_INVALID_ARGS); +} diff --git a/tests/ut/cpp/a5/test_hbg_scheduler_ready.cpp b/tests/ut/cpp/a5/test_hbg_scheduler_ready.cpp new file mode 100644 index 0000000000..2959b8925f --- /dev/null +++ b/tests/ut/cpp/a5/test_hbg_scheduler_ready.cpp @@ -0,0 +1,1002 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "scheduler/scheduler_ready.h" +#include "runtime_types.h" + +namespace { + +class SchedulerStateBuffer { +public: + explicit SchedulerStateBuffer(const AicoreSchedulerLayout &layout) : + base_(std::aligned_alloc(SCHEDULER_STATE_ALIGNMENT, layout.total_size)) { + EXPECT_NE(base_, nullptr); + if (base_ != nullptr) EXPECT_TRUE(scheduler_init_data_from_layout(base_, layout)); + } + ~SchedulerStateBuffer() { std::free(base_); } + void *base() const { return base_; } + +private: + void *base_{nullptr}; +}; + +class GraphBuffer { +public: + explicit GraphBuffer(size_t task_count) : + task_count_(task_count), + image_(std::make_unique()) { + while (capacity_ < std::max(task_count, 1)) + capacity_ <<= 1; + if (capacity_ > kMaxTaskCount) throw std::invalid_argument("test graph exceeds GraphBuffer capacity"); + descriptors_ = image_->descriptors.data(); + payloads_ = image_->payloads.data(); + fanins_ = image_->fanins.data(); + for (size_t task = 0; task < capacity_; ++task) { + descriptors_[task].task_id = TaskId{static_cast(task)}; + payloads_[task].bind_regions( + nullptr, nullptr, fanins_ + task * static_cast(SCHEDULER_GRAPH_MAX_FANIN) + ); + if (payloads_[task].fanin_data() == nullptr) { + throw std::logic_error("test graph fanin region must share its contiguous image"); + } + for (int slot = 0; slot < 3; ++slot) + descriptors_[task].kernel_id[slot] = INVALID_KERNEL_ID; + } + } + + void executable(size_t task, uint8_t subtask_slot, std::vector fanins = {}) { + ASSERT_LT(task, task_count_); + ASSERT_LT(subtask_slot, 3); + ASSERT_LE(fanins.size(), static_cast(SCHEDULER_GRAPH_MAX_FANIN)); + descriptors_[task].kernel_id[subtask_slot] = 1; + payloads_[task].fanin_count = static_cast(fanins.size()); + ASSERT_TRUE(fanins.empty() || payloads_[task].fanin_data() != nullptr); + std::copy(fanins.begin(), fanins.end(), payloads_[task].fanin_data()); + } + + void mixed(size_t task, uint8_t active_mask) { + ASSERT_LT(task, task_count_); + for (uint8_t subtask_slot = 0; subtask_slot < 3; ++subtask_slot) { + if ((active_mask & (1U << subtask_slot)) != 0) descriptors_[task].kernel_id[subtask_slot] = 1; + } + payloads_[task].fanin_count = 0; + } + + void predicate(size_t task, uint64_t addr, uint8_t elem_size, uint8_t op, int64_t target = 0) { + ASSERT_LT(task, task_count_); + payloads_[task].predicate.addr = addr; + payloads_[task].predicate.target = target; + payloads_[task].predicate.elem_size = elem_size; + payloads_[task].predicate.op = static_cast(op); + } + + SchedulerGraphView graph() const { + return { + reinterpret_cast(descriptors_), + reinterpret_cast(payloads_), + task_count_, + capacity_ - 1, + }; + } + +private: + static constexpr size_t kMaxTaskCount = 8192; + struct alignas(64) GraphImage { + std::array descriptors{}; + std::array payloads{}; + std::array fanins{}; + }; + + size_t task_count_; + size_t capacity_{1}; + std::unique_ptr image_; + TaskDescriptor *descriptors_{nullptr}; + TaskPayload *payloads_{nullptr}; + int32_t *fanins_{nullptr}; +}; + +struct FixtureStorage { + explicit FixtureStorage(uint64_t task_count, uint64_t workers = 2) { + EXPECT_TRUE(scheduler_plan_layout(task_count, task_count, 0, &layout)); + scheduler_state = std::make_unique(layout); + run_control = scheduler_state_at(scheduler_state->base(), layout.run_control_offset); + contexts = scheduler_state_at(scheduler_state->base(), layout.worker_contexts_offset); + run_control->aiv_active_worker_count = workers; + run_control->resolver_count = workers; + for (uint64_t worker = 0; worker < workers; ++worker) { + SchedulerWorkerContext &context = contexts[worker]; + context.core_type = static_cast(CoreType::AIV); + context.active = 1; + context.task_controls_offset = layout.task_controls_offset; + context.task_metadata_offset = layout.task_metadata_offset; + context.completion_inboxes_offset = layout.completion_inboxes_offset; + context.ready_inboxes_offset = layout.ready_inboxes_offset; + context.ready_owner_states_offset = layout.ready_owner_states_offset; + context.ready_directory_offset = layout.ready_directory_offset; + context.trace_cells_offset = layout.trace_cells_offset; + context.worker_contexts_offset = layout.worker_contexts_offset; + context.dispatch_slots_offset = layout.dispatch_slots_offset; + context.callable_addresses_offset = layout.callable_addresses_offset; + context.gang_coordinator_offset = layout.gang_coordinator_offset; + context.gang_cohorts_offset = layout.gang_cohorts_offset; + context.gang_participants_offset = layout.gang_participants_offset; + context.gang_commands_offset = layout.gang_commands_offset; + context.dispatch_payload_offset = + layout.dispatch_payloads_offset + worker * SCHEDULER_PENDING_SLOT_COUNT * sizeof(DispatchPayload); + context.graph_task_count = task_count; + context.runtime_worker_count = workers; + context.worker_index = worker; + context.inbox_index = worker; + } + metadata = scheduler_state_at(scheduler_state->base(), layout.task_metadata_offset); + callable_addresses = scheduler_state_at(scheduler_state->base(), layout.callable_addresses_offset); + callable_addresses[1] = UINT64_C(0x1000); + for (uint64_t task = 0; task < task_count; ++task) { + metadata[task].kernel_ids[0] = 1; + metadata[task].kernel_ids[1] = UINT16_MAX; + metadata[task].kernel_ids[2] = UINT16_MAX; + metadata[task].active_mask = 1; + metadata[task].logical_block_num = 1; + metadata[task].total_required_subtasks = 1; + metadata[task].flags = SCHEDULER_TASK_EXECUTABLE; + } + } + + AicoreSchedulerLayout layout{}; + std::unique_ptr scheduler_state; + SchedulerRunControl *run_control{nullptr}; + SchedulerWorkerContext *contexts{nullptr}; + SchedulerTaskMetadata *metadata{nullptr}; + uint64_t *callable_addresses{nullptr}; +}; + +TEST(SchedulerBootstrap, RegistersOnlyOnFirstExecutableProducer) { + FixtureStorage storage(4, 2); + GraphBuffer graph(4); + graph.executable(0, 0); + graph.executable(1, 1, {0}); + graph.executable(3, 1, {2, 1}); + storage.metadata[1].active_mask = 2; + storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + storage.metadata[2].flags = 0; + storage.metadata[3].active_mask = 2; + storage.metadata[3].flags |= SCHEDULER_TASK_HAS_FANIN; + auto *controls = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.task_controls_offset); + controls[2].state = static_cast(SchedulerTaskState::DONE); + controls[2].wake_list_head = SCHEDULER_WAKE_LIST_CLOSED; + + SchedulerWakeStats stats{}; + EXPECT_EQ( + scheduler_bootstrap_route_task( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 1, &stats + ), + SchedulerRouteResult::WAITING + ); + EXPECT_EQ(controls[0].wake_list_head, 1); + EXPECT_EQ(controls[1].next_waiter, SCHEDULER_WAKE_LIST_OPEN); + EXPECT_EQ(controls[1].waiting_producer, 0); + + EXPECT_EQ( + scheduler_bootstrap_route_task( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[1], storage.run_control, 3, &stats + ), + SchedulerRouteResult::WAITING + ); + EXPECT_EQ(controls[1].wake_list_head, 3); + EXPECT_EQ(controls[3].next_fanin_index, 1); + EXPECT_EQ(controls[3].waiting_producer, 1); + EXPECT_EQ(stats.wake_register_count, 2u); + EXPECT_EQ(stats.fanin_state_load_count, 0u); + EXPECT_EQ(stats.wake_cas_retry_count, 0u); +} + +TEST(SchedulerBootstrap, PublishesExclusiveInboxAndAggregatesDirectory) { + FixtureStorage storage(2, 2); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0); + storage.contexts[0].inbox_index = 1; + storage.contexts[1].inbox_index = 0; + SchedulerReadyBatch batch{}; + SchedulerReadyStats stats{}; + ASSERT_TRUE( + scheduler_bootstrap_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[1], 0, &batch, &stats) + ); + ASSERT_TRUE( + scheduler_bootstrap_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[1], 1, &batch, &stats) + ); + scheduler_cache_barrier(); + uint64_t ready_types = 0; + ASSERT_TRUE(scheduler_bootstrap_ready_batch_publish( + storage.scheduler_state->base(), &storage.contexts[1], 0, 0, &batch, &stats, &ready_types + )); + auto *directory = scheduler_ready_directory_at(storage.scheduler_state->base(), &storage.contexts[1]); + directory->bootstrap_ready_types[0] = ready_types; + scheduler_bootstrap_ready_directory_publish(storage.scheduler_state->base(), &storage.contexts[1], 2); + + auto *controls = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.task_controls_offset); + EXPECT_EQ(controls[0].next_waiter, 1); + EXPECT_EQ(controls[1].next_waiter, SCHEDULER_INBOX_EMPTY); + EXPECT_EQ(scheduler_ready_inbox_at(storage.scheduler_state->base(), &storage.contexts[1], 0, 0)->head, 0); + EXPECT_EQ(directory->core_types[0][0].bits, 1u); + EXPECT_EQ(directory->core_types[1][0].bits, 0u); + EXPECT_EQ(stats.enqueue_count, 2u); + EXPECT_EQ(stats.batch_count, 1u); +} + +TEST(SchedulerReadyInbox, BatchPushAndOwnerMaintenancePreserveFifoAndDirectory) { + constexpr uint64_t kTasks = 4; + FixtureStorage storage(kTasks, 1); + GraphBuffer graph(kTasks); + for (uint64_t task = 0; task < kTasks; ++task) + graph.executable(task, 0); + SchedulerReadyBatch batch{}; + SchedulerReadyOwnerState owner_state{}; + SchedulerReadyStats stats{}; + for (uint64_t task = 0; task < kTasks; ++task) + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], task, &batch, &stats) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &batch, &stats, &owner_state + )); + + auto *directory = scheduler_state_at( + storage.scheduler_state->base(), storage.layout.ready_directory_offset + ); + EXPECT_NE(directory->core_types[0][0].bits & 1, 0u); + for (uint64_t index = 0; index < kTasks; ++index) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, + &stats + )); + EXPECT_EQ(task, static_cast(index)); + } + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &stats + )); + EXPECT_EQ(task, SCHEDULER_TASK_ID_INVALID); + EXPECT_NE(directory->core_types[0][0].bits & 1, 0u); + ASSERT_TRUE( + scheduler_ready_owner_maintain_type(storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state) + ); + EXPECT_EQ(directory->core_types[0][0].bits & 1, 0u); + EXPECT_EQ(stats.pop_count, kTasks); +} + +TEST(SchedulerReadyInbox, OwnerStateInitializationRestoresEmptySentinels) { + SchedulerReadyOwnerState owner_state; + __builtin_memset(&owner_state, 0, sizeof(owner_state)); + + scheduler_ready_owner_init(&owner_state); + + for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + EXPECT_EQ(owner_state.queues[type].pending_endpoints, SCHEDULER_READY_PENDING_EMPTY); + EXPECT_EQ(owner_state.queues[type].advertised, 0u); + } +} + +TEST(SchedulerReadyInbox, PackedOwnerEndpointsRoundTripAsOneWord) { + EXPECT_EQ(scheduler_ready_pending_pack(SCHEDULER_INBOX_EMPTY, SCHEDULER_INBOX_EMPTY), UINT64_MAX); + const uint64_t endpoints = scheduler_ready_pending_pack(INT32_MAX - 1, INT32_MAX); + EXPECT_EQ(scheduler_ready_pending_head(endpoints), INT32_MAX - 1); + EXPECT_EQ(scheduler_ready_pending_tail(endpoints), INT32_MAX); +} + +TEST(SchedulerReadyInbox, OwnerPromotesPendingBankAfterPublishedBankDrains) { + constexpr uint64_t kTasks = 4; + FixtureStorage storage(kTasks, 1); + GraphBuffer graph(kTasks); + for (uint64_t task = 0; task < kTasks; ++task) + graph.executable(task, 0); + SchedulerReadyOwnerState owner_state{}; + SchedulerReadyStats stats{}; + SchedulerReadyBatch published{}; + SchedulerReadyBatch pending{}; + for (int64_t task = 0; task < 2; ++task) + ASSERT_TRUE(scheduler_ready_batch_append( + storage.scheduler_state->base(), &storage.contexts[0], task, &published, &stats + )); + for (int64_t task = 2; task < 4; ++task) + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], task, &pending, &stats) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &published, &stats, &owner_state + )); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &pending, &stats, &owner_state + )); + const uint64_t endpoints = owner_state.queues[0].pending_endpoints; + EXPECT_EQ(scheduler_ready_pending_head(endpoints), 2); + EXPECT_EQ(scheduler_ready_pending_tail(endpoints), 3); + + for (int64_t expected = 0; expected < 2; ++expected) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, + &stats + )); + EXPECT_EQ(task, expected); + } + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &stats + )); + EXPECT_EQ(task, SCHEDULER_TASK_ID_INVALID); + ASSERT_TRUE( + scheduler_ready_owner_maintain_type(storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state) + ); + for (int64_t expected = 2; expected < 4; ++expected) { + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, + &stats + )); + EXPECT_EQ(task, expected); + } +} + +TEST(SchedulerReadyInbox, OlderPendingBankPrecedesBatchArrivingAfterDrain) { + constexpr uint64_t kTasks = 3; + FixtureStorage storage(kTasks, 1); + GraphBuffer graph(kTasks); + for (uint64_t task = 0; task < kTasks; ++task) + graph.executable(task, 0); + SchedulerReadyOwnerState owner_state{}; + SchedulerReadyStats stats{}; + for (int64_t task = 0; task < 2; ++task) { + SchedulerReadyBatch batch{}; + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], task, &batch, &stats) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &batch, &stats, &owner_state + )); + } + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &stats + )); + ASSERT_EQ(task, 0); + + SchedulerReadyBatch arriving{}; + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], 2, &arriving, &stats) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &arriving, &stats, &owner_state + )); + EXPECT_EQ(scheduler_ready_pending_head(owner_state.queues[0].pending_endpoints), 2); + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &stats + )); + EXPECT_EQ(task, 1); + ASSERT_TRUE( + scheduler_ready_owner_maintain_type(storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state) + ); + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &stats + )); + EXPECT_EQ(task, 2); +} + +TEST(SchedulerReadyInbox, ThiefCannotObserveOrPromoteOwnerPendingBank) { + FixtureStorage storage(2, 2); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0); + SchedulerReadyOwnerState owner_state{}; + SchedulerReadyStats stats{}; + for (int64_t task = 0; task < 2; ++task) { + SchedulerReadyBatch batch{}; + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[1], task, &batch, &stats) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[1], 0, 1, &batch, &stats, &owner_state + )); + } + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 1, &task, &stats + )); + ASSERT_EQ(task, 0); + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 1, &task, &stats + )); + EXPECT_EQ(task, SCHEDULER_TASK_ID_INVALID); + EXPECT_EQ(scheduler_ready_pending_head(owner_state.queues[0].pending_endpoints), 1); + ASSERT_TRUE( + scheduler_ready_owner_maintain_type(storage.scheduler_state->base(), &storage.contexts[1], 0, &owner_state) + ); + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 1, &task, &stats + )); + EXPECT_EQ(task, 1); +} + +TEST(SchedulerReadyInbox, StealsOnlyFromMarkedVictim) { + FixtureStorage storage(1, 2); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerReadyBatch batch{}; + SchedulerReadyStats stats{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[1], 0, &batch, &stats)); + ASSERT_TRUE( + scheduler_ready_batch_push(storage.scheduler_state->base(), &storage.contexts[1], 0, 1, &batch, &stats) + ); + + uint64_t cursor = 1; + SchedulerReadyClaim claim{}; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 2, 0, &cursor, + &stats, &claim + )); + EXPECT_EQ(claim.task_id, 0); + EXPECT_EQ(claim.inbox_index, 1u); + EXPECT_EQ(claim.source, SchedulerReadySource::STOLEN); + EXPECT_EQ(stats.steal_count, 1u); +} + +TEST(SchedulerReadyInbox, DirectoryShardIgnoresResolverTail) { + FixtureStorage storage(1, 9); + auto *directory = scheduler_ready_directory_at(storage.scheduler_state->base(), &storage.contexts[0]); + directory->core_types[0][1].bits = UINT64_C(1) << 6; + EXPECT_EQ(scheduler_load_ready_directory_shard(directory, 9, 0, 7), 0u); + + directory->core_types[0][1].bits = UINT64_C(1) << 1; + EXPECT_EQ(scheduler_load_ready_directory_shard(directory, 9, 0, 7), UINT64_C(1) << 1); +} + +TEST(SchedulerReadyInbox, BootstrapPublishesIndependentDirectoryShards) { + FixtureStorage storage(1, 14); + auto *directory = scheduler_ready_directory_at(storage.scheduler_state->base(), &storage.contexts[0]); + directory->bootstrap_ready_types[0] = UINT64_C(1) << 0; + directory->bootstrap_ready_types[6] = UINT64_C(1) << 0; + directory->bootstrap_ready_types[7] = UINT64_C(1) << 1; + directory->bootstrap_ready_types[13] = (UINT64_C(1) << 0) | (UINT64_C(1) << 1); + + scheduler_bootstrap_ready_directory_publish(storage.scheduler_state->base(), &storage.contexts[0], 14); + + EXPECT_EQ(directory->core_types[0][0].bits, (UINT64_C(1) << 0) | (UINT64_C(1) << 6)); + EXPECT_EQ(directory->core_types[1][0].bits, 0u); + EXPECT_EQ(directory->core_types[0][1].bits, UINT64_C(1) << 6); + EXPECT_EQ(directory->core_types[1][1].bits, (UINT64_C(1) << 0) | (UINT64_C(1) << 6)); +} + +TEST(SchedulerReadyInbox, SparseDirectoryWrapsWithinShard) { + FixtureStorage storage(2, 14); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0); + SchedulerReadyStats stats{}; + SchedulerReadyBatch high_batch{}; + SchedulerReadyBatch low_batch{}; + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[13], 0, &high_batch, &stats) + ); + ASSERT_TRUE( + scheduler_ready_batch_push(storage.scheduler_state->base(), &storage.contexts[13], 0, 13, &high_batch, &stats) + ); + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[8], 1, &low_batch, &stats) + ); + ASSERT_TRUE( + scheduler_ready_batch_push(storage.scheduler_state->base(), &storage.contexts[8], 0, 8, &low_batch, &stats) + ); + + uint64_t cursor = 12; + SchedulerReadyClaim claim{}; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[7], storage.run_control, 14, 0, &cursor, + &stats, &claim + )); + EXPECT_EQ(claim.task_id, 0); + EXPECT_EQ(claim.inbox_index, 13u); + EXPECT_EQ(cursor, 7u); + + ASSERT_TRUE(scheduler_claim_ready_for_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[7], storage.run_control, 14, 0, &cursor, + &stats, &claim + )); + EXPECT_EQ(claim.task_id, 1); + EXPECT_EQ(claim.inbox_index, 8u); + EXPECT_EQ(cursor, 9u); +} + +TEST(SchedulerReadyInbox, DoesNotStealAcrossDirectoryShards) { + FixtureStorage storage(1, 14); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerReadyBatch batch{}; + SchedulerReadyStats stats{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[8], 0, &batch, &stats)); + ASSERT_TRUE( + scheduler_ready_batch_push(storage.scheduler_state->base(), &storage.contexts[8], 0, 8, &batch, &stats) + ); + + uint64_t cursor = 1; + SchedulerReadyClaim claim{}; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 14, 0, &cursor, + &stats, &claim + )); + EXPECT_EQ(claim.task_id, SCHEDULER_TASK_ID_INVALID); + + cursor = 8; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[7], storage.run_control, 14, 0, &cursor, + &stats, &claim + )); + EXPECT_EQ(claim.task_id, 0); + EXPECT_EQ(claim.inbox_index, 8u); + EXPECT_EQ(claim.source, SchedulerReadySource::STOLEN); +} + +TEST(SchedulerDispatch, RejectsKernelIdBeforeCallableTableAccess) { + FixtureStorage storage(1, 1); + GraphBuffer graph(1); + graph.executable(0, 0); + storage.contexts[0].core_type = static_cast(CoreType::AIC); + storage.metadata[0].kernel_ids[0] = static_cast(SCHEDULER_CALLABLE_CAPACITY); + SchedulerFreeSlotClaim slot_claim{0, 0, 0}; + SchedulerReadyClaim ready_claim{}; + ready_claim.task_id = 0; + + EXPECT_FALSE(scheduler_fill_dispatch_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, slot_claim, + ready_claim + )); + EXPECT_EQ(storage.run_control->scheduler_error, static_cast(SchedulerGraphResult::INVALID_CALLABLE)); + EXPECT_EQ(storage.run_control->error_site, static_cast(SchedulerErrorSite::DISPATCH_INVALID_CALLABLE)); +} + +TEST(SchedulerDispatch, AcceptsLastCallableAndInlineSentinel) { + FixtureStorage storage(2, 1); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0); + storage.contexts[0].core_type = static_cast(CoreType::AIC); + storage.metadata[0].kernel_ids[0] = static_cast(SCHEDULER_CALLABLE_CAPACITY - 1); + storage.callable_addresses[SCHEDULER_CALLABLE_CAPACITY - 1] = UINT64_C(0x2000); + SchedulerFreeSlotClaim slot_claim{0, 0, 0}; + SchedulerReadyClaim ready_claim{}; + ready_claim.task_id = 0; + ASSERT_TRUE(scheduler_fill_dispatch_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, slot_claim, + ready_claim + )); + auto *payload = scheduler_state_at( + storage.scheduler_state->base(), storage.contexts[0].dispatch_payload_offset + ); + EXPECT_EQ(payload->function_bin_addr, UINT64_C(0x2000)); + + storage.metadata[1].kernel_ids[0] = UINT16_MAX; + storage.metadata[1].flags |= SCHEDULER_TASK_INLINE; + slot_claim.slot_index = 1; + ready_claim.task_id = 1; + ASSERT_TRUE(scheduler_fill_dispatch_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, slot_claim, + ready_claim + )); + payload = scheduler_state_at( + storage.scheduler_state->base(), storage.contexts[0].dispatch_payload_offset + sizeof(DispatchPayload) + ); + EXPECT_EQ(payload->function_bin_addr, 0u); + EXPECT_EQ(storage.run_control->scheduler_error, 0u); +} + +TEST(SchedulerPredicate, DistinguishesFailedAndMalformedPredicates) { + GraphBuffer graph(1); + graph.executable(0, 0); + alignas(8) int64_t value = 3; + graph.predicate(0, reinterpret_cast(&value), 8, static_cast(PredicateOp::GT), 4); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::FAIL); + graph.predicate(0, reinterpret_cast(&value), 8, static_cast(PredicateOp::LE), 4); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::PASS); + + graph.predicate(0, 0, 8, static_cast(PredicateOp::GT)); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::MALFORMED); + graph.predicate(0, reinterpret_cast(&value), 3, static_cast(PredicateOp::GT)); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::MALFORMED); + graph.predicate(0, reinterpret_cast(&value), 8, UINT8_C(0xff)); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::MALFORMED); + graph.predicate(0, reinterpret_cast(&value) + 1, 8, static_cast(PredicateOp::GT)); + EXPECT_EQ(scheduler_evaluate_task_predicate(graph.graph(), 0), SchedulerPredicateResult::MALFORMED); +} + +TEST(SchedulerPredicate, MalformedPredicateStopsDispatchWithoutPublishingSlot) { + FixtureStorage storage(1, 1); + GraphBuffer graph(1); + graph.executable(0, 0); + graph.predicate(0, 0, 4, static_cast(PredicateOp::GT)); + storage.contexts[0].core_type = static_cast(CoreType::AIC); + storage.metadata[0].flags |= SCHEDULER_TASK_HAS_PREDICATE; + SchedulerFreeSlotClaim slot_claim{0, 0, 0}; + SchedulerReadyClaim ready_claim{}; + ready_claim.task_id = 0; + + EXPECT_FALSE(scheduler_fill_dispatch_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, slot_claim, + ready_claim + )); + EXPECT_EQ(storage.run_control->scheduler_error, static_cast(SchedulerGraphResult::INVALID_ARGUMENTS)); + EXPECT_EQ(storage.run_control->error_site, static_cast(SchedulerErrorSite::DISPATCH_INVALID_PREDICATE)); + const auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &storage.contexts[0], 0, 0); + EXPECT_NE(scheduler_dispatch_state(slot->publication), SchedulerDispatchSlotState::READY); +} + +TEST(SchedulerPredicate, FailedPredicatePublishesDependencyOnlyDispatch) { + FixtureStorage storage(1, 1); + GraphBuffer graph(1); + graph.executable(0, 0); + alignas(4) int32_t value = 0; + graph.predicate(0, reinterpret_cast(&value), 4, static_cast(PredicateOp::GT)); + storage.contexts[0].core_type = static_cast(CoreType::AIC); + storage.metadata[0].flags |= SCHEDULER_TASK_HAS_PREDICATE; + SchedulerFreeSlotClaim slot_claim{0, 0, 0}; + SchedulerReadyClaim ready_claim{}; + ready_claim.task_id = 0; + + ASSERT_TRUE(scheduler_fill_dispatch_slot( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, slot_claim, + ready_claim + )); + const auto *payload = scheduler_state_at( + storage.scheduler_state->base(), storage.contexts[0].dispatch_payload_offset + ); + EXPECT_EQ(payload->function_bin_addr, 0u); + EXPECT_EQ(storage.run_control->scheduler_error, 0u); +} + +TEST(SchedulerReadyInbox, ConcurrentConsumersNeverDuplicateTask) { + constexpr uint64_t kTasks = 8192; + constexpr uint64_t kConsumerCount = 8; + FixtureStorage storage(kTasks, kConsumerCount); + GraphBuffer graph(kTasks); + SchedulerReadyBatch batch{}; + for (uint64_t task = 0; task < kTasks; ++task) { + graph.executable(task, 0); + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], task, &batch, nullptr) + ); + } + ASSERT_TRUE( + scheduler_ready_batch_push(storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &batch, nullptr) + ); + std::vector> seen(kTasks); + std::atomic claimed{0}; + std::atomic failed{false}; + auto consume = [&](uint64_t worker) { + while (!failed.load(std::memory_order_relaxed) && claimed.load(std::memory_order_relaxed) < kTasks) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + if (!scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[worker], storage.run_control, 0, + 0, &task, nullptr + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + if (task >= 0) { + seen[static_cast(task)].fetch_add(1, std::memory_order_relaxed); + claimed.fetch_add(1, std::memory_order_relaxed); + } + std::this_thread::yield(); + } + }; + std::vector consumers; + consumers.reserve(kConsumerCount); + for (uint64_t worker = 0; worker < kConsumerCount; ++worker) + consumers.emplace_back(consume, worker); + for (auto &consumer : consumers) + consumer.join(); + EXPECT_FALSE(failed.load()); + EXPECT_EQ(claimed.load(), kTasks); + for (const auto &count : seen) + EXPECT_EQ(count.load(), 1u); +} + +TEST(SchedulerReadyInbox, OwnerPushRacesThiefWithoutLosingTasks) { + constexpr uint64_t kTasks = 2048; + constexpr uint64_t kSpinLimit = 10000000; + FixtureStorage storage(kTasks, 2); + GraphBuffer graph(kTasks); + for (uint64_t task = 0; task < kTasks; ++task) + graph.executable(task, 0); + SchedulerReadyOwnerState owner_state{}; + scheduler_ready_owner_init(&owner_state); + std::vector> seen(kTasks); + std::atomic claimed{0}; + std::atomic failed{false}; + + std::thread owner([&] { + for (uint64_t task = 0; task < kTasks && !failed.load(std::memory_order_relaxed); ++task) { + SchedulerReadyBatch batch{}; + if (!scheduler_ready_batch_append( + storage.scheduler_state->base(), &storage.contexts[0], task, &batch, nullptr + ) || + !scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &batch, nullptr, &owner_state + ) || + !scheduler_ready_owner_maintain_type( + storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + std::this_thread::yield(); + } + for (uint64_t spin = 0; spin < kSpinLimit && claimed.load(std::memory_order_relaxed) < kTasks; ++spin) { + if (!scheduler_ready_owner_maintain_type( + storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + std::this_thread::yield(); + } + if (claimed.load(std::memory_order_relaxed) != kTasks) failed.store(true, std::memory_order_relaxed); + }); + std::thread thief([&] { + for (uint64_t spin = 0; spin < kSpinLimit && !failed.load(std::memory_order_relaxed) && + claimed.load(std::memory_order_relaxed) < kTasks; + ++spin) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + if (!scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[1], storage.run_control, 0, 0, + &task, nullptr + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + if (task >= 0) { + seen[static_cast(task)].fetch_add(1, std::memory_order_relaxed); + claimed.fetch_add(1, std::memory_order_relaxed); + } + std::this_thread::yield(); + } + }); + owner.join(); + thief.join(); + + EXPECT_FALSE(failed.load()); + EXPECT_EQ(claimed.load(), kTasks); + for (const auto &count : seen) + EXPECT_EQ(count.load(), 1u); +} + +TEST(SchedulerReadyInbox, PendingPromotionRacesThiefWithoutReplayingTasks) { + constexpr uint64_t kTasks = 2048; + constexpr uint64_t kSpinLimit = 10000000; + FixtureStorage storage(kTasks, 2); + GraphBuffer graph(kTasks); + for (uint64_t task = 0; task < kTasks; ++task) + graph.executable(task, 0); + SchedulerReadyOwnerState owner_state{}; + scheduler_ready_owner_init(&owner_state); + SchedulerReadyBatch published{}; + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], 0, &published, nullptr) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &published, nullptr, &owner_state + )); + SchedulerReadyBatch pending{}; + for (uint64_t task = 1; task < kTasks; ++task) + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], task, &pending, nullptr) + ); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &pending, nullptr, &owner_state + )); + ASSERT_NE(scheduler_ready_pending_head(owner_state.queues[0].pending_endpoints), SCHEDULER_INBOX_EMPTY); + + std::vector> seen(kTasks); + std::atomic claimed{0}; + std::atomic failed{false}; + std::thread owner([&] { + for (uint64_t spin = 0; spin < kSpinLimit && !failed.load(std::memory_order_relaxed) && + claimed.load(std::memory_order_relaxed) < kTasks; + ++spin) { + if (!scheduler_ready_owner_maintain_type( + storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + std::this_thread::yield(); + } + }); + std::thread thief([&] { + for (uint64_t spin = 0; spin < kSpinLimit && !failed.load(std::memory_order_relaxed) && + claimed.load(std::memory_order_relaxed) < kTasks; + ++spin) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + if (!scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[1], storage.run_control, 0, 0, + &task, nullptr + )) { + failed.store(true, std::memory_order_relaxed); + return; + } + if (task >= 0) { + seen[static_cast(task)].fetch_add(1, std::memory_order_relaxed); + claimed.fetch_add(1, std::memory_order_relaxed); + } + std::this_thread::yield(); + } + }); + owner.join(); + thief.join(); + if (claimed.load(std::memory_order_relaxed) != kTasks) failed.store(true, std::memory_order_relaxed); + + EXPECT_FALSE(failed.load()); + EXPECT_EQ(claimed.load(), kTasks); + for (const auto &count : seen) + EXPECT_EQ(count.load(), 1u); +} + +TEST(SchedulerReadyWake, ConcurrentRegistrationAndCloseResolveEveryConsumerExactlyOnce) { + constexpr uint64_t kConsumers = 64; + FixtureStorage storage(kConsumers + 1, 1); + GraphBuffer graph(kConsumers + 1); + graph.executable(0, 0); + for (uint64_t task = 1; task <= kConsumers; ++task) { + graph.executable(task, 0, {0}); + storage.metadata[task].flags |= SCHEDULER_TASK_HAS_FANIN; + } + auto *controls = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.task_controls_offset); + std::vector> seen(kConsumers + 1); + std::atomic started{0}; + std::atomic failed{false}; + std::vector consumers; + consumers.reserve(kConsumers); + for (uint64_t task = 1; task <= kConsumers; ++task) { + consumers.emplace_back([&, task] { + started.fetch_add(1, std::memory_order_release); + const SchedulerRouteResult route = scheduler_route_task( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, + static_cast(task), nullptr + ); + if (route == SchedulerRouteResult::READY_TO_ENQUEUE) { + seen[task].fetch_add(1, std::memory_order_relaxed); + } else if (route != SchedulerRouteResult::WAITING) { + failed.store(true, std::memory_order_relaxed); + } + }); + } + std::thread closer([&] { + while (started.load(std::memory_order_acquire) != kConsumers) + std::this_thread::yield(); + scheduler_gm_store(controls[0].state, static_cast(SchedulerTaskState::DONE)); + if (!scheduler_resolve_completion( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, nullptr, + nullptr, nullptr + )) + failed.store(true, std::memory_order_relaxed); + }); + for (auto &consumer : consumers) + consumer.join(); + closer.join(); + + while (!failed.load(std::memory_order_relaxed)) { + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, + nullptr + )); + if (task == SCHEDULER_TASK_ID_INVALID) break; + seen[static_cast(task)].fetch_add(1, std::memory_order_relaxed); + } + EXPECT_FALSE(failed.load()); + EXPECT_EQ(storage.run_control->scheduler_error, 0u); + for (uint64_t task = 1; task <= kConsumers; ++task) { + EXPECT_EQ(seen[task].load(), 1u) << "consumer=" << task; + EXPECT_EQ(controls[task].next_fanin_index, 1) << "consumer=" << task; + } +} + +TEST(SchedulerReadyWake, WakeResolvePublishesConsumerToResolverLocalInbox) { + FixtureStorage storage(2, 1); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0, {0}); + storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + SchedulerWakeStats wake{}; + SchedulerReadyStats ready{}; + SchedulerCompletionStats completion{}; + EXPECT_EQ( + scheduler_route_task( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 1, &wake + ), + SchedulerRouteResult::WAITING + ); + auto *controls = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.task_controls_offset); + controls[0].state = static_cast(SchedulerTaskState::DONE); + ASSERT_TRUE(scheduler_resolve_completion( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, &wake, &ready, + &completion + )); + EXPECT_EQ(completion.resolve_count, 1u); + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &ready + )); + EXPECT_EQ(task, 1); + EXPECT_EQ(wake.wake_register_count, 1u); + EXPECT_EQ(wake.wake_migrate_count, 1u); +} + +TEST(SchedulerReadyWake, WakeResolveQueuesBehindOlderPublishedWork) { + FixtureStorage storage(3, 1); + GraphBuffer graph(3); + graph.executable(0, 0); + graph.executable(1, 0, {0}); + graph.executable(2, 0); + storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + SchedulerWakeStats wake{}; + SchedulerReadyStats ready{}; + SchedulerCompletionStats completion{}; + SchedulerReadyOwnerState owner_state{}; + EXPECT_EQ( + scheduler_route_task( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 1, &wake + ), + SchedulerRouteResult::WAITING + ); + SchedulerReadyBatch older{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), &storage.contexts[0], 2, &older, &ready)); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &storage.contexts[0], 0, 0, &older, &ready, &owner_state + )); + auto *controls = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.task_controls_offset); + controls[0].state = static_cast(SchedulerTaskState::DONE); + ASSERT_TRUE(scheduler_resolve_completion( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, &wake, &ready, + &completion, false, true, nullptr, &owner_state + )); + EXPECT_EQ(completion.resolve_count, 1u); + EXPECT_EQ(scheduler_ready_pending_head(owner_state.queues[0].pending_endpoints), 1); + + int64_t task = SCHEDULER_TASK_ID_INVALID; + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &ready + )); + ASSERT_EQ(task, 2); + ASSERT_TRUE( + scheduler_ready_owner_maintain_type(storage.scheduler_state->base(), &storage.contexts[0], 0, &owner_state) + ); + ASSERT_TRUE(scheduler_ready_pop_from_inbox( + graph.graph(), storage.scheduler_state->base(), &storage.contexts[0], storage.run_control, 0, 0, &task, &ready + )); + EXPECT_EQ(task, 1); +} + +} // namespace