From 3c49a1803f999fcdef2b9372cb63504f7b3bb4ee Mon Sep 17 00:00:00 2001 From: ChaoZheng109 Date: Mon, 13 Jul 2026 00:08:27 -0700 Subject: [PATCH] =?UTF-8?q?feat(runtime/pto2):=20predicated=20dispatch=20?= =?UTF-8?q?=E2=80=94=20scheduler=20evaluates=20a=20per-task=20predicate=20?= =?UTF-8?q?at=20the=20dispatch=20point?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a per-task dispatch predicate that is resolved NOT in orchestration but by the scheduler at the dispatch point, so a statically-materialized task can be retired without dispatch based on a value only a prior task produces — without the orchestrator stalling to read it. Motivation: MoE (DeepSeek-V4 expert_routed) reads its per-expert row count inside the incore kernel; orchestration only has the tensor address, not the value. So "skip empty experts" cannot be a value known at submit — it must be a predicate the scheduler evaluates later. Reading the value in orchestration stalls on wait_for_tensor_ready (the producer must finish first); deferring the read to the dispatch point overlaps it with other work (pipeline). Design: - PredicateOp {NONE,EQ,NE,GT,LT,GE,LE}. A DispatchPredicate {addr,target,elem_size, op} lives on PTO2TaskPayload at a fixed offset (cache line 9, before tensors, so it never moves when MAX_TENSOR_ARGS / MAX_SCALAR_ARGS change). AICore never reads it. - Arg carries an L0TaskPredicate {src tensor, indices, op, target}, set via args.set_predicate(pred); submit resolves it into the payload's DispatchPredicate (an absolute GM address + comparison). - The scheduler evaluates it at every ready-routing point (push_ready_routed and both route_ready_once copies): the task is ready => the predicate's producer has completed => the GM read is current, without the wait get_tensor_data pays. PASS => dispatch normally (SPMD: all blocks); FAIL => dummy_ready_queue, retired inline via on_task_complete/on_task_release (reused dep-only path). Predicated tasks are excluded from early-dispatch staging. - Fast path: an active_mask bit (PTO2_SUBTASK_FLAG_HAS_PREDICATE) gates the evaluation, so a task without a predicate is routed without dereferencing its payload or touching the predicate cache line. Contract: the predicate tensor's producer MUST be a dependency of the predicated task, so the value is current once the task becomes ready. Test: tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch drives gate_producer -> x_producer -> predicated clobber (depends on gate; predicate gate[0] > 0) -> consumer. The clobber writes a poison value if dispatched: gate=0 => predicate false => not dispatched => sentinel survives and the consumer still unlocks; gate=1 => dispatched => poison propagates. Passes on a2a3sim and onboard a2a3. Frontend (pypto pl.spmd(predicate=...)) lands separately. --- .../runtime/pto2_dispatch_payload.h | 5 +- .../runtime/pto_orchestrator.cpp | 23 ++++ .../runtime/pto_runtime2_types.h | 24 +++- .../runtime/pto_submit_types.h | 63 ++++++++- .../runtime/pto_types.h | 32 +++++ .../runtime/scheduler/pto_scheduler.h | 20 ++- .../runtime/pto_orchestrator.cpp | 23 ++++ .../runtime/pto_runtime2_types.h | 24 +++- .../runtime/pto_submit_types.h | 63 ++++++++- .../runtime/pto_types.h | 32 +++++ .../runtime/scheduler/pto_scheduler.h | 6 +- .../kernels/aic/kernel_clobber.cpp | 53 ++++++++ .../kernels/aic/kernel_copy_first.cpp | 57 ++++++++ .../kernels/aic/kernel_write_const.cpp | 51 +++++++ .../kernels/aic/kernel_write_gate.cpp | 55 ++++++++ .../predicated_dispatch_orch.cpp | 112 ++++++++++++++++ .../test_predicated_dispatch.py | 124 ++++++++++++++++++ .../kernels/aic/kernel_clobber.cpp | 53 ++++++++ .../kernels/aic/kernel_copy_first.cpp | 57 ++++++++ .../kernels/aic/kernel_write_const.cpp | 51 +++++++ .../kernels/aic/kernel_write_gate.cpp | 55 ++++++++ .../predicated_dispatch_orch.cpp | 112 ++++++++++++++++ .../test_predicated_dispatch.py | 124 ++++++++++++++++++ 23 files changed, 1189 insertions(+), 30 deletions(-) create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h index 899de63f0..e382b5588 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h @@ -66,8 +66,9 @@ static_assert( // PTO2TaskPayload layout drift fails the build rather than corrupting args[]. constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET = 0; constexpr uint32_t PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET = 4; -constexpr uint32_t PTO2_TASKPAYLOAD_TENSORS_OFFSET = 576; -constexpr uint32_t PTO2_TASKPAYLOAD_SCALARS_OFFSET = 4672; +// Cache line 9 (byte 576) holds the AICPU-only DispatchPredicate; tensors follow it. +constexpr uint32_t PTO2_TASKPAYLOAD_TENSORS_OFFSET = 640; +constexpr uint32_t PTO2_TASKPAYLOAD_SCALARS_OFFSET = 4736; constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_STRIDE = 128; // sizeof(Tensor) /** diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 7c868bcd2..aa754cbcc 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -952,6 +952,25 @@ static TaskOutputTensors submit_task_common( payload.init(args, result, prepared.alloc_result, layout); cur_slot_state.set_allow_early_resolve(args.allow_early_resolve()); + + // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address + // now so the scheduler can read it at the dispatch point with a single load, + // no Arg/Tensor access. Both branches write predicate.op explicitly because + // payload slots are ring-reused; op == NONE means "always dispatch". + { + const L0TaskPredicate &pred = args.predicate(); + if (pred.op != PredicateOp::NONE && pred.operand.tensor != nullptr && pred.operand.tensor->buffer.addr != 0) { + uint64_t elem_size = get_element_size(pred.operand.tensor->dtype); + uint64_t flat_offset = pred.operand.tensor->compute_flat_offset(pred.operand.indices, pred.operand.ndims); + payload.predicate.addr = pred.operand.tensor->buffer.addr + flat_offset * elem_size; + payload.predicate.target = pred.target; + payload.predicate.elem_size = static_cast(elem_size); + payload.predicate.op = pred.op; + } else { + payload.predicate.addr = 0; + payload.predicate.op = PredicateOp::NONE; + } + } #if SIMPLER_DFX if (is_dump_args_enabled()) { if (args.scalar_count() > 0) { @@ -1069,6 +1088,10 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k active_mask.set_sync_start(); } + if (args.predicate().op != PredicateOp::NONE) { + active_mask.set_has_predicate(); + } + return submit_task_common( orch, args, active_mask, normalized.aic_kernel_id, normalized.aiv0_kernel_id, normalized.aiv1_kernel_id ); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index 01fd6704e..887d351a1 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -305,9 +305,15 @@ struct PTO2TaskPayload { // reinitialization. READY records that producer release observed OWNER; // only cancellation clears OWNER during the current task lifetime. std::atomic early_sync_drain_state{PTO2_EARLY_SYNC_DRAIN_NONE}; - // === Cache lines 9-72 (4096B) — tensors (alignas(64) forces alignment) === + // === Cache line 9 (byte 576) — dispatch predicate (AICPU-only) === + // Offset is a fixed 576, independent of MAX_TENSOR_ARGS / MAX_SCALAR_ARGS. + // AICore never reads it — args are materialized from the tensor_count / tensors + // / scalars offsets only. Resolved at submit; evaluated by the scheduler at + // dispatch. + alignas(64) DispatchPredicate predicate; + // === Cache lines 10-73 (4096B) — tensors (alignas(64) forces alignment) === Tensor tensors[MAX_TENSOR_ARGS]; - // === Cache lines 73-74 (128B) — scalars === + // === Cache lines 74-75 (128B) — scalars === uint64_t scalars[MAX_SCALAR_ARGS]; // Layout verification (size checks that don't need offsetof). @@ -399,14 +405,20 @@ static_assert(offsetof(PTO2TaskPayload, fanin_spill_pool) == 16, "spill pool poi static_assert( offsetof(PTO2TaskPayload, fanin_inline_slot_states) == 24, "inline fanin array must follow spill metadata" ); -static_assert(offsetof(PTO2TaskPayload, tensors) == 576, "tensors must start at byte 576 (cache line 9)"); static_assert( - offsetof(PTO2TaskPayload, scalars) == 576 + MAX_TENSOR_ARGS * sizeof(Tensor), + offsetof(PTO2TaskPayload, predicate) == 576, + "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" +); +static_assert( + offsetof(PTO2TaskPayload, tensors) == 640, "tensors must start at byte 640 (cache line 10, after predicate)" +); +static_assert( + offsetof(PTO2TaskPayload, scalars) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor), "scalars must immediately follow tensors" ); static_assert( - sizeof(PTO2TaskPayload) == 576 + MAX_TENSOR_ARGS * sizeof(Tensor) + MAX_SCALAR_ARGS * sizeof(uint64_t), - "PTO2TaskPayload size must stay on the baseline cache-line footprint" + sizeof(PTO2TaskPayload) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor) + MAX_SCALAR_ARGS * sizeof(uint64_t), + "PTO2TaskPayload size = metadata(576) + predicate cache line(64) + tensors + scalars" ); /** diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h index ae3f42991..56e6cab04 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h @@ -39,10 +39,60 @@ enum class PTO2SubtaskSlot : uint8_t { /** * Subtask mask bits (for ActiveMask) */ -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0); // 0x1 -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1); // 0x2 -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2); // 0x4 -inline constexpr uint8_t PTO2_SUBTASK_FLAG_SYNC_START = (1u << 3); // 0x8: all blocks must launch atomically +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0); // 0x1 +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1); // 0x2 +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2); // 0x4 +inline constexpr uint8_t PTO2_SUBTASK_FLAG_SYNC_START = (1u << 3); // 0x8: all blocks must launch atomically +inline constexpr uint8_t PTO2_SUBTASK_FLAG_HAS_PREDICATE = (1u << 4); // 0x10: task carries a dispatch predicate + +// Dispatch-predicate comparison operator. The scheduler evaluates the predicate +// at the dispatch point — the task is ready (fanin satisfied), so the predicate +// address' producer has completed and the value read is current, without the +// wait_for_tensor_ready() stall that get_tensor_data() pays in orchestration. +// PASS => dispatch normally; FAIL => retire inline via the dep-only path. +enum class PredicateOp : uint8_t { NONE = 0, EQ, NE, GT, LT, GE, LE }; + +// Resolved dispatch predicate stored on a task's payload (AICPU-side only): the +// absolute GM address of the predicate element + the comparison. op == NONE +// means "no predicate — always dispatch". Populated at submit from an +// L0TaskPredicate; evaluated by the scheduler at the dispatch point via pass(). +// Layout is 18 bytes (8-aligned). +struct DispatchPredicate { + uint64_t addr{0}; // absolute GM address of the predicate element (0 when op == NONE) + int64_t target{0}; // value compared against + uint8_t elem_size{0}; // width of the predicate element in bytes (1/2/4/8) + PredicateOp op{PredicateOp::NONE}; + + // true => dispatch, false => retire inline. Reads elem_size bytes at addr and + // sign-extends to 64 bits before comparing to target. Safe to call only when + // the owning task is ready (its producer has written the value). + bool pass() const { + if (op == PredicateOp::NONE) return true; + int64_t v = 0; + __builtin_memcpy(&v, reinterpret_cast(addr), elem_size); + uint32_t bits = static_cast(elem_size) * 8u; + if (bits < 64u) { + int64_t shift = static_cast(64u - bits); + v = (v << shift) >> shift; + } + switch (op) { + case PredicateOp::EQ: + return v == target; + case PredicateOp::NE: + return v != target; + case PredicateOp::GT: + return v > target; + case PredicateOp::LT: + return v < target; + case PredicateOp::GE: + return v >= target; + case PredicateOp::LE: + return v <= target; + default: + return true; + } + } +}; /** * Resource shape — classifies a MixedKernels into one of 3 scheduling buckets. @@ -86,6 +136,9 @@ class ActiveMask { bool requires_sync_start() const { return (raw_ & PTO2_SUBTASK_FLAG_SYNC_START) != 0; } + // True when the task carries a dispatch predicate. + bool has_predicate() const { return (raw_ & PTO2_SUBTASK_FLAG_HAS_PREDICATE) != 0; } + PTO2ResourceShape to_shape() const { uint8_t cmask = core_mask(); if (cmask == 0) return PTO2ResourceShape::DUMMY; @@ -97,6 +150,8 @@ class ActiveMask { void set_sync_start() { raw_ |= PTO2_SUBTASK_FLAG_SYNC_START; } + void set_has_predicate() { raw_ |= PTO2_SUBTASK_FLAG_HAS_PREDICATE; } + bool operator==(ActiveMask other) const { return raw_ == other.raw_; } bool operator!=(ActiveMask other) const { return raw_ != other.raw_; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index 190323ae8..7f747fcd4 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -198,6 +198,27 @@ class TensorRef { * TaskOutputTensors outs = rt_submit_aic_task(kernel_id, args); * const Tensor& y = outs.get_ref(0); */ + +// Operand of a dispatch predicate (L0 layer): locates one element of a tensor — +// tensor + ndims + indices, mirroring get_tensor_data. The tensor is borrowed and +// must outlive submit; its buffer must be allocated by then, and its producer must +// be a dependency of the predicated task so the value is current at dispatch. +struct L0PredicateOperand { + const Tensor *tensor{nullptr}; + uint32_t ndims{0}; + uint32_t indices[MAX_TENSOR_DIMS]{}; +}; + +// Dispatch predicate carried on an Arg: operand OP target (e.g. count[i] > 0). +// op == NONE means "no predicate — always dispatch". Submit resolves the operand +// into the payload's DispatchPredicate (an absolute GM address). Read in-process; +// never crosses the wire. +struct L0TaskPredicate { + L0PredicateOperand operand; + PredicateOp op{PredicateOp::NONE}; + int64_t target{0}; +}; + template struct Arg : TaskArgsTpl { using Base = TaskArgsTpl; @@ -231,6 +252,16 @@ struct Arg : TaskArgsTpl { void set_allow_early_resolve(bool v = true) { allow_early_resolve_ = v; } bool allow_early_resolve() const { return allow_early_resolve_; } + // Dispatch predicate (codegen-author set; default op == NONE = always + // dispatch). A FALSE result at the dispatch point retires the task inline + // through the dep-only path — never dispatched to an AICore — while still + // resolving fanin/fanout so consumers unlock. The predicate tensor's producer + // MUST be a dependency of this task so the value is current when the task + // becomes ready. Read in-process; never crosses the wire. + L0TaskPredicate predicate_; + void set_predicate(const L0TaskPredicate &pred) { predicate_ = pred; } + const L0TaskPredicate &predicate() const { return predicate_; } + void clear() { Base::clear(); #if SIMPLER_DFX @@ -239,6 +270,7 @@ struct Arg : TaskArgsTpl { explicit_deps_ = nullptr; explicit_dep_count_ = 0; allow_early_resolve_ = false; + predicate_ = L0TaskPredicate{}; } void reset() { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 3182d6b50..93bddf21c 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -526,12 +526,15 @@ struct PTO2SchedulerState { // Inline hot-path methods // ========================================================================= - // Route a ready slot to the right global queue. Dummy tasks (empty active_mask) - // live in dummy_ready_queue; a ready sync_start cohort goes to the per-shape - // ready_sync_queues[] (drained as Tier-0); everything else to ready_queues[]. + // Route a ready slot to the right global queue. Dep-only tasks — DUMMY-shaped + // (empty active_mask) or a task whose dispatch predicate fails — live in + // dummy_ready_queue and are retired inline; a ready sync_start cohort goes to + // the per-shape ready_sync_queues[] (drained as Tier-0); everything else to + // ready_queues[]. void push_ready_routed(PTO2TaskSlotState *slot_state) { PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { + if (shape == PTO2ResourceShape::DUMMY || + (slot_state->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) { dummy_ready_queue.push(slot_state); } else if (slot_state->active_mask.requires_sync_start()) { ready_sync_queues[static_cast(shape)].push(slot_state); @@ -896,6 +899,7 @@ struct PTO2SchedulerState { p.unlock_fanout(); for (; edge != nullptr; edge = edge->next) { PTO2TaskSlotState *c = edge->slot_state; + if (c->active_mask.has_predicate()) continue; // predicated consumers never early-dispatch // Compare to fanin_actual_count (the real producer-edge count), NOT // fanin_count: fanin_count = fanin_actual_count + 1 (a self/wiring +1 that // ready_fanin gets but dispatch_fanin does not). dispatch_fanin starts at @@ -944,6 +948,8 @@ struct PTO2SchedulerState { }; inline bool try_early_dispatch_release(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { + // A predicated task is evaluated at the ready point, never early-dispatched. + if (slot_state.active_mask.has_predicate()) return false; // Never staged => CAS NONE->DISPATCHED wins => dispatch normally. uint8_t expect = PTO2_EARLY_DISPATCH_NONE; if (slot_state.payload->early_dispatch_state.compare_exchange_strong( @@ -1015,7 +1021,8 @@ struct PTO2SchedulerState { } PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { + if (shape == PTO2ResourceShape::DUMMY || + (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { dummy_ready_queue.push(&slot_state); } else if (slot_state.active_mask.requires_sync_start()) { ready_sync_queues[static_cast(shape)].push(&slot_state); @@ -1055,7 +1062,8 @@ struct PTO2SchedulerState { } PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { + if (shape == PTO2ResourceShape::DUMMY || + (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { dummy_ready_queue.push(&slot_state, atomic_count, push_wait); } else if (slot_state.active_mask.requires_sync_start()) { ready_sync_queues[static_cast(shape)].push(&slot_state, atomic_count, push_wait); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index a02488869..3c474d662 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -935,6 +935,25 @@ static TaskOutputTensors submit_task_common( } payload.init(args, result, prepared.alloc_result, layout); + + // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address + // now so the scheduler can read it at the dispatch point with a single load, + // no Arg/Tensor access. Both branches write predicate.op explicitly because + // payload slots are ring-reused; op == NONE means "always dispatch". + { + const L0TaskPredicate &pred = args.predicate(); + if (pred.op != PredicateOp::NONE && pred.operand.tensor != nullptr && pred.operand.tensor->buffer.addr != 0) { + uint64_t elem_size = get_element_size(pred.operand.tensor->dtype); + uint64_t flat_offset = pred.operand.tensor->compute_flat_offset(pred.operand.indices, pred.operand.ndims); + payload.predicate.addr = pred.operand.tensor->buffer.addr + flat_offset * elem_size; + payload.predicate.target = pred.target; + payload.predicate.elem_size = static_cast(elem_size); + payload.predicate.op = pred.op; + } else { + payload.predicate.addr = 0; + payload.predicate.op = PredicateOp::NONE; + } + } #if SIMPLER_DFX if (is_dump_args_enabled()) { if (args.scalar_count() > 0) { @@ -1049,6 +1068,10 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k active_mask.set_sync_start(); } + if (args.predicate().op != PredicateOp::NONE) { + active_mask.set_has_predicate(); + } + return submit_task_common( orch, args, active_mask, normalized.aic_kernel_id, normalized.aiv0_kernel_id, normalized.aiv1_kernel_id ); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index 5cee99ec3..e821d3e52 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -220,9 +220,15 @@ struct PTO2TaskPayload { int32_t fanin_spill_start{0}; // Linear start index in fanin spill pool (0 = no spill) PTO2FaninPool *fanin_spill_pool{nullptr}; PTO2TaskSlotState *fanin_inline_slot_states[PTO2_FANIN_INLINE_CAP]; - // === Cache lines 9-72 (4096B) — tensors (alignas(64) forces alignment) === + // === Cache line 9 (byte 576) — dispatch predicate (AICPU-only) === + // Offset is a fixed 576, independent of MAX_TENSOR_ARGS / MAX_SCALAR_ARGS. + // AICore never reads it — args are materialized from the tensor_count / tensors + // / scalars offsets only. Resolved at submit; evaluated by the scheduler at + // dispatch. + alignas(64) DispatchPredicate predicate; + // === Cache lines 10-73 (4096B) — tensors (alignas(64) forces alignment) === Tensor tensors[MAX_TENSOR_ARGS]; - // === Cache lines 73-74 (128B) — scalars === + // === Cache lines 74-75 (128B) — scalars === uint64_t scalars[MAX_SCALAR_ARGS]; // Layout verification (size checks that don't need offsetof). @@ -270,14 +276,20 @@ static_assert(offsetof(PTO2TaskPayload, fanin_spill_pool) == 16, "spill pool poi static_assert( offsetof(PTO2TaskPayload, fanin_inline_slot_states) == 24, "inline fanin array must follow spill metadata" ); -static_assert(offsetof(PTO2TaskPayload, tensors) == 576, "tensors must start at byte 576 (cache line 9)"); static_assert( - offsetof(PTO2TaskPayload, scalars) == 576 + MAX_TENSOR_ARGS * sizeof(Tensor), + offsetof(PTO2TaskPayload, predicate) == 576, + "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" +); +static_assert( + offsetof(PTO2TaskPayload, tensors) == 640, "tensors must start at byte 640 (cache line 10, after predicate)" +); +static_assert( + offsetof(PTO2TaskPayload, scalars) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor), "scalars must immediately follow tensors" ); static_assert( - sizeof(PTO2TaskPayload) == 576 + MAX_TENSOR_ARGS * sizeof(Tensor) + MAX_SCALAR_ARGS * sizeof(uint64_t), - "PTO2TaskPayload size must stay on the baseline cache-line footprint" + sizeof(PTO2TaskPayload) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor) + MAX_SCALAR_ARGS * sizeof(uint64_t), + "PTO2TaskPayload size = metadata(576) + predicate cache line(64) + tensors + scalars" ); /** diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h index f7dba9232..360ae47b9 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h @@ -39,10 +39,60 @@ enum class PTO2SubtaskSlot : uint8_t { /** * Subtask mask bits (for ActiveMask) */ -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0); // 0x1 -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1); // 0x2 -inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2); // 0x4 -inline constexpr uint8_t PTO2_SUBTASK_FLAG_SYNC_START = (1u << 3); // 0x8: all blocks must launch atomically +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIC = (1u << 0); // 0x1 +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV0 = (1u << 1); // 0x2 +inline constexpr uint8_t PTO2_SUBTASK_MASK_AIV1 = (1u << 2); // 0x4 +inline constexpr uint8_t PTO2_SUBTASK_FLAG_SYNC_START = (1u << 3); // 0x8: all blocks must launch atomically +inline constexpr uint8_t PTO2_SUBTASK_FLAG_HAS_PREDICATE = (1u << 4); // 0x10: task carries a dispatch predicate + +// Dispatch-predicate comparison operator. The scheduler evaluates the predicate +// at the dispatch point — the task is ready (fanin satisfied), so the predicate +// address' producer has completed and the value read is current, without the +// wait_for_tensor_ready() stall that get_tensor_data() pays in orchestration. +// PASS => dispatch normally; FAIL => retire inline via the dep-only path. +enum class PredicateOp : uint8_t { NONE = 0, EQ, NE, GT, LT, GE, LE }; + +// Resolved dispatch predicate stored on a task's payload (AICPU-side only): the +// absolute GM address of the predicate element + the comparison. op == NONE +// means "no predicate — always dispatch". Populated at submit from an +// L0TaskPredicate; evaluated by the scheduler at the dispatch point via pass(). +// Layout is 18 bytes (8-aligned). +struct DispatchPredicate { + uint64_t addr{0}; // absolute GM address of the predicate element (0 when op == NONE) + int64_t target{0}; // value compared against + uint8_t elem_size{0}; // width of the predicate element in bytes (1/2/4/8) + PredicateOp op{PredicateOp::NONE}; + + // true => dispatch, false => retire inline. Reads elem_size bytes at addr and + // sign-extends to 64 bits before comparing to target. Safe to call only when + // the owning task is ready (its producer has written the value). + bool pass() const { + if (op == PredicateOp::NONE) return true; + int64_t v = 0; + __builtin_memcpy(&v, reinterpret_cast(addr), elem_size); + uint32_t bits = static_cast(elem_size) * 8u; + if (bits < 64u) { + int64_t shift = static_cast(64u - bits); + v = (v << shift) >> shift; + } + switch (op) { + case PredicateOp::EQ: + return v == target; + case PredicateOp::NE: + return v != target; + case PredicateOp::GT: + return v > target; + case PredicateOp::LT: + return v < target; + case PredicateOp::GE: + return v >= target; + case PredicateOp::LE: + return v <= target; + default: + return true; + } + } +}; /** * Resource shape — classifies a MixedKernels into one of 3 scheduling buckets. @@ -86,6 +136,9 @@ class ActiveMask { bool requires_sync_start() const { return (raw_ & PTO2_SUBTASK_FLAG_SYNC_START) != 0; } + // True when the task carries a dispatch predicate. + bool has_predicate() const { return (raw_ & PTO2_SUBTASK_FLAG_HAS_PREDICATE) != 0; } + PTO2ResourceShape to_shape() const { uint8_t cmask = core_mask(); if (cmask == 0) return PTO2ResourceShape::DUMMY; @@ -97,6 +150,8 @@ class ActiveMask { void set_sync_start() { raw_ |= PTO2_SUBTASK_FLAG_SYNC_START; } + void set_has_predicate() { raw_ |= PTO2_SUBTASK_FLAG_HAS_PREDICATE; } + bool operator==(ActiveMask other) const { return raw_ == other.raw_; } bool operator!=(ActiveMask other) const { return raw_ != other.raw_; } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index db8247c7c..9d819852f 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -196,6 +196,27 @@ class TensorRef { * TaskOutputTensors outs = rt_submit_aic_task(kernel_id, args); * const Tensor& y = outs.get_ref(0); */ + +// Operand of a dispatch predicate (L0 layer): locates one element of a tensor — +// tensor + ndims + indices, mirroring get_tensor_data. The tensor is borrowed and +// must outlive submit; its buffer must be allocated by then, and its producer must +// be a dependency of the predicated task so the value is current at dispatch. +struct L0PredicateOperand { + const Tensor *tensor{nullptr}; + uint32_t ndims{0}; + uint32_t indices[MAX_TENSOR_DIMS]{}; +}; + +// Dispatch predicate carried on an Arg: operand OP target (e.g. count[i] > 0). +// op == NONE means "no predicate — always dispatch". Submit resolves the operand +// into the payload's DispatchPredicate (an absolute GM address). Read in-process; +// never crosses the wire. +struct L0TaskPredicate { + L0PredicateOperand operand; + PredicateOp op{PredicateOp::NONE}; + int64_t target{0}; +}; + template struct Arg : TaskArgsTpl { using Base = TaskArgsTpl; @@ -220,6 +241,16 @@ struct Arg : TaskArgsTpl { const char *error_msg{nullptr}; PTO2LaunchSpec launch_spec; // SPMD launch parameters (block_num, etc.) + // Dispatch predicate (codegen-author set; default op == NONE = always + // dispatch). A FALSE result at the dispatch point retires the task inline + // through the dep-only path — never dispatched to an AICore — while still + // resolving fanin/fanout so consumers unlock. The predicate tensor's producer + // MUST be a dependency of this task so the value is current when the task + // becomes ready. Read in-process; never crosses the wire. + L0TaskPredicate predicate_; + void set_predicate(const L0TaskPredicate &pred) { predicate_ = pred; } + const L0TaskPredicate &predicate() const { return predicate_; } + void clear() { Base::clear(); #if SIMPLER_DFX @@ -227,6 +258,7 @@ struct Arg : TaskArgsTpl { #endif explicit_deps_ = nullptr; explicit_dep_count_ = 0; + predicate_ = L0TaskPredicate{}; } void reset() { diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 2089a50eb..4b69e5c6f 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -506,7 +506,8 @@ struct PTO2SchedulerState { // per-shape ready_queues[]. void push_ready_routed(PTO2TaskSlotState *slot_state) { PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { + if (shape == PTO2ResourceShape::DUMMY || + (slot_state->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) { dummy_ready_queue.push(slot_state); } else { ready_queues[static_cast(shape)].push(slot_state); @@ -645,7 +646,8 @@ struct PTO2SchedulerState { // ready_queues[]. Use the profiling-aware push so atomic_count / push_wait // stay consistent with the non-dummy path. PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { + if (shape == PTO2ResourceShape::DUMMY || + (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { dummy_ready_queue.push(&slot_state, atomic_count, push_wait); } else { ready_queues[static_cast(shape)].push(&slot_state, atomic_count, push_wait); diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp new file mode 100644 index 000000000..a3c7b389c --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp @@ -0,0 +1,53 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Clobber kernel for the predicated_dispatch scene tests: overwrites args[0][0] + * with a poison value (999.0f). It is the body of the predicated task — if the + * predicate correctly suppresses dispatch, this kernel never runs and the + * producer's sentinel (42.0f) survives; if it ever runs, the downstream copy + * reads 999.0f and the test fails. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = 999.0f; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp new file mode 100644 index 000000000..dd6804f00 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp @@ -0,0 +1,57 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Consumer kernel for the predicated_dispatch scene tests: copies args[0][0] -> args[1][0]. + * + * The producer writes the sentinel (42.0f) to args[0]; a predicated clobber sits + * between the producer and this consumer. When the predicate suppresses the + * clobber's dispatch, args[0] is untouched and args[1][0] equals 42.0f; when the + * clobber dispatches, args[1][0] is the poison value instead. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *in_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + + __gm__ float *in = reinterpret_cast<__gm__ float *>(in_tensor->buffer.addr) + in_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = in[0]; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp new file mode 100644 index 000000000..5b579cfaa --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp @@ -0,0 +1,51 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Writes a fixed sentinel (42.0f) to args[0][0]. The producer in the + * predicated_dispatch scene tests; a downstream consumer verifies the dependency + * chain (producer -> predicated clobber -> consumer) propagates the value intact. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = 42.0f; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp new file mode 100644 index 000000000..2c00413f6 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp @@ -0,0 +1,55 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Producer of the predicate value for the predicated_dispatch scene tests. + * Writes the scalar arg into gate[0] (INT32). A downstream task carries a + * dispatch predicate on gate[0]; the scheduler reads this value at the dispatch + * point (after this producer has completed) to decide whether to dispatch. + * + * args[0] = gate tensor (INOUT, INT32); args[1] = scalar gate value. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + int64_t gate_value = args[1]; // scalar arg follows the tensor args + + __gm__ int32_t *out = reinterpret_cast<__gm__ int32_t *>(out_tensor->buffer.addr) + out_tensor->start_offset; + out[0] = static_cast(gate_value); + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp new file mode 100644 index 000000000..9cba6a725 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp @@ -0,0 +1,112 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * predicated_dispatch orchestration: delayed-evaluation dispatch predicate. + * + * The predicate value is produced by a prior task and read by the scheduler at + * the dispatch point — never in orchestration — so the orchestrator does not + * stall on it: + * + * gate_producer (WRITE_GATE) writes gate[0] = gate_value (INT32) + * x_producer (WRITE_CONST) writes X[0] = 42.0 + * clobber (CLOBBER) would write X[0] = 999.0, but carries + * set_predicate(gate[0] > 0) and depends on + * gate_producer. The scheduler reads gate[0] when + * the clobber becomes ready (gate already written) + * and dispatches only if gate[0] > 0; otherwise the + * task is retired inline without dispatch. + * consumer (COPY_FIRST) copies X[0] -> Y[0] + * + * case=1: gate_value = 0 -> predicate FALSE -> clobber not dispatched -> + * X stays 42.0 -> Y = 42.0. Proves non-dispatch + consumer still unlocks. + * case=2: gate_value = 1 -> predicate TRUE -> clobber dispatched -> + * X = 999.0 -> Y = 999.0. Proves the dispatch path is taken. + * + * Args layout: [X, Y, gate] + case scalar. + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_WRITE_CONST 0 +#define FUNC_COPY_FIRST 1 +#define FUNC_CLOBBER 2 +#define FUNC_WRITE_GATE 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 4, // 3 tensors + 1 case scalar + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_X = orch_args.tensor(0).ref(); + const Tensor &ext_Y = orch_args.tensor(1).ref(); + const Tensor &ext_gate = orch_args.tensor(2).ref(); + + uint64_t case_id = orch_args.scalar(0); + if (case_id != 1 && case_id != 2) { + rt_report_fatal(PTO2_ERROR_INVALID_ARGS, "unsupported case_id=%llu", static_cast(case_id)); + return; + } + // case 1 => gate 0 (predicate FALSE, skip); case 2 => gate 1 (predicate TRUE, dispatch). + int64_t gate_value = (case_id == 1) ? 0 : 1; + + // gate producer: gate[0] = gate_value + PTO2TaskId gate_tid; + { + L0TaskArgs args; + args.add_inout(ext_gate); + args.add_scalar(gate_value); + gate_tid = rt_submit_aic_task(FUNC_WRITE_GATE, args).task_id(); + } + + // x producer: X[0] = 42.0 + { + L0TaskArgs args; + args.add_inout(ext_X); + rt_submit_aic_task(FUNC_WRITE_CONST, args); + } + + // predicated clobber: would write X[0] = 999.0 if dispatched. Depends on the + // gate producer so gate[0] is written by the time this task is ready; the + // scheduler reads gate[0] at the dispatch point and dispatches only if > 0. + { + L0TaskArgs args; + args.add_inout(ext_X); + PTO2TaskId deps[] = {gate_tid}; + args.set_dependencies(deps, 1); + // predicate: gate[0] > 0 (operand op target), built level by level. + L0TaskPredicate pred; + pred.operand.tensor = &ext_gate; + pred.operand.ndims = 1; + pred.operand.indices[0] = 0; + pred.op = PredicateOp::GT; + pred.target = 0; + args.set_predicate(pred); + rt_submit_aic_task(FUNC_CLOBBER, args); + } + + // consumer: Y[0] = X[0] + { + L0TaskArgs args; + args.add_input(ext_X); + args.add_inout(ext_Y); + rt_submit_aic_task(FUNC_COPY_FIRST, args); + } +} + +} // extern "C" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py new file mode 100644 index 000000000..01d895ded --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py @@ -0,0 +1,124 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""predicated_dispatch: a dispatch predicate is evaluated by the scheduler at the +dispatch point (delayed evaluation), not in orchestration. + +The predicate value is produced by a prior task; the scheduler reads it only +when the predicated task becomes ready (its producer done), so orchestration +never stalls on it. Chain: + + gate_producer writes gate[0] = gate_value (INT32) + x_producer writes X[0] = 42.0 + clobber would write X[0] = 999.0, but carries set_predicate(gate[0] > 0) + and depends on gate_producer; the scheduler reads gate[0] at the + dispatch point and dispatches only if gate[0] > 0. + consumer copies X[0] -> Y[0] + + case=1 (gate_value = 0): predicate FALSE -> clobber NOT dispatched -> X stays + 42.0 -> Y = 42.0. Proves non-dispatch AND that the retired task still + unlocks the consumer. + case=2 (gate_value = 1): predicate TRUE -> clobber dispatched -> X = 999.0 -> + Y = 999.0. Proves the dispatch path is taken when the predicate holds. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +SENTINEL = 42.0 +POISON = 999.0 # what the clobber writes if the predicate lets it dispatch +INIT_VAL = -1.0 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPredicatedDispatch(SceneTestCase): + """predicated_dispatch: scheduler evaluates the dispatch predicate at dispatch time.""" + + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/predicated_dispatch_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT, D.INOUT], # X, Y, gate + }, + "incores": [ + { + "func_id": 0, + "name": "WRITE_CONST", + "source": "kernels/aic/kernel_write_const.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "COPY_FIRST", + "source": "kernels/aic/kernel_copy_first.cpp", + "core_type": "aic", + "signature": [D.IN, D.INOUT], + }, + { + "func_id": 2, + "name": "CLOBBER", + "source": "kernels/aic/kernel_clobber.cpp", + "core_type": "aic", + # Body of the predicated task; runs only if the predicate holds. + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "WRITE_GATE", + "source": "kernels/aic/kernel_write_gate.cpp", + "core_type": "aic", + # One INOUT tensor (gate); the gate value rides as a trailing scalar. + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "PredicateFalseSkips", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 2, "block_dim": 1}, + "params": {"case": 1}, + }, + { + "name": "PredicateTrueDispatches", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 2, "block_dim": 1}, + "params": {"case": 2}, + }, + ] + + def generate_args(self, params): + x = torch.full((16,), INIT_VAL, dtype=torch.float32) + y = torch.full((16,), INIT_VAL, dtype=torch.float32) + gate = torch.full((16,), -1, dtype=torch.int32) + return TaskArgsBuilder( + Tensor("x", x), + Tensor("y", y), + Tensor("gate", gate), + Scalar("case", int(params["case"])), + ) + + def compute_golden(self, args, params): + gate_value = 0 if params["case"] == 1 else 1 + args.gate[0] = gate_value + # x_producer writes 42.0; the clobber (X = 999.0) runs only if gate > 0. + x_final = SENTINEL if gate_value == 0 else POISON + args.x[0] = x_final + args.y[0] = x_final + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp new file mode 100644 index 000000000..a3c7b389c --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_clobber.cpp @@ -0,0 +1,53 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Clobber kernel for the predicated_dispatch scene tests: overwrites args[0][0] + * with a poison value (999.0f). It is the body of the predicated task — if the + * predicate correctly suppresses dispatch, this kernel never runs and the + * producer's sentinel (42.0f) survives; if it ever runs, the downstream copy + * reads 999.0f and the test fails. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = 999.0f; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp new file mode 100644 index 000000000..dd6804f00 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_copy_first.cpp @@ -0,0 +1,57 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Consumer kernel for the predicated_dispatch scene tests: copies args[0][0] -> args[1][0]. + * + * The producer writes the sentinel (42.0f) to args[0]; a predicated clobber sits + * between the producer and this consumer. When the predicate suppresses the + * clobber's dispatch, args[0] is untouched and args[1][0] equals 42.0f; when the + * clobber dispatches, args[1][0] is the poison value instead. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *in_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + + __gm__ float *in = reinterpret_cast<__gm__ float *>(in_tensor->buffer.addr) + in_tensor->start_offset; + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = in[0]; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp new file mode 100644 index 000000000..5b579cfaa --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_const.cpp @@ -0,0 +1,51 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Writes a fixed sentinel (42.0f) to args[0][0]. The producer in the + * predicated_dispatch scene tests; a downstream consumer verifies the dependency + * chain (producer -> predicated clobber -> consumer) propagates the value intact. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + out[0] = 42.0f; + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp new file mode 100644 index 000000000..2c00413f6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/aic/kernel_write_gate.cpp @@ -0,0 +1,55 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Producer of the predicate value for the predicated_dispatch scene tests. + * Writes the scalar arg into gate[0] (INT32). A downstream task carries a + * dispatch predicate on gate[0]; the scheduler reads this value at the dispatch + * point (after this producer has completed) to decide whether to dispatch. + * + * args[0] = gate tensor (INOUT, INT32); args[1] = scalar gate value. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + int64_t gate_value = args[1]; // scalar arg follows the tensor args + + __gm__ int32_t *out = reinterpret_cast<__gm__ int32_t *>(out_tensor->buffer.addr) + out_tensor->start_offset; + out[0] = static_cast(gate_value); + dcci(&out[0], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp new file mode 100644 index 000000000..9cba6a725 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/kernels/orchestration/predicated_dispatch_orch.cpp @@ -0,0 +1,112 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * predicated_dispatch orchestration: delayed-evaluation dispatch predicate. + * + * The predicate value is produced by a prior task and read by the scheduler at + * the dispatch point — never in orchestration — so the orchestrator does not + * stall on it: + * + * gate_producer (WRITE_GATE) writes gate[0] = gate_value (INT32) + * x_producer (WRITE_CONST) writes X[0] = 42.0 + * clobber (CLOBBER) would write X[0] = 999.0, but carries + * set_predicate(gate[0] > 0) and depends on + * gate_producer. The scheduler reads gate[0] when + * the clobber becomes ready (gate already written) + * and dispatches only if gate[0] > 0; otherwise the + * task is retired inline without dispatch. + * consumer (COPY_FIRST) copies X[0] -> Y[0] + * + * case=1: gate_value = 0 -> predicate FALSE -> clobber not dispatched -> + * X stays 42.0 -> Y = 42.0. Proves non-dispatch + consumer still unlocks. + * case=2: gate_value = 1 -> predicate TRUE -> clobber dispatched -> + * X = 999.0 -> Y = 999.0. Proves the dispatch path is taken. + * + * Args layout: [X, Y, gate] + case scalar. + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_WRITE_CONST 0 +#define FUNC_COPY_FIRST 1 +#define FUNC_CLOBBER 2 +#define FUNC_WRITE_GATE 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 4, // 3 tensors + 1 case scalar + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_X = orch_args.tensor(0).ref(); + const Tensor &ext_Y = orch_args.tensor(1).ref(); + const Tensor &ext_gate = orch_args.tensor(2).ref(); + + uint64_t case_id = orch_args.scalar(0); + if (case_id != 1 && case_id != 2) { + rt_report_fatal(PTO2_ERROR_INVALID_ARGS, "unsupported case_id=%llu", static_cast(case_id)); + return; + } + // case 1 => gate 0 (predicate FALSE, skip); case 2 => gate 1 (predicate TRUE, dispatch). + int64_t gate_value = (case_id == 1) ? 0 : 1; + + // gate producer: gate[0] = gate_value + PTO2TaskId gate_tid; + { + L0TaskArgs args; + args.add_inout(ext_gate); + args.add_scalar(gate_value); + gate_tid = rt_submit_aic_task(FUNC_WRITE_GATE, args).task_id(); + } + + // x producer: X[0] = 42.0 + { + L0TaskArgs args; + args.add_inout(ext_X); + rt_submit_aic_task(FUNC_WRITE_CONST, args); + } + + // predicated clobber: would write X[0] = 999.0 if dispatched. Depends on the + // gate producer so gate[0] is written by the time this task is ready; the + // scheduler reads gate[0] at the dispatch point and dispatches only if > 0. + { + L0TaskArgs args; + args.add_inout(ext_X); + PTO2TaskId deps[] = {gate_tid}; + args.set_dependencies(deps, 1); + // predicate: gate[0] > 0 (operand op target), built level by level. + L0TaskPredicate pred; + pred.operand.tensor = &ext_gate; + pred.operand.ndims = 1; + pred.operand.indices[0] = 0; + pred.op = PredicateOp::GT; + pred.target = 0; + args.set_predicate(pred); + rt_submit_aic_task(FUNC_CLOBBER, args); + } + + // consumer: Y[0] = X[0] + { + L0TaskArgs args; + args.add_input(ext_X); + args.add_inout(ext_Y); + rt_submit_aic_task(FUNC_COPY_FIRST, args); + } +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py new file mode 100644 index 000000000..a4025afa6 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/predicated_dispatch/test_predicated_dispatch.py @@ -0,0 +1,124 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""predicated_dispatch: a dispatch predicate is evaluated by the scheduler at the +dispatch point (delayed evaluation), not in orchestration. + +The predicate value is produced by a prior task; the scheduler reads it only +when the predicated task becomes ready (its producer done), so orchestration +never stalls on it. Chain: + + gate_producer writes gate[0] = gate_value (INT32) + x_producer writes X[0] = 42.0 + clobber would write X[0] = 999.0, but carries set_predicate(gate[0] > 0) + and depends on gate_producer; the scheduler reads gate[0] at the + dispatch point and dispatches only if gate[0] > 0. + consumer copies X[0] -> Y[0] + + case=1 (gate_value = 0): predicate FALSE -> clobber NOT dispatched -> X stays + 42.0 -> Y = 42.0. Proves non-dispatch AND that the retired task still + unlocks the consumer. + case=2 (gate_value = 1): predicate TRUE -> clobber dispatched -> X = 999.0 -> + Y = 999.0. Proves the dispatch path is taken when the predicate holds. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +SENTINEL = 42.0 +POISON = 999.0 # what the clobber writes if the predicate lets it dispatch +INIT_VAL = -1.0 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestPredicatedDispatch(SceneTestCase): + """predicated_dispatch: scheduler evaluates the dispatch predicate at dispatch time.""" + + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/predicated_dispatch_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT, D.INOUT, D.INOUT], # X, Y, gate + }, + "incores": [ + { + "func_id": 0, + "name": "WRITE_CONST", + "source": "kernels/aic/kernel_write_const.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "COPY_FIRST", + "source": "kernels/aic/kernel_copy_first.cpp", + "core_type": "aic", + "signature": [D.IN, D.INOUT], + }, + { + "func_id": 2, + "name": "CLOBBER", + "source": "kernels/aic/kernel_clobber.cpp", + "core_type": "aic", + # Body of the predicated task; runs only if the predicate holds. + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "WRITE_GATE", + "source": "kernels/aic/kernel_write_gate.cpp", + "core_type": "aic", + # One INOUT tensor (gate); the gate value rides as a trailing scalar. + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "PredicateFalseSkips", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 2, "block_dim": 1}, + "params": {"case": 1}, + }, + { + "name": "PredicateTrueDispatches", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 2, "block_dim": 1}, + "params": {"case": 2}, + }, + ] + + def generate_args(self, params): + x = torch.full((16,), INIT_VAL, dtype=torch.float32) + y = torch.full((16,), INIT_VAL, dtype=torch.float32) + gate = torch.full((16,), -1, dtype=torch.int32) + return TaskArgsBuilder( + Tensor("x", x), + Tensor("y", y), + Tensor("gate", gate), + Scalar("case", int(params["case"])), + ) + + def compute_golden(self, args, params): + gate_value = 0 if params["case"] == 1 else 1 + args.gate[0] = gate_value + # x_producer writes 42.0; the clobber (X = 999.0) runs only if gate > 0. + x_final = SENTINEL if gate_value == 0 else POISON + args.x[0] = x_final + args.y[0] = x_final + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__)