Skip to content

Add: dual-slot AICPU dispatch payload and two-phase pipelining scheduler #477

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhusy54 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add: dual-slot AICPU dispatch payload and two-phase pipelining scheduler #477

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,16 +112,19 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in
{
uint32_t task_id = reg_val; // Decode: register holds task_id directly

// Select dual-buffer slot: same bit as AICPU used when writing payload
__gm__ PTO2DispatchPayload *exec_payload = payload + (task_id & 1u);

// Invalidate payload buffer (AICPU updates its content each dispatch)
dcci(payload, ENTIRE_DATA_CACHE);
dcci(exec_payload, ENTIRE_DATA_CACHE);

write_reg(RegId::COND, MAKE_ACK_VALUE(task_id));

// Performance profiling: record start time
uint64_t start_time = get_sys_cnt_aicore();

// Execute the task
execute_task(payload);
execute_task(exec_payload);

// Performance profiling: record task execution
if (profiling_enabled) {
Expand Down
553 changes: 413 additions & 140 deletions src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp
View file
Open in desktop

Large diffs are not rendered by default.

7 changes: 5 additions & 2 deletions src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,16 +112,19 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in
{
uint32_t task_id = reg_val; // Decode: register holds task_id directly

// Select dual-buffer slot: same bit as AICPU used when writing payload
__gm__ PTO2DispatchPayload *exec_payload = payload + (task_id & 1u);

// Invalidate payload buffer (AICPU updates its content each dispatch)
dcci(payload, ENTIRE_DATA_CACHE);
dcci(exec_payload, ENTIRE_DATA_CACHE);

write_reg(RegId::COND, MAKE_ACK_VALUE(task_id));

// Performance profiling: record start time
uint64_t start_time = get_sys_cnt_aicore();

// Execute the task
execute_task(payload);
execute_task(exec_payload);

// Performance profiling: record task execution
if (profiling_enabled) {
Expand Down
540 changes: 404 additions & 136 deletions src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp
View file
Open in desktop

Large diffs are not rendered by default.

3 changes: 3 additions & 0 deletions tests/st/a2a3/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_pv_matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,9 @@ static __aicore__ void pv_matmul_batch_impl(
pipe_barrier(PIPE_ALL);
}
}

set_flag(PIPE_FIX, PIPE_S, EVENT_ID7);
wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7);
}

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
Expand Down
3 changes: 3 additions & 0 deletions tests/st/a2a3/tensormap_and_ringbuffer/batch_paged_attention/kernels/aic/aic_qk_matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,9 @@ static __aicore__ void qk_matmul_batch_impl(
pipe_barrier(PIPE_ALL);
}
}

set_flag(PIPE_FIX, PIPE_S, EVENT_ID7);
wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7);
}

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
Expand Down
3 changes: 3 additions & 0 deletions .../st/a2a3/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_online_update.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,9 @@ static __aicore__ void online_update_batch_impl(
pipe_barrier(PIPE_ALL);
}
}

set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
}

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
Expand Down
3 changes: 3 additions & 0 deletions ...t/a2a3/tensormap_and_ringbuffer/batch_paged_attention/kernels/aiv/aiv_softmax_prepare.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,9 @@ static __aicore__ void softmax_prepare_batch_impl(
pipe_barrier(PIPE_ALL);
}
}

set_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID7);
}

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
Expand Down
Loading