From 9b1ce5ba279c389c2b02713e54f2c8008b90f76e Mon Sep 17 00:00:00 2001 From: zhusy54 Date: Thu, 27 Aug 2026 10:44:08 +0800 Subject: [PATCH] Add: complete A5 HBG single-lane scheduling - Service ordinary completion generations and refill core slots directly - Dispatch normal AIC/AIV work while deferring Resolver AIV reservations - Keep dispatch progress separate from fatal fill failure reporting - Validate topology and deferred reservations before scheduler-state access - Preserve completion tracing and name every device-side protocol error - Cover dispatch, completion, reservation, refill, and peer priority with portable C++ tests --- .../runtime/scheduler/scheduler_completion.h | 188 ++++ .../runtime/scheduler/scheduler_dispatch.h | 431 +++++++++ .../runtime/scheduler/scheduler_types.h | 6 + tests/ut/cpp/CMakeLists.txt | 1 + .../ut/cpp/a5/test_hbg_scheduler_dispatch.cpp | 896 ++++++++++++++++++ 5 files changed, 1522 insertions(+) create mode 100644 src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h create mode 100644 src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h create mode 100644 tests/ut/cpp/a5/test_hbg_scheduler_dispatch.cpp diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h new file mode 100644 index 0000000000..cd7fdb864f --- /dev/null +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h @@ -0,0 +1,188 @@ +/* + * 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 + +// AICore-resident completion library; scheduler_completion.cpp is the independent AICPU scheduler implementation. + +#include "scheduler_ready.h" + +inline __aicore__ bool scheduler_service_cluster_completion_slot( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, uint32_t cluster_lane, uint32_t pending_slot, + uint32_t completed_generation, SchedulerWakeStats *wake_stats, SchedulerReadyStats *ready_stats, + SchedulerCompletionStats *completion_stats, uint64_t *ready_victim_cursors, bool trace_enabled, + const SchedulerReadyClaim *replacement_ready, bool *direct_refilled, + SchedulerCompletionServiceTiming *timing = nullptr, __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + if (direct_refilled != nullptr) *direct_refilled = false; + if (cluster_lane >= PLATFORM_CORES_PER_BLOCKDIM || pending_slot >= SCHEDULER_PENDING_SLOT_COUNT || + completed_generation == 0) + return false; + const uint64_t worker_id = resolver->cluster_worker_ids[cluster_lane]; + if (worker_id >= resolver->runtime_worker_count) return false; + __gm__ SchedulerWorkerContext *target = scheduler_worker_context_at(scheduler_state_base, resolver, worker_id); + if (target->active == 0) return false; + __gm__ SchedulerCompletionInbox *completion_line = + scheduler_completion_inbox_at(scheduler_state_base, resolver, worker_id); + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, worker_id, pending_slot); + const uint64_t publication = scheduler_gm_query(slot->publication); + if (scheduler_dispatch_state(publication) != SchedulerDispatchSlotState::READY || + scheduler_dispatch_generation(publication) != completed_generation) { + scheduler_record_error( + run_control, slot->task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::COMPLETION_GENERATION_MISMATCH + ); + return false; + } + + const bool record_timeline = timing != nullptr; + uint64_t operation_start = record_timeline ? scheduler_cycles() : 0; + scheduler_observe_cache_line(slot); + const int64_t task_id = slot->task_id; + if (slot->gang != 0) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::UNSUPPORTED_SHAPE, &graph, resolver, + SchedulerErrorSite::COMPLETION_UNEXPECTED_GANG_SLOT + ); + return false; + } + const uint8_t completed_subtask_slot = slot->subtask_slot; + scheduler_gm_store(completion_line->completed_generations[pending_slot], UINT32_C(0)); + uint64_t operation_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) timing->consume_cycles += operation_end - operation_start; + operation_start = operation_end; + uint64_t ready_publish_cycles = 0; + uint64_t refill_cycles = 0; + uint64_t finalize_cycles = 0; + uint64_t refill_start_cycles = 0; + uint64_t refill_end_cycles = 0; + bool refilled = false; + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, resolver, task_id); + scheduler_gm_store(control->state, static_cast(SchedulerTaskState::DONE)); + if (!scheduler_resolve_completion( + graph, scheduler_state_base, resolver, run_control, task_id, wake_stats, ready_stats, completion_stats, + owner_state, trace_enabled, false, timing == nullptr ? nullptr : &ready_publish_cycles + )) + return false; + if (timing != nullptr) timing->ready_publish_cycles += ready_publish_cycles; + uint64_t resolved_count_start = timing == nullptr ? 0 : scheduler_cycles(); + scheduler_gm_fetch_add(run_control->resolved_task_count, UINT64_C(1)); + if (timing != nullptr) { + finalize_cycles = scheduler_cycles() - resolved_count_start; + timing->finalize_cycles += finalize_cycles; + } + refill_start_cycles = record_timeline ? scheduler_cycles() : 0; + SchedulerReadyClaim ready{}; + bool ready_available = replacement_ready != nullptr; + if (ready_available) { + ready = *replacement_ready; + } else if (ready_victim_cursors != nullptr && worker_id != resolver->worker_index) { + // A normal AIV task is never refilled directly onto the Resolver. + // Its completed slot becomes capacity for late binding instead. + const uint32_t core_type = scheduler_metadata_core_type_index(completed_subtask_slot); + if (!scheduler_claim_ready_for_slot( + graph, scheduler_state_base, resolver, run_control, resolver->resolver_count, core_type, + &ready_victim_cursors[core_type], ready_stats, &ready, owner_state, trace_enabled + )) + return false; + ready_available = ready.task_id >= 0; + } + if (ready_available) { + SchedulerFreeSlotClaim claim{worker_id, pending_slot, slot->generation}; + if (!scheduler_fill_dispatch_slot( + graph, scheduler_state_base, resolver, run_control, claim, ready, trace_enabled + )) + return false; + refilled = true; + } + refill_end_cycles = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) { + refill_cycles = refill_end_cycles - refill_start_cycles; + timing->refill_cycles += refill_cycles; + } + operation_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) { + uint64_t resolve_total = operation_end - operation_start; + uint64_t excluded = ready_publish_cycles + refill_cycles + finalize_cycles; + timing->resolve_cycles += resolve_total > excluded ? resolve_total - excluded : 0; + } + operation_start = operation_end; + if (!refilled) { + slot->task_id = SCHEDULER_TASK_ID_INVALID; + scheduler_writeback_cache_line(slot); + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::FREE) + ); + } + const uint64_t completion_end = record_timeline ? scheduler_cycles() : 0; + if (timing != nullptr) timing->finalize_cycles += completion_end - operation_start; + if (direct_refilled != nullptr) *direct_refilled = refilled; + return true; +} + +inline __aicore__ uint32_t scheduler_completion_catchup_mask(uint32_t initial_completion_mask) { + constexpr uint32_t all_pending_slots_mask = (1U << SCHEDULER_PENDING_SLOT_COUNT) - 1; + return initial_completion_mask == 0 ? 0 : all_pending_slots_mask & ~initial_completion_mask; +} + +inline __aicore__ bool scheduler_service_cluster_completions( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, SchedulerWakeStats *wake_stats, SchedulerReadyStats *ready_stats, + SchedulerCompletionStats *completion_stats, uint64_t *ready_victim_cursors = nullptr, bool trace_enabled = false, + uint64_t *direct_refilled_slot_mask = nullptr, SchedulerCompletionServiceTiming *timing = nullptr, + __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + if (resolver->is_resolver == 0) return false; + if (direct_refilled_slot_mask != nullptr) *direct_refilled_slot_mask = 0; + bool progress = false; + for (uint32_t cluster_lane = 0; cluster_lane < PLATFORM_CORES_PER_BLOCKDIM; ++cluster_lane) { + const uint64_t worker_id = resolver->cluster_worker_ids[cluster_lane]; + if (worker_id >= resolver->runtime_worker_count) continue; + __gm__ SchedulerWorkerContext *target = scheduler_worker_context_at(scheduler_state_base, resolver, worker_id); + if (target->active == 0) continue; + __gm__ SchedulerCompletionInbox *completion_line = + scheduler_completion_inbox_at(scheduler_state_base, resolver, worker_id); + uint64_t completed_generations = scheduler_gm_query_u32_pair(completion_line->completed_generations); + uint32_t initial_completion_mask = 0; + for (uint32_t pending_slot = 0; pending_slot < SCHEDULER_PENDING_SLOT_COUNT; ++pending_slot) { + if (static_cast(completed_generations >> (pending_slot * 32)) != 0) + initial_completion_mask |= 1U << pending_slot; + } + uint32_t scan_mask = initial_completion_mask; + // Completion processing can be much slower than the sibling kernel. Refresh the packed completion line once + // for slots that were incomplete in the initial scan, excluding newly refilled slots from the catch-up pass. + for (uint32_t scan_pass = 0; scan_pass < 2 && scan_mask != 0; ++scan_pass) { + for (uint32_t pending_slot = 0; pending_slot < SCHEDULER_PENDING_SLOT_COUNT; ++pending_slot) { + if ((scan_mask & (1U << pending_slot)) == 0) continue; + const uint32_t completed_generation = + static_cast(completed_generations >> (pending_slot * 32)); + if (completed_generation == 0) continue; + bool direct_refilled = false; + if (!scheduler_service_cluster_completion_slot( + graph, scheduler_state_base, resolver, run_control, cluster_lane, pending_slot, + completed_generation, wake_stats, ready_stats, completion_stats, ready_victim_cursors, + trace_enabled, nullptr, &direct_refilled, timing, owner_state + )) + return false; + if (direct_refilled && direct_refilled_slot_mask != nullptr) + *direct_refilled_slot_mask |= UINT64_C(1) + << (cluster_lane * SCHEDULER_PENDING_SLOT_COUNT + pending_slot); + progress = true; + } + scan_mask = scan_pass == 0 ? scheduler_completion_catchup_mask(initial_completion_mask) : 0; + if (scan_mask != 0) + completed_generations = scheduler_gm_query_u32_pair(completion_line->completed_generations); + } + } + return progress; +} diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h new file mode 100644 index 0000000000..9d7bbdcd51 --- /dev/null +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h @@ -0,0 +1,431 @@ +/* + * 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 + +// AICore-resident dispatch library; scheduler_dispatch.cpp is the independent AICPU scheduler implementation. + +#include "scheduler_completion.h" + +struct SchedulerDeferredAivDispatch { + SchedulerReadyClaim ready{}; + SchedulerFreeSlotClaim reserved_slot{}; +}; + +struct SchedulerDeferredAivQueue { + // Every entry owns one Resolver slot held in FILLING, so a peer miss can + // always fall back to local execution without another capacity decision. + SchedulerDeferredAivDispatch entries[SCHEDULER_PENDING_SLOT_COUNT]{}; + uint32_t count{0}; +}; + +inline __aicore__ void scheduler_deferred_aiv_pop_front(SchedulerDeferredAivQueue *queue) { + if (queue == nullptr || queue->count == 0) return; + for (uint32_t index = 1; index < queue->count; ++index) + queue->entries[index - 1] = queue->entries[index]; + --queue->count; + queue->entries[queue->count] = {}; +} + +struct SchedulerNormalDispatchTiming { + uint64_t probe_cycles[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t claim_cycles[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t prepare_cycles[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t materialize_cycles[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t publish_cycles[SCHEDULER_CORE_TYPE_COUNT]{}; +}; + +inline __aicore__ uint64_t +scheduler_normal_dispatch_detail_cycles(const SchedulerNormalDispatchTiming &timing, uint32_t core_type) { + return timing.claim_cycles[core_type] + timing.prepare_cycles[core_type] + timing.materialize_cycles[core_type] + + timing.publish_cycles[core_type]; +} + +inline __aicore__ void scheduler_finish_normal_dispatch_stage( + SchedulerNormalDispatchTiming *timing, uint32_t core_type, uint64_t stage_start, uint64_t detail_start +) { + if (timing == nullptr) return; + uint64_t stage_cycles = scheduler_cycles() - stage_start; + uint64_t detail_cycles = scheduler_normal_dispatch_detail_cycles(*timing, core_type) - detail_start; + timing->probe_cycles[core_type] += stage_cycles > detail_cycles ? stage_cycles - detail_cycles : 0; +} + +inline __aicore__ bool scheduler_normal_aiv_worker_precedes( + uint32_t candidate_occupied_slots, bool candidate_is_resolver, uint32_t selected_occupied_slots, + bool selected_is_resolver +) { + if (candidate_is_resolver != selected_is_resolver) return !candidate_is_resolver; + return candidate_occupied_slots < selected_occupied_slots; +} + +// The return value reports whether this pass made progress; failed independently reports an aborted pass. +inline __aicore__ bool scheduler_fill_cluster_normal_slots( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, uint64_t *ready_victim_cursors, SchedulerReadyStats *ready_stats, + bool trace_enabled, uint64_t skip_slot_mask = 0, SchedulerNormalDispatchTiming *timing = nullptr, + SchedulerDeferredAivQueue *deferred_aiv = nullptr, __gm__ SchedulerReadyOwnerState *owner_state = nullptr, + bool *failed = nullptr +) { + if (failed != nullptr) *failed = false; + if (resolver->is_resolver == 0) return false; + const uint32_t aic_core_type = static_cast(CoreType::AIC); + uint64_t stage_start = timing == nullptr ? 0 : scheduler_cycles(); + uint64_t detail_start = timing == nullptr ? 0 : scheduler_normal_dispatch_detail_cycles(*timing, aic_core_type); + bool progress = false; + + // AIC has no peer lane in its Cluster, so preserve the existing slot order. + if (scheduler_ready_directory_nonempty(scheduler_state_base, resolver, resolver->resolver_count, aic_core_type)) { + bool aic_ready_available = true; + for (uint32_t cluster_lane = 0; cluster_lane < PLATFORM_CORES_PER_BLOCKDIM && aic_ready_available; + ++cluster_lane) { + const uint64_t worker_id = resolver->cluster_worker_ids[cluster_lane]; + if (worker_id >= resolver->runtime_worker_count) continue; + __gm__ SchedulerWorkerContext *target = + scheduler_worker_context_at(scheduler_state_base, resolver, worker_id); + if (target->active == 0 || target->core_type != static_cast(CoreType::AIC)) continue; + for (uint32_t pending_slot = 0; pending_slot < SCHEDULER_PENDING_SLOT_COUNT; ++pending_slot) { + if ((skip_slot_mask & (UINT64_C(1) << (cluster_lane * SCHEDULER_PENDING_SLOT_COUNT + pending_slot))) != + 0) + continue; + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, worker_id, pending_slot); + const uint64_t publication = scheduler_gm_query(slot->publication); + if (scheduler_dispatch_state(publication) != SchedulerDispatchSlotState::FREE) continue; + SchedulerReadyClaim ready{}; + if (!scheduler_claim_ready_for_slot( + graph, scheduler_state_base, resolver, run_control, resolver->resolver_count, aic_core_type, + &ready_victim_cursors[aic_core_type], ready_stats, &ready, owner_state, trace_enabled + )) { + if (failed != nullptr) *failed = true; + return progress; + } + if (timing != nullptr && ready.claim_end_cycles >= ready.claim_start_cycles) + timing->claim_cycles[aic_core_type] += ready.claim_end_cycles - ready.claim_start_cycles; + if (ready.task_id < 0) { + aic_ready_available = false; + break; + } + SchedulerFreeSlotClaim claim{ + worker_id, + pending_slot, + scheduler_dispatch_generation(publication), + }; + scheduler_gm_store( + slot->publication, + scheduler_dispatch_publication(claim.generation, SchedulerDispatchSlotState::FILLING) + ); + SchedulerDispatchFillTiming fill_timing{}; + if (!scheduler_fill_dispatch_slot( + graph, scheduler_state_base, resolver, run_control, claim, ready, trace_enabled, + timing == nullptr ? nullptr : &fill_timing + )) { + if (failed != nullptr) *failed = true; + return progress; + } + if (timing != nullptr) { + timing->prepare_cycles[aic_core_type] += fill_timing.prepare_cycles; + timing->materialize_cycles[aic_core_type] += fill_timing.materialize_cycles; + timing->publish_cycles[aic_core_type] += fill_timing.publish_cycles; + } + progress = true; + } + } + } + scheduler_finish_normal_dispatch_stage(timing, aic_core_type, stage_start, detail_start); + + // A Resolver shares its AIV with Executor work. Exhaust the non-Resolver + // peer's free slots first, then claim more work only against reserved + // Resolver capacity. The caller decides the reserved work's owner after + // the rest of this scheduling round completes. + struct AivWorkerSlots { + uint64_t worker_id{UINT64_MAX}; + uint64_t publications[SCHEDULER_PENDING_SLOT_COUNT]{}; + uint32_t free_mask{0}; + uint32_t occupied_slots{0}; + bool is_resolver{false}; + }; + const uint32_t aiv_core_type = static_cast(CoreType::AIV); + stage_start = timing == nullptr ? 0 : scheduler_cycles(); + detail_start = timing == nullptr ? 0 : scheduler_normal_dispatch_detail_cycles(*timing, aiv_core_type); + if (scheduler_ready_directory_nonempty(scheduler_state_base, resolver, resolver->resolver_count, aiv_core_type)) { + AivWorkerSlots aiv_workers[PLATFORM_AIV_CORES_PER_BLOCKDIM]{}; + uint32_t aiv_worker_count = 0; + for (uint32_t cluster_lane = 0; cluster_lane < PLATFORM_CORES_PER_BLOCKDIM; ++cluster_lane) { + const uint64_t worker_id = resolver->cluster_worker_ids[cluster_lane]; + if (worker_id >= resolver->runtime_worker_count) continue; + __gm__ SchedulerWorkerContext *target = + scheduler_worker_context_at(scheduler_state_base, resolver, worker_id); + if (target->active == 0 || target->core_type != static_cast(CoreType::AIV)) continue; + if (aiv_worker_count >= PLATFORM_AIV_CORES_PER_BLOCKDIM) { + scheduler_record_error( + run_control, SCHEDULER_TASK_ID_INVALID, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::NORMAL_DISPATCH_INVALID_TOPOLOGY + ); + if (failed != nullptr) *failed = true; + return progress; + } + AivWorkerSlots &worker = aiv_workers[aiv_worker_count++]; + worker.worker_id = worker_id; + worker.is_resolver = worker_id == resolver->worker_index; + for (uint32_t pending_slot = 0; pending_slot < SCHEDULER_PENDING_SLOT_COUNT; ++pending_slot) { + if ((skip_slot_mask & (UINT64_C(1) << (cluster_lane * SCHEDULER_PENDING_SLOT_COUNT + pending_slot))) != + 0) { + ++worker.occupied_slots; + continue; + } + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, worker_id, pending_slot); + const uint64_t publication = scheduler_gm_query(slot->publication); + worker.publications[pending_slot] = publication; + if (scheduler_dispatch_state(publication) == SchedulerDispatchSlotState::FREE) + worker.free_mask |= 1U << pending_slot; + else ++worker.occupied_slots; + } + } + for (uint32_t attempt = 0; attempt < aiv_worker_count * SCHEDULER_PENDING_SLOT_COUNT; ++attempt) { + uint32_t selected = aiv_worker_count; + for (uint32_t worker_index = 0; worker_index < aiv_worker_count; ++worker_index) { + const AivWorkerSlots &candidate = aiv_workers[worker_index]; + if (candidate.free_mask == 0) continue; + if (selected == aiv_worker_count || + scheduler_normal_aiv_worker_precedes( + candidate.occupied_slots, candidate.is_resolver, aiv_workers[selected].occupied_slots, + aiv_workers[selected].is_resolver + )) + selected = worker_index; + } + if (selected == aiv_worker_count) break; + AivWorkerSlots &worker = aiv_workers[selected]; + const uint32_t pending_slot = static_cast(__builtin_ctz(worker.free_mask)); + worker.free_mask &= ~(1U << pending_slot); + if (worker.is_resolver && (deferred_aiv == nullptr || deferred_aiv->count >= SCHEDULER_PENDING_SLOT_COUNT)) + break; + const uint64_t publication = worker.publications[pending_slot]; + SchedulerFreeSlotClaim claim{ + worker.worker_id, + pending_slot, + scheduler_dispatch_generation(publication), + }; + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, worker.worker_id, pending_slot); + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(claim.generation, SchedulerDispatchSlotState::FILLING) + ); + SchedulerReadyClaim ready{}; + if (!scheduler_claim_ready_for_slot( + graph, scheduler_state_base, resolver, run_control, resolver->resolver_count, aiv_core_type, + &ready_victim_cursors[aiv_core_type], ready_stats, &ready, owner_state, trace_enabled + )) { + scheduler_gm_store( + slot->publication, + scheduler_dispatch_publication(claim.generation, SchedulerDispatchSlotState::FREE) + ); + if (failed != nullptr) *failed = true; + return progress; + } + if (timing != nullptr && ready.claim_end_cycles >= ready.claim_start_cycles) + timing->claim_cycles[aiv_core_type] += ready.claim_end_cycles - ready.claim_start_cycles; + if (ready.task_id < 0) { + scheduler_gm_store( + slot->publication, + scheduler_dispatch_publication(claim.generation, SchedulerDispatchSlotState::FREE) + ); + break; + } + if (worker.is_resolver) { + deferred_aiv->entries[deferred_aiv->count++] = {ready, claim}; + ++worker.occupied_slots; + progress = true; + continue; + } + SchedulerDispatchFillTiming fill_timing{}; + if (!scheduler_fill_dispatch_slot( + graph, scheduler_state_base, resolver, run_control, claim, ready, trace_enabled, + timing == nullptr ? nullptr : &fill_timing + )) { + if (failed != nullptr) *failed = true; + return progress; + } + if (timing != nullptr) { + timing->prepare_cycles[aiv_core_type] += fill_timing.prepare_cycles; + timing->materialize_cycles[aiv_core_type] += fill_timing.materialize_cycles; + timing->publish_cycles[aiv_core_type] += fill_timing.publish_cycles; + } + ++worker.occupied_slots; + progress = true; + } + } + scheduler_finish_normal_dispatch_stage(timing, aiv_core_type, stage_start, detail_start); + return progress; +} + +inline __aicore__ bool scheduler_release_deferred_aiv_reservation( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, const SchedulerFreeSlotClaim &reservation +) { + if (reservation.worker_id != resolver->worker_index || reservation.slot_index >= SCHEDULER_PENDING_SLOT_COUNT) { + scheduler_record_error( + run_control, SCHEDULER_TASK_ID_INVALID, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::DEFERRED_RESERVATION_INVALID_OWNER + ); + return false; + } + __gm__ SchedulerDispatchSlot *slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, reservation.worker_id, reservation.slot_index); + const uint64_t publication = scheduler_gm_query(slot->publication); + if (scheduler_dispatch_state(publication) != SchedulerDispatchSlotState::FILLING || + scheduler_dispatch_generation(publication) != reservation.generation || + slot->task_id != SCHEDULER_TASK_ID_INVALID) { + scheduler_record_error( + run_control, slot->task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::DEFERRED_RESERVATION_INVALID_STATE + ); + return false; + } + scheduler_gm_publish( + slot->publication, scheduler_dispatch_publication(reservation.generation, SchedulerDispatchSlotState::FREE) + ); + return true; +} + +inline __aicore__ int32_t +scheduler_deferred_aiv_peer_lane(__gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver) { + for (uint32_t cluster_lane = 0; cluster_lane < PLATFORM_CORES_PER_BLOCKDIM; ++cluster_lane) { + const uint64_t worker_id = resolver->cluster_worker_ids[cluster_lane]; + if (worker_id >= resolver->runtime_worker_count) continue; + if (worker_id == resolver->worker_index) continue; + __gm__ SchedulerWorkerContext *target = scheduler_worker_context_at(scheduler_state_base, resolver, worker_id); + scheduler_observe_cache_line(target); + if (target->active != 0 && target->core_type == static_cast(CoreType::AIV)) + return static_cast(cluster_lane); + } + return -1; +} + +inline __aicore__ bool scheduler_drain_deferred_aiv_to_peer( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, SchedulerDeferredAivQueue *queue, SchedulerWakeStats *wake_stats, + SchedulerReadyStats *ready_stats, SchedulerCompletionStats *completion_stats, bool trace_enabled, + SchedulerCompletionServiceTiming *completion_timing = nullptr, + SchedulerNormalDispatchTiming *dispatch_timing = nullptr, __gm__ SchedulerReadyOwnerState *owner_state = nullptr +) { + if (queue == nullptr || queue->count == 0) return true; + const int32_t peer_lane = scheduler_deferred_aiv_peer_lane(scheduler_state_base, resolver); + // A Resolver may be the only active AIV in its Cluster (for example, a + // single-root AIV graph). There is then nothing to drain to; leave the + // reservation queued so the caller can publish it on the Resolver itself. + if (peer_lane < 0) return true; + const uint64_t peer_worker_id = resolver->cluster_worker_ids[static_cast(peer_lane)]; + __gm__ SchedulerCompletionInbox *completion_line = + scheduler_completion_inbox_at(scheduler_state_base, resolver, peer_worker_id); + const uint32_t aiv_core_type = static_cast(CoreType::AIV); + + for (uint32_t pass = 0; pass < 2 && queue->count != 0; ++pass) { + const uint64_t completed_generations = + pass == 1 ? scheduler_gm_query_u32_pair(completion_line->completed_generations) : 0; + for (uint32_t pending_slot = 0; pending_slot < SCHEDULER_PENDING_SLOT_COUNT && queue->count != 0; + ++pending_slot) { + __gm__ SchedulerDispatchSlot *peer_slot = + scheduler_dispatch_slot_at(scheduler_state_base, resolver, peer_worker_id, pending_slot); + const uint64_t publication = scheduler_gm_query(peer_slot->publication); + const SchedulerDispatchSlotState state = scheduler_dispatch_state(publication); + const uint32_t generation = scheduler_dispatch_generation(publication); + bool refilled = false; + if (pass == 0) { + if (state != SchedulerDispatchSlotState::FREE) continue; + scheduler_gm_store( + peer_slot->publication, + scheduler_dispatch_publication(generation, SchedulerDispatchSlotState::FILLING) + ); + SchedulerDispatchFillTiming fill_timing{}; + if (!scheduler_fill_dispatch_slot( + graph, scheduler_state_base, resolver, run_control, + SchedulerFreeSlotClaim{peer_worker_id, pending_slot, generation}, queue->entries[0].ready, + trace_enabled, dispatch_timing == nullptr ? nullptr : &fill_timing + )) + return false; + if (dispatch_timing != nullptr) { + dispatch_timing->prepare_cycles[aiv_core_type] += fill_timing.prepare_cycles; + dispatch_timing->materialize_cycles[aiv_core_type] += fill_timing.materialize_cycles; + dispatch_timing->publish_cycles[aiv_core_type] += fill_timing.publish_cycles; + } + refilled = true; + } else { + if (state != SchedulerDispatchSlotState::READY) continue; + const uint32_t completed_generation = + static_cast(completed_generations >> (pending_slot * 32)); + if (completed_generation != generation) continue; + scheduler_observe_cache_line(peer_slot); + if (peer_slot->gang != 0) continue; + if (!scheduler_service_cluster_completion_slot( + graph, scheduler_state_base, resolver, run_control, static_cast(peer_lane), + pending_slot, completed_generation, wake_stats, ready_stats, completion_stats, nullptr, + trace_enabled, &queue->entries[0].ready, &refilled, completion_timing, owner_state + ) || + !refilled) + return false; + } + if (!scheduler_release_deferred_aiv_reservation( + graph, scheduler_state_base, resolver, run_control, queue->entries[0].reserved_slot + )) + return false; + scheduler_deferred_aiv_pop_front(queue); + } + } + return true; +} + +inline __aicore__ bool scheduler_publish_deferred_aiv_to_resolver( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, __gm__ SchedulerWorkerContext *resolver, + __gm__ SchedulerRunControl *run_control, SchedulerDeferredAivQueue *queue, bool trace_enabled, + uint32_t *published_slot, SchedulerNormalDispatchTiming *timing = nullptr +) { + if (published_slot != nullptr) *published_slot = UINT32_MAX; + if (queue == nullptr || queue->count == 0) return true; + const SchedulerDeferredAivDispatch &entry = queue->entries[0]; + if (entry.reserved_slot.worker_id != resolver->worker_index || + entry.reserved_slot.slot_index >= SCHEDULER_PENDING_SLOT_COUNT) { + scheduler_record_error( + run_control, entry.ready.task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::DEFERRED_PUBLISH_INVALID_RESERVATION + ); + return false; + } + __gm__ SchedulerDispatchSlot *slot = scheduler_dispatch_slot_at( + scheduler_state_base, resolver, entry.reserved_slot.worker_id, entry.reserved_slot.slot_index + ); + const uint64_t publication = scheduler_gm_query(slot->publication); + scheduler_observe_cache_line(slot); + if (scheduler_dispatch_state(publication) != SchedulerDispatchSlotState::FILLING || + scheduler_dispatch_generation(publication) != entry.reserved_slot.generation || + slot->task_id != SCHEDULER_TASK_ID_INVALID) { + scheduler_record_error( + run_control, entry.ready.task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, resolver, + SchedulerErrorSite::DEFERRED_PUBLISH_INVALID_RESERVATION + ); + return false; + } + const uint32_t aiv_core_type = static_cast(CoreType::AIV); + SchedulerDispatchFillTiming fill_timing{}; + if (!scheduler_fill_dispatch_slot( + graph, scheduler_state_base, resolver, run_control, entry.reserved_slot, entry.ready, trace_enabled, + timing == nullptr ? nullptr : &fill_timing + )) + return false; + if (timing != nullptr) { + timing->prepare_cycles[aiv_core_type] += fill_timing.prepare_cycles; + timing->materialize_cycles[aiv_core_type] += fill_timing.materialize_cycles; + timing->publish_cycles[aiv_core_type] += fill_timing.publish_cycles; + } + if (published_slot != nullptr) *published_slot = entry.reserved_slot.slot_index; + scheduler_deferred_aiv_pop_front(queue); + 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 33a7e30d8e..9461db9d89 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 @@ -914,6 +914,12 @@ enum class SchedulerErrorSite : uint64_t { COMPLETION_READY_APPEND_FAILED = 65, COMPLETION_READY_PUBLISH_FAILED = 66, COMPLETION_INVALID_SHAPE = 67, + COMPLETION_UNEXPECTED_GANG_SLOT = 73, + COMPLETION_GENERATION_MISMATCH = 74, + DEFERRED_RESERVATION_INVALID_OWNER = 75, + DEFERRED_RESERVATION_INVALID_STATE = 76, + DEFERRED_PUBLISH_INVALID_RESERVATION = 77, + NORMAL_DISPATCH_INVALID_TOPOLOGY = 78, }; struct alignas(128) SchedulerRunControl { diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index 7aefab1a97..e6e6c0af38 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -940,6 +940,7 @@ target_sources(test_hbg_submit_poison PRIVATE ) 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_a5_hbg_runtime_test(test_a5_hbg_scheduler_dispatch a5/test_hbg_scheduler_dispatch.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/a5/test_hbg_scheduler_dispatch.cpp b/tests/ut/cpp/a5/test_hbg_scheduler_dispatch.cpp new file mode 100644 index 0000000000..ac5c3f6552 --- /dev/null +++ b/tests/ut/cpp/a5/test_hbg_scheduler_dispatch.cpp @@ -0,0 +1,896 @@ +/* + * 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 "scheduler/scheduler_dispatch.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"); + storage_ = image_->storage.data(); + fanins_ = image_->fanins.data(); + for (size_t task = 0; task < capacity_; ++task) { + storage_[task].task.task_id = TaskId{static_cast(task)}; + storage_[task].payload.bind_regions( + nullptr, nullptr, fanins_ + task * static_cast(SCHEDULER_GRAPH_MAX_FANIN) + ); + if (storage_[task].payload.fanin_data() == nullptr) { + throw std::logic_error("test graph fanin region must share its contiguous image"); + } + for (int slot = 0; slot < 3; ++slot) + storage_[task].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)); + storage_[task].task.kernel_id[subtask_slot] = 1; + storage_[task].payload.fanin_count = static_cast(fanins.size()); + ASSERT_TRUE(fanins.empty() || storage_[task].payload.fanin_data() != nullptr); + std::copy(fanins.begin(), fanins.end(), storage_[task].payload.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) storage_[task].task.kernel_id[subtask_slot] = 1; + } + storage_[task].payload.fanin_count = 0; + } + + SchedulerGraphView graph() const { + return { + reinterpret_cast(storage_), + 0, + task_count_, + capacity_ - 1, + }; + } + +private: + static constexpr size_t kMaxTaskCount = 8192; + struct alignas(64) GraphImage { + std::array storage{}; + std::array fanins{}; + }; + + size_t task_count_; + size_t capacity_{1}; + std::unique_ptr image_; + ChipTaskStorage *storage_{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); + owner_states = + scheduler_state_at(scheduler_state->base(), layout.ready_owner_states_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; + scheduler_ready_owner_init(&owner_states[worker]); + } + metadata = scheduler_state_at(scheduler_state->base(), layout.task_metadata_offset); + 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}; + SchedulerReadyOwnerState *owner_states{nullptr}; + SchedulerTaskMetadata *metadata{nullptr}; +}; + +void configure_normal_aiv_cluster(FixtureStorage &storage, uint64_t task_count) { + storage.contexts[0].core_type = static_cast(CoreType::AIC); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_index = 0; + resolver.resolver_count = 1; + resolver.inbox_index = 0; + resolver.cluster_worker_ids[0] = 0; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = 2; + storage.run_control->resolver_count = 1; + auto *callables = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.callable_addresses_offset); + callables[1] = 0x1000; + for (uint64_t worker = 0; worker < 3; ++worker) { + for (uint32_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) + scheduler_initialize_free_slot( + scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, worker, slot) + ); + } + for (uint64_t task = 0; task < task_count; ++task) { + storage.metadata[task].kernel_ids[0] = UINT16_MAX; + storage.metadata[task].kernel_ids[1] = 1; + storage.metadata[task].active_mask = 2; + storage.metadata[task].flags = SCHEDULER_TASK_EXECUTABLE; + } +} + +void enqueue_normal_aiv_tasks( + FixtureStorage &storage, SchedulerWorkerContext &resolver, uint64_t task_begin, uint64_t task_end +) { + SchedulerReadyBatch batch{}; + SchedulerReadyStats ready_stats{}; + for (uint64_t task = task_begin; task < task_end; ++task) { + auto *control = + scheduler_task_control_at(storage.scheduler_state->base(), &resolver, static_cast(task)); + control->state = static_cast(SchedulerTaskState::READY); + ASSERT_TRUE(scheduler_ready_batch_append( + storage.scheduler_state->base(), &resolver, static_cast(task), &batch, &ready_stats + )); + } + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &resolver, 1, 0, &batch, &ready_stats, + &storage.owner_states[resolver.inbox_index] + )); +} + +void occupy_normal_slot( + FixtureStorage &storage, SchedulerWorkerContext &resolver, uint64_t worker_id, uint32_t pending_slot, + int64_t task_id +) { + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, worker_id, pending_slot); + slot->task_id = task_id; + slot->subtask_slot = 1; + slot->gang = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::READY) + ); +} + +TEST(SchedulerCompletionInbox, PacksBothGenerationSlotsInOneDeviceWord) { + alignas(uint64_t) volatile uint32_t generations[SCHEDULER_PENDING_SLOT_COUNT] = { + UINT32_C(0x11223344), UINT32_C(0x55667788) + }; + EXPECT_EQ(scheduler_gm_query_u32_pair(generations), UINT64_C(0x5566778811223344)); + scheduler_gm_store(generations[1], UINT32_C(0)); + EXPECT_EQ(scheduler_gm_query_u32_pair(generations), UINT64_C(0x0000000011223344)); +} + +TEST(SchedulerClusterCompletion, SpscGenerationCompletesNormalTask) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_index = 0; + resolver.resolver_count = 1; + resolver.cluster_worker_ids[0] = 0; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = UINT64_MAX; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + slot->task_id = 0; + slot->gang = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::READY) + ); + auto *completion_line = scheduler_completion_inbox_at(storage.scheduler_state->base(), &resolver, 0); + completion_line->completed_generations[0] = slot->generation; + auto *control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + control->state = static_cast(SchedulerTaskState::READY); + SchedulerWakeStats wake_stats{}; + SchedulerReadyStats ready_stats{}; + SchedulerCompletionStats completion_stats{}; + ASSERT_TRUE(scheduler_service_cluster_completions( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &wake_stats, &ready_stats, + &completion_stats, nullptr, false, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(completion_line->completed_generations[0], 0u); + EXPECT_EQ(scheduler_dispatch_state(slot->publication), SchedulerDispatchSlotState::FREE); + EXPECT_EQ(control->state, static_cast(SchedulerTaskState::DONE)); + EXPECT_EQ(control->wake_list_head, SCHEDULER_WAKE_LIST_CLOSED); + EXPECT_EQ(storage.run_control->resolved_task_count, 1u); +} + +TEST(SchedulerClusterCompletion, RejectsStaleCompletionGenerationAtNamedSite) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.cluster_worker_ids[0] = 0; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + slot->task_id = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::READY) + ); + + SchedulerWakeStats wake_stats{}; + SchedulerReadyStats ready_stats{}; + SchedulerCompletionStats completion_stats{}; + EXPECT_FALSE(scheduler_service_cluster_completion_slot( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, 0, 0, slot->generation + 1, + &wake_stats, &ready_stats, &completion_stats, nullptr, false, nullptr, nullptr, nullptr, + &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::COMPLETION_GENERATION_MISMATCH) + ); +} + +TEST(SchedulerClusterCompletion, RejectsUnexpectedGangSlotAtNamedSite) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.cluster_worker_ids[0] = 0; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + slot->task_id = 0; + slot->gang = 1; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::READY) + ); + + SchedulerWakeStats wake_stats{}; + SchedulerReadyStats ready_stats{}; + SchedulerCompletionStats completion_stats{}; + EXPECT_FALSE(scheduler_service_cluster_completion_slot( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, 0, 0, slot->generation, + &wake_stats, &ready_stats, &completion_stats, nullptr, false, nullptr, nullptr, nullptr, + &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::COMPLETION_UNEXPECTED_GANG_SLOT) + ); +} + +TEST(SchedulerClusterCompletion, PropagatesTraceToCompletionAndWokenTask) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0, {0}); + storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_count = 1; + resolver.cluster_worker_ids[0] = 0; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = 2; + auto *producer = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + producer->state = static_cast(SchedulerTaskState::READY); + SchedulerWakeStats wake_stats{}; + ASSERT_EQ( + scheduler_bootstrap_route_task( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, 1, &wake_stats + ), + SchedulerRouteResult::WAITING + ); + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + slot->task_id = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::READY) + ); + auto *completion_line = scheduler_completion_inbox_at(storage.scheduler_state->base(), &resolver, 0); + completion_line->completed_generations[0] = slot->generation; + auto *traces = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.trace_cells_offset); + producer->completion_resolve_start_cycles = 1; + producer->completion_resolve_end_cycles = 1; + producer->resolver_worker_id = UINT64_MAX; + traces[1].ready_transition_cycles = 1; + + SchedulerReadyStats ready_stats{}; + SchedulerCompletionStats completion_stats{}; + ASSERT_TRUE(scheduler_service_cluster_completions( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &wake_stats, &ready_stats, + &completion_stats, nullptr, true, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(producer->completion_resolve_start_cycles, 0u); + EXPECT_EQ(producer->completion_resolve_end_cycles, 0u); + EXPECT_EQ(producer->resolver_worker_id, resolver.worker_index); + EXPECT_EQ(traces[1].ready_transition_cycles, 0u); + auto *waiter = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 1); + EXPECT_EQ(waiter->state, static_cast(SchedulerTaskState::READY)); +} + +TEST(SchedulerClusterCompletion, CatchupRefreshIsBoundedToInitiallyEmptySibling) { + EXPECT_EQ(scheduler_completion_catchup_mask(0), UINT32_C(0)); + EXPECT_EQ(scheduler_completion_catchup_mask(1), UINT32_C(2)); + EXPECT_EQ(scheduler_completion_catchup_mask(2), UINT32_C(1)); + EXPECT_EQ(scheduler_completion_catchup_mask(3), UINT32_C(0)); +} + +TEST(SchedulerClusterCompletion, DirectlyRefillsCompletedSlotWhenReadyTaskExists) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.executable(1, 0); + storage.contexts[0].core_type = static_cast(CoreType::AIC); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_index = 0; + resolver.resolver_count = 1; + resolver.inbox_index = 0; + resolver.cluster_worker_ids[0] = 0; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = 2; + auto *callables = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.callable_addresses_offset); + callables[1] = 0x1000; + + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + const uint32_t completed_generation = slot->generation; + slot->task_id = 0; + slot->subtask_slot = 0; + slot->gang = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(completed_generation, SchedulerDispatchSlotState::READY) + ); + auto *completion_line = scheduler_completion_inbox_at(storage.scheduler_state->base(), &resolver, 0); + completion_line->completed_generations[0] = completed_generation; + auto *completed_control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + completed_control->state = static_cast(SchedulerTaskState::READY); + auto *ready_control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 1); + ready_control->state = static_cast(SchedulerTaskState::READY); + SchedulerReadyBatch batch{}; + SchedulerReadyStats ready_stats{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), &resolver, 1, &batch, &ready_stats)); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &resolver, 0, 0, &batch, &ready_stats, + &storage.owner_states[resolver.inbox_index] + )); + + SchedulerWakeStats wake_stats{}; + SchedulerCompletionStats completion_stats{}; + uint64_t ready_victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t direct_refilled_slot_mask = 0; + ASSERT_TRUE(scheduler_service_cluster_completions( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &wake_stats, &ready_stats, + &completion_stats, ready_victim_cursors, false, &direct_refilled_slot_mask, nullptr, + &storage.owner_states[resolver.inbox_index] + )); + + EXPECT_EQ(completion_line->completed_generations[0], 0u); + EXPECT_EQ(slot->task_id, 1); + EXPECT_EQ(slot->generation, completed_generation + 1); + EXPECT_EQ(scheduler_dispatch_state(slot->publication), SchedulerDispatchSlotState::READY); + EXPECT_EQ(completed_control->state, static_cast(SchedulerTaskState::DONE)); + EXPECT_EQ(storage.run_control->resolved_task_count, 1u); + EXPECT_EQ(direct_refilled_slot_mask, 1u); +} + +TEST(SchedulerNormalDispatch, FillsFreshAicSlot) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 0); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_count = 1; + resolver.inbox_index = 0; + resolver.cluster_worker_ids[0] = UINT64_MAX; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = 0; + storage.contexts[0].core_type = static_cast(CoreType::AIC); + auto *callables = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.callable_addresses_offset); + callables[1] = 0x1000; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + scheduler_initialize_free_slot(slot); + auto *control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + control->state = static_cast(SchedulerTaskState::READY); + SchedulerReadyBatch batch{}; + SchedulerReadyStats ready_stats{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), &resolver, 0, &batch, &ready_stats)); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &resolver, 0, 0, &batch, &ready_stats, + &storage.owner_states[resolver.inbox_index] + )); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(slot->task_id, 0); + EXPECT_EQ(scheduler_dispatch_state(slot->publication), SchedulerDispatchSlotState::READY); +} + +TEST(SchedulerNormalDispatch, PreservesProgressWhenALaterFillFails) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + graph.executable(0, 0); + graph.mixed(1, 3); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.is_resolver = 1; + resolver.resolver_count = 1; + resolver.inbox_index = 0; + resolver.cluster_worker_ids[0] = 0; + resolver.cluster_worker_ids[1] = 1; + resolver.cluster_worker_ids[2] = 2; + storage.contexts[0].core_type = static_cast(CoreType::AIC); + storage.metadata[1].active_mask = 3; + auto *callables = + scheduler_state_at(storage.scheduler_state->base(), storage.layout.callable_addresses_offset); + callables[1] = 0x1000; + for (uint32_t slot_index = 0; slot_index < SCHEDULER_PENDING_SLOT_COUNT; ++slot_index) + scheduler_initialize_free_slot( + scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, slot_index) + ); + SchedulerReadyBatch batch{}; + SchedulerReadyStats ready_stats{}; + for (int64_t task_id = 0; task_id < 2; ++task_id) { + auto *control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, task_id); + control->state = static_cast(SchedulerTaskState::READY); + ASSERT_TRUE( + scheduler_ready_batch_append(storage.scheduler_state->base(), &resolver, task_id, &batch, &ready_stats) + ); + } + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), &resolver, 0, 0, &batch, &ready_stats, + &storage.owner_states[resolver.inbox_index] + )); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + bool failed = false; + EXPECT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, nullptr, &storage.owner_states[resolver.inbox_index], &failed + )); + EXPECT_TRUE(failed); + auto *first_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 0); + EXPECT_EQ(first_slot->task_id, 0); + EXPECT_EQ(scheduler_dispatch_state(first_slot->publication), SchedulerDispatchSlotState::READY); + auto *failed_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 0, 1); + EXPECT_EQ(failed_slot->task_id, SCHEDULER_TASK_ID_INVALID); + EXPECT_EQ(scheduler_dispatch_state(failed_slot->publication), SchedulerDispatchSlotState::FILLING); + EXPECT_NE(storage.run_control->scheduler_error, 0u); + EXPECT_EQ(storage.run_control->error_site, static_cast(SchedulerErrorSite::DISPATCH_INVALID_SHAPE)); +} + +TEST(SchedulerNormalDispatch, RejectsExcessAivWorkersAtNamedSite) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + storage.contexts[0].core_type = static_cast(CoreType::AIV); + enqueue_normal_aiv_tasks(storage, resolver, 0, 1); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + bool failed = false; + EXPECT_FALSE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index], &failed + )); + EXPECT_TRUE(failed); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::NORMAL_DISPATCH_INVALID_TOPOLOGY) + ); +} + +TEST(SchedulerNormalDispatch, PublishesOrdinaryAivToPeerBeforeResolver) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + enqueue_normal_aiv_tasks(storage, resolver, 0, 1); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + auto *peer_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 2, 0); + EXPECT_EQ(peer_slot->task_id, 0); + EXPECT_EQ(scheduler_dispatch_state(peer_slot->publication), SchedulerDispatchSlotState::READY); + EXPECT_EQ(deferred.count, 0u); +} + +TEST(SchedulerDeferredAiv, ReservesOnlyAvailableResolverSlotsBeforeClaiming) { + FixtureStorage storage(3, 3); + GraphBuffer graph(3); + for (uint64_t task = 0; task < 3; ++task) + graph.executable(task, 1); + configure_normal_aiv_cluster(storage, 3); + SchedulerWorkerContext &resolver = storage.contexts[1]; + occupy_normal_slot(storage, resolver, 2, 0, SCHEDULER_TASK_ID_INVALID); + occupy_normal_slot(storage, resolver, 2, 1, SCHEDULER_TASK_ID_INVALID); + enqueue_normal_aiv_tasks(storage, resolver, 0, 3); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + + ASSERT_EQ(deferred.count, SCHEDULER_PENDING_SLOT_COUNT); + for (uint32_t slot_index = 0; slot_index < SCHEDULER_PENDING_SLOT_COUNT; ++slot_index) { + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, slot_index); + EXPECT_EQ(scheduler_dispatch_state(slot->publication), SchedulerDispatchSlotState::FILLING); + EXPECT_EQ(slot->task_id, SCHEDULER_TASK_ID_INVALID); + } + auto *ready_inbox = scheduler_ready_inbox_at(storage.scheduler_state->base(), &resolver, 1, 0); + EXPECT_NE(ready_inbox->head, SCHEDULER_INBOX_EMPTY); +} + +TEST(SchedulerDeferredAiv, RejectsInvalidReservationBeforeSlotAccess) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + SchedulerDeferredAivQueue deferred{}; + deferred.count = 1; + deferred.entries[0].ready.task_id = 0; + deferred.entries[0].reserved_slot = {UINT64_MAX, UINT32_MAX, 1}; + + uint32_t published_slot = 0; + EXPECT_FALSE(scheduler_publish_deferred_aiv_to_resolver( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, false, + &published_slot + )); + EXPECT_EQ(published_slot, UINT32_MAX); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::DEFERRED_PUBLISH_INVALID_RESERVATION) + ); +} + +TEST(SchedulerDeferredAiv, RejectsReservationWithMaterializedTask) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, 0); + slot->task_id = 0; + scheduler_gm_store( + slot->publication, scheduler_dispatch_publication(slot->generation, SchedulerDispatchSlotState::FILLING) + ); + SchedulerDeferredAivQueue deferred{}; + deferred.count = 1; + deferred.entries[0].ready.task_id = 0; + deferred.entries[0].reserved_slot = {resolver.worker_index, 0, slot->generation}; + + uint32_t published_slot = 0; + EXPECT_FALSE(scheduler_publish_deferred_aiv_to_resolver( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, false, + &published_slot + )); + EXPECT_EQ(published_slot, UINT32_MAX); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::DEFERRED_PUBLISH_INVALID_RESERVATION) + ); +} + +TEST(SchedulerDeferredAiv, RejectsReservationOwnedByPeerAtNamedSite) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + + EXPECT_FALSE(scheduler_release_deferred_aiv_reservation( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, SchedulerFreeSlotClaim{2, 0, 1} + )); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::DEFERRED_RESERVATION_INVALID_OWNER) + ); +} + +TEST(SchedulerDeferredAiv, RejectsReservationWithInvalidStateAtNamedSite) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + auto *slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, 0); + + EXPECT_FALSE(scheduler_release_deferred_aiv_reservation( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, + SchedulerFreeSlotClaim{resolver.worker_index, 0, slot->generation} + )); + EXPECT_EQ( + storage.run_control->error_site, static_cast(SchedulerErrorSite::DEFERRED_RESERVATION_INVALID_STATE) + ); +} + +TEST(SchedulerDeferredAiv, DoesNotClaimWithoutResolverReservation) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + for (uint64_t worker = 1; worker <= 2; ++worker) { + for (uint32_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) + occupy_normal_slot(storage, resolver, worker, slot, SCHEDULER_TASK_ID_INVALID); + } + enqueue_normal_aiv_tasks(storage, resolver, 0, 1); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + bool failed = true; + EXPECT_FALSE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index], &failed + )); + EXPECT_FALSE(failed); + EXPECT_EQ(deferred.count, 0u); + auto *ready_inbox = scheduler_ready_inbox_at(storage.scheduler_state->base(), &resolver, 1, 0); + EXPECT_EQ(ready_inbox->head, 0); +} + +TEST(SchedulerDeferredAiv, KeepsReservationForResolverWhenNoPeerIsActive) { + FixtureStorage storage(1, 3); + GraphBuffer graph(1); + graph.executable(0, 1); + configure_normal_aiv_cluster(storage, 1); + SchedulerWorkerContext &resolver = storage.contexts[1]; + resolver.cluster_worker_ids[2] = UINT64_MAX; + enqueue_normal_aiv_tasks(storage, resolver, 0, 1); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + ASSERT_EQ(deferred.count, 1u); + + SchedulerWakeStats wake_stats{}; + SchedulerCompletionStats completion_stats{}; + EXPECT_TRUE(scheduler_drain_deferred_aiv_to_peer( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, &wake_stats, + &ready_stats, &completion_stats, false, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(deferred.count, 1u); + + uint32_t self_slot = UINT32_MAX; + ASSERT_TRUE(scheduler_publish_deferred_aiv_to_resolver( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, false, &self_slot + )); + EXPECT_EQ(deferred.count, 0u); + EXPECT_LT(self_slot, SCHEDULER_PENDING_SLOT_COUNT); +} + +TEST(SchedulerDeferredAiv, PrefersNewPeerCapacityAndSelfPublishesOnlyOne) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + for (uint64_t task = 0; task < 2; ++task) + graph.executable(task, 1); + configure_normal_aiv_cluster(storage, 2); + SchedulerWorkerContext &resolver = storage.contexts[1]; + occupy_normal_slot(storage, resolver, 2, 0, SCHEDULER_TASK_ID_INVALID); + occupy_normal_slot(storage, resolver, 2, 1, SCHEDULER_TASK_ID_INVALID); + enqueue_normal_aiv_tasks(storage, resolver, 0, 2); + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + ASSERT_EQ(deferred.count, 2u); + const SchedulerFreeSlotClaim first_reservation = deferred.entries[0].reserved_slot; + + auto *peer_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 2, 0); + peer_slot->task_id = SCHEDULER_TASK_ID_INVALID; + scheduler_gm_store( + peer_slot->publication, scheduler_dispatch_publication(peer_slot->generation, SchedulerDispatchSlotState::FREE) + ); + SchedulerWakeStats wake_stats{}; + SchedulerCompletionStats completion_stats{}; + ASSERT_TRUE(scheduler_drain_deferred_aiv_to_peer( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, &wake_stats, + &ready_stats, &completion_stats, false, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + ASSERT_EQ(deferred.count, 1u); + EXPECT_EQ(scheduler_dispatch_state(peer_slot->publication), SchedulerDispatchSlotState::READY); + auto *released = scheduler_dispatch_slot_at( + storage.scheduler_state->base(), &resolver, first_reservation.worker_id, first_reservation.slot_index + ); + EXPECT_EQ(scheduler_dispatch_state(released->publication), SchedulerDispatchSlotState::FREE); + + uint32_t self_slot = UINT32_MAX; + ASSERT_TRUE(scheduler_publish_deferred_aiv_to_resolver( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, false, &self_slot + )); + EXPECT_EQ(deferred.count, 0u); + ASSERT_LT(self_slot, SCHEDULER_PENDING_SLOT_COUNT); + auto *published = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, self_slot); + EXPECT_EQ(scheduler_dispatch_state(published->publication), SchedulerDispatchSlotState::READY); +} + +TEST(SchedulerDeferredAiv, KeepsSecondReservationAfterOneSelfPublish) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + for (uint64_t task = 0; task < 2; ++task) + graph.executable(task, 1); + configure_normal_aiv_cluster(storage, 2); + SchedulerWorkerContext &resolver = storage.contexts[1]; + occupy_normal_slot(storage, resolver, 2, 0, SCHEDULER_TASK_ID_INVALID); + occupy_normal_slot(storage, resolver, 2, 1, SCHEDULER_TASK_ID_INVALID); + enqueue_normal_aiv_tasks(storage, resolver, 0, 2); + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + ASSERT_EQ(deferred.count, 2u); + + uint32_t self_slot = UINT32_MAX; + ASSERT_TRUE(scheduler_publish_deferred_aiv_to_resolver( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, false, &self_slot + )); + ASSERT_EQ(deferred.count, 1u); + auto *published = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, self_slot); + EXPECT_EQ(scheduler_dispatch_state(published->publication), SchedulerDispatchSlotState::READY); + const SchedulerFreeSlotClaim remaining = deferred.entries[0].reserved_slot; + auto *reserved = scheduler_dispatch_slot_at( + storage.scheduler_state->base(), &resolver, remaining.worker_id, remaining.slot_index + ); + EXPECT_EQ(scheduler_dispatch_state(reserved->publication), SchedulerDispatchSlotState::FILLING); + EXPECT_EQ(reserved->task_id, SCHEDULER_TASK_ID_INVALID); +} + +TEST(SchedulerDeferredAiv, RetiresCompletedPeerAndRefillsWithoutFreeDecision) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + graph.executable(0, 1); + graph.executable(1, 1); + configure_normal_aiv_cluster(storage, 2); + SchedulerWorkerContext &resolver = storage.contexts[1]; + occupy_normal_slot(storage, resolver, 2, 0, 0); + occupy_normal_slot(storage, resolver, 2, 1, SCHEDULER_TASK_ID_INVALID); + occupy_normal_slot(storage, resolver, 1, 1, SCHEDULER_TASK_ID_INVALID); + auto *completed_control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + completed_control->state = static_cast(SchedulerTaskState::READY); + enqueue_normal_aiv_tasks(storage, resolver, 1, 2); + + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyStats ready_stats{}; + SchedulerDeferredAivQueue deferred{}; + ASSERT_TRUE(scheduler_fill_cluster_normal_slots( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, victim_cursors, &ready_stats, + false, 0, nullptr, &deferred, &storage.owner_states[resolver.inbox_index] + )); + ASSERT_EQ(deferred.count, 1u); + auto *peer_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 2, 0); + const uint32_t completed_generation = peer_slot->generation; + auto *completion_line = scheduler_completion_inbox_at(storage.scheduler_state->base(), &resolver, 2); + completion_line->completed_generations[0] = completed_generation; + + SchedulerWakeStats wake_stats{}; + SchedulerCompletionStats completion_stats{}; + ASSERT_TRUE(scheduler_drain_deferred_aiv_to_peer( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &deferred, &wake_stats, + &ready_stats, &completion_stats, false, nullptr, nullptr, &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(deferred.count, 0u); + EXPECT_EQ(completion_line->completed_generations[0], 0u); + EXPECT_EQ(completed_control->state, static_cast(SchedulerTaskState::DONE)); + EXPECT_EQ(storage.run_control->resolved_task_count, 1u); + EXPECT_EQ(peer_slot->task_id, 1); + EXPECT_EQ(peer_slot->generation, completed_generation + 1); + EXPECT_EQ(scheduler_dispatch_state(peer_slot->publication), SchedulerDispatchSlotState::READY); +} + +TEST(SchedulerDeferredAiv, ResolverCompletionDoesNotDirectRefillItself) { + FixtureStorage storage(2, 3); + GraphBuffer graph(2); + graph.executable(0, 1); + graph.executable(1, 1); + configure_normal_aiv_cluster(storage, 2); + SchedulerWorkerContext &resolver = storage.contexts[1]; + occupy_normal_slot(storage, resolver, 1, 0, 0); + auto *completed_slot = scheduler_dispatch_slot_at(storage.scheduler_state->base(), &resolver, 1, 0); + auto *completion_line = scheduler_completion_inbox_at(storage.scheduler_state->base(), &resolver, 1); + completion_line->completed_generations[0] = completed_slot->generation; + auto *completed_control = scheduler_task_control_at(storage.scheduler_state->base(), &resolver, 0); + completed_control->state = static_cast(SchedulerTaskState::READY); + enqueue_normal_aiv_tasks(storage, resolver, 1, 2); + + SchedulerWakeStats wake_stats{}; + SchedulerReadyStats ready_stats{}; + SchedulerCompletionStats completion_stats{}; + uint64_t victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t direct_refilled_slot_mask = 0; + ASSERT_TRUE(scheduler_service_cluster_completions( + graph.graph(), storage.scheduler_state->base(), &resolver, storage.run_control, &wake_stats, &ready_stats, + &completion_stats, victim_cursors, false, &direct_refilled_slot_mask, nullptr, + &storage.owner_states[resolver.inbox_index] + )); + EXPECT_EQ(direct_refilled_slot_mask, 0u); + EXPECT_EQ(completed_slot->task_id, SCHEDULER_TASK_ID_INVALID); + EXPECT_EQ(scheduler_dispatch_state(completed_slot->publication), SchedulerDispatchSlotState::FREE); + auto *ready_inbox = scheduler_ready_inbox_at(storage.scheduler_state->base(), &resolver, 1, 0); + EXPECT_EQ(ready_inbox->head, 1); +} + +} // namespace