diff --git a/src/a5/runtime/host_build_graph/aicore/aicore_executor.cpp b/src/a5/runtime/host_build_graph/aicore/aicore_executor.cpp index c56fe7c1e7..c3959d1d0e 100644 --- a/src/a5/runtime/host_build_graph/aicore/aicore_executor.cpp +++ b/src/a5/runtime/host_build_graph/aicore/aicore_executor.cpp @@ -12,7 +12,7 @@ #include "aicore/aicore.h" #include "aicore/aicore_profiling_state.h" // AICore dependency scheduling uses one device-side protocol. -#include "scheduler/scheduler_dispatch.h" +#include "scheduler/scheduler_mix.h" #include "common/platform_config.h" #include "dispatch_payload.h" #include "runtime.h" @@ -149,7 +149,7 @@ __aicore__ bool bootstrap_ready_graph( return false; const bool phase_timing_enabled = scheduler_phase_timing_enabled(profiling_level); if (phase_timing_enabled) stats->bootstrap_start_cycles = scheduler_cycles(); - SchedulerReadyBatch batches[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyBatch batches[SCHEDULER_READY_QUEUE_COUNT]{}; uint64_t tasks_per_scheduler = graph.task_count / scheduler_count; uint64_t remainder = graph.task_count % scheduler_count; uint64_t task_begin = @@ -170,8 +170,7 @@ __aicore__ bool bootstrap_ready_graph( SchedulerRouteResult::READY_TO_ENQUEUE; if (route == SchedulerRouteResult::ERROR) return false; if (route == SchedulerRouteResult::READY_TO_ENQUEUE) { - const uint32_t core_type = - scheduler_metadata_core_type_index(scheduler_metadata_single_subtask_slot(metadata->active_mask)); + const uint32_t core_type = scheduler_task_ready_queue(metadata->flags, metadata->active_mask); if (!scheduler_bootstrap_ready_batch_append( scheduler_state_base, scheduler, static_cast(task_id), &batches[core_type], phase_timing_enabled ? &stats->ready : nullptr, profiling_level @@ -182,7 +181,7 @@ __aicore__ bool bootstrap_ready_graph( } scheduler_cache_barrier(); uint64_t ready_types = 0; - for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + for (uint32_t type = 0; type < SCHEDULER_READY_QUEUE_COUNT; ++type) { if (!scheduler_bootstrap_ready_batch_publish( scheduler_state_base, scheduler, type, scheduler->config.scheduler_index, &batches[type], phase_timing_enabled ? &stats->ready : nullptr, &ready_types @@ -233,11 +232,16 @@ __aicore__ bool bootstrap_ready_graph( // Prepare the first executable wave while the sole DMB launch gate is // still closed. - uint64_t ready_victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{ + uint64_t ready_victim_cursors[SCHEDULER_READY_QUEUE_COUNT]{ (scheduler->config.scheduler_index + 1) % scheduler_count, (scheduler->config.scheduler_index + 1) % scheduler_count, }; bool fill_failed = false; + if (!scheduler_fill_cluster_mix_slots( + graph, scheduler_state_base, scheduler, run_control, &ready_victim_cursors[SCHEDULER_MIX_QUEUE], + phase_timing_enabled ? &stats->ready : nullptr, profiling_level, true, ssbuf_region, nullptr + )) + return false; (void)scheduler_fill_cluster_normal_slots( graph, scheduler_state_base, scheduler, run_control, ready_victim_cursors, phase_timing_enabled ? &stats->ready : nullptr, profiling_level, 0, deferred_aiv, &fill_failed, ssbuf_region @@ -276,13 +280,13 @@ __aicore__ bool run_ready_dispatch_loop( bool scheduler_worker = context->is_scheduler(); if (scheduler_count == 0) return false; if (ssbuf_region == nullptr || cluster_lane >= PLATFORM_CORES_PER_BLOCKDIM) return false; - uint64_t ready_victim_cursors[SCHEDULER_CORE_TYPE_COUNT]{ + uint64_t ready_victim_cursors[SCHEDULER_READY_QUEUE_COUNT]{ scheduler_worker ? (context->config.scheduler_index + 1) % scheduler_count : 0, scheduler_worker ? (context->config.scheduler_index + 1) % scheduler_count : 0, }; - uint32_t seen_generation[SCHEDULER_PENDING_SLOT_COUNT]{}; uint64_t completion_publication = 0; uint32_t scan_start = 0; + uint64_t next_dispatch_sequence = 1; uint32_t backoff_iterations = kInitialBackoffIterations; uint32_t scheduler_error_poll_count = 0; if (phase_timing_enabled) context->profiling->loop_iter = UINT32_MAX; @@ -317,7 +321,7 @@ __aicore__ bool run_ready_dispatch_loop( return false; } uint32_t published_local_slot = UINT32_MAX; - if (scheduler_worker && deferred_aiv != nullptr && deferred_aiv->count != 0) { + if (scheduler_worker && !context->has_mix && deferred_aiv != nullptr && deferred_aiv->count != 0) { const uint32_t deferred_before = deferred_aiv->count; if (!scheduler_drain_deferred_aiv_to_peer( graph, scheduler_state_base, context, run_control, deferred_aiv, @@ -334,12 +338,26 @@ __aicore__ bool run_ready_dispatch_loop( } if (scheduler_worker && published_local_slot == UINT32_MAX) { uint64_t direct_refilled_slot_mask = 0; + SchedulerRefillCandidates refill_candidates(context->has_mix); const bool completion_progress = scheduler_service_cluster_completions( graph, scheduler_state_base, context, run_control, phase_timing_enabled ? &stats->wake : nullptr, phase_timing_enabled ? &stats->ready : nullptr, phase_timing_enabled ? &stats->completion : nullptr, - ready_victim_cursors, profiling_level, &direct_refilled_slot_mask, ssbuf_region + ready_victim_cursors, profiling_level, &direct_refilled_slot_mask, ssbuf_region, + context->has_mix ? &refill_candidates : nullptr ); scheduler_progress = completion_progress; + if (!scheduler_fill_mix_after_completions( + graph, scheduler_state_base, context, run_control, &refill_candidates.mix_ready, + &ready_victim_cursors[SCHEDULER_MIX_QUEUE], phase_timing_enabled ? &stats->ready : nullptr, + profiling_level, deferred_aiv == nullptr || deferred_aiv->count == 0, ssbuf_region, + &scheduler_progress + )) + return false; + if (context->has_mix && !scheduler_publish_refill_candidates( + graph, scheduler_state_base, context, run_control, &refill_candidates, + phase_timing_enabled ? &stats->ready : nullptr, profiling_level, ssbuf_region + )) + return false; bool fill_failed = false; const bool dispatch_progress = scheduler_fill_cluster_normal_slots( graph, scheduler_state_base, context, run_control, ready_victim_cursors, @@ -377,10 +395,10 @@ __aicore__ bool run_ready_dispatch_loop( if (scheduler_worker) { uint32_t local_slot = UINT32_MAX; uint64_t local_generation = 0; - if (scheduler_local_ready_pop(context, scan_start, &local_slot, &local_generation)) { + if (scheduler_local_ready_pop(context, &local_slot, &local_generation)) { if (scheduler_dispatch_state(local_generation) != SchedulerDispatchSlotState::READY || scheduler_dispatch_generation(local_generation) == 0 || - scheduler_dispatch_generation(local_generation) == seen_generation[local_slot]) { + scheduler_dispatch_generation(local_generation) != next_dispatch_sequence) { scheduler_record_error( run_control, context->slots[cluster_lane][local_slot].task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, @@ -398,7 +416,7 @@ __aicore__ bool run_ready_dispatch_loop( scheduler_ssbuf_load_relaxed(&ssbuf_region->lanes[cluster_lane].dispatch[slot_index].publication); const uint32_t generation = static_cast(publication); if (phase_timing_enabled) ++stats->task_state_poll_count; - if (generation != 0 && generation != seen_generation[slot_index]) { + if (generation == next_dispatch_sequence) { ready_slot = static_cast(slot_index); ready_generation = generation; ready_publication = publication; @@ -424,9 +442,11 @@ __aicore__ bool run_ready_dispatch_loop( scheduler_append_idle_activity(scheduler_state_base, context, idle_start_cycles, ready_observe); idle_active = false; } - scheduler_observe_dispatch_payload_control(payload); - scheduler_observe_dispatch_payload_arguments(payload); - scheduler_observe_dispatch_payload_barrier(); + if (!scheduler_worker) { + scheduler_observe_dispatch_payload_control(payload); + scheduler_observe_dispatch_payload_arguments(payload); + scheduler_observe_dispatch_payload_barrier(); + } const int64_t task_id = scheduler_worker ? local_slot->task_id : ssbuf_control->task_id; if (task_id < 0 || static_cast(task_id) >= graph.task_count || ready_generation == 0) { scheduler_record_error( @@ -435,7 +455,7 @@ __aicore__ bool run_ready_dispatch_loop( ); return false; } - seen_generation[slot_index] = ready_generation; + ++next_dispatch_sequence; SchedulerExecutorTaskTrace execution_trace{}; if (commit_scheduler_trace || commit_task_timing) { uint64_t local_completion_index = stats->completion.enqueue_count; diff --git a/src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp index 86364eabed..d023a70068 100644 --- a/src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp @@ -179,16 +179,23 @@ static void publish_aicore_task_timing(Runtime *runtime) { auto *metadata = scheduler_state_at(scheduler_state_base, context->task_metadata_offset); auto *traces = scheduler_state_at(scheduler_state_base, context->trace_cells_offset); cache_invalidate_range(metadata, static_cast(context->graph_task_count) * sizeof(*metadata)); - cache_invalidate_range(traces, static_cast(context->graph_task_count) * sizeof(*traces)); for (uint64_t task_id = 0; task_id < context->graph_task_count; ++task_id) { const int32_t slot = metadata[task_id].timing_slot; if (slot < 0 || slot >= NUM_TASK_TIMING_SLOTS) continue; - const uint64_t start = traces[task_id].kernel_start_cycles; - const uint64_t end = traces[task_id].kernel_end_cycles; - if (start == 0 || end <= start) continue; - if (start < records[slot].dispatch_cycle) records[slot].dispatch_cycle = start; - if (end > records[slot].finish_cycle) records[slot].finish_cycle = end; + if (!scheduler_task_is_executable(metadata[task_id].flags)) continue; + for (uint8_t subtask = 0; subtask < 3; ++subtask) { + if ((metadata[task_id].active_mask & (1U << subtask)) == 0) continue; + auto &trace = traces[scheduler_task_trace_index( + metadata[task_id].trace_index_base, metadata[task_id].active_mask, subtask + )]; + cache_invalidate_range(&trace, sizeof(trace)); + const uint64_t start = trace.kernel_start_cycles; + const uint64_t end = trace.kernel_end_cycles; + if (start == 0 || end <= start) continue; + if (start < records[slot].dispatch_cycle) records[slot].dispatch_cycle = start; + if (end > records[slot].finish_cycle) records[slot].finish_cycle = end; + } } aicpu_publish_task_timing_tail_usage(1); } diff --git a/src/a5/runtime/host_build_graph/common/intrinsic.h b/src/a5/runtime/host_build_graph/common/intrinsic.h index 98f8289f54..089b80e2fe 100644 --- a/src/a5/runtime/host_build_graph/common/intrinsic.h +++ b/src/a5/runtime/host_build_graph/common/intrinsic.h @@ -62,8 +62,8 @@ * issue #900 (PR #899 spmd_paged_attention_highperf); the kernel * compiled, ran without error, and produced wrong output. Use * `get_sub_block_id(args)` instead, which reads from the runtime's - * `GlobalContext.sub_block_id`. Resident graph materialization sets it - * from the selected AIV task subslot before every dispatch publication. + * `GlobalContext.sub_block_id`. Resident Mix dispatch sets it + * from the final physical AIV placement before publication. * * - `get_block_idx()` and `get_block_num()` are not redirected to * simpler's LocalContext either — use the `(args)` variants below @@ -107,7 +107,7 @@ static constexpr int32_t PAYLOAD_GLOBAL_CONTEXT_INDEX = SPMD_GLOBAL_CONTEXT_INDE /** * Per-dispatch global context, stored in DispatchPayload. Resident graph - * materialization sets sub_block_id from the selected AIV task subslot; the + * Mix dispatch sets sub_block_id from the final physical AIV placement; the * legacy scheduler seeds the equivalent per-core value during startup. */ struct GlobalContext { diff --git a/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index 5c352bf687..b44fc457c5 100644 --- a/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -420,7 +420,7 @@ READY with the same generation until the local Executor claims it, so the ready token is reconstructed from the slot. Completion generation validation still prevents stale notifications from freeing or refilling a pending slot. -The local configuration occupies 88 bytes and the base local state 256 bytes +The local configuration occupies 88 bytes and the base local state 376 bytes under the 64-bit ABI, including its optional profiling pointer. A separate 240-byte profiling state holds six timing slots, two self-execution traces, worker trace caches, profiling offsets, and the loop counter/valid mask. @@ -430,7 +430,7 @@ The host records whether any task requests sampled timing in the run control. Only a run with chip profiling or sampled timing enters the resident function specialization that allocates profiling state; the plain specialization does not allocate it. These functions do not inline into the common entry. The -combined local state with profiling is 496 bytes. These sizes exclude other +combined local state with profiling is 616 bytes. These sizes exclude other function locals, worker statistics and compiler spills. Compile-time assertions anchor the 64-bit configuration, slot, base and profiling state sizes; they do not establish the dynamic AICore stack high-water mark. @@ -453,6 +453,10 @@ writeback and ready publication share a release barrier. Tokens are SPSC and use no SSBUF read-modify-write atomics. The Scheduler initializes every token before publishing the header, and each invocation validates the region. Self-execution uses local notifications, completion generations and trace storage. +Each slot's embedded context pointers and constant single-block fields are +initialized in the host-created resident payload image. Self-execution consumes +its own payload writes after the publication barrier, without an additional +payload invalidation on pickup. Remote Executors retain payload invalidation. The Scheduler and Executor share the SSBUF structure definitions in the same runtime build. Each run initializes the region before use. @@ -488,6 +492,39 @@ therefore avoids periodic publication on the busy path; the continuous-busy timeout is an accepted limit outside that expected duration, not evidence that busy execution has stopped making progress. +Single-block ordinary Mix uses the resident path. Each participating lane +owns an independent GM DispatchPayload and one slot. A local six-entry tracker +joins completions: each lane releases its slot immediately; only the final lane +marks the task done and resolves dependencies. Kernels implement their own +cross-lane synchronization. There is no cluster waiting-count gate or ACK state. + +Ordinary and Mix dispatches share a per-lane sequence in the publication word. +Preparation reserves resources without consuming a sequence; publication commits +it after the final lane is known. Executors consume dispatch order across both +slots, including Scheduler self-execution. Mix prepares every participating lane +before publishing any lane, and finishes the group's publications before the +next Mix. Published ordinary tasks retain their position. +One local metadata snapshot and one argument-region parse serve all lanes of a +Mix. Shared scalar values and tensor descriptor addresses are read or computed +once, then written into each lane's independent payload. After all payloads +are written back, one barrier precedes their individual READY publications. + +Mix admission checks all participating slots and a tracker before claiming the +ready inbox head. With ordinary deferred reservations, local Mix may be claimed +but remote Mix cannot be stolen. Mix never enters the deferred queue. On graphs +containing Mix, a scheduling pass consumes completions, attempts the unlocked +direct Mix candidate, admits queued Mix, and then refills ordinary work. +Displaced direct successors return to ready queues; +ordinary-only graphs keep immediate direct refill. Task traces use metadata's +32-bit base index plus the logical subtask offset, independently of placement. +An unlocked Mix may remain a candidate within the current scheduling pass when +no Mix work was queued at its ready-queue snapshot. Complete lane and tracker +capacity permit direct preparation before the next queue claim; otherwise the +candidate joins its ready queue first. Queue arrivals between that snapshot +and a failed placement can precede the candidate after it joins the queue. +Phase profiling records ready transitions for every active subtask. With +profiling disabled, refill candidates do not compute trace indices. + Directory queries are skipped when no unreserved FREE slot can accept work. Idle polling backs off from 8 to at most 32 iterations. Dispatch trace GM writes follow ready publication; Host tooling computes ready-to-kernel diff --git a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp index bbd7d6a16b..c50529dbf7 100644 --- a/src/a5/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a5/runtime/host_build_graph/host/runtime_maker.cpp @@ -481,10 +481,11 @@ bool publish_aicore_scheduler_profiling(Runtime *runtime, const HostApi *api) { // activity streams -- every section after the one that gave up -- while // leaving the already-published AicoreTasks in the one shape the reader // rejects. - std::vector emitted_tasks; - emitted_tasks.reserve(static_cast(owner.layout.task_count)); - for (uint64_t task_id = 0; task_id < owner.layout.task_count; ++task_id) { - const SchedulerTaskTrace &trace = traces[task_id]; + std::vector emitted_traces; + emitted_traces.reserve(static_cast(owner.layout.executable_subtask_count)); + for (uint64_t trace_index = 0; trace_index < owner.layout.executable_subtask_count; ++trace_index) { + const SchedulerTaskTrace &trace = traces[trace_index]; + const uint64_t task_id = trace.task_id; if (trace.valid == 0 || trace.kernel_start_cycles == 0 || trace.kernel_end_cycles < trace.kernel_start_cycles || trace.worker_id >= SCHEDULER_WORKER_CAPACITY) continue; @@ -497,15 +498,16 @@ bool publish_aicore_scheduler_profiling(Runtime *runtime, const HostApi *api) { ); continue; } - emitted_tasks.push_back(task_id); + emitted_traces.push_back(trace_index); } std::ostringstream tasks_json; tasks_json << "["; bool first_task = true; const uint64_t swimlane_run_epoch = api->run_epoch(); - for (uint64_t task_id : emitted_tasks) { - const SchedulerTaskTrace &trace = traces[task_id]; + for (uint64_t trace_index : emitted_traces) { + const SchedulerTaskTrace &trace = traces[trace_index]; + const uint64_t task_id = trace.task_id; const uint64_t receive_to_start = trace.ready_observe_cycles != 0 && trace.kernel_start_cycles >= trace.ready_observe_cycles ? trace.kernel_start_cycles - trace.ready_observe_cycles : @@ -530,8 +532,9 @@ bool publish_aicore_scheduler_profiling(Runtime *runtime, const HostApi *api) { std::ostringstream scheduler_tasks_json; scheduler_tasks_json << "{\n \"producer\": \"aicore\",\n \"records\": ["; bool first_scheduler_task = true; - for (uint64_t task_id : emitted_tasks) { - const SchedulerTaskTrace &trace = traces[task_id]; + for (uint64_t trace_index : emitted_traces) { + const SchedulerTaskTrace &trace = traces[trace_index]; + const uint64_t task_id = trace.task_id; if (!first_scheduler_task) scheduler_tasks_json << ","; scheduler_tasks_json << "\n [" << trace.worker_id << ", " << task_id << ", " << trace.dispatch_end_cycles << ", " << trace.complete_start_cycles << ", " @@ -597,8 +600,9 @@ bool publish_aicore_scheduler_profiling(Runtime *runtime, const HostApi *api) { SchedPhaseKind::Bootstrap, context.bootstrap_task_count, 0, false ); } - for (uint64_t task_id = 0; task_id < owner.layout.task_count; ++task_id) { - const SchedulerTaskTrace &trace = traces[task_id]; + for (uint64_t trace_index = 0; trace_index < owner.layout.executable_subtask_count; ++trace_index) { + const SchedulerTaskTrace &trace = traces[trace_index]; + const uint64_t task_id = trace.task_id; if (trace.state_probe_scheduler_worker_id < SCHEDULER_WORKER_CAPACITY) { append_scheduler_record( &records[trace.state_probe_scheduler_worker_id], trace.state_probe_start_cycles, @@ -628,6 +632,8 @@ bool publish_aicore_scheduler_profiling(Runtime *runtime, const HostApi *api) { trace.refill_loop_iter, SchedPhaseKind::Refill, 1, trace.refill_task_id ); } + } + for (uint64_t task_id = 0; task_id < owner.layout.task_count; ++task_id) { const SchedulerTaskControl &control = controls[task_id]; if (control.scheduler_worker_id < SCHEDULER_WORKER_CAPACITY) { append_scheduler_record( @@ -1043,6 +1049,7 @@ bool create_scheduler_state( uint64_t executable_task_count = 0; uint64_t executable_subtask_count = 0; uint64_t gang_task_count = 0; + uint64_t mix_task_count = 0; uint64_t aic_worker_demand = 0; uint64_t aiv_worker_demand = 0; int64_t legacy_shape_task_id = -1; @@ -1156,8 +1163,7 @@ bool create_scheduler_state( metadata.flags = scheduler_task_metadata_flags_from_submit_state( metadata_active_mask, slot.task_attrs, slot.logical_block_num, fanin_count != 0, inline_dispatch_task ); - metadata.logical_block_num = static_cast(logical_block_num); - metadata.total_required_subtasks = static_cast(expected_subtasks); + metadata.trace_index_base = static_cast(executable_subtask_count); metadata.timing_slot = slot.task_attrs.timing_slot(); sampled_task_timing_enabled = sampled_task_timing_enabled || @@ -1171,6 +1177,7 @@ bool create_scheduler_state( ++aiv_task_count; aiv_worker_demand = std::max(aiv_worker_demand, logical_block_num * active_aiv_subtasks); } + if (scheduler_task_is_mix(metadata.flags)) ++mix_task_count; if (scheduler_task_is_gang(metadata.flags)) ++gang_task_count; executable_subtask_count += expected_subtasks; ++executable_task_count; @@ -1183,7 +1190,7 @@ bool create_scheduler_state( if (legacy_shape_task_id >= 0) { select_legacy_scheduler(runtime, SCHEDULER_RUNTIME_MODE_LEGACY_UNSUPPORTED_SHAPE); LOG_INFO( - "A5 HBG: retaining AICPU scheduling for task id=%" PRId64 " with MIX, SPMD, or sync-start shape", + "A5 HBG: retaining AICPU scheduling for task id=%" PRId64 " with SPMD or sync-start shape", legacy_shape_task_id ); return true; @@ -1192,7 +1199,8 @@ bool create_scheduler_state( AicoreSchedulerLayout layout{}; if (!scheduler_plan_layout( static_cast(total_tasks), aic_task_count, aiv_task_count, &layout, - api->chip_swimlane_level() >= static_cast(ChipSwimlaneLevel::SCHED_PHASES) + api->chip_swimlane_level() >= static_cast(ChipSwimlaneLevel::SCHED_PHASES), + executable_subtask_count ) || layout.total_size > std::numeric_limits::max() - (SCHEDULER_STATE_ALIGNMENT - 1)) { LOG_ERROR("A5 HBG AICore scheduler: scheduler state layout overflow"); @@ -1262,6 +1270,7 @@ bool create_scheduler_state( run_control->error_core_type = UINT64_MAX; auto *gang_coordinator = scheduler_state_at(host_base, layout.gang_coordinator_offset); gang_coordinator->gang_task_count = gang_task_count; + gang_coordinator->mix_task_count = mix_task_count; auto *contexts = scheduler_state_at(host_base, layout.worker_contexts_offset); int32_t aic_rank = 0; @@ -1303,6 +1312,17 @@ bool create_scheduler_state( context.task_window_last_index = task_window_size - 1; context.graph_task_count = static_cast(total_tasks); context.worker_index = static_cast(i); + for (uint32_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) { + const uint64_t offset = context.dispatch_payload_offset + slot * sizeof(DispatchPayload); + auto *payload = scheduler_state_at(host_base, offset); + const uint64_t device_payload = aligned_address + offset; + payload->args[PAYLOAD_LOCAL_CONTEXT_INDEX] = device_payload + offsetof(DispatchPayload, local_context); + payload->args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = device_payload + offsetof(DispatchPayload, global_context); + payload->local_context.block_idx = 0; + payload->local_context.block_num = 1; + payload->local_context.async_ctx.task_token = TaskId::invalid(); + payload->src_payload = 0; + } } runtime->publish_scheduler_bootstrap( diff --git a/src/a5/runtime/host_build_graph/runtime/dispatch_payload.h b/src/a5/runtime/host_build_graph/runtime/dispatch_payload.h index 795c49afc9..b98fb654b4 100644 --- a/src/a5/runtime/host_build_graph/runtime/dispatch_payload.h +++ b/src/a5/runtime/host_build_graph/runtime/dispatch_payload.h @@ -89,18 +89,21 @@ constexpr uint32_t TASKPAYLOAD_TENSOR_STRIDE = 128; // sizeof(simpler::hbg::Ten * concurrently dispatched cores. */ struct alignas(64) DispatchPayload { - // === Cache line 0 (64B): control block, the only line written per dispatch === - // function_bin_addr, local_context.{block_idx,block_num,async_ctx.task_token} - // and src_payload are the per-dispatch writes; async_ctx's slab pointers + - // capacity are cold (prefilled once at init) but ride this hot line for free. + // === Cache line 0 (64B): control block === + // The legacy path writes function_bin_addr, local_context's block fields and + // task token, and src_payload per dispatch. The SSBUF ready path initializes + // its fixed fields in the per-run image and updates function_bin_addr (and + // src_payload for inline tasks) during dispatch. + // async_ctx's slab pointers + capacity are cold but ride this hot line. // Sized to exactly 64B so both dispatch paths write one control line: the // ready path (src_payload = 0) then also fills args[0..num_args); the gated // path (src_payload = &TaskPayload) leaves args[] to the idle AICore. uint64_t function_bin_addr; /**< Kernel entry address in GM (set by Scheduler). */ - /** Per-dispatch context: block_idx/block_num (hot) + async_ctx (task_token hot, - * slab pointers + capacity prefilled once at init). args[SPMD_LOCAL_CONTEXT_INDEX] - * points here. */ + /** Context for the dispatched kernel. The legacy path updates block fields + * and task_token per dispatch; the SSBUF path keeps fixed values per run. + * Slab pointers and capacity are prefilled once at init. + * args[SPMD_LOCAL_CONTEXT_INDEX] points here. */ LocalContext local_context; /** Early-dispatch gate AND source pointer, folded into one field. 0 = ready: 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 index c16afe5d9f..29700dfaee 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.h @@ -38,13 +38,26 @@ inline __aicore__ void scheduler_account_failed_completion( ++scheduler->pending_completed; } +struct SchedulerRefillCandidates { + int64_t tasks[PLATFORM_CORES_PER_BLOCKDIM][SCHEDULER_PENDING_SLOT_COUNT]; + // Read an index only for a valid task in the same slot with phase timing enabled. + uint32_t completed_trace_indices[PLATFORM_CORES_PER_BLOCKDIM][SCHEDULER_PENDING_SLOT_COUNT]; + SchedulerReadyClaim mix_ready{}; + inline __aicore__ explicit SchedulerRefillCandidates(bool initialize = true) { + if (!initialize) return; + for (uint32_t lane = 0; lane < PLATFORM_CORES_PER_BLOCKDIM; ++lane) + for (uint32_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) + tasks[lane][slot] = SCHEDULER_TASK_ID_INVALID; + } +}; + inline __aicore__ bool scheduler_service_cluster_completion_slot( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, __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, uint64_t profiling_level, const SchedulerReadyClaim *replacement_ready, bool *direct_refilled, - SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, SchedulerRefillCandidates *candidates = nullptr ) { if (direct_refilled != nullptr) *direct_refilled = false; if (ssbuf_region == nullptr || scheduler == nullptr || cluster_lane >= PLATFORM_CORES_PER_BLOCKDIM || @@ -76,6 +89,24 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( scheduler_record_error(run_control, task_id, SchedulerGraphResult::INVALID_TASK_ID, &graph, scheduler); return false; } + SchedulerMixTracker *mix_tracker = nullptr; + bool valid_subtask = local_slot->subtask_slot < 3; + if (local_slot->mix_tracker != UINT8_MAX) { + if (local_slot->mix_tracker >= SCHEDULER_MIX_TRACKER_COUNT) valid_subtask = false; + else { + mix_tracker = &scheduler->mix_trackers[local_slot->mix_tracker]; + const uint8_t bit = valid_subtask ? 1U << local_slot->subtask_slot : 0; + valid_subtask = valid_subtask && mix_tracker->task_id == task_id && (mix_tracker->active_mask & bit) != 0 && + (mix_tracker->completed_mask & bit) == 0; + } + } + if (!valid_subtask) { + scheduler_record_error( + run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, scheduler, + SchedulerErrorSite::COMPLETION_GENERATION_MISMATCH + ); + return false; + } const bool sampled_task_timing_enabled = scheduler->sampled_task_timing(cluster_lane, pending_slot); SchedulerExecutorTaskTrace executor_trace{}; if (chip_task_timing_enabled || sampled_task_timing_enabled) { @@ -96,27 +127,34 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( if (chip_task_timing_enabled || sampled_task_timing_enabled) { __gm__ SchedulerTaskTrace *traces = scheduler_state_at(scheduler_state_base, scheduler->profiling->trace_cells_offset); - completed_trace = &traces[task_id]; + completed_trace = + &traces[scheduler_trace_index_at(scheduler_state_base, scheduler, task_id, local_slot->subtask_slot)]; } if (chip_task_timing_enabled) { - scheduler_observe_cache_line(completed_trace); - scheduler_observe_cache_line(&completed_trace->kernel_start_cycles); + // Finish all trace invalidations before any cached trace write. + scheduler_invalidate_cache_line(completed_trace); + scheduler_invalidate_cache_line(&completed_trace->kernel_start_cycles); + if (phase_timing_enabled) { + scheduler_invalidate_cache_line(&completed_trace->ready_transition_cycles); + scheduler_invalidate_cache_line(&completed_trace->dispatch_start_cycles); + scheduler_invalidate_cache_line(&completed_trace->refill_scheduler_worker_id); + scheduler_invalidate_cache_line(&completed_trace->descriptor_cache_observed_cycles); + if ((scheduler->profiling->worker_trace_valid_mask & (1U << cluster_lane)) == 0) { + scheduler_invalidate_cache_line(&target->trace_aicore_entry_cycles); + scheduler_invalidate_cache_line(&target->trace_register_release_cycles); + } + } + scheduler_cache_barrier(); completed_trace->kernel_start_cycles = executor_trace.kernel_start_cycles; completed_trace->kernel_end_cycles = executor_trace.kernel_end_cycles; completed_trace->ready_observe_cycles = executor_trace.ready_observe_cycles; if (schedule_timing_enabled) completed_trace->complete_start_cycles = completion_observe; if (phase_timing_enabled) { - scheduler_observe_cache_line(&completed_trace->ready_transition_cycles); - scheduler_observe_cache_line(&completed_trace->dispatch_start_cycles); - scheduler_observe_cache_line(&completed_trace->refill_scheduler_worker_id); - scheduler_observe_cache_line(&completed_trace->descriptor_cache_observed_cycles); completed_trace->ready_scan_start_cycles = executor_trace.ready_scan_start_cycles; completed_trace->completion_id = executor_trace.completion_id; completed_trace->completion_inbox_index = executor_trace.completion_inbox_index; SchedulerWorkerTraceCache *worker_trace = &scheduler->profiling->worker_traces[cluster_lane]; if ((scheduler->profiling->worker_trace_valid_mask & (1U << cluster_lane)) == 0) { - scheduler_observe_cache_line(&target->trace_aicore_entry_cycles); - scheduler_observe_cache_line(&target->trace_register_release_cycles); worker_trace->descriptor_cache_observed_cycles = target->trace_descriptor_cache_observed_cycles; worker_trace->aicore_entry_cycles = target->trace_aicore_entry_cycles; worker_trace->handshake_publish_cycles = target->trace_handshake_publish_cycles; @@ -139,29 +177,51 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( 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, scheduler, task_id); - scheduler_gm_store(control->state, static_cast(SchedulerTaskState::DONE)); + bool task_complete = true; + if (mix_tracker != nullptr) { + mix_tracker->completed_mask |= 1U << completed_subtask_slot; + task_complete = mix_tracker->completed_mask == mix_tracker->active_mask; + if (task_complete) mix_tracker->task_id = SCHEDULER_TASK_ID_INVALID; + } + local_slot->task_id = SCHEDULER_TASK_ID_INVALID; + local_slot->state = SchedulerDispatchSlotState::FREE; + local_slot->subtask_slot = UINT8_MAX; + local_slot->mix_tracker = UINT8_MAX; + scheduler->set_timing_slot(cluster_lane, pending_slot, -1); const uint64_t completion_phase_end = phase_timing_enabled ? scheduler_cycles() : 0; SchedulerReadyClaim resolved_ready{}; - SchedulerReadyClaim *direct_ready = - replacement_ready == nullptr && worker_id != scheduler->worker_id() ? &resolved_ready : nullptr; - if (!scheduler_resolve_completion( - graph, scheduler_state_base, scheduler, run_control, task_id, wake_stats, ready_stats, completion_stats, - profiling_level, false, direct_ready, scheduler_metadata_core_type_index(completed_subtask_slot) - )) { - scheduler_account_failed_completion( - graph, scheduler, run_control, task_id, SchedulerErrorSite::COMPLETION_RESOLVE_FAILED - ); - return false; + if (task_complete) { + __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, scheduler, task_id); + scheduler_gm_store(control->state, static_cast(SchedulerTaskState::DONE)); + SchedulerReadyClaim *direct_ready = replacement_ready == nullptr && worker_id != scheduler->worker_id() && + (!scheduler->has_mix || candidates != nullptr) ? + &resolved_ready : + nullptr; + if (!scheduler_resolve_completion( + graph, scheduler_state_base, scheduler, run_control, task_id, wake_stats, ready_stats, completion_stats, + profiling_level, false, direct_ready, scheduler_metadata_core_type_index(completed_subtask_slot), + scheduler->has_mix && candidates != nullptr ? &candidates->mix_ready : nullptr + )) { + scheduler_account_failed_completion( + graph, scheduler, run_control, task_id, SchedulerErrorSite::COMPLETION_RESOLVE_FAILED + ); + return false; + } + } + if (scheduler->has_mix && resolved_ready.task_id >= 0 && candidates != nullptr) { + candidates->tasks[cluster_lane][pending_slot] = resolved_ready.task_id; + if (phase_timing_enabled) + candidates->completed_trace_indices[cluster_lane][pending_slot] = + scheduler_trace_index_at(scheduler_state_base, scheduler, task_id, completed_subtask_slot); } SchedulerReadyClaim ready{}; - bool ready_available = replacement_ready != nullptr; + bool ready_available = !scheduler->has_mix && replacement_ready != nullptr; if (ready_available) { ready = *replacement_ready; - } else if (resolved_ready.task_id >= 0) { + } else if (!scheduler->has_mix && resolved_ready.task_id >= 0) { ready = resolved_ready; ready_available = true; - } else if (ready_victim_cursors != nullptr && worker_id != scheduler->worker_id()) { + } else if (!scheduler->has_mix && ready_victim_cursors != nullptr && worker_id != scheduler->worker_id()) { // A normal AIV task is never refilled directly onto the Scheduler. // Its completed slot becomes capacity for late binding instead. const uint32_t core_type = scheduler_metadata_core_type_index(completed_subtask_slot); @@ -191,6 +251,7 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( local_slot->generation, cluster_lane, }; + local_slot->state = SchedulerDispatchSlotState::FILLING; if (!scheduler_fill_dispatch_slot( graph, scheduler_state_base, scheduler, run_control, claim, ready, profiling_level, ssbuf_region )) { @@ -203,12 +264,6 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( } refill_end_cycles = phase_timing_enabled ? scheduler_cycles() : 0; if (refill_start_cycles == 0) refill_start_cycles = refill_end_cycles; - if (!refilled) { - local_slot->task_id = SCHEDULER_TASK_ID_INVALID; - local_slot->state = SchedulerDispatchSlotState::FREE; - local_slot->subtask_slot = UINT8_MAX; - scheduler->set_timing_slot(cluster_lane, pending_slot, -1); - } if (chip_task_timing_enabled) { if (phase_timing_enabled) { completed_trace->complete_end_cycles = completion_phase_end; @@ -227,15 +282,13 @@ inline __aicore__ bool scheduler_service_cluster_completion_slot( if (phase_timing_enabled) scheduler_writeback_cache_line(&completed_trace->refill_scheduler_worker_id); if (phase_timing_enabled) scheduler_writeback_cache_line(&completed_trace->descriptor_cache_observed_cycles); scheduler_cache_barrier(); - scheduler_gm_publish(completed_trace->valid, UINT64_C(1)); - // AICPU treats resolved_task_count as the graph-completion token. Keep - // valid globally ordered before that token so host collection cannot - // race the final trace publication. - scheduler_cache_barrier(); + scheduler_gm_store(completed_trace->valid, UINT64_C(1)); + // The preceding barrier publishes trace lines before valid. The store + // drains valid before pending_completed reaches resolved_task_count. } else if (completed_trace != nullptr) { scheduler_publish_cache_line(&completed_trace->kernel_start_cycles); } - ++scheduler->pending_completed; + if (task_complete) ++scheduler->pending_completed; if (direct_refilled != nullptr) *direct_refilled = refilled; return true; } @@ -244,7 +297,8 @@ inline __aicore__ bool scheduler_service_cluster_completions( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, __gm__ SchedulerRunControl *run_control, SchedulerWakeStats *wake_stats, SchedulerReadyStats *ready_stats, SchedulerCompletionStats *completion_stats, uint64_t *ready_victim_cursors, uint64_t profiling_level, - uint64_t *direct_refilled_slot_mask, SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region + uint64_t *direct_refilled_slot_mask, SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, + SchedulerRefillCandidates *candidates = nullptr ) { if (scheduler == nullptr || !scheduler->is_scheduler() || ssbuf_region == nullptr) return false; if (direct_refilled_slot_mask != nullptr) *direct_refilled_slot_mask = 0; @@ -269,7 +323,7 @@ inline __aicore__ bool scheduler_service_cluster_completions( if (!scheduler_service_cluster_completion_slot( graph, scheduler_state_base, scheduler, run_control, cluster_lane, pending_slot, completed_generation, wake_stats, ready_stats, completion_stats, ready_victim_cursors, - profiling_level, nullptr, &direct_refilled, ssbuf_region + profiling_level, nullptr, &direct_refilled, ssbuf_region, candidates )) { return false; } 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 index 7ff4f7a81c..cccbd77fae 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.h @@ -293,7 +293,9 @@ inline __aicore__ bool scheduler_drain_deferred_aiv_to_peer( const int32_t peer_lane = scheduler_deferred_aiv_peer_lane(scheduler_state_base, scheduler); if (peer_lane < 0) return true; const uint64_t peer_worker_id = scheduler->config.worker_ids[static_cast(peer_lane)]; - for (uint32_t pass = 0; pass < 2 && queue->count != 0; ++pass) { + // Mix completions enter the main completion scan; pass 1's direct + // replacement is valid only for ordinary-only graphs. + for (uint32_t pass = 0; pass < (scheduler->has_mix ? 1U : 2U) && queue->count != 0; ++pass) { const uint64_t publication = pass == 1 ? scheduler_ssbuf_load_relaxed( &ssbuf_region->lanes[static_cast(peer_lane)].completion.publication diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_graph.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_graph.h index cdfa3130a1..92a4726ca6 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_graph.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_graph.h @@ -57,7 +57,7 @@ enum class SchedulerGraphResult : uint64_t { inline constexpr bool scheduler_resident_v0_task_shape_supported(uint32_t active_subtasks, uint32_t logical_block_num, bool sync_start) { - return active_subtasks == 1 && logical_block_num == 1 && !sync_start; + return active_subtasks >= 1 && active_subtasks <= 3 && logical_block_num == 1 && !sync_start; } struct SchedulerGraphView { @@ -179,19 +179,22 @@ scheduler_classify_task_shape(const SchedulerGraphView &graph, int64_t task_id, return SchedulerGraphResult::OK; } -inline __aicore__ SchedulerGraphResult scheduler_materialize_task_payload_resolved( - const SchedulerGraphView &graph, const SchedulerTaskInfo &task, uint64_t function_bin_address, - __gm__ DispatchPayload *dispatch_payload, int32_t block_idx = 0, int32_t block_num = 1 +struct SchedulerTaskPayloadArguments { + __gm__ uint8_t *tensors{nullptr}; + __gm__ uint64_t *scalars{nullptr}; + int32_t tensor_count{0}; + int32_t scalar_count{0}; +}; + +inline __aicore__ SchedulerGraphResult scheduler_parse_task_payload_arguments( + const SchedulerGraphView &graph, int64_t task_id, SchedulerTaskPayloadArguments *arguments ) { - if (dispatch_payload == nullptr || function_bin_address == 0 || block_idx < 0 || block_num <= 0 || - block_idx >= block_num) { - return SchedulerGraphResult::INVALID_CALLABLE; - } + if (arguments == nullptr) return SchedulerGraphResult::INVALID_ARGUMENTS; if (graph.reserved != 0) return SchedulerGraphResult::INVALID_ARGUMENTS; - if (graph.storage_address == 0 || task.task_id < 0 || static_cast(task.task_id) >= graph.task_count) { + if (graph.storage_address == 0 || task_id < 0 || static_cast(task_id) >= graph.task_count) { return SchedulerGraphResult::INVALID_TASK_COUNT; } - __gm__ uint8_t *payload = scheduler_graph_payload(graph, task.task_id); + __gm__ uint8_t *payload = scheduler_graph_payload(graph, task_id); int32_t tensor_count = *reinterpret_cast<__gm__ int32_t *>(payload + TASKPAYLOAD_TENSOR_COUNT_OFFSET); int32_t scalar_count = *reinterpret_cast<__gm__ int32_t *>(payload + TASKPAYLOAD_SCALAR_COUNT_OFFSET); if (tensor_count < 0 || tensor_count > MAX_TENSOR_ARGS || scalar_count < 0 || scalar_count > MAX_SCALAR_ARGS || @@ -206,26 +209,54 @@ inline __aicore__ SchedulerGraphResult scheduler_materialize_task_payload_resolv if ((tensor_count > 0 && tensors_delta == 0) || (scalar_count > 0 && scalars_delta == 0)) { return SchedulerGraphResult::INVALID_ARGUMENTS; } + arguments->tensors = tensors_field + tensors_delta; + arguments->scalars = reinterpret_cast<__gm__ uint64_t *>(scalars_field + scalars_delta); + arguments->tensor_count = tensor_count; + arguments->scalar_count = scalar_count; + return SchedulerGraphResult::OK; +} + +inline __aicore__ SchedulerGraphResult scheduler_materialize_task_payload_resolved( + const SchedulerGraphView &graph, const SchedulerTaskInfo &task, uint64_t function_bin_address, + __gm__ DispatchPayload *dispatch_payload, const SchedulerTaskPayloadArguments *shared_arguments = nullptr, + bool write_arguments = true +) { + if (dispatch_payload == nullptr || function_bin_address == 0) { + return SchedulerGraphResult::INVALID_CALLABLE; + } + SchedulerTaskPayloadArguments parsed{}; + if (shared_arguments == nullptr) { + const SchedulerGraphResult status = scheduler_parse_task_payload_arguments(graph, task.task_id, &parsed); + if (status != SchedulerGraphResult::OK) return status; + shared_arguments = &parsed; + } dispatch_payload->function_bin_addr = function_bin_address; - __gm__ uint8_t *tensors = tensors_field + tensors_delta; - __gm__ uint64_t *scalars = reinterpret_cast<__gm__ uint64_t *>(scalars_field + scalars_delta); - int32_t n = 0; - for (int32_t i = 0; i < tensor_count; ++i) { - dispatch_payload->args[n++] = - reinterpret_cast(tensors + static_cast(i) * TASKPAYLOAD_TENSOR_STRIDE); + if (write_arguments) { + int32_t n = 0; + for (int32_t i = 0; i < shared_arguments->tensor_count; ++i) { + dispatch_payload->args[n++] = reinterpret_cast( + shared_arguments->tensors + static_cast(i) * TASKPAYLOAD_TENSOR_STRIDE + ); + } + for (int32_t i = 0; i < shared_arguments->scalar_count; ++i) + dispatch_payload->args[n++] = shared_arguments->scalars[i]; } - for (int32_t i = 0; i < scalar_count; ++i) - dispatch_payload->args[n++] = scalars[i]; - - dispatch_payload->src_payload = 0; - dispatch_payload->local_context.block_idx = block_idx; - dispatch_payload->local_context.block_num = block_num; - // The AICore scheduler has no deferred-completion slab. Mark the context - // non-deferred so async backend adapters take their synchronous fallback. - TaskId::assign(dispatch_payload->local_context.async_ctx.task_token, TaskId::invalid()); - dispatch_payload->args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload->local_context); - dispatch_payload->args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = - reinterpret_cast(&dispatch_payload->global_context); - dispatch_payload->global_context.sub_block_id = task.subtask_slot == 2 ? 1 : 0; + return SchedulerGraphResult::OK; } + +inline __aicore__ void scheduler_materialize_mix_arguments( + const SchedulerTaskPayloadArguments &arguments, __gm__ DispatchPayload *const *targets, uint8_t count +) { + for (int32_t i = 0; i < arguments.tensor_count; ++i) { + const uint64_t value = + reinterpret_cast(arguments.tensors + static_cast(i) * TASKPAYLOAD_TENSOR_STRIDE); + for (uint8_t lane = 0; lane < count; ++lane) + targets[lane]->args[i] = value; + } + for (int32_t i = 0; i < arguments.scalar_count; ++i) { + const uint64_t value = arguments.scalars[i]; + for (uint8_t lane = 0; lane < count; ++lane) + targets[lane]->args[arguments.tensor_count + i] = value; + } +} diff --git a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_mix.h b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_mix.h new file mode 100644 index 0000000000..ce56b65706 --- /dev/null +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_mix.h @@ -0,0 +1,194 @@ +/* + * 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_dispatch.h" + +inline __aicore__ bool scheduler_dispatch_mix_ready( + const SchedulerGraphView &graph, __gm__ void *base, SchedulerLocalState *local, __gm__ SchedulerRunControl *run, + const SchedulerReadyClaim &ready, const SchedulerMixPlacement &placement, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf +) { + if (ready.task_id < 0 || placement.count < 2 || placement.tracker >= SCHEDULER_MIX_TRACKER_COUNT) return false; + const SchedulerTaskMetadata metadata = + scheduler_load_dispatch_metadata(base, local, ready.task_id, profiling_level); + SchedulerTaskPayloadArguments arguments{}; + const bool inline_task = scheduler_task_is_inline(metadata.flags); + if (!inline_task) { + const SchedulerGraphResult status = scheduler_parse_task_payload_arguments(graph, ready.task_id, &arguments); + if (status != SchedulerGraphResult::OK) { + scheduler_record_error( + run, ready.task_id, status, &graph, local, SchedulerErrorSite::DISPATCH_MATERIALIZE_FAILED + ); + return false; + } + } + auto &tracker = local->mix_trackers[placement.tracker]; + tracker.task_id = ready.task_id; + tracker.active_mask = metadata.active_mask; + tracker.completed_mask = 0; + for (uint8_t i = 0; i < placement.count; ++i) { + const auto &claim = placement.slots[i]; + auto &slot = local->slots[claim.cluster_lane][claim.slot_index]; + slot.state = SchedulerDispatchSlotState::FILLING; + slot.mix_tracker = placement.tracker; + } + uint64_t prepare_start_cycles[3]{}; + __gm__ DispatchPayload *payloads[3]{}; + for (uint8_t i = 0; i < placement.count; ++i) { + if (!scheduler_prepare_dispatch_slot( + graph, base, local, run, placement.slots[i], ready, metadata, profiling_level, ssbuf, + placement.subtasks[i], &prepare_start_cycles[i], inline_task ? nullptr : &arguments, false + )) { + for (uint8_t j = 0; j < placement.count; ++j) { + const auto &claim = placement.slots[j]; + scheduler_initialize_free_slot(&local->slots[claim.cluster_lane][claim.slot_index]); + } + tracker.task_id = SCHEDULER_TASK_ID_INVALID; + scheduler_record_error(run, ready.task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, local); + return false; + } + payloads[i] = scheduler_state_at( + base, local->dispatch_payload_offset(placement.slots[i].cluster_lane, placement.slots[i].slot_index) + ); + } + if (!inline_task) scheduler_materialize_mix_arguments(arguments, payloads, placement.count); + for (uint8_t i = 0; i < placement.count; ++i) + scheduler_stage_dispatch_slot(base, local, placement.slots[i], ready, metadata, ssbuf, placement.subtasks[i]); + scheduler_cache_barrier(); + for (uint8_t i = 0; i < placement.count; ++i) + scheduler_commit_dispatch_slot( + base, local, placement.slots[i], ready, metadata, profiling_level, ssbuf, placement.subtasks[i], + prepare_start_cycles[i] + ); + return true; +} + +inline __aicore__ bool scheduler_fill_cluster_mix_slots( + const SchedulerGraphView &graph, __gm__ void *base, SchedulerLocalState *local, __gm__ SchedulerRunControl *run, + uint64_t *victim_cursor, SchedulerReadyStats *stats, uint64_t profiling_level, bool allow_steal, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf, bool *progress +) { + if (!local->has_mix) return true; + for (uint32_t attempt = 0; attempt < SCHEDULER_MIX_TRACKER_COUNT; ++attempt) { + SchedulerMixPlacement placement{}; + SchedulerReadyClaim ready{}; + if (!scheduler_claim_ready_for_slot( + graph, base, local, run, local->config.scheduler_count, SCHEDULER_MIX_QUEUE, victim_cursor, stats, + &ready, &placement, allow_steal + )) + return false; + if (ready.task_id < 0) return true; + if (!scheduler_dispatch_mix_ready(graph, base, local, run, ready, placement, profiling_level, ssbuf)) + return false; + if (progress != nullptr) *progress = true; + } + return true; +} + +inline __aicore__ bool scheduler_publish_direct_mix_candidate( + const SchedulerGraphView &graph, __gm__ void *base, SchedulerLocalState *local, __gm__ SchedulerRunControl *run, + SchedulerReadyClaim *candidate, SchedulerReadyStats *stats, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf, bool *progress +) { + if (candidate == nullptr || candidate->task_id < 0) return true; + // The ready-queue snapshot permits a direct attempt before the next queue + // claim. A candidate without complete placement joins the queue first. + SchedulerMixPlacement placement{}; + if (scheduler_plan_mix_placement(base, local, candidate->task_id, &placement)) { + if (!scheduler_dispatch_mix_ready(graph, base, local, run, *candidate, placement, profiling_level, ssbuf)) + return false; + if (progress != nullptr) *progress = true; + } else { + SchedulerReadyBatch batch{}; + if (!scheduler_ready_batch_append(base, local, candidate->task_id, &batch, stats, profiling_level) || + !scheduler_ready_batch_push(base, local, SCHEDULER_MIX_QUEUE, &batch, stats)) + return false; + } + *candidate = {}; + return true; +} + +inline __aicore__ bool scheduler_fill_mix_after_completions( + const SchedulerGraphView &graph, __gm__ void *base, SchedulerLocalState *local, __gm__ SchedulerRunControl *run, + SchedulerReadyClaim *candidate, uint64_t *victim_cursor, SchedulerReadyStats *stats, uint64_t profiling_level, + bool allow_steal, SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf, bool *progress +) { + if (local->has_mix && !scheduler_publish_direct_mix_candidate( + graph, base, local, run, candidate, stats, profiling_level, ssbuf, progress + )) + return false; + return scheduler_fill_cluster_mix_slots( + graph, base, local, run, victim_cursor, stats, profiling_level, allow_steal, ssbuf, progress + ); +} + +inline __aicore__ bool scheduler_publish_refill_candidates( + const SchedulerGraphView &graph, __gm__ void *base, SchedulerLocalState *local, __gm__ SchedulerRunControl *run, + SchedulerRefillCandidates *candidates, SchedulerReadyStats *stats, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf +) { + SchedulerReadyBatch batches[SCHEDULER_READY_QUEUE_COUNT]{}; + for (uint32_t lane = 0; lane < PLATFORM_CORES_PER_BLOCKDIM; ++lane) { + for (uint32_t index = 0; index < SCHEDULER_PENDING_SLOT_COUNT; ++index) { + const int64_t task = candidates->tasks[lane][index]; + if (task < 0) continue; + const auto &metadata = *scheduler_task_metadata_at(base, local, task); + SchedulerFreeSlotClaim claim{}; + for (uint32_t offset = 0; offset < PLATFORM_CORES_PER_BLOCKDIM && claim.slot_index == UINT32_MAX; + ++offset) { + const uint32_t target_lane = (lane + offset) % PLATFORM_CORES_PER_BLOCKDIM; + const uint64_t worker = local->config.worker_ids[target_lane]; + if (worker >= local->config.runtime_worker_count || target_lane == local->config.self_lane) continue; + const auto *target = scheduler_worker_context_at(base, local, worker); + if (scheduler_core_type_index(target->core_type) != + scheduler_task_ready_queue(metadata.flags, metadata.active_mask)) + continue; + for (uint32_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) { + const auto &state = local->slots[target_lane][slot]; + if (state.state == SchedulerDispatchSlotState::FREE) { + claim = {worker, slot, state.generation, target_lane}; + break; + } + } + } + if (claim.slot_index != UINT32_MAX) { + SchedulerReadyClaim ready{}; + ready.task_id = task; + ready.source = SchedulerReadySource::DIRECT_RESOLVE; + ready.publication_mode = SchedulerPublicationMode::REFILL; + const uint64_t start = scheduler_phase_timing_enabled(profiling_level) ? scheduler_cycles() : 0; + if (!scheduler_fill_dispatch_slot(graph, base, local, run, claim, ready, profiling_level, ssbuf)) + return false; + if (scheduler_phase_timing_enabled(profiling_level)) { + auto *traces = scheduler_state_at(base, local->profiling->trace_cells_offset); + auto &completed = traces[candidates->completed_trace_indices[lane][index]]; + completed.refill_scheduler_worker_id = local->worker_id(); + completed.refill_start_cycles = start; + completed.refill_end_cycles = scheduler_cycles(); + completed.refill_task_id = static_cast(task); + completed.refill_loop_iter = local->profiling->loop_iter; + scheduler_publish_cache_line(&completed.refill_scheduler_worker_id); + } + } else { + if (!scheduler_ready_batch_append( + base, local, task, &batches[scheduler_task_ready_queue(metadata.flags, metadata.active_mask)], + stats, 0 + )) + return false; + } + candidates->tasks[lane][index] = SCHEDULER_TASK_ID_INVALID; + } + } + for (uint32_t type = 0; type < SCHEDULER_READY_QUEUE_COUNT; ++type) + if (!scheduler_ready_batch_push(base, local, type, &batches[type], stats)) return false; + return true; +} 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 index 43f1102228..bb08dfe6b9 100644 --- a/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_ready.h +++ b/src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_ready.h @@ -102,6 +102,7 @@ struct SchedulerLocalSlotState { uint32_t generation{0}; SchedulerDispatchSlotState state{SchedulerDispatchSlotState::EMPTY}; uint8_t subtask_slot{UINT8_MAX}; + uint8_t mix_tracker{UINT8_MAX}; }; struct SchedulerWorkerTraceCache { @@ -154,14 +155,24 @@ struct SchedulerLocalProfilingState { } }; +inline constexpr uint32_t SCHEDULER_MIX_TRACKER_COUNT = PLATFORM_CORES_PER_BLOCKDIM * SCHEDULER_PENDING_SLOT_COUNT; +struct SchedulerMixTracker { + int64_t task_id{SCHEDULER_TASK_ID_INVALID}; + uint8_t active_mask{0}; + uint8_t completed_mask{0}; +}; + struct SchedulerLocalState { SchedulerLocalConfig config{}; SchedulerLocalProfilingState *profiling{nullptr}; uint64_t pending_completed{0}; + SchedulerMixTracker mix_trackers[SCHEDULER_MIX_TRACKER_COUNT]{}; + uint32_t dispatch_sequences[PLATFORM_CORES_PER_BLOCKDIM]{}; + bool has_mix{false}; inline __aicore__ explicit SchedulerLocalState(SchedulerLocalProfilingState *profile = nullptr) : profiling(profile) { - for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) + for (uint32_t type = 0; type < SCHEDULER_READY_QUEUE_COUNT; ++type) owner_pending_endpoints[type] = SCHEDULER_READY_PENDING_EMPTY; } inline __aicore__ bool is_scheduler() const { @@ -190,7 +201,7 @@ struct SchedulerLocalState { SchedulerLocalSlotState slots[PLATFORM_CORES_PER_BLOCKDIM][SCHEDULER_PENDING_SLOT_COUNT]{}; uint64_t consumed_completion_generations[PLATFORM_CORES_PER_BLOCKDIM]{}; uint32_t local_completed_generations[SCHEDULER_PENDING_SLOT_COUNT]{}; - uint64_t owner_pending_endpoints[SCHEDULER_CORE_TYPE_COUNT]{}; + uint64_t owner_pending_endpoints[SCHEDULER_READY_QUEUE_COUNT]{}; uint8_t local_ready_mask{0}; uint8_t owner_ready_queue_mask{0}; }; @@ -199,12 +210,12 @@ static_assert(SCHEDULER_WORKER_CAPACITY < UINT16_MAX, "worker IDs must fit the l static_assert(SCHEDULER_CAPACITY < UINT16_MAX, "scheduler IDs must fit the local config"); static_assert(PLATFORM_CORES_PER_BLOCKDIM <= 8, "worker trace mask must fit one byte"); static_assert(SCHEDULER_PENDING_SLOT_COUNT <= 8, "local ready mask must fit one byte"); -static_assert(SCHEDULER_CORE_TYPE_COUNT <= 8, "owner ready mask must fit one byte"); +static_assert(SCHEDULER_READY_QUEUE_COUNT <= 8, "owner ready mask must fit one byte"); static_assert(alignof(SchedulerLocalSlotState) == alignof(uint64_t)); static_assert(alignof(SchedulerLocalState) == alignof(uint64_t)); static_assert(sizeof(void *) != 8 || sizeof(SchedulerLocalConfig) == 88, "64-bit local config size changed"); static_assert(sizeof(void *) != 8 || sizeof(SchedulerLocalSlotState) == 16, "64-bit local slot size changed"); -static_assert(sizeof(void *) != 8 || sizeof(SchedulerLocalState) == 256, "64-bit local state size changed"); +static_assert(sizeof(void *) != 8 || sizeof(SchedulerLocalState) == 376, "64-bit local state size changed"); static_assert( sizeof(void *) != 8 || sizeof(SchedulerLocalProfilingState) == 240, "64-bit local profiling size changed" ); @@ -249,6 +260,9 @@ inline __aicore__ bool scheduler_initialize_local_config( config.callable_addresses_address = context->callable_addresses_address; config.callable_addresses_count = static_cast(context->callable_addresses_count); config.gang_coordinator_offset = context->gang_coordinator_offset; + auto *coordinator = scheduler_state_at(base, config.gang_coordinator_offset); + scheduler_observe_cache_line(coordinator); + local->has_mix = coordinator->mix_task_count != 0; const uint64_t runtime_worker_count = context->runtime_worker_count; const uint64_t scheduler_count = context->scheduler_count; const uint64_t scheduler_index = context->scheduler_index; @@ -317,21 +331,25 @@ inline __aicore__ void scheduler_local_ready_publish(SchedulerLocalState *schedu scheduler_local_state->local_ready_mask |= UINT32_C(1) << slot_index; } -inline __aicore__ bool scheduler_local_ready_pop( - SchedulerLocalState *scheduler_local_state, uint32_t scan_start, uint32_t *slot_index, uint64_t *publication -) { +inline __aicore__ bool +scheduler_local_ready_pop(SchedulerLocalState *scheduler_local_state, uint32_t *slot_index, uint64_t *publication) { if (slot_index == nullptr || publication == nullptr || scheduler_local_state->local_ready_mask == 0 || scheduler_local_state->config.self_lane >= PLATFORM_CORES_PER_BLOCKDIM) return false; - for (uint32_t offset = 0; offset < SCHEDULER_PENDING_SLOT_COUNT; ++offset) { - const uint32_t selected = (scan_start + offset) % SCHEDULER_PENDING_SLOT_COUNT; - const uint32_t selected_mask = UINT32_C(1) << selected; - if ((scheduler_local_state->local_ready_mask & selected_mask) == 0) continue; - scheduler_local_state->local_ready_mask &= ~selected_mask; - *slot_index = selected; + uint32_t oldest = UINT32_MAX; + for (uint32_t selected = 0; selected < SCHEDULER_PENDING_SLOT_COUNT; ++selected) { + if ((scheduler_local_state->local_ready_mask & (1U << selected)) == 0) continue; const auto &slot = scheduler_local_state->slots[scheduler_local_state->config.self_lane][selected]; - // A pending notification owns this READY slot until the same core consumes it. - // Preserve an invalid state in the token so the Executor rejects it. + if (oldest == UINT32_MAX || + slot.generation < scheduler_local_state->slots[scheduler_local_state->config.self_lane][oldest].generation) + oldest = selected; + } + if (oldest != UINT32_MAX) { + scheduler_local_state->local_ready_mask &= ~(1U << oldest); + *slot_index = oldest; + const auto &slot = scheduler_local_state->slots[scheduler_local_state->config.self_lane][oldest]; + // A local notification owns this slot until it is consumed here. Keep + // an invalid state in the token so the Executor rejects it. *publication = (static_cast(slot.generation) << 8) | static_cast(slot.state); return true; } @@ -340,13 +358,13 @@ inline __aicore__ bool scheduler_local_ready_pop( inline __aicore__ void scheduler_owner_queue_activate(SchedulerLocalState *scheduler_local_state, uint32_t core_type_index) { - if (core_type_index >= SCHEDULER_CORE_TYPE_COUNT) return; + if (core_type_index >= SCHEDULER_READY_QUEUE_COUNT) return; scheduler_local_state->owner_ready_queue_mask |= UINT32_C(1) << core_type_index; } inline __aicore__ void scheduler_owner_queue_deactivate(SchedulerLocalState *scheduler_local_state, uint32_t core_type_index) { - if (core_type_index >= SCHEDULER_CORE_TYPE_COUNT) return; + if (core_type_index >= SCHEDULER_READY_QUEUE_COUNT) return; scheduler_local_state->owner_ready_queue_mask &= ~(UINT32_C(1) << core_type_index); } @@ -694,6 +712,12 @@ inline __aicore__ SchedulerRouteResult scheduler_bootstrap_route_task( return SchedulerRouteResult::READY_TO_ENQUEUE; } +inline __aicore__ uint32_t +scheduler_trace_index_at(__gm__ void *base, SchedulerLocalState *local, int64_t task_id, uint8_t subtask) { + const auto *metadata = scheduler_task_metadata_at(base, local, task_id); + return scheduler_task_trace_index(metadata->trace_index_base, metadata->active_mask, subtask); +} + inline __aicore__ bool scheduler_bootstrap_ready_batch_append( __gm__ void *scheduler_state_base, SchedulerLocalState *context, int64_t task_id, SchedulerReadyBatch *batch, SchedulerReadyStats *stats, uint64_t profiling_level = 0 @@ -715,8 +739,14 @@ inline __aicore__ bool scheduler_bootstrap_ready_batch_append( if (scheduler_phase_timing_enabled(profiling_level)) { __gm__ SchedulerTaskTrace *cells = scheduler_state_at(scheduler_state_base, context->profiling->trace_cells_offset); - cells[task_id].ready_transition_cycles = scheduler_cycles(); - scheduler_writeback_cache_line(&cells[task_id].ready_transition_cycles); + const auto &metadata = *scheduler_task_metadata_at(scheduler_state_base, context, task_id); + const uint64_t now = scheduler_cycles(); + for (uint8_t subtask = 0; subtask < 3; ++subtask) { + if ((metadata.active_mask & (1U << subtask)) == 0) continue; + auto &trace = cells[scheduler_task_trace_index(metadata.trace_index_base, metadata.active_mask, subtask)]; + trace.ready_transition_cycles = now; + scheduler_writeback_cache_line(&trace.ready_transition_cycles); + } } ++batch->count; if (stats != nullptr) ++stats->enqueue_count; @@ -728,7 +758,7 @@ inline __aicore__ bool scheduler_bootstrap_ready_batch_publish( 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_CAPACITY || batch->tail < 0 || + if (core_type_index >= SCHEDULER_READY_QUEUE_COUNT || inbox_index >= SCHEDULER_CAPACITY || batch->tail < 0 || ready_types == nullptr) return false; __gm__ SchedulerReadyInbox *inbox = @@ -752,7 +782,8 @@ inline __aicore__ bool scheduler_bootstrap_ready_directory_publish( uint32_t shard_count = static_cast( (scheduler_count + SCHEDULER_READY_DIRECTORY_OWNERS_PER_SHARD - 1) / SCHEDULER_READY_DIRECTORY_OWNERS_PER_SHARD ); - for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + const uint32_t queue_count = context->has_mix ? SCHEDULER_READY_QUEUE_COUNT : SCHEDULER_CORE_TYPE_COUNT; + for (uint32_t type = 0; type < queue_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_OWNERS_PER_SHARD; @@ -791,10 +822,17 @@ inline __aicore__ bool scheduler_ready_batch_append( if (scheduler_phase_timing_enabled(profiling_level)) { __gm__ SchedulerTaskTrace *cells = scheduler_state_at(scheduler_state_base, context->profiling->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); + const auto *metadata = scheduler_task_metadata_at(scheduler_state_base, context, task_id); + const uint8_t mask = metadata->active_mask; + const uint32_t trace_base = metadata->trace_index_base; + const uint64_t transition_cycles = scheduler_cycles(); + for (uint8_t subtask = 0; subtask < 3; ++subtask) { + if ((mask & (1U << subtask)) == 0) continue; + __gm__ SchedulerTaskTrace *trace = &cells[scheduler_task_trace_index(trace_base, mask, subtask)]; + scheduler_observe_cache_line(&trace->ready_transition_cycles); + trace->ready_transition_cycles = transition_cycles; + scheduler_publish_cache_line(&trace->ready_transition_cycles); + } } ++batch->count; if (stats != nullptr) ++stats->enqueue_count; @@ -864,9 +902,10 @@ inline __aicore__ bool scheduler_ready_owner_pending_append( // Publish the owner's pending list when its shared inbox drains, and keep // the directory bit consistent so other Schedulers can discover stealable work. inline __aicore__ bool scheduler_refresh_ready_inbox_type( - __gm__ void *scheduler_state_base, SchedulerLocalState *context, uint32_t core_type_index + __gm__ void *scheduler_state_base, SchedulerLocalState *context, uint32_t core_type_index, + int64_t *head_snapshot = nullptr ) { - if (context == nullptr || core_type_index >= SCHEDULER_CORE_TYPE_COUNT || + if (context == nullptr || core_type_index >= SCHEDULER_READY_QUEUE_COUNT || context->config.scheduler_index >= SCHEDULER_CAPACITY) return false; uint64_t *pending_endpoints = &context->owner_pending_endpoints[core_type_index]; @@ -880,6 +919,7 @@ inline __aicore__ bool scheduler_refresh_ready_inbox_type( scheduler_ready_directory_set(directory, core_type_index, context->config.scheduler_index); scheduler_owner_queue_activate(context, core_type_index); } + if (head_snapshot != nullptr) *head_snapshot = head; return true; } const uint64_t pending = scheduler_ready_owner_pending_load(pending_endpoints); @@ -895,6 +935,7 @@ inline __aicore__ bool scheduler_refresh_ready_inbox_type( scheduler_ready_directory_set(directory, core_type_index, context->config.scheduler_index); scheduler_owner_queue_activate(context, core_type_index); } + if (head_snapshot != nullptr) *head_snapshot = pending_head; return true; } else if (pending_tail != SCHEDULER_INBOX_EMPTY) { return false; @@ -903,12 +944,13 @@ inline __aicore__ bool scheduler_refresh_ready_inbox_type( scheduler_ready_directory_clear(directory, core_type_index, context->config.scheduler_index); scheduler_owner_queue_deactivate(context, core_type_index); } + if (head_snapshot != nullptr) *head_snapshot = SCHEDULER_INBOX_EMPTY; return true; } inline __aicore__ bool scheduler_refresh_ready_inbox(__gm__ void *scheduler_state_base, SchedulerLocalState *context) { const uint32_t active_mask = context->owner_ready_queue_mask; - for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + for (uint32_t type = 0; type < SCHEDULER_READY_QUEUE_COUNT; ++type) { if ((active_mask & (UINT32_C(1) << type)) == 0) continue; if (!scheduler_refresh_ready_inbox_type(scheduler_state_base, context, type)) return false; } @@ -919,7 +961,7 @@ inline __aicore__ bool scheduler_ready_batch_push( __gm__ void *scheduler_state_base, SchedulerLocalState *context, uint32_t core_type_index, SchedulerReadyBatch *batch, SchedulerReadyStats *stats ) { - if (batch == nullptr || context == nullptr || core_type_index >= SCHEDULER_CORE_TYPE_COUNT || + if (batch == nullptr || context == nullptr || core_type_index >= SCHEDULER_READY_QUEUE_COUNT || context->config.scheduler_index >= SCHEDULER_CAPACITY) return false; if (batch->head == SCHEDULER_INBOX_EMPTY) return true; @@ -976,17 +1018,69 @@ inline __aicore__ bool scheduler_ready_batch_push( return true; } +struct SchedulerMixPlacement { + SchedulerFreeSlotClaim slots[3]{}; + uint8_t subtasks[3]{}; + uint8_t count{0}; + uint8_t tracker{UINT8_MAX}; +}; + +inline __aicore__ bool scheduler_plan_mix_placement( + __gm__ void *base, SchedulerLocalState *local, int64_t task_id, SchedulerMixPlacement *placement +) { + *placement = {}; + const auto &metadata = *scheduler_task_metadata_at(base, local, task_id); + if (!scheduler_task_is_mix(metadata.flags) || scheduler_task_is_gang(metadata.flags)) return false; + for (uint8_t i = 0; i < SCHEDULER_MIX_TRACKER_COUNT; ++i) { + if (local->mix_trackers[i].task_id < 0) { + placement->tracker = i; + break; + } + } + if (placement->tracker == UINT8_MAX) return false; + uint8_t used_lanes = 0; + for (uint8_t subtask = 0; subtask < 3; ++subtask) { + if ((metadata.active_mask & (1U << subtask)) == 0) continue; + bool found = false; + const bool single_aiv = (metadata.active_mask & 6U) != 6U; + const uint8_t pass_count = subtask != 0 && single_aiv ? 2 : 1; + for (uint8_t pass = 0; pass < pass_count && !found; ++pass) { + for (uint8_t lane = 0; lane < PLATFORM_CORES_PER_BLOCKDIM && !found; ++lane) { + if ((used_lanes & (1U << lane)) != 0) continue; + const uint64_t worker = local->config.worker_ids[lane]; + if (worker >= local->config.runtime_worker_count) continue; + const auto *target = scheduler_worker_context_at(base, local, worker); + if (target->core_type != static_cast(subtask == 0 ? CoreType::AIC : CoreType::AIV)) continue; + if (subtask != 0 && !single_aiv && lane != subtask) continue; + if (subtask != 0 && single_aiv && ((lane == local->config.self_lane) != (pass == 1))) continue; + for (uint8_t slot = 0; slot < SCHEDULER_PENDING_SLOT_COUNT; ++slot) { + const auto &state = local->slots[lane][slot]; + if (state.state != SchedulerDispatchSlotState::FREE) continue; + placement->slots[placement->count] = {worker, slot, state.generation, lane}; + placement->subtasks[placement->count++] = subtask; + used_lanes |= 1U << lane; + found = true; + break; + } + } + } + if (!found) return false; + } + return placement->count >= 2; +} + inline __aicore__ bool scheduler_ready_pop_from_inbox( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *context, __gm__ SchedulerRunControl *run_control, uint32_t core_type_index, uint64_t inbox_index, int64_t *task_id, - SchedulerReadyStats *stats + SchedulerReadyStats *stats, SchedulerMixPlacement *placement = nullptr, const int64_t *first_head_snapshot = nullptr ) { 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); + int64_t head = + attempt == 0 && first_head_snapshot != nullptr ? *first_head_snapshot : scheduler_gm_query(inbox->head); if (head == SCHEDULER_INBOX_EMPTY) return true; if (head < 0 || static_cast(head) >= graph.task_count) { scheduler_record_error( @@ -995,6 +1089,8 @@ inline __aicore__ bool scheduler_ready_pop_from_inbox( ); return false; } + if (placement != nullptr && !scheduler_plan_mix_placement(scheduler_state_base, context, head, placement)) + return true; __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, head); int64_t next = scheduler_observe_next_waiter(control); if (next < SCHEDULER_INBOX_EMPTY) { @@ -1032,7 +1128,8 @@ inline __aicore__ uint64_t scheduler_load_ready_directory_shard( inline __aicore__ bool scheduler_steal_ready_from_shard( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *context, __gm__ SchedulerRunControl *run_control, uint32_t core_type_index, uint64_t shard_begin, uint64_t shard_end, - uint64_t start, uint64_t bits, SchedulerReadyStats *stats, SchedulerReadyClaim *claim + uint64_t start, uint64_t bits, SchedulerReadyStats *stats, SchedulerReadyClaim *claim, + SchedulerMixPlacement *placement = nullptr ) { int64_t task_id = SCHEDULER_TASK_ID_INVALID; bits &= ~(UINT64_C(1) << (context->config.scheduler_index - shard_begin)); @@ -1048,7 +1145,8 @@ inline __aicore__ bool scheduler_steal_ready_from_shard( candidates &= candidates - 1; uint64_t victim = shard_begin + bit_index; if (!scheduler_ready_pop_from_inbox( - graph, scheduler_state_base, context, run_control, core_type_index, victim, &task_id, stats + graph, scheduler_state_base, context, run_control, core_type_index, victim, &task_id, stats, + placement )) return false; if (task_id >= 0) { @@ -1066,18 +1164,23 @@ inline __aicore__ bool scheduler_steal_ready_from_shard( inline __aicore__ bool scheduler_claim_ready_for_slot( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *context, __gm__ SchedulerRunControl *run_control, uint64_t scheduler_count, uint32_t core_type_index, - uint64_t *victim_cursor, SchedulerReadyStats *stats, SchedulerReadyClaim *claim + uint64_t *victim_cursor, SchedulerReadyStats *stats, SchedulerReadyClaim *claim, + SchedulerMixPlacement *placement = nullptr, bool allow_steal = true ) { if (victim_cursor == nullptr || claim == nullptr || context == nullptr || scheduler_count == 0 || scheduler_count > SCHEDULER_CAPACITY || context->config.scheduler_index >= scheduler_count || - core_type_index >= SCHEDULER_CORE_TYPE_COUNT) + core_type_index >= SCHEDULER_READY_QUEUE_COUNT) return false; *claim = {}; - if (!scheduler_refresh_ready_inbox_type(scheduler_state_base, context, core_type_index)) return false; + int64_t owner_head_snapshot = SCHEDULER_INBOX_EMPTY; + if (!scheduler_refresh_ready_inbox_type(scheduler_state_base, context, core_type_index, &owner_head_snapshot)) + return false; + // This Scheduler is the only publisher of its empty inbox. A thief can + // change a nonempty head, and the pop CAS retries against that change. 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->config.scheduler_index, - &task_id, stats + &task_id, stats, placement, &owner_head_snapshot )) return false; if (task_id >= 0) { @@ -1086,6 +1189,7 @@ inline __aicore__ bool scheduler_claim_ready_for_slot( return true; } + if (!allow_steal) return true; __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); uint64_t shard_begin = context->config.scheduler_index / SCHEDULER_READY_DIRECTORY_OWNERS_PER_SHARD * SCHEDULER_READY_DIRECTORY_OWNERS_PER_SHARD; @@ -1098,7 +1202,7 @@ inline __aicore__ bool scheduler_claim_ready_for_slot( ); if (bits != 0 && !scheduler_steal_ready_from_shard( graph, scheduler_state_base, context, run_control, core_type_index, shard_begin, shard_end, - start, bits, stats, claim + start, bits, stats, claim, placement )) return false; const uint64_t cursor_base = claim->task_id >= 0 ? claim->inbox_index : start; @@ -1110,7 +1214,7 @@ inline __aicore__ bool scheduler_ready_directory_nonempty( __gm__ void *scheduler_state_base, SchedulerLocalState *context, uint64_t scheduler_count, uint32_t core_type_index ) { if (scheduler_count == 0 || scheduler_count > SCHEDULER_CAPACITY || - context->config.scheduler_index >= scheduler_count || core_type_index >= SCHEDULER_CORE_TYPE_COUNT) + context->config.scheduler_index >= scheduler_count || core_type_index >= SCHEDULER_READY_QUEUE_COUNT) return false; __gm__ SchedulerReadyDirectory *directory = scheduler_ready_directory_at(scheduler_state_base, context); return scheduler_load_ready_directory_shard( @@ -1119,43 +1223,45 @@ inline __aicore__ bool scheduler_ready_directory_nonempty( } inline __aicore__ void scheduler_initialize_free_slot(SchedulerLocalSlotState *local_slot) { - uint32_t generation = local_slot->generation + 1; - if (generation == 0) generation = 1; local_slot->task_id = SCHEDULER_TASK_ID_INVALID; - local_slot->generation = generation; local_slot->state = SchedulerDispatchSlotState::FREE; + local_slot->mix_tracker = UINT8_MAX; local_slot->subtask_slot = UINT8_MAX; } -inline __aicore__ bool scheduler_fill_dispatch_slot( +inline __aicore__ SchedulerTaskMetadata scheduler_load_dispatch_metadata( + __gm__ void *base, SchedulerLocalState *scheduler, int64_t task_id, uint64_t profiling_level +) { + const auto *source = scheduler_task_metadata_at(base, scheduler, task_id); + SchedulerTaskMetadata metadata{}; + metadata.kernel_ids[0] = source->kernel_ids[0]; + metadata.kernel_ids[1] = source->kernel_ids[1]; + metadata.kernel_ids[2] = source->kernel_ids[2]; + metadata.active_mask = source->active_mask; + metadata.flags = source->flags; + metadata.timing_slot = source->timing_slot; + if (scheduler_task_timing_enabled(profiling_level)) metadata.trace_index_base = source->trace_index_base; + return metadata; +} + +inline __aicore__ bool scheduler_prepare_dispatch_slot( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, __gm__ SchedulerRunControl *run_control, const SchedulerFreeSlotClaim &slot_claim, - const SchedulerReadyClaim &ready_claim, uint64_t profiling_level, SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region + const SchedulerReadyClaim &ready_claim, const SchedulerTaskMetadata &metadata, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, uint8_t subtask_slot, uint64_t *prepare_start_cycles = nullptr, + const SchedulerTaskPayloadArguments *arguments = nullptr, bool write_arguments = true ) { if (ssbuf_region == nullptr || scheduler == nullptr || ready_claim.task_id < 0 || static_cast(ready_claim.task_id) >= graph.task_count || slot_claim.worker_id >= scheduler->config.runtime_worker_count || slot_claim.slot_index >= SCHEDULER_PENDING_SLOT_COUNT || slot_claim.cluster_lane >= PLATFORM_CORES_PER_BLOCKDIM) return false; - const bool task_timing_enabled = scheduler_task_timing_enabled(profiling_level); - const bool schedule_timing_enabled = scheduler_schedule_timing_enabled(profiling_level); - const bool phase_timing_enabled = scheduler_phase_timing_enabled(profiling_level); - const uint64_t dispatch_start_cycles = phase_timing_enabled ? scheduler_cycles() : 0; - __gm__ SchedulerTaskMetadata *metadata_source = - scheduler_task_metadata_at(scheduler_state_base, scheduler, ready_claim.task_id); - 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); + if (prepare_start_cycles != nullptr) + *prepare_start_cycles = scheduler_phase_timing_enabled(profiling_level) ? scheduler_cycles() : 0; + if (subtask_slot == UINT8_MAX) subtask_slot = scheduler_metadata_single_subtask_slot(metadata.active_mask); __gm__ SchedulerWorkerContext *target = scheduler_worker_context_at(scheduler_state_base, scheduler, slot_claim.worker_id); - if (subtask_slot == UINT8_MAX || + if (subtask_slot >= 3 || (metadata.active_mask & (1U << subtask_slot)) == 0 || (target->core_type != static_cast(CoreType::AIC) && target->core_type != static_cast(CoreType::AIV)) || !scheduler_task_is_executable(metadata.flags) || scheduler_task_is_gang(metadata.flags) || @@ -1167,8 +1273,11 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( return false; } const uint16_t kernel_id = metadata.kernel_ids[subtask_slot]; - uint32_t generation = slot_claim.generation + 1; - if (generation == 0) generation = 1; + if (scheduler->dispatch_sequences[slot_claim.cluster_lane] == UINT32_MAX) return false; + auto &reserved = scheduler->slots[slot_claim.cluster_lane][slot_claim.slot_index]; + if (reserved.state != SchedulerDispatchSlotState::FILLING || reserved.generation != slot_claim.generation || + scheduler->config.worker_ids[slot_claim.cluster_lane] != slot_claim.worker_id) + return false; __gm__ uint64_t *callable_addresses = reinterpret_cast<__gm__ uint64_t *>(scheduler->config.callable_addresses_address); const bool inline_task = scheduler_task_is_inline(metadata.flags); @@ -1199,7 +1308,9 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( payload->function_bin_addr = 0; payload->src_payload = 0; } else { - status = scheduler_materialize_task_payload_resolved(graph, task, callable_address, payload); + status = scheduler_materialize_task_payload_resolved( + graph, task, callable_address, payload, arguments, write_arguments + ); 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) { @@ -1218,6 +1329,20 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( ); return false; } + payload->global_context.sub_block_id = + scheduler_task_is_mix(metadata.flags) ? (slot_claim.cluster_lane == 2 ? 1 : 0) : (subtask_slot == 2 ? 1 : 0); + return true; +} + +inline __aicore__ void scheduler_stage_dispatch_slot( + __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, const SchedulerFreeSlotClaim &slot_claim, + const SchedulerReadyClaim &ready_claim, const SchedulerTaskMetadata &metadata, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, uint8_t subtask_slot +) { + const uint32_t generation = ++scheduler->dispatch_sequences[slot_claim.cluster_lane]; + auto *payload = scheduler_state_at( + scheduler_state_base, scheduler->dispatch_payload_offset(slot_claim.cluster_lane, slot_claim.slot_index) + ); SchedulerLocalSlotState *local_slot = &scheduler->slots[slot_claim.cluster_lane][slot_claim.slot_index]; local_slot->task_id = ready_claim.task_id; scheduler->set_timing_slot(slot_claim.cluster_lane, slot_claim.slot_index, metadata.timing_slot); @@ -1229,7 +1354,21 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( &ssbuf_region->lanes[slot_claim.cluster_lane].dispatch[slot_claim.slot_index]; if (remote) dispatch_control->task_id = ready_claim.task_id; scheduler_writeback_dispatch_payload(payload); - scheduler_cache_barrier(); +} + +inline __aicore__ void scheduler_commit_dispatch_slot( + __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, const SchedulerFreeSlotClaim &slot_claim, + const SchedulerReadyClaim &ready_claim, const SchedulerTaskMetadata &metadata, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, uint8_t subtask_slot, uint64_t prepare_start_cycles +) { + const bool task_timing_enabled = scheduler_task_timing_enabled(profiling_level); + const bool schedule_timing_enabled = scheduler_schedule_timing_enabled(profiling_level); + const bool phase_timing_enabled = scheduler_phase_timing_enabled(profiling_level); + const uint64_t dispatch_start_cycles = prepare_start_cycles; + const uint32_t generation = scheduler->slots[slot_claim.cluster_lane][slot_claim.slot_index].generation; + const bool remote = slot_claim.worker_id != scheduler->worker_id(); + SCHEDULER_SSBUF SchedulerSsbufDispatchControl *dispatch_control = + &ssbuf_region->lanes[slot_claim.cluster_lane].dispatch[slot_claim.slot_index]; const uint64_t ready_publish_cycles = schedule_timing_enabled ? scheduler_cycles() : 0; if (remote) { scheduler_ssbuf_store_relaxed( @@ -1241,7 +1380,8 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( if (task_timing_enabled) { __gm__ SchedulerTaskTrace *traces = scheduler_state_at(scheduler_state_base, scheduler->profiling->trace_cells_offset); - __gm__ SchedulerTaskTrace *trace = &traces[ready_claim.task_id]; + __gm__ SchedulerTaskTrace *trace = + &traces[scheduler_task_trace_index(metadata.trace_index_base, metadata.active_mask, subtask_slot)]; trace->worker_id = slot_claim.worker_id; trace->task_id = static_cast(ready_claim.task_id); if (phase_timing_enabled) { @@ -1266,6 +1406,47 @@ inline __aicore__ bool scheduler_fill_dispatch_slot( if (schedule_timing_enabled) trace->dispatch_end_cycles = ready_publish_cycles; if (schedule_timing_enabled) scheduler_publish_cache_line(&trace->dispatch_start_cycles); } +} + +inline __aicore__ void scheduler_publish_dispatch_slot( + __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, const SchedulerFreeSlotClaim &slot_claim, + const SchedulerReadyClaim &ready_claim, const SchedulerTaskMetadata &metadata, uint64_t profiling_level, + SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region, uint8_t subtask_slot, uint64_t prepare_start_cycles = 0 +) { + scheduler_stage_dispatch_slot( + scheduler_state_base, scheduler, slot_claim, ready_claim, metadata, ssbuf_region, subtask_slot + ); + scheduler_cache_barrier(); + scheduler_commit_dispatch_slot( + scheduler_state_base, scheduler, slot_claim, ready_claim, metadata, profiling_level, ssbuf_region, subtask_slot, + prepare_start_cycles + ); +} + +inline __aicore__ bool scheduler_fill_dispatch_slot( + const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *scheduler, + __gm__ SchedulerRunControl *run_control, const SchedulerFreeSlotClaim &slot_claim, + const SchedulerReadyClaim &ready_claim, uint64_t profiling_level, SCHEDULER_SSBUF SchedulerSsbufRegion *ssbuf_region +) { + if (scheduler == nullptr || slot_claim.cluster_lane >= PLATFORM_CORES_PER_BLOCKDIM || + slot_claim.slot_index >= SCHEDULER_PENDING_SLOT_COUNT) + return false; + auto &slot = scheduler->slots[slot_claim.cluster_lane][slot_claim.slot_index]; + if (slot.state == SchedulerDispatchSlotState::FREE && slot.generation == slot_claim.generation) + slot.state = SchedulerDispatchSlotState::FILLING; + const SchedulerTaskMetadata metadata = + scheduler_load_dispatch_metadata(scheduler_state_base, scheduler, ready_claim.task_id, profiling_level); + const uint8_t subtask_slot = scheduler_metadata_single_subtask_slot(metadata.active_mask); + uint64_t prepare_start_cycles = 0; + if (!scheduler_prepare_dispatch_slot( + graph, scheduler_state_base, scheduler, run_control, slot_claim, ready_claim, metadata, profiling_level, + ssbuf_region, subtask_slot, &prepare_start_cycles + )) + return false; + scheduler_publish_dispatch_slot( + scheduler_state_base, scheduler, slot_claim, ready_claim, metadata, profiling_level, ssbuf_region, subtask_slot, + prepare_start_cycles + ); return true; } @@ -1273,7 +1454,8 @@ inline __aicore__ bool scheduler_resolve_completion( const SchedulerGraphView &graph, __gm__ void *scheduler_state_base, SchedulerLocalState *context, __gm__ SchedulerRunControl *run_control, int64_t task_id, SchedulerWakeStats *wake_stats, SchedulerReadyStats *ready_stats, SchedulerCompletionStats *completion_stats, uint64_t profiling_level, - bool validate_done_state, SchedulerReadyClaim *direct_ready, uint32_t direct_core_type + bool validate_done_state, SchedulerReadyClaim *direct_ready, uint32_t direct_core_type, + SchedulerReadyClaim *direct_mix_ready = nullptr ) { if (context == nullptr) return false; __gm__ SchedulerTaskControl *control = scheduler_task_control_at(scheduler_state_base, context, task_id); @@ -1301,7 +1483,7 @@ inline __aicore__ bool scheduler_resolve_completion( return false; } if (wake_stats != nullptr) ++wake_stats->wake_close_count; - SchedulerReadyBatch batches[SCHEDULER_CORE_TYPE_COUNT]{}; + SchedulerReadyBatch batches[SCHEDULER_READY_QUEUE_COUNT]{}; while (waiter >= 0) { if (static_cast(waiter) >= graph.task_count) { scheduler_record_error( @@ -1336,24 +1518,59 @@ inline __aicore__ bool scheduler_resolve_completion( scheduler_publish_gang_ready(scheduler_state_base, context, metadata->flags); } else { const uint8_t subtask_slot = scheduler_metadata_single_subtask_slot(metadata->active_mask); - if (subtask_slot == UINT8_MAX) { + if (subtask_slot == UINT8_MAX && !scheduler_task_is_mix(metadata->flags)) { scheduler_record_error( run_control, waiter, SchedulerGraphResult::UNSUPPORTED_SHAPE, &graph, context, SchedulerErrorSite::COMPLETION_INVALID_SHAPE ); return false; } - const uint32_t core_type = scheduler_metadata_core_type_index(subtask_slot); - if (direct_ready != nullptr && direct_ready->task_id < 0 && core_type == direct_core_type) { - direct_ready->task_id = waiter; - direct_ready->source = SchedulerReadySource::DIRECT_RESOLVE; - direct_ready->publication_mode = SchedulerPublicationMode::REFILL; + const uint32_t core_type = scheduler_task_ready_queue(metadata->flags, metadata->active_mask); + if (core_type == SCHEDULER_MIX_QUEUE && direct_mix_ready != nullptr && direct_mix_ready->task_id >= 0) { + if (!scheduler_ready_batch_append( + scheduler_state_base, context, direct_mix_ready->task_id, &batches[core_type], ready_stats, + profiling_level + )) + return false; + *direct_mix_ready = {}; + } + bool direct_mix_available = false; + if (core_type == SCHEDULER_MIX_QUEUE && direct_mix_ready != nullptr && + batches[core_type].head == SCHEDULER_INBOX_EMPTY) { + __gm__ SchedulerReadyInbox *inbox = scheduler_ready_inbox_at( + scheduler_state_base, context, core_type, context->config.scheduler_index + ); + direct_mix_available = + scheduler_gm_query(inbox->head) == SCHEDULER_INBOX_EMPTY && + scheduler_ready_pending_head( + scheduler_ready_owner_pending_load(&context->owner_pending_endpoints[core_type]) + ) == SCHEDULER_INBOX_EMPTY && + !scheduler_ready_directory_nonempty( + scheduler_state_base, context, context->config.scheduler_count, core_type + ); + } + SchedulerReadyClaim *direct_claim = + direct_mix_available ? direct_mix_ready : + direct_ready != nullptr && direct_ready->task_id < 0 && core_type == direct_core_type ? + direct_ready : + nullptr; + if (direct_claim != nullptr) { + direct_claim->task_id = waiter; + direct_claim->source = SchedulerReadySource::DIRECT_RESOLVE; + direct_claim->publication_mode = SchedulerPublicationMode::REFILL; if (phase_timing_enabled) { __gm__ SchedulerTaskTrace *cells = scheduler_state_at( scheduler_state_base, context->profiling->trace_cells_offset ); - cells[waiter].ready_transition_cycles = scheduler_cycles(); - scheduler_writeback_cache_line(&cells[waiter].ready_transition_cycles); + const uint64_t transition_cycles = scheduler_cycles(); + for (uint8_t subtask = 0; subtask < 3; ++subtask) { + if ((metadata->active_mask & (1U << subtask)) == 0) continue; + auto &trace = cells[scheduler_task_trace_index( + metadata->trace_index_base, metadata->active_mask, subtask + )]; + trace.ready_transition_cycles = transition_cycles; + scheduler_writeback_cache_line(&trace.ready_transition_cycles); + } } } else if (!scheduler_ready_batch_append( scheduler_state_base, context, waiter, &batches[core_type], ready_stats, profiling_level @@ -1368,7 +1585,7 @@ inline __aicore__ bool scheduler_resolve_completion( } waiter = next; } - for (uint32_t type = 0; type < SCHEDULER_CORE_TYPE_COUNT; ++type) { + for (uint32_t type = 0; type < SCHEDULER_READY_QUEUE_COUNT; ++type) { if (!scheduler_ready_batch_push(scheduler_state_base, context, type, &batches[type], ready_stats)) { scheduler_record_error( run_control, task_id, SchedulerGraphResult::INVALID_ARGUMENTS, &graph, context, 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 016e3c0226..548b3d8348 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 @@ -632,6 +632,8 @@ inline constexpr uint64_t SCHEDULER_STATE_ALIGNMENT = 128; inline constexpr uint64_t SCHEDULER_WORKER_CAPACITY = 108; inline constexpr uint32_t SCHEDULER_PENDING_SLOT_COUNT = 2; inline constexpr uint32_t SCHEDULER_CORE_TYPE_COUNT = 2; +inline constexpr uint32_t SCHEDULER_READY_QUEUE_COUNT = 3; +inline constexpr uint32_t SCHEDULER_MIX_QUEUE = 2; inline constexpr uint32_t SCHEDULER_CLUSTER_CAPACITY = SCHEDULER_WORKER_CAPACITY / 3; inline constexpr uint32_t SCHEDULER_CAPACITY = SCHEDULER_CLUSTER_CAPACITY; inline constexpr uint32_t SCHEDULER_GANG_COHORT_COUNT = 2; @@ -738,11 +740,17 @@ struct alignas(16) SchedulerTaskMetadata { uint16_t kernel_ids[3]; uint8_t active_mask; uint8_t flags; - uint16_t logical_block_num; - uint16_t total_required_subtasks; + uint32_t trace_index_base; int32_t timing_slot; }; +inline __aicore__ uint32_t scheduler_task_trace_index(uint32_t trace_index_base, uint8_t active_mask, uint8_t subtask) { + return trace_index_base + static_cast(__builtin_popcount(active_mask & ((1U << subtask) - 1))); +} +inline __aicore__ uint32_t scheduler_task_ready_queue(uint8_t flags, uint8_t active_mask) { + return (flags & SCHEDULER_TASK_MIX) != 0 ? SCHEDULER_MIX_QUEUE : ((active_mask & 1U) != 0 ? 0U : 1U); +} + inline __aicore__ bool scheduler_task_is_executable(uint8_t flags) { return (flags & SCHEDULER_TASK_EXECUTABLE) != 0; } inline __aicore__ bool scheduler_task_has_fanin(uint8_t flags) { return (flags & SCHEDULER_TASK_HAS_FANIN) != 0; } @@ -762,7 +770,7 @@ inline __aicore__ bool scheduler_task_has_predicate(uint8_t flags) { } inline __aicore__ bool scheduler_task_is_gang(uint8_t flags) { - return (flags & (SCHEDULER_TASK_MIX | SCHEDULER_TASK_SPMD)) != 0; + return (flags & (SCHEDULER_TASK_SYNC_START | SCHEDULER_TASK_SPMD)) != 0; } inline __aicore__ uint32_t scheduler_task_priority_bit(uint8_t flags) { @@ -816,7 +824,7 @@ struct alignas(128) SchedulerGangCoordinator { uint64_t gang_task_count; uint64_t scheduler_count; uint64_t cohort_count; - uint64_t reserved0; + uint64_t mix_task_count; uint64_t owner_reserved; uint64_t admitted_count; @@ -885,7 +893,7 @@ struct alignas(64) SchedulerReadyDirectoryShard { }; struct alignas(128) SchedulerReadyDirectory { - SchedulerReadyDirectoryShard core_types[SCHEDULER_CORE_TYPE_COUNT][SCHEDULER_READY_DIRECTORY_SHARD_COUNT]; + SchedulerReadyDirectoryShard core_types[SCHEDULER_READY_QUEUE_COUNT][SCHEDULER_READY_DIRECTORY_SHARD_COUNT]; volatile uint64_t bootstrap_ready_types[SCHEDULER_WORKER_CAPACITY]; }; @@ -1193,7 +1201,7 @@ static_assert( ); static_assert( offsetof(SchedulerReadyDirectory, bootstrap_ready_types) == - SCHEDULER_CORE_TYPE_COUNT * SCHEDULER_READY_DIRECTORY_SHARD_COUNT * 64, + SCHEDULER_READY_QUEUE_COUNT * SCHEDULER_READY_DIRECTORY_SHARD_COUNT * 64, "bootstrap flags must follow the ready directory shards" ); static_assert( @@ -1317,7 +1325,7 @@ inline bool scheduler_layout_reserve(uint64_t *cursor, uint64_t size, uint64_t a inline bool scheduler_plan_layout( uint64_t task_count, uint64_t aic_task_count, uint64_t aiv_task_count, AicoreSchedulerLayout *layout, - bool enable_activity_profiling = false + bool enable_activity_profiling = false, uint64_t trace_count = 0 ) { if (layout == nullptr || aic_task_count > task_count || aiv_task_count > task_count) return false; AicoreSchedulerLayout next{}; @@ -1342,7 +1350,7 @@ inline bool scheduler_plan_layout( !SCHEDULER_RESERVE_ARRAY(task_count, SchedulerTaskMetadata, task_metadata_offset) || !SCHEDULER_RESERVE_ARRAY(task_count, SchedulerTaskControl, task_controls_offset) || !SCHEDULER_RESERVE_ARRAY( - SCHEDULER_CORE_TYPE_COUNT * SCHEDULER_WORKER_CAPACITY, SchedulerReadyInbox, ready_inboxes_offset + SCHEDULER_READY_QUEUE_COUNT * SCHEDULER_WORKER_CAPACITY, SchedulerReadyInbox, ready_inboxes_offset ) || !scheduler_layout_reserve( &cursor, sizeof(SchedulerReadyDirectory), alignof(SchedulerReadyDirectory), &next.ready_directory_offset @@ -1355,7 +1363,7 @@ inline bool scheduler_plan_layout( SCHEDULER_GANG_COHORT_COUNT * SCHEDULER_CLUSTER_CAPACITY, SchedulerGangParticipant, gang_participants_offset ) || !SCHEDULER_RESERVE_ARRAY(SCHEDULER_CLUSTER_CAPACITY, SchedulerGangCommand, gang_commands_offset) || - !SCHEDULER_RESERVE_ARRAY(task_count, SchedulerTaskTrace, trace_cells_offset) || + !SCHEDULER_RESERVE_ARRAY(trace_count == 0 ? task_count : trace_count, SchedulerTaskTrace, trace_cells_offset) || (enable_activity_profiling && !SCHEDULER_RESERVE_ARRAY(SCHEDULER_CLUSTER_CAPACITY, SchedulerActivityBuffer, activity_buffers_offset)) || !scheduler_layout_checked_align(cursor, SCHEDULER_STATE_ALIGNMENT, &next.total_size)) { @@ -1378,7 +1386,7 @@ inline bool scheduler_init_data_from_layout(void *base, const AicoreSchedulerLay controls[i].waiting_producer = static_cast(SCHEDULER_TASK_ID_INVALID); } auto *ready = scheduler_state_at(base, layout.ready_inboxes_offset); - for (uint64_t i = 0; i < SCHEDULER_CORE_TYPE_COUNT * SCHEDULER_WORKER_CAPACITY; ++i) + for (uint64_t i = 0; i < SCHEDULER_READY_QUEUE_COUNT * SCHEDULER_WORKER_CAPACITY; ++i) ready[i].head = SCHEDULER_INBOX_EMPTY; auto *contexts = scheduler_state_at(base, layout.worker_contexts_offset); for (uint64_t worker = 0; worker < SCHEDULER_WORKER_CAPACITY; ++worker) { diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/rendezvous_orch.cpp b/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/rendezvous_orch.cpp new file mode 100644 index 0000000000..89a04c15af --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/rendezvous_orch.cpp @@ -0,0 +1,34 @@ +/* + * 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 "orchestration_api.h" + +extern "C" __attribute__((visibility("default"))) OrchestrationConfig aicpu_orchestration_config(const ChipTaskArgs &) { + return OrchestrationConfig{.expected_arg_count = 3}; +} + +extern "C" __attribute__((visibility("default"))) void aicpu_orchestration_entry(const ChipTaskArgs &orch_args) { + const auto mask = static_cast(orch_args.scalar(0)); + const uint32_t shape[1] = {48}; + for (uint32_t task = 0; task < static_cast(orch_args.scalar(1)); ++task) { + const uint32_t offset[1] = {task * 48}; + auto state = orch_args.tensor(0).ref().view(shape, offset); + CoreTaskArgs args; + args.add_inout(state); + args.add_scalar(static_cast(mask)); + MixedKernels kernels; + kernels.aic_kernel_id = (mask & 1U) != 0 ? 0 : INVALID_KERNEL_ID; + kernels.aiv0_kernel_id = 1; + kernels.aiv1_kernel_id = 2; + rt_submit_task(kernels, args); + } +} diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/single_block_mix_orch.cpp b/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/single_block_mix_orch.cpp new file mode 100644 index 0000000000..8bec324f11 --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/orchestration/single_block_mix_orch.cpp @@ -0,0 +1,87 @@ +/* + * 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 "orchestration_api.h" // NOLINT(build/include_subdir) + +namespace { + +constexpr int32_t kMatmulFuncId = 0; +constexpr int32_t kAddFuncId = 1; +constexpr int32_t kMulFuncId = 2; +constexpr int32_t kTaskCount = 16; +constexpr uint32_t kTileElems = 128 * 128; + +TaskId submit_mix( + const ChipTaskArgs &orch_args, int32_t task_index, uint8_t mask, const TaskId *dependency, uint32_t dependency_count +) { + const uint32_t shape[1] = {kTileElems}; + const uint32_t offset[1] = {static_cast(task_index) * kTileElems}; + simpler::hbg::Tensor c = orch_args.tensor(2).ref().view(shape, offset); + simpler::hbg::Tensor f = orch_args.tensor(5).ref().view(shape, offset); + simpler::hbg::Tensor i = orch_args.tensor(8).ref().view(shape, offset); + + const uint32_t previous_offset[1] = {static_cast(task_index == 0 ? 0 : task_index - 1) * kTileElems}; + const bool chained = dependency_count != 0; + auto a = chained && (mask & 1U) != 0 ? orch_args.tensor(2).ref().view(shape, previous_offset) : + orch_args.tensor(0).ref(); + auto d = chained && (mask & 2U) != 0 ? orch_args.tensor(5).ref().view(shape, previous_offset) : + orch_args.tensor(3).ref(); + auto g = chained && (mask & 4U) != 0 ? orch_args.tensor(8).ref().view(shape, previous_offset) : + orch_args.tensor(6).ref(); + CoreTaskArgs args; + args.add_input(a); + args.add_input(orch_args.tensor(1).ref()); + args.add_output(c); + args.add_input(d); + args.add_input(orch_args.tensor(4).ref()); + args.add_output(f); + args.add_input(g); + args.add_input(orch_args.tensor(7).ref()); + args.add_output(i); + args.set_dependencies(dependency, dependency_count); + + MixedKernels kernels; + kernels.aic_kernel_id = (mask & 1U) != 0 ? kMatmulFuncId : INVALID_KERNEL_ID; + kernels.aiv0_kernel_id = (mask & 2U) != 0 ? kAddFuncId : INVALID_KERNEL_ID; + kernels.aiv1_kernel_id = (mask & 4U) != 0 ? kMulFuncId : INVALID_KERNEL_ID; + return rt_submit_task(kernels, args).task_id(); +} + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) OrchestrationConfig aicpu_orchestration_config(const ChipTaskArgs &orch_args) { + (void)orch_args; + return OrchestrationConfig{.expected_arg_count = 11}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const ChipTaskArgs &orch_args) { + const int64_t graph_case = orch_args.scalar(0); + if (graph_case < 0 || graph_case > 2) { + rt_report_fatal(SIMPLER_ERROR_INVALID_ARGS, "unsupported graph case %ld", graph_case); + return; + } + + const uint8_t mask = static_cast(orch_args.scalar(1)); + TaskId previous = TaskId::invalid(); + for (int32_t task_index = 0; task_index < kTaskCount; ++task_index) { + const bool chained = graph_case == 0 && task_index != 0; + previous = submit_mix( + orch_args, task_index, graph_case == 2 && task_index % 2 ? (task_index % 4 == 1 ? 1U : 2U) : mask, + chained ? &previous : nullptr, chained ? 1U : 0U + ); + } +} + +} // extern "C" diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous.h b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous.h new file mode 100644 index 0000000000..7b6ff2a450 --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous.h @@ -0,0 +1,63 @@ +/* + * 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 +#include +#ifdef PTO_CPUSTUB_HPP +#include +#endif +#include "intrinsic.h" +#include "tensor.h" + +struct alignas(128) RendezvousFlag { + volatile int64_t value; + int64_t padding[15]; +}; + +static_assert(sizeof(RendezvousFlag) == 128); + +static __aicore__ void rendezvous(__gm__ int64_t *args, uint32_t lane) { + auto *tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + auto *state = reinterpret_cast<__gm__ RendezvousFlag *>( + reinterpret_cast<__gm__ int64_t *>(tensor->buffer.addr) + tensor->start_offset + ); + const uint32_t mask = static_cast(args[1]); + state[lane].value = 1; + dcci(&state[lane], cache_line_t::SINGLE_CACHE_LINE, dcci_dst_t::CACHELINE_OUT); + dsb((mem_dsb_t)0); + for (uint32_t poll = 0; poll < 10000000; ++poll) { + bool arrived = true; + for (uint32_t peer = 0; peer < 3; ++peer) { + if (peer == lane || (mask & (1U << peer)) == 0) continue; + dcci(&state[peer], cache_line_t::SINGLE_CACHE_LINE); + dsb((mem_dsb_t)0); + arrived = arrived && state[peer].value >= 1; + } + if (!arrived) { +#ifdef PTO_CPUSTUB_HPP + std::this_thread::yield(); +#endif + continue; + } + if (lane == 1) { + for (volatile uint32_t delay = 0; delay < 2048; ++delay) {} + } + state[lane].value = 2; + dcci(&state[lane], cache_line_t::SINGLE_CACHE_LINE, dcci_dst_t::CACHELINE_OUT); + dsb((mem_dsb_t)0); + return; + } + state[lane].value = -1; + dcci(&state[lane], cache_line_t::SINGLE_CACHE_LINE, dcci_dst_t::CACHELINE_OUT); + dsb((mem_dsb_t)0); +} diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_0.cpp b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_0.cpp new file mode 100644 index 0000000000..c6e1c5a9ac --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_0.cpp @@ -0,0 +1,14 @@ +/* + * 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 "rendezvous.h" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { rendezvous(args, 0); } diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_1.cpp b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_1.cpp new file mode 100644 index 0000000000..5f1f7816b3 --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_1.cpp @@ -0,0 +1,14 @@ +/* + * 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 "rendezvous.h" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { rendezvous(args, 1); } diff --git a/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_2.cpp b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_2.cpp new file mode 100644 index 0000000000..0ae930264d --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/kernels/rendezvous_2.cpp @@ -0,0 +1,14 @@ +/* + * 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 "rendezvous.h" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { rendezvous(args, 2); } diff --git a/tests/st/a5/host_build_graph/single_block_mix/test_single_block_mix.py b/tests/st/a5/host_build_graph/single_block_mix/test_single_block_mix.py new file mode 100644 index 0000000000..dffe8baca8 --- /dev/null +++ b/tests/st/a5/host_build_graph/single_block_mix/test_single_block_mix.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +# 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. +# ----------------------------------------------------------------------------------------------------------- +"""Single-block Mix scheduling coverage for dependent and independent DAGs.""" + +import ctypes + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, TensorArg, scene_test + +MATMUL_SIZE = 128 +TILE_ELEMS = MATMUL_SIZE * MATMUL_SIZE +TASK_COUNT = 16 + + +@scene_test(level=2, runtime="host_build_graph") +class TestSingleBlockMixHostBuildGraphA5(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/single_block_mix_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.INOUT, D.IN, D.IN, D.INOUT, D.IN, D.IN, D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "MATMUL", + "source": "../../tensormap_and_ringbuffer/mixed_example/kernels/aic/kernel_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "ADD", + "source": "../../tensormap_and_ringbuffer/mixed_example/kernels/aiv/kernel_add.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT], + }, + { + "func_id": 2, + "name": "MUL", + "source": "../../tensormap_and_ringbuffer/mixed_example/kernels/aiv/kernel_mul.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT, D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + {"name": f"{kind}16_mask{mask}", "platforms": ["a5sim", "a5"], "params": {"graph_case": case, "mask": mask}} + for kind, case in [("chain", 0), ("burst", 1)] + for mask in [3, 5, 6, 7] + ] + [{"name": "ordinary_mix_interleaved", "platforms": ["a5sim", "a5"], "params": {"graph_case": 2, "mask": 7}}] + + def generate_args(self, params): + torch.manual_seed(42) + a = torch.randn(MATMUL_SIZE, MATMUL_SIZE, dtype=torch.float32) * 0.01 + b = torch.randn(MATMUL_SIZE, MATMUL_SIZE, dtype=torch.float32) * 0.01 + d = torch.randn(TILE_ELEMS, dtype=torch.float32) * 0.01 + e = torch.randn(TILE_ELEMS, dtype=torch.float32) * 0.01 + g = torch.randn(TILE_ELEMS, dtype=torch.float32) * 0.01 + h = torch.randn(TILE_ELEMS, dtype=torch.float32) * 0.01 + + def output(): + return torch.zeros(TASK_COUNT * TILE_ELEMS, dtype=torch.float32) + + return TaskArgsBuilder( + TensorArg("a", a.flatten()), + TensorArg("b", b.flatten()), + TensorArg("c", output()), + TensorArg("d", d), + TensorArg("e", e), + TensorArg("f", output()), + TensorArg("g", g), + TensorArg("h", h), + TensorArg("i", output()), + Scalar("graph_case", ctypes.c_int64(params["graph_case"])), + Scalar("mask", ctypes.c_int64(params["mask"])), + ) + + def compute_golden(self, args, params): + c = args.c.reshape(TASK_COUNT, TILE_ELEMS) + f = args.f.reshape(TASK_COUNT, TILE_ELEMS) + i = args.i.reshape(TASK_COUNT, TILE_ELEMS) + for task in range(TASK_COUNT): + mask = params["mask"] + if params["graph_case"] == 2 and task % 2: + mask = 1 if task % 4 == 1 else 2 + chained = params["graph_case"] == 0 and task != 0 + if mask & 1: + a = c[task - 1] if chained else args.a + c[task] = torch.matmul( + a.reshape(MATMUL_SIZE, MATMUL_SIZE), args.b.reshape(MATMUL_SIZE, MATMUL_SIZE) + ).flatten() + if mask & 2: + f[task] = (f[task - 1] if chained else args.d) + args.e + if mask & 4: + i[task] = (i[task - 1] if chained else args.g) * args.h + + +@scene_test(level=2, runtime="host_build_graph") +class TestMixKernelRendezvous(SceneTestCase): + RTOL = 0 + ATOL = 0 + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/rendezvous_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": lane, + "name": f"RENDEZVOUS_{lane}", + "source": f"kernels/rendezvous_{lane}.cpp", + "core_type": "aic" if lane == 0 else "aiv", + "signature": [D.INOUT], + } + for lane in range(3) + ], + } + CASES = [ + {"name": f"burst256_mask{mask}", "platforms": ["a5sim", "a5"], "params": {"mask": mask, "task_count": 256}} + for mask in [6, 7] + ] + + def generate_args(self, params): + return TaskArgsBuilder( + TensorArg("state", torch.zeros(params["task_count"] * 48, dtype=torch.int64)), + Scalar("mask", ctypes.c_int64(params["mask"])), + Scalar("task_count", ctypes.c_int64(params["task_count"])), + ) + + def compute_golden(self, args, params): + for lane in range(3): + if params["mask"] & (1 << lane): + args.state.reshape(params["task_count"], 48)[:, lane * 16] = 2 + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/a5/runtime/host_build_graph/hbg_scheduler_test_support.h b/tests/ut/cpp/a5/runtime/host_build_graph/hbg_scheduler_test_support.h index 964742bf91..2f94382615 100644 --- a/tests/ut/cpp/a5/runtime/host_build_graph/hbg_scheduler_test_support.h +++ b/tests/ut/cpp/a5/runtime/host_build_graph/hbg_scheduler_test_support.h @@ -127,12 +127,19 @@ using GraphBuffer = BasicGraphBuffer<8192>; inline constexpr uint64_t kFixtureCallableCount = 1024; struct FixtureStorage { - explicit FixtureStorage(uint64_t task_count, uint64_t workers = 2) : + explicit FixtureStorage(uint64_t task_count, uint64_t workers = 2, uint64_t traces_per_task = 1) : test_profiles(workers), test_contexts(workers), test_graphs(workers), local_states(workers) { - EXPECT_TRUE(scheduler_plan_layout(task_count, task_count, 0, &layout)); + for (auto &local : test_contexts) + for (auto &lane : local.slots) + for (auto &slot : lane) + scheduler_initialize_free_slot(&slot); + for (auto &lane : scheduler_local_state.slots) + for (auto &slot : lane) + scheduler_initialize_free_slot(&slot); + EXPECT_TRUE(scheduler_plan_layout(task_count, task_count, 0, &layout, false, task_count * traces_per_task)); 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); @@ -167,8 +174,7 @@ struct FixtureStorage { 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].trace_index_base = task; metadata[task].flags = SCHEDULER_TASK_EXECUTABLE; } } diff --git a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_contracts.cpp b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_contracts.cpp index a275112a34..234e0323e6 100644 --- a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_contracts.cpp +++ b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_contracts.cpp @@ -70,9 +70,9 @@ TEST(AicoreSchedulerState, DistinguishesResidentAndExplicitLegacyModes) { EXPECT_TRUE(aicore_scheduler_runtime_mode_is_explicit_legacy(SCHEDULER_RUNTIME_MODE_LEGACY_UNSUPPORTED_SHAPE)); } -TEST(AicoreSchedulerState, ResidentV0AcceptsOnlySingleLaneSingleBlockTasks) { +TEST(AicoreSchedulerState, ResidentAcceptsSingleBlockOrdinaryMix) { EXPECT_TRUE(scheduler_resident_v0_task_shape_supported(1, 1, false)); - EXPECT_FALSE(scheduler_resident_v0_task_shape_supported(2, 1, false)); + EXPECT_TRUE(scheduler_resident_v0_task_shape_supported(2, 1, false)); EXPECT_FALSE(scheduler_resident_v0_task_shape_supported(1, 2, false)); EXPECT_FALSE(scheduler_resident_v0_task_shape_supported(1, 1, true)); } @@ -300,25 +300,30 @@ TEST(SchedulerGraph, RejectsAViewWithANonZeroReservedWord) { ); } -TEST(SchedulerDispatchPayload, DisablesDeferredCompletionWithoutASlab) { +TEST(SchedulerDispatchPayload, PreservesInitializedSlotContextDuringMaterialization) { GraphBuffer graph(1); graph.executable(0, 0); DispatchPayload payload{}; - payload.local_context.async_ctx.task_token = TaskId::make_global(17); + payload.local_context.async_ctx.task_token = TaskId::invalid(); + payload.global_context.sub_block_id = 7; + payload.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&payload.local_context); + payload.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&payload.global_context); SchedulerTaskInfo task{0, 1, 0, CoreType::AIC}; ASSERT_EQ( scheduler_materialize_task_payload_resolved(graph.graph(), task, 0x1000, &payload), SchedulerGraphResult::OK ); EXPECT_FALSE(payload.local_context.async_ctx.task_token.is_valid()); - EXPECT_EQ(payload.global_context.sub_block_id, 0); + EXPECT_EQ(payload.global_context.sub_block_id, 7); + EXPECT_EQ(payload.args[PAYLOAD_LOCAL_CONTEXT_INDEX], reinterpret_cast(&payload.local_context)); + EXPECT_EQ(payload.args[PAYLOAD_GLOBAL_CONTEXT_INDEX], reinterpret_cast(&payload.global_context)); task.subtask_slot = 2; task.core_type = CoreType::AIV; ASSERT_EQ( scheduler_materialize_task_payload_resolved(graph.graph(), task, 0x1000, &payload), SchedulerGraphResult::OK ); - EXPECT_EQ(payload.global_context.sub_block_id, 1); + EXPECT_EQ(payload.global_context.sub_block_id, 7); } TEST(SchedulerDispatchPayload, RejectsInvalidGraphBoundsBeforeReadingPayload) { diff --git a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_dispatch.cpp b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_dispatch.cpp index ea81004758..6635738989 100644 --- a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_dispatch.cpp +++ b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_dispatch.cpp @@ -19,7 +19,7 @@ #include #include -#include "scheduler/scheduler_dispatch.h" +#include "scheduler/scheduler_mix.h" #include "runtime_types.h" #include "hbg_scheduler_test_support.h" @@ -111,6 +111,7 @@ void occupy_normal_slot( local_slot.task_id = task_id; local_slot.subtask_slot = 1; local_slot.state = SchedulerDispatchSlotState::READY; + local_slot.generation = ++storage.scheduler_local_state.dispatch_sequences[test_cluster_lane(scheduler, worker_id)]; if (worker_id == scheduler.worker_index) { scheduler_local_ready_publish(&storage.scheduler_local_state, pending_slot); } else { @@ -134,6 +135,7 @@ prepare_completed_normal_slot(FixtureStorage &storage, SchedulerWorkerContext &s local_slot.task_id = 0; local_slot.subtask_slot = 0; local_slot.state = SchedulerDispatchSlotState::READY; + local_slot.generation = ++storage.scheduler_local_state.dispatch_sequences[test_cluster_lane(scheduler, worker_id)]; scheduler_ssbuf_store_relaxed( &test_ssbuf_control(storage, scheduler, worker_id, 0)->publication, local_slot.generation ); @@ -191,6 +193,7 @@ TEST(SchedulerClusterCompletion, AccumulatesBatchOnSuccessAndBothKindsOfErrorExi second.task_id = 1; second.subtask_slot = 0; second.state = SchedulerDispatchSlotState::READY; + second.generation = ++storage.scheduler_local_state.dispatch_sequences[0]; auto *control = scheduler_task_control_at(storage.scheduler_state->base(), storage.local_context(&scheduler), 1); if (failure == 1) control->wake_list_head = SCHEDULER_WAKE_LIST_CLOSED; @@ -230,6 +233,7 @@ TEST(SchedulerClusterCompletion, SpscGenerationCompletesNormalTask) { local_slot->subtask_slot = 0; storage.scheduler_local_state.set_timing_slot(0, 0, 0); local_slot->state = SchedulerDispatchSlotState::READY; + local_slot->generation = ++storage.scheduler_local_state.dispatch_sequences[0]; scheduler_ssbuf_store_relaxed(&test_ssbuf_control(storage, scheduler, 0, 0)->publication, local_slot->generation); auto *executor_trace = test_ssbuf_trace(storage, scheduler, 0, 0); executor_trace->kernel_start_cycles = 100; @@ -284,7 +288,7 @@ TEST(SchedulerClusterCompletion, UnprofiledSelfHandoffNeedsNoProfilingStorage) { )); uint32_t selected = UINT32_MAX; uint64_t publication = 0; - ASSERT_TRUE(scheduler_local_ready_pop(local, 0, &selected, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(local, &selected, &publication)); ASSERT_EQ(selected, 0u); EXPECT_EQ(local->timing_slot(1, selected), -1); local->local_completed_generations[selected] = scheduler_dispatch_generation(publication); @@ -316,6 +320,8 @@ TEST(SchedulerClusterCompletion, SelfSlotsKeepDistinctSampledTraces) { slot.subtask_slot = 1; local->set_timing_slot(1, index, index); slot.state = SchedulerDispatchSlotState::READY; + slot.generation = ++local->dispatch_sequences[1]; + storage.metadata[index].active_mask = 2; local->profiling->executor_traces[index].kernel_start_cycles = 100 + index; local->profiling->executor_traces[index].kernel_end_cycles = 200 + index; local->local_completed_generations[index] = slot.generation; @@ -350,6 +356,7 @@ TEST(SchedulerClusterCompletion, RejectsStaleCompletionGenerationAtNamedSite) { slot->task_id = 0; local_slot->task_id = 0; local_slot->state = SchedulerDispatchSlotState::READY; + local_slot->generation = ++storage.scheduler_local_state.dispatch_sequences[0]; SchedulerWakeStats wake_stats{}; SchedulerReadyStats ready_stats{}; @@ -461,6 +468,7 @@ TEST(SchedulerClusterCompletion, PropagatesTraceToCompletionAndWokenTask) { local_slot->task_id = 0; local_slot->subtask_slot = 0; local_slot->state = SchedulerDispatchSlotState::READY; + local_slot->generation = ++storage.scheduler_local_state.dispatch_sequences[0]; auto *executor_trace = test_ssbuf_trace(storage, scheduler, 0, 0); executor_trace->kernel_start_cycles = 100; executor_trace->kernel_end_cycles = 200; @@ -515,6 +523,7 @@ TEST(SchedulerClusterCompletion, DirectlyRefillsCompletedSlotWhenReadyTaskExists auto *slot = test_slot(storage, scheduler, 0, 0); SchedulerLocalSlotState *local_slot = &storage.scheduler_local_state.slots[0][0]; scheduler_initialize_free_slot(local_slot); + local_slot->generation = ++storage.scheduler_local_state.dispatch_sequences[0]; const uint32_t completed_generation = local_slot->generation; slot->task_id = 0; local_slot->task_id = 0; @@ -1145,7 +1154,7 @@ TEST(SchedulerDeferredAiv, LocalReservationsAndReadyHandoffsDoNotPollSharedState for (uint32_t expected_slot = 0; expected_slot < SCHEDULER_PENDING_SLOT_COUNT; ++expected_slot) { uint32_t pending_slot = UINT32_MAX; uint64_t publication = 0; - ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, expected_slot, &pending_slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, &pending_slot, &publication)); EXPECT_EQ(pending_slot, expected_slot); EXPECT_EQ(scheduler_dispatch_state(publication), SchedulerDispatchSlotState::READY); EXPECT_EQ(scheduler_dispatch_generation(publication), expected_generations[expected_slot]); @@ -1153,7 +1162,7 @@ TEST(SchedulerDeferredAiv, LocalReservationsAndReadyHandoffsDoNotPollSharedState EXPECT_EQ(scheduler_local_state.local_ready_mask, 0u); } -TEST(SchedulerLocalReady, RotatesPastRepublishedLowerSlot) { +TEST(SchedulerLocalReady, OlderDispatchPrecedesReusedLowerSlot) { SchedulerLocalState scheduler_local_state{}; scheduler_local_state.config.self_lane = 1; scheduler_local_state.slots[1][0].generation = 1; @@ -1167,18 +1176,19 @@ TEST(SchedulerLocalReady, RotatesPastRepublishedLowerSlot) { uint32_t pending_slot = UINT32_MAX; uint64_t publication = 0; - ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, 0, &pending_slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, &pending_slot, &publication)); EXPECT_EQ(pending_slot, 0u); EXPECT_EQ(publication, slot_0_publication); + scheduler_local_state.slots[1][0].generation = 3; scheduler_local_ready_publish(&scheduler_local_state, 0); - ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, 1, &pending_slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, &pending_slot, &publication)); EXPECT_EQ(pending_slot, 1u); EXPECT_EQ(publication, slot_1_publication); - ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, 0, &pending_slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&scheduler_local_state, &pending_slot, &publication)); EXPECT_EQ(pending_slot, 0u); - EXPECT_EQ(publication, slot_0_publication); + EXPECT_EQ(publication, scheduler_dispatch_publication(3, SchedulerDispatchSlotState::READY)); EXPECT_EQ(scheduler_local_state.local_ready_mask, 0u); } @@ -1190,9 +1200,9 @@ TEST(SchedulerLocalReady, InvalidSlotStateIsNotReconstructedAsReady) { scheduler_local_ready_publish(&local, 0); uint32_t slot = UINT32_MAX; uint64_t publication = 0; - ASSERT_TRUE(scheduler_local_ready_pop(&local, 0, &slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&local, &slot, &publication)); EXPECT_NE(scheduler_dispatch_state(publication), SchedulerDispatchSlotState::READY); - EXPECT_FALSE(scheduler_local_ready_pop(&local, 0, &slot, &publication)); + EXPECT_FALSE(scheduler_local_ready_pop(&local, &slot, &publication)); } TEST(SchedulerDeferredAiv, PrefersNewPeerCapacityAndSelfPublishesOnlyOne) { @@ -1331,7 +1341,7 @@ TEST(SchedulerDeferredAiv, RetiresCompletedPeerAndRefillsWithoutFreeDecision) { EXPECT_EQ(completed_control->state, static_cast(SchedulerTaskState::DONE)); EXPECT_EQ(storage.scheduler_local_state.pending_completed, 1u); EXPECT_EQ(peer_slot->task_id, 1); - EXPECT_EQ(peer_slot->generation, completed_generation + 1); + EXPECT_EQ(peer_slot->generation, 3u); EXPECT_EQ(peer_slot->state, SchedulerDispatchSlotState::READY); auto *peer_dispatch = test_ssbuf_control(storage, scheduler, 2, 0); const uint64_t peer_publication = scheduler_ssbuf_load_relaxed(&peer_dispatch->publication); @@ -1432,6 +1442,31 @@ TEST(SchedulerLocalConfig, SnapshotsReadyPublicationAndReloadsOnNextRun) { } } +TEST(SchedulerLocalConfig, DetectsMixIndependentlyOfGangTasks) { + FixtureStorage storage(1, 3); + GraphBuffer graph_buffer(1); + const SchedulerGraphView graph = graph_buffer.graph(); + configure_cached_cluster(storage, 1); + auto *coordinator = scheduler_state_at( + storage.scheduler_state->base(), storage.layout.gang_coordinator_offset + ); + coordinator->gang_task_count = 1; + coordinator->mix_task_count = 0; + SchedulerLocalState gang_only{}; + ASSERT_TRUE( + scheduler_initialize_local_config(storage.scheduler_state->base(), &storage.contexts[1], &graph, &gang_only) + ); + EXPECT_FALSE(gang_only.has_mix); + + coordinator->gang_task_count = 0; + coordinator->mix_task_count = 1; + SchedulerLocalState mix_only{}; + ASSERT_TRUE( + scheduler_initialize_local_config(storage.scheduler_state->base(), &storage.contexts[1], &graph, &mix_only) + ); + EXPECT_TRUE(mix_only.has_mix); +} + TEST(SchedulerLocalConfig, RejectsTruncatedOffsetsAndInconsistentPayloadRoutes) { FixtureStorage storage(1, 3); GraphBuffer graph_buffer(1); @@ -1575,3 +1610,444 @@ TEST(SchedulerNormalDispatch, NoUsableCapacityDoesNotAccessDirectory) { } } // namespace + +namespace { +struct MixFixture { + FixtureStorage storage{6, 3, 3}; + GraphBuffer graph{6}; + SchedulerLocalState *local; + uint64_t cursor{0}; + MixFixture() { + configure_normal_aiv_cluster(storage, 6); + local = storage.local_context(&storage.contexts[1], &storage.scheduler_local_state); + local->has_mix = true; + for (uint32_t task = 0; task < 6; ++task) + set_task(task, 7); + } + void set_task(uint32_t task, uint8_t mask) { + graph.mixed(task, mask); + auto &metadata = storage.metadata[task]; + metadata.active_mask = mask; + metadata.flags = SCHEDULER_TASK_EXECUTABLE | (__builtin_popcount(mask) > 1 ? SCHEDULER_TASK_MIX : 0); + metadata.trace_index_base = task * 3; + metadata.timing_slot = -1; + for (uint8_t subtask = 0; subtask < 3; ++subtask) + metadata.kernel_ids[subtask] = (mask & (1U << subtask)) != 0 ? 1 : UINT16_MAX; + } + void enqueue(uint32_t task) { + SchedulerReadyBatch batch{}; + ASSERT_TRUE(scheduler_ready_batch_append(storage.scheduler_state->base(), local, task, &batch, nullptr)); + ASSERT_TRUE(scheduler_ready_batch_push( + storage.scheduler_state->base(), local, + scheduler_task_ready_queue(storage.metadata[task].flags, storage.metadata[task].active_mask), &batch, + nullptr + )); + } + bool dispatch(bool steal = true) { + return scheduler_fill_cluster_mix_slots( + graph.graph(), storage.scheduler_state->base(), local, storage.run_control, &cursor, nullptr, 0, steal, + storage.ssbuf_region, nullptr + ); + } + bool complete(uint32_t lane, uint32_t slot, SchedulerRefillCandidates *candidates = nullptr) { + local->local_ready_mask &= lane == local->config.self_lane ? ~(1U << slot) : UINT8_MAX; + return scheduler_service_cluster_completion_slot( + graph.graph(), storage.scheduler_state->base(), local, storage.run_control, lane, slot, + local->slots[lane][slot].generation, nullptr, nullptr, nullptr, nullptr, 0, nullptr, nullptr, + storage.ssbuf_region, candidates + ); + } +}; +} // namespace + +TEST(SchedulerMix, MasksHaveIndependentPayloadsAndStablePhysicalContext) { + for (uint8_t mask : {3, 5, 6, 7}) { + MixFixture f; + f.set_task(0, mask); + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + uint32_t published = 0; + uint64_t payload_addresses[3]{}; + for (uint32_t lane = 0; lane < 3; ++lane) { + const auto &slot = f.local->slots[lane][0]; + if (slot.state != SchedulerDispatchSlotState::READY) continue; + ++published; + EXPECT_EQ(slot.task_id, 0); + EXPECT_EQ(slot.generation, 1u); + auto *payload = scheduler_state_at( + f.storage.scheduler_state->base(), f.local->dispatch_payload_offset(lane, 0) + ); + payload_addresses[lane] = reinterpret_cast(payload); + EXPECT_EQ(payload->function_bin_addr, 0x1000u); + EXPECT_EQ(payload->global_context.sub_block_id, lane == 2 ? 1 : 0); + } + EXPECT_EQ(published, static_cast(__builtin_popcount(mask))); + for (uint32_t a = 0; a < 3; ++a) + for (uint32_t b = a + 1; b < 3; ++b) + if (payload_addresses[a] && payload_addresses[b]) EXPECT_NE(payload_addresses[a], payload_addresses[b]); + if (mask == 3 || mask == 5) EXPECT_EQ(f.local->slots[1][0].state, SchedulerDispatchSlotState::FREE); + } +} + +TEST(SchedulerMix, TwoPublishedMixesReleaseEachLaneBeforeWholeTask) { + MixFixture f; + f.enqueue(0); + f.enqueue(1); + ASSERT_TRUE(f.dispatch()); + for (uint32_t lane = 0; lane < 3; ++lane) { + EXPECT_EQ(f.local->slots[lane][0].task_id, 0); + EXPECT_EQ(f.local->slots[lane][1].task_id, 1); + EXPECT_EQ(f.local->slots[lane][1].generation, 2u); + } + ASSERT_TRUE(f.complete(0, 0)); + EXPECT_EQ(f.local->slots[0][0].state, SchedulerDispatchSlotState::FREE); + EXPECT_EQ(f.local->pending_completed, 0u); + ASSERT_TRUE(f.complete(0, 1)); + EXPECT_EQ(f.local->pending_completed, 0u); + EXPECT_EQ(f.local->slots[1][0].task_id, 0); + ASSERT_TRUE(f.complete(1, 0)); + EXPECT_EQ(f.local->pending_completed, 0u); + ASSERT_TRUE(f.complete(2, 0)); + EXPECT_EQ(f.local->pending_completed, 1u); + EXPECT_EQ(f.local->mix_trackers[0].task_id, SCHEDULER_TASK_ID_INVALID); + EXPECT_EQ(f.local->mix_trackers[1].task_id, 1); + EXPECT_FALSE(f.complete(2, 0)); + EXPECT_EQ(f.local->pending_completed, 1u); +} + +TEST(SchedulerMix, CapacityCheckLeavesTaskAvailableAndPreservesReservations) { + MixFixture f; + f.local->slots[0][0].state = SchedulerDispatchSlotState::FILLING; + f.local->slots[0][1].state = SchedulerDispatchSlotState::READY; + SchedulerReadyClaim candidate{}; + candidate.task_id = 0; + candidate.source = SchedulerReadySource::DIRECT_RESOLVE; + bool progress = false; + ASSERT_TRUE(scheduler_publish_direct_mix_candidate( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidate, nullptr, 0, + f.storage.ssbuf_region, &progress + )); + EXPECT_FALSE(progress); + EXPECT_EQ(candidate.task_id, SCHEDULER_TASK_ID_INVALID); + ASSERT_TRUE(f.dispatch()); + EXPECT_EQ(f.local->dispatch_sequences[1], 0u); + EXPECT_EQ(f.local->slots[0][0].state, SchedulerDispatchSlotState::FILLING); + f.local->slots[0][1].state = SchedulerDispatchSlotState::FREE; + ASSERT_TRUE(f.dispatch()); + EXPECT_EQ(f.local->slots[0][1].task_id, 0); + EXPECT_EQ(f.local->slots[0][0].state, SchedulerDispatchSlotState::FILLING); +} + +TEST(SchedulerMix, FailedPreparationPublishesNothingAndConsumesNoSequence) { + MixFixture f; + f.storage.metadata[0].kernel_ids[2] = 2; + f.enqueue(0); + EXPECT_FALSE(f.dispatch()); + EXPECT_NE(f.storage.run_control->scheduler_error, 0u); + for (uint32_t lane = 0; lane < 3; ++lane) { + EXPECT_EQ(f.local->dispatch_sequences[lane], 0u); + EXPECT_EQ(f.local->slots[lane][0].state, SchedulerDispatchSlotState::FREE); + EXPECT_EQ(f.storage.ssbuf_region->lanes[lane].dispatch[0].publication, 0u); + } + EXPECT_EQ(f.local->local_ready_mask, 0u); + EXPECT_EQ(f.local->mix_trackers[0].task_id, SCHEDULER_TASK_ID_INVALID); +} + +TEST(SchedulerMix, DisplacedDirectSuccessorReturnsToReadyQueue) { + MixFixture f; + f.set_task(2, 1); + SchedulerRefillCandidates candidates; + candidates.tasks[0][0] = 2; + f.local->slots[0][1].state = SchedulerDispatchSlotState::FILLING; + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + ASSERT_TRUE(scheduler_publish_refill_candidates( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidates, nullptr, 0, + f.storage.ssbuf_region + )); + EXPECT_EQ(f.local->slots[0][0].task_id, 0); + SchedulerReadyClaim claim; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, 0, &f.cursor, nullptr, + &claim + )); + EXPECT_EQ(claim.task_id, 2); + EXPECT_EQ(candidates.tasks[0][0], SCHEDULER_TASK_ID_INVALID); +} + +TEST(SchedulerMix, DirectCandidateDoesNotOvertakeQueuedMix) { + MixFixture f; + f.graph.executable(2, 0, {0}); + f.storage.metadata[2].flags |= SCHEDULER_TASK_HAS_FANIN; + ASSERT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 2, nullptr + ), + SchedulerRouteResult::WAITING + ); + f.enqueue(1); + SchedulerReadyClaim candidate{}; + ASSERT_TRUE(scheduler_resolve_completion( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 0, nullptr, nullptr, + nullptr, 0, false, nullptr, SCHEDULER_MIX_QUEUE, &candidate + )); + EXPECT_EQ(candidate.task_id, SCHEDULER_TASK_ID_INVALID); + ASSERT_TRUE(f.dispatch()); + for (uint32_t lane = 0; lane < 3; ++lane) { + EXPECT_EQ(f.local->slots[lane][0].task_id, 1); + EXPECT_EQ(f.local->slots[lane][1].task_id, 2); + } +} + +TEST(SchedulerMix, DirectCandidatePrecedesMixQueuedAfterResolve) { + MixFixture f; + f.graph.executable(1, 0, {0}); + f.storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + ASSERT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, nullptr + ), + SchedulerRouteResult::WAITING + ); + SchedulerReadyClaim candidate{}; + ASSERT_TRUE(scheduler_resolve_completion( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 0, nullptr, nullptr, + nullptr, 0, false, nullptr, SCHEDULER_MIX_QUEUE, &candidate + )); + ASSERT_EQ(candidate.task_id, 1); + f.enqueue(2); + + bool progress = false; + ASSERT_TRUE(scheduler_fill_mix_after_completions( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidate, &f.cursor, + nullptr, 0, true, f.storage.ssbuf_region, &progress + )); + EXPECT_TRUE(progress); + EXPECT_EQ(candidate.task_id, SCHEDULER_TASK_ID_INVALID); + for (uint32_t lane = 0; lane < 3; ++lane) { + EXPECT_EQ(f.local->slots[lane][0].task_id, 1); + EXPECT_EQ(f.local->slots[lane][1].task_id, 2); + } +} + +TEST(SchedulerMix, DirectCandidateQueuesWhenRequiredLaneIsBusy) { + MixFixture f; + f.graph.executable(1, 0, {0}); + f.storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + ASSERT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, nullptr + ), + SchedulerRouteResult::WAITING + ); + f.set_task(2, 3); + for (auto &slot : f.local->slots[2]) + slot.state = SchedulerDispatchSlotState::FILLING; + SchedulerReadyClaim candidate{}; + ASSERT_TRUE(scheduler_resolve_completion( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 0, nullptr, nullptr, + nullptr, 0, false, nullptr, SCHEDULER_MIX_QUEUE, &candidate + )); + ASSERT_EQ(candidate.task_id, 1); + f.enqueue(2); + + bool progress = false; + ASSERT_TRUE(scheduler_fill_mix_after_completions( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidate, &f.cursor, + nullptr, 0, true, f.storage.ssbuf_region, &progress + )); + EXPECT_TRUE(progress); + EXPECT_EQ(candidate.task_id, SCHEDULER_TASK_ID_INVALID); + EXPECT_EQ(f.local->slots[0][0].task_id, 2); + EXPECT_EQ(f.local->slots[1][0].task_id, 2); + for (auto &slot : f.local->slots[2]) + slot.state = SchedulerDispatchSlotState::FREE; + ASSERT_TRUE(f.dispatch()); + EXPECT_EQ(f.local->slots[0][1].task_id, 1); + EXPECT_EQ(f.local->slots[1][1].task_id, 1); + EXPECT_EQ(f.local->slots[2][0].task_id, 1); +} + +TEST(SchedulerMix, CompletedMixPassesUnlockedMixToDirectCandidate) { + MixFixture f; + f.graph.executable(1, 0, {0}); + f.storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + ASSERT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, nullptr + ), + SchedulerRouteResult::WAITING + ); + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + SchedulerRefillCandidates candidates; + ASSERT_TRUE(f.complete(0, 0, &candidates)); + ASSERT_TRUE(f.complete(1, 0, &candidates)); + EXPECT_EQ(candidates.mix_ready.task_id, SCHEDULER_TASK_ID_INVALID); + ASSERT_TRUE(f.complete(2, 0, &candidates)); + EXPECT_EQ(candidates.mix_ready.task_id, 1); + ASSERT_TRUE(scheduler_publish_direct_mix_candidate( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidates.mix_ready, + nullptr, 0, f.storage.ssbuf_region, nullptr + )); + for (uint32_t lane = 0; lane < 3; ++lane) + EXPECT_EQ(f.local->slots[lane][0].task_id, 1); +} + +TEST(SchedulerMix, MultipleUnlockedMixesKeepWakeOrder) { + MixFixture f; + for (uint32_t task : {1U, 2U}) { + f.graph.executable(task, 0, {0}); + f.storage.metadata[task].flags |= SCHEDULER_TASK_HAS_FANIN; + ASSERT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, task, nullptr + ), + SchedulerRouteResult::WAITING + ); + } + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + SchedulerRefillCandidates candidates; + ASSERT_TRUE(f.complete(0, 0, &candidates)); + ASSERT_TRUE(f.complete(1, 0, &candidates)); + ASSERT_TRUE(f.complete(2, 0, &candidates)); + EXPECT_EQ(candidates.mix_ready.task_id, SCHEDULER_TASK_ID_INVALID); + ASSERT_TRUE(f.dispatch()); + for (uint32_t lane = 0; lane < 3; ++lane) { + EXPECT_EQ(f.local->slots[lane][0].task_id, 2); + EXPECT_EQ(f.local->slots[lane][1].task_id, 1); + } +} + +TEST(SchedulerMix, LastLaneAloneResolvesDependency) { + MixFixture f; + f.set_task(1, 1); + f.graph.executable(1, 0, {0}); + f.storage.metadata[1].flags |= SCHEDULER_TASK_HAS_FANIN; + EXPECT_EQ( + scheduler_route_task( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, nullptr + ), + SchedulerRouteResult::WAITING + ); + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + ASSERT_TRUE(f.complete(0, 0)); + ASSERT_TRUE(f.complete(1, 0)); + auto *control = scheduler_task_control_at(f.storage.scheduler_state->base(), f.local, 0); + EXPECT_NE(control->state, static_cast(SchedulerTaskState::DONE)); + SchedulerRefillCandidates candidates; + ASSERT_TRUE(f.complete(2, 0, &candidates)); + EXPECT_EQ(control->state, static_cast(SchedulerTaskState::DONE)); + SchedulerReadyClaim claim; + ASSERT_TRUE(scheduler_claim_ready_for_slot( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, 1, 0, &f.cursor, nullptr, + &claim + )); + EXPECT_EQ(claim.task_id, 1); +} + +TEST(SchedulerMix, PartialCompletionsAllowMoreThanTwoLiveMixes) { + MixFixture f; + f.set_task(0, 3); + f.set_task(1, 6); + f.set_task(2, 5); + f.set_task(3, 6); + f.enqueue(0); + f.enqueue(1); + ASSERT_TRUE(f.dispatch()); + ASSERT_TRUE(f.complete(2, 0)); + ASSERT_TRUE(f.complete(2, 1)); + EXPECT_EQ(f.local->pending_completed, 0u); + f.enqueue(2); + f.enqueue(3); + ASSERT_TRUE(f.dispatch()); + uint32_t live = 0; + for (const auto &tracker : f.local->mix_trackers) + live += tracker.task_id >= 0; + EXPECT_EQ(live, 4u); +} + +TEST(SchedulerMix, DeferredPolicyDisallowsStealingButStillAdmitsLocalMix) { + MixFixture f; + f.local->config.scheduler_count = 2; + SchedulerLocalState publisher; + f.storage.local_context(&f.storage.contexts[1], &publisher); + publisher.config.scheduler_index = 1; + publisher.config.scheduler_count = 2; + SchedulerReadyBatch batch; + ASSERT_TRUE(scheduler_ready_batch_append(f.storage.scheduler_state->base(), &publisher, 0, &batch, nullptr)); + ASSERT_TRUE( + scheduler_ready_batch_push(f.storage.scheduler_state->base(), &publisher, SCHEDULER_MIX_QUEUE, &batch, nullptr) + ); + ASSERT_TRUE(f.dispatch(false)); + EXPECT_EQ(f.local->dispatch_sequences[0], 0u); + f.enqueue(1); + ASSERT_TRUE(f.dispatch(false)); + EXPECT_EQ(f.local->slots[0][0].task_id, 1); + ASSERT_TRUE(f.dispatch(true)); + EXPECT_EQ(f.local->slots[0][1].task_id, 0); +} + +TEST(SchedulerMix, DirectSuccessorUsesOtherRemainingSlot) { + MixFixture f; + f.set_task(2, 1); + SchedulerRefillCandidates candidates; + candidates.tasks[0][0] = 2; + f.enqueue(0); + ASSERT_TRUE(f.dispatch()); + ASSERT_TRUE(scheduler_publish_refill_candidates( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &candidates, nullptr, 0, + f.storage.ssbuf_region + )); + EXPECT_EQ(f.local->slots[0][0].task_id, 0); + EXPECT_EQ(f.local->slots[0][1].task_id, 2); + EXPECT_EQ(f.local->slots[0][1].generation, 2u); +} + +TEST(SchedulerMix, CompletionTracesUseLogicalSubtaskRatherThanPhysicalLane) { + MixFixture f; + f.set_task(0, 3); + f.enqueue(0); + bool progress = false; + ASSERT_TRUE(scheduler_fill_cluster_mix_slots( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, &f.cursor, nullptr, + SCHEDULER_PROFILING_SCHED_PHASES_LEVEL, true, f.storage.ssbuf_region, &progress + )); + for (uint32_t lane : {0, 2}) { + auto &source = f.storage.ssbuf_region->lanes[lane].traces[0].payload; + source.kernel_start_cycles = 100 + lane; + source.kernel_end_cycles = 200 + lane; + ASSERT_TRUE(scheduler_service_cluster_completion_slot( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, lane, 0, + f.local->slots[lane][0].generation, nullptr, nullptr, nullptr, nullptr, + SCHEDULER_PROFILING_SCHED_PHASES_LEVEL, nullptr, nullptr, f.storage.ssbuf_region + )); + } + auto *traces = + scheduler_state_at(f.storage.scheduler_state->base(), f.storage.layout.trace_cells_offset); + EXPECT_EQ(traces[0].kernel_start_cycles, 100u); + EXPECT_EQ(traces[1].kernel_start_cycles, 102u); + EXPECT_EQ(traces[1].worker_id, 2u); + EXPECT_EQ(traces[1].valid, 1u); + EXPECT_EQ(traces[2].valid, 0u); + EXPECT_EQ(f.local->pending_completed, 1u); +} + +TEST(SchedulerDispatch, SequenceExhaustionDoesNotPublishOrWrap) { + MixFixture f; + f.set_task(0, 1); + f.local->dispatch_sequences[0] = UINT32_MAX; + f.local->slots[0][0].state = SchedulerDispatchSlotState::FILLING; + SchedulerReadyClaim ready; + ready.task_id = 0; + const auto metadata = + scheduler_load_dispatch_metadata(f.storage.scheduler_state->base(), f.local, ready.task_id, 0); + EXPECT_FALSE(scheduler_prepare_dispatch_slot( + f.graph.graph(), f.storage.scheduler_state->base(), f.local, f.storage.run_control, {0, 0, 0, 0}, ready, + metadata, 0, f.storage.ssbuf_region, 0 + )); + EXPECT_EQ(f.local->dispatch_sequences[0], UINT32_MAX); + EXPECT_EQ(f.storage.ssbuf_region->lanes[0].dispatch[0].publication, 0u); +} diff --git a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_ready.cpp b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_ready.cpp index aae3184ab1..5033c8ab62 100644 --- a/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_ready.cpp +++ b/tests/ut/cpp/a5/runtime/host_build_graph/test_hbg_scheduler_ready.cpp @@ -541,6 +541,18 @@ TEST(SchedulerReadyInbox, BootstrapPublishesIndependentDirectoryShards) { EXPECT_EQ(directory->core_types[1][1].bits, (UINT64_C(1) << 0) | (UINT64_C(1) << 6)); } +TEST(SchedulerReadyInbox, BootstrapPublishesMixDirectoryWhenPresent) { + FixtureStorage storage(1, 2); + auto *local = storage.local_context(&storage.contexts[0]); + local->has_mix = true; + auto *directory = scheduler_ready_directory_at(storage.scheduler_state->base(), local); + directory->bootstrap_ready_types[1] = UINT64_C(1) << SCHEDULER_MIX_QUEUE; + ASSERT_TRUE(scheduler_bootstrap_ready_directory_publish(storage.scheduler_state->base(), local, 2)); + EXPECT_EQ(directory->core_types[SCHEDULER_MIX_QUEUE][0].bits, UINT64_C(1) << 1); + EXPECT_EQ(directory->core_types[0][0].bits, 0u); + EXPECT_EQ(directory->core_types[1][0].bits, 0u); +} + TEST(SchedulerReadyInbox, SparseDirectoryWrapsWithinShard) { FixtureStorage storage(2, 14); GraphBuffer graph(2); @@ -694,13 +706,13 @@ TEST(SchedulerDispatch, RejectsUnknownTargetCoreType) { EXPECT_EQ(storage.run_control->error_site, static_cast(SchedulerErrorSite::DISPATCH_INVALID_SHAPE)); } -TEST(SchedulerDispatch, WrapsGenerationAndRejectsZeroCallable) { +TEST(SchedulerDispatch, DispatchSequenceIsIndependentOfSlotGenerationAndRejectsZeroCallable) { FixtureStorage storage(2, 1); GraphBuffer graph(2); graph.executable(0, 0); graph.executable(1, 0); storage.contexts[0].core_type = static_cast(CoreType::AIC); - SchedulerFreeSlotClaim slot_claim{0, 0, UINT32_MAX, 0}; + SchedulerFreeSlotClaim slot_claim{0, 0, 0, 0}; SchedulerReadyClaim ready_claim{}; ready_claim.task_id = 0; @@ -712,7 +724,7 @@ TEST(SchedulerDispatch, WrapsGenerationAndRejectsZeroCallable) { EXPECT_EQ(storage.scheduler_local_state.slots[0][0].generation, 1u); uint32_t pending_slot = UINT32_MAX; uint64_t publication = 0; - ASSERT_TRUE(scheduler_local_ready_pop(&storage.scheduler_local_state, 0, &pending_slot, &publication)); + ASSERT_TRUE(scheduler_local_ready_pop(&storage.scheduler_local_state, &pending_slot, &publication)); EXPECT_EQ(pending_slot, 0u); EXPECT_EQ(publication, scheduler_dispatch_publication(1, SchedulerDispatchSlotState::READY)); diff --git a/tests/ut/cpp/common/host_build_graph/test_hbg_bind_ledger.cpp b/tests/ut/cpp/common/host_build_graph/test_hbg_bind_ledger.cpp index 42bcd5bcca..54da047929 100644 --- a/tests/ut/cpp/common/host_build_graph/test_hbg_bind_ledger.cpp +++ b/tests/ut/cpp/common/host_build_graph/test_hbg_bind_ledger.cpp @@ -1024,11 +1024,11 @@ TEST_F(HbgBindLedgerTest, SchedulerModeChangesPublishOnlyThisRunsSources) { const bool a5 = std::strcmp(get_platform(), "a5sim") == 0; uint32_t resident_mode = 0; uint32_t graph_mode = 0; - // A5 selects resident for ordinary tasks and legacy for GRAPH/MIX. A2/A3 + // A5 selects resident for ordinary/Mix tasks and legacy for GRAPH. A2/A3 // uses AICPU scheduling throughout, with no resident scheduler region. for (TestOrchEntryFunc entry : {ordinary_orch_entry, recording_orch_entry, mixed_orch_entry, ordinary_orch_entry}) { SCOPED_TRACE(entry == recording_orch_entry ? "graph" : entry == mixed_orch_entry ? "mixed" : "ordinary"); - const bool resident = a5 && entry == ordinary_orch_entry; + const bool resident = a5 && entry != recording_orch_entry; const bool definitions = entry == recording_orch_entry; fake_.copy_count = 0; fake_.copies.clear(); @@ -1116,7 +1116,7 @@ TEST_F(HbgBindLedgerTest, SchedulerPublicationFailureAllowsFreshModeSelection) { for (TestOrchEntryFunc entry : {ordinary_orch_entry, recording_orch_entry, mixed_orch_entry}) { SCOPED_TRACE(entry == recording_orch_entry ? "graph" : entry == mixed_orch_entry ? "mixed" : "ordinary"); const size_t regions = 1 + static_cast(entry == recording_orch_entry) + - static_cast(a5 && entry == ordinary_orch_entry); + static_cast(a5 && entry != recording_orch_entry); for (size_t failure = 1; failure <= regions; ++failure) { SCOPED_TRACE(failure); fake_.copy_count = 0; @@ -1142,7 +1142,7 @@ TEST_F(HbgBindLedgerTest, SchedulerPublicationFailureAllowsFreshModeSelection) { EXPECT_EQ(runtime.dev.scheduler_bootstrap.worker_context_base, 0u); } fake_.fail_copy_on = 0; - const bool definitions = entry == ordinary_orch_entry; + const bool definitions = entry != recording_orch_entry; const bool resident = a5 && !definitions; eps_ = {definitions ? recording_orch_entry : ordinary_orch_entry, capture_orch_bind}; ASSERT_EQ(bind(runtime, args, nullptr, 0), 0);