diff --git a/examples/a2a3/host_build_graph/paged_attention/README.md b/examples/a2a3/host_build_graph/paged_attention/README.md new file mode 100644 index 0000000000..ef624e642f --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/README.md @@ -0,0 +1,65 @@ +# paged_attention — host_build_graph + +Online-softmax paged attention in bfloat16, split across AIC and AIV, on the +`host_build_graph` runtime. The kernels and the orchestration source are the +same ones the `tensormap_and_ringbuffer` sibling compiles. + +## Why the same orchestration compiles on both runtimes + +`orchestration/pto_orchestration_api.h` is identical under +`src/a2a3/runtime/host_build_graph/` and +`src/a2a3/runtime/tensormap_and_ringbuffer/` — same entry points, same macros. +`KernelCompiler.compile_orchestration(runtime, source)` picks the include dirs +of whichever runtime `@scene_test(runtime=...)` names, so switching runtime is a +one-line change in the test class and nothing in the C++. + +## The four-task loop + +Per (batch, head) and per KV block the orchestration submits four tasks: + +| Task | Core | Computation | +| ---- | ---- | ----------- | +| `QK` | AIC | `qi @ K^T` for the block | +| `SF` | AIV | softmax prepare — running `mi`, `li` | +| `PV` | AIC | `P @ V` | +| `UP` | AIV | online-softmax accumulation into the running output | + +Ordering is not written down: the tasks sit in a plain `PTO2_SCOPE()` with a +`PTO2_SCOPE_GUARD()` and the dependency graph is derived from the tensors they +share. See `../paged_attention_manual_scope/` for the variant that wires the +same edges by hand. + +## Ring sizing is the one place host_build_graph differs + +`host_build_graph` builds the entire task graph on the host before the device +starts scheduling, so no ring slot can be reclaimed mid-orchestration — the ring +window and GM heap must hold every task of a case at once. Task count is + +```text +tasks = batch * (4 * ceil(context_len / block_size) + 1) +``` + +so `Case1` needs 65 792 slots and `Case2` needs 32 832, both above the default +window of 16384. Each carries its own `runtime_env` in `CASES[*]["config"]` +(`ring_task_window` must be a power of two in `[4, INT32_MAX]`; `ring_heap` at +least 1024), so no `PTO2_RING_*` environment variable is needed. + +## Cases + +`Case1` (65 792 tasks) and `Case2` (32 832) at production scale, `CaseSmall1` / +`CaseSmall2`, and `CaseVarSeq2` / `CaseVarSeq4` for ragged sequence lengths. +All but `CaseSmall1` are `manual`, so the default run executes `CaseSmall1` +only; add `--manual include` to reach the rest. + +The upstream `tensormap_and_ringbuffer` example also carries a `Case3` +(`head_dim: 256`). It is absent here because it does not produce correct +results on either runtime — see `KNOWN_ISSUES.md`. + +## Run + +```bash +python examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py \ + -p a2a3 -d 0 # CaseSmall1, golden checked +python examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py \ + -p a2a3 -d 0 --manual include --case Case2 --rounds 2 +``` diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..0220a6bbb3 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,114 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// PV Matmul Kernel: pij(M, K) @ vj(K, N) -> oi_new(M, N) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// pij is bfloat16 (converted from fp32 in softmax_prepare via TCVT). +// vj is stored as (K, N) = (block_size, head_dim) in row-major (ND) layout. +// Standard non-transposed B pattern: ND GlobalB + ColMajor/RowMajor TileMatB. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void pv_matmul_impl(__gm__ Tensor *pij, __gm__ Tensor *vj, __gm__ Tensor *oi) { + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ bfloat16_t *vj_addr = reinterpret_cast<__gm__ bfloat16_t *>(vj->buffer.addr); + __gm__ float *oi_addr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + + // pij (M, K) bf16, vj (K, N) bf16 in ND (row-major), oi_new (M, N) fp32 + using GlobalA = GlobalTensor, Stride>; + using GlobalB = GlobalTensor, Stride>; + using GlobalOut = GlobalTensor, Stride>; + + GlobalA pijGlobal(pij_addr + pij->start_offset); + GlobalB vjGlobal(vj_addr + vj->start_offset); + GlobalOut oiGlobal(oi_addr + oi->start_offset); + + // L1 Mat tiles: standard ND pattern for both A and B + using TileMatA = Tile; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Load pij and vj to L1 with separate events for pipeline overlap + TLOAD(aMatTile, pijGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // A load done + TLOAD(bMatTile, vjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // B load done + + // Move A to L0A as soon as A load completes (B may still be loading) + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + // Move B to L0B after B load completes + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + // Single matmul: (M,K) x (K,N) -> (M,N) + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(oiGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *vj = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + uint64_t q_tile_size = static_cast(pij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + pv_matmul_impl<16, 16, 16>(pij, vj, oi_new); + } else if (q_tile_size == 16) { + pv_matmul_impl<16, 128, 128>(pij, vj, oi_new); + } else { + pv_matmul_impl<64, 64, 128>(pij, vj, oi_new); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..efd423bd6e --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,115 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// QK Matmul Kernel: qi(M, K) @ kj.T(K, N) -> sij(M, N) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// kj is stored as (N, K) = (block_size, head_dim) in row-major memory. +// This is equivalent to (K, N) in column-major (DN) layout. +// Using DN GlobalB + RowMajor/ColMajor TileMatB to handle the transposed B pattern. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void qk_matmul_impl(__gm__ Tensor *qi, __gm__ Tensor *kj, __gm__ Tensor *sij) { + __gm__ bfloat16_t *qi_addr = reinterpret_cast<__gm__ bfloat16_t *>(qi->buffer.addr); + __gm__ bfloat16_t *kj_addr = reinterpret_cast<__gm__ bfloat16_t *>(kj->buffer.addr); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + + // qi (M, K) bf16 in ND (row-major) layout + using GlobalA = GlobalTensor, Stride>; + // kj stored as (N, K) row-major = (K, N) column-major -> DN layout + using GlobalB = GlobalTensor, Stride, Layout::DN>; + using GlobalOut = GlobalTensor, Stride>; + + GlobalA qiGlobal(qi_addr + qi->start_offset); + GlobalB kjGlobal(kj_addr + kj->start_offset); + GlobalOut sijGlobal(sij_addr + sij->start_offset); + + // L1 Mat tiles: A is standard ND, B uses transposed-B pattern (RowMajor/ColMajor) + using TileMatA = Tile; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Load A and B to L1 with separate events for pipeline overlap + TLOAD(aMatTile, qiGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // A load done + TLOAD(bMatTile, kjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // B load done + + // Move A to L0A as soon as A load completes (B may still be loading) + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + // Move B to L0B after B load completes + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + // Matmul + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *qi = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *kj = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[2]); + uint64_t q_tile_size = static_cast(qi->shapes[0]); + + if (q_tile_size == 16 && qi->shapes[1] <= 16) { + qk_matmul_impl<16, 16, 16>(qi, kj, sij); + } else if (q_tile_size == 16) { + qk_matmul_impl<16, 128, 128>(qi, kj, sij); + } else { + qk_matmul_impl<64, 128, 64>(qi, kj, sij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..ded4dcad87 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,256 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Data (M, N) RowMajor + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalDataMxN oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalDataMxN oiGlobal(oi_ptr + oi->start_offset); + GlobalDataMxN dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to the same UB as DN tiles for storing as ND format + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + pipe_barrier(PIPE_V); + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + pipe_barrier(PIPE_V); + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + pipe_barrier(PIPE_V); + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + pipe_barrier(PIPE_V); + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + pipe_barrier(PIPE_V); + TMUL(tmpRow, alphaRow, liRow); // alpha * li + pipe_barrier(PIPE_V); + TMUL(liNewRow, betaRow, lijRow); // beta * lij + pipe_barrier(PIPE_V); + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + pipe_barrier(PIPE_V); + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + pipe_barrier(PIPE_V); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t q_tile_size = static_cast(mij->shapes[0]); + + if (q_tile_size == 16 && oi_new->shapes[1] <= 16) { + online_update_impl<16, 16>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8f0c41775d --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,156 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Softmax Preparation Kernel (AIV) with partial block masking +// +// Operates on (M, N) tile where M=q_tile_size, N=block_size: +// Case1: sij is (16, 128) +// Case2: sij is (64, 64) +// +// For partial blocks (valid_len < N), positions [valid_len, N) in sij are +// filled with -inf via TFILLPAD_INPLACE before softmax, ensuring exp(-inf)=0 +// so that invalid key positions contribute zero attention weight. +// +// Computes: +// sij_masked = TFILLPAD(sij, valid_len, pad=-inf) +// sij_scale = sij_masked * scale +// mij = row_max(sij_scale) -> (M, 1) +// pij = exp(sij_scale - mij) -> (M, N) +// lij = row_sum(pij) -> (M, 1) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_impl( + __gm__ Tensor *sij, float scale_value, __gm__ Tensor *pij, __gm__ Tensor *mij, __gm__ Tensor *lij +) { + uint64_t valid_len = static_cast(sij->shapes[1]); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + GlobalDataMxN sijGlobal(sij_addr + sij->start_offset); + GlobalDataMxN_bf16 pijGlobal(pij_addr + pij->start_offset); + GlobalScalarDN mijGlobal(mij_addr + mij->start_offset); + GlobalScalarDN lijGlobal(lij_addr + lij->start_offset); + + // Dynamic-cols tile: marks which columns are valid for TFILLPAD boundary + using TileSijDyn = Tile; + // Padded tile: TFILLPAD_INPLACE fills positions [valid_len, N) with -inf + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijDyn sijDynTile(static_cast(valid_len)); + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + // All sij tiles share UB address 0x0 (in-place masking) + TASSIGN(sijTile, 0x0); + TASSIGN(sijDynTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + // Load full sij (M, N) tile from GM - all N columns including garbage for partial blocks + // printf("sij addr incore %x\n", sij->buffer.addr); + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Mask columns [valid_len, N) with -inf. sijDynTile provides the valid boundary, + // sijPadTile provides PadValue::Min as the fill value. No-op when valid_len == N. + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWMAX(maxTile, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first + pipe_barrier(PIPE_V); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // pij bf16 ready, can store early + + // Continue computing: bf16 → f32 and rowsum while pij store proceeds in parallel + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TROWSUM(sumTile, pijTile, tmpTile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); // sum ready + + // Store pij (overlaps with TCVT + TROWSUM above) + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Store max and sum + TSTORE(mijGlobal, maxTile); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(lijGlobal, sumTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t q_tile_size = static_cast(sij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + softmax_prepare_impl<16, 16>(sij, scale_value, pij, mij, lij); + } else if (q_tile_size == 16) { + softmax_prepare_impl<16, 128>(sij, scale_value, pij, mij, lij); + } else { + softmax_prepare_impl<64, 64>(sij, scale_value, pij, mij, lij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp b/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..e87df0bf14 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,292 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Paged Attention Orchestration Function - 16x16 Version + * + * Simplified for 16x16 framework-generated matmul kernels. + * Each block processes a single 16x16 matmul operation. + * + * Memory Layout: + * Query: (batch, 16, 16) - one 16x16 tile per batch + * Key: (total_blocks, 16, 16) - stored as K^T for direct matmul + * Value: (total_blocks, 16, 16) - direct format + */ + +#include +#include +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +constexpr uint64_t PLATFORM_PROF_SYS_CNT_FREQ = 50000000; // 50 MHz + +inline double cycles_to_us(uint64_t cycles) { + return (static_cast(cycles) / PLATFORM_PROF_SYS_CNT_FREQ) * 1000000.0; +} + +inline uint64_t get_sys_cnt_aicpu() { +#if defined(__aarch64__) + uint64_t ticks; + asm volatile("mrs %0, cntvct_el0" : "=r"(ticks)); + return ticks; +#elif defined(__x86_64__) + return 0; +#else + return 0; +#endif +} + +#ifdef ENABLE_PROFILING +#define CYCLE_COUNT_START() uint64_t _t0 = get_sys_cnt_aicpu(), _t1 +#define CYCLE_COUNT_LAP(acc) \ + do { \ + _t1 = get_sys_cnt_aicpu(); \ + acc += (_t1 - _t0); \ + _t0 = _t1; \ + } while (0) +#define PROF_INC(counter, n) (counter) += (n) +#else +#define CYCLE_COUNT_START() (void)0 +#define CYCLE_COUNT_LAP(acc) (void)0 +#define PROF_INC(counter, n) (void)0 +#endif + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { +#ifdef ENABLE_PROFILING + uint64_t prof_param_extract = 0; + uint64_t prof_ext_tensor = 0; + uint64_t prof_scope = 0; + uint64_t prof_make_tensor = 0; + uint64_t prof_tensor_view = 0; + uint64_t prof_param_setup = 0; + uint64_t prof_submit_task = 0; + int prof_submit_count = 0; + int prof_make_count = 0; + int prof_view_count = 0; +#endif + + CYCLE_COUNT_START(); + + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_head_num = num_heads; + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (q_head_num + q_tile - 1) / q_tile; + CYCLE_COUNT_LAP(prof_param_extract); + + LOG_INFO(">>>>>> batch = %" PRIu64, batch); + + // Reshape tensors for kernel consumption (2D flattened) + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint32_t key_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t value_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32); + CYCLE_COUNT_LAP(prof_ext_tensor); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + // Create infos are loop-invariant — shapes depend only on q_tile/head_dim/block_size + uint32_t tile2d_shapes[2] = {static_cast(q_tile), static_cast(head_dim)}; + uint32_t scalar_shapes[1] = {static_cast(q_tile)}; + uint32_t sij_shapes[2] = {static_cast(q_tile), static_cast(block_size)}; + TensorCreateInfo tile2d_ci(tile2d_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32); + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_f16_ci(sij_shapes, 2, data_type); + + PROF_INC(prof_make_count, 4); + CYCLE_COUNT_LAP(prof_make_tensor); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint32_t cl_idx[1] = {static_cast(b_idx)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE() { + CYCLE_COUNT_LAP(prof_scope); + uint64_t cur_offset = b_idx * q_head_num + q_idx * q_tile; + + uint32_t qi_offsets[2] = {static_cast(cur_offset), 0}; + Tensor qi = query.view(tile2d_shapes, qi_offsets); + uint32_t out_view_offsets[2] = {static_cast(cur_offset), 0}; + Tensor out_view = out.view(tile2d_shapes, out_view_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors alloc_outs = alloc_tensors(tile2d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + for (uint64_t bn = 0; bn < bn_this_batch; bn++) { + PTO2_SCOPE_GUARD(); + + uint32_t bt_idx[2] = {static_cast(b_idx), static_cast(bn)}; + uint64_t cur_block_idx = static_cast(get_tensor_data(block_table, 2, bt_idx)); + uint64_t valid_len = std::min(block_size, cur_seq - bn * block_size); + CYCLE_COUNT_LAP(prof_param_extract); + + uint32_t kv_shapes[2] = {static_cast(block_size), static_cast(head_dim)}; + uint32_t kv_offsets[2] = {static_cast(cur_block_idx * block_size), 0}; + Tensor kj = key_cache.view(kv_shapes, kv_offsets); + Tensor vj = value_cache.view(kv_shapes, kv_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_qk; + params_qk.add_input(qi); + params_qk.add_input(kj); + params_qk.add_output(sij_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij = qk_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint32_t sij_valid_shapes[2] = {static_cast(q_tile), static_cast(valid_len)}; + uint32_t sij_valid_offsets[2] = {0, 0}; + Tensor sij_valid = sij.view(sij_valid_shapes, sij_valid_offsets); + PROF_INC(prof_view_count, 1); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_sf; + params_sf.add_input(sij_valid); + params_sf.add_output(pij_f16_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + params_sf.add_scalar(scale_value); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_f16 = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + L0TaskArgs params_pv; + params_pv.add_input(pij_f16); + params_pv.add_input(vj); + params_pv.add_output(tile2d_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_tmp = pv_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == bn_this_batch - 1) ? 1 : 0; + CYCLE_COUNT_LAP(prof_param_extract); + + L0TaskArgs params_up; + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_tmp); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + CYCLE_COUNT_LAP(prof_param_setup); + rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + } + } + CYCLE_COUNT_LAP(prof_scope); + } + } + +#ifdef ENABLE_PROFILING + uint64_t total = prof_param_extract + prof_ext_tensor + prof_make_tensor + prof_tensor_view + prof_param_setup + + prof_submit_task + prof_scope; + LOG_INFO( + "=== PagedAttn Orch Profiling: %d submits, %d makes, %d views, total=%.3fus ===", prof_submit_count, + prof_make_count, prof_view_count, cycles_to_us(total) + ); + if (total > 0) { + LOG_INFO( + " param_extract : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_extract), + prof_param_extract * 100.0 / total + ); + LOG_INFO( + " ext_tensor(x4) : %7.3fus (%5.1f%%)", cycles_to_us(prof_ext_tensor), prof_ext_tensor * 100.0 / total + ); + LOG_INFO( + " create_info(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_make_count, cycles_to_us(prof_make_tensor), + prof_make_tensor * 100.0 / total, + prof_make_count > 0 ? cycles_to_us(prof_make_tensor) / prof_make_count : 0.0 + ); + LOG_INFO( + " tensor_view(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_view_count, cycles_to_us(prof_tensor_view), + prof_tensor_view * 100.0 / total, + prof_view_count > 0 ? cycles_to_us(prof_tensor_view) / prof_view_count : 0.0 + ); + LOG_INFO( + " param_setup : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_setup), prof_param_setup * 100.0 / total + ); + LOG_INFO(" scope : %7.3fus (%5.1f%%)", cycles_to_us(prof_scope), prof_scope * 100.0 / total); + LOG_INFO( + " submit_task(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_submit_count, cycles_to_us(prof_submit_task), + prof_submit_task * 100.0 / total, + prof_submit_count > 0 ? cycles_to_us(prof_submit_task) / prof_submit_count : 0.0 + ); + } +#endif +} + +} // extern "C" diff --git a/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py b/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py new file mode 100644 index 0000000000..490f6f18b2 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention/test_paged_attention.py @@ -0,0 +1,202 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention: online softmax with AIC/AIV subgraph splitting (bfloat16), host_build_graph runtime. + +The orchestration and kernel sources are the same as the tensormap_and_ringbuffer +variant: `pto_orchestration_api.h` is identical between the two runtimes, and the +framework compiles the orchestration against the include dirs of whichever +runtime `@scene_test` names. + +host_build_graph populates the whole task graph on the host before the device +begins scheduling, so a ring slot cannot be reclaimed mid-orchestration: the +ring window and GM heap must hold every task of a case at once. Cases above the +default window carry explicit `runtime_env` sizing below. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + +# tasks = batch * (4 * ceil(context_len / block_size) + 1). ring_task_window must +# be a power of two in [4, INT32_MAX] and is rejected otherwise by the runtime. +_RING_65K = {"ring_task_window": 131072, "ring_heap": 2 * 1024 * 1024 * 1024} +_RING_33K = {"ring_task_window": 65536, "ring_heap": 1024 * 1024 * 1024} + + +@scene_test(level=2, runtime="host_build_graph") +class TestPagedAttentionHostBuildGraph(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + # 65 792 tasks. + "name": "Case1", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_65K}, + "manual": True, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + # 32 832 tasks. + "name": "Case2", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_33K}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + # context_lens_list makes the per-batch block counts differ, so the + # graph is ragged rather than a uniform batch * blocks grid. + "name": "CaseVarSeq2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md b/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md new file mode 100644 index 0000000000..781c03c47f --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/README.md @@ -0,0 +1,49 @@ +# paged_attention_manual_scope — host_build_graph + +Same computation as `../paged_attention/`, with the automatic same-scope +dependency wiring replaced by explicit task-to-task edges inside +`PTO2_SCOPE(PTO2ScopeMode::MANUAL)`. Read it against the baseline to see what +the automatic mode was deriving. + +## The two dependency APIs it demonstrates + +| API | Shape | Suited to | +| --- | ----- | --------- | +| `Arg::set_dependencies(buf, n)` | caller owns the buffer, `Arg` stores `(ptr, count)` | codegen, fixed dep sets | +| `L0TaskArgsWithDeps<>::add_dep(id)` | wrapper owns a stack-sized buffer, incremental | hand-written orch, deps assembled across branches | + +`SF` and `PV` use the primitive form; `UP` uses the convenience form because its +dep set is conditional — it always takes the `PV` edge, adds the previous +`UP` task when there is one, and adds the alloc task on the last block so the +scratch buffers outlive their final consumer. + +Both APIs exist unchanged on `host_build_graph`: the orchestration header is +byte-identical between the two runtimes, so this variant is the same C++ as its +`tensormap_and_ringbuffer` sibling. + +## Ring sizing + +`host_build_graph` submits the whole graph before the device schedules, so every +task of a case is live at once and the ring must hold all of them — `Case1` +65 792 and `Case2` 32 832, both above the default 16384 window. Each case +carries its own `runtime_env` sizing in `CASES[*]["config"]`. + +This is a stronger constraint than the `tensormap_and_ringbuffer` variant faces, +where slots retire as orchestration proceeds and only the *in-flight* depth of a +single MANUAL scope has to fit. + +## Cases + +`Case1`, `Case2`, `CaseSmall1`, `CaseSmall2`, `CaseVarSeq2`, `CaseVarSeq4`. All +but `CaseSmall1` are `manual`. The upstream `Case3` (`head_dim: 256`) is absent +because it does not produce correct results on either runtime — see +`KNOWN_ISSUES.md`. + +## Run + +```bash +python examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py \ + -p a2a3 -d 0 # CaseSmall1, golden checked +python examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py \ + -p a2a3 -d 0 --manual include --case Case1 --rounds 2 +``` diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp new file mode 100644 index 0000000000..0220a6bbb3 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_pv_matmul.cpp @@ -0,0 +1,114 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// PV Matmul Kernel: pij(M, K) @ vj(K, N) -> oi_new(M, N) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128) -> (16, 128) +// Case2: (64, 64) @ ( 64, 128) -> (64, 128) +// +// pij is bfloat16 (converted from fp32 in softmax_prepare via TCVT). +// vj is stored as (K, N) = (block_size, head_dim) in row-major (ND) layout. +// Standard non-transposed B pattern: ND GlobalB + ColMajor/RowMajor TileMatB. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void pv_matmul_impl(__gm__ Tensor *pij, __gm__ Tensor *vj, __gm__ Tensor *oi) { + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ bfloat16_t *vj_addr = reinterpret_cast<__gm__ bfloat16_t *>(vj->buffer.addr); + __gm__ float *oi_addr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + + // pij (M, K) bf16, vj (K, N) bf16 in ND (row-major), oi_new (M, N) fp32 + using GlobalA = GlobalTensor, Stride>; + using GlobalB = GlobalTensor, Stride>; + using GlobalOut = GlobalTensor, Stride>; + + GlobalA pijGlobal(pij_addr + pij->start_offset); + GlobalB vjGlobal(vj_addr + vj->start_offset); + GlobalOut oiGlobal(oi_addr + oi->start_offset); + + // L1 Mat tiles: standard ND pattern for both A and B + using TileMatA = Tile; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Load pij and vj to L1 with separate events for pipeline overlap + TLOAD(aMatTile, pijGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // A load done + TLOAD(bMatTile, vjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // B load done + + // Move A to L0A as soon as A load completes (B may still be loading) + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + // Move B to L0B after B load completes + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + // Single matmul: (M,K) x (K,N) -> (M,N) + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(oiGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *vj = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + uint64_t q_tile_size = static_cast(pij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + pv_matmul_impl<16, 16, 16>(pij, vj, oi_new); + } else if (q_tile_size == 16) { + pv_matmul_impl<16, 128, 128>(pij, vj, oi_new); + } else { + pv_matmul_impl<64, 64, 128>(pij, vj, oi_new); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp new file mode 100644 index 0000000000..efd423bd6e --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aic/aic_qk_matmul.cpp @@ -0,0 +1,115 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// QK Matmul Kernel: qi(M, K) @ kj.T(K, N) -> sij(M, N) +// +// Supports two tile configurations via runtime dispatch: +// Case1: (16, 128) @ (128, 128).T -> (16, 128) +// Case2: (64, 128) @ (128, 64).T -> (64, 64) +// +// kj is stored as (N, K) = (block_size, head_dim) in row-major memory. +// This is equivalent to (K, N) in column-major (DN) layout. +// Using DN GlobalB + RowMajor/ColMajor TileMatB to handle the transposed B pattern. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void qk_matmul_impl(__gm__ Tensor *qi, __gm__ Tensor *kj, __gm__ Tensor *sij) { + __gm__ bfloat16_t *qi_addr = reinterpret_cast<__gm__ bfloat16_t *>(qi->buffer.addr); + __gm__ bfloat16_t *kj_addr = reinterpret_cast<__gm__ bfloat16_t *>(kj->buffer.addr); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + + // qi (M, K) bf16 in ND (row-major) layout + using GlobalA = GlobalTensor, Stride>; + // kj stored as (N, K) row-major = (K, N) column-major -> DN layout + using GlobalB = GlobalTensor, Stride, Layout::DN>; + using GlobalOut = GlobalTensor, Stride>; + + GlobalA qiGlobal(qi_addr + qi->start_offset); + GlobalB kjGlobal(kj_addr + kj->start_offset); + GlobalOut sijGlobal(sij_addr + sij->start_offset); + + // L1 Mat tiles: A is standard ND, B uses transposed-B pattern (RowMajor/ColMajor) + using TileMatA = Tile; + using TileMatB = Tile; + + // L0 tiles + using LeftTile = TileLeft; + using RightTile = TileRight; + using AccTile = TileAcc; + + TileMatA aMatTile; + TileMatB bMatTile; + TASSIGN(aMatTile, 0x0); + TASSIGN(bMatTile, 0x20000); + + LeftTile aTile; + RightTile bTile; + AccTile cTile; + TASSIGN(aTile, 0x0); + TASSIGN(bTile, 0x0); + TASSIGN(cTile, 0x0); + + // Load A and B to L1 with separate events for pipeline overlap + TLOAD(aMatTile, qiGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); // A load done + TLOAD(bMatTile, kjGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); // B load done + + // Move A to L0A as soon as A load completes (B may still be loading) + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aTile, aMatTile); + // Move B to L0B after B load completes + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TMOV(bTile, bMatTile); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + // Matmul + TMATMUL(cTile, aTile, bTile); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + TSTORE(sijGlobal, cTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *qi = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *kj = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[2]); + uint64_t q_tile_size = static_cast(qi->shapes[0]); + + if (q_tile_size == 16 && qi->shapes[1] <= 16) { + qk_matmul_impl<16, 16, 16>(qi, kj, sij); + } else if (q_tile_size == 16) { + qk_matmul_impl<16, 128, 128>(qi, kj, sij); + } else { + qk_matmul_impl<64, 128, 64>(qi, kj, sij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp new file mode 100644 index 0000000000..ded4dcad87 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_online_update.cpp @@ -0,0 +1,256 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Online Softmax Update + Normalize Kernel (AIV) +// +// Operates on full tiles where M=q_tile_size, N=head_dim (128): +// Case1: oi/oi_new are (16, 128), mij/lij/mi/li are 16-element vectors +// Case2: oi/oi_new are (64, 128), mij/lij/mi/li are 64-element vectors +// +// Scalar layout strategy using TRESHAPE (zero-copy UB reshape): +// Scalars loaded as DN ColMajor (M, 1) for TROWEXPANDMUL/TROWEXPANDDIV. +// For element-wise ops (TMAX, TSUB, TEXP, etc.), TRESHAPE to RowMajor (1, M). +// After arithmetic, TRESHAPE back to ColMajor (M, 1) for row-broadcast ops. +// This eliminates the GM round-trip (TSTORE ND → TLOAD DN) used in the original. + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void online_update_impl( + __gm__ Tensor *mij, __gm__ Tensor *lij, __gm__ Tensor *oi_new, __gm__ Tensor *mi, __gm__ Tensor *li, + __gm__ Tensor *oi, uint64_t is_first, uint64_t is_last, __gm__ Tensor *dst +) { + __gm__ float *mij_ptr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_ptr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + __gm__ float *oi_new_ptr = reinterpret_cast<__gm__ float *>(oi_new->buffer.addr); + __gm__ float *mi_ptr = reinterpret_cast<__gm__ float *>(mi->buffer.addr); + __gm__ float *li_ptr = reinterpret_cast<__gm__ float *>(li->buffer.addr); + __gm__ float *oi_ptr = reinterpret_cast<__gm__ float *>(oi->buffer.addr); + __gm__ float *dst_ptr = reinterpret_cast<__gm__ float *>(dst->buffer.addr); + + // Aligned rows for ColMajor DN tiles (32-byte alignment) + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + // --- GlobalTensor types --- + + // Data (M, N) RowMajor + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + + // Scalar DN: M contiguous floats as (kAlignedRows, 1) ColMajor for TROWEXPAND ops and loading + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + // Scalar ND: for storing mi_new and li_new back to GM + constexpr int kScalarCols = 32 / sizeof(float); + constexpr int kScalarRows = M / kScalarCols; + using GlobalScalarND = + GlobalTensor, Stride<1, 1, 1, kScalarCols, 1>>; + + // --- GlobalTensor instances --- + + GlobalDataMxN oiNewGlobal(oi_new_ptr + oi_new->start_offset); + GlobalDataMxN oiGlobal(oi_ptr + oi->start_offset); + GlobalDataMxN dstGlobal(dst_ptr + dst->start_offset); + + // DN globals for loading scalars as ColMajor + GlobalScalarDN mijGlobalDN(mij_ptr + mij->start_offset); + GlobalScalarDN lijGlobalDN(lij_ptr + lij->start_offset); + GlobalScalarDN miGlobalDN(mi_ptr + mi->start_offset); + GlobalScalarDN liGlobalDN(li_ptr + li->start_offset); + + // ND globals for storing scalar results + GlobalScalarND miGlobalND(mi_ptr + mi->start_offset); + GlobalScalarND liGlobalND(li_ptr + li->start_offset); + + // --- Tile types --- + + using TileDataMxN = Tile; + using TileScalarDN = Tile; + + // RowMajor (1, M) tiles for element-wise arithmetic via TRESHAPE + using TileScalarRow = Tile; + + // ND tile for storing back to GM + using TileScalarND = + Tile; + + // --- UB memory layout --- + + constexpr int kDataBytes = M * N * sizeof(float); + constexpr int kScalarDNBytes = kAlignedRows * sizeof(float); + + // Data tiles + TileDataMxN oiNewTile; + TileDataMxN oiTile; + + // Scalar DN tiles loaded from GM (ColMajor) + TileScalarDN mijDN, lijDN, miDN, liDN; + + // Temporary DN tiles for results + TileScalarDN miNewDN, alphaDN, betaDN, liNewDN, tmpDN; + + TASSIGN(oiNewTile, 0); + TASSIGN(oiTile, kDataBytes); + TASSIGN(mijDN, 2 * kDataBytes); + TASSIGN(lijDN, 2 * kDataBytes + kScalarDNBytes); + TASSIGN(miDN, 2 * kDataBytes + 2 * kScalarDNBytes); + TASSIGN(liDN, 2 * kDataBytes + 3 * kScalarDNBytes); + TASSIGN(miNewDN, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaDN, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaDN, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewDN, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpDN, 2 * kDataBytes + 8 * kScalarDNBytes); + + if (is_first) { + // --- First block: copy inputs to accumulators --- + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Store mi = mij, li = lij, oi = oi_new + // Alias ND tiles to the same UB as DN tiles for storing as ND format + TileScalarND mijND, lijND; + TASSIGN(mijND, 2 * kDataBytes); // alias same UB as mijDN + TASSIGN(lijND, 2 * kDataBytes + kScalarDNBytes); // alias same UB as lijDN + + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, mijND); // mi = mij + TSTORE(liGlobalND, lijND); // li = lij + TSTORE(oiGlobal, oiNewTile); // oi = oi_new + + if (is_last) { + // Single block: normalize dst = oi_new / lij + // lijDN already in ColMajor DN format, use directly for TROWEXPANDDIV + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TROWEXPANDDIV(oiNewTile, oiNewTile, lijDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(dstGlobal, oiNewTile); + } + } else { + // --- Subsequent blocks: accumulate --- + + // Load all inputs + TLOAD(oiNewTile, oiNewGlobal); + TLOAD(oiTile, oiGlobal); + TLOAD(mijDN, mijGlobalDN); + TLOAD(lijDN, lijGlobalDN); + TLOAD(miDN, miGlobalDN); + TLOAD(liDN, liGlobalDN); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // TRESHAPE: ColMajor(M,1) → RowMajor(1,M) for element-wise arithmetic + TileScalarRow miRow, mijRow, liRow, lijRow; + TRESHAPE(miRow, miDN); + TRESHAPE(mijRow, mijDN); + TRESHAPE(liRow, liDN); + TRESHAPE(lijRow, lijDN); + + // Scalar arithmetic in RowMajor (1, M) layout + TileScalarRow miNewRow, alphaRow, betaRow, liNewRow, tmpRow; + TASSIGN(miNewRow, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(alphaRow, 2 * kDataBytes + 5 * kScalarDNBytes); + TASSIGN(betaRow, 2 * kDataBytes + 6 * kScalarDNBytes); + TASSIGN(liNewRow, 2 * kDataBytes + 7 * kScalarDNBytes); + TASSIGN(tmpRow, 2 * kDataBytes + 8 * kScalarDNBytes); + + TMAX(miNewRow, miRow, mijRow); // mi_new = max(mi, mij) + pipe_barrier(PIPE_V); + TSUB(alphaRow, miRow, miNewRow); // alpha_exp = mi - mi_new + pipe_barrier(PIPE_V); + TEXP(alphaRow, alphaRow); // alpha = exp(mi - mi_new) + pipe_barrier(PIPE_V); + TSUB(betaRow, mijRow, miNewRow); // beta_exp = mij - mi_new + pipe_barrier(PIPE_V); + TEXP(betaRow, betaRow); // beta = exp(mij - mi_new) + pipe_barrier(PIPE_V); + TMUL(tmpRow, alphaRow, liRow); // alpha * li + pipe_barrier(PIPE_V); + TMUL(liNewRow, betaRow, lijRow); // beta * lij + pipe_barrier(PIPE_V); + TADD(liNewRow, tmpRow, liNewRow); // li_new = alpha*li + beta*lij + + // TRESHAPE back: RowMajor(1,M) → ColMajor(M,1) for TROWEXPANDMUL + TRESHAPE(alphaDN, alphaRow); + TRESHAPE(betaDN, betaRow); + + // Scale data tiles using row-broadcast multiply + TROWEXPANDMUL(oiTile, oiTile, alphaDN); // oi *= alpha + TROWEXPANDMUL(oiNewTile, oiNewTile, betaDN); // oi_new *= beta + pipe_barrier(PIPE_V); + TADD(oiTile, oiTile, oiNewTile); // oi = alpha*oi + beta*oi_new + + // Store mi_new and li_new to GM (ND format) + // Alias ND tiles to the same UB locations as miNewRow and liNewRow + TileScalarND miNewND, liNewND; + TASSIGN(miNewND, 2 * kDataBytes + 4 * kScalarDNBytes); + TASSIGN(liNewND, 2 * kDataBytes + 7 * kScalarDNBytes); + + if (is_last) { + // Normalize and output: dst = oi / li_new + TRESHAPE(liNewDN, liNewRow); + pipe_barrier(PIPE_V); + TROWEXPANDDIV(oiTile, oiTile, liNewDN); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(dstGlobal, oiTile); + } else { + // Store updated accumulators + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(miGlobalND, miNewND); // persist mi_new + TSTORE(liGlobalND, liNewND); // persist li_new + TSTORE(oiGlobal, oiTile); + } + } + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *oi_new = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *mi = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ Tensor *li = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ Tensor *oi = reinterpret_cast<__gm__ Tensor *>(args[5]); + __gm__ Tensor *dst = reinterpret_cast<__gm__ Tensor *>(args[6]); + uint64_t is_first = static_cast(args[7]); + uint64_t is_last = static_cast(args[8]); + uint64_t q_tile_size = static_cast(mij->shapes[0]); + + if (q_tile_size == 16 && oi_new->shapes[1] <= 16) { + online_update_impl<16, 16>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else if (q_tile_size == 16) { + online_update_impl<16, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } else { + online_update_impl<64, 128>(mij, lij, oi_new, mi, li, oi, is_first, is_last, dst); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp new file mode 100644 index 0000000000..8f0c41775d --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/aiv/aiv_softmax_prepare.cpp @@ -0,0 +1,156 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Softmax Preparation Kernel (AIV) with partial block masking +// +// Operates on (M, N) tile where M=q_tile_size, N=block_size: +// Case1: sij is (16, 128) +// Case2: sij is (64, 64) +// +// For partial blocks (valid_len < N), positions [valid_len, N) in sij are +// filled with -inf via TFILLPAD_INPLACE before softmax, ensuring exp(-inf)=0 +// so that invalid key positions contribute zero attention weight. +// +// Computes: +// sij_masked = TFILLPAD(sij, valid_len, pad=-inf) +// sij_scale = sij_masked * scale +// mij = row_max(sij_scale) -> (M, 1) +// pij = exp(sij_scale - mij) -> (M, N) +// lij = row_sum(pij) -> (M, 1) + +#include +#include + +#include "tensor.h" + +using namespace pto; + +#include "pipe_sync.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +template +static __aicore__ void softmax_prepare_impl( + __gm__ Tensor *sij, float scale_value, __gm__ Tensor *pij, __gm__ Tensor *mij, __gm__ Tensor *lij +) { + uint64_t valid_len = static_cast(sij->shapes[1]); + __gm__ float *sij_addr = reinterpret_cast<__gm__ float *>(sij->buffer.addr); + __gm__ bfloat16_t *pij_addr = reinterpret_cast<__gm__ bfloat16_t *>(pij->buffer.addr); + __gm__ float *mij_addr = reinterpret_cast<__gm__ float *>(mij->buffer.addr); + __gm__ float *lij_addr = reinterpret_cast<__gm__ float *>(lij->buffer.addr); + + constexpr int kAlignedRows = ((M * sizeof(float) + 31) / 32) * (32 / sizeof(float)); + + using GlobalDataMxN = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalDataMxN_bf16 = GlobalTensor, Stride<1, 1, 1, N, 1>>; + using GlobalScalarDN = GlobalTensor, Stride<1, 1, 1, 1, 1>, Layout::DN>; + + GlobalDataMxN sijGlobal(sij_addr + sij->start_offset); + GlobalDataMxN_bf16 pijGlobal(pij_addr + pij->start_offset); + GlobalScalarDN mijGlobal(mij_addr + mij->start_offset); + GlobalScalarDN lijGlobal(lij_addr + lij->start_offset); + + // Dynamic-cols tile: marks which columns are valid for TFILLPAD boundary + using TileSijDyn = Tile; + // Padded tile: TFILLPAD_INPLACE fills positions [valid_len, N) with -inf + using TileSijPad = Tile; + + using TileVecMxN = Tile; + using TileVecMxN_bf16 = Tile; + using TileScalarDN = Tile; + + TileVecMxN sijTile; + TileSijDyn sijDynTile(static_cast(valid_len)); + TileSijPad sijPadTile; + TileVecMxN pijTile; + TileVecMxN tmpTile; + TileScalarDN maxTile; + TileScalarDN sumTile; + TileVecMxN_bf16 pijBf16Tile; + + // All sij tiles share UB address 0x0 (in-place masking) + TASSIGN(sijTile, 0x0); + TASSIGN(sijDynTile, 0x0); + TASSIGN(sijPadTile, 0x0); + TASSIGN(pijTile, M * N * sizeof(float)); + TASSIGN(tmpTile, 2 * M * N * sizeof(float)); + TASSIGN(maxTile, 3 * M * N * sizeof(float)); + TASSIGN(sumTile, 3 * M * N * sizeof(float) + kAlignedRows * sizeof(float)); + TASSIGN(pijBf16Tile, 3 * M * N * sizeof(float) + 2 * kAlignedRows * sizeof(float)); + + // Load full sij (M, N) tile from GM - all N columns including garbage for partial blocks + // printf("sij addr incore %x\n", sij->buffer.addr); + TLOAD(sijTile, sijGlobal); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + + // Mask columns [valid_len, N) with -inf. sijDynTile provides the valid boundary, + // sijPadTile provides PadValue::Min as the fill value. No-op when valid_len == N. + TFILLPAD_INPLACE(sijPadTile, sijDynTile); + pipe_barrier(PIPE_V); + + TMULS(sijTile, sijTile, scale_value); + pipe_barrier(PIPE_V); + TROWMAX(maxTile, sijTile, tmpTile); + pipe_barrier(PIPE_V); + TROWEXPANDSUB(pijTile, sijTile, maxTile); + pipe_barrier(PIPE_V); + TEXP(pijTile, pijTile); + // Truncate pij to bf16 first + pipe_barrier(PIPE_V); + TCVT(pijBf16Tile, pijTile, RoundMode::CAST_ROUND); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // pij bf16 ready, can store early + + // Continue computing: bf16 → f32 and rowsum while pij store proceeds in parallel + pipe_barrier(PIPE_V); + TCVT(pijTile, pijBf16Tile, RoundMode::CAST_ROUND); + pipe_barrier(PIPE_V); + TROWSUM(sumTile, pijTile, tmpTile); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); // sum ready + + // Store pij (overlaps with TCVT + TROWSUM above) + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(pijGlobal, pijBf16Tile); + + // Store max and sum + TSTORE(mijGlobal, maxTile); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(lijGlobal, sumTile); + + pipe_sync(); +} + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *sij = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *pij = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *mij = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *lij = reinterpret_cast<__gm__ Tensor *>(args[3]); + union { + uint64_t u; + float f; + } scale_conv; + scale_conv.u = static_cast(args[4]); + float scale_value = scale_conv.f; + uint64_t q_tile_size = static_cast(sij->shapes[0]); + + if (q_tile_size == 16 && pij->shapes[1] <= 16) { + softmax_prepare_impl<16, 16>(sij, scale_value, pij, mij, lij); + } else if (q_tile_size == 16) { + softmax_prepare_impl<16, 128>(sij, scale_value, pij, mij, lij); + } else { + softmax_prepare_impl<64, 64>(sij, scale_value, pij, mij, lij); + } +} diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp new file mode 100644 index 0000000000..37d545dd17 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/kernels/orchestration/paged_attention_orch.cpp @@ -0,0 +1,311 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Paged Attention Orchestration Function - manual-scope variant + * + * Matches the small-case paged_attention orchestration shape while replacing + * the automatic same-scope dependency wiring with explicit task-to-task deps + * inside PTO2_SCOPE(PTO2ScopeMode::MANUAL). + */ + +#include +#include +#include +#include + +#include "pto_orchestration_api.h" + +#define FUNC_QK_MATMUL 0 +#define FUNC_SOFTMAX_PREPARE 1 +#define FUNC_PV_MATMUL 2 +#define FUNC_ONLINE_UPDATE 3 +constexpr uint64_t PLATFORM_PROF_SYS_CNT_FREQ = 50000000; // 50 MHz + +inline double cycles_to_us(uint64_t cycles) { + return (static_cast(cycles) / PLATFORM_PROF_SYS_CNT_FREQ) * 1000000.0; +} + +inline uint64_t get_sys_cnt_aicpu() { +#if defined(__aarch64__) + uint64_t ticks; + asm volatile("mrs %0, cntvct_el0" : "=r"(ticks)); + return ticks; +#elif defined(__x86_64__) + return 0; +#else + return 0; +#endif +} + +#ifdef ENABLE_PROFILING +#define CYCLE_COUNT_START() uint64_t _t0 = get_sys_cnt_aicpu(), _t1 +#define CYCLE_COUNT_LAP(acc) \ + do { \ + _t1 = get_sys_cnt_aicpu(); \ + acc += (_t1 - _t0); \ + _t0 = _t1; \ + } while (0) +#define PROF_INC(counter, n) (counter) += (n) +#else +#define CYCLE_COUNT_START() (void)0 +#define CYCLE_COUNT_LAP(acc) (void)0 +#define PROF_INC(counter, n) (void)0 +#endif + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 7, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { +#ifdef ENABLE_PROFILING + uint64_t prof_param_extract = 0; + uint64_t prof_ext_tensor = 0; + uint64_t prof_scope = 0; + uint64_t prof_make_tensor = 0; + uint64_t prof_tensor_view = 0; + uint64_t prof_param_setup = 0; + uint64_t prof_submit_task = 0; + int prof_submit_count = 0; + int prof_make_count = 0; + int prof_view_count = 0; +#endif + + CYCLE_COUNT_START(); + + // Read dimensions from tensor metadata + uint64_t batch = orch_args.tensor(0).ref().shapes[0]; + uint64_t num_heads = orch_args.tensor(0).ref().shapes[1]; + uint64_t head_dim = orch_args.tensor(0).ref().shapes[2]; + DataType data_type = orch_args.tensor(0).ref().dtype; + + uint64_t block_size = orch_args.tensor(1).ref().shapes[1]; + uint64_t block_num = orch_args.tensor(3).ref().shapes[1]; + + uint64_t scale_value = orch_args.scalar(0); + + uint64_t q_head_num = num_heads; + uint64_t q_tile = std::min(num_heads, static_cast(128)); + uint64_t q_loop = (q_head_num + q_tile - 1) / q_tile; + CYCLE_COUNT_LAP(prof_param_extract); + + LOG_INFO(">>>>>> batch = %" PRIu64, batch); + + // Reshape tensors for kernel consumption (2D flattened) + void *query_ptr = orch_args.tensor(0).ref().data_as(); + void *kc_ptr = orch_args.tensor(1).ref().data_as(); + void *vc_ptr = orch_args.tensor(2).ref().data_as(); + void *out_ptr = orch_args.tensor(5).ref().data_as(); + + uint64_t total_blocks_count = orch_args.tensor(1).ref().shapes[0]; + + uint32_t query_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + uint32_t key_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t value_cache_shapes[2] = { + static_cast(total_blocks_count * block_size), static_cast(head_dim) + }; + uint32_t out_shapes[2] = {static_cast(batch * num_heads), static_cast(head_dim)}; + Tensor query = make_tensor_external(query_ptr, query_shapes, 2, data_type); + Tensor key_cache = make_tensor_external(kc_ptr, key_cache_shapes, 2, data_type); + Tensor value_cache = make_tensor_external(vc_ptr, value_cache_shapes, 2, data_type); + Tensor out = make_tensor_external(out_ptr, out_shapes, 2, DataType::FLOAT32); + CYCLE_COUNT_LAP(prof_ext_tensor); + + uint32_t bt_shapes[2] = {static_cast(batch), static_cast(block_num)}; + Tensor block_table = + make_tensor_external(orch_args.tensor(3).ref().data_as(), bt_shapes, 2, DataType::INT32, false); + uint32_t cl_shapes[1] = {static_cast(batch)}; + Tensor context_lens = + make_tensor_external(orch_args.tensor(4).ref().data_as(), cl_shapes, 1, DataType::INT32, false); + + // Create infos are loop-invariant — shapes depend only on q_tile/head_dim/block_size + uint32_t tile2d_shapes[2] = {static_cast(q_tile), static_cast(head_dim)}; + uint32_t scalar_shapes[1] = {static_cast(q_tile)}; + uint32_t sij_shapes[2] = {static_cast(q_tile), static_cast(block_size)}; + TensorCreateInfo tile2d_ci(tile2d_shapes, 2, DataType::FLOAT32); + TensorCreateInfo scalar_ci(scalar_shapes, 1, DataType::FLOAT32); + TensorCreateInfo sij_ci(sij_shapes, 2, DataType::FLOAT32); + TensorCreateInfo pij_f16_ci(sij_shapes, 2, data_type); + + PROF_INC(prof_make_count, 4); + CYCLE_COUNT_LAP(prof_make_tensor); + + for (uint64_t b_idx = 0; b_idx < batch; b_idx++) { + uint32_t cl_idx[1] = {static_cast(b_idx)}; + uint64_t cur_seq = static_cast(get_tensor_data(context_lens, 1, cl_idx)); + uint64_t bn_this_batch = (cur_seq + block_size - 1) / block_size; + for (uint64_t q_idx = 0; q_idx < q_loop; q_idx++) { + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + CYCLE_COUNT_LAP(prof_scope); + uint64_t cur_offset = b_idx * q_head_num + q_idx * q_tile; + + uint32_t qi_offsets[2] = {static_cast(cur_offset), 0}; + Tensor qi = query.view(tile2d_shapes, qi_offsets); + uint32_t out_view_offsets[2] = {static_cast(cur_offset), 0}; + Tensor out_view = out.view(tile2d_shapes, out_view_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors alloc_outs = alloc_tensors(tile2d_ci, scalar_ci, scalar_ci); + const Tensor &oi = alloc_outs.get_ref(0); + const Tensor &li_update = alloc_outs.get_ref(1); + const Tensor &mi_update = alloc_outs.get_ref(2); + PTO2TaskId alloc_task = alloc_outs.task_id(); + PTO2TaskId prev_update_task = PTO2TaskId::invalid(); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + for (uint64_t bn = 0; bn < bn_this_batch; bn++) { + uint32_t bt_idx[2] = {static_cast(b_idx), static_cast(bn)}; + uint64_t cur_block_idx = static_cast(get_tensor_data(block_table, 2, bt_idx)); + uint64_t valid_len = std::min(block_size, cur_seq - bn * block_size); + CYCLE_COUNT_LAP(prof_param_extract); + + uint32_t kv_shapes[2] = {static_cast(block_size), static_cast(head_dim)}; + uint32_t kv_offsets[2] = {static_cast(cur_block_idx * block_size), 0}; + Tensor kj = key_cache.view(kv_shapes, kv_offsets); + Tensor vj = value_cache.view(kv_shapes, kv_offsets); + PROF_INC(prof_view_count, 2); + CYCLE_COUNT_LAP(prof_tensor_view); + + L0TaskArgs params_qk; + params_qk.add_input(qi); + params_qk.add_input(kj); + params_qk.add_output(sij_ci); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors qk_outs = rt_submit_aic_task(FUNC_QK_MATMUL, params_qk); + const Tensor &sij = qk_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint32_t sij_valid_shapes[2] = {static_cast(q_tile), static_cast(valid_len)}; + uint32_t sij_valid_offsets[2] = {0, 0}; + Tensor sij_valid = sij.view(sij_valid_shapes, sij_valid_offsets); + PROF_INC(prof_view_count, 1); + CYCLE_COUNT_LAP(prof_tensor_view); + + // --- Primitive dep API (Arg + set_dependencies) --- + // Caller owns the deps buffer; Arg stores (ptr, count). + // Suited for codegen and for cases with a fixed dep set. + L0TaskArgs params_sf; + params_sf.add_input(sij_valid); + params_sf.add_output(pij_f16_ci); + params_sf.add_output(scalar_ci); + params_sf.add_output(scalar_ci); + PTO2TaskId sf_deps[] = {qk_outs.task_id()}; + params_sf.set_dependencies(sf_deps, 1); + params_sf.add_scalar(scale_value); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors sf_outs = rt_submit_aiv_task(FUNC_SOFTMAX_PREPARE, params_sf); + const Tensor &pij_f16 = sf_outs.get_ref(0); + const Tensor &mi = sf_outs.get_ref(1); + const Tensor &li = sf_outs.get_ref(2); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + L0TaskArgs params_pv; + params_pv.add_input(pij_f16); + params_pv.add_input(vj); + params_pv.add_output(tile2d_ci); + PTO2TaskId pv_deps[] = {sf_outs.task_id()}; + params_pv.set_dependencies(pv_deps, 1); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors pv_outs = rt_submit_aic_task(FUNC_PV_MATMUL, params_pv); + const Tensor &oi_tmp = pv_outs.get_ref(0); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + + uint64_t is_first = (bn == 0) ? 1 : 0; + uint64_t is_last = (bn == bn_this_batch - 1) ? 1 : 0; + CYCLE_COUNT_LAP(prof_param_extract); + + // --- Convenience dep API (L0TaskArgsWithDeps + add_dep) --- + // Wrapper owns a stack-sized deps buffer and accepts + // incremental add_dep() calls; the submit overload binds + // them to the underlying Arg via set_dependencies(...). + // Suited for hand-written orch where the dep set is + // assembled conditionally across branches. + L0TaskArgsWithDeps<> params_up; + params_up.add_input(mi); + params_up.add_input(li); + params_up.add_input(oi_tmp); + params_up.add_inout(mi_update); + params_up.add_inout(li_update); + params_up.add_inout(oi); + params_up.add_inout(out_view); + // UP reads SF's mi/li, but SF -> PV -> UP already orders it; only the PV edge is explicit. + params_up.add_dep(pv_outs.task_id()); + if (prev_update_task.is_valid()) { + params_up.add_dep(prev_update_task); + } + // alloc completes inline; this dep only keeps the scratch buffers alive until the last consumer. + if (is_last) { + params_up.add_dep(alloc_task); + } + params_up.add_scalar(is_first); + params_up.add_scalar(is_last); + CYCLE_COUNT_LAP(prof_param_setup); + TaskOutputTensors up_outs = rt_submit_aiv_task(FUNC_ONLINE_UPDATE, params_up); + prev_update_task = up_outs.task_id(); + PROF_INC(prof_submit_count, 1); + CYCLE_COUNT_LAP(prof_submit_task); + } + } + CYCLE_COUNT_LAP(prof_scope); + } + } + +#ifdef ENABLE_PROFILING + uint64_t total = prof_param_extract + prof_ext_tensor + prof_make_tensor + prof_tensor_view + prof_param_setup + + prof_submit_task + prof_scope; + LOG_INFO( + "=== PagedAttn Orch Profiling: %d submits, %d makes, %d views, total=%.3fus ===", prof_submit_count, + prof_make_count, prof_view_count, cycles_to_us(total) + ); + if (total > 0) { + LOG_INFO( + " param_extract : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_extract), + prof_param_extract * 100.0 / total + ); + LOG_INFO( + " ext_tensor(x4) : %7.3fus (%5.1f%%)", cycles_to_us(prof_ext_tensor), prof_ext_tensor * 100.0 / total + ); + LOG_INFO( + " create_info(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_make_count, cycles_to_us(prof_make_tensor), + prof_make_tensor * 100.0 / total, + prof_make_count > 0 ? cycles_to_us(prof_make_tensor) / prof_make_count : 0.0 + ); + LOG_INFO( + " tensor_view(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_view_count, cycles_to_us(prof_tensor_view), + prof_tensor_view * 100.0 / total, + prof_view_count > 0 ? cycles_to_us(prof_tensor_view) / prof_view_count : 0.0 + ); + LOG_INFO( + " param_setup : %7.3fus (%5.1f%%)", cycles_to_us(prof_param_setup), prof_param_setup * 100.0 / total + ); + LOG_INFO(" scope : %7.3fus (%5.1f%%)", cycles_to_us(prof_scope), prof_scope * 100.0 / total); + LOG_INFO( + " submit_task(x%d) : %7.3fus (%5.1f%%) avg=%.3fus", prof_submit_count, cycles_to_us(prof_submit_task), + prof_submit_task * 100.0 / total, + prof_submit_count > 0 ? cycles_to_us(prof_submit_task) / prof_submit_count : 0.0 + ); + } +#endif +} + +} // extern "C" diff --git a/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py b/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py new file mode 100644 index 0000000000..7c4cdb8ca8 --- /dev/null +++ b/examples/a2a3/host_build_graph/paged_attention_manual_scope/test_paged_attention_manual_scope.py @@ -0,0 +1,200 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Paged attention manual-scope wrapper for A2A3 host_build_graph. + +Replaces the automatic same-scope dependency wiring with explicit task-to-task +deps inside PTO2_SCOPE(PTO2ScopeMode::MANUAL). The orchestration source is the +same one the tensormap_and_ringbuffer variant compiles; only the runtime named +by @scene_test differs. + +host_build_graph submits the whole graph before the device schedules, so the +ring window must hold every task of a case at once — larger than the +tensormap_and_ringbuffer variant needs, where slots retire as orchestration +proceeds. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.goldens.paged_attention import compute_golden as _pa_compute_golden +from simpler_setup.goldens.paged_attention import generate_inputs as _pa_generate_inputs + +# tasks = batch * (4 * ceil(context_len / block_size) + 1). ring_task_window must +# be a power of two in [4, INT32_MAX] and is rejected otherwise by the runtime. +_RING_65K = {"ring_task_window": 131072, "ring_heap": 2 * 1024 * 1024 * 1024} +_RING_33K = {"ring_task_window": 65536, "ring_heap": 1024 * 1024 * 1024} + + +@scene_test(level=2, runtime="host_build_graph") +class TestPagedAttentionManualScopeHostBuildGraph(SceneTestCase): + RTOL = 1e-3 + ATOL = 1e-3 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/paged_attention_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "QK", + "source": "kernels/aic/aic_qk_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "SF", + "source": "kernels/aiv/aiv_softmax_prepare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.OUT, D.OUT, D.OUT], + }, + { + "func_id": 2, + "name": "PV", + "source": "kernels/aic/aic_pv_matmul.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 3, + "name": "UP", + "source": "kernels/aiv/aiv_online_update.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + ], + } + + CASES = [ + { + # 65 792 tasks, all live at once, in a single MANUAL scope. + "name": "Case1", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_65K}, + "manual": True, + "params": { + "batch": 256, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 128, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + # 32 832 tasks. + "name": "Case2", + "platforms": ["a2a3"], + "config": {"aicpu_thread_num": 4, "runtime_env": _RING_33K}, + "manual": True, + "params": { + "batch": 64, + "num_heads": 64, + "kv_head_num": 1, + "head_dim": 128, + "block_size": 64, + "context_len": 8192, + "max_model_len": 32768, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseSmall2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 1, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq2", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 2, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 33, + "context_lens_list": [33, 17], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + { + "name": "CaseVarSeq4", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4}, + "manual": True, + "params": { + "batch": 4, + "num_heads": 16, + "kv_head_num": 1, + "head_dim": 16, + "block_size": 16, + "context_len": 128, + "context_lens_list": [33, 64, 128, 15], + "max_model_len": 256, + "dtype": "bfloat16", + }, + }, + ] + + def generate_args(self, params): + result = _pa_generate_inputs(params) + specs = [] + for name, value in result: + if isinstance(value, torch.Tensor): + specs.append(Tensor(name, value)) + else: + specs.append(Scalar(name, value)) + return TaskArgsBuilder(*specs) + + def compute_golden(self, args, params): + tensors = {s.name: s.value for s in args.specs if isinstance(s, Tensor)} + _pa_compute_golden(tensors, params) + for s in args.specs: + if isinstance(s, Tensor) and s.name in tensors: + getattr(args, s.name)[:] = tensors[s.name] + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/README.md b/examples/a2a3/host_build_graph/qwen3_14b_decode/README.md new file mode 100644 index 0000000000..3e91143eb5 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/README.md @@ -0,0 +1,206 @@ +# `qwen3_14b_decode/` — Qwen3-14B 40-layer decode (CANN fused attention), host_build_graph + +Self-contained SceneTestCase port of pypto-lib +`models/qwen3/14b/decode_fwd.py` entry `decode_fwd_layers` with +`_CHUNK_NLAYERS == 40`: **the whole Qwen3-14B decode stack as one fused +dispatch** (hidden → hidden, no LM head), with the FP32 inter-layer residual +carry. A simpler developer builds and runs it directly — no descent through +pypto-lib / the JIT, no auto-built intermediate artifacts. + +The layer loop is a real loop in the generated orchestration +(`for (int64_t i = 0; i < 40; i += 1)`), not an unrolled one, so a 40-layer +chunk costs one extra literal over a 2-layer chunk rather than 20× the kernels. + +## Parameter regime — matches `stress_profile.py` + +The fixture mirrors the vLLM serving stress run (`stress_profile.py`): + +| Param | Value | Source | +| ----- | ----- | ------ | +| `BATCH` | 16 | `CONCURRENCY` (aligned with decode kernel BATCH=16) | +| `MAX_SEQ` | 5500 | `max_model_len` (KV-pool / RoPE sizing) | +| decode `seq_len` | 3500 | the ~3500-token prompt | +| layers | 40 | full model (`decode_fwd_layers` N=40) | + +Per the lib's const-layer-0 stacked-fwd reference, every layer reuses layer-0 +weights (weights + paged KV pool are stacked ×40 along dim 0, one slice per +layer); each layer still reads and writes its own KV pool. + +Footprint at this regime, bf16: + +| component | size | +| --------- | ---: | +| weights (×40) | 24.61 GiB | +| paged KV pool (×40) | 13.44 GiB | +| **total fixture** | **38.05 GiB** | + +Die HBM is 64 GiB, so it fits with ~26 GiB of headroom for ring heap and +workspace. + +**This runtime needs the ring heap raised to 512 MiB**, which the case carries in +its own `runtime_env`. On `tensormap_and_ringbuffer` each layer's intermediates +live inside that iteration's scope and are freed at its end, so the live set does +not grow with layer count and the defaults carry the graph. `host_build_graph` +builds the entire graph on the host before the device schedules anything, so no +task has completed while the graph is being built, the heap tail never advances +off 0, and all 40 layers' intermediates are live at once. Measured on a2a3: + +| `ring_heap` | outcome | +| ----------- | ------- | +| 256 MiB (default) | `Task Allocator Deadlock - Heap Exhausted`, tail=0, ~254 MiB of 256 used | +| 512 MiB | passes | + +The task window is *not* the constraint: the graph is ~10.6K tasks, inside the +16384 default. + +## Dataflow per layer (`_decode_layer`) + +input RMSNorm → split-K SPMD Q/K/V (seed + atomic-add) → **`paged_attention_rope_cce`** +→ split-K out_proj + residual → post-RMSNorm → SwiGLU FFN → `dcr_xgamma`. + +`copy_hidden` embeds the bf16 input; the FP32 residual is carried between +layers; `copy_out` does the single FP32→bf16 round at the chunk tail. + +The attention stage is one **CANN `FusedInferAttentionScore` extern** that +subsumes what used to be seven generated kernels — it folds per-head Q/K +RMS-norm, RoPE, the paged KV write, the flash-attention inner loop and the +online softmax into a single mixed (AIC + 2×AIV) task, gated by an +`AscendC::SyncAll()` FFTS barrier. `paged_attention_tiling_cce` builds +its runtime tiling metadata first. + +Paged KV uses vLLM's **BSND** layout: a page holds `[BLOCK_SIZE, KV_HIDDEN]` +ordered `[page, token, kv_head, dim]`, so `slot_mapping[b]` is directly the +row index. (The previous harvest used NSND; the golden was updated to match.) + +## Provenance — how the C++ was produced + +| component | source | +| --------- | ------ | +| pypto-lib | `45be52c` | +| pypto | `d64380cb` | +| ptoas | `v0.48` | +| pto-isa | `83d01313d9bfc247c4b7c8bcf969d1019f0d106f` (`pto_isa.pin`) | + +`kernels/orchestration/` + `kernels/aic/` (18) + `kernels/aiv/` (16) are +harvested pypto codegen for `decode_fwd_layers` (`_CHUNK_NLAYERS=40`, +`PTO2_MANUAL_MAX_SEQ=5500`) — license header prepended, otherwise verbatim. +The `CALLABLE` is transcribed from that run's `kernel_config.py`, and +`simpler_setup/goldens/qwen3_14b_decode.py` ports the per-layer +`golden_decode_layer` math (RoPE θ=1e4, controlled scales, FP32 residual, bf16 +cast points) composed over 40 layers with FP32 carry + per-layer KV pools. + +**There are no hand-edits.** The previous harvest patched `fa_fused_aiv` to work +around a `[[block_local]] static`; that kernel no longer exists (attention is +the extern now), and the current codegen emits no such construct. + +### `kernels/vendor/paged_attention_cce/` — the attention extern + +Copied verbatim from pypto-lib +`models/qwen3/14b/kernels/paged_attention_cce/`. The tree must stay intact: +`kernel/fai_body.hpp` reaches its dependencies through relative includes, so +splitting it would mean patching the source and re-patching on every refresh. + +- `attention_rope/`, `tiling/`, `kernel/`, `generated/` — PyPTO-authored glue. +- `vendor/fused_infer_attention_score/` — **CANN `FusedInferAttentionScore`, + Copyright (c) 2025 Huawei Technologies, CANN Open Software License 2.0** + (~16 k LOC), upstream's own vendored copy, left where upstream put it. +- `attention/` is the non-RoPE variant of the extern. `decode_fwd_layers` does + not use it; it is kept so a refresh is a plain directory copy. + +**Why it sits under `kernels/vendor/`.** Nothing below a `vendor/` directory is +ours to reformat: the repo's header, formatting and language lint all skip that +path (`.pre-commit-config.yaml`, `tests/lint/check_headers.py`). Without that, +`clang-format` rewrites the glue files and `end-of-file-fixer` touches the CANN +headers, and the next refresh diffs against *our* reformatting instead of +against upstream — which is how the drift this example is meant to expose starts. +The carve-out keys on the directory, not on this operator's name, so harvesting +another extern is a matter of dropping it in `kernels/vendor/` with no lint +change at all. + +Building these needs **CANN devkit headers** (`$ASCEND_HOME_PATH/aarch64-linux/asc/…`, +`tikcpp/…`), declared per-incore via `extra_include_dirs` in the `CALLABLE`. +`$ASCEND_HOME_PATH` keeps them machine-independent; paths a given CANN layout +does not ship are dropped rather than failing the build. + +They also depend on simpler linking incore objects before extracting `.text` +([#1497](https://github.com/hw-native-sys/simpler/pull/1497)): AscendC declares +`g_vecTPipePtr` / `g_kfcClient` as block-local globals, whose relocations no +amount of inlining removes. + +### To regenerate + +From a simpler worktree with pypto + pypto-lib cloned under `build/` (see the +[`multi-repo-setup`](../../../../.claude/skills/multi-repo-setup/SKILL.md) +skill) and `eval "$(pypto-setup --export)"`: + +```python +# PTO2_MANUAL_MAX_SEQ must be set before importing decode_fwd (read at import). +os.environ["PTO2_MANUAL_MAX_SEQ"] = "5500" +D = +D._CHUNK_NLAYERS = 40 # read at trace time; rebind before the first call +D.decode_fwd_layers(*inputs, out, config=RunConfig( + platform="a2a3", codegen_only=True, save_kernels=True, save_kernels_dir=OUT)) +``` + +`codegen_only` needs no device. Then copy `OUT/orchestration/`, `OUT/kernels/` +and `models/qwen3/14b/kernels/paged_attention_cce/` into `kernels/vendor/` here, and +re-transcribe `CALLABLE` from `OUT/kernel_config.py` (which already records +`func_id`, `core_type`, per-kernel `signature`, and `extra_include_dirs`). + +One deliberate deviation from `kernel_config.py`: `decode_fwd_layers` declares +`k_cache` / `v_cache` as plain inputs, but the extern writes the current token's +KV into them. The `CALLABLE` marks them `INOUT` so simpler copies the pools back +and the golden can verify all 40 layers' KV writes, not just the hidden output. + +## Running + +```bash +# pytest (hardware; wrap in task-submit on shared boxes) +pytest examples/a2a3/host_build_graph/qwen3_14b_decode \ + --platform a2a3 --device ${DEVICE} + +# standalone +python examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py -p a2a3 -d ${DEVICE} +``` + +DFX is opt-in via the existing flags — no kernel changes needed: + +```bash +pytest .../qwen3_14b_decode --platform a2a3 --device ${DEVICE} \ + --enable-l2-swimlane 1 --enable-dep-gen +``` + +Note that `--enable-dep-gen` / `--enable-l2-swimlane` on the full 40-layer graph +can overflow the per-run SHM record buffer ("records dropped"); pypto-lib warns +about the same thing for `decode_fwd.py --fwd-layers`. Capture on a smaller +harvest if you need a clean trace. + +## Status — PASSING + +Passes on device: output **and all 40 layers' KV caches** match the torch +reference at `RTOL=5e-2 / ATOL=1e-1`. + +## Cost + +It runs in the general `st-onboard-a2a3` scene-test sweep like any other case. +Measured on this repo's a2a3 box, one device: + +| Phase | Wall | +| ----- | ---- | +| kernel compilation — 36 incores + 1 orchestration | **59 s** | +| `generate_inputs` (the 38 GiB fixture) | **13 s** | +| `compute_golden` (40 layers, torch, thread-capped) | **~43 s** | +| host→device upload, device run, comparison | the remainder; the device itself is busy for tens of ms | +| | **~115 s** | + +Two details behind those numbers. The 38 GiB is mostly replication — +`torch.cat([w] * 40)` of one layer's weights — so only ~2 GB is actually +generated; and the vendor FAI kernel is the compile outlier at 8.4 s for its AIV +variant against ~1.2 s for an ordinary incore. + +The golden is thread-capped by the scene-test framework, and that matters more +than its size: its 3584 small slice operations per layer took **6.35 s per layer +at 320 torch threads against 1.05 s at 4**, so on a many-core host the untamed +thread pool cost 5-8x the work it was doing. Before the cap this case held a +device for **406 s** (ci run 30507320146, job 90761014912) and had to be kept out +of the sweep with a `--ignore` and a step of its own. diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/down_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/down_proj.cpp new file mode 100644 index 0000000000..774948fce5 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/down_proj.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: down_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int64_t v6) { + const int64_t v7 = 960; + const int64_t v8 = 2; + const int64_t v9 = 15; + const int64_t v10 = 32; + const int64_t v11 = 64; + const int64_t v12 = 5120; + const int64_t v13 = 1; + const int64_t v14 = 17408; + const int64_t v15 = 16; + const int64_t v16 = 512; + const int64_t v17 = 2048; + const int64_t v18 = 32768; + const int64_t v19 = 1536; + const int64_t v20 = 1024; + const int64_t v21 = 0; + const int64_t v22 = 135168; + const int64_t v23 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v24 = (size_t)v11; + size_t v25 = (size_t)v21; + size_t v26 = (size_t)v10; + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v27 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v28 = (uint64_t)v23; + TASSIGN(v27, v28); + pto::Shape<1, 1, 1, 16, 64> v29 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v30 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v31 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + v4 * v13), v29, v30 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v27, v31); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v32 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v33 = (uint64_t)v22; + TASSIGN(v32, v33); + int64_t v34 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v35 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v36 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v21 + v34 * v12 + v6 * v13), v35, v36 + ); + TLOAD(v32, v37); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v38 = v25; v38 < v24; v38 += v26) { + int64_t v39 = (int64_t)v38; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v41 = (uint64_t)v20; + TASSIGN(v40, v41); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v40, v27, v21, v38); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + TEXTRACT(v42, v32, v38, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v45 = (uint64_t)v19; + TASSIGN(v44, v45); + int64_t v46 = (int64_t)((uint64_t)v39 + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v44, v27, v21, v46); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v48 = (uint64_t)v18; + TASSIGN(v47, v48); + TEXTRACT(v47, v32, v46, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v39 == v21) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + pipe_barrier(PIPE_M); + TMATMUL(v49, v40, v42); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v52 = (uint64_t)v21; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v51, v51, v40, v42); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v54 = (uint64_t)v21; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v53, v53, v44, v47); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v55 = (size_t)v13; v55 < ((size_t)v9); v55 += (size_t)v8) { + int64_t v56 = (int64_t)((uint64_t)((int64_t)v55) * (uint64_t)v11); + int64_t v57 = (int64_t)((uint64_t)v56 + (uint64_t)v11); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v59 = (uint64_t)v21; + TASSIGN(v58, v59); + pto::Shape<1, 1, 1, 16, 64> v60 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v61 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v62 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, + pto::Layout::ND>(v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v13), v60, v61); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v58, v62); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v64 = (uint64_t)v22; + TASSIGN(v63, v64); + pto::Shape<1, 1, 1, 64, 1024> v65 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v66 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v56) * v12 + v6 * v13), v65, v66); + TLOAD(v63, v67); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v69 = (uint64_t)v17; + TASSIGN(v68, v69); + pto::Shape<1, 1, 1, 16, 64> v70 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v71 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v72 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, + pto::Layout::ND>(v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v57) * v13), v70, v71); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v68, v72); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v74 = (uint64_t)v23; + TASSIGN(v73, v74); + pto::Shape<1, 1, 1, 64, 1024> v75 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v76 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v77 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v57) * v12 + v6 * v13), v75, v76); + TLOAD(v73, v77); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v78 = v25; v78 < v24; v78 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v79 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v80 = (uint64_t)v20; + TASSIGN(v79, v80); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v79, v58, v21, v78); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v81 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + TEXTRACT(v81, v63, v78, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v84 = (uint64_t)v19; + TASSIGN(v83, v84); + int64_t v85 = (int64_t)((uint64_t)((int64_t)v78) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v83, v58, v21, v85); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v86 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v87 = (uint64_t)v18; + TASSIGN(v86, v87); + TEXTRACT(v86, v63, v85, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v88 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v89 = (uint64_t)v21; + TASSIGN(v88, v89); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v88, v88, v79, v81); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v91 = (uint64_t)v21; + TASSIGN(v90, v91); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v90, v90, v83, v86); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v92 = v25; v92 < v24; v92 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v93 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v94 = (uint64_t)v21; + TASSIGN(v93, v94); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v93, v68, v21, v92); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v95 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v96 = (uint64_t)v21; + TASSIGN(v95, v96); + TEXTRACT(v95, v73, v92, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v98 = (uint64_t)v16; + TASSIGN(v97, v98); + int64_t v99 = (int64_t)((uint64_t)((int64_t)v92) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v97, v68, v21, v99); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v100 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v101 = (uint64_t)v18; + TASSIGN(v100, v101); + TEXTRACT(v100, v73, v99, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v102 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v103 = (uint64_t)v21; + TASSIGN(v102, v103); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v102, v102, v93, v95); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v105 = (uint64_t)v21; + TASSIGN(v104, v105); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v104, v104, v97, v100); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + pto::Shape<1, 1, 1, 16, 64> v108 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v109 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v110 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v108, v109 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v106, v110); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v112 = (uint64_t)v22; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 64, 1024> v113 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v114 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v7) * v12 + v6 * v13), v113, v114 + ); + TLOAD(v111, v115); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v116 = v25; v116 < v24; v116 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v117 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v118 = (uint64_t)v20; + TASSIGN(v117, v118); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v117, v106, v21, v116); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + TEXTRACT(v119, v111, v116, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v122 = (uint64_t)v19; + TASSIGN(v121, v122); + int64_t v123 = (int64_t)((uint64_t)((int64_t)v116) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v121, v106, v21, v123); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v125 = (uint64_t)v18; + TASSIGN(v124, v125); + TEXTRACT(v124, v111, v123, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v126, v126, v117, v119); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v129 = (uint64_t)v21; + TASSIGN(v128, v129); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v128, v128, v121, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v131 = (uint64_t)v21; + TASSIGN(v130, v131); + pto::Shape<1, 1, 1, 16, 1024> v132 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v133 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v134 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v21 + v21 * v12 + v6 * v13), v132, v133 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v134, v130); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_tile_inline149__rv_v2 + __gm__ Tensor *mlp_tile_inline149__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_tile_inline149__rv_v2 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_tile_inline149__rv_v2_tensor->buffer.addr) + + mlp_tile_inline149__rv_v2_tensor->start_offset; + + // Unpack tensor: w_down__ssa_v0 + __gm__ Tensor *w_down__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_down__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_down__ssa_v0_tensor->buffer.addr) + w_down__ssa_v0_tensor->start_offset; + + // Unpack tensor: down_acc_all_inline168__iter_v6 + __gm__ Tensor *down_acc_all_inline168__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *down_acc_all_inline168__iter_v6 = + reinterpret_cast<__gm__ float *>(down_acc_all_inline168__iter_v6_tensor->buffer.addr) + + down_acc_all_inline168__iter_v6_tensor->start_offset; + + // Unpack scalar: k0_inline113__ssa_v8 + union { + uint64_t u64; + int64_t val; + } k0_inline113__ssa_v8_conv; + k0_inline113__ssa_v8_conv.u64 = args[3]; + int64_t k0_inline113__ssa_v8 = k0_inline113__ssa_v8_conv.val; + + // Unpack scalar: layer_inter_base_inline107__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_inter_base_inline107__ssa_v0_conv; + layer_inter_base_inline107__ssa_v0_conv.u64 = args[4]; + int64_t layer_inter_base_inline107__ssa_v0 = layer_inter_base_inline107__ssa_v0_conv.val; + + // Unpack scalar: n0_inline122__ssa_v8 + union { + uint64_t u64; + int64_t val; + } n0_inline122__ssa_v8_conv; + n0_inline122__ssa_v8_conv.u64 = args[5]; + int64_t n0_inline122__ssa_v8 = n0_inline122__ssa_v8_conv.val; + + // Forward to ptoas-generated function + down_proj( + mlp_tile_inline149__rv_v2, w_down__ssa_v0, down_acc_all_inline168__iter_v6, k0_inline113__ssa_v8, + layer_inter_base_inline107__ssa_v0, n0_inline122__ssa_v8 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj.cpp new file mode 100644 index 0000000000..de4335aeac --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void gate_proj( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__rv_v2 + __gm__ Tensor *gate_acc_all_inline203__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__rv_v2 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__rv_v2_tensor->buffer.addr) + + gate_acc_all_inline203__rv_v2_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v0 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v0_conv; + gu_k0_inline131__ssa_v0_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v0 = gu_k0_inline131__ssa_v0_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + gate_proj( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__rv_v2, gu_k0_inline131__ssa_v0, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp new file mode 100644 index 0000000000..e05dc60546 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj_0 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void gate_proj_0( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__ssa_v4 + __gm__ Tensor *gate_acc_all_inline203__ssa_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__ssa_v4 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__ssa_v4_tensor->buffer.addr) + + gate_acc_all_inline203__ssa_v4_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v1 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v1_conv; + gu_k0_inline131__ssa_v1_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v1 = gu_k0_inline131__ssa_v1_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + gate_proj_0( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__ssa_v4, gu_k0_inline131__ssa_v1, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp new file mode 100644 index 0000000000..a5688a2059 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj_1 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void gate_proj_1( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__ssa_v5 + __gm__ Tensor *gate_acc_all_inline203__ssa_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__ssa_v5 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__ssa_v5_tensor->buffer.addr) + + gate_acc_all_inline203__ssa_v5_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v2 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v2_conv; + gu_k0_inline131__ssa_v2_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v2 = gu_k0_inline131__ssa_v2_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + gate_proj_1( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__ssa_v5, gu_k0_inline131__ssa_v2, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp new file mode 100644 index 0000000000..1b7ecba2cc --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj_2 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void gate_proj_2( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__ssa_v6 + __gm__ Tensor *gate_acc_all_inline203__ssa_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__ssa_v6 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__ssa_v6_tensor->buffer.addr) + + gate_acc_all_inline203__ssa_v6_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v3 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v3_conv; + gu_k0_inline131__ssa_v3_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v3 = gu_k0_inline131__ssa_v3_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + gate_proj_2( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__ssa_v6, gu_k0_inline131__ssa_v3, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp new file mode 100644 index 0000000000..4012cd6415 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj_3 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void gate_proj_3( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__ssa_v7 + __gm__ Tensor *gate_acc_all_inline203__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__ssa_v7 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__ssa_v7_tensor->buffer.addr) + + gate_acc_all_inline203__ssa_v7_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v4 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v4_conv; + gu_k0_inline131__ssa_v4_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v4 = gu_k0_inline131__ssa_v4_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + gate_proj_3( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__ssa_v7, gu_k0_inline131__ssa_v4, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp new file mode 100644 index 0000000000..caf3aa32c0 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp @@ -0,0 +1,622 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: gate_proj_4 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +gate_proj_4(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int64_t v6) { + const int64_t v7 = 960; + const int64_t v8 = 2; + const int64_t v9 = 15; + const int64_t v10 = 32; + const int64_t v11 = 64; + const int64_t v12 = 17408; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 512; + const int64_t v17 = 2048; + const int64_t v18 = 32768; + const int64_t v19 = 1536; + const int64_t v20 = 1024; + const int64_t v21 = 0; + const int64_t v22 = 135168; + const int64_t v23 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v24 = (size_t)v11; + size_t v25 = (size_t)v21; + size_t v26 = (size_t)v10; + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v27 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v28 = (uint64_t)v23; + TASSIGN(v27, v28); + pto::Shape<1, 1, 1, 16, 64> v29 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v31 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + v4 * v13), v29, v30 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v27, v31); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v32 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v33 = (uint64_t)v22; + TASSIGN(v32, v33); + int64_t v34 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v35 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v36 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + v34 * v12 + v6 * v13), v35, v36); + TLOAD(v32, v37); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v38 = v25; v38 < v24; v38 += v26) { + int64_t v39 = (int64_t)v38; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v41 = (uint64_t)v20; + TASSIGN(v40, v41); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v40, v27, v21, v38); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + TEXTRACT(v42, v32, v38, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v45 = (uint64_t)v19; + TASSIGN(v44, v45); + int64_t v46 = (int64_t)((uint64_t)v39 + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v44, v27, v21, v46); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v48 = (uint64_t)v18; + TASSIGN(v47, v48); + TEXTRACT(v47, v32, v46, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v39 == v21) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + pipe_barrier(PIPE_M); + TMATMUL(v49, v40, v42); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v52 = (uint64_t)v21; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v51, v51, v40, v42); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v54 = (uint64_t)v21; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v53, v53, v44, v47); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v55 = (size_t)v13; v55 < ((size_t)v9); v55 += (size_t)v8) { + int64_t v56 = (int64_t)((uint64_t)((int64_t)v55) * (uint64_t)v11); + int64_t v57 = (int64_t)((uint64_t)v56 + (uint64_t)v11); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v59 = (uint64_t)v21; + TASSIGN(v58, v59); + pto::Shape<1, 1, 1, 16, 64> v60 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v61 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v62 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v13), v60, v61 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v58, v62); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v64 = (uint64_t)v22; + TASSIGN(v63, v64); + pto::Shape<1, 1, 1, 64, 1024> v65 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v66 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v56) * v12 + v6 * v13), v65, v66); + TLOAD(v63, v67); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v69 = (uint64_t)v17; + TASSIGN(v68, v69); + pto::Shape<1, 1, 1, 16, 64> v70 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v71 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v72 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v57) * v13), v70, v71 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v68, v72); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v74 = (uint64_t)v23; + TASSIGN(v73, v74); + pto::Shape<1, 1, 1, 64, 1024> v75 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v76 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v77 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v57) * v12 + v6 * v13), v75, v76); + TLOAD(v73, v77); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v78 = v25; v78 < v24; v78 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v79 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v80 = (uint64_t)v20; + TASSIGN(v79, v80); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v79, v58, v21, v78); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v81 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + TEXTRACT(v81, v63, v78, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v84 = (uint64_t)v19; + TASSIGN(v83, v84); + int64_t v85 = (int64_t)((uint64_t)((int64_t)v78) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v83, v58, v21, v85); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v86 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v87 = (uint64_t)v18; + TASSIGN(v86, v87); + TEXTRACT(v86, v63, v85, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v88 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v89 = (uint64_t)v21; + TASSIGN(v88, v89); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v88, v88, v79, v81); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v91 = (uint64_t)v21; + TASSIGN(v90, v91); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v90, v90, v83, v86); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v92 = v25; v92 < v24; v92 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v93 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v94 = (uint64_t)v21; + TASSIGN(v93, v94); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v93, v68, v21, v92); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v95 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v96 = (uint64_t)v21; + TASSIGN(v95, v96); + TEXTRACT(v95, v73, v92, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v98 = (uint64_t)v16; + TASSIGN(v97, v98); + int64_t v99 = (int64_t)((uint64_t)((int64_t)v92) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v97, v68, v21, v99); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v100 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v101 = (uint64_t)v18; + TASSIGN(v100, v101); + TEXTRACT(v100, v73, v99, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v102 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v103 = (uint64_t)v21; + TASSIGN(v102, v103); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v102, v102, v93, v95); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v105 = (uint64_t)v21; + TASSIGN(v104, v105); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v104, v104, v97, v100); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + pto::Shape<1, 1, 1, 16, 64> v108 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v109 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v110 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v108, v109 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v106, v110); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v112 = (uint64_t)v22; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 64, 1024> v113 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v114 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v7) * v12 + v6 * v13), v113, v114); + TLOAD(v111, v115); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v116 = v25; v116 < v24; v116 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v117 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v118 = (uint64_t)v20; + TASSIGN(v117, v118); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v117, v106, v21, v116); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + TEXTRACT(v119, v111, v116, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v122 = (uint64_t)v19; + TASSIGN(v121, v122); + int64_t v123 = (int64_t)((uint64_t)((int64_t)v116) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v121, v106, v21, v123); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v125 = (uint64_t)v18; + TASSIGN(v124, v125); + TEXTRACT(v124, v111, v123, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v126, v126, v117, v119); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v129 = (uint64_t)v21; + TASSIGN(v128, v129); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v128, v128, v121, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v131 = (uint64_t)v21; + TASSIGN(v130, v131); + pto::Shape<1, 1, 1, 16, 1024> v132 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v133 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v134 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v21 + v21 * v12 + v6 * v13), v132, v133 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v134, v130); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_gate__ssa_v0 + __gm__ Tensor *w_gate__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_gate__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_gate__ssa_v0_tensor->buffer.addr) + w_gate__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__iter_v11 + __gm__ Tensor *gate_acc_all_inline203__iter_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__iter_v11 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__iter_v11_tensor->buffer.addr) + + gate_acc_all_inline203__iter_v11_tensor->start_offset; + + // Unpack scalar: k0_inline113__ssa_v7 + union { + uint64_t u64; + int64_t val; + } k0_inline113__ssa_v7_conv; + k0_inline113__ssa_v7_conv.u64 = args[3]; + int64_t k0_inline113__ssa_v7 = k0_inline113__ssa_v7_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Unpack scalar: n0_inline122__ssa_v6 + union { + uint64_t u64; + int64_t val; + } n0_inline122__ssa_v6_conv; + n0_inline122__ssa_v6_conv.u64 = args[5]; + int64_t n0_inline122__ssa_v6 = n0_inline122__ssa_v6_conv.val; + + // Forward to ptoas-generated function + gate_proj_4( + mlp_norm_in_inline71__rv_v14, w_gate__ssa_v0, gate_acc_all_inline203__iter_v11, k0_inline113__ssa_v7, + layer_hidden_base_inline151__ssa_v0, n0_inline122__ssa_v6 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/k_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/k_proj.cpp new file mode 100644 index 0000000000..79666ceefb --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/k_proj.cpp @@ -0,0 +1,617 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: k_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +k_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v4, int32_t v5, int32_t v6) { + const int64_t v7 = 768; + const int64_t v8 = 64; + const int64_t v9 = 128; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 512; + const int64_t v13 = 5; + const int64_t v14 = 5120; + const int64_t v15 = 1; + const int64_t v16 = 1024; + const int64_t v17 = 16; + const int64_t v18 = 2048; + const int64_t v19 = 8192; + const int64_t v20 = 32768; + const int64_t v21 = 6144; + const int64_t v22 = 4096; + const int64_t v23 = 0; + const int64_t v24 = 147456; + const int64_t v25 = 16384; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v26 = (size_t)v23; + size_t v27 = (size_t)v10; + size_t v28 = (size_t)v9; + int64_t v29 = (int64_t)v5; + int64_t v30 = (int64_t)((uint64_t)(v29 % v13) * (uint64_t)v16); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + for (size_t v31 = v26; v31 < ((size_t)v11); v31 += (size_t)v15) { + int64_t v32 = (int64_t)((uint64_t)((int64_t)(uint64_t)(v29 / v13) * (uint64_t)v12) + + (uint64_t)((int64_t)(uint64_t)((int64_t)v31) * (uint64_t)v10)); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v34 = (uint64_t)v25; + TASSIGN(v33, v34); + pto::Shape<1, 1, 1, 16, 256> v35 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v36 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v30 * v15), v35, v36 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + TLOAD(v33, v37); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v39 = (uint64_t)v24; + TASSIGN(v38, v39); + pto::Shape<1, 1, 1, 256, 256> v40 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v41 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v42 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v30) * v16 + v32 * v15), v40, v41); + TLOAD(v38, v42); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v43 = v26; v43 < v27; v43 += v28) { + int64_t v44 = (int64_t)v43; + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v45 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v46 = (uint64_t)v22; + TASSIGN(v45, v46); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v45, v33, v23, v43); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v48 = (uint64_t)v23; + TASSIGN(v47, v48); + TEXTRACT(v47, v38, v43, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + int64_t v51 = (int64_t)((uint64_t)v44 + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v49, v33, v23, v51); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v52 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v53 = (uint64_t)v20; + TASSIGN(v52, v53); + TEXTRACT(v52, v38, v51, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v44 == v23) { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v54 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); + pipe_barrier(PIPE_M); + TMATMUL(v54, v45, v47); + } else { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v56, v56, v45, v47); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v59 = (uint64_t)v23; + TASSIGN(v58, v59); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v58, v58, v49, v52); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + int64_t v60 = (int64_t)((uint64_t)v30 + (uint64_t)v10); + int64_t v61 = (int64_t)((uint64_t)v30 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v62 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v63 = (uint64_t)v23; + TASSIGN(v62, v63); + pto::Shape<1, 1, 1, 16, 256> v64 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v65 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v66 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v60 * v15), v64, v65 + ); + TLOAD(v62, v66); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v68 = (uint64_t)v24; + TASSIGN(v67, v68); + pto::Shape<1, 1, 1, 256, 256> v69 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v70 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v71 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v60) * v16 + v32 * v15), v69, v70); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v67, v71); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v73 = (uint64_t)v19; + TASSIGN(v72, v73); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v61 * v15), v74, v75 + ); + TLOAD(v72, v76); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v78 = (uint64_t)v25; + TASSIGN(v77, v78); + pto::Shape<1, 1, 1, 256, 256> v79 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v80 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v81 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v61) * v16 + v32 * v15), v79, v80); + TLOAD(v77, v81); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v82 = v26; v82 < v27; v82 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v83, v62, v23, v82); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v86 = (uint64_t)v23; + TASSIGN(v85, v86); + TEXTRACT(v85, v67, v82, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v87 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v88 = (uint64_t)v21; + TASSIGN(v87, v88); + int64_t v89 = (int64_t)((uint64_t)((int64_t)v82) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v87, v62, v23, v89); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v91 = (uint64_t)v20; + TASSIGN(v90, v91); + TEXTRACT(v90, v67, v89, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v93 = (uint64_t)v23; + TASSIGN(v92, v93); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v92, v92, v83, v85); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v94 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v95 = (uint64_t)v23; + TASSIGN(v94, v95); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v94, v94, v87, v90); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v96 = v26; v96 < v27; v96 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v98 = (uint64_t)v23; + TASSIGN(v97, v98); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v97, v72, v23, v96); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v100 = (uint64_t)v23; + TASSIGN(v99, v100); + TEXTRACT(v99, v77, v96, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v101 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v102 = (uint64_t)v18; + TASSIGN(v101, v102); + int64_t v103 = (int64_t)((uint64_t)((int64_t)v96) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v101, v72, v23, v103); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v105 = (uint64_t)v20; + TASSIGN(v104, v105); + TEXTRACT(v104, v77, v103, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v106, v106, v97, v99); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v109 = (uint64_t)v23; + TASSIGN(v108, v109); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v108, v108, v101, v104); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + int64_t v110 = (int64_t)((uint64_t)v30 + (uint64_t)v7); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v112 = (uint64_t)v25; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 16, 256> v113 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v114 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v110 * v15), v113, v114 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v111, v115); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v116 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v117 = (uint64_t)v24; + TASSIGN(v116, v117); + pto::Shape<1, 1, 1, 256, 256> v118 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v119 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v120 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v110) * v16 + v32 * v15), v118, v119); + TLOAD(v116, v120); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v121 = v26; v121 < v27; v121 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v122 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v123 = (uint64_t)v22; + TASSIGN(v122, v123); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v122, v111, v23, v121); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v125 = (uint64_t)v23; + TASSIGN(v124, v125); + TEXTRACT(v124, v116, v121, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + int64_t v128 = (int64_t)((uint64_t)((int64_t)v121) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v126, v111, v23, v128); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v129 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v130 = (uint64_t)v20; + TASSIGN(v129, v130); + TEXTRACT(v129, v116, v128, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v131 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v132 = (uint64_t)v23; + TASSIGN(v131, v132); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v131, v131, v122, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v133 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v134 = (uint64_t)v23; + TASSIGN(v133, v134); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v133, v133, v126, v129); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v135 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v136 = (uint64_t)v23; + TASSIGN(v135, v136); + pto::Shape<1, 1, 1, 16, 256> v137 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<16384, 16384, 16384, 1024, 1> v138 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> + v139 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + pipe_barrier(PIPE_FIX); + TSTORE< + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v139, v135); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: k_proj_inline135__ssa_v1 + __gm__ Tensor *k_proj_inline135__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *k_proj_inline135__ssa_v1 = + reinterpret_cast<__gm__ float *>(k_proj_inline135__ssa_v1_tensor->buffer.addr) + + k_proj_inline135__ssa_v1_tensor->start_offset; + + // Unpack tensor: normed__iter_v4 + __gm__ Tensor *normed__iter_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *normed__iter_v4 = reinterpret_cast<__gm__ bfloat16_t *>(normed__iter_v4_tensor->buffer.addr) + + normed__iter_v4_tensor->start_offset; + + // Unpack tensor: wk__ssa_v0 + __gm__ Tensor *wk__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ bfloat16_t *wk__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(wk__ssa_v0_tensor->buffer.addr) + wk__ssa_v0_tensor->start_offset; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[3]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + k_proj( + k_proj_inline135__ssa_v1, normed__iter_v4, wk__ssa_v0, layer_hidden_base_inline151__ssa_v0, + __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj.cpp new file mode 100644 index 0000000000..1d35867e0e --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj.cpp @@ -0,0 +1,577 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: out_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int64_t v6) { + const int64_t v7 = 960; + const int64_t v8 = 2; + const int64_t v9 = 15; + const int64_t v10 = 32; + const int64_t v11 = 512; + const int64_t v12 = 64; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 1024; + const int64_t v17 = 32768; + const int64_t v18 = 3072; + const int64_t v19 = 2048; + const int64_t v20 = 0; + const int64_t v21 = 69632; + const int64_t v22 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v23 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v12); + uint64_t v24 = (uint64_t)v22; + TASSIGN(v23, v24); + pto::Shape<1, 1, 1, 16, 64> v25 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v26 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v27 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v20 + v20 * v14 + v4 * v13), v25, v26 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TLOAD(v23, v27); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v28 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v11); + uint64_t v29 = (uint64_t)v21; + TASSIGN(v28, v29); + int64_t v30 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 512> v31 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v32 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v20 + v30 * v14 + v6 * v13), v31, v32 + ); + TLOAD(v28, v33); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v35 = (uint64_t)v19; + TASSIGN(v34, v35); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v34, v23, v20, v20); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v37 = (uint64_t)v20; + TASSIGN(v36, v37); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v36, v28, v20, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v39 = (uint64_t)v18; + TASSIGN(v38, v39); + TEXTRACT(v38, v23, v20, v10); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v41 = (uint64_t)v17; + TASSIGN(v40, v41); + TEXTRACT(v40, v28, v10, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v43 = (uint64_t)v20; + TASSIGN(v42, v43); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + TMATMUL(v42, v34, v36); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v45 = (uint64_t)v20; + TASSIGN(v44, v45); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v44, v44, v38, v40); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + for (size_t v46 = (size_t)v13; v46 < ((size_t)v9); v46 += (size_t)v8) { + int64_t v47 = (int64_t)((uint64_t)((int64_t)v46) * (uint64_t)v12); + int64_t v48 = (int64_t)((uint64_t)v47 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v12); + uint64_t v50 = (uint64_t)v20; + TASSIGN(v49, v50); + pto::Shape<1, 1, 1, 16, 64> v51 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v52 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v53 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v47) * v13), v51, v52 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v49, v53); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v54 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v11); + uint64_t v55 = (uint64_t)v21; + TASSIGN(v54, v55); + pto::Shape<1, 1, 1, 64, 512> v56 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v57 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v58 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v30 + (uint64_t)v47) * v14 + v6 * v13), v56, v57); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v54, v58); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v59 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v12); + uint64_t v60 = (uint64_t)v19; + TASSIGN(v59, v60); + pto::Shape<1, 1, 1, 16, 64> v61 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v62 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v63 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v48) * v13), v61, v62 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v59, v63); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v64 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v11); + uint64_t v65 = (uint64_t)v22; + TASSIGN(v64, v65); + pto::Shape<1, 1, 1, 64, 512> v66 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v67 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v68 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v30 + (uint64_t)v48) * v14 + v6 * v13), v66, v67); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + TLOAD(v64, v68); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v69 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v70 = (uint64_t)v19; + TASSIGN(v69, v70); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v69, v49, v20, v20); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v71 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v72 = (uint64_t)v20; + TASSIGN(v71, v72); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v71, v54, v20, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v74 = (uint64_t)v18; + TASSIGN(v73, v74); + TEXTRACT(v73, v49, v20, v10); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v76 = (uint64_t)v17; + TASSIGN(v75, v76); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v75, v54, v10, v20); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v78 = (uint64_t)v20; + TASSIGN(v77, v78); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v77, v77, v69, v71); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v79 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v80 = (uint64_t)v20; + TASSIGN(v79, v80); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v79, v79, v73, v75); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v82 = (uint64_t)v20; + TASSIGN(v81, v82); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v59, v20, v20); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v84 = (uint64_t)v20; + TASSIGN(v83, v84); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v83, v64, v20, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v86 = (uint64_t)v16; + TASSIGN(v85, v86); + TEXTRACT(v85, v59, v20, v10); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v87 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v88 = (uint64_t)v17; + TASSIGN(v87, v88); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v87, v64, v10, v20); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v89 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v90 = (uint64_t)v20; + TASSIGN(v89, v90); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v89, v89, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v91 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v92 = (uint64_t)v20; + TASSIGN(v91, v92); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v91, v91, v85, v87); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + } + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v93 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v12); + uint64_t v94 = (uint64_t)v22; + TASSIGN(v93, v94); + pto::Shape<1, 1, 1, 16, 64> v95 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v96 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v97 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v95, v96 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5); + TLOAD(v93, v97); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v98 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v11); + uint64_t v99 = (uint64_t)v21; + TASSIGN(v98, v99); + pto::Shape<1, 1, 1, 64, 512> v100 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v101 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v102 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v20 + (int64_t)((uint64_t)v30 + (uint64_t)v7) * v14 + v6 * v13), v100, v101 + ); + TLOAD(v98, v102); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v103 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v104 = (uint64_t)v19; + TASSIGN(v103, v104); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6); + TEXTRACT(v103, v93, v20, v20); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v105 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v106 = (uint64_t)v20; + TASSIGN(v105, v106); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v105, v98, v20, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v107 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v108 = (uint64_t)v18; + TASSIGN(v107, v108); + TEXTRACT(v107, v93, v20, v10); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v109 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v11); + uint64_t v110 = (uint64_t)v17; + TASSIGN(v109, v110); + TEXTRACT(v109, v98, v10, v20); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v112 = (uint64_t)v20; + TASSIGN(v111, v112); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v111, v111, v103, v105); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v114 = (uint64_t)v20; + TASSIGN(v113, v114); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v113, v113, v107, v109); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v115 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v116 = (uint64_t)v20; + TASSIGN(v115, v116); + pto::Shape<1, 1, 1, 16, 512> v117 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v118 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v119 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v20 + v20 * v14 + v6 * v13), v117, v118 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v119, v115); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: attn_out_inline282__ssa_v4 + __gm__ Tensor *attn_out_inline282__ssa_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *attn_out_inline282__ssa_v4 = + reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline282__ssa_v4_tensor->buffer.addr) + + attn_out_inline282__ssa_v4_tensor->start_offset; + + // Unpack tensor: wo__ssa_v0 + __gm__ Tensor *wo__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *wo__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(wo__ssa_v0_tensor->buffer.addr) + wo__ssa_v0_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__iter_v4 + __gm__ Tensor *attn_proj_fp32_inline220__iter_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__iter_v4 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__iter_v4_tensor->buffer.addr) + + attn_proj_fp32_inline220__iter_v4_tensor->start_offset; + + // Unpack scalar: k_op_inline266__ssa_v0 + union { + uint64_t u64; + int64_t val; + } k_op_inline266__ssa_v0_conv; + k_op_inline266__ssa_v0_conv.u64 = args[3]; + int64_t k_op_inline266__ssa_v0 = k_op_inline266__ssa_v0_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Unpack scalar: n_op_inline64__ssa_v0 + union { + uint64_t u64; + int64_t val; + } n_op_inline64__ssa_v0_conv; + n_op_inline64__ssa_v0_conv.u64 = args[5]; + int64_t n_op_inline64__ssa_v0 = n_op_inline64__ssa_v0_conv.val; + + // Forward to ptoas-generated function + out_proj( + attn_out_inline282__ssa_v4, wo__ssa_v0, attn_proj_fp32_inline220__iter_v4, k_op_inline266__ssa_v0, + layer_hidden_base_inline151__ssa_v0, n_op_inline64__ssa_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj_0.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj_0.cpp new file mode 100644 index 0000000000..50b5cc3582 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/out_proj_0.cpp @@ -0,0 +1,579 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: out_proj_0 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void out_proj_0( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 512; + const int64_t v14 = 5; + const int64_t v15 = 1; + const int64_t v16 = 5120; + const int64_t v17 = 16; + const int64_t v18 = 1024; + const int64_t v19 = 32768; + const int64_t v20 = 3072; + const int64_t v21 = 2048; + const int64_t v22 = 0; + const int64_t v23 = 69632; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + int64_t v25 = (int64_t)((uint64_t)v4 + (uint64_t)((int64_t)v6)); + int64_t v26 = (int64_t)((uint64_t)(v25 / v14) * (uint64_t)v13); + int64_t v27 = (int64_t)((uint64_t)(v25 % v14) * (uint64_t)v18); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v28 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v12); + uint64_t v29 = (uint64_t)v24; + TASSIGN(v28, v29); + pto::Shape<1, 1, 1, 16, 64> v30 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v31 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v32 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v16 + v27 * v15), v30, v31 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TLOAD(v28, v32); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v13); + uint64_t v34 = (uint64_t)v23; + TASSIGN(v33, v34); + int64_t v35 = (int64_t)((uint64_t)v5 + (uint64_t)v27); + pto::Shape<1, 1, 1, 64, 512> v36 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v37 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v38 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v22 + v35 * v16 + v26 * v15), v36, v37 + ); + TLOAD(v33, v38); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v39 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v40 = (uint64_t)v21; + TASSIGN(v39, v40); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v39, v28, v22, v22); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v42 = (uint64_t)v22; + TASSIGN(v41, v42); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v41, v33, v22, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v43 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v44 = (uint64_t)v20; + TASSIGN(v43, v44); + TEXTRACT(v43, v28, v22, v11); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v45 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v46 = (uint64_t)v19; + TASSIGN(v45, v46); + TEXTRACT(v45, v33, v11, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v48 = (uint64_t)v22; + TASSIGN(v47, v48); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + TMATMUL(v47, v39, v41); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v50 = (uint64_t)v22; + TASSIGN(v49, v50); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v49, v49, v43, v45); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + for (size_t v51 = (size_t)v15; v51 < ((size_t)v10); v51 += (size_t)v9) { + int64_t v52 = (int64_t)((uint64_t)((int64_t)v51) * (uint64_t)v12); + int64_t v53 = (int64_t)((uint64_t)v52 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v54 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v12); + uint64_t v55 = (uint64_t)v22; + TASSIGN(v54, v55); + pto::Shape<1, 1, 1, 16, 64> v56 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v57 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v58 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v16 + (int64_t)((uint64_t)v27 + (uint64_t)v52) * v15), v56, v57 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v54, v58); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v59 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v13); + uint64_t v60 = (uint64_t)v23; + TASSIGN(v59, v60); + pto::Shape<1, 1, 1, 64, 512> v61 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v62 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v63 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v35 + (uint64_t)v52) * v16 + v26 * v15), v61, v62); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v59, v63); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v64 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v12); + uint64_t v65 = (uint64_t)v21; + TASSIGN(v64, v65); + pto::Shape<1, 1, 1, 16, 64> v66 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v67 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v68 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v16 + (int64_t)((uint64_t)v27 + (uint64_t)v53) * v15), v66, v67 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v64, v68); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v69 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v13); + uint64_t v70 = (uint64_t)v24; + TASSIGN(v69, v70); + pto::Shape<1, 1, 1, 64, 512> v71 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v72 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v73 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v35 + (uint64_t)v53) * v16 + v26 * v15), v71, v72); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + TLOAD(v69, v73); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v74 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v75 = (uint64_t)v21; + TASSIGN(v74, v75); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v74, v54, v22, v22); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v76 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v77 = (uint64_t)v22; + TASSIGN(v76, v77); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v76, v59, v22, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v78 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v79 = (uint64_t)v20; + TASSIGN(v78, v79); + TEXTRACT(v78, v54, v22, v11); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v80 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v81 = (uint64_t)v19; + TASSIGN(v80, v81); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v80, v59, v11, v22); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v82 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v83 = (uint64_t)v22; + TASSIGN(v82, v83); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v82, v82, v74, v76); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v84 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v85 = (uint64_t)v22; + TASSIGN(v84, v85); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v84, v84, v78, v80); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v86 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v87 = (uint64_t)v22; + TASSIGN(v86, v87); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v86, v64, v22, v22); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v89 = (uint64_t)v22; + TASSIGN(v88, v89); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v88, v69, v22, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v91 = (uint64_t)v18; + TASSIGN(v90, v91); + TEXTRACT(v90, v64, v22, v11); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v93 = (uint64_t)v19; + TASSIGN(v92, v93); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v92, v69, v11, v22); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v94 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v95 = (uint64_t)v22; + TASSIGN(v94, v95); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v94, v94, v86, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v96 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v97 = (uint64_t)v22; + TASSIGN(v96, v97); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v96, v96, v90, v92); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + } + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v98 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v12); + uint64_t v99 = (uint64_t)v24; + TASSIGN(v98, v99); + pto::Shape<1, 1, 1, 16, 64> v100 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v101 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v102 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v16 + (int64_t)((uint64_t)v27 + (uint64_t)v8) * v15), v100, v101 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5); + TLOAD(v98, v102); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6); + Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v103 = Tile< + TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v13); + uint64_t v104 = (uint64_t)v23; + TASSIGN(v103, v104); + pto::Shape<1, 1, 1, 64, 512> v105 = pto::Shape<1, 1, 1, 64, 512>(); + pto::Stride<327680, 327680, 327680, 5120, 1> v106 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND> + v107 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( + v2 + (v22 + (int64_t)((uint64_t)v35 + (uint64_t)v8) * v16 + v26 * v15), v105, v106 + ); + TLOAD(v103, v107); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v109 = (uint64_t)v21; + TASSIGN(v108, v109); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6); + TEXTRACT(v108, v98, v22, v22); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v110 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v111 = (uint64_t)v22; + TASSIGN(v110, v111); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v110, v103, v22, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v112 = Tile< + TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v11); + uint64_t v113 = (uint64_t)v20; + TASSIGN(v112, v113); + TEXTRACT(v112, v98, v22, v11); + Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v114 = Tile< + TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v115 = (uint64_t)v19; + TASSIGN(v114, v115); + TEXTRACT(v114, v103, v11, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v116 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v117 = (uint64_t)v22; + TASSIGN(v116, v117); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v116, v116, v108, v110); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v118 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v119 = (uint64_t)v22; + TASSIGN(v118, v119); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v118, v118, v112, v114); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v120 = Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v13); + uint64_t v121 = (uint64_t)v22; + TASSIGN(v120, v121); + pto::Shape<1, 1, 1, 16, 512> v122 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v123 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v124 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v16 + v26 * v15), v122, v123 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v124, v120); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: attn_out_inline282__ssa_v4 + __gm__ Tensor *attn_out_inline282__ssa_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *attn_out_inline282__ssa_v4 = + reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline282__ssa_v4_tensor->buffer.addr) + + attn_out_inline282__ssa_v4_tensor->start_offset; + + // Unpack tensor: wo__ssa_v0 + __gm__ Tensor *wo__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *wo__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(wo__ssa_v0_tensor->buffer.addr) + wo__ssa_v0_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__rv_v5 + __gm__ Tensor *attn_proj_fp32_inline220__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__rv_v5 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__rv_v5_tensor->buffer.addr) + + attn_proj_fp32_inline220__rv_v5_tensor->start_offset; + + // Unpack scalar: N_OUT_DIRECT_inline61__ssa_v0 + union { + uint64_t u64; + int64_t val; + } N_OUT_DIRECT_inline61__ssa_v0_conv; + N_OUT_DIRECT_inline61__ssa_v0_conv.u64 = args[3]; + int64_t N_OUT_DIRECT_inline61__ssa_v0 = N_OUT_DIRECT_inline61__ssa_v0_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + out_proj_0( + attn_out_inline282__ssa_v4, wo__ssa_v0, attn_proj_fp32_inline220__rv_v5, N_OUT_DIRECT_inline61__ssa_v0, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/q_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/q_proj.cpp new file mode 100644 index 0000000000..cc978d8d64 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/q_proj.cpp @@ -0,0 +1,617 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: q_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v4, int32_t v5, int32_t v6) { + const int64_t v7 = 768; + const int64_t v8 = 64; + const int64_t v9 = 128; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 1024; + const int64_t v13 = 512; + const int64_t v14 = 5; + const int64_t v15 = 1; + const int64_t v16 = 5120; + const int64_t v17 = 16; + const int64_t v18 = 2048; + const int64_t v19 = 8192; + const int64_t v20 = 32768; + const int64_t v21 = 6144; + const int64_t v22 = 4096; + const int64_t v23 = 0; + const int64_t v24 = 147456; + const int64_t v25 = 16384; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v26 = (size_t)v23; + size_t v27 = (size_t)v10; + size_t v28 = (size_t)v9; + int64_t v29 = (int64_t)v5; + int64_t v30 = (int64_t)((uint64_t)(v29 % v14) * (uint64_t)v12); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + for (size_t v31 = v26; v31 < ((size_t)v11); v31 += (size_t)v15) { + int64_t v32 = (int64_t)((uint64_t)((int64_t)(uint64_t)(v29 / v14) * (uint64_t)v13) + + (uint64_t)((int64_t)(uint64_t)((int64_t)v31) * (uint64_t)v10)); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v34 = (uint64_t)v25; + TASSIGN(v33, v34); + pto::Shape<1, 1, 1, 16, 256> v35 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v36 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v16 + v30 * v15), v35, v36 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + TLOAD(v33, v37); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v39 = (uint64_t)v24; + TASSIGN(v38, v39); + pto::Shape<1, 1, 1, 256, 256> v40 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<1310720, 1310720, 1310720, 5120, 1> v41 = pto::Stride<1310720, 1310720, 1310720, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, pto::Layout::ND> + v42 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v30) * v16 + v32 * v15), v40, v41); + TLOAD(v38, v42); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v43 = v26; v43 < v27; v43 += v28) { + int64_t v44 = (int64_t)v43; + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v45 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v46 = (uint64_t)v22; + TASSIGN(v45, v46); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v45, v33, v23, v43); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v48 = (uint64_t)v23; + TASSIGN(v47, v48); + TEXTRACT(v47, v38, v43, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + int64_t v51 = (int64_t)((uint64_t)v44 + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v49, v33, v23, v51); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v52 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v53 = (uint64_t)v20; + TASSIGN(v52, v53); + TEXTRACT(v52, v38, v51, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v44 == v23) { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v54 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); + pipe_barrier(PIPE_M); + TMATMUL(v54, v45, v47); + } else { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v56, v56, v45, v47); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v59 = (uint64_t)v23; + TASSIGN(v58, v59); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v58, v58, v49, v52); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + int64_t v60 = (int64_t)((uint64_t)v30 + (uint64_t)v10); + int64_t v61 = (int64_t)((uint64_t)v30 + (uint64_t)v13); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v62 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v63 = (uint64_t)v23; + TASSIGN(v62, v63); + pto::Shape<1, 1, 1, 16, 256> v64 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v65 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v66 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v16 + v60 * v15), v64, v65 + ); + TLOAD(v62, v66); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v68 = (uint64_t)v24; + TASSIGN(v67, v68); + pto::Shape<1, 1, 1, 256, 256> v69 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<1310720, 1310720, 1310720, 5120, 1> v70 = pto::Stride<1310720, 1310720, 1310720, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, pto::Layout::ND> + v71 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v60) * v16 + v32 * v15), v69, v70); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v67, v71); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v73 = (uint64_t)v19; + TASSIGN(v72, v73); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v16 + v61 * v15), v74, v75 + ); + TLOAD(v72, v76); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v78 = (uint64_t)v25; + TASSIGN(v77, v78); + pto::Shape<1, 1, 1, 256, 256> v79 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<1310720, 1310720, 1310720, 5120, 1> v80 = pto::Stride<1310720, 1310720, 1310720, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, pto::Layout::ND> + v81 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v61) * v16 + v32 * v15), v79, v80); + TLOAD(v77, v81); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v82 = v26; v82 < v27; v82 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v83, v62, v23, v82); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v86 = (uint64_t)v23; + TASSIGN(v85, v86); + TEXTRACT(v85, v67, v82, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v87 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v88 = (uint64_t)v21; + TASSIGN(v87, v88); + int64_t v89 = (int64_t)((uint64_t)((int64_t)v82) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v87, v62, v23, v89); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v91 = (uint64_t)v20; + TASSIGN(v90, v91); + TEXTRACT(v90, v67, v89, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v93 = (uint64_t)v23; + TASSIGN(v92, v93); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v92, v92, v83, v85); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v94 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v95 = (uint64_t)v23; + TASSIGN(v94, v95); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v94, v94, v87, v90); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v96 = v26; v96 < v27; v96 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v98 = (uint64_t)v23; + TASSIGN(v97, v98); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v97, v72, v23, v96); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v100 = (uint64_t)v23; + TASSIGN(v99, v100); + TEXTRACT(v99, v77, v96, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v101 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v102 = (uint64_t)v18; + TASSIGN(v101, v102); + int64_t v103 = (int64_t)((uint64_t)((int64_t)v96) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v101, v72, v23, v103); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v105 = (uint64_t)v20; + TASSIGN(v104, v105); + TEXTRACT(v104, v77, v103, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v106, v106, v97, v99); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v109 = (uint64_t)v23; + TASSIGN(v108, v109); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v108, v108, v101, v104); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + int64_t v110 = (int64_t)((uint64_t)v30 + (uint64_t)v7); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v112 = (uint64_t)v25; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 16, 256> v113 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v114 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v16 + v110 * v15), v113, v114 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v111, v115); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v116 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v117 = (uint64_t)v24; + TASSIGN(v116, v117); + pto::Shape<1, 1, 1, 256, 256> v118 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<1310720, 1310720, 1310720, 5120, 1> v119 = pto::Stride<1310720, 1310720, 1310720, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, pto::Layout::ND> + v120 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v110) * v16 + v32 * v15), v118, v119); + TLOAD(v116, v120); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v121 = v26; v121 < v27; v121 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v122 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v123 = (uint64_t)v22; + TASSIGN(v122, v123); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v122, v111, v23, v121); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v125 = (uint64_t)v23; + TASSIGN(v124, v125); + TEXTRACT(v124, v116, v121, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + int64_t v128 = (int64_t)((uint64_t)((int64_t)v121) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v126, v111, v23, v128); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v129 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v130 = (uint64_t)v20; + TASSIGN(v129, v130); + TEXTRACT(v129, v116, v128, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v131 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v132 = (uint64_t)v23; + TASSIGN(v131, v132); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v131, v131, v122, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v133 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v134 = (uint64_t)v23; + TASSIGN(v133, v134); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v133, v133, v126, v129); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v135 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v136 = (uint64_t)v23; + TASSIGN(v135, v136); + pto::Shape<1, 1, 1, 16, 256> v137 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v138 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v139 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + pipe_barrier(PIPE_FIX); + TSTORE< + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v139, v135); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: q_proj_inline139__rv_v2 + __gm__ Tensor *q_proj_inline139__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *q_proj_inline139__rv_v2 = + reinterpret_cast<__gm__ float *>(q_proj_inline139__rv_v2_tensor->buffer.addr) + + q_proj_inline139__rv_v2_tensor->start_offset; + + // Unpack tensor: normed__iter_v4 + __gm__ Tensor *normed__iter_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *normed__iter_v4 = reinterpret_cast<__gm__ bfloat16_t *>(normed__iter_v4_tensor->buffer.addr) + + normed__iter_v4_tensor->start_offset; + + // Unpack tensor: wq__ssa_v0 + __gm__ Tensor *wq__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ bfloat16_t *wq__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(wq__ssa_v0_tensor->buffer.addr) + wq__ssa_v0_tensor->start_offset; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[3]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + q_proj( + q_proj_inline139__rv_v2, normed__iter_v4, wq__ssa_v0, layer_hidden_base_inline151__ssa_v0, + __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj.cpp new file mode 100644 index 0000000000..09f48d3335 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void up_proj( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__rv_v2 + __gm__ Tensor *up_acc_all_inline303__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__rv_v2 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__rv_v2_tensor->buffer.addr) + + up_acc_all_inline303__rv_v2_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v0 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v0_conv; + gu_k0_inline131__ssa_v0_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v0 = gu_k0_inline131__ssa_v0_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + up_proj( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__rv_v2, gu_k0_inline131__ssa_v0, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_0.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_0.cpp new file mode 100644 index 0000000000..01a120e8df --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_0.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj_0 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void up_proj_0( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__ssa_v4 + __gm__ Tensor *up_acc_all_inline303__ssa_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__ssa_v4 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__ssa_v4_tensor->buffer.addr) + + up_acc_all_inline303__ssa_v4_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v1 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v1_conv; + gu_k0_inline131__ssa_v1_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v1 = gu_k0_inline131__ssa_v1_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + up_proj_0( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__ssa_v4, gu_k0_inline131__ssa_v1, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_1.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_1.cpp new file mode 100644 index 0000000000..5238de9792 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_1.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj_1 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void up_proj_1( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__ssa_v5 + __gm__ Tensor *up_acc_all_inline303__ssa_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__ssa_v5 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__ssa_v5_tensor->buffer.addr) + + up_acc_all_inline303__ssa_v5_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v2 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v2_conv; + gu_k0_inline131__ssa_v2_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v2 = gu_k0_inline131__ssa_v2_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + up_proj_1( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__ssa_v5, gu_k0_inline131__ssa_v2, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_2.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_2.cpp new file mode 100644 index 0000000000..e7e043d882 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_2.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj_2 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void up_proj_2( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__ssa_v6 + __gm__ Tensor *up_acc_all_inline303__ssa_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__ssa_v6 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__ssa_v6_tensor->buffer.addr) + + up_acc_all_inline303__ssa_v6_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v3 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v3_conv; + gu_k0_inline131__ssa_v3_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v3 = gu_k0_inline131__ssa_v3_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + up_proj_2( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__ssa_v6, gu_k0_inline131__ssa_v3, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_3.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_3.cpp new file mode 100644 index 0000000000..8f85c9a488 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_3.cpp @@ -0,0 +1,621 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj_3 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void up_proj_3( + __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int32_t v6, int32_t v7 +) { + const int64_t v8 = 960; + const int64_t v9 = 2; + const int64_t v10 = 15; + const int64_t v11 = 32; + const int64_t v12 = 64; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 512; + const int64_t v18 = 2048; + const int64_t v19 = 32768; + const int64_t v20 = 1536; + const int64_t v21 = 1024; + const int64_t v22 = 0; + const int64_t v23 = 135168; + const int64_t v24 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v25 = (size_t)v12; + size_t v26 = (size_t)v22; + size_t v27 = (size_t)v11; + int64_t v28 = (int64_t)((uint64_t)((int64_t)v6) * (uint64_t)v21); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v30 = (uint64_t)v24; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 16, 64> v31 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + v4 * v14), v31, v32 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v29, v33); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v35 = (uint64_t)v23; + TASSIGN(v34, v35); + int64_t v36 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v37 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v38 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v39 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + v36 * v13 + v28 * v14), v37, v38); + TLOAD(v34, v39); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v40 = v26; v40 < v25; v40 += v27) { + int64_t v41 = (int64_t)v40; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v42, v29, v22, v40); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v45 = (uint64_t)v22; + TASSIGN(v44, v45); + TEXTRACT(v44, v34, v40, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v47 = (uint64_t)v20; + TASSIGN(v46, v47); + int64_t v48 = (int64_t)((uint64_t)v41 + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v46, v29, v22, v48); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v50 = (uint64_t)v19; + TASSIGN(v49, v50); + TEXTRACT(v49, v34, v48, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v41 == v22) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL(v51, v42, v44); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v53, v53, v42, v44); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v56 = (uint64_t)v22; + TASSIGN(v55, v56); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v55, v55, v46, v49); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v57 = (size_t)v14; v57 < ((size_t)v10); v57 += (size_t)v9) { + int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) * (uint64_t)v12); + int64_t v59 = (int64_t)((uint64_t)v58 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v61 = (uint64_t)v22; + TASSIGN(v60, v61); + pto::Shape<1, 1, 1, 16, 64> v62 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v60, v64); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v66 = (uint64_t)v23; + TASSIGN(v65, v66); + pto::Shape<1, 1, 1, 64, 1024> v67 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v68 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v69 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v58) * v13 + v28 * v14), v67, v68); + TLOAD(v65, v69); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v71 = (uint64_t)v18; + TASSIGN(v70, v71); + pto::Shape<1, 1, 1, 16, 64> v72 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v73 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v74 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v70, v74); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v76 = (uint64_t)v24; + TASSIGN(v75, v76); + pto::Shape<1, 1, 1, 64, 1024> v77 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v78 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v59) * v13 + v28 * v14), v77, v78); + TLOAD(v75, v79); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v80 = v26; v80 < v25; v80 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v81, v60, v22, v80); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v83 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + TEXTRACT(v83, v65, v80, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v86 = (uint64_t)v20; + TASSIGN(v85, v86); + int64_t v87 = (int64_t)((uint64_t)((int64_t)v80) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v85, v60, v22, v87); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v88 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v89 = (uint64_t)v19; + TASSIGN(v88, v89); + TEXTRACT(v88, v65, v87, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v91 = (uint64_t)v22; + TASSIGN(v90, v91); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v90, v90, v81, v83); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v93 = (uint64_t)v22; + TASSIGN(v92, v93); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v92, v92, v85, v88); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v94 = v26; v94 < v25; v94 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v95, v70, v22, v94); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v97 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v98 = (uint64_t)v22; + TASSIGN(v97, v98); + TEXTRACT(v97, v75, v94, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v16); + uint64_t v100 = (uint64_t)v17; + TASSIGN(v99, v100); + int64_t v101 = (int64_t)((uint64_t)((int64_t)v94) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v99, v70, v22, v101); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v102 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v103 = (uint64_t)v19; + TASSIGN(v102, v103); + TEXTRACT(v102, v75, v101, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v105 = (uint64_t)v22; + TASSIGN(v104, v105); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v104, v104, v95, v97); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v107 = (uint64_t)v22; + TASSIGN(v106, v107); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v106, v106, v99, v102); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v12); + uint64_t v109 = (uint64_t)v24; + TASSIGN(v108, v109); + pto::Shape<1, 1, 1, 16, 64> v110 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v111 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v112 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v108, v112); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v113 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v12, v21); + uint64_t v114 = (uint64_t)v23; + TASSIGN(v113, v114); + pto::Shape<1, 1, 1, 64, 1024> v115 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v116 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v117 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v22 + (int64_t)((uint64_t)v36 + (uint64_t)v8) * v13 + v28 * v14), v115, v116); + TLOAD(v113, v117); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v118 = v26; v118 < v25; v118 += v27) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v119, v108, v22, v118); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v122 = (uint64_t)v22; + TASSIGN(v121, v122); + TEXTRACT(v121, v113, v118, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v123 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v16, v16); + uint64_t v124 = (uint64_t)v20; + TASSIGN(v123, v124); + int64_t v125 = (int64_t)((uint64_t)((int64_t)v118) + (uint64_t)v16); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v123, v108, v22, v125); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v16, v21); + uint64_t v127 = (uint64_t)v19; + TASSIGN(v126, v127); + TEXTRACT(v126, v113, v125, v22); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v129 = (uint64_t)v22; + TASSIGN(v128, v129); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v128, v128, v119, v121); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v131 = (uint64_t)v22; + TASSIGN(v130, v131); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v130, v130, v123, v126); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v132 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v16, v21); + uint64_t v133 = (uint64_t)v22; + TASSIGN(v132, v133); + pto::Shape<1, 1, 1, 16, 1024> v134 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v135 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v136 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v136, v132); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__ssa_v7 + __gm__ Tensor *up_acc_all_inline303__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__ssa_v7 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__ssa_v7_tensor->buffer.addr) + + up_acc_all_inline303__ssa_v7_tensor->start_offset; + + // Unpack scalar: gu_k0_inline131__ssa_v4 + union { + uint64_t u64; + int64_t val; + } gu_k0_inline131__ssa_v4_conv; + gu_k0_inline131__ssa_v4_conv.u64 = args[3]; + int64_t gu_k0_inline131__ssa_v4 = gu_k0_inline131__ssa_v4_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + up_proj_3( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__ssa_v7, gu_k0_inline131__ssa_v4, + layer_hidden_base_inline151__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_4.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_4.cpp new file mode 100644 index 0000000000..136ab89c18 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/up_proj_4.cpp @@ -0,0 +1,622 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: up_proj_4 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +up_proj_4(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int64_t v5, int64_t v6) { + const int64_t v7 = 960; + const int64_t v8 = 2; + const int64_t v9 = 15; + const int64_t v10 = 32; + const int64_t v11 = 64; + const int64_t v12 = 17408; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 512; + const int64_t v17 = 2048; + const int64_t v18 = 32768; + const int64_t v19 = 1536; + const int64_t v20 = 1024; + const int64_t v21 = 0; + const int64_t v22 = 135168; + const int64_t v23 = 4096; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v24 = (size_t)v11; + size_t v25 = (size_t)v21; + size_t v26 = (size_t)v10; + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v27 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v28 = (uint64_t)v23; + TASSIGN(v27, v28); + pto::Shape<1, 1, 1, 16, 64> v29 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v31 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + v4 * v13), v29, v30 + ); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v27, v31); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v32 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v33 = (uint64_t)v22; + TASSIGN(v32, v33); + int64_t v34 = (int64_t)((uint64_t)v5 + (uint64_t)v4); + pto::Shape<1, 1, 1, 64, 1024> v35 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v36 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + v34 * v12 + v6 * v13), v35, v36); + TLOAD(v32, v37); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v38 = v25; v38 < v24; v38 += v26) { + int64_t v39 = (int64_t)v38; + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v41 = (uint64_t)v20; + TASSIGN(v40, v41); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v40, v27, v21, v38); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v42 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); + TEXTRACT(v42, v32, v38, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v45 = (uint64_t)v19; + TASSIGN(v44, v45); + int64_t v46 = (int64_t)((uint64_t)v39 + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v44, v27, v21, v46); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v48 = (uint64_t)v18; + TASSIGN(v47, v48); + TEXTRACT(v47, v32, v46, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v39 == v21) { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + pipe_barrier(PIPE_M); + TMATMUL(v49, v40, v42); + } else { + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v52 = (uint64_t)v21; + TASSIGN(v51, v52); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v51, v51, v40, v42); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v54 = (uint64_t)v21; + TASSIGN(v53, v54); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v53, v53, v44, v47); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + for (size_t v55 = (size_t)v13; v55 < ((size_t)v9); v55 += (size_t)v8) { + int64_t v56 = (int64_t)((uint64_t)((int64_t)v55) * (uint64_t)v11); + int64_t v57 = (int64_t)((uint64_t)v56 + (uint64_t)v11); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v59 = (uint64_t)v21; + TASSIGN(v58, v59); + pto::Shape<1, 1, 1, 16, 64> v60 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v61 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v62 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v13), v60, v61 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v58, v62); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v64 = (uint64_t)v22; + TASSIGN(v63, v64); + pto::Shape<1, 1, 1, 64, 1024> v65 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v66 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v56) * v12 + v6 * v13), v65, v66); + TLOAD(v63, v67); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v69 = (uint64_t)v17; + TASSIGN(v68, v69); + pto::Shape<1, 1, 1, 16, 64> v70 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v71 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v72 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v57) * v13), v70, v71 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v68, v72); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v74 = (uint64_t)v23; + TASSIGN(v73, v74); + pto::Shape<1, 1, 1, 64, 1024> v75 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v76 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND> + v77 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v57) * v12 + v6 * v13), v75, v76); + TLOAD(v73, v77); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v78 = v25; v78 < v24; v78 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v79 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v80 = (uint64_t)v20; + TASSIGN(v79, v80); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v79, v58, v21, v78); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v81 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); + TEXTRACT(v81, v63, v78, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v84 = (uint64_t)v19; + TASSIGN(v83, v84); + int64_t v85 = (int64_t)((uint64_t)((int64_t)v78) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v83, v58, v21, v85); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v86 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v87 = (uint64_t)v18; + TASSIGN(v86, v87); + TEXTRACT(v86, v63, v85, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v88 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v89 = (uint64_t)v21; + TASSIGN(v88, v89); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v88, v88, v79, v81); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v91 = (uint64_t)v21; + TASSIGN(v90, v91); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v90, v90, v83, v86); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v92 = v25; v92 < v24; v92 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v93 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v94 = (uint64_t)v21; + TASSIGN(v93, v94); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v93, v68, v21, v92); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v95 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v96 = (uint64_t)v21; + TASSIGN(v95, v96); + TEXTRACT(v95, v73, v92, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v98 = (uint64_t)v16; + TASSIGN(v97, v98); + int64_t v99 = (int64_t)((uint64_t)((int64_t)v92) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v97, v68, v21, v99); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null> + v100 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v101 = (uint64_t)v18; + TASSIGN(v100, v101); + TEXTRACT(v100, v73, v99, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v102 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v103 = (uint64_t)v21; + TASSIGN(v102, v103); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v102, v102, v93, v95); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v105 = (uint64_t)v21; + TASSIGN(v104, v105); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v104, v104, v97, v100); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v11); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + pto::Shape<1, 1, 1, 16, 64> v108 = pto::Shape<1, 1, 1, 16, 64>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v109 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v110 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v108, v109 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + TLOAD(v106, v110); + Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v11, v20); + uint64_t v112 = (uint64_t)v22; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 64, 1024> v113 = pto::Shape<1, 1, 1, 64, 1024>(); + pto::Stride<1114112, 1114112, 1114112, 17408, 1> v114 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, + pto::Layout::ND>(v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v7) * v12 + v6 * v13), v113, v114); + TLOAD(v111, v115); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + for (size_t v116 = v25; v116 < v24; v116 += v26) { + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v117 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v118 = (uint64_t)v20; + TASSIGN(v117, v118); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + TEXTRACT(v117, v106, v21, v116); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v119 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); + TEXTRACT(v119, v111, v116, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v121 = Tile< + TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v15, v15); + uint64_t v122 = (uint64_t)v19; + TASSIGN(v121, v122); + int64_t v123 = (int64_t)((uint64_t)((int64_t)v116) + (uint64_t)v15); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v121, v106, v21, v123); + Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v15, v20); + uint64_t v125 = (uint64_t)v18; + TASSIGN(v124, v125); + TEXTRACT(v124, v111, v123, v21); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v126, v126, v117, v119); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v128 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v129 = (uint64_t)v21; + TASSIGN(v128, v129); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v128, v128, v121, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v130 = Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v15, v20); + uint64_t v131 = (uint64_t)v21; + TASSIGN(v130, v131); + pto::Shape<1, 1, 1, 16, 1024> v132 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v133 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v134 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v21 + v21 * v12 + v6 * v13), v132, v133 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE< + Tile< + TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v134, v130); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v14 + __gm__ Tensor *mlp_norm_in_inline71__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v14 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v14_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v14_tensor->start_offset; + + // Unpack tensor: w_up__ssa_v0 + __gm__ Tensor *w_up__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *w_up__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(w_up__ssa_v0_tensor->buffer.addr) + w_up__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__iter_v11 + __gm__ Tensor *up_acc_all_inline303__iter_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__iter_v11 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__iter_v11_tensor->buffer.addr) + + up_acc_all_inline303__iter_v11_tensor->start_offset; + + // Unpack scalar: k0_inline113__ssa_v7 + union { + uint64_t u64; + int64_t val; + } k0_inline113__ssa_v7_conv; + k0_inline113__ssa_v7_conv.u64 = args[3]; + int64_t k0_inline113__ssa_v7 = k0_inline113__ssa_v7_conv.val; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[4]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Unpack scalar: n0_inline122__ssa_v6 + union { + uint64_t u64; + int64_t val; + } n0_inline122__ssa_v6_conv; + n0_inline122__ssa_v6_conv.u64 = args[5]; + int64_t n0_inline122__ssa_v6 = n0_inline122__ssa_v6_conv.val; + + // Forward to ptoas-generated function + up_proj_4( + mlp_norm_in_inline71__rv_v14, w_up__ssa_v0, up_acc_all_inline303__iter_v11, k0_inline113__ssa_v7, + layer_hidden_base_inline151__ssa_v0, n0_inline122__ssa_v6 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/v_proj.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/v_proj.cpp new file mode 100644 index 0000000000..6b98eb96c3 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aic/v_proj.cpp @@ -0,0 +1,617 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: v_proj +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void +v_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v4, int32_t v5, int32_t v6) { + const int64_t v7 = 768; + const int64_t v8 = 64; + const int64_t v9 = 128; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 512; + const int64_t v13 = 5; + const int64_t v14 = 5120; + const int64_t v15 = 1; + const int64_t v16 = 1024; + const int64_t v17 = 16; + const int64_t v18 = 2048; + const int64_t v19 = 8192; + const int64_t v20 = 32768; + const int64_t v21 = 6144; + const int64_t v22 = 4096; + const int64_t v23 = 0; + const int64_t v24 = 147456; + const int64_t v25 = 16384; + using T = float; + +#if defined(__DAV_CUBE__) + size_t v26 = (size_t)v23; + size_t v27 = (size_t)v10; + size_t v28 = (size_t)v9; + int64_t v29 = (int64_t)v5; + int64_t v30 = (int64_t)((uint64_t)(v29 % v13) * (uint64_t)v16); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + for (size_t v31 = v26; v31 < ((size_t)v11); v31 += (size_t)v15) { + int64_t v32 = (int64_t)((uint64_t)((int64_t)(uint64_t)(v29 / v13) * (uint64_t)v12) + + (uint64_t)((int64_t)(uint64_t)((int64_t)v31) * (uint64_t)v10)); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v34 = (uint64_t)v25; + TASSIGN(v33, v34); + pto::Shape<1, 1, 1, 16, 256> v35 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v36 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v37 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v30 * v15), v35, v36 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + TLOAD(v33, v37); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v39 = (uint64_t)v24; + TASSIGN(v38, v39); + pto::Shape<1, 1, 1, 256, 256> v40 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v41 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v42 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v30) * v16 + v32 * v15), v40, v41); + TLOAD(v38, v42); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v43 = v26; v43 < v27; v43 += v28) { + int64_t v44 = (int64_t)v43; + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v45 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v46 = (uint64_t)v22; + TASSIGN(v45, v46); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + pipe_barrier(PIPE_MTE1); + TEXTRACT(v45, v33, v23, v43); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v48 = (uint64_t)v23; + TASSIGN(v47, v48); + TEXTRACT(v47, v38, v43, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); + int64_t v51 = (int64_t)((uint64_t)v44 + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v49, v33, v23, v51); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v52 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v53 = (uint64_t)v20; + TASSIGN(v52, v53); + TEXTRACT(v52, v38, v51, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + if (v44 == v23) { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v54 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); + pipe_barrier(PIPE_M); + TMATMUL(v54, v45, v47); + } else { + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v56, v56, v45, v47); + }; + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v59 = (uint64_t)v23; + TASSIGN(v58, v59); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); + TMATMUL_ACC(v58, v58, v49, v52); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + int64_t v60 = (int64_t)((uint64_t)v30 + (uint64_t)v10); + int64_t v61 = (int64_t)((uint64_t)v30 + (uint64_t)v12); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v62 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v63 = (uint64_t)v23; + TASSIGN(v62, v63); + pto::Shape<1, 1, 1, 16, 256> v64 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v65 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v66 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v60 * v15), v64, v65 + ); + TLOAD(v62, v66); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v68 = (uint64_t)v24; + TASSIGN(v67, v68); + pto::Shape<1, 1, 1, 256, 256> v69 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v70 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v71 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v60) * v16 + v32 * v15), v69, v70); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); + TLOAD(v67, v71); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v73 = (uint64_t)v19; + TASSIGN(v72, v73); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v61 * v15), v74, v75 + ); + TLOAD(v72, v76); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v78 = (uint64_t)v25; + TASSIGN(v77, v78); + pto::Shape<1, 1, 1, 256, 256> v79 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v80 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v81 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v61) * v16 + v32 * v15), v79, v80); + TLOAD(v77, v81); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + for (size_t v82 = v26; v82 < v27; v82 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v84 = (uint64_t)v22; + TASSIGN(v83, v84); + pipe_barrier(PIPE_MTE1); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + TEXTRACT(v83, v62, v23, v82); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v86 = (uint64_t)v23; + TASSIGN(v85, v86); + TEXTRACT(v85, v67, v82, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v87 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v88 = (uint64_t)v21; + TASSIGN(v87, v88); + int64_t v89 = (int64_t)((uint64_t)((int64_t)v82) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + TEXTRACT(v87, v62, v23, v89); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v91 = (uint64_t)v20; + TASSIGN(v90, v91); + TEXTRACT(v90, v67, v89, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v92 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v93 = (uint64_t)v23; + TASSIGN(v92, v93); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v92, v92, v83, v85); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v94 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v95 = (uint64_t)v23; + TASSIGN(v94, v95); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); + TMATMUL_ACC(v94, v94, v87, v90); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + for (size_t v96 = v26; v96 < v27; v96 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v98 = (uint64_t)v23; + TASSIGN(v97, v98); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + TEXTRACT(v97, v72, v23, v96); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v100 = (uint64_t)v23; + TASSIGN(v99, v100); + TEXTRACT(v99, v77, v96, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v101 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v102 = (uint64_t)v18; + TASSIGN(v101, v102); + int64_t v103 = (int64_t)((uint64_t)((int64_t)v96) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v101, v72, v23, v103); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v104 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v105 = (uint64_t)v20; + TASSIGN(v104, v105); + TEXTRACT(v104, v77, v103, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v106 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v107 = (uint64_t)v23; + TASSIGN(v106, v107); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v106, v106, v97, v99); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v108 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v109 = (uint64_t)v23; + TASSIGN(v108, v109); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); + TMATMUL_ACC(v108, v108, v101, v104); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + int64_t v110 = (int64_t)((uint64_t)v30 + (uint64_t)v7); + Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v111 = Tile< + TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v112 = (uint64_t)v25; + TASSIGN(v111, v112); + pto::Shape<1, 1, 1, 16, 256> v113 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v114 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v115 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v23 + v23 * v14 + v110 * v15), v113, v114 + ); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); + TLOAD(v111, v115); + Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v116 = Tile< + TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null>(v10, v10); + uint64_t v117 = (uint64_t)v24; + TASSIGN(v116, v117); + pto::Shape<1, 1, 1, 256, 256> v118 = pto::Shape<1, 1, 1, 256, 256>(); + pto::Stride<262144, 262144, 262144, 1024, 1> v119 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, pto::Layout::ND> + v120 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, + pto::Layout::ND>(v3 + (v23 + (int64_t)((uint64_t)v4 + (uint64_t)v110) * v16 + v32 * v15), v118, v119); + TLOAD(v116, v120); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + for (size_t v121 = v26; v121 < v27; v121 += v28) { + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v122 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v123 = (uint64_t)v22; + TASSIGN(v122, v123); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + TEXTRACT(v122, v111, v23, v121); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v124 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v125 = (uint64_t)v23; + TASSIGN(v124, v125); + TEXTRACT(v124, v116, v121, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + CompactMode::Null> + v126 = Tile< + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, + PadValue::Null, CompactMode::Null>(v17, v8); + uint64_t v127 = (uint64_t)v21; + TASSIGN(v126, v127); + int64_t v128 = (int64_t)((uint64_t)((int64_t)v121) + (uint64_t)v8); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + TEXTRACT(v126, v111, v23, v128); + Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, + CompactMode::Null> + v129 = Tile< + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, + PadValue::Null, CompactMode::Null>(v8, v10); + uint64_t v130 = (uint64_t)v20; + TASSIGN(v129, v130); + TEXTRACT(v129, v116, v128, v23); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v131 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v132 = (uint64_t)v23; + TASSIGN(v131, v132); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); + pipe_barrier(PIPE_M); + TMATMUL_ACC(v131, v131, v122, v124); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v133 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v134 = (uint64_t)v23; + TASSIGN(v133, v134); + pipe_barrier(PIPE_M); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); + TMATMUL_ACC(v133, v133, v126, v129); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + }; + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); + set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null> + v135 = Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>(v17, v10); + uint64_t v136 = (uint64_t)v23; + TASSIGN(v135, v136); + pto::Shape<1, 1, 1, 16, 256> v137 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<16384, 16384, 16384, 1024, 1> v138 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> + v139 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 + ); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + pipe_barrier(PIPE_FIX); + TSTORE< + Tile< + TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, + CompactMode::Null>, + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>, + AtomicType::AtomicAdd>(v139, v135); + set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); + } + wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); +#endif // __DAV_CUBE__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: v_proj_inline255__ssa_v1 + __gm__ Tensor *v_proj_inline255__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *v_proj_inline255__ssa_v1 = + reinterpret_cast<__gm__ float *>(v_proj_inline255__ssa_v1_tensor->buffer.addr) + + v_proj_inline255__ssa_v1_tensor->start_offset; + + // Unpack tensor: normed__iter_v4 + __gm__ Tensor *normed__iter_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *normed__iter_v4 = reinterpret_cast<__gm__ bfloat16_t *>(normed__iter_v4_tensor->buffer.addr) + + normed__iter_v4_tensor->start_offset; + + // Unpack tensor: wv__ssa_v0 + __gm__ Tensor *wv__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ bfloat16_t *wv__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(wv__ssa_v0_tensor->buffer.addr) + wv__ssa_v0_tensor->start_offset; + + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 + union { + uint64_t u64; + int64_t val; + } layer_hidden_base_inline151__ssa_v0_conv; + layer_hidden_base_inline151__ssa_v0_conv.u64 = args[3]; + int64_t layer_hidden_base_inline151__ssa_v0 = layer_hidden_base_inline151__ssa_v0_conv.val; + + // Forward to ptoas-generated function + v_proj( + v_proj_inline255__ssa_v1, normed__iter_v4, wv__ssa_v0, layer_hidden_base_inline151__ssa_v0, + __pypto_spmd_block_idx, __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp new file mode 100644 index 0000000000..4fc5a9778d --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp @@ -0,0 +1,73 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: attn_out_seed +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void attn_out_seed(__gm__ bfloat16_t *v1) { + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); +#endif // __DAV_VEC__ + + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: attn_out_inline282__ssa_v0 + __gm__ Tensor *attn_out_inline282__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *attn_out_inline282__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline282__ssa_v0_tensor->buffer.addr) + + attn_out_inline282__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + attn_out_seed(attn_out_inline282__ssa_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp new file mode 100644 index 0000000000..b743ef0c9e --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp @@ -0,0 +1,145 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: copy_hidden +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void copy_hidden(__gm__ float *v1, __gm__ bfloat16_t *v2, int64_t v3) { + SaturationMode v4 = SaturationMode::OFF; + RoundMode v5 = RoundMode::CAST_ROUND; + const int64_t v6 = 256; + const int64_t v7 = 20; + const int64_t v8 = 1; + const int64_t v9 = 5120; + const int64_t v10 = 16; + const int64_t v11 = 8192; + const int64_t v12 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + for (size_t v13 = (size_t)v12; v13 < ((size_t)v7); v13 += (size_t)v8) { + int64_t v14 = (int64_t)((uint64_t)((int64_t)v13) * (uint64_t)v6); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v15 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v6); + uint64_t v16 = (uint64_t)v12; + TASSIGN(v15, v16); + pto::Shape<1, 1, 1, 16, 256> v17 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v18 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v19 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v12 + v3 * v9 + v14 * v8), v17, v18 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v15, v19); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v20 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v6); + uint64_t v21 = (uint64_t)v11; + TASSIGN(v20, v21); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TCVT(v20, v15, v5, v4); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v22 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v23 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v24 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v12 + v3 * v9 + v14 * v8), v22, v23 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v24, v20); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: cur__iter_v1 + __gm__ Tensor *cur__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *cur__iter_v1 = + reinterpret_cast<__gm__ float *>(cur__iter_v1_tensor->buffer.addr) + cur__iter_v1_tensor->start_offset; + + // Unpack tensor: hidden_states__ssa_v0 + __gm__ Tensor *hidden_states__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *hidden_states__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(hidden_states__ssa_v0_tensor->buffer.addr) + + hidden_states__ssa_v0_tensor->start_offset; + + // Unpack scalar: cb0__idx_v0 + union { + uint64_t u64; + int64_t val; + } cb0__idx_v0_conv; + cb0__idx_v0_conv.u64 = args[2]; + int64_t cb0__idx_v0 = cb0__idx_v0_conv.val; + + // Forward to ptoas-generated function + copy_hidden(cur__iter_v1, hidden_states__ssa_v0, cb0__idx_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_out.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_out.cpp new file mode 100644 index 0000000000..16110bc219 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/copy_out.cpp @@ -0,0 +1,139 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: copy_out +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void copy_out(__gm__ bfloat16_t *v1, __gm__ float *v2, int64_t v3) { + SaturationMode v4 = SaturationMode::OFF; + RoundMode v5 = RoundMode::CAST_ROUND; + const int64_t v6 = 256; + const int64_t v7 = 20; + const int64_t v8 = 1; + const int64_t v9 = 5120; + const int64_t v10 = 16; + const int64_t v11 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + for (size_t v12 = (size_t)v11; v12 < ((size_t)v7); v12 += (size_t)v8) { + int64_t v13 = (int64_t)((uint64_t)((int64_t)v12) * (uint64_t)v6); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v14 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v6); + uint64_t v15 = (uint64_t)v11; + TASSIGN(v14, v15); + pto::Shape<1, 1, 1, 16, 256> v16 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v17 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v18 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v11 + v3 * v9 + v13 * v8), v16, v17 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + TLOAD(v14, v18); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v19 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v6); + uint64_t v20 = (uint64_t)v11; + TASSIGN(v19, v20); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCVT(v19, v14, v5, v4); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v21 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v22 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v23 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v11 + v3 * v9 + v13 * v8), v21, v22 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v23, v19); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: out__iter_v1 + __gm__ Tensor *out__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *out__iter_v1 = + reinterpret_cast<__gm__ bfloat16_t *>(out__iter_v1_tensor->buffer.addr) + out__iter_v1_tensor->start_offset; + + // Unpack tensor: cur__rv_v7 + __gm__ Tensor *cur__rv_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *cur__rv_v7 = + reinterpret_cast<__gm__ float *>(cur__rv_v7_tensor->buffer.addr) + cur__rv_v7_tensor->start_offset; + + // Unpack scalar: ob0__idx_v0 + union { + uint64_t u64; + int64_t val; + } ob0__idx_v0_conv; + ob0__idx_v0_conv.u64 = args[2]; + int64_t ob0__idx_v0 = ob0__idx_v0_conv.val; + + // Forward to ptoas-generated function + copy_out(out__iter_v1, cur__rv_v7, ob0__idx_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp new file mode 100644 index 0000000000..7af717198c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp @@ -0,0 +1,226 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: dcr_xgamma +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void dcr_xgamma( + __gm__ float *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ bfloat16_t *v5, int64_t v6, + int32_t v7, int32_t v8 +) { + SaturationMode v9 = SaturationMode::OFF; + RoundMode v10 = RoundMode::CAST_ROUND; + const int64_t v11 = 1024; + const int64_t v12 = 1; + const int64_t v13 = 5120; + const int64_t v14 = 16; + const int64_t v15 = 65536; + const int64_t v16 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + int64_t v17 = (int64_t)((uint64_t)((int64_t)v7) * (uint64_t)v11); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v18 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v11); + uint64_t v19 = (uint64_t)v16; + TASSIGN(v18, v19); + pto::Shape<1, 1, 1, 16, 1024> v20 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v21 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v22 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v13 + v17 * v12), v20, v21 + ); + TLOAD(v18, v22); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v23 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v11); + uint64_t v24 = (uint64_t)v15; + TASSIGN(v23, v24); + pto::Shape<1, 1, 1, 16, 1024> v25 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v26 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v27 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v13 + v17 * v12), v25, v26 + ); + TLOAD(v23, v27); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v28 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v11); + uint64_t v29 = (uint64_t)v16; + TASSIGN(v28, v29); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v28, v18, v23); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 1024> v30 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v31 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v32 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v13 + v17 * v12), v30, v31 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v32, v28); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Vec, float, 1, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v11); + uint64_t v34 = (uint64_t)v15; + TASSIGN(v33, v34); + pto::Shape<1, 1, 1, 1, 1024> v35 = pto::Shape<1, 1, 1, 1, 1024>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v36 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v37 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v6 * v13 + v17 * v12), v35, v36 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v33, v37); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v11); + uint64_t v39 = (uint64_t)v16; + TASSIGN(v38, v39); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TCOLEXPANDMUL(v38, v28, v33); + Tile< + TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v11); + uint64_t v41 = (uint64_t)v16; + TASSIGN(v40, v41); + pipe_barrier(PIPE_V); + TCVT(v40, v38, v10, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 1024> v42 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v43 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v44 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v16 * v13 + v17 * v12), v42, v43 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v44, v40); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: down_acc_all_inline168__rv_v5 + __gm__ Tensor *down_acc_all_inline168__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *down_acc_all_inline168__rv_v5 = + reinterpret_cast<__gm__ float *>(down_acc_all_inline168__rv_v5_tensor->buffer.addr) + + down_acc_all_inline168__rv_v5_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__rv_v14 + __gm__ Tensor *post_norm_partial_inline118__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__rv_v14 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__rv_v14_tensor->buffer.addr) + + post_norm_partial_inline118__rv_v14_tensor->start_offset; + + // Unpack tensor: next_hidden__ssa_v0 + __gm__ Tensor *next_hidden__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *next_hidden__ssa_v0 = reinterpret_cast<__gm__ float *>(next_hidden__ssa_v0_tensor->buffer.addr) + + next_hidden__ssa_v0_tensor->start_offset; + + // Unpack tensor: input_rms_weight__ssa_v0 + __gm__ Tensor *input_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *input_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(input_rms_weight__ssa_v0_tensor->buffer.addr) + + input_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack tensor: next_normed__ssa_v0 + __gm__ Tensor *next_normed__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ bfloat16_t *next_normed__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(next_normed__ssa_v0_tensor->buffer.addr) + + next_normed__ssa_v0_tensor->start_offset; + + // Unpack scalar: next_gamma_idx__ssa_v0 + union { + uint64_t u64; + int64_t val; + } next_gamma_idx__ssa_v0_conv; + next_gamma_idx__ssa_v0_conv.u64 = args[5]; + int64_t next_gamma_idx__ssa_v0 = next_gamma_idx__ssa_v0_conv.val; + + // Forward to ptoas-generated function + dcr_xgamma( + down_acc_all_inline168__rv_v5, post_norm_partial_inline118__rv_v14, next_hidden__ssa_v0, + input_rms_weight__ssa_v0, next_normed__ssa_v0, next_gamma_idx__ssa_v0, __pypto_spmd_block_idx, + __pypto_spmd_block_num + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/kv_seed.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/kv_seed.cpp new file mode 100644 index 0000000000..632e4bb663 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/kv_seed.cpp @@ -0,0 +1,123 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: kv_seed +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void kv_seed(__gm__ float *v1, __gm__ float *v2) { + const float v3 = 0.0f; + const int64_t v4 = 1; + const int64_t v5 = 1024; + const int64_t v6 = 16; + const int64_t v7 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v8 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v6, v5); + uint64_t v9 = (uint64_t)v7; + TASSIGN(v8, v9); + TEXPANDS(v8, v3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 1024> v10 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<16384, 16384, 16384, 1024, 1> v11 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> v12 = + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( + v1 + (v7 + v7 * v5 + v7 * v4), v10, v11 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v12, v8); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v13 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v6, v5); + uint64_t v14 = (uint64_t)v7; + TASSIGN(v13, v14); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v13, v3); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 1024> v15 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<16384, 16384, 16384, 1024, 1> v16 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> v17 = + GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( + v2 + (v7 + v7 * v5 + v7 * v4), v15, v16 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v17, v13); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: k_proj_inline135__ssa_v0 + __gm__ Tensor *k_proj_inline135__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *k_proj_inline135__ssa_v0 = + reinterpret_cast<__gm__ float *>(k_proj_inline135__ssa_v0_tensor->buffer.addr) + + k_proj_inline135__ssa_v0_tensor->start_offset; + + // Unpack tensor: v_proj_inline255__ssa_v0 + __gm__ Tensor *v_proj_inline255__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *v_proj_inline255__ssa_v0 = + reinterpret_cast<__gm__ float *>(v_proj_inline255__ssa_v0_tensor->buffer.addr) + + v_proj_inline255__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + kv_seed(k_proj_inline135__ssa_v0, v_proj_inline255__ssa_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp new file mode 100644 index 0000000000..d528272138 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp @@ -0,0 +1,391 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: mlp_out_seed +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void mlp_out_seed(__gm__ float *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4) { + const int64_t v5 = 16384; + const int64_t v6 = 4096; + const int64_t v7 = 512; + const int64_t v8 = 10; + const float v9 = 0.0f; + const int64_t v10 = 1024; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 17408; + const int64_t v14 = 1; + const int64_t v15 = 5120; + const int64_t v16 = 16; + const int64_t v17 = 65536; + const int64_t v18 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + size_t v19 = (size_t)v16; + size_t v20 = (size_t)v18; + size_t v21 = (size_t)v11; + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + for (size_t v22 = v20; v22 < ((size_t)v12); v22 += v21) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v24 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v25 = (uint64_t)v18; + TASSIGN(v24, v25); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v24, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 1024> v26 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v27 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v28 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v18 + v18 * v15 + v23 * v14), v26, v27 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v28, v24); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v30 = (uint64_t)v17; + TASSIGN(v29, v30); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TEXPANDS(v29, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 1024> v31 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v32 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v33 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v18 + v18 * v15 + (int64_t)((uint64_t)v23 + (uint64_t)v10) * v14), v31, v32 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pipe_barrier(PIPE_MTE3); + TSTORE(v33, v29); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v35 = (uint64_t)v18; + TASSIGN(v34, v35); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); + TEXPANDS(v34, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 1024> v36 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v37 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v38 = + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v18 + v18 * v15 + v6 * v14), v36, v37 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v38, v34); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID3); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID3); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID5); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID4); + for (size_t v39 = v20; v39 < v19; v39 += v21) { + int64_t v40 = (int64_t)((uint64_t)((int64_t)v39) * (uint64_t)v10); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID4); + TEXPANDS(v41, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 1024> v43 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v44 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v2 + (v18 + v18 * v13 + v40 * v14), v43, v44 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pipe_barrier(PIPE_MTE3); + TSTORE(v45, v41); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID4); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID5); + TEXPANDS(v46, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); + pto::Shape<1, 1, 1, 16, 1024> v48 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v49 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v2 + (v18 + v18 * v13 + (int64_t)((uint64_t)v40 + (uint64_t)v10) * v14), v48, v49 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); + pipe_barrier(PIPE_MTE3); + TSTORE(v50, v46); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID5); + } + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID4); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID5); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID6); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v52 = (uint64_t)v18; + TASSIGN(v51, v52); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID6); + TEXPANDS(v51, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + pto::Shape<1, 1, 1, 16, 1024> v53 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v54 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v55 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v2 + (v18 + v18 * v13 + v5 * v14), v53, v54 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + pipe_barrier(PIPE_MTE3); + TSTORE(v55, v51); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID7); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID7); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + for (size_t v56 = v20; v56 < v19; v56 += v21) { + int64_t v57 = (int64_t)((uint64_t)((int64_t)v56) * (uint64_t)v10); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v59 = (uint64_t)v18; + TASSIGN(v58, v59); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v58, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID6); + pto::Shape<1, 1, 1, 16, 1024> v60 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v61 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v62 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v13 + v57 * v14), v60, v61 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID6); + pipe_barrier(PIPE_MTE3); + TSTORE(v62, v58); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v64 = (uint64_t)v17; + TASSIGN(v63, v64); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TEXPANDS(v63, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID7); + pto::Shape<1, 1, 1, 16, 1024> v65 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v66 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v67 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v13 + (int64_t)((uint64_t)v57 + (uint64_t)v10) * v14), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID7); + pipe_barrier(PIPE_MTE3); + TSTORE(v67, v63); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v68, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 1024> v70 = pto::Shape<1, 1, 1, 16, 1024>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v71 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v72 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v13 + v5 * v14), v70, v71 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v72, v68); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + for (size_t v73 = v20; v73 < ((size_t)v8); v73 += v21) { + int64_t v74 = (int64_t)((uint64_t)((int64_t)v73) * (uint64_t)v7); + Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v7); + uint64_t v76 = (uint64_t)v18; + TASSIGN(v75, v76); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v75, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 512> v77 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v18 + v18 * v15 + v74 * v14), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v79, v75); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v80 = Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v16, v7); + uint64_t v81 = (uint64_t)v17; + TASSIGN(v80, v81); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TEXPANDS(v80, v9); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 512> v82 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v83 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v84 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v18 + v18 * v15 + (int64_t)((uint64_t)v74 + (uint64_t)v7) * v14), v82, v83 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v84, v80); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: down_acc_all_inline168__ssa_v0 + __gm__ Tensor *down_acc_all_inline168__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *down_acc_all_inline168__ssa_v0 = + reinterpret_cast<__gm__ float *>(down_acc_all_inline168__ssa_v0_tensor->buffer.addr) + + down_acc_all_inline168__ssa_v0_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__ssa_v0 + __gm__ Tensor *gate_acc_all_inline203__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *gate_acc_all_inline203__ssa_v0 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__ssa_v0_tensor->buffer.addr) + + gate_acc_all_inline203__ssa_v0_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__ssa_v0 + __gm__ Tensor *up_acc_all_inline303__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *up_acc_all_inline303__ssa_v0 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__ssa_v0_tensor->buffer.addr) + + up_acc_all_inline303__ssa_v0_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v0 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *attn_proj_fp32_inline220__ssa_v0 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v0_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + mlp_out_seed( + down_acc_all_inline168__ssa_v0, gate_acc_all_inline203__ssa_v0, up_acc_all_inline303__ssa_v0, + attn_proj_fp32_inline220__ssa_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp new file mode 100644 index 0000000000..0c236440d2 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp @@ -0,0 +1,353 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: post_rms_reduce +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void post_rms_reduce(__gm__ float *v1, __gm__ float *v2, __gm__ float *v3) { + const float v4 = 9.99999997E-7f; + const float v5 = 1.95312503E-4f; + const int64_t v6 = 256; + const int64_t v7 = 2; + const int64_t v8 = 20; + const float v9 = 0.0f; + const int64_t v10 = 1; + const int64_t v11 = 5120; + const int64_t v12 = 16; + const int64_t v13 = 64; + const int64_t v14 = 0; + const int64_t v15 = 49344; + const int64_t v16 = 32960; + const int64_t v17 = 16576; + const int64_t v18 = 192; + const int64_t v19 = 128; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v20 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v21 = (uint64_t)v19; + TASSIGN(v20, v21); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TEXPANDS(v20, v9); + for (size_t v22 = (size_t)v14; v22 < ((size_t)v8); v22 += (size_t)v7) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v6); + int64_t v24 = (int64_t)((uint64_t)v23 + (uint64_t)v6); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v25 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v26 = (uint64_t)v18; + TASSIGN(v25, v26); + pto::Shape<1, 1, 1, 16, 256> v27 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v28 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v29 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v14 + v14 * v11 + v23 * v10), v27, v28 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v25, v29); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v30 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v31 = (uint64_t)v17; + TASSIGN(v30, v31); + pto::Shape<1, 1, 1, 16, 256> v32 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v33 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v34 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v14 + v14 * v11 + v23 * v10), v32, v33 + ); + TLOAD(v30, v34); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v35 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v36 = (uint64_t)v16; + TASSIGN(v35, v36); + pto::Shape<1, 1, 1, 16, 256> v37 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v38 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v39 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v14 + v14 * v11 + v24 * v10), v37, v38 + ); + TLOAD(v35, v39); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v40 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v41 = (uint64_t)v15; + TASSIGN(v40, v41); + pto::Shape<1, 1, 1, 16, 256> v42 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v43 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v44 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v14 + v14 * v11 + v24 * v10), v42, v43 + ); + TLOAD(v40, v44); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v45 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v46 = (uint64_t)v18; + TASSIGN(v45, v46); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v45, v25, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v47 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v48 = (uint64_t)v18; + TASSIGN(v47, v48); + pipe_barrier(PIPE_V); + TMUL(v47, v45, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v49 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v50 = (uint64_t)v17; + TASSIGN(v49, v50); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v10); + uint64_t v52 = (uint64_t)v14; + TASSIGN(v51, v52); + pipe_barrier(PIPE_V); + TROWSUM(v51, v47, v49); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v54 = (uint64_t)v14; + TASSIGN(v53, v54); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v56 = (uint64_t)v18; + TASSIGN(v55, v56); + pipe_barrier(PIPE_V); + TADD(v55, v20, v53); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v57 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v58 = (uint64_t)v16; + TASSIGN(v57, v58); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TADD(v57, v35, v40); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v59 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v60 = (uint64_t)v16; + TASSIGN(v59, v60); + pipe_barrier(PIPE_V); + TMUL(v59, v57, v57); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v61 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v6); + uint64_t v62 = (uint64_t)v15; + TASSIGN(v61, v62); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v10); + uint64_t v64 = (uint64_t)v13; + TASSIGN(v63, v64); + pipe_barrier(PIPE_V); + TROWSUM(v63, v59, v61); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v66 = (uint64_t)v13; + TASSIGN(v65, v66); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v68 = (uint64_t)v19; + TASSIGN(v67, v68); + pipe_barrier(PIPE_V); + TADD(v67, v55, v65); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + } + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v69 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v70 = (uint64_t)v18; + TASSIGN(v69, v70); + pipe_barrier(PIPE_V); + TMULS(v69, v20, v5); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v71 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v72 = (uint64_t)v18; + TASSIGN(v71, v72); + pipe_barrier(PIPE_V); + TADDS(v71, v69, v4); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v74 = (uint64_t)v18; + TASSIGN(v73, v74); + pipe_barrier(PIPE_V); + TSQRT(v73, v71); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v10, v12); + uint64_t v76 = (uint64_t)v17; + TASSIGN(v75, v76); + pipe_barrier(PIPE_V); + TRECIP(v75, v73); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v10); + uint64_t v78 = (uint64_t)v17; + TASSIGN(v77, v78); + pto::Shape<1, 1, 1, 16, 1> v79 = pto::Shape<1, 1, 1, 16, 1>(); + pto::Stride<16, 16, 16, 1, 16> v80 = pto::Stride<16, 16, 16, 1, 16>(); + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v81 = + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN>( + v3 + (v14 + v14 * v10 + v14 * v12), v79, v80 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v81, v77); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: inv_rms_tile_inline126__ssa_v0 + __gm__ Tensor *inv_rms_tile_inline126__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *inv_rms_tile_inline126__ssa_v0 = + reinterpret_cast<__gm__ float *>(inv_rms_tile_inline126__ssa_v0_tensor->buffer.addr) + + inv_rms_tile_inline126__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + post_rms_reduce(attn_proj_fp32_inline220__ssa_v7, cur__iter_v6, inv_rms_tile_inline126__ssa_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/q_seed.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/q_seed.cpp new file mode 100644 index 0000000000..f0495a0ddc --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/q_seed.cpp @@ -0,0 +1,134 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: q_seed +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void q_seed(__gm__ float *v1) { + const float v2 = 0.0f; + const int64_t v3 = 512; + const int64_t v4 = 2; + const int64_t v5 = 10; + const int64_t v6 = 1; + const int64_t v7 = 5120; + const int64_t v8 = 16; + const int64_t v9 = 32768; + const int64_t v10 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + for (size_t v11 = (size_t)v10; v11 < ((size_t)v5); v11 += (size_t)v4) { + Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v12 = Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v8, v3); + uint64_t v13 = (uint64_t)v10; + TASSIGN(v12, v13); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TEXPANDS(v12, v2); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + int64_t v14 = (int64_t)((uint64_t)((int64_t)v11) * (uint64_t)v3); + pto::Shape<1, 1, 1, 16, 512> v15 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v16 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v17 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v10 + v10 * v7 + v14 * v6), v15, v16 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v17, v12); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v18 = Tile< + TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v8, v3); + uint64_t v19 = (uint64_t)v9; + TASSIGN(v18, v19); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TEXPANDS(v18, v2); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 512> v20 = pto::Shape<1, 1, 1, 16, 512>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v21 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v22 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v10 + v10 * v7 + (int64_t)((uint64_t)v14 + (uint64_t)v3) * v6), v20, v21 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pipe_barrier(PIPE_MTE3); + TSTORE(v22, v18); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: q_proj_inline139__ssa_v0 + __gm__ Tensor *q_proj_inline139__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *q_proj_inline139__ssa_v0 = + reinterpret_cast<__gm__ float *>(q_proj_inline139__ssa_v0_tensor->buffer.addr) + + q_proj_inline139__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + q_seed(q_proj_inline139__ssa_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp new file mode 100644 index 0000000000..5d44c597bd --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp @@ -0,0 +1,366 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: residual_rms_cast +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void residual_rms_cast( + __gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int64_t v6, + int64_t v7 +) { + SaturationMode v8 = SaturationMode::OFF; + RoundMode v9 = RoundMode::CAST_ROUND; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 0; + const int64_t v17 = 51200; + const int64_t v18 = 34816; + const int64_t v19 = 33792; + const int64_t v20 = 17408; + const int64_t v21 = 1024; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + for (size_t v22 = (size_t)v16; v22 < ((size_t)v12); v22 += (size_t)v11) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + int64_t v24 = (int64_t)((uint64_t)v6 + (uint64_t)v23); + int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v23 + (uint64_t)v10)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v26 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v27 = (uint64_t)v21; + TASSIGN(v26, v27); + pto::Shape<1, 1, 1, 16, 256> v28 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v30 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v32 = (uint64_t)v20; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v24 * v13), v33, v34 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v37 = (uint64_t)v19; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 1, 256> v38 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v39 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v40 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 + ); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v25 * v13), v48, v49 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + pto::Shape<1, 1, 1, 1, 256> v53 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v54 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v55 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 + ); + TLOAD(v51, v55); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v56, v26, v31); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v58, v56, v36); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); + pipe_barrier(PIPE_V); + TCVT(v60, v58, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v62 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v24 * v13), v62, v63 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v64, v56); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v65 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v66 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v24 * v13), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v67, v60); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TADD(v68, v41, v46); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TCOLEXPANDMUL(v70, v68, v51); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); + pipe_barrier(PIPE_V); + TCVT(v72, v70, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v25 * v13), v74, v75 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v76, v68); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 256> v77 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + TSTORE(v79, v72); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__ssa_v0 + __gm__ Tensor *mlp_norm_in_inline71__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__ssa_v0_tensor->buffer.addr) + + mlp_norm_in_inline71__ssa_v0_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__ssa_v0 + __gm__ Tensor *post_norm_partial_inline118__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__ssa_v0_tensor->buffer.addr) + + post_norm_partial_inline118__ssa_v0_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: post_rms_weight__ssa_v0 + __gm__ Tensor *post_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ float *post_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + + post_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack scalar: k_base_inline111__ssa_v0 + union { + uint64_t u64; + int64_t val; + } k_base_inline111__ssa_v0_conv; + k_base_inline111__ssa_v0_conv.u64 = args[5]; + int64_t k_base_inline111__ssa_v0 = k_base_inline111__ssa_v0_conv.val; + + // Unpack scalar: i__idx_v0 + union { + uint64_t u64; + int64_t val; + } i__idx_v0_conv; + i__idx_v0_conv.u64 = args[6]; + int64_t i__idx_v0 = i__idx_v0_conv.val; + + // Forward to ptoas-generated function + residual_rms_cast( + mlp_norm_in_inline71__ssa_v0, post_norm_partial_inline118__ssa_v0, attn_proj_fp32_inline220__ssa_v7, + cur__iter_v6, post_rms_weight__ssa_v0, k_base_inline111__ssa_v0, i__idx_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp new file mode 100644 index 0000000000..d3dc242f7c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp @@ -0,0 +1,366 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: residual_rms_cast_0 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void residual_rms_cast_0( + __gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int64_t v6, + int64_t v7 +) { + SaturationMode v8 = SaturationMode::OFF; + RoundMode v9 = RoundMode::CAST_ROUND; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 0; + const int64_t v17 = 51200; + const int64_t v18 = 34816; + const int64_t v19 = 33792; + const int64_t v20 = 17408; + const int64_t v21 = 1024; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + for (size_t v22 = (size_t)v16; v22 < ((size_t)v12); v22 += (size_t)v11) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + int64_t v24 = (int64_t)((uint64_t)v6 + (uint64_t)v23); + int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v23 + (uint64_t)v10)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v26 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v27 = (uint64_t)v21; + TASSIGN(v26, v27); + pto::Shape<1, 1, 1, 16, 256> v28 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v30 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v32 = (uint64_t)v20; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v24 * v13), v33, v34 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v37 = (uint64_t)v19; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 1, 256> v38 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v39 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v40 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 + ); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v25 * v13), v48, v49 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + pto::Shape<1, 1, 1, 1, 256> v53 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v54 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v55 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 + ); + TLOAD(v51, v55); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v56, v26, v31); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v58, v56, v36); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); + pipe_barrier(PIPE_V); + TCVT(v60, v58, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v62 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v24 * v13), v62, v63 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v64, v56); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v65 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v66 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v24 * v13), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v67, v60); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TADD(v68, v41, v46); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TCOLEXPANDMUL(v70, v68, v51); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); + pipe_barrier(PIPE_V); + TCVT(v72, v70, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v25 * v13), v74, v75 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v76, v68); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 256> v77 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + TSTORE(v79, v72); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v2 + __gm__ Tensor *mlp_norm_in_inline71__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v2 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v2_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v2_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__rv_v2 + __gm__ Tensor *post_norm_partial_inline118__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__rv_v2 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__rv_v2_tensor->buffer.addr) + + post_norm_partial_inline118__rv_v2_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: post_rms_weight__ssa_v0 + __gm__ Tensor *post_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ float *post_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + + post_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack scalar: k_base_inline111__ssa_v1 + union { + uint64_t u64; + int64_t val; + } k_base_inline111__ssa_v1_conv; + k_base_inline111__ssa_v1_conv.u64 = args[5]; + int64_t k_base_inline111__ssa_v1 = k_base_inline111__ssa_v1_conv.val; + + // Unpack scalar: i__idx_v0 + union { + uint64_t u64; + int64_t val; + } i__idx_v0_conv; + i__idx_v0_conv.u64 = args[6]; + int64_t i__idx_v0 = i__idx_v0_conv.val; + + // Forward to ptoas-generated function + residual_rms_cast_0( + mlp_norm_in_inline71__rv_v2, post_norm_partial_inline118__rv_v2, attn_proj_fp32_inline220__ssa_v7, cur__iter_v6, + post_rms_weight__ssa_v0, k_base_inline111__ssa_v1, i__idx_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp new file mode 100644 index 0000000000..0a54a73863 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp @@ -0,0 +1,366 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: residual_rms_cast_1 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void residual_rms_cast_1( + __gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int64_t v6, + int64_t v7 +) { + SaturationMode v8 = SaturationMode::OFF; + RoundMode v9 = RoundMode::CAST_ROUND; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 0; + const int64_t v17 = 51200; + const int64_t v18 = 34816; + const int64_t v19 = 33792; + const int64_t v20 = 17408; + const int64_t v21 = 1024; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + for (size_t v22 = (size_t)v16; v22 < ((size_t)v12); v22 += (size_t)v11) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + int64_t v24 = (int64_t)((uint64_t)v6 + (uint64_t)v23); + int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v23 + (uint64_t)v10)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v26 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v27 = (uint64_t)v21; + TASSIGN(v26, v27); + pto::Shape<1, 1, 1, 16, 256> v28 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v30 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v32 = (uint64_t)v20; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v24 * v13), v33, v34 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v37 = (uint64_t)v19; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 1, 256> v38 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v39 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v40 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 + ); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v25 * v13), v48, v49 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + pto::Shape<1, 1, 1, 1, 256> v53 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v54 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v55 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 + ); + TLOAD(v51, v55); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v56, v26, v31); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v58, v56, v36); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); + pipe_barrier(PIPE_V); + TCVT(v60, v58, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v62 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v24 * v13), v62, v63 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v64, v56); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v65 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v66 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v24 * v13), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v67, v60); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TADD(v68, v41, v46); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TCOLEXPANDMUL(v70, v68, v51); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); + pipe_barrier(PIPE_V); + TCVT(v72, v70, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v25 * v13), v74, v75 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v76, v68); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 256> v77 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + TSTORE(v79, v72); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v5 + __gm__ Tensor *mlp_norm_in_inline71__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v5 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v5_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v5_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__rv_v5 + __gm__ Tensor *post_norm_partial_inline118__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__rv_v5 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__rv_v5_tensor->buffer.addr) + + post_norm_partial_inline118__rv_v5_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: post_rms_weight__ssa_v0 + __gm__ Tensor *post_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ float *post_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + + post_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack scalar: k_base_inline111__ssa_v2 + union { + uint64_t u64; + int64_t val; + } k_base_inline111__ssa_v2_conv; + k_base_inline111__ssa_v2_conv.u64 = args[5]; + int64_t k_base_inline111__ssa_v2 = k_base_inline111__ssa_v2_conv.val; + + // Unpack scalar: i__idx_v0 + union { + uint64_t u64; + int64_t val; + } i__idx_v0_conv; + i__idx_v0_conv.u64 = args[6]; + int64_t i__idx_v0 = i__idx_v0_conv.val; + + // Forward to ptoas-generated function + residual_rms_cast_1( + mlp_norm_in_inline71__rv_v5, post_norm_partial_inline118__rv_v5, attn_proj_fp32_inline220__ssa_v7, cur__iter_v6, + post_rms_weight__ssa_v0, k_base_inline111__ssa_v2, i__idx_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp new file mode 100644 index 0000000000..3b85e537b8 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp @@ -0,0 +1,366 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: residual_rms_cast_2 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void residual_rms_cast_2( + __gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int64_t v6, + int64_t v7 +) { + SaturationMode v8 = SaturationMode::OFF; + RoundMode v9 = RoundMode::CAST_ROUND; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 0; + const int64_t v17 = 51200; + const int64_t v18 = 34816; + const int64_t v19 = 33792; + const int64_t v20 = 17408; + const int64_t v21 = 1024; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + for (size_t v22 = (size_t)v16; v22 < ((size_t)v12); v22 += (size_t)v11) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + int64_t v24 = (int64_t)((uint64_t)v6 + (uint64_t)v23); + int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v23 + (uint64_t)v10)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v26 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v27 = (uint64_t)v21; + TASSIGN(v26, v27); + pto::Shape<1, 1, 1, 16, 256> v28 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v30 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v32 = (uint64_t)v20; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v24 * v13), v33, v34 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v37 = (uint64_t)v19; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 1, 256> v38 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v39 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v40 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 + ); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v25 * v13), v48, v49 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + pto::Shape<1, 1, 1, 1, 256> v53 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v54 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v55 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 + ); + TLOAD(v51, v55); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v56, v26, v31); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v58, v56, v36); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); + pipe_barrier(PIPE_V); + TCVT(v60, v58, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v62 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v24 * v13), v62, v63 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v64, v56); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v65 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v66 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v24 * v13), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v67, v60); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TADD(v68, v41, v46); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TCOLEXPANDMUL(v70, v68, v51); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); + pipe_barrier(PIPE_V); + TCVT(v72, v70, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v25 * v13), v74, v75 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v76, v68); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 256> v77 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + TSTORE(v79, v72); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v8 + __gm__ Tensor *mlp_norm_in_inline71__rv_v8_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v8 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v8_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v8_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__rv_v8 + __gm__ Tensor *post_norm_partial_inline118__rv_v8_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__rv_v8 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__rv_v8_tensor->buffer.addr) + + post_norm_partial_inline118__rv_v8_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: post_rms_weight__ssa_v0 + __gm__ Tensor *post_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ float *post_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + + post_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack scalar: k_base_inline111__ssa_v3 + union { + uint64_t u64; + int64_t val; + } k_base_inline111__ssa_v3_conv; + k_base_inline111__ssa_v3_conv.u64 = args[5]; + int64_t k_base_inline111__ssa_v3 = k_base_inline111__ssa_v3_conv.val; + + // Unpack scalar: i__idx_v0 + union { + uint64_t u64; + int64_t val; + } i__idx_v0_conv; + i__idx_v0_conv.u64 = args[6]; + int64_t i__idx_v0 = i__idx_v0_conv.val; + + // Forward to ptoas-generated function + residual_rms_cast_2( + mlp_norm_in_inline71__rv_v8, post_norm_partial_inline118__rv_v8, attn_proj_fp32_inline220__ssa_v7, cur__iter_v6, + post_rms_weight__ssa_v0, k_base_inline111__ssa_v3, i__idx_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp new file mode 100644 index 0000000000..9532317302 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp @@ -0,0 +1,366 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: residual_rms_cast_3 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void residual_rms_cast_3( + __gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int64_t v6, + int64_t v7 +) { + SaturationMode v8 = SaturationMode::OFF; + RoundMode v9 = RoundMode::CAST_ROUND; + const int64_t v10 = 256; + const int64_t v11 = 2; + const int64_t v12 = 4; + const int64_t v13 = 1; + const int64_t v14 = 5120; + const int64_t v15 = 16; + const int64_t v16 = 0; + const int64_t v17 = 51200; + const int64_t v18 = 34816; + const int64_t v19 = 33792; + const int64_t v20 = 17408; + const int64_t v21 = 1024; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + for (size_t v22 = (size_t)v16; v22 < ((size_t)v12); v22 += (size_t)v11) { + int64_t v23 = (int64_t)((uint64_t)((int64_t)v22) * (uint64_t)v10); + int64_t v24 = (int64_t)((uint64_t)v6 + (uint64_t)v23); + int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v23 + (uint64_t)v10)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v26 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v27 = (uint64_t)v21; + TASSIGN(v26, v27); + pto::Shape<1, 1, 1, 16, 256> v28 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v30 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v32 = (uint64_t)v20; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v24 * v13), v33, v34 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v37 = (uint64_t)v19; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 1, 256> v38 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v39 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v40 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 + ); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v42 = (uint64_t)v18; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v47 = (uint64_t)v17; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v4 + (v16 + v16 * v14 + v25 * v13), v48, v49 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v10); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + pto::Shape<1, 1, 1, 1, 256> v53 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v54 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v55 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 + ); + TLOAD(v51, v55); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v56 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TADD(v56, v26, v31); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v58 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v58, v56, v36); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v60 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); + pipe_barrier(PIPE_V); + TCVT(v60, v58, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v62 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v63 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v64 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v24 * v13), v62, v63 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v64, v56); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v65 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v66 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v67 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v24 * v13), v65, v66 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v67, v60); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v68 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TADD(v68, v41, v46); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v70 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TCOLEXPANDMUL(v70, v68, v51); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v72 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v15, v10); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); + pipe_barrier(PIPE_V); + TCVT(v72, v70, v9, v8); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pto::Shape<1, 1, 1, 16, 256> v74 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v75 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v76 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v16 + v16 * v14 + v25 * v13), v74, v75 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + pipe_barrier(PIPE_MTE3); + TSTORE(v76, v68); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + pto::Shape<1, 1, 1, 16, 256> v77 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v78 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v79 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + TSTORE(v79, v72); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: mlp_norm_in_inline71__rv_v11 + __gm__ Tensor *mlp_norm_in_inline71__rv_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *mlp_norm_in_inline71__rv_v11 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline71__rv_v11_tensor->buffer.addr) + + mlp_norm_in_inline71__rv_v11_tensor->start_offset; + + // Unpack tensor: post_norm_partial_inline118__rv_v11 + __gm__ Tensor *post_norm_partial_inline118__rv_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *post_norm_partial_inline118__rv_v11 = + reinterpret_cast<__gm__ float *>(post_norm_partial_inline118__rv_v11_tensor->buffer.addr) + + post_norm_partial_inline118__rv_v11_tensor->start_offset; + + // Unpack tensor: attn_proj_fp32_inline220__ssa_v7 + __gm__ Tensor *attn_proj_fp32_inline220__ssa_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *attn_proj_fp32_inline220__ssa_v7 = + reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline220__ssa_v7_tensor->buffer.addr) + + attn_proj_fp32_inline220__ssa_v7_tensor->start_offset; + + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: post_rms_weight__ssa_v0 + __gm__ Tensor *post_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); + __gm__ float *post_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + + post_rms_weight__ssa_v0_tensor->start_offset; + + // Unpack scalar: k_base_inline111__ssa_v4 + union { + uint64_t u64; + int64_t val; + } k_base_inline111__ssa_v4_conv; + k_base_inline111__ssa_v4_conv.u64 = args[5]; + int64_t k_base_inline111__ssa_v4 = k_base_inline111__ssa_v4_conv.val; + + // Unpack scalar: i__idx_v0 + union { + uint64_t u64; + int64_t val; + } i__idx_v0_conv; + i__idx_v0_conv.u64 = args[6]; + int64_t i__idx_v0 = i__idx_v0_conv.val; + + // Forward to ptoas-generated function + residual_rms_cast_3( + mlp_norm_in_inline71__rv_v11, post_norm_partial_inline118__rv_v11, attn_proj_fp32_inline220__ssa_v7, + cur__iter_v6, post_rms_weight__ssa_v0, k_base_inline111__ssa_v4, i__idx_v0 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/rms_recip.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/rms_recip.cpp new file mode 100644 index 0000000000..a4a9666bcb --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/rms_recip.cpp @@ -0,0 +1,452 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: rms_recip +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void rms_recip(__gm__ float *v1, __gm__ float *v2) { + const float v3 = 9.99999997E-7f; + const float v4 = 1.95312503E-4f; + const int64_t v5 = 768; + const int64_t v6 = 512; + const int64_t v7 = 256; + const int64_t v8 = 4; + const int64_t v9 = 20; + const float v10 = 0.0f; + const int64_t v11 = 1; + const int64_t v12 = 5120; + const int64_t v13 = 16; + const int64_t v14 = 49344; + const int64_t v15 = 32960; + const int64_t v16 = 32832; + const int64_t v17 = 16448; + const int64_t v18 = 16384; + const int64_t v19 = 0; + const int64_t v20 = 131328; + const int64_t v21 = 114944; + const int64_t v22 = 98560; + const int64_t v23 = 82176; + const int64_t v24 = 65792; + const int64_t v25 = 49408; + const int64_t v26 = 32896; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v27 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v28 = (uint64_t)v26; + TASSIGN(v27, v28); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + TEXPANDS(v27, v10); + for (size_t v29 = (size_t)v19; v29 < ((size_t)v9); v29 += (size_t)v8) { + int64_t v30 = (int64_t)((uint64_t)((int64_t)v29) * (uint64_t)v7); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v32 = (uint64_t)v25; + TASSIGN(v31, v32); + pto::Shape<1, 1, 1, 16, 256> v33 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v34 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v35 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v19 + v19 * v12 + v30 * v11), v33, v34 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v31, v35); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v36 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v37 = (uint64_t)v24; + TASSIGN(v36, v37); + pto::Shape<1, 1, 1, 16, 256> v38 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v39 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v40 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v19 + v19 * v12 + (int64_t)((uint64_t)v30 + (uint64_t)v7) * v11), v38, v39 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v36, v40); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v41 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v42 = (uint64_t)v23; + TASSIGN(v41, v42); + pto::Shape<1, 1, 1, 16, 256> v43 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v44 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v45 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v19 + v19 * v12 + (int64_t)((uint64_t)v30 + (uint64_t)v6) * v11), v43, v44 + ); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + TLOAD(v41, v45); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v47 = (uint64_t)v22; + TASSIGN(v46, v47); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v19 + v19 * v12 + (int64_t)((uint64_t)v30 + (uint64_t)v5) * v11), v48, v49 + ); + TLOAD(v46, v50); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v52 = (uint64_t)v25; + TASSIGN(v51, v52); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TMUL(v51, v31, v31); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v54 = (uint64_t)v21; + TASSIGN(v53, v54); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v11); + uint64_t v56 = (uint64_t)v20; + TASSIGN(v55, v56); + pipe_barrier(PIPE_V); + TROWSUM(v55, v51, v53); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v57 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v58 = (uint64_t)v20; + TASSIGN(v57, v58); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v59 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v60 = (uint64_t)v25; + TASSIGN(v59, v60); + pipe_barrier(PIPE_V); + TADD(v59, v27, v57); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v61 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v62 = (uint64_t)v24; + TASSIGN(v61, v62); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TMUL(v61, v36, v36); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v64 = (uint64_t)v19; + TASSIGN(v63, v64); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v11); + uint64_t v66 = (uint64_t)v18; + TASSIGN(v65, v66); + pipe_barrier(PIPE_V); + TROWSUM(v65, v61, v63); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v68 = (uint64_t)v18; + TASSIGN(v67, v68); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v69 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v70 = (uint64_t)v24; + TASSIGN(v69, v70); + pipe_barrier(PIPE_V); + TADD(v69, v59, v67); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v71 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v72 = (uint64_t)v23; + TASSIGN(v71, v72); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TMUL(v71, v41, v41); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v73 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v74 = (uint64_t)v17; + TASSIGN(v73, v74); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v75 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v11); + uint64_t v76 = (uint64_t)v16; + TASSIGN(v75, v76); + pipe_barrier(PIPE_V); + TROWSUM(v75, v71, v73); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v77 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v78 = (uint64_t)v16; + TASSIGN(v77, v78); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v79 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v80 = (uint64_t)v23; + TASSIGN(v79, v80); + pipe_barrier(PIPE_V); + TADD(v79, v69, v77); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v81 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v82 = (uint64_t)v22; + TASSIGN(v81, v82); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TMUL(v81, v46, v46); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v83 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v7); + uint64_t v84 = (uint64_t)v15; + TASSIGN(v83, v84); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v85 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v11); + uint64_t v86 = (uint64_t)v14; + TASSIGN(v85, v86); + pipe_barrier(PIPE_V); + TROWSUM(v85, v81, v83); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v87 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v88 = (uint64_t)v14; + TASSIGN(v87, v88); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v89 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v90 = (uint64_t)v26; + TASSIGN(v89, v90); + pipe_barrier(PIPE_V); + TADD(v89, v79, v87); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + } + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v91 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v92 = (uint64_t)v25; + TASSIGN(v91, v92); + pipe_barrier(PIPE_V); + TMULS(v91, v27, v4); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v93 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v94 = (uint64_t)v25; + TASSIGN(v93, v94); + pipe_barrier(PIPE_V); + TADDS(v93, v91, v3); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v95 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v96 = (uint64_t)v25; + TASSIGN(v95, v96); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v97 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v98 = (uint64_t)v25; + TASSIGN(v97, v98); + pipe_barrier(PIPE_V); + TSQRT(v97, v95); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v99 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v100 = (uint64_t)v25; + TASSIGN(v99, v100); + Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v101 = Tile< + TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v11, v13); + uint64_t v102 = (uint64_t)v24; + TASSIGN(v101, v102); + pipe_barrier(PIPE_V); + TRECIP(v101, v99); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v103 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v13, v11); + uint64_t v104 = (uint64_t)v24; + TASSIGN(v103, v104); + pto::Shape<1, 1, 1, 16, 1> v105 = pto::Shape<1, 1, 1, 16, 1>(); + pto::Stride<16, 16, 16, 1, 16> v106 = pto::Stride<16, 16, 16, 1, 16>(); + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v107 = + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN>( + v2 + (v19 + v19 * v11 + v19 * v13), v105, v106 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v107, v103); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: cur__iter_v6 + __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *cur__iter_v6 = + reinterpret_cast<__gm__ float *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + + // Unpack tensor: inv_rms_states_inline176__ssa_v0 + __gm__ Tensor *inv_rms_states_inline176__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *inv_rms_states_inline176__ssa_v0 = + reinterpret_cast<__gm__ float *>(inv_rms_states_inline176__ssa_v0_tensor->buffer.addr) + + inv_rms_states_inline176__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + rms_recip(cur__iter_v6, inv_rms_states_inline176__ssa_v0); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/silu.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/silu.cpp new file mode 100644 index 0000000000..327ce0009c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/silu.cpp @@ -0,0 +1,423 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: silu +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void silu(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ float *v4, int64_t v5) { + SaturationMode v6 = SaturationMode::OFF; + RoundMode v7 = RoundMode::CAST_ROUND; + const float v8 = 1.0f; + const int64_t v9 = 256; + const int64_t v10 = 2; + const int64_t v11 = 4; + const int64_t v12 = 17408; + const int64_t v13 = 1; + const int64_t v14 = 16; + const int64_t v15 = 49152; + const int64_t v16 = 32768; + const int64_t v17 = 16384; + const int64_t v18 = 0; + const int64_t v19 = 114752; + const int64_t v20 = 98368; + const int64_t v21 = 81984; + const int64_t v22 = 65600; + const int64_t v23 = 65536; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v24 = Tile< + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v13); + uint64_t v25 = (uint64_t)v23; + TASSIGN(v24, v25); + pto::Shape<1, 1, 1, 16, 1> v26 = pto::Shape<1, 1, 1, 16, 1>(); + pto::Stride<16, 16, 16, 1, 16> v27 = pto::Stride<16, 16, 16, 1, 16>(); + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v28 = + GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN>( + v1 + (v18 + v18 * v13 + v18 * v14), v26, v27 + ); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v24, v28); + for (size_t v29 = (size_t)v18; v29 < ((size_t)v11); v29 += (size_t)v10) { + int64_t v30 = (int64_t)((uint64_t)((int64_t)v29) * (uint64_t)v9); + int64_t v31 = (int64_t)((uint64_t)v5 + (uint64_t)v30); + int64_t v32 = (int64_t)((uint64_t)v5 + (uint64_t)((int64_t)(uint64_t)v30 + (uint64_t)v9)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v33 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v34 = (uint64_t)v22; + TASSIGN(v33, v34); + pto::Shape<1, 1, 1, 16, 256> v35 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v36 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v37 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v12 + v31 * v13), v35, v36 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + TLOAD(v33, v37); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v38 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v39 = (uint64_t)v21; + TASSIGN(v38, v39); + pto::Shape<1, 1, 1, 16, 256> v40 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v41 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v42 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v4 + (v18 + v18 * v12 + v31 * v13), v40, v41 + ); + TLOAD(v38, v42); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v43 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v44 = (uint64_t)v20; + TASSIGN(v43, v44); + pto::Shape<1, 1, 1, 16, 256> v45 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v46 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v47 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v12 + v32 * v13), v45, v46 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v43, v47); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v48 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v49 = (uint64_t)v19; + TASSIGN(v48, v49); + pto::Shape<1, 1, 1, 16, 256> v50 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v51 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v52 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( + v4 + (v18 + v18 * v12 + v32 * v13), v50, v51 + ); + TLOAD(v48, v52); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TROWEXPANDMUL(v53, v33, v24); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v55 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v56 = (uint64_t)v21; + TASSIGN(v55, v56); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TROWEXPANDMUL(v55, v38, v24); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v57 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v58 = (uint64_t)v18; + TASSIGN(v57, v58); + pipe_barrier(PIPE_V); + TNEG(v57, v53); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v59 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v60 = (uint64_t)v18; + TASSIGN(v59, v60); + pipe_barrier(PIPE_V); + TEXP(v59, v57); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v61 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v62 = (uint64_t)v18; + TASSIGN(v61, v62); + pipe_barrier(PIPE_V); + TADDS(v61, v59, v8); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v63 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v64 = (uint64_t)v17; + TASSIGN(v63, v64); + pipe_barrier(PIPE_V); + TRECIP(v63, v61); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v65 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v66 = (uint64_t)v22; + TASSIGN(v65, v66); + pipe_barrier(PIPE_V); + TMUL(v65, v53, v63); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v67 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v68 = (uint64_t)v22; + TASSIGN(v67, v68); + pipe_barrier(PIPE_V); + TMUL(v67, v65, v55); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v69 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v70 = (uint64_t)v22; + TASSIGN(v69, v70); + pipe_barrier(PIPE_V); + TCVT(v69, v67, v7, v6); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v71 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v72 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v73 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, + pto::Layout::ND>(v2 + (v18 + v18 * v12 + v31 * v13), v71, v72); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v73, v69); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v74 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v75 = (uint64_t)v20; + TASSIGN(v74, v75); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TROWEXPANDMUL(v74, v43, v24); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v76 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v77 = (uint64_t)v19; + TASSIGN(v76, v77); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + TROWEXPANDMUL(v76, v48, v24); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v78 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v79 = (uint64_t)v16; + TASSIGN(v78, v79); + pipe_barrier(PIPE_V); + TNEG(v78, v74); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v80 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v81 = (uint64_t)v16; + TASSIGN(v80, v81); + pipe_barrier(PIPE_V); + TEXP(v80, v78); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v82 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v83 = (uint64_t)v16; + TASSIGN(v82, v83); + pipe_barrier(PIPE_V); + TADDS(v82, v80, v8); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v84 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v85 = (uint64_t)v15; + TASSIGN(v84, v85); + pipe_barrier(PIPE_V); + TRECIP(v84, v82); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v86 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v87 = (uint64_t)v20; + TASSIGN(v86, v87); + pipe_barrier(PIPE_V); + TMUL(v86, v74, v84); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v88 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v89 = (uint64_t)v20; + TASSIGN(v88, v89); + pipe_barrier(PIPE_V); + TMUL(v88, v86, v76); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v90 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v9); + uint64_t v91 = (uint64_t)v20; + TASSIGN(v90, v91); + pipe_barrier(PIPE_V); + TCVT(v90, v88, v7, v6); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v92 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<278528, 278528, 278528, 17408, 1> v93 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> + v94 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, + pto::Layout::ND>(v2 + (v18 + v18 * v12 + v32 * v13), v92, v93); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pipe_barrier(PIPE_MTE3); + TSTORE(v94, v90); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Unpack tensor: inv_rms_tile_inline126__ssa_v1 + __gm__ Tensor *inv_rms_tile_inline126__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *inv_rms_tile_inline126__ssa_v1 = + reinterpret_cast<__gm__ float *>(inv_rms_tile_inline126__ssa_v1_tensor->buffer.addr) + + inv_rms_tile_inline126__ssa_v1_tensor->start_offset; + + // Unpack tensor: mlp_tile_inline149__iter_v1 + __gm__ Tensor *mlp_tile_inline149__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ bfloat16_t *mlp_tile_inline149__iter_v1 = + reinterpret_cast<__gm__ bfloat16_t *>(mlp_tile_inline149__iter_v1_tensor->buffer.addr) + + mlp_tile_inline149__iter_v1_tensor->start_offset; + + // Unpack tensor: gate_acc_all_inline203__rv_v10 + __gm__ Tensor *gate_acc_all_inline203__rv_v10_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *gate_acc_all_inline203__rv_v10 = + reinterpret_cast<__gm__ float *>(gate_acc_all_inline203__rv_v10_tensor->buffer.addr) + + gate_acc_all_inline203__rv_v10_tensor->start_offset; + + // Unpack tensor: up_acc_all_inline303__rv_v10 + __gm__ Tensor *up_acc_all_inline303__rv_v10_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + __gm__ float *up_acc_all_inline303__rv_v10 = + reinterpret_cast<__gm__ float *>(up_acc_all_inline303__rv_v10_tensor->buffer.addr) + + up_acc_all_inline303__rv_v10_tensor->start_offset; + + // Unpack scalar: n0_inline122__ssa_v7 + union { + uint64_t u64; + int64_t val; + } n0_inline122__ssa_v7_conv; + n0_inline122__ssa_v7_conv.u64 = args[4]; + int64_t n0_inline122__ssa_v7 = n0_inline122__ssa_v7_conv.val; + + // Forward to ptoas-generated function + silu( + inv_rms_tile_inline126__ssa_v1, mlp_tile_inline149__iter_v1, gate_acc_all_inline203__rv_v10, + up_acc_all_inline303__rv_v10, n0_inline122__ssa_v7 + ); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp new file mode 100644 index 0000000000..268481d684 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp @@ -0,0 +1,244 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Kernel Function: x_gamma0 +// Generated by PyPTO IR Compiler (PTO backend) + +#include + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#if defined(__CPU_SIM) +#define __aicore__ +#else +#define __aicore__ [aicore] +#endif +#endif + +#include +#include "tensor.h" +#include "intrinsic.h" + +using namespace pto; + +// --- ptoas-generated code --- + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +static __aicore__ void x_gamma0(__gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ float *v3, int32_t v4, int32_t v5) { + SaturationMode v6 = SaturationMode::OFF; + RoundMode v7 = RoundMode::CAST_ROUND; + const int64_t v8 = 256; + const int64_t v9 = 2; + const int64_t v10 = 4; + const int64_t v11 = 1024; + const int64_t v12 = 1; + const int64_t v13 = 5120; + const int64_t v14 = 16; + const int64_t v15 = 33792; + const int64_t v16 = 17408; + const int64_t v17 = 16384; + const int64_t v18 = 0; + using T = float; + +#if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + int64_t v19 = (int64_t)((uint64_t)((int64_t)v4) * (uint64_t)v11); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + for (size_t v20 = (size_t)v18; v20 < ((size_t)v10); v20 += (size_t)v9) { + int64_t v21 = (int64_t)((uint64_t)((int64_t)v20) * (uint64_t)v8); + int64_t v22 = (int64_t)((uint64_t)v19 + (uint64_t)v21); + int64_t v23 = (int64_t)((uint64_t)v19 + (uint64_t)((int64_t)(uint64_t)v21 + (uint64_t)v8)); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v24 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v25 = (uint64_t)v18; + TASSIGN(v24, v25); + pto::Shape<1, 1, 1, 16, 256> v26 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v27 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v28 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v18 + v18 * v13 + v22 * v12), v26, v27 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + TLOAD(v24, v28); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v29 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v8); + uint64_t v30 = (uint64_t)v17; + TASSIGN(v29, v30); + pto::Shape<1, 1, 1, 1, 256> v31 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v32 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v33 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v13 + v22 * v12), v31, v32 + ); + TLOAD(v29, v33); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v34 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v35 = (uint64_t)v16; + TASSIGN(v34, v35); + pto::Shape<1, 1, 1, 16, 256> v36 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v37 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v38 = GlobalTensor< + float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v2 + (v18 + v18 * v13 + v23 * v12), v36, v37 + ); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v34, v38); + Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v39 = Tile< + TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v12, v8); + uint64_t v40 = (uint64_t)v15; + TASSIGN(v39, v40); + pto::Shape<1, 1, 1, 1, 256> v41 = pto::Shape<1, 1, 1, 1, 256>(); + pto::Stride<5120, 5120, 5120, 5120, 1> v42 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v43 = + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( + v3 + (v18 + v18 * v13 + v23 * v12), v41, v42 + ); + TLOAD(v39, v43); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v44 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v45 = (uint64_t)v18; + TASSIGN(v44, v45); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + TCOLEXPANDMUL(v44, v24, v29); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v46 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v47 = (uint64_t)v18; + TASSIGN(v46, v47); + pipe_barrier(PIPE_V); + TCVT(v46, v44, v7, v6); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pto::Shape<1, 1, 1, 16, 256> v48 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v49 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v50 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v18 + v18 * v13 + v22 * v12), v48, v49 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v50, v46); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v51 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v52 = (uint64_t)v16; + TASSIGN(v51, v52); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TCOLEXPANDMUL(v51, v34, v39); + Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null> + v53 = Tile< + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + CompactMode::Null>(v14, v8); + uint64_t v54 = (uint64_t)v16; + TASSIGN(v53, v54); + pipe_barrier(PIPE_V); + TCVT(v53, v51, v7, v6); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pto::Shape<1, 1, 1, 16, 256> v55 = pto::Shape<1, 1, 1, 16, 256>(); + pto::Stride<81920, 81920, 81920, 5120, 1> v56 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> + v57 = GlobalTensor< + bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( + v1 + (v18 + v18 * v13 + v23 * v12), v55, v56 + ); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + pipe_barrier(PIPE_MTE3); + TSTORE(v57, v53); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + } + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); +#endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} +// --- Kernel entry point --- +extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { + // Read logical SPMD block identity from runtime dispatch payload + int32_t __pypto_spmd_block_idx = get_block_idx(args); + int32_t __pypto_spmd_block_num = get_block_num(args); + + // Unpack tensor: normed__ssa_v0 + __gm__ Tensor *normed__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ bfloat16_t *normed__ssa_v0 = + reinterpret_cast<__gm__ bfloat16_t *>(normed__ssa_v0_tensor->buffer.addr) + normed__ssa_v0_tensor->start_offset; + + // Unpack tensor: cur__rv_v2 + __gm__ Tensor *cur__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ float *cur__rv_v2 = + reinterpret_cast<__gm__ float *>(cur__rv_v2_tensor->buffer.addr) + cur__rv_v2_tensor->start_offset; + + // Unpack tensor: input_rms_weight__ssa_v0 + __gm__ Tensor *input_rms_weight__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ float *input_rms_weight__ssa_v0 = + reinterpret_cast<__gm__ float *>(input_rms_weight__ssa_v0_tensor->buffer.addr) + + input_rms_weight__ssa_v0_tensor->start_offset; + + // Forward to ptoas-generated function + x_gamma0(normed__ssa_v0, cur__rv_v2, input_rms_weight__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num); +} diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp new file mode 100644 index 0000000000..db08d899aa --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp @@ -0,0 +1,1980 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +// Orchestration Function: decode_fwd_layers +// Generated by PyPTO IR Compiler + +#include "runtime.h" +#include + +#include +#include +#include + +#include "pto_orchestration_api.h" + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 20, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + // External tensors + const Tensor &ext_hidden_states = orch_args.tensor(0).ref(); + const Tensor &ext_input_rms_weight = orch_args.tensor(1).ref(); + const Tensor &ext_wq = orch_args.tensor(2).ref(); + const Tensor &ext_wk = orch_args.tensor(3).ref(); + const Tensor &ext_wv = orch_args.tensor(4).ref(); + const Tensor &ext_q_norm_weight = orch_args.tensor(5).ref(); + const Tensor &ext_k_norm_weight = orch_args.tensor(6).ref(); + const Tensor &ext_seq_lens = orch_args.tensor(7).ref(); + const Tensor &ext_block_table = orch_args.tensor(8).ref(); + const Tensor &ext_slot_mapping = orch_args.tensor(9).ref(); + const Tensor &ext_rope_cos = orch_args.tensor(10).ref(); + const Tensor &ext_rope_sin = orch_args.tensor(11).ref(); + const Tensor &ext_k_cache = orch_args.tensor(12).ref(); + const Tensor &ext_v_cache = orch_args.tensor(13).ref(); + const Tensor &ext_wo = orch_args.tensor(14).ref(); + const Tensor &ext_w_gate = orch_args.tensor(15).ref(); + const Tensor &ext_w_up = orch_args.tensor(16).ref(); + const Tensor &ext_w_down = orch_args.tensor(17).ref(); + const Tensor &ext_post_rms_weight = orch_args.tensor(18).ref(); + const Tensor &ext_out = orch_args.tensor(19).ref(); + + // Dynamic-dim symbols (extent of the declaring argument) + int64_t BLOCK_TABLE_FLAT_DYN = (int64_t)orch_args.tensor(8).ref().shapes[0]; + int64_t KV_CACHE_ROWS_DYN = (int64_t)orch_args.tensor(12).ref().shapes[0]; + + PTO2_SCOPE() { + uint32_t pa_metadata_ci_shapes[1] = {27840}; + TensorCreateInfo pa_metadata_ci(pa_metadata_ci_shapes, 1, DataType::UINT8); + uint32_t pa_workspace_ci_shapes[1] = {66132544}; + TensorCreateInfo pa_workspace_ci(pa_workspace_ci_shapes, 1, DataType::UINT8); + uint32_t cur_ci_shapes[2] = {16, 5120}; + TensorCreateInfo cur_ci(cur_ci_shapes, 2, DataType::FLOAT32); + uint32_t normed_ci_shapes[2] = {16, 5120}; + TensorCreateInfo normed_ci(normed_ci_shapes, 2, DataType::BFLOAT16); + TaskOutputTensors alloc_0 = alloc_tensors(pa_metadata_ci, pa_workspace_ci, cur_ci, normed_ci); + const Tensor &pa_metadata = alloc_0.get_ref(0); + const Tensor &pa_workspace = alloc_0.get_ref(1); + const Tensor &cur = alloc_0.get_ref(2); + const Tensor &normed = alloc_0.get_ref(3); + int64_t pa_num_layers = 40; + int64_t pa_num_pages = (KV_CACHE_ROWS_DYN / (pa_num_layers * 1024)); + int64_t pa_max_blocks = (BLOCK_TABLE_FLAT_DYN / 16); + int32_t pa_num_pages_i32 = static_cast(pa_num_pages); + int32_t pa_max_blocks_i32 = static_cast(pa_max_blocks); + + // Spmd pa_tiling: paged_attention_tiling_cce + L0TaskArgs params_t0; + params_t0.add_input(ext_seq_lens); + params_t0.add_output(pa_metadata); + params_t0.add_scalar(pa_max_blocks_i32); + params_t0.add_scalar(pa_num_pages_i32); + params_t0.launch_spec.set_block_num(1); + params_t0.set_allow_early_resolve(true); + TaskOutputTensors task_0_outs = rt_submit_aiv_task(0, params_t0); + PTO2TaskId tiling_tid_inline0 = task_0_outs.task_id(); + PTO2TaskId pa_tiling_tid = tiling_tid_inline0; + PTO2TaskId prev_out_tid[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + prev_out_tid[__init_i] = PTO2TaskId::invalid(); + + // Phase-fence barrier 0: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_0; + TaskOutputTensors phase_fence_barrier_0_outs = rt_submit_dummy_task(params_phase_fence_barrier_0); + PTO2TaskId t = phase_fence_barrier_0_outs.task_id(); + prev_out_tid[0] = t; + for (int64_t cb0 = 0; cb0 < 16; cb0 += 16) { + PTO2_SCOPE() { + // Task 1: copy_hidden + L0TaskArgs params_t1; + params_t1.add_output(cur); + params_t1.add_input(ext_hidden_states); + params_t1.add_scalar(cb0); + TaskOutputTensors task_1_outs = rt_submit_aiv_task(1, params_t1); + PTO2TaskId ch_tid = task_1_outs.task_id(); + prev_out_tid[0] = ch_tid; + } + } + PTO2TaskId prev_normed_tid[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + prev_normed_tid[__init_i] = PTO2TaskId::invalid(); + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + PTO2TaskId _submit_deps_buf[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v3 = prev_out_tid[0]; + _submit_deps_buf[0] = t__tmp_v3; + + // Spmd x_gamma0_spmd: x_gamma0 + L0TaskArgs params_t2; + params_t2.add_output(normed); + params_t2.add_input(cur); + params_t2.add_input(ext_input_rms_weight); + params_t2.launch_spec.set_block_num(5); + params_t2.set_allow_early_resolve(true); + PTO2TaskId params_t2_deps[1]; + uint32_t params_t2_deps_count = 0; + if (_submit_deps_buf[0].is_valid()) params_t2_deps[params_t2_deps_count++] = _submit_deps_buf[0]; + params_t2.set_dependencies(params_t2_deps, params_t2_deps_count); + TaskOutputTensors task_2_outs = rt_submit_aiv_task(2, params_t2); + PTO2TaskId xgamma_tid = task_2_outs.task_id(); + prev_normed_tid[0] = xgamma_tid; + } + Tensor cur__rv_v7 = cur; + Tensor normed__rv_v5 = normed; + for (int64_t i = 0; i < 40; i += 1) { + PTO2_SCOPE() { + uint32_t next_hidden_ci_shapes[2] = {16, 5120}; + TensorCreateInfo next_hidden_ci(next_hidden_ci_shapes, 2, DataType::FLOAT32); + uint32_t next_normed_ci_shapes[2] = {16, 5120}; + TensorCreateInfo next_normed_ci(next_normed_ci_shapes, 2, DataType::BFLOAT16); + uint32_t inv_rms_states_inline176_ci_shapes[2] = {16, 1}; + TensorCreateInfo inv_rms_states_inline176_ci(inv_rms_states_inline176_ci_shapes, 2, DataType::FLOAT32); + uint32_t q_proj_inline139_ci_shapes[2] = {16, 5120}; + TensorCreateInfo q_proj_inline139_ci(q_proj_inline139_ci_shapes, 2, DataType::FLOAT32); + uint32_t k_proj_inline135_ci_shapes[2] = {16, 1024}; + TensorCreateInfo k_proj_inline135_ci(k_proj_inline135_ci_shapes, 2, DataType::FLOAT32); + uint32_t v_proj_inline255_ci_shapes[2] = {16, 1024}; + TensorCreateInfo v_proj_inline255_ci(v_proj_inline255_ci_shapes, 2, DataType::FLOAT32); + uint32_t q_tnd_flat_inline127_ci_shapes[2] = {640, 128}; + TensorCreateInfo q_tnd_flat_inline127_ci(q_tnd_flat_inline127_ci_shapes, 2, DataType::BFLOAT16); + uint32_t attn_out_inline282_ci_shapes[2] = {16, 5120}; + TensorCreateInfo attn_out_inline282_ci(attn_out_inline282_ci_shapes, 2, DataType::BFLOAT16); + TaskOutputTensors alloc_1 = alloc_tensors( + next_hidden_ci, next_normed_ci, inv_rms_states_inline176_ci, q_proj_inline139_ci, + k_proj_inline135_ci, v_proj_inline255_ci, q_tnd_flat_inline127_ci, attn_out_inline282_ci + ); + const Tensor &next_hidden = alloc_1.get_ref(0); + const Tensor &next_normed = alloc_1.get_ref(1); + const Tensor &inv_rms_states_inline176 = alloc_1.get_ref(2); + const Tensor &q_proj_inline139 = alloc_1.get_ref(3); + const Tensor &k_proj_inline135 = alloc_1.get_ref(4); + const Tensor &v_proj_inline255 = alloc_1.get_ref(5); + const Tensor &q_tnd_flat_inline127 = alloc_1.get_ref(6); + const Tensor &attn_out_inline282 = alloc_1.get_ref(7); + int64_t next_gamma_idx = std::min((i + 1), 39); + int64_t layer_hidden_base_inline151 = (static_cast(i) * 5120); + int64_t layer_inter_base_inline107 = (static_cast(i) * 17408); + int64_t num_layers_actual_inline152 = 40; + int64_t t__tmp_v6 = (int64_t)orch_args.tensor(12).ref().shapes[0]; + int64_t layer_cache_rows_inline128 = (t__tmp_v6 / num_layers_actual_inline152); + int64_t layer_cache_base_inline193 = (static_cast(i) * layer_cache_rows_inline128); + uint32_t q_norm_w_inline124_offsets[2] = {static_cast(i), 0}; + uint32_t q_norm_w_inline124_shapes[2] = { + (q_norm_w_inline124_offsets[0] >= ext_q_norm_weight.shapes[0] ? + 0u : + std::min(1, ext_q_norm_weight.shapes[0] - q_norm_w_inline124_offsets[0])), + (q_norm_w_inline124_offsets[1] >= ext_q_norm_weight.shapes[1] ? + 0u : + std::min(128, ext_q_norm_weight.shapes[1] - q_norm_w_inline124_offsets[1])) + }; + Tensor q_norm_w_inline124 = + ext_q_norm_weight.view(q_norm_w_inline124_shapes, q_norm_w_inline124_offsets); + uint32_t k_norm_w_inline114_offsets[2] = {static_cast(i), 0}; + uint32_t k_norm_w_inline114_shapes[2] = { + (k_norm_w_inline114_offsets[0] >= ext_k_norm_weight.shapes[0] ? + 0u : + std::min(1, ext_k_norm_weight.shapes[0] - k_norm_w_inline114_offsets[0])), + (k_norm_w_inline114_offsets[1] >= ext_k_norm_weight.shapes[1] ? + 0u : + std::min(128, ext_k_norm_weight.shapes[1] - k_norm_w_inline114_offsets[1])) + }; + Tensor k_norm_w_inline114 = + ext_k_norm_weight.view(k_norm_w_inline114_shapes, k_norm_w_inline114_offsets); + PTO2TaskId down_tids_inline156[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + down_tids_inline156[__init_i] = PTO2TaskId::invalid(); + uint32_t down_acc_all_inline168_ci_shapes[2] = {16, 5120}; + TensorCreateInfo down_acc_all_inline168_ci(down_acc_all_inline168_ci_shapes, 2, DataType::FLOAT32); + uint32_t gate_acc_all_inline203_ci_shapes[2] = {16, 17408}; + TensorCreateInfo gate_acc_all_inline203_ci(gate_acc_all_inline203_ci_shapes, 2, DataType::FLOAT32); + uint32_t up_acc_all_inline303_ci_shapes[2] = {16, 17408}; + TensorCreateInfo up_acc_all_inline303_ci(up_acc_all_inline303_ci_shapes, 2, DataType::FLOAT32); + uint32_t attn_proj_fp32_inline220_ci_shapes[2] = {16, 5120}; + TensorCreateInfo attn_proj_fp32_inline220_ci(attn_proj_fp32_inline220_ci_shapes, 2, DataType::FLOAT32); + uint32_t post_norm_partial_inline118_ci_shapes[2] = {16, 5120}; + TensorCreateInfo post_norm_partial_inline118_ci( + post_norm_partial_inline118_ci_shapes, 2, DataType::FLOAT32 + ); + uint32_t mlp_norm_in_inline71_ci_shapes[2] = {16, 5120}; + TensorCreateInfo mlp_norm_in_inline71_ci(mlp_norm_in_inline71_ci_shapes, 2, DataType::BFLOAT16); + uint32_t inv_rms_tile_inline126_ci_shapes[2] = {16, 1}; + TensorCreateInfo inv_rms_tile_inline126_ci(inv_rms_tile_inline126_ci_shapes, 2, DataType::FLOAT32); + uint32_t mlp_tile_inline149_ci_shapes[2] = {16, 17408}; + TensorCreateInfo mlp_tile_inline149_ci(mlp_tile_inline149_ci_shapes, 2, DataType::BFLOAT16); + TaskOutputTensors alloc_2 = alloc_tensors( + down_acc_all_inline168_ci, gate_acc_all_inline203_ci, up_acc_all_inline303_ci, + attn_proj_fp32_inline220_ci, post_norm_partial_inline118_ci, mlp_norm_in_inline71_ci, + inv_rms_tile_inline126_ci, mlp_tile_inline149_ci + ); + const Tensor &down_acc_all_inline168 = alloc_2.get_ref(0); + const Tensor &gate_acc_all_inline203 = alloc_2.get_ref(1); + const Tensor &up_acc_all_inline303 = alloc_2.get_ref(2); + const Tensor &attn_proj_fp32_inline220 = alloc_2.get_ref(3); + const Tensor &post_norm_partial_inline118 = alloc_2.get_ref(4); + const Tensor &mlp_norm_in_inline71 = alloc_2.get_ref(5); + const Tensor &inv_rms_tile_inline126 = alloc_2.get_ref(6); + const Tensor &mlp_tile_inline149 = alloc_2.get_ref(7); + PTO2_SCOPE(PTO2ScopeMode::MANUAL) { + // Phase-fence barrier 1: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_1; + TaskOutputTensors phase_fence_barrier_1_outs = rt_submit_dummy_task(params_phase_fence_barrier_1); + PTO2TaskId seed_dummy_inline49 = phase_fence_barrier_1_outs.task_id(); + PTO2TaskId prev_normed_seed_deps_inline120[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + prev_normed_seed_deps_inline120[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v7 = prev_normed_tid[0]; + prev_normed_seed_deps_inline120[0] = t__tmp_v7; + prev_normed_seed_deps_inline120[1] = seed_dummy_inline49; + + // Task 3: attn_out_seed + L0TaskArgs params_t3; + params_t3.add_input(attn_out_inline282); + params_t3.set_allow_early_resolve(true); + TaskOutputTensors task_3_outs = rt_submit_aiv_task(3, params_t3); + PTO2TaskId attn_out_seed_tid_inline116 = task_3_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline42[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline42[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v8 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline42[0] = t__tmp_v8; + PTO2TaskId t__tmp_v9 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline42[1] = t__tmp_v9; + + // Task 4: rms_recip + L0TaskArgs params_t4; + params_t4.add_input(cur__rv_v7); + params_t4.add_inout(inv_rms_states_inline176); + PTO2TaskId params_t4_deps[2]; + uint32_t params_t4_deps_count = 0; + if (_submit_deps_buf_inline42[0].is_valid()) + params_t4_deps[params_t4_deps_count++] = _submit_deps_buf_inline42[0]; + if (_submit_deps_buf_inline42[1].is_valid()) + params_t4_deps[params_t4_deps_count++] = _submit_deps_buf_inline42[1]; + params_t4.set_dependencies(params_t4_deps, params_t4_deps_count); + params_t4.set_allow_early_resolve(true); + TaskOutputTensors task_4_outs = rt_submit_aiv_task(4, params_t4); + PTO2TaskId rms_tid_inline148 = task_4_outs.task_id(); + + // Task 5: q_seed + L0TaskArgs params_t5; + params_t5.add_inout(q_proj_inline139); + params_t5.set_allow_early_resolve(true); + TaskOutputTensors task_5_outs = rt_submit_aiv_task(5, params_t5); + PTO2TaskId q_seed_tid_inline162 = task_5_outs.task_id(); + PTO2TaskId prev_normed_q_deps_inline105[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + prev_normed_q_deps_inline105[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v17 = prev_normed_tid[0]; + prev_normed_q_deps_inline105[0] = t__tmp_v17; + prev_normed_q_deps_inline105[1] = q_seed_tid_inline162; + PTO2TaskId _submit_deps_buf_inline182[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline182[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v18 = prev_normed_q_deps_inline105[0]; + _submit_deps_buf_inline182[0] = t__tmp_v18; + PTO2TaskId t__tmp_v19 = prev_normed_q_deps_inline105[1]; + _submit_deps_buf_inline182[1] = t__tmp_v19; + + // Spmd q_proj_spmd: q_proj + L0TaskArgs params_t6; + params_t6.add_inout(q_proj_inline139); + params_t6.add_input(normed__rv_v5); + params_t6.add_input(ext_wq); + params_t6.add_scalar(layer_hidden_base_inline151); + params_t6.launch_spec.set_block_num(50); + params_t6.set_allow_early_resolve(true); + PTO2TaskId params_t6_deps[2]; + uint32_t params_t6_deps_count = 0; + if (_submit_deps_buf_inline182[0].is_valid()) + params_t6_deps[params_t6_deps_count++] = _submit_deps_buf_inline182[0]; + if (_submit_deps_buf_inline182[1].is_valid()) + params_t6_deps[params_t6_deps_count++] = _submit_deps_buf_inline182[1]; + params_t6.set_dependencies(params_t6_deps, params_t6_deps_count); + TaskOutputTensors task_6_outs = rt_submit_aic_task(6, params_t6); + PTO2TaskId q_proj_tid_inline183 = task_6_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline261[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline261[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v24 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline261[0] = t__tmp_v24; + PTO2TaskId t__tmp_v25 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline261[1] = t__tmp_v25; + + // Task 7: kv_seed + L0TaskArgs params_t7; + params_t7.add_inout(k_proj_inline135); + params_t7.add_inout(v_proj_inline255); + PTO2TaskId params_t7_deps[2]; + uint32_t params_t7_deps_count = 0; + if (_submit_deps_buf_inline261[0].is_valid()) + params_t7_deps[params_t7_deps_count++] = _submit_deps_buf_inline261[0]; + if (_submit_deps_buf_inline261[1].is_valid()) + params_t7_deps[params_t7_deps_count++] = _submit_deps_buf_inline261[1]; + params_t7.set_dependencies(params_t7_deps, params_t7_deps_count); + TaskOutputTensors task_7_outs = rt_submit_aiv_task(7, params_t7); + PTO2TaskId kv_seed_tid_inline238 = task_7_outs.task_id(); + PTO2TaskId _submit_deps_buf_inline267[2]; + for (int64_t __init_i = 0; __init_i < 2; ++__init_i) + _submit_deps_buf_inline267[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v28 = prev_normed_seed_deps_inline120[0]; + _submit_deps_buf_inline267[0] = t__tmp_v28; + PTO2TaskId t__tmp_v29 = prev_normed_seed_deps_inline120[1]; + _submit_deps_buf_inline267[1] = t__tmp_v29; + + // Task 8: mlp_out_seed + L0TaskArgs params_t8; + params_t8.add_inout(down_acc_all_inline168); + params_t8.add_inout(gate_acc_all_inline203); + params_t8.add_inout(up_acc_all_inline303); + params_t8.add_inout(attn_proj_fp32_inline220); + PTO2TaskId params_t8_deps[2]; + uint32_t params_t8_deps_count = 0; + if (_submit_deps_buf_inline267[0].is_valid()) + params_t8_deps[params_t8_deps_count++] = _submit_deps_buf_inline267[0]; + if (_submit_deps_buf_inline267[1].is_valid()) + params_t8_deps[params_t8_deps_count++] = _submit_deps_buf_inline267[1]; + params_t8.set_dependencies(params_t8_deps, params_t8_deps_count); + params_t8.set_allow_early_resolve(true); + TaskOutputTensors task_8_outs = rt_submit_aiv_task(8, params_t8); + PTO2TaskId mlp_out_seed_tid_inline206 = task_8_outs.task_id(); + + // Spmd k_proj_spmd: k_proj + L0TaskArgs params_t9; + params_t9.add_inout(k_proj_inline135); + params_t9.add_input(normed__rv_v5); + params_t9.add_input(ext_wk); + params_t9.add_scalar(layer_hidden_base_inline151); + params_t9.launch_spec.set_block_num(10); + params_t9.set_allow_early_resolve(true); + PTO2TaskId params_t9_deps[1]; + uint32_t params_t9_deps_count = 0; + params_t9_deps[params_t9_deps_count++] = kv_seed_tid_inline238; + params_t9.set_dependencies(params_t9_deps, params_t9_deps_count); + TaskOutputTensors task_9_outs = rt_submit_aic_task(9, params_t9); + PTO2TaskId k_proj_tid_inline136 = task_9_outs.task_id(); + + // Spmd v_proj_spmd: v_proj + L0TaskArgs params_t10; + params_t10.add_inout(v_proj_inline255); + params_t10.add_input(normed__rv_v5); + params_t10.add_input(ext_wv); + params_t10.add_scalar(layer_hidden_base_inline151); + params_t10.launch_spec.set_block_num(10); + params_t10.set_allow_early_resolve(true); + PTO2TaskId params_t10_deps[1]; + uint32_t params_t10_deps_count = 0; + params_t10_deps[params_t10_deps_count++] = kv_seed_tid_inline238; + params_t10.set_dependencies(params_t10_deps, params_t10_deps_count); + TaskOutputTensors task_10_outs = rt_submit_aic_task(10, params_t10); + PTO2TaskId v_proj_tid_inline63 = task_10_outs.task_id(); + uint32_t q_tnd_inline191_shapes[3] = {16, 40, 128}; + Tensor q_tnd_inline191 = q_tnd_flat_inline127.reshape(q_tnd_inline191_shapes, 3); + uint32_t attn_out_tnd_inline79_shapes[3] = {16, 40, 128}; + Tensor attn_out_tnd_inline79 = attn_out_inline282.reshape(attn_out_tnd_inline79_shapes, 3); + int64_t attention_core_num_inline188 = 24; + + // Group paged_attention_rope_cce: MixedKernels (AIC + AIV lanes) + L0TaskArgs params_t11; + params_t11.add_inout(attn_out_tnd_inline79); + params_t11.add_inout(q_tnd_inline191); + params_t11.add_inout(ext_k_cache); + params_t11.add_inout(ext_v_cache); + params_t11.add_input(ext_block_table); + params_t11.add_inout(pa_workspace); + params_t11.add_inout(pa_metadata); + params_t11.add_input(q_proj_inline139); + params_t11.add_input(k_proj_inline135); + params_t11.add_input(v_proj_inline255); + params_t11.add_input(q_norm_w_inline124); + params_t11.add_input(k_norm_w_inline114); + params_t11.add_input(ext_rope_cos); + params_t11.add_input(ext_rope_sin); + params_t11.add_input(inv_rms_states_inline176); + params_t11.add_input(ext_slot_mapping); + params_t11.add_input(ext_seq_lens); + params_t11.add_scalar(layer_cache_base_inline193); + MixedKernels mixed_11 = {11, 12, 12}; + params_t11.launch_spec.set_block_num(attention_core_num_inline188); + params_t11.launch_spec.set_require_sync_start(true); + params_t11.set_allow_early_resolve(true); + PTO2TaskId params_t11_deps[7]; + uint32_t params_t11_deps_count = 0; + params_t11_deps[params_t11_deps_count++] = q_proj_tid_inline183; + params_t11_deps[params_t11_deps_count++] = k_proj_tid_inline136; + params_t11_deps[params_t11_deps_count++] = v_proj_tid_inline63; + params_t11_deps[params_t11_deps_count++] = rms_tid_inline148; + params_t11_deps[params_t11_deps_count++] = tiling_tid_inline0; + params_t11_deps[params_t11_deps_count++] = attn_out_seed_tid_inline116; + params_t11_deps[params_t11_deps_count++] = mlp_out_seed_tid_inline206; + params_t11.set_dependencies(params_t11_deps, params_t11_deps_count); + TaskOutputTensors task_11_outs = rt_submit_task(mixed_11, params_t11); + const Tensor &attn_out_tnd_inline79__ssa_v1 = attn_out_tnd_inline79; + PTO2TaskId attn_done_tid_inline78 = task_11_outs.task_id(); + uint32_t attn_out_inline282__ssa_v4_shapes[2] = {16, 5120}; + Tensor attn_out_inline282__ssa_v4 = + attn_out_tnd_inline79__ssa_v1.reshape(attn_out_inline282__ssa_v4_shapes, 2); + PTO2TaskId silu_tids_inline265[17]; + for (int64_t __init_i = 0; __init_i < 17; ++__init_i) + silu_tids_inline265[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId gate_tids_inline56[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + gate_tids_inline56[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId up_tids_inline310[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + up_tids_inline310[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId cast_tids_inline88[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + cast_tids_inline88[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId gate_late_tids_inline249[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + gate_late_tids_inline249[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId up_late_tids_inline69[5]; + for (int64_t __init_i = 0; __init_i < 5; ++__init_i) + up_late_tids_inline69[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId out_tids_inline271[50]; + for (int64_t __init_i = 0; __init_i < 50; ++__init_i) + out_tids_inline271[__init_i] = PTO2TaskId::invalid(); + + // Phase-fence barrier 2: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_2; + PTO2TaskId params_phase_fence_barrier_2_deps[1]; + uint32_t params_phase_fence_barrier_2_deps_count = 0; + params_phase_fence_barrier_2_deps[params_phase_fence_barrier_2_deps_count++] = + attn_done_tid_inline78; + params_phase_fence_barrier_2.set_dependencies( + params_phase_fence_barrier_2_deps, params_phase_fence_barrier_2_deps_count + ); + PTO2TaskId out_proj_dummy_inline257 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_2_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_2_outs = + rt_submit_dummy_task(params_phase_fence_barrier_2); + out_proj_dummy_inline257 = phase_fence_barrier_2_outs.task_id(); + } + int64_t N_OUT_DIRECT_inline61 = 26; + for (int64_t out_idx_inline74 = 0; out_idx_inline74 < N_OUT_DIRECT_inline61; + out_idx_inline74 += 1) { + int64_t n_out_proj_inline185 = (out_idx_inline74 / 5); + int64_t k_split_out_inline66 = (out_idx_inline74 % 5); + int64_t n_op_inline64 = (n_out_proj_inline185 * 512); + int64_t k_op_inline266 = (k_split_out_inline66 * 1024); + + // Task 12: out_proj + L0TaskArgs params_t12; + params_t12.add_input(attn_out_inline282__ssa_v4); + params_t12.add_input(ext_wo); + params_t12.add_inout(attn_proj_fp32_inline220); + params_t12.add_scalar(k_op_inline266); + params_t12.add_scalar(layer_hidden_base_inline151); + params_t12.add_scalar(n_op_inline64); + PTO2TaskId params_t12_deps[1]; + uint32_t params_t12_deps_count = 0; + if (out_proj_dummy_inline257.is_valid()) + params_t12_deps[params_t12_deps_count++] = out_proj_dummy_inline257; + params_t12.set_dependencies(params_t12_deps, params_t12_deps_count); + TaskOutputTensors task_12_outs = rt_submit_aic_task(13, params_t12); + PTO2TaskId out_tid_inline141 = task_12_outs.task_id(); + out_tids_inline271[out_idx_inline74] = out_tid_inline141; + } + + // Spmd out_proj_spmd: out_proj_0 + L0TaskArgs params_t13; + params_t13.add_input(attn_out_inline282__ssa_v4); + params_t13.add_input(ext_wo); + params_t13.add_inout(attn_proj_fp32_inline220); + params_t13.add_scalar(N_OUT_DIRECT_inline61); + params_t13.add_scalar(layer_hidden_base_inline151); + params_t13.launch_spec.set_block_num(24); + PTO2TaskId params_t13_deps[1]; + uint32_t params_t13_deps_count = 0; + params_t13_deps[params_t13_deps_count++] = attn_done_tid_inline78; + params_t13.set_dependencies(params_t13_deps, params_t13_deps_count); + TaskOutputTensors task_13_outs = rt_submit_aic_task(14, params_t13); + PTO2TaskId out_proj_direct_tid_inline70 = task_13_outs.task_id(); + out_tids_inline271[N_OUT_DIRECT_inline61] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 1)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 2)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 3)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 4)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 5)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 6)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 7)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 8)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 9)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 10)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 11)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 12)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 13)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 14)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 15)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 16)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 17)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 18)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 19)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 20)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 21)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 22)] = out_proj_direct_tid_inline70; + out_tids_inline271[(N_OUT_DIRECT_inline61 + 23)] = out_proj_direct_tid_inline70; + int64_t k_base_inline111 = 0; + int64_t n_split_base_inline163 = 0; + PTO2TaskId _submit_deps_buf_inline165[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v39 = out_tids_inline271[(n_split_base_inline163 * 5)]; + _submit_deps_buf_inline165[0] = t__tmp_v39; + PTO2TaskId t__tmp_v40 = out_tids_inline271[((n_split_base_inline163 * 5) + 1)]; + _submit_deps_buf_inline165[1] = t__tmp_v40; + PTO2TaskId t__tmp_v41 = out_tids_inline271[((n_split_base_inline163 * 5) + 2)]; + _submit_deps_buf_inline165[2] = t__tmp_v41; + PTO2TaskId t__tmp_v42 = out_tids_inline271[((n_split_base_inline163 * 5) + 3)]; + _submit_deps_buf_inline165[3] = t__tmp_v42; + PTO2TaskId t__tmp_v43 = out_tids_inline271[((n_split_base_inline163 * 5) + 4)]; + _submit_deps_buf_inline165[4] = t__tmp_v43; + PTO2TaskId t__tmp_v44 = out_tids_inline271[((n_split_base_inline163 * 5) + 5)]; + _submit_deps_buf_inline165[5] = t__tmp_v44; + PTO2TaskId t__tmp_v45 = out_tids_inline271[((n_split_base_inline163 * 5) + 6)]; + _submit_deps_buf_inline165[6] = t__tmp_v45; + PTO2TaskId t__tmp_v46 = out_tids_inline271[((n_split_base_inline163 * 5) + 7)]; + _submit_deps_buf_inline165[7] = t__tmp_v46; + PTO2TaskId t__tmp_v47 = out_tids_inline271[((n_split_base_inline163 * 5) + 8)]; + _submit_deps_buf_inline165[8] = t__tmp_v47; + PTO2TaskId t__tmp_v48 = out_tids_inline271[((n_split_base_inline163 * 5) + 9)]; + _submit_deps_buf_inline165[9] = t__tmp_v48; + + // Task 14: residual_rms_cast + L0TaskArgs params_t14; + params_t14.add_inout(mlp_norm_in_inline71); + params_t14.add_inout(post_norm_partial_inline118); + params_t14.add_input(attn_proj_fp32_inline220); + params_t14.add_input(cur__rv_v7); + params_t14.add_input(ext_post_rms_weight); + params_t14.add_scalar(k_base_inline111); + params_t14.add_scalar(i); + PTO2TaskId params_t14_deps[10]; + uint32_t params_t14_deps_count = 0; + if (_submit_deps_buf_inline165[0].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[0]; + if (_submit_deps_buf_inline165[1].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[1]; + if (_submit_deps_buf_inline165[2].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[2]; + if (_submit_deps_buf_inline165[3].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[3]; + if (_submit_deps_buf_inline165[4].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[4]; + if (_submit_deps_buf_inline165[5].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[5]; + if (_submit_deps_buf_inline165[6].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[6]; + if (_submit_deps_buf_inline165[7].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[7]; + if (_submit_deps_buf_inline165[8].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[8]; + if (_submit_deps_buf_inline165[9].is_valid()) + params_t14_deps[params_t14_deps_count++] = _submit_deps_buf_inline165[9]; + params_t14.set_dependencies(params_t14_deps, params_t14_deps_count); + params_t14.set_allow_early_resolve(true); + TaskOutputTensors task_14_outs = rt_submit_aiv_task(15, params_t14); + PTO2TaskId cast_tid_k_inline76 = task_14_outs.task_id(); + cast_tids_inline88[0] = cast_tid_k_inline76; + int64_t k_base_inline111__ssa_v1 = 1024; + int64_t n_split_base_inline163__ssa_v1 = 2; + PTO2TaskId _submit_deps_buf_inline165__ssa_v1[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v51 = out_tids_inline271[(n_split_base_inline163__ssa_v1 * 5)]; + _submit_deps_buf_inline165__ssa_v1[0] = t__tmp_v51; + PTO2TaskId t__tmp_v52 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v1[1] = t__tmp_v52; + PTO2TaskId t__tmp_v53 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v1[2] = t__tmp_v53; + PTO2TaskId t__tmp_v54 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v1[3] = t__tmp_v54; + PTO2TaskId t__tmp_v55 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v1[4] = t__tmp_v55; + PTO2TaskId t__tmp_v56 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v1[5] = t__tmp_v56; + PTO2TaskId t__tmp_v57 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v1[6] = t__tmp_v57; + PTO2TaskId t__tmp_v58 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v1[7] = t__tmp_v58; + PTO2TaskId t__tmp_v59 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v1[8] = t__tmp_v59; + PTO2TaskId t__tmp_v60 = out_tids_inline271[((n_split_base_inline163__ssa_v1 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v1[9] = t__tmp_v60; + + // Task 15: residual_rms_cast_0 + L0TaskArgs params_t15; + params_t15.add_inout(mlp_norm_in_inline71); + params_t15.add_inout(post_norm_partial_inline118); + params_t15.add_input(attn_proj_fp32_inline220); + params_t15.add_input(cur__rv_v7); + params_t15.add_input(ext_post_rms_weight); + params_t15.add_scalar(k_base_inline111__ssa_v1); + params_t15.add_scalar(i); + PTO2TaskId params_t15_deps[10]; + uint32_t params_t15_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v1[0].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[0]; + if (_submit_deps_buf_inline165__ssa_v1[1].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[1]; + if (_submit_deps_buf_inline165__ssa_v1[2].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[2]; + if (_submit_deps_buf_inline165__ssa_v1[3].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[3]; + if (_submit_deps_buf_inline165__ssa_v1[4].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[4]; + if (_submit_deps_buf_inline165__ssa_v1[5].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[5]; + if (_submit_deps_buf_inline165__ssa_v1[6].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[6]; + if (_submit_deps_buf_inline165__ssa_v1[7].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[7]; + if (_submit_deps_buf_inline165__ssa_v1[8].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[8]; + if (_submit_deps_buf_inline165__ssa_v1[9].is_valid()) + params_t15_deps[params_t15_deps_count++] = _submit_deps_buf_inline165__ssa_v1[9]; + params_t15.set_dependencies(params_t15_deps, params_t15_deps_count); + params_t15.set_allow_early_resolve(true); + TaskOutputTensors task_15_outs = rt_submit_aiv_task(16, params_t15); + PTO2TaskId cast_tid_k_inline76__ssa_v1 = task_15_outs.task_id(); + cast_tids_inline88[1] = cast_tid_k_inline76__ssa_v1; + int64_t k_base_inline111__ssa_v2 = 2048; + int64_t n_split_base_inline163__ssa_v2 = 4; + PTO2TaskId _submit_deps_buf_inline165__ssa_v2[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v63 = out_tids_inline271[(n_split_base_inline163__ssa_v2 * 5)]; + _submit_deps_buf_inline165__ssa_v2[0] = t__tmp_v63; + PTO2TaskId t__tmp_v64 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v2[1] = t__tmp_v64; + PTO2TaskId t__tmp_v65 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v2[2] = t__tmp_v65; + PTO2TaskId t__tmp_v66 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v2[3] = t__tmp_v66; + PTO2TaskId t__tmp_v67 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v2[4] = t__tmp_v67; + PTO2TaskId t__tmp_v68 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v2[5] = t__tmp_v68; + PTO2TaskId t__tmp_v69 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v2[6] = t__tmp_v69; + PTO2TaskId t__tmp_v70 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v2[7] = t__tmp_v70; + PTO2TaskId t__tmp_v71 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v2[8] = t__tmp_v71; + PTO2TaskId t__tmp_v72 = out_tids_inline271[((n_split_base_inline163__ssa_v2 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v2[9] = t__tmp_v72; + + // Task 16: residual_rms_cast_1 + L0TaskArgs params_t16; + params_t16.add_inout(mlp_norm_in_inline71); + params_t16.add_inout(post_norm_partial_inline118); + params_t16.add_input(attn_proj_fp32_inline220); + params_t16.add_input(cur__rv_v7); + params_t16.add_input(ext_post_rms_weight); + params_t16.add_scalar(k_base_inline111__ssa_v2); + params_t16.add_scalar(i); + PTO2TaskId params_t16_deps[10]; + uint32_t params_t16_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v2[0].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[0]; + if (_submit_deps_buf_inline165__ssa_v2[1].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[1]; + if (_submit_deps_buf_inline165__ssa_v2[2].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[2]; + if (_submit_deps_buf_inline165__ssa_v2[3].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[3]; + if (_submit_deps_buf_inline165__ssa_v2[4].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[4]; + if (_submit_deps_buf_inline165__ssa_v2[5].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[5]; + if (_submit_deps_buf_inline165__ssa_v2[6].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[6]; + if (_submit_deps_buf_inline165__ssa_v2[7].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[7]; + if (_submit_deps_buf_inline165__ssa_v2[8].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[8]; + if (_submit_deps_buf_inline165__ssa_v2[9].is_valid()) + params_t16_deps[params_t16_deps_count++] = _submit_deps_buf_inline165__ssa_v2[9]; + params_t16.set_dependencies(params_t16_deps, params_t16_deps_count); + params_t16.set_allow_early_resolve(true); + TaskOutputTensors task_16_outs = rt_submit_aiv_task(17, params_t16); + PTO2TaskId cast_tid_k_inline76__ssa_v2 = task_16_outs.task_id(); + cast_tids_inline88[2] = cast_tid_k_inline76__ssa_v2; + int64_t k_base_inline111__ssa_v3 = 3072; + int64_t n_split_base_inline163__ssa_v3 = 6; + PTO2TaskId _submit_deps_buf_inline165__ssa_v3[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v75 = out_tids_inline271[(n_split_base_inline163__ssa_v3 * 5)]; + _submit_deps_buf_inline165__ssa_v3[0] = t__tmp_v75; + PTO2TaskId t__tmp_v76 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v3[1] = t__tmp_v76; + PTO2TaskId t__tmp_v77 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v3[2] = t__tmp_v77; + PTO2TaskId t__tmp_v78 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v3[3] = t__tmp_v78; + PTO2TaskId t__tmp_v79 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v3[4] = t__tmp_v79; + PTO2TaskId t__tmp_v80 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v3[5] = t__tmp_v80; + PTO2TaskId t__tmp_v81 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v3[6] = t__tmp_v81; + PTO2TaskId t__tmp_v82 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v3[7] = t__tmp_v82; + PTO2TaskId t__tmp_v83 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v3[8] = t__tmp_v83; + PTO2TaskId t__tmp_v84 = out_tids_inline271[((n_split_base_inline163__ssa_v3 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v3[9] = t__tmp_v84; + + // Task 17: residual_rms_cast_2 + L0TaskArgs params_t17; + params_t17.add_inout(mlp_norm_in_inline71); + params_t17.add_inout(post_norm_partial_inline118); + params_t17.add_input(attn_proj_fp32_inline220); + params_t17.add_input(cur__rv_v7); + params_t17.add_input(ext_post_rms_weight); + params_t17.add_scalar(k_base_inline111__ssa_v3); + params_t17.add_scalar(i); + PTO2TaskId params_t17_deps[10]; + uint32_t params_t17_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v3[0].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[0]; + if (_submit_deps_buf_inline165__ssa_v3[1].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[1]; + if (_submit_deps_buf_inline165__ssa_v3[2].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[2]; + if (_submit_deps_buf_inline165__ssa_v3[3].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[3]; + if (_submit_deps_buf_inline165__ssa_v3[4].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[4]; + if (_submit_deps_buf_inline165__ssa_v3[5].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[5]; + if (_submit_deps_buf_inline165__ssa_v3[6].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[6]; + if (_submit_deps_buf_inline165__ssa_v3[7].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[7]; + if (_submit_deps_buf_inline165__ssa_v3[8].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[8]; + if (_submit_deps_buf_inline165__ssa_v3[9].is_valid()) + params_t17_deps[params_t17_deps_count++] = _submit_deps_buf_inline165__ssa_v3[9]; + params_t17.set_dependencies(params_t17_deps, params_t17_deps_count); + params_t17.set_allow_early_resolve(true); + TaskOutputTensors task_17_outs = rt_submit_aiv_task(18, params_t17); + PTO2TaskId cast_tid_k_inline76__ssa_v3 = task_17_outs.task_id(); + cast_tids_inline88[3] = cast_tid_k_inline76__ssa_v3; + int64_t k_base_inline111__ssa_v4 = 4096; + int64_t n_split_base_inline163__ssa_v4 = 8; + PTO2TaskId _submit_deps_buf_inline165__ssa_v4[10]; + for (int64_t __init_i = 0; __init_i < 10; ++__init_i) + _submit_deps_buf_inline165__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v87 = out_tids_inline271[(n_split_base_inline163__ssa_v4 * 5)]; + _submit_deps_buf_inline165__ssa_v4[0] = t__tmp_v87; + PTO2TaskId t__tmp_v88 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 1)]; + _submit_deps_buf_inline165__ssa_v4[1] = t__tmp_v88; + PTO2TaskId t__tmp_v89 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 2)]; + _submit_deps_buf_inline165__ssa_v4[2] = t__tmp_v89; + PTO2TaskId t__tmp_v90 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 3)]; + _submit_deps_buf_inline165__ssa_v4[3] = t__tmp_v90; + PTO2TaskId t__tmp_v91 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 4)]; + _submit_deps_buf_inline165__ssa_v4[4] = t__tmp_v91; + PTO2TaskId t__tmp_v92 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 5)]; + _submit_deps_buf_inline165__ssa_v4[5] = t__tmp_v92; + PTO2TaskId t__tmp_v93 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 6)]; + _submit_deps_buf_inline165__ssa_v4[6] = t__tmp_v93; + PTO2TaskId t__tmp_v94 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 7)]; + _submit_deps_buf_inline165__ssa_v4[7] = t__tmp_v94; + PTO2TaskId t__tmp_v95 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 8)]; + _submit_deps_buf_inline165__ssa_v4[8] = t__tmp_v95; + PTO2TaskId t__tmp_v96 = out_tids_inline271[((n_split_base_inline163__ssa_v4 * 5) + 9)]; + _submit_deps_buf_inline165__ssa_v4[9] = t__tmp_v96; + + // Task 18: residual_rms_cast_3 + L0TaskArgs params_t18; + params_t18.add_inout(mlp_norm_in_inline71); + params_t18.add_inout(post_norm_partial_inline118); + params_t18.add_input(attn_proj_fp32_inline220); + params_t18.add_input(cur__rv_v7); + params_t18.add_input(ext_post_rms_weight); + params_t18.add_scalar(k_base_inline111__ssa_v4); + params_t18.add_scalar(i); + PTO2TaskId params_t18_deps[10]; + uint32_t params_t18_deps_count = 0; + if (_submit_deps_buf_inline165__ssa_v4[0].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[0]; + if (_submit_deps_buf_inline165__ssa_v4[1].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[1]; + if (_submit_deps_buf_inline165__ssa_v4[2].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[2]; + if (_submit_deps_buf_inline165__ssa_v4[3].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[3]; + if (_submit_deps_buf_inline165__ssa_v4[4].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[4]; + if (_submit_deps_buf_inline165__ssa_v4[5].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[5]; + if (_submit_deps_buf_inline165__ssa_v4[6].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[6]; + if (_submit_deps_buf_inline165__ssa_v4[7].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[7]; + if (_submit_deps_buf_inline165__ssa_v4[8].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[8]; + if (_submit_deps_buf_inline165__ssa_v4[9].is_valid()) + params_t18_deps[params_t18_deps_count++] = _submit_deps_buf_inline165__ssa_v4[9]; + params_t18.set_dependencies(params_t18_deps, params_t18_deps_count); + params_t18.set_allow_early_resolve(true); + TaskOutputTensors task_18_outs = rt_submit_aiv_task(19, params_t18); + PTO2TaskId cast_tid_k_inline76__ssa_v4 = task_18_outs.task_id(); + cast_tids_inline88[4] = cast_tid_k_inline76__ssa_v4; + + // Task 19: post_rms_reduce + L0TaskArgs params_t19; + params_t19.add_input(attn_proj_fp32_inline220); + params_t19.add_input(cur__rv_v7); + params_t19.add_inout(inv_rms_tile_inline126); + PTO2TaskId params_t19_deps[50]; + uint32_t params_t19_deps_count = 0; + if (out_tids_inline271[0].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[0]; + if (out_tids_inline271[1].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[1]; + if (out_tids_inline271[2].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[2]; + if (out_tids_inline271[3].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[3]; + if (out_tids_inline271[4].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[4]; + if (out_tids_inline271[5].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[5]; + if (out_tids_inline271[6].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[6]; + if (out_tids_inline271[7].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[7]; + if (out_tids_inline271[8].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[8]; + if (out_tids_inline271[9].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[9]; + if (out_tids_inline271[10].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[10]; + if (out_tids_inline271[11].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[11]; + if (out_tids_inline271[12].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[12]; + if (out_tids_inline271[13].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[13]; + if (out_tids_inline271[14].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[14]; + if (out_tids_inline271[15].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[15]; + if (out_tids_inline271[16].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[16]; + if (out_tids_inline271[17].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[17]; + if (out_tids_inline271[18].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[18]; + if (out_tids_inline271[19].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[19]; + if (out_tids_inline271[20].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[20]; + if (out_tids_inline271[21].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[21]; + if (out_tids_inline271[22].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[22]; + if (out_tids_inline271[23].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[23]; + if (out_tids_inline271[24].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[24]; + if (out_tids_inline271[25].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[25]; + if (out_tids_inline271[26].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[26]; + if (out_tids_inline271[27].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[27]; + if (out_tids_inline271[28].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[28]; + if (out_tids_inline271[29].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[29]; + if (out_tids_inline271[30].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[30]; + if (out_tids_inline271[31].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[31]; + if (out_tids_inline271[32].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[32]; + if (out_tids_inline271[33].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[33]; + if (out_tids_inline271[34].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[34]; + if (out_tids_inline271[35].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[35]; + if (out_tids_inline271[36].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[36]; + if (out_tids_inline271[37].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[37]; + if (out_tids_inline271[38].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[38]; + if (out_tids_inline271[39].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[39]; + if (out_tids_inline271[40].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[40]; + if (out_tids_inline271[41].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[41]; + if (out_tids_inline271[42].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[42]; + if (out_tids_inline271[43].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[43]; + if (out_tids_inline271[44].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[44]; + if (out_tids_inline271[45].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[45]; + if (out_tids_inline271[46].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[46]; + if (out_tids_inline271[47].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[47]; + if (out_tids_inline271[48].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[48]; + if (out_tids_inline271[49].is_valid()) + params_t19_deps[params_t19_deps_count++] = out_tids_inline271[49]; + params_t19.set_dependencies(params_t19_deps, params_t19_deps_count); + TaskOutputTensors task_19_outs = rt_submit_aiv_task(20, params_t19); + PTO2TaskId reduce_tid_inline226 = task_19_outs.task_id(); + int64_t gu_k0_inline131 = 0; + PTO2TaskId _submit_deps_buf_inline236[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v105 = cast_tids_inline88[0]; + _submit_deps_buf_inline236[0] = t__tmp_v105; + + // Phase-fence barrier 3: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_3; + PTO2TaskId params_phase_fence_barrier_3_deps[1]; + uint32_t params_phase_fence_barrier_3_deps_count = 0; + if (_submit_deps_buf_inline236[0].is_valid()) + params_phase_fence_barrier_3_deps[params_phase_fence_barrier_3_deps_count++] = + _submit_deps_buf_inline236[0]; + params_phase_fence_barrier_3.set_dependencies( + params_phase_fence_barrier_3_deps, params_phase_fence_barrier_3_deps_count + ); + PTO2TaskId t__tmp_v106 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_3_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_3_outs = + rt_submit_dummy_task(params_phase_fence_barrier_3); + t__tmp_v106 = phase_fence_barrier_3_outs.task_id(); + } + gate_late_tids_inline249[0] = t__tmp_v106; + PTO2TaskId _submit_deps_buf_inline225[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v107 = cast_tids_inline88[0]; + _submit_deps_buf_inline225[0] = t__tmp_v107; + + // Phase-fence barrier 4: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_4; + PTO2TaskId params_phase_fence_barrier_4_deps[1]; + uint32_t params_phase_fence_barrier_4_deps_count = 0; + if (_submit_deps_buf_inline225[0].is_valid()) + params_phase_fence_barrier_4_deps[params_phase_fence_barrier_4_deps_count++] = + _submit_deps_buf_inline225[0]; + params_phase_fence_barrier_4.set_dependencies( + params_phase_fence_barrier_4_deps, params_phase_fence_barrier_4_deps_count + ); + PTO2TaskId t__tmp_v108 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_4_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_4_outs = + rt_submit_dummy_task(params_phase_fence_barrier_4); + t__tmp_v108 = phase_fence_barrier_4_outs.task_id(); + } + up_late_tids_inline69[0] = t__tmp_v108; + PTO2TaskId _submit_deps_buf_inline237[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v109 = cast_tids_inline88[0]; + _submit_deps_buf_inline237[0] = t__tmp_v109; + + // Spmd gate_proj_spmd: gate_proj + L0TaskArgs params_t20; + params_t20.add_input(mlp_norm_in_inline71); + params_t20.add_input(ext_w_gate); + params_t20.add_inout(gate_acc_all_inline203); + params_t20.add_scalar(gu_k0_inline131); + params_t20.add_scalar(layer_hidden_base_inline151); + params_t20.launch_spec.set_block_num(6); + PTO2TaskId params_t20_deps[1]; + uint32_t params_t20_deps_count = 0; + if (_submit_deps_buf_inline237[0].is_valid()) + params_t20_deps[params_t20_deps_count++] = _submit_deps_buf_inline237[0]; + params_t20.set_dependencies(params_t20_deps, params_t20_deps_count); + TaskOutputTensors task_20_outs = rt_submit_aic_task(21, params_t20); + PTO2TaskId gate_spmd_tid_inline245 = task_20_outs.task_id(); + gate_tids_inline56[0] = gate_spmd_tid_inline245; + gate_tids_inline56[5] = gate_spmd_tid_inline245; + gate_tids_inline56[10] = gate_spmd_tid_inline245; + gate_tids_inline56[15] = gate_spmd_tid_inline245; + gate_tids_inline56[20] = gate_spmd_tid_inline245; + gate_tids_inline56[25] = gate_spmd_tid_inline245; + PTO2TaskId _submit_deps_buf_inline260[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v110 = cast_tids_inline88[0]; + _submit_deps_buf_inline260[0] = t__tmp_v110; + + // Spmd up_proj_spmd: up_proj + L0TaskArgs params_t21; + params_t21.add_input(mlp_norm_in_inline71); + params_t21.add_input(ext_w_up); + params_t21.add_inout(up_acc_all_inline303); + params_t21.add_scalar(gu_k0_inline131); + params_t21.add_scalar(layer_hidden_base_inline151); + params_t21.launch_spec.set_block_num(6); + PTO2TaskId params_t21_deps[1]; + uint32_t params_t21_deps_count = 0; + if (_submit_deps_buf_inline260[0].is_valid()) + params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline260[0]; + params_t21.set_dependencies(params_t21_deps, params_t21_deps_count); + TaskOutputTensors task_21_outs = rt_submit_aic_task(22, params_t21); + PTO2TaskId up_spmd_tid_inline264 = task_21_outs.task_id(); + up_tids_inline310[0] = up_spmd_tid_inline264; + up_tids_inline310[5] = up_spmd_tid_inline264; + up_tids_inline310[10] = up_spmd_tid_inline264; + up_tids_inline310[15] = up_spmd_tid_inline264; + up_tids_inline310[20] = up_spmd_tid_inline264; + up_tids_inline310[25] = up_spmd_tid_inline264; + int64_t gu_k0_inline131__ssa_v1 = 1024; + PTO2TaskId _submit_deps_buf_inline236__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v111 = cast_tids_inline88[1]; + _submit_deps_buf_inline236__ssa_v1[0] = t__tmp_v111; + + // Phase-fence barrier 5: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_5; + PTO2TaskId params_phase_fence_barrier_5_deps[1]; + uint32_t params_phase_fence_barrier_5_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v1[0].is_valid()) + params_phase_fence_barrier_5_deps[params_phase_fence_barrier_5_deps_count++] = + _submit_deps_buf_inline236__ssa_v1[0]; + params_phase_fence_barrier_5.set_dependencies( + params_phase_fence_barrier_5_deps, params_phase_fence_barrier_5_deps_count + ); + PTO2TaskId t__tmp_v112 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_5_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_5_outs = + rt_submit_dummy_task(params_phase_fence_barrier_5); + t__tmp_v112 = phase_fence_barrier_5_outs.task_id(); + } + gate_late_tids_inline249[1] = t__tmp_v112; + PTO2TaskId _submit_deps_buf_inline225__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v113 = cast_tids_inline88[1]; + _submit_deps_buf_inline225__ssa_v1[0] = t__tmp_v113; + + // Phase-fence barrier 6: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_6; + PTO2TaskId params_phase_fence_barrier_6_deps[1]; + uint32_t params_phase_fence_barrier_6_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v1[0].is_valid()) + params_phase_fence_barrier_6_deps[params_phase_fence_barrier_6_deps_count++] = + _submit_deps_buf_inline225__ssa_v1[0]; + params_phase_fence_barrier_6.set_dependencies( + params_phase_fence_barrier_6_deps, params_phase_fence_barrier_6_deps_count + ); + PTO2TaskId t__tmp_v114 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_6_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_6_outs = + rt_submit_dummy_task(params_phase_fence_barrier_6); + t__tmp_v114 = phase_fence_barrier_6_outs.task_id(); + } + up_late_tids_inline69[1] = t__tmp_v114; + PTO2TaskId _submit_deps_buf_inline237__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v115 = cast_tids_inline88[1]; + _submit_deps_buf_inline237__ssa_v1[0] = t__tmp_v115; + + // Spmd gate_proj_spmd_0: gate_proj_0 + L0TaskArgs params_t22; + params_t22.add_input(mlp_norm_in_inline71); + params_t22.add_input(ext_w_gate); + params_t22.add_inout(gate_acc_all_inline203); + params_t22.add_scalar(gu_k0_inline131__ssa_v1); + params_t22.add_scalar(layer_hidden_base_inline151); + params_t22.launch_spec.set_block_num(6); + PTO2TaskId params_t22_deps[1]; + uint32_t params_t22_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v1[0].is_valid()) + params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline237__ssa_v1[0]; + params_t22.set_dependencies(params_t22_deps, params_t22_deps_count); + TaskOutputTensors task_22_outs = rt_submit_aic_task(23, params_t22); + PTO2TaskId gate_spmd_tid_inline245__ssa_v1 = task_22_outs.task_id(); + gate_tids_inline56[1] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[6] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[11] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[16] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[21] = gate_spmd_tid_inline245__ssa_v1; + gate_tids_inline56[26] = gate_spmd_tid_inline245__ssa_v1; + PTO2TaskId _submit_deps_buf_inline260__ssa_v1[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v1[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v116 = cast_tids_inline88[1]; + _submit_deps_buf_inline260__ssa_v1[0] = t__tmp_v116; + + // Spmd up_proj_spmd_0: up_proj_0 + L0TaskArgs params_t23; + params_t23.add_input(mlp_norm_in_inline71); + params_t23.add_input(ext_w_up); + params_t23.add_inout(up_acc_all_inline303); + params_t23.add_scalar(gu_k0_inline131__ssa_v1); + params_t23.add_scalar(layer_hidden_base_inline151); + params_t23.launch_spec.set_block_num(6); + PTO2TaskId params_t23_deps[1]; + uint32_t params_t23_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v1[0].is_valid()) + params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline260__ssa_v1[0]; + params_t23.set_dependencies(params_t23_deps, params_t23_deps_count); + TaskOutputTensors task_23_outs = rt_submit_aic_task(24, params_t23); + PTO2TaskId up_spmd_tid_inline264__ssa_v1 = task_23_outs.task_id(); + up_tids_inline310[1] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[6] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[11] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[16] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[21] = up_spmd_tid_inline264__ssa_v1; + up_tids_inline310[26] = up_spmd_tid_inline264__ssa_v1; + int64_t gu_k0_inline131__ssa_v2 = 2048; + PTO2TaskId _submit_deps_buf_inline236__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v117 = cast_tids_inline88[2]; + _submit_deps_buf_inline236__ssa_v2[0] = t__tmp_v117; + + // Phase-fence barrier 7: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_7; + PTO2TaskId params_phase_fence_barrier_7_deps[1]; + uint32_t params_phase_fence_barrier_7_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v2[0].is_valid()) + params_phase_fence_barrier_7_deps[params_phase_fence_barrier_7_deps_count++] = + _submit_deps_buf_inline236__ssa_v2[0]; + params_phase_fence_barrier_7.set_dependencies( + params_phase_fence_barrier_7_deps, params_phase_fence_barrier_7_deps_count + ); + PTO2TaskId t__tmp_v118 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_7_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_7_outs = + rt_submit_dummy_task(params_phase_fence_barrier_7); + t__tmp_v118 = phase_fence_barrier_7_outs.task_id(); + } + gate_late_tids_inline249[2] = t__tmp_v118; + PTO2TaskId _submit_deps_buf_inline225__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v119 = cast_tids_inline88[2]; + _submit_deps_buf_inline225__ssa_v2[0] = t__tmp_v119; + + // Phase-fence barrier 8: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_8; + PTO2TaskId params_phase_fence_barrier_8_deps[1]; + uint32_t params_phase_fence_barrier_8_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v2[0].is_valid()) + params_phase_fence_barrier_8_deps[params_phase_fence_barrier_8_deps_count++] = + _submit_deps_buf_inline225__ssa_v2[0]; + params_phase_fence_barrier_8.set_dependencies( + params_phase_fence_barrier_8_deps, params_phase_fence_barrier_8_deps_count + ); + PTO2TaskId t__tmp_v120 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_8_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_8_outs = + rt_submit_dummy_task(params_phase_fence_barrier_8); + t__tmp_v120 = phase_fence_barrier_8_outs.task_id(); + } + up_late_tids_inline69[2] = t__tmp_v120; + PTO2TaskId _submit_deps_buf_inline237__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v121 = cast_tids_inline88[2]; + _submit_deps_buf_inline237__ssa_v2[0] = t__tmp_v121; + + // Spmd gate_proj_spmd_1: gate_proj_1 + L0TaskArgs params_t24; + params_t24.add_input(mlp_norm_in_inline71); + params_t24.add_input(ext_w_gate); + params_t24.add_inout(gate_acc_all_inline203); + params_t24.add_scalar(gu_k0_inline131__ssa_v2); + params_t24.add_scalar(layer_hidden_base_inline151); + params_t24.launch_spec.set_block_num(6); + PTO2TaskId params_t24_deps[1]; + uint32_t params_t24_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v2[0].is_valid()) + params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline237__ssa_v2[0]; + params_t24.set_dependencies(params_t24_deps, params_t24_deps_count); + TaskOutputTensors task_24_outs = rt_submit_aic_task(25, params_t24); + PTO2TaskId gate_spmd_tid_inline245__ssa_v2 = task_24_outs.task_id(); + gate_tids_inline56[2] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[7] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[12] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[17] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[22] = gate_spmd_tid_inline245__ssa_v2; + gate_tids_inline56[27] = gate_spmd_tid_inline245__ssa_v2; + PTO2TaskId _submit_deps_buf_inline260__ssa_v2[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v2[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v122 = cast_tids_inline88[2]; + _submit_deps_buf_inline260__ssa_v2[0] = t__tmp_v122; + + // Spmd up_proj_spmd_1: up_proj_1 + L0TaskArgs params_t25; + params_t25.add_input(mlp_norm_in_inline71); + params_t25.add_input(ext_w_up); + params_t25.add_inout(up_acc_all_inline303); + params_t25.add_scalar(gu_k0_inline131__ssa_v2); + params_t25.add_scalar(layer_hidden_base_inline151); + params_t25.launch_spec.set_block_num(6); + PTO2TaskId params_t25_deps[1]; + uint32_t params_t25_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v2[0].is_valid()) + params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline260__ssa_v2[0]; + params_t25.set_dependencies(params_t25_deps, params_t25_deps_count); + TaskOutputTensors task_25_outs = rt_submit_aic_task(26, params_t25); + PTO2TaskId up_spmd_tid_inline264__ssa_v2 = task_25_outs.task_id(); + up_tids_inline310[2] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[7] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[12] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[17] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[22] = up_spmd_tid_inline264__ssa_v2; + up_tids_inline310[27] = up_spmd_tid_inline264__ssa_v2; + int64_t gu_k0_inline131__ssa_v3 = 3072; + PTO2TaskId _submit_deps_buf_inline236__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v123 = cast_tids_inline88[3]; + _submit_deps_buf_inline236__ssa_v3[0] = t__tmp_v123; + + // Phase-fence barrier 9: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_9; + PTO2TaskId params_phase_fence_barrier_9_deps[1]; + uint32_t params_phase_fence_barrier_9_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v3[0].is_valid()) + params_phase_fence_barrier_9_deps[params_phase_fence_barrier_9_deps_count++] = + _submit_deps_buf_inline236__ssa_v3[0]; + params_phase_fence_barrier_9.set_dependencies( + params_phase_fence_barrier_9_deps, params_phase_fence_barrier_9_deps_count + ); + PTO2TaskId t__tmp_v124 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_9_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_9_outs = + rt_submit_dummy_task(params_phase_fence_barrier_9); + t__tmp_v124 = phase_fence_barrier_9_outs.task_id(); + } + gate_late_tids_inline249[3] = t__tmp_v124; + PTO2TaskId _submit_deps_buf_inline225__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v125 = cast_tids_inline88[3]; + _submit_deps_buf_inline225__ssa_v3[0] = t__tmp_v125; + + // Phase-fence barrier 10: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_10; + PTO2TaskId params_phase_fence_barrier_10_deps[1]; + uint32_t params_phase_fence_barrier_10_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v3[0].is_valid()) + params_phase_fence_barrier_10_deps[params_phase_fence_barrier_10_deps_count++] = + _submit_deps_buf_inline225__ssa_v3[0]; + params_phase_fence_barrier_10.set_dependencies( + params_phase_fence_barrier_10_deps, params_phase_fence_barrier_10_deps_count + ); + PTO2TaskId t__tmp_v126 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_10_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_10_outs = + rt_submit_dummy_task(params_phase_fence_barrier_10); + t__tmp_v126 = phase_fence_barrier_10_outs.task_id(); + } + up_late_tids_inline69[3] = t__tmp_v126; + PTO2TaskId _submit_deps_buf_inline237__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v127 = cast_tids_inline88[3]; + _submit_deps_buf_inline237__ssa_v3[0] = t__tmp_v127; + + // Spmd gate_proj_spmd_2: gate_proj_2 + L0TaskArgs params_t26; + params_t26.add_input(mlp_norm_in_inline71); + params_t26.add_input(ext_w_gate); + params_t26.add_inout(gate_acc_all_inline203); + params_t26.add_scalar(gu_k0_inline131__ssa_v3); + params_t26.add_scalar(layer_hidden_base_inline151); + params_t26.launch_spec.set_block_num(6); + PTO2TaskId params_t26_deps[1]; + uint32_t params_t26_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v3[0].is_valid()) + params_t26_deps[params_t26_deps_count++] = _submit_deps_buf_inline237__ssa_v3[0]; + params_t26.set_dependencies(params_t26_deps, params_t26_deps_count); + TaskOutputTensors task_26_outs = rt_submit_aic_task(27, params_t26); + PTO2TaskId gate_spmd_tid_inline245__ssa_v3 = task_26_outs.task_id(); + gate_tids_inline56[3] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[8] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[13] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[18] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[23] = gate_spmd_tid_inline245__ssa_v3; + gate_tids_inline56[28] = gate_spmd_tid_inline245__ssa_v3; + PTO2TaskId _submit_deps_buf_inline260__ssa_v3[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v3[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v128 = cast_tids_inline88[3]; + _submit_deps_buf_inline260__ssa_v3[0] = t__tmp_v128; + + // Spmd up_proj_spmd_2: up_proj_2 + L0TaskArgs params_t27; + params_t27.add_input(mlp_norm_in_inline71); + params_t27.add_input(ext_w_up); + params_t27.add_inout(up_acc_all_inline303); + params_t27.add_scalar(gu_k0_inline131__ssa_v3); + params_t27.add_scalar(layer_hidden_base_inline151); + params_t27.launch_spec.set_block_num(6); + PTO2TaskId params_t27_deps[1]; + uint32_t params_t27_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v3[0].is_valid()) + params_t27_deps[params_t27_deps_count++] = _submit_deps_buf_inline260__ssa_v3[0]; + params_t27.set_dependencies(params_t27_deps, params_t27_deps_count); + TaskOutputTensors task_27_outs = rt_submit_aic_task(28, params_t27); + PTO2TaskId up_spmd_tid_inline264__ssa_v3 = task_27_outs.task_id(); + up_tids_inline310[3] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[8] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[13] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[18] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[23] = up_spmd_tid_inline264__ssa_v3; + up_tids_inline310[28] = up_spmd_tid_inline264__ssa_v3; + int64_t gu_k0_inline131__ssa_v4 = 4096; + PTO2TaskId _submit_deps_buf_inline236__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline236__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v129 = cast_tids_inline88[4]; + _submit_deps_buf_inline236__ssa_v4[0] = t__tmp_v129; + + // Phase-fence barrier 11: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_11; + PTO2TaskId params_phase_fence_barrier_11_deps[1]; + uint32_t params_phase_fence_barrier_11_deps_count = 0; + if (_submit_deps_buf_inline236__ssa_v4[0].is_valid()) + params_phase_fence_barrier_11_deps[params_phase_fence_barrier_11_deps_count++] = + _submit_deps_buf_inline236__ssa_v4[0]; + params_phase_fence_barrier_11.set_dependencies( + params_phase_fence_barrier_11_deps, params_phase_fence_barrier_11_deps_count + ); + PTO2TaskId t__tmp_v130 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_11_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_11_outs = + rt_submit_dummy_task(params_phase_fence_barrier_11); + t__tmp_v130 = phase_fence_barrier_11_outs.task_id(); + } + gate_late_tids_inline249[4] = t__tmp_v130; + PTO2TaskId _submit_deps_buf_inline225__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline225__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v131 = cast_tids_inline88[4]; + _submit_deps_buf_inline225__ssa_v4[0] = t__tmp_v131; + + // Phase-fence barrier 12: dependency-only dummy task + L0TaskArgs params_phase_fence_barrier_12; + PTO2TaskId params_phase_fence_barrier_12_deps[1]; + uint32_t params_phase_fence_barrier_12_deps_count = 0; + if (_submit_deps_buf_inline225__ssa_v4[0].is_valid()) + params_phase_fence_barrier_12_deps[params_phase_fence_barrier_12_deps_count++] = + _submit_deps_buf_inline225__ssa_v4[0]; + params_phase_fence_barrier_12.set_dependencies( + params_phase_fence_barrier_12_deps, params_phase_fence_barrier_12_deps_count + ); + PTO2TaskId t__tmp_v132 = PTO2TaskId::invalid(); + if (params_phase_fence_barrier_12_deps_count > 0) { + TaskOutputTensors phase_fence_barrier_12_outs = + rt_submit_dummy_task(params_phase_fence_barrier_12); + t__tmp_v132 = phase_fence_barrier_12_outs.task_id(); + } + up_late_tids_inline69[4] = t__tmp_v132; + PTO2TaskId _submit_deps_buf_inline237__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline237__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v133 = cast_tids_inline88[4]; + _submit_deps_buf_inline237__ssa_v4[0] = t__tmp_v133; + + // Spmd gate_proj_spmd_3: gate_proj_3 + L0TaskArgs params_t28; + params_t28.add_input(mlp_norm_in_inline71); + params_t28.add_input(ext_w_gate); + params_t28.add_inout(gate_acc_all_inline203); + params_t28.add_scalar(gu_k0_inline131__ssa_v4); + params_t28.add_scalar(layer_hidden_base_inline151); + params_t28.launch_spec.set_block_num(6); + PTO2TaskId params_t28_deps[1]; + uint32_t params_t28_deps_count = 0; + if (_submit_deps_buf_inline237__ssa_v4[0].is_valid()) + params_t28_deps[params_t28_deps_count++] = _submit_deps_buf_inline237__ssa_v4[0]; + params_t28.set_dependencies(params_t28_deps, params_t28_deps_count); + TaskOutputTensors task_28_outs = rt_submit_aic_task(29, params_t28); + PTO2TaskId gate_spmd_tid_inline245__ssa_v4 = task_28_outs.task_id(); + gate_tids_inline56[4] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[9] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[14] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[19] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[24] = gate_spmd_tid_inline245__ssa_v4; + gate_tids_inline56[29] = gate_spmd_tid_inline245__ssa_v4; + PTO2TaskId _submit_deps_buf_inline260__ssa_v4[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline260__ssa_v4[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v134 = cast_tids_inline88[4]; + _submit_deps_buf_inline260__ssa_v4[0] = t__tmp_v134; + + // Spmd up_proj_spmd_3: up_proj_3 + L0TaskArgs params_t29; + params_t29.add_input(mlp_norm_in_inline71); + params_t29.add_input(ext_w_up); + params_t29.add_inout(up_acc_all_inline303); + params_t29.add_scalar(gu_k0_inline131__ssa_v4); + params_t29.add_scalar(layer_hidden_base_inline151); + params_t29.launch_spec.set_block_num(6); + PTO2TaskId params_t29_deps[1]; + uint32_t params_t29_deps_count = 0; + if (_submit_deps_buf_inline260__ssa_v4[0].is_valid()) + params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline260__ssa_v4[0]; + params_t29.set_dependencies(params_t29_deps, params_t29_deps_count); + TaskOutputTensors task_29_outs = rt_submit_aic_task(30, params_t29); + PTO2TaskId up_spmd_tid_inline264__ssa_v4 = task_29_outs.task_id(); + up_tids_inline310[4] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[9] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[14] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[19] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[24] = up_spmd_tid_inline264__ssa_v4; + up_tids_inline310[29] = up_spmd_tid_inline264__ssa_v4; + for (int64_t n_out_inline275 = 6; n_out_inline275 < 17; n_out_inline275 += 1) { + int64_t n0_inline122 = (n_out_inline275 * 1024); + for (int64_t k_split_inline276 = 0; k_split_inline276 < 5; k_split_inline276 += 1) { + int64_t k0_inline113 = (k_split_inline276 * 1024); + PTO2TaskId _submit_deps_buf_inline102[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline102[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v135 = gate_late_tids_inline249[k_split_inline276]; + _submit_deps_buf_inline102[0] = t__tmp_v135; + + // Task 30: gate_proj_4 + L0TaskArgs params_t30; + params_t30.add_input(mlp_norm_in_inline71); + params_t30.add_input(ext_w_gate); + params_t30.add_inout(gate_acc_all_inline203); + params_t30.add_scalar(k0_inline113); + params_t30.add_scalar(layer_hidden_base_inline151); + params_t30.add_scalar(n0_inline122); + PTO2TaskId params_t30_deps[1]; + uint32_t params_t30_deps_count = 0; + if (_submit_deps_buf_inline102[0].is_valid()) + params_t30_deps[params_t30_deps_count++] = _submit_deps_buf_inline102[0]; + params_t30.set_dependencies(params_t30_deps, params_t30_deps_count); + TaskOutputTensors task_30_outs = rt_submit_aic_task(31, params_t30); + PTO2TaskId gate_tid_inline277 = task_30_outs.task_id(); + gate_tids_inline56[((n_out_inline275 * 5) + k_split_inline276)] = gate_tid_inline277; + PTO2TaskId _submit_deps_buf_inline246[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline246[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v136 = up_late_tids_inline69[k_split_inline276]; + _submit_deps_buf_inline246[0] = t__tmp_v136; + + // Task 31: up_proj_4 + L0TaskArgs params_t31; + params_t31.add_input(mlp_norm_in_inline71); + params_t31.add_input(ext_w_up); + params_t31.add_inout(up_acc_all_inline303); + params_t31.add_scalar(k0_inline113); + params_t31.add_scalar(layer_hidden_base_inline151); + params_t31.add_scalar(n0_inline122); + PTO2TaskId params_t31_deps[1]; + uint32_t params_t31_deps_count = 0; + if (_submit_deps_buf_inline246[0].is_valid()) + params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline246[0]; + params_t31.set_dependencies(params_t31_deps, params_t31_deps_count); + TaskOutputTensors task_31_outs = rt_submit_aic_task(32, params_t31); + PTO2TaskId up_tid_inline290 = task_31_outs.task_id(); + up_tids_inline310[((n_out_inline275 * 5) + k_split_inline276)] = up_tid_inline290; + } + } + for (int64_t n_out_inline292 = 0; n_out_inline292 < 17; n_out_inline292 += 1) { + int64_t n0_inline122__ssa_v7 = (n_out_inline292 * 1024); + PTO2TaskId _submit_deps_buf_inline167[11]; + for (int64_t __init_i = 0; __init_i < 11; ++__init_i) + _submit_deps_buf_inline167[__init_i] = PTO2TaskId::invalid(); + _submit_deps_buf_inline167[0] = reduce_tid_inline226; + PTO2TaskId t__tmp_v137 = gate_tids_inline56[(n_out_inline292 * 5)]; + _submit_deps_buf_inline167[1] = t__tmp_v137; + PTO2TaskId t__tmp_v138 = gate_tids_inline56[((n_out_inline292 * 5) + 1)]; + _submit_deps_buf_inline167[2] = t__tmp_v138; + PTO2TaskId t__tmp_v139 = gate_tids_inline56[((n_out_inline292 * 5) + 2)]; + _submit_deps_buf_inline167[3] = t__tmp_v139; + PTO2TaskId t__tmp_v140 = gate_tids_inline56[((n_out_inline292 * 5) + 3)]; + _submit_deps_buf_inline167[4] = t__tmp_v140; + PTO2TaskId t__tmp_v141 = gate_tids_inline56[((n_out_inline292 * 5) + 4)]; + _submit_deps_buf_inline167[5] = t__tmp_v141; + PTO2TaskId t__tmp_v142 = up_tids_inline310[(n_out_inline292 * 5)]; + _submit_deps_buf_inline167[6] = t__tmp_v142; + PTO2TaskId t__tmp_v143 = up_tids_inline310[((n_out_inline292 * 5) + 1)]; + _submit_deps_buf_inline167[7] = t__tmp_v143; + PTO2TaskId t__tmp_v144 = up_tids_inline310[((n_out_inline292 * 5) + 2)]; + _submit_deps_buf_inline167[8] = t__tmp_v144; + PTO2TaskId t__tmp_v145 = up_tids_inline310[((n_out_inline292 * 5) + 3)]; + _submit_deps_buf_inline167[9] = t__tmp_v145; + PTO2TaskId t__tmp_v146 = up_tids_inline310[((n_out_inline292 * 5) + 4)]; + _submit_deps_buf_inline167[10] = t__tmp_v146; + + // Task 32: silu + L0TaskArgs params_t32; + params_t32.add_input(inv_rms_tile_inline126); + params_t32.add_inout(mlp_tile_inline149); + params_t32.add_input(gate_acc_all_inline203); + params_t32.add_input(up_acc_all_inline303); + params_t32.add_scalar(n0_inline122__ssa_v7); + PTO2TaskId params_t32_deps[11]; + uint32_t params_t32_deps_count = 0; + if (_submit_deps_buf_inline167[0].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[0]; + if (_submit_deps_buf_inline167[1].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[1]; + if (_submit_deps_buf_inline167[2].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[2]; + if (_submit_deps_buf_inline167[3].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[3]; + if (_submit_deps_buf_inline167[4].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[4]; + if (_submit_deps_buf_inline167[5].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[5]; + if (_submit_deps_buf_inline167[6].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[6]; + if (_submit_deps_buf_inline167[7].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[7]; + if (_submit_deps_buf_inline167[8].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[8]; + if (_submit_deps_buf_inline167[9].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[9]; + if (_submit_deps_buf_inline167[10].is_valid()) + params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline167[10]; + params_t32.set_dependencies(params_t32_deps, params_t32_deps_count); + TaskOutputTensors task_32_outs = rt_submit_aiv_task(33, params_t32); + PTO2TaskId silu_tid_inline80 = task_32_outs.task_id(); + silu_tids_inline265[n_out_inline292] = silu_tid_inline80; + } + for (int64_t n_out_inline195 = 0; n_out_inline195 < 5; n_out_inline195 += 1) { + int64_t n0_inline122__ssa_v8 = (n_out_inline195 * 1024); + for (int64_t k_split_inline302 = 0; k_split_inline302 < 17; k_split_inline302 += 1) { + int64_t k0_inline113__ssa_v8 = (k_split_inline302 * 1024); + PTO2TaskId _submit_deps_buf_inline229[1]; + for (int64_t __init_i = 0; __init_i < 1; ++__init_i) + _submit_deps_buf_inline229[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v152 = silu_tids_inline265[k_split_inline302]; + _submit_deps_buf_inline229[0] = t__tmp_v152; + + // Task 33: down_proj + L0TaskArgs params_t33; + params_t33.add_input(mlp_tile_inline149); + params_t33.add_input(ext_w_down); + params_t33.add_inout(down_acc_all_inline168); + params_t33.add_scalar(k0_inline113__ssa_v8); + params_t33.add_scalar(layer_inter_base_inline107); + params_t33.add_scalar(n0_inline122__ssa_v8); + PTO2TaskId params_t33_deps[1]; + uint32_t params_t33_deps_count = 0; + if (_submit_deps_buf_inline229[0].is_valid()) + params_t33_deps[params_t33_deps_count++] = _submit_deps_buf_inline229[0]; + params_t33.set_dependencies(params_t33_deps, params_t33_deps_count); + params_t33.set_allow_early_resolve(true); + TaskOutputTensors task_33_outs = rt_submit_aic_task(34, params_t33); + PTO2TaskId down_tid_inline210 = task_33_outs.task_id(); + down_tids_inline156[((n_out_inline195 * 17) + k_split_inline302)] = down_tid_inline210; + } + } + } + PTO2TaskId _submit_deps_buf_inline123[85]; + for (int64_t __init_i = 0; __init_i < 85; ++__init_i) + _submit_deps_buf_inline123[__init_i] = PTO2TaskId::invalid(); + PTO2TaskId t__tmp_v153 = down_tids_inline156[0]; + _submit_deps_buf_inline123[0] = t__tmp_v153; + PTO2TaskId t__tmp_v154 = down_tids_inline156[1]; + _submit_deps_buf_inline123[1] = t__tmp_v154; + PTO2TaskId t__tmp_v155 = down_tids_inline156[2]; + _submit_deps_buf_inline123[2] = t__tmp_v155; + PTO2TaskId t__tmp_v156 = down_tids_inline156[3]; + _submit_deps_buf_inline123[3] = t__tmp_v156; + PTO2TaskId t__tmp_v157 = down_tids_inline156[4]; + _submit_deps_buf_inline123[4] = t__tmp_v157; + PTO2TaskId t__tmp_v158 = down_tids_inline156[5]; + _submit_deps_buf_inline123[5] = t__tmp_v158; + PTO2TaskId t__tmp_v159 = down_tids_inline156[6]; + _submit_deps_buf_inline123[6] = t__tmp_v159; + PTO2TaskId t__tmp_v160 = down_tids_inline156[7]; + _submit_deps_buf_inline123[7] = t__tmp_v160; + PTO2TaskId t__tmp_v161 = down_tids_inline156[8]; + _submit_deps_buf_inline123[8] = t__tmp_v161; + PTO2TaskId t__tmp_v162 = down_tids_inline156[9]; + _submit_deps_buf_inline123[9] = t__tmp_v162; + PTO2TaskId t__tmp_v163 = down_tids_inline156[10]; + _submit_deps_buf_inline123[10] = t__tmp_v163; + PTO2TaskId t__tmp_v164 = down_tids_inline156[11]; + _submit_deps_buf_inline123[11] = t__tmp_v164; + PTO2TaskId t__tmp_v165 = down_tids_inline156[12]; + _submit_deps_buf_inline123[12] = t__tmp_v165; + PTO2TaskId t__tmp_v166 = down_tids_inline156[13]; + _submit_deps_buf_inline123[13] = t__tmp_v166; + PTO2TaskId t__tmp_v167 = down_tids_inline156[14]; + _submit_deps_buf_inline123[14] = t__tmp_v167; + PTO2TaskId t__tmp_v168 = down_tids_inline156[15]; + _submit_deps_buf_inline123[15] = t__tmp_v168; + PTO2TaskId t__tmp_v169 = down_tids_inline156[16]; + _submit_deps_buf_inline123[16] = t__tmp_v169; + PTO2TaskId t__tmp_v170 = down_tids_inline156[17]; + _submit_deps_buf_inline123[17] = t__tmp_v170; + PTO2TaskId t__tmp_v171 = down_tids_inline156[18]; + _submit_deps_buf_inline123[18] = t__tmp_v171; + PTO2TaskId t__tmp_v172 = down_tids_inline156[19]; + _submit_deps_buf_inline123[19] = t__tmp_v172; + PTO2TaskId t__tmp_v173 = down_tids_inline156[20]; + _submit_deps_buf_inline123[20] = t__tmp_v173; + PTO2TaskId t__tmp_v174 = down_tids_inline156[21]; + _submit_deps_buf_inline123[21] = t__tmp_v174; + PTO2TaskId t__tmp_v175 = down_tids_inline156[22]; + _submit_deps_buf_inline123[22] = t__tmp_v175; + PTO2TaskId t__tmp_v176 = down_tids_inline156[23]; + _submit_deps_buf_inline123[23] = t__tmp_v176; + PTO2TaskId t__tmp_v177 = down_tids_inline156[24]; + _submit_deps_buf_inline123[24] = t__tmp_v177; + PTO2TaskId t__tmp_v178 = down_tids_inline156[25]; + _submit_deps_buf_inline123[25] = t__tmp_v178; + PTO2TaskId t__tmp_v179 = down_tids_inline156[26]; + _submit_deps_buf_inline123[26] = t__tmp_v179; + PTO2TaskId t__tmp_v180 = down_tids_inline156[27]; + _submit_deps_buf_inline123[27] = t__tmp_v180; + PTO2TaskId t__tmp_v181 = down_tids_inline156[28]; + _submit_deps_buf_inline123[28] = t__tmp_v181; + PTO2TaskId t__tmp_v182 = down_tids_inline156[29]; + _submit_deps_buf_inline123[29] = t__tmp_v182; + PTO2TaskId t__tmp_v183 = down_tids_inline156[30]; + _submit_deps_buf_inline123[30] = t__tmp_v183; + PTO2TaskId t__tmp_v184 = down_tids_inline156[31]; + _submit_deps_buf_inline123[31] = t__tmp_v184; + PTO2TaskId t__tmp_v185 = down_tids_inline156[32]; + _submit_deps_buf_inline123[32] = t__tmp_v185; + PTO2TaskId t__tmp_v186 = down_tids_inline156[33]; + _submit_deps_buf_inline123[33] = t__tmp_v186; + PTO2TaskId t__tmp_v187 = down_tids_inline156[34]; + _submit_deps_buf_inline123[34] = t__tmp_v187; + PTO2TaskId t__tmp_v188 = down_tids_inline156[35]; + _submit_deps_buf_inline123[35] = t__tmp_v188; + PTO2TaskId t__tmp_v189 = down_tids_inline156[36]; + _submit_deps_buf_inline123[36] = t__tmp_v189; + PTO2TaskId t__tmp_v190 = down_tids_inline156[37]; + _submit_deps_buf_inline123[37] = t__tmp_v190; + PTO2TaskId t__tmp_v191 = down_tids_inline156[38]; + _submit_deps_buf_inline123[38] = t__tmp_v191; + PTO2TaskId t__tmp_v192 = down_tids_inline156[39]; + _submit_deps_buf_inline123[39] = t__tmp_v192; + PTO2TaskId t__tmp_v193 = down_tids_inline156[40]; + _submit_deps_buf_inline123[40] = t__tmp_v193; + PTO2TaskId t__tmp_v194 = down_tids_inline156[41]; + _submit_deps_buf_inline123[41] = t__tmp_v194; + PTO2TaskId t__tmp_v195 = down_tids_inline156[42]; + _submit_deps_buf_inline123[42] = t__tmp_v195; + PTO2TaskId t__tmp_v196 = down_tids_inline156[43]; + _submit_deps_buf_inline123[43] = t__tmp_v196; + PTO2TaskId t__tmp_v197 = down_tids_inline156[44]; + _submit_deps_buf_inline123[44] = t__tmp_v197; + PTO2TaskId t__tmp_v198 = down_tids_inline156[45]; + _submit_deps_buf_inline123[45] = t__tmp_v198; + PTO2TaskId t__tmp_v199 = down_tids_inline156[46]; + _submit_deps_buf_inline123[46] = t__tmp_v199; + PTO2TaskId t__tmp_v200 = down_tids_inline156[47]; + _submit_deps_buf_inline123[47] = t__tmp_v200; + PTO2TaskId t__tmp_v201 = down_tids_inline156[48]; + _submit_deps_buf_inline123[48] = t__tmp_v201; + PTO2TaskId t__tmp_v202 = down_tids_inline156[49]; + _submit_deps_buf_inline123[49] = t__tmp_v202; + PTO2TaskId t__tmp_v203 = down_tids_inline156[50]; + _submit_deps_buf_inline123[50] = t__tmp_v203; + PTO2TaskId t__tmp_v204 = down_tids_inline156[51]; + _submit_deps_buf_inline123[51] = t__tmp_v204; + PTO2TaskId t__tmp_v205 = down_tids_inline156[52]; + _submit_deps_buf_inline123[52] = t__tmp_v205; + PTO2TaskId t__tmp_v206 = down_tids_inline156[53]; + _submit_deps_buf_inline123[53] = t__tmp_v206; + PTO2TaskId t__tmp_v207 = down_tids_inline156[54]; + _submit_deps_buf_inline123[54] = t__tmp_v207; + PTO2TaskId t__tmp_v208 = down_tids_inline156[55]; + _submit_deps_buf_inline123[55] = t__tmp_v208; + PTO2TaskId t__tmp_v209 = down_tids_inline156[56]; + _submit_deps_buf_inline123[56] = t__tmp_v209; + PTO2TaskId t__tmp_v210 = down_tids_inline156[57]; + _submit_deps_buf_inline123[57] = t__tmp_v210; + PTO2TaskId t__tmp_v211 = down_tids_inline156[58]; + _submit_deps_buf_inline123[58] = t__tmp_v211; + PTO2TaskId t__tmp_v212 = down_tids_inline156[59]; + _submit_deps_buf_inline123[59] = t__tmp_v212; + PTO2TaskId t__tmp_v213 = down_tids_inline156[60]; + _submit_deps_buf_inline123[60] = t__tmp_v213; + PTO2TaskId t__tmp_v214 = down_tids_inline156[61]; + _submit_deps_buf_inline123[61] = t__tmp_v214; + PTO2TaskId t__tmp_v215 = down_tids_inline156[62]; + _submit_deps_buf_inline123[62] = t__tmp_v215; + PTO2TaskId t__tmp_v216 = down_tids_inline156[63]; + _submit_deps_buf_inline123[63] = t__tmp_v216; + PTO2TaskId t__tmp_v217 = down_tids_inline156[64]; + _submit_deps_buf_inline123[64] = t__tmp_v217; + PTO2TaskId t__tmp_v218 = down_tids_inline156[65]; + _submit_deps_buf_inline123[65] = t__tmp_v218; + PTO2TaskId t__tmp_v219 = down_tids_inline156[66]; + _submit_deps_buf_inline123[66] = t__tmp_v219; + PTO2TaskId t__tmp_v220 = down_tids_inline156[67]; + _submit_deps_buf_inline123[67] = t__tmp_v220; + PTO2TaskId t__tmp_v221 = down_tids_inline156[68]; + _submit_deps_buf_inline123[68] = t__tmp_v221; + PTO2TaskId t__tmp_v222 = down_tids_inline156[69]; + _submit_deps_buf_inline123[69] = t__tmp_v222; + PTO2TaskId t__tmp_v223 = down_tids_inline156[70]; + _submit_deps_buf_inline123[70] = t__tmp_v223; + PTO2TaskId t__tmp_v224 = down_tids_inline156[71]; + _submit_deps_buf_inline123[71] = t__tmp_v224; + PTO2TaskId t__tmp_v225 = down_tids_inline156[72]; + _submit_deps_buf_inline123[72] = t__tmp_v225; + PTO2TaskId t__tmp_v226 = down_tids_inline156[73]; + _submit_deps_buf_inline123[73] = t__tmp_v226; + PTO2TaskId t__tmp_v227 = down_tids_inline156[74]; + _submit_deps_buf_inline123[74] = t__tmp_v227; + PTO2TaskId t__tmp_v228 = down_tids_inline156[75]; + _submit_deps_buf_inline123[75] = t__tmp_v228; + PTO2TaskId t__tmp_v229 = down_tids_inline156[76]; + _submit_deps_buf_inline123[76] = t__tmp_v229; + PTO2TaskId t__tmp_v230 = down_tids_inline156[77]; + _submit_deps_buf_inline123[77] = t__tmp_v230; + PTO2TaskId t__tmp_v231 = down_tids_inline156[78]; + _submit_deps_buf_inline123[78] = t__tmp_v231; + PTO2TaskId t__tmp_v232 = down_tids_inline156[79]; + _submit_deps_buf_inline123[79] = t__tmp_v232; + PTO2TaskId t__tmp_v233 = down_tids_inline156[80]; + _submit_deps_buf_inline123[80] = t__tmp_v233; + PTO2TaskId t__tmp_v234 = down_tids_inline156[81]; + _submit_deps_buf_inline123[81] = t__tmp_v234; + PTO2TaskId t__tmp_v235 = down_tids_inline156[82]; + _submit_deps_buf_inline123[82] = t__tmp_v235; + PTO2TaskId t__tmp_v236 = down_tids_inline156[83]; + _submit_deps_buf_inline123[83] = t__tmp_v236; + PTO2TaskId t__tmp_v237 = down_tids_inline156[84]; + _submit_deps_buf_inline123[84] = t__tmp_v237; + + // Spmd dcr_xgamma_spmd: dcr_xgamma + L0TaskArgs params_t34; + params_t34.add_input(down_acc_all_inline168); + params_t34.add_input(post_norm_partial_inline118); + params_t34.add_inout(next_hidden); + params_t34.add_input(ext_input_rms_weight); + params_t34.add_inout(next_normed); + params_t34.add_scalar(next_gamma_idx); + params_t34.launch_spec.set_block_num(5); + params_t34.set_allow_early_resolve(true); + PTO2TaskId params_t34_deps[85]; + uint32_t params_t34_deps_count = 0; + if (_submit_deps_buf_inline123[0].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[0]; + if (_submit_deps_buf_inline123[1].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[1]; + if (_submit_deps_buf_inline123[2].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[2]; + if (_submit_deps_buf_inline123[3].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[3]; + if (_submit_deps_buf_inline123[4].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[4]; + if (_submit_deps_buf_inline123[5].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[5]; + if (_submit_deps_buf_inline123[6].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[6]; + if (_submit_deps_buf_inline123[7].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[7]; + if (_submit_deps_buf_inline123[8].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[8]; + if (_submit_deps_buf_inline123[9].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[9]; + if (_submit_deps_buf_inline123[10].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[10]; + if (_submit_deps_buf_inline123[11].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[11]; + if (_submit_deps_buf_inline123[12].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[12]; + if (_submit_deps_buf_inline123[13].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[13]; + if (_submit_deps_buf_inline123[14].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[14]; + if (_submit_deps_buf_inline123[15].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[15]; + if (_submit_deps_buf_inline123[16].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[16]; + if (_submit_deps_buf_inline123[17].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[17]; + if (_submit_deps_buf_inline123[18].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[18]; + if (_submit_deps_buf_inline123[19].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[19]; + if (_submit_deps_buf_inline123[20].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[20]; + if (_submit_deps_buf_inline123[21].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[21]; + if (_submit_deps_buf_inline123[22].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[22]; + if (_submit_deps_buf_inline123[23].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[23]; + if (_submit_deps_buf_inline123[24].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[24]; + if (_submit_deps_buf_inline123[25].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[25]; + if (_submit_deps_buf_inline123[26].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[26]; + if (_submit_deps_buf_inline123[27].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[27]; + if (_submit_deps_buf_inline123[28].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[28]; + if (_submit_deps_buf_inline123[29].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[29]; + if (_submit_deps_buf_inline123[30].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[30]; + if (_submit_deps_buf_inline123[31].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[31]; + if (_submit_deps_buf_inline123[32].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[32]; + if (_submit_deps_buf_inline123[33].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[33]; + if (_submit_deps_buf_inline123[34].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[34]; + if (_submit_deps_buf_inline123[35].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[35]; + if (_submit_deps_buf_inline123[36].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[36]; + if (_submit_deps_buf_inline123[37].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[37]; + if (_submit_deps_buf_inline123[38].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[38]; + if (_submit_deps_buf_inline123[39].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[39]; + if (_submit_deps_buf_inline123[40].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[40]; + if (_submit_deps_buf_inline123[41].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[41]; + if (_submit_deps_buf_inline123[42].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[42]; + if (_submit_deps_buf_inline123[43].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[43]; + if (_submit_deps_buf_inline123[44].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[44]; + if (_submit_deps_buf_inline123[45].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[45]; + if (_submit_deps_buf_inline123[46].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[46]; + if (_submit_deps_buf_inline123[47].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[47]; + if (_submit_deps_buf_inline123[48].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[48]; + if (_submit_deps_buf_inline123[49].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[49]; + if (_submit_deps_buf_inline123[50].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[50]; + if (_submit_deps_buf_inline123[51].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[51]; + if (_submit_deps_buf_inline123[52].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[52]; + if (_submit_deps_buf_inline123[53].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[53]; + if (_submit_deps_buf_inline123[54].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[54]; + if (_submit_deps_buf_inline123[55].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[55]; + if (_submit_deps_buf_inline123[56].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[56]; + if (_submit_deps_buf_inline123[57].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[57]; + if (_submit_deps_buf_inline123[58].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[58]; + if (_submit_deps_buf_inline123[59].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[59]; + if (_submit_deps_buf_inline123[60].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[60]; + if (_submit_deps_buf_inline123[61].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[61]; + if (_submit_deps_buf_inline123[62].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[62]; + if (_submit_deps_buf_inline123[63].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[63]; + if (_submit_deps_buf_inline123[64].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[64]; + if (_submit_deps_buf_inline123[65].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[65]; + if (_submit_deps_buf_inline123[66].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[66]; + if (_submit_deps_buf_inline123[67].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[67]; + if (_submit_deps_buf_inline123[68].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[68]; + if (_submit_deps_buf_inline123[69].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[69]; + if (_submit_deps_buf_inline123[70].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[70]; + if (_submit_deps_buf_inline123[71].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[71]; + if (_submit_deps_buf_inline123[72].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[72]; + if (_submit_deps_buf_inline123[73].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[73]; + if (_submit_deps_buf_inline123[74].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[74]; + if (_submit_deps_buf_inline123[75].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[75]; + if (_submit_deps_buf_inline123[76].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[76]; + if (_submit_deps_buf_inline123[77].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[77]; + if (_submit_deps_buf_inline123[78].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[78]; + if (_submit_deps_buf_inline123[79].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[79]; + if (_submit_deps_buf_inline123[80].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[80]; + if (_submit_deps_buf_inline123[81].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[81]; + if (_submit_deps_buf_inline123[82].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[82]; + if (_submit_deps_buf_inline123[83].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[83]; + if (_submit_deps_buf_inline123[84].is_valid()) + params_t34_deps[params_t34_deps_count++] = _submit_deps_buf_inline123[84]; + params_t34.set_dependencies(params_t34_deps, params_t34_deps_count); + TaskOutputTensors task_34_outs = rt_submit_aiv_task(35, params_t34); + PTO2TaskId dcr_tid_inline58 = task_34_outs.task_id(); + prev_out_tid[0] = dcr_tid_inline58; + prev_normed_tid[0] = dcr_tid_inline58; + Tensor cur__ssa_v8 = next_hidden; + Tensor normed__ssa_v6 = next_normed; + cur__rv_v7 = cur__ssa_v8; + normed__rv_v5 = normed__ssa_v6; + } + } + for (int64_t ob0 = 0; ob0 < 16; ob0 += 16) { + PTO2_SCOPE() { + // Task 35: copy_out + L0TaskArgs params_t35; + params_t35.add_output(ext_out); + params_t35.add_input(cur__rv_v7); + params_t35.add_scalar(ob0); + rt_submit_aiv_task(36, params_t35); + } + } + } +} + +} // extern "C" diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp new file mode 100644 index 0000000000..cadbef80ef --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp @@ -0,0 +1,44 @@ +/* + * 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 "tensor.h" + +#ifdef __CPU_SIM +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; } + +#else + +#include "../kernel/fai_body.hpp" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + GM_ADDR metadata = tensor_data(args, 6); + acquire_qwen_fai_metadata(metadata); + __gm__ const FAInferTilingData *tiling = reinterpret_cast<__gm__ const FAInferTilingData *>(metadata); + if (tiling->needCoreNum != 0) { + uint64_t raw_barrier = reinterpret_cast(metadata + qwen_fai_metadata::kBarrierAlignmentOffset); + uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & + ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); + run_qwen_fai(args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier)); + } else { + run_qwen_fai(args); + } +} + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp new file mode 100644 index 0000000000..c087c9f179 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp @@ -0,0 +1,50 @@ +/* + * 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 "tensor.h" + +#ifdef __CPU_SIM +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; } + +#else + +#include "../kernel/fai_body.hpp" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + GM_ADDR metadata = tensor_data(args, 6); + acquire_qwen_fai_metadata(metadata); + __gm__ const FAInferTilingData *tiling = + reinterpret_cast<__gm__ const FAInferTilingData *>(metadata); + if (tiling->needCoreNum != 0) { + uint64_t raw_barrier = reinterpret_cast( + metadata + qwen_fai_metadata::kBarrierAlignmentOffset); + uint64_t aligned_barrier = + (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & + ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); + run_qwen_fai( + args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier)); + } else { + run_qwen_fai(args); + } +} + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/generated/kernel_tiling/kernel_tiling.h b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/generated/kernel_tiling/kernel_tiling.h new file mode 100644 index 0000000000..6ab3cf249d --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/generated/kernel_tiling/kernel_tiling.h @@ -0,0 +1,94 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * Modifications 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef PYPTO_QWEN_FAI_KERNEL_TILING_H +#define PYPTO_QWEN_FAI_KERNEL_TILING_H + +#include +#include + +// Device-side mirror of flash_attention_infer_tiling.h. +struct coreNode { + int32_t startBIdx[26]; + int32_t startN1Idx[26]; + int32_t startS1Idx[26]; + int32_t startS2Idx[26]; + int32_t endBIdx[26]; + int32_t endN1Idx[26]; + int32_t endS1Idx[26]; + int32_t endS2Idx[26]; + int64_t firstSplitKVTaskLseOffset[26]; + int64_t firstSplitKVTaskOOffset[26]; +}; + +struct splitNode { + int32_t batchIdx[26]; + int32_t headStartIdx[26]; + int32_t headEndIdx[26]; + int32_t qStartIdx[26]; + int32_t qEndIdx[26]; + int32_t splitNum[26]; + int64_t lseTaskOffset[26]; + int64_t oTaskOffset[26]; +}; + +struct FAInferTilingData { + uint32_t numHeads; + uint32_t embeddingSize; + uint32_t embeddingSizeV; + uint32_t numBlocks; + uint32_t blockSize; + uint32_t maxQSeqlen; + uint32_t maxKvSeqlen; + uint32_t kvHeads; + uint32_t batch; + uint32_t maxNumBlocksPerBatch; + uint32_t firstBatchTaskNum; + uint32_t totalTaskNum; + uint32_t maskType; + uint32_t _pad_before_workspace_sizes; + uint64_t mm1OutSize; + uint64_t smOnlineOutSize; + uint64_t mm2OutSize; + uint64_t UpdateSize; + uint64_t workSpaceSize; + float scaleValue; + uint32_t _pad_before_pse; + uint64_t pseQ; + uint64_t pseKv; + uint32_t padding3; + uint32_t _pad_before_tokens; + int64_t preToken; + int64_t nextToken; + uint32_t sparseMode; + uint32_t _pad_before_split_sizes; + uint64_t splitLseTotalSize; + uint64_t splitOTotalSize; + uint32_t totalSplitNodeNum; + uint32_t needCoreNum; + uint32_t mainLoopTaskNum; + uint32_t tailLoopTaskNum; + uint32_t tailStartBatch; + uint32_t tailStartN2; + uint32_t tailKvNBlockTile; + uint32_t _pad_before_nodes; + coreNode coreInfo; + splitNode splitInfo; +}; + +static_assert(sizeof(coreNode) == 1248, "coreNode ABI mismatch"); +static_assert(sizeof(splitNode) == 1040, "splitNode ABI mismatch"); +static_assert(offsetof(FAInferTilingData, mm1OutSize) == 56, "FAInfer scalar ABI mismatch"); +static_assert(offsetof(FAInferTilingData, coreInfo) == 200, "FAInfer node ABI mismatch"); +static_assert(sizeof(FAInferTilingData) == 2488, "FAInferTilingData ABI mismatch"); + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp new file mode 100644 index 0000000000..4381fe0660 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp @@ -0,0 +1,240 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef PYPTO_QWEN_FAI_BODY_HPP +#define PYPTO_QWEN_FAI_BODY_HPP + +#include +#include + +#ifndef TILING_KEY_VAR +#define TILING_KEY_VAR 0 +#endif +#ifndef ASC_DEVKIT_MAJOR +#define ASC_DEVKIT_MAJOR 9 +#define ASC_DEVKIT_MINOR 0 +#define ASC_DEVKIT_PATCH 0 +#define ASC_DEVKIT_VERSION_NUM 90000000 +#endif + +#include "intrinsic.h" +#include "tensor.h" + +#include "../generated/kernel_tiling/kernel_tiling.h" +#include "metadata_layout.h" + +#include "../vendor/fused_infer_attention_score/flash_attention_regular.h" + +#include "rope_qkv_generated.hpp" + +constexpr uint64_t kQwenFaiHeadDim = 128; +// The generated RoPE body is specialized to the standalone 32-lane dispatch. +constexpr uint32_t kQwenRopeCores = 32; + +// Global cube<->vector barrier between phase-0 RoPE and the attention phase. +// The FFTS flag-region base is set by the simpler runtime at launch. +// AscendC::SyncAll is the fused (mixed AIC+AIV) all-core barrier; the +// default SyncAll() is AIV-only and never releases the Cube cores. +static __aicore__ __attribute__((always_inline)) void qwen_fai_syncall_mix() { + AscendC::PipeBarrier(); + // isAIVOnly=false: fused Cube+Vector whole-core barrier. + AscendC::SyncAll(); +} + +static __aicore__ __attribute__((always_inline)) void +acquire_qwen_fai_metadata(GM_ADDR metadata) { + uint64_t first_line = + reinterpret_cast(metadata) & + ~(static_cast(qwen_fai_metadata::kDcciLineBytes) - 1); + uint64_t end = reinterpret_cast(metadata) + + qwen_fai_metadata::kBarrierAlignmentOffset; + for (uint64_t line = first_line; line < end; + line += qwen_fai_metadata::kDcciLineBytes) { + dcci(reinterpret_cast<__gm__ void *>(line), SINGLE_CACHE_LINE); + } + dsb(DSB_DDR); +} + +template +static __aicore__ __attribute__((always_inline)) GM_ADDR +tensor_data(__gm__ int64_t *args, int32_t index) { + __gm__ Tensor *tensor = reinterpret_cast<__gm__ Tensor *>(args[index]); + __gm__ T *data = + reinterpret_cast<__gm__ T *>(tensor->buffer.addr) + tensor->start_offset; + return reinterpret_cast(data); +} + +template +static __aicore__ __attribute__((always_inline)) void +run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { + using namespace NpuArch; + using namespace KernelCommon; + + using ElementQ = bfloat16_t; + using ElementK = bfloat16_t; + using ElementV = bfloat16_t; + using ElementS = float; + using ElementP = bfloat16_t; + using ElementO = bfloat16_t; + using ElementLse = float; + using ElementMask = int8_t; + using ElementOTmp = float; + using ElementUpdate = float; + using ElementSink = bfloat16_t; + + using LayoutQ = layout::RowMajor; + using LayoutK = layout::ColumnMajor; + using LayoutV = layout::RowMajor; + using LayoutS = layout::RowMajor; + using LayoutP = layout::RowMajor; + using LayoutO = layout::RowMajor; + using LayoutLse = layout::RowMajor; + using LayoutMask = layout::RowMajor; + using LayoutOTmp = layout::RowMajor; + using LayoutUpdate = layout::RowMajor; + using LayoutSink = layout::RowMajor; + + using L1TileShapeQK = GemmShape; + using L0TileShapeQK = GemmShape<128, 128, 128>; + using DispatchPolicyQK = Gemm::MmadAtlasA2FAIQK; + using QType = Gemm::GemmType; + using KType = Gemm::GemmType; + using SType = Gemm::GemmType; + using SinkType = Gemm::GemmType; + using BlockMmadQK = + Gemm::Block::BlockMmad; + + using PType = Gemm::GemmType; + using MaskType = Gemm::GemmType; + using PseShiftType = Gemm::GemmType; + using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueAtlasA2OnlineSoftmax< + Epilogue::LseMode::NONE, Epilogue::SinkMode::DISABLE, + static_cast(FaiKernel::MaskType::NO_MASK), float>; + using EpilogueOnlineSoftmax = + Epilogue::Block::BlockEpilogue; + + using L1TileShapePV = GemmShape<128, 128, 256>; + using L0TileShapePV = GemmShape<128, 128, 128>; + using DispatchPolicyPV = Gemm::MmadAtlasA2FAIPV; + using VType = Gemm::GemmType; + using OTmpType = Gemm::GemmType; + using BlockMmadPV = + Gemm::Block::BlockMmad; + + using OType = Gemm::GemmType; + using OUpdateType = Gemm::GemmType; + using LseType = Gemm::GemmType; + using DispatchPolicyRescaleO = + Epilogue::EpilogueAtlasA2RescaleO; + using EpilogueRescaleO = + Epilogue::Block::BlockEpilogue; + using DispatchPolicyInitOut = + Epilogue::EpilogueAtlasA2InitOutWhenZero; + using EpilogueInitOut = + Epilogue::Block::BlockEpilogue; + using CombineScale = Epilogue::Block::CombineScale; + + using FdKernel = SplitFuse::FAInferKernel< + BlockMmadQK, BlockMmadPV, EpilogueOnlineSoftmax, EpilogueRescaleO, + EpilogueInitOut, true, FaiKernel::MaskType::NO_MASK, + FaiKernel::inputLayout::TND, CombineScale, true, true, true>; + using NonFdKernel = + SplitFuse::FAInferKernel; + using Kernel = std::conditional_t; + + GM_ADDR metadata = tensor_data(args, 6); + // pypto packs tensors first, then the sole scalar last: the rope-fused ABI + // has 17 tensors so cache_row_offset is at args[17]; the attention-only ABI + // has 7 tensors so it is at args[7]. + uint64_t cache_row_offset = + static_cast(WithRope ? args[17] : args[7]); + uint64_t cache_byte_offset = + cache_row_offset * kQwenFaiHeadDim * sizeof(uint16_t); + constexpr int32_t query_arg = WithRope ? 1 : 0; + constexpr int32_t key_arg = WithRope ? 2 : 1; + constexpr int32_t value_arg = WithRope ? 3 : 2; + constexpr int32_t block_table_arg = WithRope ? 4 : 3; + constexpr int32_t out_arg = WithRope ? 0 : 4; + GM_ADDR key = tensor_data(args, key_arg) + cache_byte_offset; + GM_ADDR value = tensor_data(args, value_arg) + cache_byte_offset; + + FAIKernelParams params{tensor_data(args, query_arg), + key, + value, + nullptr, + nullptr, + tensor_data(args, block_table_arg), + metadata + qwen_fai_metadata::kCumulativeQOffset, + metadata + qwen_fai_metadata::kKvLengthsOffset, + tensor_data(args, out_arg), + nullptr, + tensor_data(args, 5), + metadata + qwen_fai_metadata::kTilingOffset, + nullptr}; + + uint32_t sub_block_idx = 0; +#ifdef __DAV_C220_VEC__ + sub_block_idx = static_cast(get_sub_block_id(args)); +#endif + uint32_t block_idx = static_cast(get_block_idx(args)); + uint32_t block_num = static_cast(get_block_num(args)); + + // Fold QK-norm + RoPE in as phase 0: the AIV lanes rotate Q/K and publish + // paged K plus projected V, then a global cube<->vec FFTS barrier makes those + // GM writes visible to every core before the attention phase reads them. + if constexpr (WithRope) { +#ifdef __DAV_C220_VEC__ + // Drive the golden-correct pypto-generated rope_qkv (copied verbatim). The + // fused ABI packs 17 tensors then the sole scalar; map them to the + // generated parameter order (k_cache, q_tnd, v_cache, seq_lens, inv_rms, + // slot_mapping, rope_cos, rope_sin, k_proj, k_norm_w, v_proj, q_proj, + // q_norm_w, layer_cache_base, KV_CACHE_ROWS, block_idx, block_num). + int64_t kv_cache_rows = static_cast( + reinterpret_cast<__gm__ Tensor *>(args[2])->shapes[0]); + uint32_t rope_lane = block_idx * 2 + sub_block_idx; + if (rope_lane < kQwenRopeCores) { + qwen_rope_gen::rope_qkv( + reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, 2)), + reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, 1)), + reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, 3)), + reinterpret_cast<__gm__ int32_t *>(tensor_data(args, 16)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 14)), + reinterpret_cast<__gm__ int32_t *>(tensor_data(args, 15)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 12)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 13)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 8)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 11)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 9)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 7)), + reinterpret_cast<__gm__ float *>(tensor_data(args, 10)), + static_cast(args[17]), kv_cache_rows, + static_cast(rope_lane), + static_cast(kQwenRopeCores)); + } +#endif + qwen_fai_syncall_mix(); + } + + Arch::PtoTopology topology{block_idx, block_num, sub_block_idx, 2}; + Kernel kernel; + kernel(params, topology, barrier_state); +} + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h new file mode 100644 index 0000000000..b21f1991ca --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h @@ -0,0 +1,40 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef PYPTO_QWEN_FAI_METADATA_LAYOUT_H +#define PYPTO_QWEN_FAI_METADATA_LAYOUT_H + +#include + +namespace qwen_fai_metadata { + +constexpr uint32_t kTilingOffset = 0; +constexpr uint32_t kTilingBytes = 2488; +constexpr uint32_t kCumulativeQOffset = 2488; +constexpr uint32_t kLengthArrayBytes = 16 * sizeof(int64_t); +constexpr uint32_t kKvLengthsOffset = kCumulativeQOffset + kLengthArrayBytes; +constexpr uint32_t kBarrierAlignmentOffset = kKvLengthsOffset + kLengthArrayBytes; +constexpr uint32_t kDcciLineBytes = 64; +constexpr uint32_t kBarrierAlignmentBytes = 512; +constexpr uint32_t kBarrierSlotBytes = 512; +constexpr uint32_t kBarrierSlotWords = kBarrierSlotBytes / sizeof(int32_t); +constexpr uint32_t kBarrierSlotCount = 48; +constexpr uint32_t kBarrierBytes = kBarrierSlotCount * kBarrierSlotBytes; +constexpr uint32_t kMetadataBytes = 27840; + +static_assert( + kBarrierAlignmentOffset + kBarrierAlignmentBytes - 1 + kBarrierBytes <= kMetadataBytes, + "metadata buffer does not cover the aligned barrier region" +); + +} // namespace qwen_fai_metadata + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/rope_qkv_generated.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/rope_qkv_generated.hpp new file mode 100644 index 0000000000..66ea7fedb4 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/rope_qkv_generated.hpp @@ -0,0 +1,1246 @@ +// Copyright (c) PyPTO Contributors. CANN Open Software License Agreement v2.0. +// RoPE + QK-norm prologue: the pypto/ptoas-GENERATED rope_qkv kernel, copied +// verbatim from the decode_fwd rope_qkv scope codegen (golden-correct) and +// wrapped for reuse inside the fused paged_attention_rope_cce extern. Do not +// hand-edit the generated body; regenerate from decode_fwd rope_qkv if the +// math changes. VEC-only (pto Vec tiles); the caller guards the invocation. +#ifndef PYPTO_QWEN_ROPE_QKV_GENERATED_HPP +#define PYPTO_QWEN_ROPE_QKV_GENERATED_HPP + +#include + +#ifdef __DAV_C220_VEC__ +#include +#include "tensor.h" +#include "intrinsic.h" + +namespace qwen_rope_gen { +using namespace pto; + +enum class PTOAutoSyncTailMode : int { + kBarrierAll = 0, + kSetWaitMte3ToSEvent0 = 1, +}; + +static __aicore__ inline void ptoas_auto_sync_tail( + PTOAutoSyncTailMode mode = PTOAutoSyncTailMode::kBarrierAll) { + switch (mode) { + case PTOAutoSyncTailMode::kSetWaitMte3ToSEvent0: + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID0); + break; + case PTOAutoSyncTailMode::kBarrierAll: + default: + pipe_barrier(PIPE_ALL); + break; + } +} + +template +static __aicore__ inline void PTOAS__DCCI_SINGLE_CACHE_LINE(Ptr ptr) { + dcci((__gm__ void*)ptr, cache_line_t::SINGLE_CACHE_LINE); +} + +static __aicore__ void rope_qkv(__gm__ bfloat16_t* v1, __gm__ bfloat16_t* v2, __gm__ bfloat16_t* v3, __gm__ int32_t* v4, __gm__ float* v5, __gm__ int32_t* v6, __gm__ float* v7, __gm__ float* v8, __gm__ float* v9, __gm__ float* v10, __gm__ float* v11, __gm__ float* v12, __gm__ float* v13, int64_t v14, int64_t v15, int32_t v16, int32_t v17) { + SaturationMode v18 = SaturationMode::OFF; + RoundMode v19 = RoundMode::CAST_ROUND; + const int64_t v20 = 2048; + const float v21 = 9.99999997E-7f; + const float v22 = 0.0078125f; + const int64_t v23 = 64; + const int64_t v24 = 40; + const int64_t v25 = 5; + const int64_t v26 = 8; + const int64_t v27 = 32; + const int64_t v28 = 2; + const int64_t v29 = 4; + const int64_t v30 = 896; + const float v31 = 0.0f; + const int64_t v32 = 1408; + const int64_t v33 = 5120; + const int64_t v34 = 1024; + const int64_t v35 = 16; + const int64_t v36 = 640; + const int64_t v37 = 1; + const int64_t v38 = 128; + const int64_t v39 = 63232; + const int64_t v40 = 62720; + const int64_t v41 = 8192; + const int64_t v42 = 62976; + const int64_t v43 = 16384; + const int64_t v44 = 62208; + const int64_t v45 = 0; + const int64_t v46 = 61952; + const int64_t v47 = 61696; + const int64_t v48 = 61440; + const int64_t v49 = 61184; + const int64_t v50 = 32768; + const int64_t v51 = 32256; + const int64_t v52 = 48896; + const int64_t v53 = 32512; + const int64_t v54 = 57088; + const int64_t v55 = 31744; + const int64_t v56 = 40704; + const int64_t v57 = 31488; + const int64_t v58 = 31232; + const int64_t v59 = 30976; + const int64_t v60 = 30720; + const int64_t v61 = 27136; + const int64_t v62 = 21504; + const int64_t v63 = 20992; + const int64_t v64 = 20480; + using T = float; + + #if defined(__DAV_VEC__) + set_mask_norm(); + set_vector_mask(-1, -1); + // pto: %k_norm_w_inline55__tile + Tile v65 = Tile(v37, v38); + // pto: %k_norm_w_inline55__tile + uint64_t v66 = (uint64_t) v64; + TASSIGN(v65, v66); + // pto: %k_norm_w_inline55__ssa_v0_pview + pto::Shape<1, 1, 1, 1, 128> v67 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %k_norm_w_inline55__ssa_v0_pview + pto::Stride<128, 128, 128, 128, 1> v68 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %k_norm_w_inline55__ssa_v0_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v69 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v10 + ((v45 + v45 * v38) + v45 * v37), v67, v68); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID4); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID5); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID6); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TLOAD(v65, v69); + // pto: %q_norm_w_inline246__tile + Tile v70 = Tile(v37, v38); + // pto: %q_norm_w_inline246__tile + uint64_t v71 = (uint64_t) v63; + TASSIGN(v70, v71); + // pto: %q_norm_w_inline246__ssa_v0_pview + pto::Shape<1, 1, 1, 1, 128> v72 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %q_norm_w_inline246__ssa_v0_pview + pto::Stride<128, 128, 128, 128, 1> v73 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %q_norm_w_inline246__ssa_v0_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v74 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v13 + ((v45 + v45 * v38) + v45 * v37), v72, v73); + TLOAD(v70, v74); + // pto: %rope_core_inline264__tile + int64_t v75 = (int64_t) v16; + // pto: %q_red_pad_inline73__tile + Tile v76 = Tile(v37, v32); + // pto: %q_red_pad_inline73__tile + uint64_t v77 = (uint64_t) v62; + TASSIGN(v76, v77); + TEXPANDS(v76, v31); + // pto: %k_red_pad_inline268__tile + Tile v78 = Tile(v37, v30); + // pto: %k_red_pad_inline268__tile + uint64_t v79 = (uint64_t) v61; + TASSIGN(v78, v79); + TEXPANDS(v78, v31); + for (int64_t i80 = v45; i80 < v29; i80 += v28) { + // pto: %106 + int64_t v81 = (int64_t) ((uint64_t) i80 * (uint64_t) v27); + // pto: %107 + int64_t v82 = (int64_t) ((uint64_t) v75 + (uint64_t) v81); + // pto: %110, %109 + int64_t v83 = (int64_t) ((uint64_t) v75 + (uint64_t) ((int64_t) ((uint64_t) v81 + (uint64_t) v27))); + // pto: %111 + if (v82 < v38) { + // pto: %112 + int64_t v84 = v82 / v35; + // pto: %114, %113 + int64_t v85 = (int64_t) ((uint64_t) v82 - (uint64_t) ((int64_t) ((uint64_t) v84 * (uint64_t) v35))); + // pto: %ctx_len_inline54__tile + int32_t v86 = v4[v85]; + // pto: %inv_rms_b_inline212__tile + float v87 = v5[v85]; + // pto: %115, %116 + int64_t v88 = (int64_t) ((uint64_t) ((int64_t) v86) - (uint64_t) v37); + // pto: %117 + int32_t v89 = v6[v85]; + // pto: %122 + int64_t v90 = (int64_t) ((uint64_t) v84 * (uint64_t) v38); + // pto: %126, %118, %125, %127 + int64_t v91 = (int64_t) ((uint64_t) ((int64_t) ((uint64_t) v14 + (uint64_t) ((int64_t) ((uint64_t) ((int64_t) v89) * (uint64_t) v26)))) + (uint64_t) v84); + // pto: %cos_lo_inline82__tile + Tile v92 = Tile(v37, v23); + // pto: %cos_lo_inline82__tile + uint64_t v93 = (uint64_t) v60; + TASSIGN(v92, v93); + // pto: %rope_cos__ssa_v0_pview + pto::Shape<1, 1, 1, 1, 64> v94 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %rope_cos__ssa_v0_pview + pto::Stride<128, 128, 128, 128, 1> v95 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %rope_cos__ssa_v0_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v96 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v7 + ((v45 + v88 * v38) + v45 * v37), v94, v95); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v92, v96); + // pto: %cos_hi_inline158__tile + Tile v97 = Tile(v37, v23); + // pto: %cos_hi_inline158__tile + uint64_t v98 = (uint64_t) v59; + TASSIGN(v97, v98); + // pto: %131 + pto::Shape<1, 1, 1, 1, 64> v99 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %131 + pto::Stride<128, 128, 128, 128, 1> v100 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %131 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v101 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v7 + ((v45 + v88 * v38) + v23 * v37), v99, v100); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v97, v101); + // pto: %sin_lo_inline57__tile + Tile v102 = Tile(v37, v23); + // pto: %sin_lo_inline57__tile + uint64_t v103 = (uint64_t) v58; + TASSIGN(v102, v103); + // pto: %rope_sin__ssa_v0_pview + pto::Shape<1, 1, 1, 1, 64> v104 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %rope_sin__ssa_v0_pview + pto::Stride<128, 128, 128, 128, 1> v105 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %rope_sin__ssa_v0_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v106 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v8 + ((v45 + v88 * v38) + v45 * v37), v104, v105); + TLOAD(v102, v106); + // pto: %sin_hi_inline88__tile + Tile v107 = Tile(v37, v23); + // pto: %sin_hi_inline88__tile + uint64_t v108 = (uint64_t) v57; + TASSIGN(v107, v108); + // pto: %132 + pto::Shape<1, 1, 1, 1, 64> v109 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %132 + pto::Stride<128, 128, 128, 128, 1> v110 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %132 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v111 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v8 + ((v45 + v88 * v38) + v23 * v37), v109, v110); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + TLOAD(v107, v111); + // pto: %t__tile + Tile v112 = Tile(v37, v38); + // pto: %t__tile + uint64_t v113 = (uint64_t) v56; + TASSIGN(v112, v113); + // pto: %k_proj_inline145__rv_v3_pview + pto::Shape<1, 1, 1, 1, 128> v114 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %k_proj_inline145__rv_v3_pview + pto::Stride<1024, 1024, 1024, 1024, 1> v115 = pto::Stride<1024, 1024, 1024, 1024, 1>(); + // pto: %k_proj_inline145__rv_v3_pview + GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v116 = GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>(v9 + ((v45 + v85 * v34) + v90 * v37), v114, v115); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + TLOAD(v112, v116); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + // pto: %0 + Tile v117 = Tile(v37, v38); + // pto: %0 + uint64_t v118 = (uint64_t) v55; + TASSIGN(v117, v118); + // pto: %v_proj_inline207__rv_v3_pview + pto::Shape<1, 1, 1, 1, 128> v119 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %v_proj_inline207__rv_v3_pview + pto::Stride<1024, 1024, 1024, 1024, 1> v120 = pto::Stride<1024, 1024, 1024, 1024, 1>(); + // pto: %v_proj_inline207__rv_v3_pview + GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v121 = GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>(v11 + ((v45 + v85 * v34) + v90 * v37), v119, v120); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + TLOAD(v117, v121); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + // pto: %1 + Tile v122 = Tile(v37, v36); + // pto: %1 + uint64_t v123 = (uint64_t) v54; + TASSIGN(v122, v123); + // pto: %q_proj_inline157__rv_v5_pview + pto::Shape<1, 1, 1, 1, 640> v124 = pto::Shape<1, 1, 1, 1, 640>(); + // pto: %q_proj_inline157__rv_v5_pview + pto::Stride<5120, 5120, 5120, 5120, 1> v125 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + // pto: %q_proj_inline157__rv_v5_pview + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v126 = GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>(v12 + ((v45 + v85 * v33) + (int64_t) ((uint64_t) v84 * (uint64_t) v36) * v37), v124, v125); + TLOAD(v122, v126); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + // pto: %2 + Tile v127 = Tile(v37, v34); + // pto: %2 + uint64_t v128 = (uint64_t) v53; + TASSIGN(v127, v128); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + TCONCAT(v127, v112, v78); + // pto: %3 + Tile v129 = Tile(v26, v38); + // pto: %3 + uint64_t v130 = (uint64_t) v53; + TASSIGN(v129, v130); + // pto: %k_raw_inline122__tile + Tile v131 = Tile(v26, v38); + // pto: %k_raw_inline122__tile + uint64_t v132 = (uint64_t) v53; + TASSIGN(v131, v132); + pipe_barrier(PIPE_V); + TMULS(v131, v129, v87); + // pto: %4 + Tile v133 = Tile(v26, v38); + // pto: %4 + uint64_t v134 = (uint64_t) v56; + TASSIGN(v133, v134); + pipe_barrier(PIPE_V); + TMUL(v133, v131, v131); + // pto: %tmp_tile + Tile v135 = Tile(v26, v38); + // pto: %tmp_tile + uint64_t v136 = (uint64_t) v52; + TASSIGN(v135, v136); + // pto: %k_ss_inline196__tile + Tile v137 = Tile(v26, v37); + // pto: %k_ss_inline196__tile + uint64_t v138 = (uint64_t) v51; + TASSIGN(v137, v138); + pipe_barrier(PIPE_V); + TROWSUM(v137, v133, v135); + // pto: %t__rm_a0_tmp_v0 + Tile v139 = Tile(v37, v26); + // pto: %t__rm_a0_tmp_v0 + uint64_t v140 = (uint64_t) v51; + TASSIGN(v139, v140); + // pto: %t__row_major_tmp_v1 + Tile v141 = Tile(v37, v26); + // pto: %t__row_major_tmp_v1 + uint64_t v142 = (uint64_t) v56; + TASSIGN(v141, v142); + pipe_barrier(PIPE_V); + TMULS(v141, v139, v22); + // pto: %t__rm_a0_tmp_v2 + Tile v143 = Tile(v37, v26); + // pto: %t__rm_a0_tmp_v2 + uint64_t v144 = (uint64_t) v56; + TASSIGN(v143, v144); + // pto: %t__row_major_tmp_v3 + Tile v145 = Tile(v37, v26); + // pto: %t__row_major_tmp_v3 + uint64_t v146 = (uint64_t) v56; + TASSIGN(v145, v146); + pipe_barrier(PIPE_V); + TADDS(v145, v143, v21); + // pto: %t__rm_a0_tmp_v4 + Tile v147 = Tile(v37, v26); + // pto: %t__rm_a0_tmp_v4 + uint64_t v148 = (uint64_t) v56; + TASSIGN(v147, v148); + // pto: %t__row_major_tmp_v5 + Tile v149 = Tile(v37, v26); + // pto: %t__row_major_tmp_v5 + uint64_t v150 = (uint64_t) v56; + TASSIGN(v149, v150); + pipe_barrier(PIPE_V); + TSQRT(v149, v147); + // pto: %k_inv_inline78__rm_a0_tmp_v6 + Tile v151 = Tile(v37, v26); + // pto: %k_inv_inline78__rm_a0_tmp_v6 + uint64_t v152 = (uint64_t) v56; + TASSIGN(v151, v152); + // pto: %k_inv_inline78__row_major_tmp_v7 + Tile v153 = Tile(v37, v26); + // pto: %k_inv_inline78__row_major_tmp_v7 + uint64_t v154 = (uint64_t) v52; + TASSIGN(v153, v154); + pipe_barrier(PIPE_V); + TRECIP(v153, v151); + // pto: %k_inv_inline78__tile + Tile v155 = Tile(v26, v37); + // pto: %k_inv_inline78__tile + uint64_t v156 = (uint64_t) v52; + TASSIGN(v155, v156); + // pto: %8 + Tile v157 = Tile(v26, v38); + // pto: %8 + uint64_t v158 = (uint64_t) v53; + TASSIGN(v157, v158); + TCOLEXPANDMUL(v157, v131, v65); + // pto: %k_normed_inline49__tile + Tile v159 = Tile(v26, v38); + // pto: %k_normed_inline49__tile + uint64_t v160 = (uint64_t) v53; + TASSIGN(v159, v160); + pipe_barrier(PIPE_V); + TROWEXPANDMUL(v159, v157, v155); + // pto: %slice_view + Tile v161; + // pto: %slice_view + Tile v162 = v161; + // pto: %slice_view + uint64_t v163 = (uint64_t) v53; + TASSIGN(v162, v163); + // pto: %k_lo_inline243__tile + Tile v164 = Tile(v37, v23); + // pto: %k_lo_inline243__tile + uint64_t v165 = (uint64_t) v53; + TASSIGN(v164, v165); + // pto: %k_hi_inline58__tile + Tile v166 = Tile(v37, v23); + // pto: %k_hi_inline58__tile + uint64_t v167 = (uint64_t) v50; + TASSIGN(v166, v167); + // pto: %9 + Tile v168 = Tile(v37, v23); + // pto: %9 + uint64_t v169 = (uint64_t) v56; + TASSIGN(v168, v169); + pipe_barrier(PIPE_V); + TEXTRACT(v164, v162, v45, v45); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v168, v164, v92); + // pto: %10 + Tile v170 = Tile(v37, v23); + // pto: %10 + uint64_t v171 = (uint64_t) v52; + TASSIGN(v170, v171); + TEXTRACT(v166, v162, v45, v23); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v170, v166, v102); + // pto: %rot_lo_inline45__tile + Tile v172 = Tile(v37, v23); + // pto: %rot_lo_inline45__tile + uint64_t v173 = (uint64_t) v56; + TASSIGN(v172, v173); + pipe_barrier(PIPE_V); + TSUB(v172, v168, v170); + // pto: %11 + Tile v174 = Tile(v37, v23); + // pto: %11 + uint64_t v175 = (uint64_t) v52; + TASSIGN(v174, v175); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v174, v166, v97); + // pto: %12 + Tile v176 = Tile(v37, v23); + // pto: %12 + uint64_t v177 = (uint64_t) v53; + TASSIGN(v176, v177); + TCOLEXPANDMUL(v176, v164, v107); + // pto: %rot_hi_inline221__tile + Tile v178 = Tile(v37, v23); + // pto: %rot_hi_inline221__tile + uint64_t v179 = (uint64_t) v52; + TASSIGN(v178, v179); + pipe_barrier(PIPE_V); + TADD(v178, v174, v176); + // pto: %13 + Tile v180 = Tile(v37, v38); + // pto: %13 + uint64_t v181 = (uint64_t) v53; + TASSIGN(v180, v181); + pipe_barrier(PIPE_V); + TCONCAT(v180, v172, v178); + // pto: %14 + Tile v182 = Tile(v37, v38); + // pto: %14 + uint64_t v183 = (uint64_t) v51; + TASSIGN(v182, v183); + pipe_barrier(PIPE_V); + TCVT(v182, v180, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + // pto: %15 + Tile v184 = Tile(v37, v38); + // pto: %15 + uint64_t v185 = (uint64_t) v53; + TASSIGN(v184, v185); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); + TMULS(v184, v117, v87); + // pto: %v_row_bf16_inline202__tile + Tile v186 = Tile(v37, v38); + // pto: %v_row_bf16_inline202__tile + uint64_t v187 = (uint64_t) v55; + TASSIGN(v186, v187); + pipe_barrier(PIPE_V); + TCVT(v186, v184, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + // pto: %16 + Tile v188 = Tile(v37, v20); + // pto: %16 + uint64_t v189 = (uint64_t) v53; + TASSIGN(v188, v189); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); + TCONCAT(v188, v122, v76); + // pto: %17 + Tile v190 = Tile(v35, v38); + // pto: %17 + uint64_t v191 = (uint64_t) v53; + TASSIGN(v190, v191); + // pto: %q_raw_inline151__tile + Tile v192 = Tile(v35, v38); + // pto: %q_raw_inline151__tile + uint64_t v193 = (uint64_t) v53; + TASSIGN(v192, v193); + pipe_barrier(PIPE_V); + TMULS(v192, v190, v87); + // pto: %18 + Tile v194 = Tile(v35, v38); + // pto: %18 + uint64_t v195 = (uint64_t) v56; + TASSIGN(v194, v195); + pipe_barrier(PIPE_V); + TMUL(v194, v192, v192); + // pto: %19 + Tile v196 = Tile(v35, v38); + // pto: %19 + uint64_t v197 = (uint64_t) v52; + TASSIGN(v196, v197); + // pto: %q_ss_inline61__tile + Tile v198 = Tile(v35, v37); + // pto: %q_ss_inline61__tile + uint64_t v199 = (uint64_t) v54; + TASSIGN(v198, v199); + pipe_barrier(PIPE_V); + TROWSUM(v198, v194, v196); + // pto: %t__rm_a0_tmp_v8 + Tile v200 = Tile(v37, v35); + // pto: %t__rm_a0_tmp_v8 + uint64_t v201 = (uint64_t) v54; + TASSIGN(v200, v201); + // pto: %t__row_major_tmp_v9 + Tile v202 = Tile(v37, v35); + // pto: %t__row_major_tmp_v9 + uint64_t v203 = (uint64_t) v56; + TASSIGN(v202, v203); + pipe_barrier(PIPE_V); + TMULS(v202, v200, v22); + // pto: %t__rm_a0_tmp_v10 + Tile v204 = Tile(v37, v35); + // pto: %t__rm_a0_tmp_v10 + uint64_t v205 = (uint64_t) v56; + TASSIGN(v204, v205); + // pto: %t__row_major_tmp_v11 + Tile v206 = Tile(v37, v35); + // pto: %t__row_major_tmp_v11 + uint64_t v207 = (uint64_t) v56; + TASSIGN(v206, v207); + pipe_barrier(PIPE_V); + TADDS(v206, v204, v21); + // pto: %t__rm_a0_tmp_v12 + Tile v208 = Tile(v37, v35); + // pto: %t__rm_a0_tmp_v12 + uint64_t v209 = (uint64_t) v56; + TASSIGN(v208, v209); + // pto: %t__row_major_tmp_v13 + Tile v210 = Tile(v37, v35); + // pto: %t__row_major_tmp_v13 + uint64_t v211 = (uint64_t) v56; + TASSIGN(v210, v211); + pipe_barrier(PIPE_V); + TSQRT(v210, v208); + // pto: %q_inv_inline192__rm_a0_tmp_v14 + Tile v212 = Tile(v37, v35); + // pto: %q_inv_inline192__rm_a0_tmp_v14 + uint64_t v213 = (uint64_t) v56; + TASSIGN(v212, v213); + // pto: %q_inv_inline192__row_major_tmp_v15 + Tile v214 = Tile(v37, v35); + // pto: %q_inv_inline192__row_major_tmp_v15 + uint64_t v215 = (uint64_t) v52; + TASSIGN(v214, v215); + pipe_barrier(PIPE_V); + TRECIP(v214, v212); + // pto: %q_inv_inline192__tile + Tile v216 = Tile(v35, v37); + // pto: %q_inv_inline192__tile + uint64_t v217 = (uint64_t) v52; + TASSIGN(v216, v217); + // pto: %23 + Tile v218 = Tile(v35, v38); + // pto: %23 + uint64_t v219 = (uint64_t) v53; + TASSIGN(v218, v219); + TCOLEXPANDMUL(v218, v192, v70); + // pto: %q_heads_inline237__tile + Tile v220 = Tile(v35, v38); + // pto: %q_heads_inline237__tile + uint64_t v221 = (uint64_t) v53; + TASSIGN(v220, v221); + pipe_barrier(PIPE_V); + TROWEXPANDMUL(v220, v218, v216); + // pto: %q_lo_inline136__tile_textract + Tile v222 = Tile(v35, v23); + // pto: %q_lo_inline136__tile_textract + uint64_t v223 = (uint64_t) v56; + TASSIGN(v222, v223); + pipe_barrier(PIPE_V); + TEXTRACT(v222, v220, v45, v45); + // pto: %24 + Tile v224 = Tile(v35, v23); + // pto: %24 + uint64_t v225 = (uint64_t) v56; + TASSIGN(v224, v225); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v224, v222, v92); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + // pto: %q_hi_inline142__tile_textract + Tile v226 = Tile(v35, v23); + // pto: %q_hi_inline142__tile_textract + uint64_t v227 = (uint64_t) v52; + TASSIGN(v226, v227); + TEXTRACT(v226, v220, v45, v23); + // pto: %25 + Tile v228 = Tile(v35, v23); + // pto: %25 + uint64_t v229 = (uint64_t) v52; + TASSIGN(v228, v229); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v228, v226, v102); + // pto: %q_rot_lo_inline285__tile + Tile v230 = Tile(v35, v23); + // pto: %q_rot_lo_inline285__tile + uint64_t v231 = (uint64_t) v56; + TASSIGN(v230, v231); + pipe_barrier(PIPE_V); + TSUB(v230, v224, v228); + // pto: %26 + Tile v232 = Tile(v35, v23); + // pto: %26 + uint64_t v233 = (uint64_t) v52; + TASSIGN(v232, v233); + pipe_barrier(PIPE_V); + TEXTRACT(v232, v220, v45, v23); + // pto: %27 + Tile v234 = Tile(v35, v23); + // pto: %27 + uint64_t v235 = (uint64_t) v52; + TASSIGN(v234, v235); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v234, v232, v97); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + // pto: %28 + Tile v236 = Tile(v35, v23); + // pto: %28 + uint64_t v237 = (uint64_t) v54; + TASSIGN(v236, v237); + TEXTRACT(v236, v220, v45, v45); + // pto: %29 + Tile v238 = Tile(v35, v23); + // pto: %29 + uint64_t v239 = (uint64_t) v53; + TASSIGN(v238, v239); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v238, v236, v107); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + // pto: %q_rot_hi_inline39__tile + Tile v240 = Tile(v35, v23); + // pto: %q_rot_hi_inline39__tile + uint64_t v241 = (uint64_t) v52; + TASSIGN(v240, v241); + pipe_barrier(PIPE_V); + TADD(v240, v234, v238); + // pto: %30 + Tile v242 = Tile(v35, v38); + // pto: %30 + uint64_t v243 = (uint64_t) v53; + TASSIGN(v242, v243); + pipe_barrier(PIPE_V); + TCONCAT(v242, v230, v240); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + // pto: %137 + Tile v244; + // pto: %137 + Tile v245 = v244; + // pto: %137 + uint64_t v246 = (uint64_t) v53; + TASSIGN(v245, v246); + // pto: %32 + Tile v247 = Tile(v25, v38); + // pto: %32 + uint64_t v248 = (uint64_t) v53; + TASSIGN(v247, v248); + pipe_barrier(PIPE_V); + TCVT(v247, v245, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + // pto: %k_cache__iter_v3_pview + pto::Shape<1, 1, 1, 1, 128> v249 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %k_cache__iter_v3_pview + pto::Stride<128, 128, 128, 128, 1> v250 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %k_cache__iter_v3_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v251 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v1 + ((v45 + v91 * v38) + v45 * v37), v249, v250); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + pipe_barrier(PIPE_MTE3); + TSTORE(v251, v182); + // pto: %v_cache__iter_v3_pview + pto::Shape<1, 1, 1, 1, 128> v252 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %v_cache__iter_v3_pview + pto::Stride<128, 128, 128, 128, 1> v253 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %v_cache__iter_v3_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v254 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v3 + ((v45 + v91 * v38) + v45 * v37), v252, v253); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); + TSTORE(v254, v186); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + // pto: %q_tnd_flat_inline134__iter_v1_pview + pto::Shape<1, 1, 1, 5, 128> v255 = pto::Shape<1, 1, 1, 5, 128>(); + // pto: %q_tnd_flat_inline134__iter_v1_pview + pto::Stride<640, 640, 640, 128, 1> v256 = pto::Stride<640, 640, 640, 128, 1>(); + // pto: %129, %130, %128, %q_tnd_flat_inline134__iter_v1_pview + GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v257 = GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>(v2 + ((v45 + (int64_t) ((uint64_t) ((int64_t) ((uint64_t) v85 * (uint64_t) v24)) + (uint64_t) ((int64_t) ((uint64_t) v84 * (uint64_t) v25))) * v38) + v45 * v37), v255, v256); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); + TSTORE(v257, v247); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + } + // pto: %138 + if (v83 < v38) { + // pto: %139 + int64_t v258 = v83 / v35; + // pto: %141, %140 + int64_t v259 = (int64_t) ((uint64_t) v83 - (uint64_t) ((int64_t) ((uint64_t) v258 * (uint64_t) v35))); + // pto: %142 + int32_t v260 = v4[v259]; + // pto: %143 + float v261 = v5[v259]; + // pto: %146, %147 + int64_t v262 = (int64_t) ((uint64_t) ((int64_t) v260) - (uint64_t) v37); + // pto: %148 + int32_t v263 = v6[v259]; + // pto: %153 + int64_t v264 = (int64_t) ((uint64_t) v258 * (uint64_t) v38); + // pto: %157, %149, %156, %158 + int64_t v265 = (int64_t) ((uint64_t) ((int64_t) ((uint64_t) v14 + (uint64_t) ((int64_t) ((uint64_t) ((int64_t) v263) * (uint64_t) v26)))) + (uint64_t) v258); + // pto: %33 + Tile v266 = Tile(v37, v23); + // pto: %33 + uint64_t v267 = (uint64_t) v49; + TASSIGN(v266, v267); + // pto: %162 + pto::Shape<1, 1, 1, 1, 64> v268 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %162 + pto::Stride<128, 128, 128, 128, 1> v269 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %162 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v270 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v7 + ((v45 + v262 * v38) + v45 * v37), v268, v269); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID4); + TLOAD(v266, v270); + // pto: %34 + Tile v271 = Tile(v37, v23); + // pto: %34 + uint64_t v272 = (uint64_t) v48; + TASSIGN(v271, v272); + // pto: %163 + pto::Shape<1, 1, 1, 1, 64> v273 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %163 + pto::Stride<128, 128, 128, 128, 1> v274 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %163 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v275 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v7 + ((v45 + v262 * v38) + v23 * v37), v273, v274); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID5); + TLOAD(v271, v275); + // pto: %35 + Tile v276 = Tile(v37, v23); + // pto: %35 + uint64_t v277 = (uint64_t) v47; + TASSIGN(v276, v277); + // pto: %164 + pto::Shape<1, 1, 1, 1, 64> v278 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %164 + pto::Stride<128, 128, 128, 128, 1> v279 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %164 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v280 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v8 + ((v45 + v262 * v38) + v45 * v37), v278, v279); + TLOAD(v276, v280); + // pto: %36 + Tile v281 = Tile(v37, v23); + // pto: %36 + uint64_t v282 = (uint64_t) v46; + TASSIGN(v281, v282); + // pto: %165 + pto::Shape<1, 1, 1, 1, 64> v283 = pto::Shape<1, 1, 1, 1, 64>(); + // pto: %165 + pto::Stride<128, 128, 128, 128, 1> v284 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %165 + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v285 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v8 + ((v45 + v262 * v38) + v23 * v37), v283, v284); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID6); + TLOAD(v281, v285); + // pto: %37 + Tile v286 = Tile(v37, v38); + // pto: %37 + uint64_t v287 = (uint64_t) v45; + TASSIGN(v286, v287); + // pto: %166 + pto::Shape<1, 1, 1, 1, 128> v288 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %166 + pto::Stride<1024, 1024, 1024, 1024, 1> v289 = pto::Stride<1024, 1024, 1024, 1024, 1>(); + // pto: %166 + GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v290 = GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>(v9 + ((v45 + v259 * v34) + v264 * v37), v288, v289); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + TLOAD(v286, v290); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + // pto: %38 + Tile v291 = Tile(v37, v38); + // pto: %38 + uint64_t v292 = (uint64_t) v44; + TASSIGN(v291, v292); + // pto: %168 + pto::Shape<1, 1, 1, 1, 128> v293 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %168 + pto::Stride<1024, 1024, 1024, 1024, 1> v294 = pto::Stride<1024, 1024, 1024, 1024, 1>(); + // pto: %168 + GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v295 = GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>(v11 + ((v45 + v259 * v34) + v264 * v37), v293, v294); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + TLOAD(v291, v295); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); + // pto: %39 + Tile v296 = Tile(v37, v36); + // pto: %39 + uint64_t v297 = (uint64_t) v43; + TASSIGN(v296, v297); + // pto: %170 + pto::Shape<1, 1, 1, 1, 640> v298 = pto::Shape<1, 1, 1, 1, 640>(); + // pto: %170 + pto::Stride<5120, 5120, 5120, 5120, 1> v299 = pto::Stride<5120, 5120, 5120, 5120, 1>(); + // pto: %170 + GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v300 = GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>(v12 + ((v45 + v259 * v33) + (int64_t) ((uint64_t) v258 * (uint64_t) v36) * v37), v298, v299); + TLOAD(v296, v300); + set_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); + // pto: %40 + Tile v301 = Tile(v37, v34); + // pto: %40 + uint64_t v302 = (uint64_t) v42; + TASSIGN(v301, v302); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + TCONCAT(v301, v286, v78); + // pto: %41 + Tile v303 = Tile(v26, v38); + // pto: %41 + uint64_t v304 = (uint64_t) v42; + TASSIGN(v303, v304); + // pto: %42 + Tile v305 = Tile(v26, v38); + // pto: %42 + uint64_t v306 = (uint64_t) v42; + TASSIGN(v305, v306); + pipe_barrier(PIPE_V); + TMULS(v305, v303, v261); + // pto: %43 + Tile v307 = Tile(v26, v38); + // pto: %43 + uint64_t v308 = (uint64_t) v45; + TASSIGN(v307, v308); + pipe_barrier(PIPE_V); + TMUL(v307, v305, v305); + // pto: %44 + Tile v309 = Tile(v26, v38); + // pto: %44 + uint64_t v310 = (uint64_t) v41; + TASSIGN(v309, v310); + // pto: %45 + Tile v311 = Tile(v26, v37); + // pto: %45 + uint64_t v312 = (uint64_t) v40; + TASSIGN(v311, v312); + pipe_barrier(PIPE_V); + TROWSUM(v311, v307, v309); + // pto: %46 + Tile v313 = Tile(v37, v26); + // pto: %46 + uint64_t v314 = (uint64_t) v40; + TASSIGN(v313, v314); + // pto: %47 + Tile v315 = Tile(v37, v26); + // pto: %47 + uint64_t v316 = (uint64_t) v45; + TASSIGN(v315, v316); + pipe_barrier(PIPE_V); + TMULS(v315, v313, v22); + // pto: %49 + Tile v317 = Tile(v37, v26); + // pto: %49 + uint64_t v318 = (uint64_t) v45; + TASSIGN(v317, v318); + // pto: %50 + Tile v319 = Tile(v37, v26); + // pto: %50 + uint64_t v320 = (uint64_t) v45; + TASSIGN(v319, v320); + pipe_barrier(PIPE_V); + TADDS(v319, v317, v21); + // pto: %52 + Tile v321 = Tile(v37, v26); + // pto: %52 + uint64_t v322 = (uint64_t) v45; + TASSIGN(v321, v322); + // pto: %53 + Tile v323 = Tile(v37, v26); + // pto: %53 + uint64_t v324 = (uint64_t) v45; + TASSIGN(v323, v324); + pipe_barrier(PIPE_V); + TSQRT(v323, v321); + // pto: %55 + Tile v325 = Tile(v37, v26); + // pto: %55 + uint64_t v326 = (uint64_t) v45; + TASSIGN(v325, v326); + // pto: %56 + Tile v327 = Tile(v37, v26); + // pto: %56 + uint64_t v328 = (uint64_t) v41; + TASSIGN(v327, v328); + pipe_barrier(PIPE_V); + TRECIP(v327, v325); + // pto: %57 + Tile v329 = Tile(v26, v37); + // pto: %57 + uint64_t v330 = (uint64_t) v41; + TASSIGN(v329, v330); + // pto: %58 + Tile v331 = Tile(v26, v38); + // pto: %58 + uint64_t v332 = (uint64_t) v42; + TASSIGN(v331, v332); + TCOLEXPANDMUL(v331, v305, v65); + // pto: %59 + Tile v333 = Tile(v26, v38); + // pto: %59 + uint64_t v334 = (uint64_t) v42; + TASSIGN(v333, v334); + pipe_barrier(PIPE_V); + TROWEXPANDMUL(v333, v331, v329); + // pto: %171 + Tile v335; + // pto: %171 + Tile v336 = v335; + // pto: %171 + uint64_t v337 = (uint64_t) v42; + TASSIGN(v336, v337); + // pto: %61 + Tile v338 = Tile(v37, v23); + // pto: %61 + uint64_t v339 = (uint64_t) v42; + TASSIGN(v338, v339); + // pto: %62 + Tile v340 = Tile(v37, v23); + // pto: %62 + uint64_t v341 = (uint64_t) v39; + TASSIGN(v340, v341); + // pto: %63 + Tile v342 = Tile(v37, v23); + // pto: %63 + uint64_t v343 = (uint64_t) v45; + TASSIGN(v342, v343); + pipe_barrier(PIPE_V); + TEXTRACT(v338, v336, v45, v45); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v342, v338, v266); + // pto: %64 + Tile v344 = Tile(v37, v23); + // pto: %64 + uint64_t v345 = (uint64_t) v41; + TASSIGN(v344, v345); + TEXTRACT(v340, v336, v45, v23); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v344, v340, v276); + // pto: %65 + Tile v346 = Tile(v37, v23); + // pto: %65 + uint64_t v347 = (uint64_t) v45; + TASSIGN(v346, v347); + pipe_barrier(PIPE_V); + TSUB(v346, v342, v344); + // pto: %66 + Tile v348 = Tile(v37, v23); + // pto: %66 + uint64_t v349 = (uint64_t) v41; + TASSIGN(v348, v349); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v348, v340, v271); + // pto: %67 + Tile v350 = Tile(v37, v23); + // pto: %67 + uint64_t v351 = (uint64_t) v42; + TASSIGN(v350, v351); + TCOLEXPANDMUL(v350, v338, v281); + // pto: %68 + Tile v352 = Tile(v37, v23); + // pto: %68 + uint64_t v353 = (uint64_t) v41; + TASSIGN(v352, v353); + pipe_barrier(PIPE_V); + TADD(v352, v348, v350); + // pto: %69 + Tile v354 = Tile(v37, v38); + // pto: %69 + uint64_t v355 = (uint64_t) v42; + TASSIGN(v354, v355); + pipe_barrier(PIPE_V); + TCONCAT(v354, v346, v352); + // pto: %70 + Tile v356 = Tile(v37, v38); + // pto: %70 + uint64_t v357 = (uint64_t) v40; + TASSIGN(v356, v357); + pipe_barrier(PIPE_V); + TCVT(v356, v354, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + // pto: %71 + Tile v358 = Tile(v37, v38); + // pto: %71 + uint64_t v359 = (uint64_t) v42; + TASSIGN(v358, v359); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); + TMULS(v358, v291, v261); + // pto: %72 + Tile v360 = Tile(v37, v38); + // pto: %72 + uint64_t v361 = (uint64_t) v44; + TASSIGN(v360, v361); + pipe_barrier(PIPE_V); + TCVT(v360, v358, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); + // pto: %73 + Tile v362 = Tile(v37, v20); + // pto: %73 + uint64_t v363 = (uint64_t) v42; + TASSIGN(v362, v363); + pipe_barrier(PIPE_V); + wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); + TCONCAT(v362, v296, v76); + // pto: %74 + Tile v364 = Tile(v35, v38); + // pto: %74 + uint64_t v365 = (uint64_t) v42; + TASSIGN(v364, v365); + // pto: %75 + Tile v366 = Tile(v35, v38); + // pto: %75 + uint64_t v367 = (uint64_t) v42; + TASSIGN(v366, v367); + pipe_barrier(PIPE_V); + TMULS(v366, v364, v261); + // pto: %76 + Tile v368 = Tile(v35, v38); + // pto: %76 + uint64_t v369 = (uint64_t) v45; + TASSIGN(v368, v369); + pipe_barrier(PIPE_V); + TMUL(v368, v366, v366); + // pto: %77 + Tile v370 = Tile(v35, v38); + // pto: %77 + uint64_t v371 = (uint64_t) v41; + TASSIGN(v370, v371); + // pto: %78 + Tile v372 = Tile(v35, v37); + // pto: %78 + uint64_t v373 = (uint64_t) v43; + TASSIGN(v372, v373); + pipe_barrier(PIPE_V); + TROWSUM(v372, v368, v370); + // pto: %79 + Tile v374 = Tile(v37, v35); + // pto: %79 + uint64_t v375 = (uint64_t) v43; + TASSIGN(v374, v375); + // pto: %80 + Tile v376 = Tile(v37, v35); + // pto: %80 + uint64_t v377 = (uint64_t) v45; + TASSIGN(v376, v377); + pipe_barrier(PIPE_V); + TMULS(v376, v374, v22); + // pto: %82 + Tile v378 = Tile(v37, v35); + // pto: %82 + uint64_t v379 = (uint64_t) v45; + TASSIGN(v378, v379); + // pto: %83 + Tile v380 = Tile(v37, v35); + // pto: %83 + uint64_t v381 = (uint64_t) v45; + TASSIGN(v380, v381); + pipe_barrier(PIPE_V); + TADDS(v380, v378, v21); + // pto: %85 + Tile v382 = Tile(v37, v35); + // pto: %85 + uint64_t v383 = (uint64_t) v45; + TASSIGN(v382, v383); + // pto: %86 + Tile v384 = Tile(v37, v35); + // pto: %86 + uint64_t v385 = (uint64_t) v45; + TASSIGN(v384, v385); + pipe_barrier(PIPE_V); + TSQRT(v384, v382); + // pto: %88 + Tile v386 = Tile(v37, v35); + // pto: %88 + uint64_t v387 = (uint64_t) v45; + TASSIGN(v386, v387); + // pto: %89 + Tile v388 = Tile(v37, v35); + // pto: %89 + uint64_t v389 = (uint64_t) v41; + TASSIGN(v388, v389); + pipe_barrier(PIPE_V); + TRECIP(v388, v386); + // pto: %90 + Tile v390 = Tile(v35, v37); + // pto: %90 + uint64_t v391 = (uint64_t) v41; + TASSIGN(v390, v391); + // pto: %91 + Tile v392 = Tile(v35, v38); + // pto: %91 + uint64_t v393 = (uint64_t) v42; + TASSIGN(v392, v393); + TCOLEXPANDMUL(v392, v366, v70); + // pto: %92 + Tile v394 = Tile(v35, v38); + // pto: %92 + uint64_t v395 = (uint64_t) v42; + TASSIGN(v394, v395); + pipe_barrier(PIPE_V); + TROWEXPANDMUL(v394, v392, v390); + // pto: %93 + Tile v396 = Tile(v35, v23); + // pto: %93 + uint64_t v397 = (uint64_t) v45; + TASSIGN(v396, v397); + pipe_barrier(PIPE_V); + TEXTRACT(v396, v394, v45, v45); + // pto: %94 + Tile v398 = Tile(v35, v23); + // pto: %94 + uint64_t v399 = (uint64_t) v45; + TASSIGN(v398, v399); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v398, v396, v266); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID4); + // pto: %95 + Tile v400 = Tile(v35, v23); + // pto: %95 + uint64_t v401 = (uint64_t) v41; + TASSIGN(v400, v401); + TEXTRACT(v400, v394, v45, v23); + // pto: %96 + Tile v402 = Tile(v35, v23); + // pto: %96 + uint64_t v403 = (uint64_t) v41; + TASSIGN(v402, v403); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v402, v400, v276); + // pto: %97 + Tile v404 = Tile(v35, v23); + // pto: %97 + uint64_t v405 = (uint64_t) v45; + TASSIGN(v404, v405); + pipe_barrier(PIPE_V); + TSUB(v404, v398, v402); + // pto: %98 + Tile v406 = Tile(v35, v23); + // pto: %98 + uint64_t v407 = (uint64_t) v41; + TASSIGN(v406, v407); + pipe_barrier(PIPE_V); + TEXTRACT(v406, v394, v45, v23); + // pto: %99 + Tile v408 = Tile(v35, v23); + // pto: %99 + uint64_t v409 = (uint64_t) v41; + TASSIGN(v408, v409); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v408, v406, v271); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID5); + // pto: %100 + Tile v410 = Tile(v35, v23); + // pto: %100 + uint64_t v411 = (uint64_t) v43; + TASSIGN(v410, v411); + TEXTRACT(v410, v394, v45, v45); + // pto: %101 + Tile v412 = Tile(v35, v23); + // pto: %101 + uint64_t v413 = (uint64_t) v42; + TASSIGN(v412, v413); + pipe_barrier(PIPE_V); + TCOLEXPANDMUL(v412, v410, v281); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID6); + // pto: %102 + Tile v414 = Tile(v35, v23); + // pto: %102 + uint64_t v415 = (uint64_t) v41; + TASSIGN(v414, v415); + pipe_barrier(PIPE_V); + TADD(v414, v408, v412); + // pto: %103 + Tile v416 = Tile(v35, v38); + // pto: %103 + uint64_t v417 = (uint64_t) v42; + TASSIGN(v416, v417); + pipe_barrier(PIPE_V); + TCONCAT(v416, v404, v414); + set_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + // pto: %174 + Tile v418; + // pto: %174 + Tile v419 = v418; + // pto: %174 + uint64_t v420 = (uint64_t) v42; + TASSIGN(v419, v420); + // pto: %105 + Tile v421 = Tile(v25, v38); + // pto: %105 + uint64_t v422 = (uint64_t) v42; + TASSIGN(v421, v422); + pipe_barrier(PIPE_V); + TCVT(v421, v419, v19, v18); + set_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + // pto: %k_cache__phi_v6_pview + pto::Shape<1, 1, 1, 1, 128> v423 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %k_cache__phi_v6_pview + pto::Stride<128, 128, 128, 128, 1> v424 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %k_cache__phi_v6_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v425 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v1 + ((v45 + v265 * v38) + v45 * v37), v423, v424); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); + pipe_barrier(PIPE_MTE3); + TSTORE(v425, v356); + // pto: %v_cache__phi_v6_pview + pto::Shape<1, 1, 1, 1, 128> v426 = pto::Shape<1, 1, 1, 1, 128>(); + // pto: %v_cache__phi_v6_pview + pto::Stride<128, 128, 128, 128, 1> v427 = pto::Stride<128, 128, 128, 128, 1>(); + // pto: %v_cache__phi_v6_pview + GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v428 = GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>(v3 + ((v45 + v265 * v38) + v45 * v37), v426, v427); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); + TSTORE(v428, v360); + set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + // pto: %q_tnd_flat_inline134__phi_v4_pview + pto::Shape<1, 1, 1, 5, 128> v429 = pto::Shape<1, 1, 1, 5, 128>(); + // pto: %q_tnd_flat_inline134__phi_v4_pview + pto::Stride<640, 640, 640, 128, 1> v430 = pto::Stride<640, 640, 640, 128, 1>(); + // pto: %160, %161, %159, %q_tnd_flat_inline134__phi_v4_pview + GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v431 = GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>(v2 + ((v45 + (int64_t) ((uint64_t) ((int64_t) ((uint64_t) v259 * (uint64_t) v24)) + (uint64_t) ((int64_t) ((uint64_t) v258 * (uint64_t) v25))) * v38) + v45 * v37), v429, v430); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); + TSTORE(v431, v421); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + } + } + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID3); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID4); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID5); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID6); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID7); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); + #endif // __DAV_VEC__ + + ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); + return; +} + +} // namespace qwen_rope_gen +#endif // __DAV_C220_VEC__ + +#endif // PYPTO_QWEN_ROPE_QKV_GENERATED_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp new file mode 100644 index 0000000000..3c48e349fa --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp @@ -0,0 +1,116 @@ +/* + * 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 "tensor.h" + +#ifdef __CPU_SIM +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; } + +#else + +#include "intrinsic.h" + +#define QWEN_FAI_TILER_FUNCTION static __aicore__ +#include "qwen_fai_runtime_tiler.hpp" + +namespace { + +template +static __aicore__ __attribute__((always_inline)) __gm__ T *tensor_data(__gm__ int64_t *args, int32_t index) { + __gm__ Tensor *tensor = reinterpret_cast<__gm__ Tensor *>(args[index]); + return reinterpret_cast<__gm__ T *>(tensor->buffer.addr) + tensor->start_offset; +} + +static __aicore__ __attribute__((always_inline)) __gm__ Tensor *tensor_desc(__gm__ int64_t *args, int32_t index) { + return reinterpret_cast<__gm__ Tensor *>(args[index]); +} + +static __aicore__ __attribute__((always_inline)) __gm__ int32_t *barrier_data(__gm__ uint8_t *metadata) { + uint64_t raw_barrier = reinterpret_cast(metadata + qwen_fai_metadata::kBarrierAlignmentOffset); + uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & + ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); + return reinterpret_cast<__gm__ int32_t *>(aligned_barrier); +} + +static __aicore__ void clear_barrier(__gm__ int32_t *barrier) { + for (uint32_t slot = 0; slot < qwen_fai_metadata::kBarrierSlotCount; ++slot) { + __gm__ int32_t *slot_data = barrier + slot * qwen_fai_metadata::kBarrierSlotWords; + slot_data[0] = 0; + dcci(slot_data, SINGLE_CACHE_LINE, CACHELINE_OUT); + } + dsb(DSB_DDR); +} + +static __aicore__ void flush_metadata_prefix(__gm__ uint8_t *metadata) { + uint64_t first_line = + reinterpret_cast(metadata) & ~(static_cast(qwen_fai_metadata::kDcciLineBytes) - 1); + uint64_t end = reinterpret_cast(metadata) + qwen_fai_metadata::kBarrierAlignmentOffset; + for (uint64_t line = first_line; line < end; line += qwen_fai_metadata::kDcciLineBytes) { + dcci(reinterpret_cast<__gm__ void *>(line), SINGLE_CACHE_LINE, CACHELINE_OUT); + } + dsb(DSB_DDR); +} + +} // namespace + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *seq_lens_desc = tensor_desc(args, 0); + __gm__ const int32_t *seq_lens = tensor_data(args, 0); + __gm__ uint8_t *metadata = tensor_data(args, 1); + __gm__ uint32_t *tiling_out = reinterpret_cast<__gm__ uint32_t *>(metadata + qwen_fai_metadata::kTilingOffset); + __gm__ int64_t *cumulative_q_out = + reinterpret_cast<__gm__ int64_t *>(metadata + qwen_fai_metadata::kCumulativeQOffset); + __gm__ int64_t *kv_lengths_out = reinterpret_cast<__gm__ int64_t *>(metadata + qwen_fai_metadata::kKvLengthsOffset); + uint32_t batch = static_cast(seq_lens_desc->shapes[0]); + uint32_t max_blocks_per_batch = static_cast(args[2]); + uint32_t num_blocks = static_cast(args[3]); + + clear_barrier(barrier_data(metadata)); + if (batch == 0 || batch > qwen_fai_tiler::kMaxBatch) { + return; + } + + uint32_t local_seq_lens[qwen_fai_tiler::kMaxBatch] = {}; + int64_t local_cumulative_q[qwen_fai_tiler::kMaxBatch] = {}; + int64_t local_kv_lengths[qwen_fai_tiler::kMaxBatch] = {}; + for (uint32_t batch_idx = 0; batch_idx < batch; ++batch_idx) { + local_seq_lens[batch_idx] = static_cast(seq_lens[batch_idx]); + } + + FAInferTilingData tiling; + bool valid = qwen_fai_tiler::build( + local_seq_lens, batch, max_blocks_per_batch, num_blocks, tiling, local_cumulative_q, local_kv_lengths + ); + if (!valid) { + return; + } + + uint32_t *tiling_words = reinterpret_cast(&tiling); + for (uint32_t word_idx = 0; word_idx < sizeof(FAInferTilingData) / sizeof(uint32_t); ++word_idx) { + tiling_out[word_idx] = tiling_words[word_idx]; + } + for (uint32_t batch_idx = 0; batch_idx < qwen_fai_tiler::kMaxBatch; ++batch_idx) { + cumulative_q_out[batch_idx] = local_cumulative_q[batch_idx]; + kv_lengths_out[batch_idx] = local_kv_lengths[batch_idx]; + } + flush_metadata_prefix(metadata); +} + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/qwen_fai_runtime_tiler.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/qwen_fai_runtime_tiler.hpp new file mode 100644 index 0000000000..5e463017ae --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/qwen_fai_runtime_tiler.hpp @@ -0,0 +1,346 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#ifndef PYPTO_QWEN_FAI_RUNTIME_TILER_HPP +#define PYPTO_QWEN_FAI_RUNTIME_TILER_HPP + +#include +#include + +#include "../generated/kernel_tiling/kernel_tiling.h" +#include "../kernel/metadata_layout.h" + +#ifndef QWEN_FAI_TILER_FUNCTION +#define QWEN_FAI_TILER_FUNCTION inline +#endif + +namespace qwen_fai_tiler { + +static_assert( + sizeof(FAInferTilingData) == qwen_fai_metadata::kTilingBytes, + "metadata tiling size does not match FAInferTilingData" +); + +constexpr uint32_t kMaxBatch = 16; +constexpr uint32_t kNumHeads = 40; +constexpr uint32_t kNumKvHeads = 8; +constexpr uint32_t kHeadDim = 128; +constexpr uint32_t kBlockSize = 128; +constexpr uint32_t kPlanningCoreNum = 24; +constexpr uint32_t kMaxCoreNum = 26; +constexpr uint32_t kQTileCeil = 128; +constexpr uint32_t kKvTile = 512; +constexpr uint32_t kPrelaunchNum = 3; +constexpr uint32_t kWorkspaceBlockSize = kQTileCeil * kKvTile; +constexpr uint32_t kUint16Bytes = 2; +constexpr uint32_t kUint32Bytes = 4; +constexpr uint64_t kBaseWorkspaceSize = 66'060'288; +constexpr int64_t kSparseTokenLimit = 2147483647; + +struct BatchParams { + uint32_t q_seqlen; + uint32_t kv_seqlen; + uint32_t qn_block_tile; + uint32_t qn_blocks_per_group; + uint32_t qn_block_num; + uint32_t qs_block_tile; + uint32_t qs_block_num; + uint32_t ks_block_tile; + uint32_t ks_block_num; +}; + +QWEN_FAI_TILER_FUNCTION uint32_t min_u32(uint32_t lhs, uint32_t rhs) { return lhs < rhs ? lhs : rhs; } + +QWEN_FAI_TILER_FUNCTION uint32_t ceil_div(uint32_t value, uint32_t divisor) { return (value + divisor - 1) / divisor; } + +QWEN_FAI_TILER_FUNCTION BatchParams get_batch_params(uint32_t batch_idx, const uint32_t *kv_seqlens) { + constexpr uint32_t group_size = kNumHeads / kNumKvHeads; + BatchParams params{}; + params.q_seqlen = 1; + params.kv_seqlen = kv_seqlens[batch_idx]; + params.qn_block_tile = min_u32(kQTileCeil, group_size); + params.qn_blocks_per_group = ceil_div(group_size, params.qn_block_tile); + params.qn_block_num = params.qn_blocks_per_group * kNumKvHeads; + params.qs_block_tile = kQTileCeil; + params.qs_block_num = ceil_div(params.q_seqlen, params.qs_block_tile); + params.ks_block_tile = kKvTile; + params.ks_block_num = ceil_div(params.kv_seqlen, params.ks_block_tile); + return params; +} + +QWEN_FAI_TILER_FUNCTION void zero_tiling(FAInferTilingData &tiling) { + uint32_t *words = reinterpret_cast(&tiling); + for (uint32_t idx = 0; idx < sizeof(FAInferTilingData) / sizeof(uint32_t); ++idx) { + words[idx] = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void +fill_basic(FAInferTilingData &tiling, uint32_t batch, uint32_t max_blocks_per_batch, uint32_t num_blocks) { + tiling.numHeads = kNumHeads; + tiling.embeddingSize = kHeadDim; + tiling.embeddingSizeV = kHeadDim; + tiling.numBlocks = num_blocks; + tiling.blockSize = kBlockSize; + tiling.kvHeads = kNumKvHeads; + tiling.batch = batch; + tiling.maxNumBlocksPerBatch = max_blocks_per_batch; + tiling.maskType = 1; + tiling.scaleValue = 0.08838834764831843F; + tiling.preToken = kSparseTokenLimit; + tiling.nextToken = kSparseTokenLimit; + tiling.sparseMode = 3; +} + +QWEN_FAI_TILER_FUNCTION void fill_task_counts(FAInferTilingData &tiling, uint32_t batch) { + constexpr uint32_t tasks_per_batch = kNumKvHeads; + tiling.firstBatchTaskNum = tasks_per_batch; + tiling.totalTaskNum = tasks_per_batch * batch; +} + +QWEN_FAI_TILER_FUNCTION void init_core_info(FAInferTilingData &tiling, uint32_t planning_core_num) { + for (uint32_t core_idx = 0; core_idx < planning_core_num; ++core_idx) { + tiling.coreInfo.startBIdx[core_idx] = 0; + tiling.coreInfo.startN1Idx[core_idx] = 0; + tiling.coreInfo.startS1Idx[core_idx] = 0; + tiling.coreInfo.startS2Idx[core_idx] = 0; + tiling.coreInfo.endBIdx[core_idx] = 0; + tiling.coreInfo.endN1Idx[core_idx] = 0; + tiling.coreInfo.endS1Idx[core_idx] = 0; + tiling.coreInfo.endS2Idx[core_idx] = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void consume_s2_blocks( + const uint32_t *kv_seqlens, int64_t &remaining_tasks, uint32_t batch_idx, uint32_t s1_idx, uint32_t &s2_idx +) { + while (s2_idx < get_batch_params(batch_idx, kv_seqlens).ks_block_num && remaining_tasks > 0) { + BatchParams params = get_batch_params(batch_idx, kv_seqlens); + uint32_t remaining_q = s1_idx < params.qs_block_num - 1 ? + params.qs_block_tile : + (params.q_seqlen - s1_idx * params.qs_block_tile) * params.qn_block_tile; + uint32_t remaining_kv = + s2_idx < params.ks_block_num - 1 ? params.ks_block_tile : params.kv_seqlen - s2_idx * params.ks_block_tile; + remaining_tasks -= static_cast(remaining_q) * remaining_kv; + ++s2_idx; + } +} + +QWEN_FAI_TILER_FUNCTION void consume_remaining_batches( + const uint32_t *kv_seqlens, uint32_t batch, int64_t &remaining_tasks, uint32_t &batch_idx, uint32_t &n1_idx, + uint32_t &s1_idx, uint32_t &s2_idx +) { + while (batch_idx < batch && remaining_tasks > 0) { + BatchParams params = get_batch_params(batch_idx, kv_seqlens); + uint32_t remaining_q = + params.q_seqlen * (kNumHeads - params.qn_block_tile * n1_idx) - s1_idx * params.qs_block_tile; + uint32_t remaining_in_batch = remaining_q * params.kv_seqlen; + if (remaining_tasks < static_cast(remaining_in_batch)) { + break; + } + remaining_tasks -= remaining_in_batch; + ++batch_idx; + n1_idx = 0; + s1_idx = 0; + s2_idx = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void consume_remaining_n1_groups( + int64_t &remaining_tasks, const BatchParams ¶ms, uint32_t &n1_idx, uint32_t &s1_idx, uint32_t &s2_idx +) { + while (n1_idx < params.qn_block_num && remaining_tasks > 0) { + uint32_t remaining_q = params.q_seqlen * params.qn_block_tile - s1_idx * params.qs_block_tile; + uint32_t remaining_in_n1 = remaining_q * params.kv_seqlen; + if (remaining_tasks < static_cast(remaining_in_n1)) { + break; + } + remaining_tasks -= remaining_in_n1; + ++n1_idx; + s1_idx = 0; + s2_idx = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void +consume_remaining_s1_groups(int64_t &remaining_tasks, const BatchParams ¶ms, uint32_t &s1_idx, uint32_t &s2_idx) { + while (s1_idx < params.qs_block_num && remaining_tasks > 0) { + uint32_t remaining_q = s1_idx < params.qs_block_num - 1 ? + params.qs_block_tile : + (params.q_seqlen - s1_idx * params.qs_block_tile) * params.qn_block_tile; + uint64_t remaining_in_s1 = static_cast(remaining_q) * params.kv_seqlen; + if (remaining_tasks < static_cast(remaining_in_s1)) { + break; + } + remaining_tasks -= remaining_in_s1; + ++s1_idx; + s2_idx = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void +finish_batch(FAInferTilingData &tiling, const uint32_t *kv_seqlens, uint32_t batch, uint32_t core_idx) { + BatchParams params = get_batch_params(batch - 1, kv_seqlens); + tiling.coreInfo.endBIdx[core_idx] = batch - 1; + tiling.coreInfo.endN1Idx[core_idx] = params.qn_block_num - 1; + tiling.coreInfo.endS1Idx[core_idx] = params.qs_block_num - 1; + tiling.coreInfo.endS2Idx[core_idx] = params.ks_block_num; + tiling.needCoreNum = core_idx + 1; +} + +QWEN_FAI_TILER_FUNCTION void +advance_counters(const BatchParams ¶ms, uint32_t &batch_idx, uint32_t &n1_idx, uint32_t &s1_idx, uint32_t &s2_idx) { + if (s2_idx == params.ks_block_num) { + ++s1_idx; + s2_idx = 0; + } + if (s1_idx == params.qs_block_num) { + ++n1_idx; + s1_idx = 0; + s2_idx = 0; + } + if (n1_idx == params.qn_block_num) { + ++batch_idx; + n1_idx = 0; + s1_idx = 0; + s2_idx = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void init_split_info(FAInferTilingData &tiling, uint32_t planning_core_num) { + for (uint32_t split_idx = 0; split_idx < planning_core_num + 1; ++split_idx) { + tiling.splitInfo.batchIdx[split_idx] = 0; + tiling.splitInfo.headStartIdx[split_idx] = 0; + tiling.splitInfo.headEndIdx[split_idx] = 0; + tiling.splitInfo.qStartIdx[split_idx] = 0; + tiling.splitInfo.qEndIdx[split_idx] = 0; + tiling.splitInfo.splitNum[split_idx] = 0; + tiling.splitInfo.lseTaskOffset[split_idx] = 0; + tiling.splitInfo.oTaskOffset[split_idx] = 0; + } +} + +QWEN_FAI_TILER_FUNCTION void process_core_split_info( + FAInferTilingData &tiling, const uint32_t *kv_seqlens, uint32_t core_idx, int32_t &split_idx, + int32_t &prev_batch_idx, int32_t &prev_n1_idx, int32_t &prev_s1_idx, int64_t ¤t_lse_offset, + int64_t ¤t_o_offset, uint32_t planning_core_num +) { + int32_t start_batch_idx = tiling.coreInfo.startBIdx[core_idx]; + int32_t start_n1_idx = tiling.coreInfo.startN1Idx[core_idx]; + int32_t start_s1_idx = tiling.coreInfo.startS1Idx[core_idx]; + int32_t start_s2_idx = tiling.coreInfo.startS2Idx[core_idx]; + int32_t end_batch_idx = tiling.coreInfo.endBIdx[core_idx]; + int32_t end_n1_idx = tiling.coreInfo.endN1Idx[core_idx]; + int32_t end_s1_idx = tiling.coreInfo.endS1Idx[core_idx]; + int32_t end_s2_idx = tiling.coreInfo.endS2Idx[core_idx]; + + tiling.coreInfo.firstSplitKVTaskLseOffset[core_idx] = 0; + tiling.coreInfo.firstSplitKVTaskOOffset[core_idx] = 0; + bool found_first_split = false; + + for (int32_t batch_idx = start_batch_idx; batch_idx <= end_batch_idx; ++batch_idx) { + BatchParams params = get_batch_params(batch_idx, kv_seqlens); + int32_t current_start_n1 = batch_idx == start_batch_idx ? start_n1_idx : 0; + int32_t current_end_n1 = batch_idx == end_batch_idx ? end_n1_idx : params.qn_block_num - 1; + for (int32_t n1_idx = current_start_n1; n1_idx <= current_end_n1; ++n1_idx) { + int32_t current_start_s1 = batch_idx == start_batch_idx && n1_idx == start_n1_idx ? start_s1_idx : 0; + int32_t current_end_s1 = + batch_idx == end_batch_idx && n1_idx == end_n1_idx ? end_s1_idx : params.qs_block_num - 1; + for (int32_t s1_idx = current_start_s1; s1_idx <= current_end_s1; ++s1_idx) { + int32_t current_start_s2 = + batch_idx == start_batch_idx && n1_idx == start_n1_idx && s1_idx == start_s1_idx ? start_s2_idx : 0; + int32_t current_end_s2 = batch_idx == end_batch_idx && n1_idx == end_n1_idx && s1_idx == end_s1_idx ? + end_s2_idx : + params.ks_block_num; + uint32_t covered_s2 = current_end_s2 - current_start_s2; + bool is_split = covered_s2 > 0 && covered_s2 < params.ks_block_num; + if (!is_split) { + continue; + } + + int64_t temporary_lse_offset = current_lse_offset; + int64_t temporary_o_offset = current_o_offset; + uint32_t n1_per_group = n1_idx % params.qn_blocks_per_group; + uint32_t kv_head_idx = n1_idx / params.qn_blocks_per_group; + uint32_t head_start = kv_head_idx * (kNumHeads / kNumKvHeads) + n1_per_group * params.qn_block_tile; + uint32_t head_end = + min_u32(head_start + params.qn_block_tile, (kv_head_idx + 1) * (kNumHeads / kNumKvHeads)); + uint32_t q_start = s1_idx * params.qs_block_tile; + uint32_t q_end = min_u32(q_start + params.qs_block_tile, params.q_seqlen); + uint32_t head_len = head_end - head_start; + uint32_t q_len = q_end - q_start; + + if (batch_idx != prev_batch_idx || n1_idx != prev_n1_idx || s1_idx != prev_s1_idx) { + ++split_idx; + if (split_idx >= 0 && split_idx < static_cast(planning_core_num + 1)) { + tiling.splitInfo.batchIdx[split_idx] = batch_idx; + tiling.splitInfo.splitNum[split_idx] = 0; + tiling.splitInfo.headStartIdx[split_idx] = head_start; + tiling.splitInfo.headEndIdx[split_idx] = head_end; + tiling.splitInfo.qStartIdx[split_idx] = q_start; + tiling.splitInfo.qEndIdx[split_idx] = q_end; + tiling.splitInfo.lseTaskOffset[split_idx] = current_lse_offset; + tiling.splitInfo.oTaskOffset[split_idx] = current_o_offset; + } + prev_batch_idx = batch_idx; + prev_n1_idx = n1_idx; + prev_s1_idx = s1_idx; + } + if (split_idx >= 0 && split_idx < static_cast(planning_core_num + 1)) { + ++tiling.splitInfo.splitNum[split_idx]; + current_lse_offset += static_cast(head_len) * q_len; + current_o_offset += static_cast(head_len) * q_len * kHeadDim; + } + if (!found_first_split) { + found_first_split = true; + tiling.coreInfo.firstSplitKVTaskLseOffset[core_idx] = temporary_lse_offset; + tiling.coreInfo.firstSplitKVTaskOOffset[core_idx] = temporary_o_offset; + } + } + } + } +} + +QWEN_FAI_TILER_FUNCTION void fill_workspace(FAInferTilingData &tiling) { + tiling.mm1OutSize = static_cast(kPlanningCoreNum) * kWorkspaceBlockSize * kUint32Bytes * kPrelaunchNum; + tiling.smOnlineOutSize = + static_cast(kPlanningCoreNum) * kWorkspaceBlockSize * kUint16Bytes * kPrelaunchNum; + tiling.mm2OutSize = tiling.mm1OutSize; + tiling.UpdateSize = tiling.mm1OutSize; + tiling.workSpaceSize = kBaseWorkspaceSize + tiling.splitLseTotalSize + tiling.splitOTotalSize; +} + +QWEN_FAI_TILER_FUNCTION void +initialize_tiling(FAInferTilingData &tiling, uint32_t batch, uint32_t max_blocks_per_batch, uint32_t num_blocks) { + zero_tiling(tiling); + fill_basic(tiling, batch, max_blocks_per_batch, num_blocks); + fill_task_counts(tiling, batch); +} + +QWEN_FAI_TILER_FUNCTION bool build( + const uint32_t *kv_seqlens, uint32_t batch, uint32_t max_blocks_per_batch, uint32_t num_blocks, + FAInferTilingData &tiling, int64_t *cumulative_q_lengths, int64_t *kv_lengths +) { + if (batch == 0 || batch > kMaxBatch) { + return false; + } + for (uint32_t batch_idx = 0; batch_idx < batch; ++batch_idx) { + cumulative_q_lengths[batch_idx] = batch_idx + 1; + kv_lengths[batch_idx] = kv_seqlens[batch_idx]; + } + initialize_tiling(tiling, batch, max_blocks_per_batch, num_blocks); + fill_workspace(tiling); + return true; +} + +} // namespace qwen_fai_tiler + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/arch.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/arch.hpp new file mode 100644 index 0000000000..245d272b80 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/arch.hpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef ARCH_ARCH_HPP +#define ARCH_ARCH_HPP + +#include "../../attn_infra/base_defs.hpp" + +namespace NpuArch::Arch +{ + +struct AtlasA2 { + static constexpr uint32_t BIAS_SIZE = 1024; + static constexpr uint32_t FIXBUF_SIZE = 7U * 1024U; + static constexpr uint32_t UB_SIZE = 192U * 1024U; + static constexpr uint32_t L1_SIZE = 512U * 1024U; + static constexpr uint32_t L0A_SIZE = 64U * 1024U; + static constexpr uint32_t L0B_SIZE = 64U * 1024U; + static constexpr uint32_t L0C_SIZE = 128U * 1024U; +}; + +struct PositionGM { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::GM; +}; + +struct PositionL1 { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::A1; +}; + +struct PositionL0A { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::A2; +}; + +struct PositionL0B { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::B2; +}; + +struct PositionL0C { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::CO1; +}; + +struct PositionUB { + static constexpr AscendC::TPosition POSITION = AscendC::TPosition::VECCALC; +}; + +} // namespace NpuArch::Arch + +#endif // ARCH_ARCH_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/cross_core_sync.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/cross_core_sync.hpp new file mode 100644 index 0000000000..156d2dd1d2 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/cross_core_sync.hpp @@ -0,0 +1,120 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef ARCH_CROSS_CORE_SYNC_HPP +#define ARCH_CROSS_CORE_SYNC_HPP + +#include "../../attn_infra/base_defs.hpp" + +namespace NpuArch::Arch +{ + +constexpr uint32_t MAX_REVERSE_DEPTH = 16; + +using FlagID = uint16_t; +constexpr FlagID AIV_INTER_BLOCK_BARRIER = 8; +constexpr FlagID AIC_INTER_BLOCK_BARRIER = 9; +constexpr FlagID AIV_INTER_SUBBLOCK_BARRIER = 10; +constexpr FlagID FFTS_MAX_FLAG = 7; + +struct CrossCoreFlag { + __aicore__ inline + CrossCoreFlag() : id(0) {} + + __aicore__ inline + CrossCoreFlag(FlagID id) : id(id) {} + + FlagID id; +}; + +template +struct CrossCoreFlagWithReverse { + __aicore__ inline + CrossCoreFlagWithReverse() : id(0), reverseId(0) {} + + __aicore__ inline + CrossCoreFlagWithReverse(FlagID id, FlagID reverseId) : id(id), reverseId(reverseId) {} + + FlagID id; + FlagID reverseId; + uint32_t count{ 0 }; +}; + +template +struct BarrierFlag { + static_assert(MODE != MODE, "Unsupported cross core barrier flag, can not find the specialization."); +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIC> { + static constexpr FlagID ID = AIC_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x1, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_SUBBLOCK_BARRIER; +}; + +template +__aicore__ inline +void CrossCoreBarrier() +{ + FlagID flagId; + if (g_coreType == AscendC::AIC) { + flagId = BarrierFlag::ID; + } else if (g_coreType == AscendC::AIV) { + flagId = BarrierFlag::ID; + } + AscendC::CrossCoreSetFlag(flagId); + AscendC::CrossCoreWaitFlag(flagId); +} + +template +__aicore__ inline +void CrossCoreSetFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); +} + +__aicore__ inline void CrossCoreWaitFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); +} + +template +__aicore__ inline +void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreWaitFlag(flag.reverseId); + flag.count = 0; + } +} + +template +__aicore__ inline +void CrossCoreWaitFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreSetFlag(flag.reverseId); + flag.count = 0; + } +} + +} // namespace NpuArch::Arch + +#endif // ARCH_CROSS_CORE_SYNC_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/local_tensor_buffer.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/local_tensor_buffer.hpp new file mode 100644 index 0000000000..467d319834 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/local_tensor_buffer.hpp @@ -0,0 +1,234 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef INCLUDE_ARCH_MEMORY_H +#define INCLUDE_ARCH_MEMORY_H + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/arch/arch.hpp" + +namespace NpuArch::Arch +{ + +struct LocalTensorBufferBase { +public: + template + __aicore__ inline + AscendC::LocalTensor GetBufferByByte(const uint32_t offset) const + { + return tensor[offset].template ReinterpretCast(); + } + +protected: + __aicore__ inline + LocalTensorBufferBase() = default; + + AscendC::LocalTensor tensor; +}; + +template < + class ArchTag, + AscendC::TPosition Position +> +struct LocalTensorBuffer { + static_assert(DEPENDENT_FALSE, "Unsupported local tensor buffer, can not find the specialization."); +}; + +/// Partial specialization for TPosition::A1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A1; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufA1; + GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE); + tensor = tbufA1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::A2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A2; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufA2; + GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE); + tensor = tbufA2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B1; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufB1; + GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE); + tensor = tbufB1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::B2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B2; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufB2; + GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE); + tensor = tbufB2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C1 +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C1; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufC1; + GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE); + tensor = tbufC1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C2 +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufC2; + GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE); + tensor = tbufC2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::CO1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufCO1; + GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE); + tensor = tbufCO1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for AtlasA2, TPosition::C2PIPE2GM +template <> +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + using ArchTag = Arch::AtlasA2; + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufC2PIPE2GM; + GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE); + tensor = tbufC2PIPE2GM.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECIN +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECIN; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufVECIN; + GetTPipePtr()->InitBuffer(tbufVECIN, ArchTag::UB_SIZE); + tensor = tbufVECIN.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECOUT +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECOUT; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufVECOUT; + GetTPipePtr()->InitBuffer(tbufVECOUT, ArchTag::UB_SIZE); + tensor = tbufVECOUT.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECCALC +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC; + + __aicore__ inline + LocalTensorBuffer() + { + AscendC::TBuf tbufVECCALC; + GetTPipePtr()->InitBuffer(tbufVECCALC, ArchTag::UB_SIZE); + tensor = tbufVECCALC.Get(); + } +}; + +} // namespace NpuArch::Arch + +#endif // INCLUDE_ARCH_MEMORY_H diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/resource.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/resource.hpp new file mode 100644 index 0000000000..7db93ee746 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/resource.hpp @@ -0,0 +1,51 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef INCLUDE_ARCH_RESOURCE_HPP +#define INCLUDE_ARCH_RESOURCE_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/arch/local_tensor_buffer.hpp" + +namespace NpuArch::Arch +{ + +struct PtoTopology { + uint32_t logicalBlockIdx; + uint32_t logicalBlockNum; + uint32_t subBlockIdx; + uint32_t lanesPerBlock; +}; + +template +struct Resource { +public: + AscendC::TPipe pipe; + + LocalTensorBuffer l1Buf; + LocalTensorBuffer l0ABuf; + LocalTensorBuffer l0BBuf; + LocalTensorBuffer btBuf; + LocalTensorBuffer l0CBuf; + LocalTensorBuffer ubBuf; + PtoTopology ptoTopology{0, 1, 0, 1}; + + __aicore__ inline + Resource() + { + // The initialization of AscendC::Tpipe will insert some synchronization interfaces, + // which may conflict with the usage by users. Therefore, the "destroy" interface is used for releasing. + pipe.Destroy(); + } +}; + +} // namespace NpuArch::Arch + +#endif // INCLUDE_ARCH_RESOURCE_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/base_defs.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/base_defs.hpp new file mode 100644 index 0000000000..0ff58ade47 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/base_defs.hpp @@ -0,0 +1,46 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! + * \file base_defs.hpp + * \brief + */ + +#ifndef HPP_HPP +#define HPP_HPP + + +#if ASC_DEVKIT_MAJOR >= 9 +#include "basic_api/kernel_basic_intf.h" +#else +#include "kernel_operator.h" +#endif + +#include "../attn_infra/detail/alignment.hpp" +#include "../attn_infra/detail/dependent_false.hpp" +#include "../attn_infra/detail/macros.hpp" + +namespace NpuArch { + +constexpr uint32_t BYTE_PER_C0 = 32; +constexpr uint32_t BYTE_PER_C2 = 64; +constexpr uint32_t C0_NUM_PER_FRACTAL = 16; +constexpr uint32_t BYTE_PER_FRACTAL = BYTE_PER_C0 * C0_NUM_PER_FRACTAL; + +constexpr uint32_t BYTE_PER_BLK = 32; +constexpr uint32_t BLK_NUM_PER_VECTOR_FRACTAL = 8; +constexpr uint32_t BYTE_PER_VECTOR_FRACTAL = BYTE_PER_BLK * BLK_NUM_PER_VECTOR_FRACTAL; + +constexpr uint64_t L2_OFFSET = 0; +constexpr uint32_t STRIDE_LIMIT = 65536; + +} // namespace NpuArch + +#endif // HPP_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/coord.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/coord.hpp new file mode 100644 index 0000000000..d0c5447d7c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/coord.hpp @@ -0,0 +1,442 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! + * \file coord.hpp + * \brief + */ + +#ifndef COORD_HPP +#define COORD_HPP + +#include "../attn_infra/base_defs.hpp" + +namespace NpuArch { + +/// Statically-sized array specifying Coords within a tensor +template < + int RANK_, ///< Logical rank of coordinate + class Index_ = uint32_t, ///< Index type used for each dimension + class LongIndex_ = int64_t ///< Long index type used for linear offsets +> +struct Coord { +public: + // Number of elements in Coord + static const int RANK = RANK_; + + // Index typen used to store elements + using Index = Index_; + + // Type used to represent linear offsets + using LongIndex = LongIndex_; + + // Default ctor initializes uniformly + HOST_DEVICE constexpr + explicit Coord(Index value = Index(0)) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = value; + } + } + + // Constructs from an array of integers + HOST_DEVICE constexpr + Coord(Index const (&idx_)[RANK]) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = idx_[i]; + } + } + + HOST_DEVICE + int Argmin() const + { + return ArgminImpl<1>(0); + } + + // Returns the index of the dimension with greatest value + HOST_DEVICE + int Argmax() const + { + return ArgmaxImpl<1>(0); + } + + // Returns true if Coord is non-zero + HOST_DEVICE + explicit operator bool() const + { + return AnyImpl<0>(); + } + + // Return true if Coord is uniformly zero. + HOST_DEVICE + bool operator!() const + { + return !AnyImpl<0>(); + } + + // Element-wise addition + HOST_DEVICE + Coord operator+(Coord const &b) const + { + Coord c; + AddCoordImpl<0>(c, b); + return c; + } + + // Add a scalar to each element + HOST_DEVICE + Coord operator+(const Index val) const + { + Coord c; + AddScalarImpl<0>(c, val); + return c; + } + + // Element-wise subtraction + HOST_DEVICE + Coord operator-(Coord const &b) const + { + Coord c; + SubCoordImpl<0>(c, b); + return c; + } + + // Subtract a scalar from each element + HOST_DEVICE + Coord operator-(Index const val) const + { + Coord c; + SubScalarImpl<0>(c, val); + return c; + } + + // Element-wise multiply + HOST_DEVICE + Coord operator*(Coord const &b) const + { + Coord c; + MulCoordImpl<0>(c, b); + return c; + } + + // Element-wise division + HOST_DEVICE + Coord operator/(Coord const &b) const + { + Coord c; + DivCoordImpl<0>(c, b); + return c; + } + + // Element-wise mod + HOST_DEVICE + Coord operator%(Coord const &b) const + { + Coord c; + ModCoordImpl<0>(c, b); + return c; + } + + // In-place addition + HOST_DEVICE + Coord &operator+=(Coord const &b) + { + PlusEqualImpl<0>(b); + return *this; + } + + // In-place equal + HOST_DEVICE + bool operator==(Coord const &b) const + { + return EqualCoordImpl<0>(b); + } + + // In-place equal + HOST_DEVICE + bool operator==(Index const val) const + { + return EqualScalarImpl<0>(val); + } + + // Member acces operator + HOST_DEVICE + Index &operator[](int dim) + { + return idx[dim]; + } + + // Member access operator + HOST_DEVICE + Index const &operator[](int dim) const + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + HOST_DEVICE + Index &At() + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + HOST_DEVICE + Index &At(int dim) + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + HOST_DEVICE + Index const &At() const + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + HOST_DEVICE + Index const &At(int dim) const + { + return idx[dim]; + } + + template + HOST_DEVICE + auto GetCoordByAxis() const + { + Index idx_[sizeof...(Is)]{idx[Is]...}; + return Coord{idx_}; + } + + HOST_DEVICE + static Coord Min(Coord const &a, Coord const &b) + { + Coord res; + for (int i = 0; i < RANK; ++i) { + res[i] = a[i] < b[i] ? a[i] : b[i]; + } + return res; + } + +private: + template + HOST_DEVICE + int ArgminImpl(int i) const + { + if constexpr (N == RANK) { + return i; + } + else { + return ArgminImpl(idx[N] < idx[i] ? N : i); + } + } + + template + HOST_DEVICE + int ArgmaxImpl(int i) const + { + if constexpr (N == RANK) { + return i; + } + else { + return ArgmaxImpl(idx[N] > idx[i] ? N : i); + } + } + + template + HOST_DEVICE + bool AnyImpl() const + { + if constexpr (N == RANK) { + return false; + } + else { + return idx[N] || AnyImpl(); + } + } + + template + HOST_DEVICE + void AddCoordImpl(Coord &c, Coord const &b) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] + b.idx[N]; + AddCoordImpl(c, b); + } + } + + template + HOST_DEVICE + void AddScalarImpl(Coord &c, Index const val) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] + val; + AddScalarImpl(c, val); + } + } + + template + HOST_DEVICE + void SubCoordImpl(Coord &c, Coord const &b) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] - b.idx[N]; + SubCoordImpl(c, b); + } + } + + template + HOST_DEVICE + void SubScalarImpl(Coord &c, Index const val) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] - val; + SubScalarImpl(c, val); + } + } + + template + HOST_DEVICE + void MulCoordImpl(Coord &c, Coord const &b) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] * b.idx[N]; + MulCoordImpl(c, b); + } + } + + template + HOST_DEVICE + void DivCoordImpl(Coord &c, Coord const &b) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] / b.idx[N]; + DivCoordImpl(c, b); + } + } + + template + HOST_DEVICE + void ModCoordImpl(Coord &c, Coord const &b) const + { + if constexpr (N < RANK) { + c.idx[N] = idx[N] % b.idx[N]; + ModCoordImpl(c, b); + } + } + + template + HOST_DEVICE + void PlusEqualImpl(Coord const &b) + { + if constexpr (N < RANK) { + idx[N] += b.idx[N]; + PlusEqualImpl(b); + } + } + + template + HOST_DEVICE + bool EqualCoordImpl(Coord const &b) const + { + if constexpr (N == RANK) { + return true; + } + else { + return idx[N] == b.idx[N] && EqualCoordImpl(b); + } + } + + template + HOST_DEVICE + bool EqualScalarImpl(Index const val) const + { + if constexpr (N == RANK) { + return true; + } + else { + return idx[N] == val && EqualScalarImpl(val); + } + } + + // Indices + Index idx[RANK]; +}; + +// Helper to make a 1-element coordinate +template +HOST_DEVICE constexpr +Coord<1, T> MakeCoord(T dim0) +{ + T values[1] = {dim0}; + return Coord<1, T>(values); +} + +/// Helper to make a 2-element coordinate +template +HOST_DEVICE constexpr +Coord<2, T> MakeCoord(T dim0, T dim1) +{ + T values[2] = {dim0, dim1}; + return Coord<2, T>(values); +} + +/// Helper to make a 3-element coordinate +template +HOST_DEVICE constexpr +Coord<3, T> MakeCoord(T dim0, T dim1, T dim2) +{ + T values[3] = {dim0, dim1, dim2}; + return Coord<3, T>(values); +} + +/// Helper to make a 4-element coordinate +template +HOST_DEVICE constexpr +Coord<4, T> MakeCoord(T dim0, T dim1, T dim2, T dim3) +{ + T values[4] = {dim0, dim1, dim2, dim3}; + return Coord<4, T>(values); +} + +/// Helper to make a 5-element coordinate +template +HOST_DEVICE constexpr +Coord<5, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4) +{ + T values[5] = {dim0, dim1, dim2, dim3, dim4}; + return Coord<5, T>(values); +} + +/// Helper to make a 6-element coordinate +template +HOST_DEVICE constexpr +Coord<6, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5) +{ + T values[6] = {dim0, dim1, dim2, dim3, dim4, dim5}; + return Coord<6, T>(values); +} + +/// Helper to make a 7-element coordinate +template +HOST_DEVICE constexpr +Coord<7, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5, T dim6) +{ + T values[7] = {dim0, dim1, dim2, dim3, dim4, dim5, dim6}; + return Coord<7, T>(values); +} + +} // namespace NpuArch + +#endif // COORD_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/alignment.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/alignment.hpp new file mode 100644 index 0000000000..4ae3201203 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/alignment.hpp @@ -0,0 +1,75 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef ALIGNMENT_HPP +#define ALIGNMENT_HPP + +#include "../../attn_infra/detail/macros.hpp" + +namespace NpuArch::Detail::Alignment +{ + +template +HOST_DEVICE +constexpr T RoundUp(const T &val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return (val + ALIGN - 1) / ALIGN * ALIGN; +} + +template +HOST_DEVICE +constexpr auto RoundUp(T const &val, U const &align) +{ + if (align == 0) { + return val; + } + return (val + align - 1) / align * align; +} + +template +HOST_DEVICE +constexpr T RoundDown(const T val) +{ + static_assert(ALIGN != 0U, "ALIGN must not be 0"); + return val / ALIGN * ALIGN; +} + +template +HOST_DEVICE +constexpr T RoundDown(const T val, const T align) +{ + if (align == 0) { + return val; + } + return val / align * align; +} + +template +HOST_DEVICE +constexpr T CeilDiv(const T dividend) +{ + static_assert(DIVISOR != 0, "DIVISOR must not be 0"); + return (dividend + DIVISOR - 1) / DIVISOR; +} + +template +HOST_DEVICE +constexpr auto CeilDiv(T const ÷nd, U const &divisor) +{ + if (divisor == 0) { + return dividend; + } + return (dividend + divisor - 1) / divisor; +} + +} + +#endif // ALIGNMENT_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/dependent_false.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/dependent_false.hpp new file mode 100644 index 0000000000..91f85dff92 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/dependent_false.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef DETAIL_DEPENDENT_FALSE_HPP +#define DETAIL_DEPENDENT_FALSE_HPP + +template +constexpr bool DEPENDENT_BOOL_VALUE = VALUE; + +template +constexpr bool DEPENDENT_FALSE = DEPENDENT_BOOL_VALUE; + +#endif // DETAIL_DEPENDENT_FALSE_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/macros.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/macros.hpp new file mode 100644 index 0000000000..3d23fe819e --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/macros.hpp @@ -0,0 +1,16 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef DETAIL_MACROS_HPP +#define DETAIL_MACROS_HPP + +#define HOST_DEVICE __host_aicore__ inline + +#endif // DETAIL_MACROS_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/CombineScale.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/CombineScale.hpp new file mode 100644 index 0000000000..393d45db5a --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/CombineScale.hpp @@ -0,0 +1,290 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_COMBINE_SCALE_HPP +#define EPILOGUE_BLOCK_COMBINE_SCALE_HPP + +#include +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "adv_api/pad/broadcast.h" +#include "adv_api/reduce/reduce.h" + +namespace NpuArch::Epilogue::Block { + +template < + class OutputType_, + class LseType_> +class CombineScale { +public: + using ElementOutput = typename OutputType_::Element; + using ElementLse = typename LseType_::Element; + using ArchTag = Arch::AtlasA2; + + static constexpr uint32_t STAGE2_UB_UINT8_BLOCK_SIZE = 6144; // 24 * 64 * 4 + static constexpr uint32_t UB_UINT8_LINE_SIZE = 32768; // 1 * 64 * 128 * 4 + + __aicore__ inline + CombineScale() {} + + __aicore__ inline + ~CombineScale() {} + + __aicore__ inline + void init(Arch::Resource &resource) { + ptoLogicalBlockIdx = resource.ptoTopology.logicalBlockIdx; + ptoLogicalBlockNum = resource.ptoTopology.logicalBlockNum; + ptoSubBlockIdx = resource.ptoTopology.subBlockIdx; + ptoLanesPerBlock = resource.ptoTopology.lanesPerBlock; + // UB Memory Allocation + constexpr uint32_t LL_UB_OFFSET = 0; // splitnum_align * (q * h)_algin + constexpr uint32_t LM_UB_OFFSET = 1 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin + constexpr uint32_t BROADCAST_OFFSET = 2 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin + constexpr uint32_t TL_UB_OFFSET = 3 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin + constexpr uint32_t RS_UB_OFFSET = 4 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin + constexpr uint32_t TS_UB_OFFSET = 5 * STAGE2_UB_UINT8_BLOCK_SIZE; // 1 * (q * h)_algin + constexpr uint32_t BROADCASTSCALE_OFFSET = 6 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin + constexpr uint32_t GL_UB_OFFSET = 7 * STAGE2_UB_UINT8_BLOCK_SIZE; // splitnum_align * (q * h)_algin + constexpr uint32_t BROADCASTO_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE; //splitnum_align * (q * h)_algin * v + constexpr uint32_t GO_UB_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE + 1 * UB_UINT8_LINE_SIZE; // (q * h)_algin * v + constexpr uint32_t GO16_UB_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE + 2 * UB_UINT8_LINE_SIZE; // (q * h)_algin * v + constexpr uint32_t tempReduceMax_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE + 3 * UB_UINT8_LINE_SIZE + 1 * STAGE2_UB_UINT8_BLOCK_SIZE; //splitnum_align * (q * h)_algin * v + constexpr uint32_t tempReduceSum_OFFSET = 8 * STAGE2_UB_UINT8_BLOCK_SIZE + 3 * UB_UINT8_LINE_SIZE + 2 * STAGE2_UB_UINT8_BLOCK_SIZE; //splitnum_align * (q * h)_algin * v + + // Buffer Init + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_OFFSET); + broadCastTensor = resource.ubBuf.template GetBufferByByte(BROADCAST_OFFSET); + tlUbTensor = resource.ubBuf.template GetBufferByByte(TL_UB_OFFSET); + rsUbTensor = resource.ubBuf.template GetBufferByByte(RS_UB_OFFSET); + tsUbTensor = resource.ubBuf.template GetBufferByByte(TS_UB_OFFSET); + broadCastScaleTensor = resource.ubBuf.template GetBufferByByte(BROADCASTSCALE_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_OFFSET); + broadCastOTensor = resource.ubBuf.template GetBufferByByte(BROADCASTO_OFFSET); + toUbTensor = resource.ubBuf.template GetBufferByByte(BROADCASTO_OFFSET); + goUbTensor = resource.ubBuf.template GetBufferByByte(GO_UB_OFFSET); + loFloatUbTensor = resource.ubBuf.template GetBufferByByte(GO16_UB_OFFSET); + go16UbTensor = resource.ubBuf.template GetBufferByByte(GO_UB_OFFSET); + + tempReduceMax = resource.ubBuf.template GetBufferByByte(tempReduceMax_OFFSET); + tempReduceSum = resource.ubBuf.template GetBufferByByte(tempReduceSum_OFFSET); + + } + + + __aicore__ inline void operator()( + uint32_t qHeads, + uint32_t kvSplitCoreNum, + uint32_t headSizeV, + __gm__ splitNode *splitInfo, + AscendC::GlobalTensor lGmTensor, + AscendC::GlobalTensor oCoreTmpGmTensor, + AscendC::GlobalTensor oGmTensor, + AscendC::GlobalTensor gActualQseqlen, + bool inputLayoutTND = true + ) { + AscendC::SetAtomicNone(); + AscendC::SetMaskNorm(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + + AscendC::SetFlag(EVENT_ID0); + + int64_t subBlockNum = ptoLogicalBlockNum * ptoLanesPerBlock; + int64_t subBlockID = ptoLogicalBlockIdx * ptoLanesPerBlock + ptoSubBlockIdx; + + for (uint32_t process = subBlockID; process < kvSplitCoreNum * 2; process += subBlockNum) { + uint32_t vectorsubBlockID = process % 2; + uint32_t batchIdx = splitInfo->batchIdx[process/2]; + uint32_t headStartIndx = splitInfo->headStartIdx[process/2]; + uint32_t headEndIndx = splitInfo->headEndIdx[process/2]; + uint32_t qStartIndx = splitInfo->qStartIdx[process/2]; + uint32_t qEndIndx = splitInfo->qEndIdx[process/2]; + + + uint32_t q_len = (qEndIndx - qStartIndx); + uint32_t n_len = (headEndIndx - headStartIndx); + + uint32_t sum = q_len * n_len; + uint32_t sum_former = q_len == 1 ? sum / 2 : (q_len / 2) * n_len; + + uint32_t addrLOffset = vectorsubBlockID == 0 ? splitInfo->lseTaskOffset[process/2] : splitInfo->lseTaskOffset[process/2] + sum_former; + uint32_t addrOOffset = vectorsubBlockID == 0 ? splitInfo->oTaskOffset[process/2] : splitInfo->oTaskOffset[process/2] + sum_former * headSizeV; + + uint32_t prevQSeqlenSum = 0; + if (inputLayoutTND) { + prevQSeqlenSum = (batchIdx == 0) ? + 0 : static_cast(gActualQseqlen.GetValue(batchIdx - 1)); + } + uint32_t baseGmOffset = prevQSeqlenSum * qHeads * headSizeV + qStartIndx * qHeads * headSizeV + headStartIndx * headSizeV; + uint32_t gmOScalar = 0; + if (q_len == 1) { + gmOScalar = vectorsubBlockID == 0 ? baseGmOffset + : baseGmOffset + sum_former * headSizeV; + } else { + uint32_t q_half = q_len / 2; + gmOScalar = vectorsubBlockID == 0 ? baseGmOffset + : baseGmOffset + q_half * qHeads * headSizeV; + } + + uint32_t splitNum = splitInfo->splitNum[process/2]; + + uint32_t splitNumAlign = (splitNum + 7) / 8 * 8; // 32b align + uint32_t lseBlock = vectorsubBlockID == 0 ? sum_former : sum - sum_former; + uint32_t lseBlockAlign = (lseBlock + 7) / 8 * 8; // 32b align + int32_t count = splitNum * lseBlockAlign; + int32_t lnCount = 1 * lseBlockAlign; + // Initialize LSE UB space + int32_t calcLen = splitNumAlign * lseBlockAlign; + int32_t oCount = lseBlock * headSizeV; + int32_t lseCount = lseBlockAlign * headSizeV; + int32_t oCount_vector = sum * headSizeV; + + + AscendC::Duplicate(llUbTensor, std::numeric_limits::lowest(), calcLen); + AscendC::Duplicate(tlUbTensor, 0.0f, calcLen); + + AscendC::WaitFlag(EVENT_ID0); + + // Copy LSE from GM to UB + uint32_t srcStride = vectorsubBlockID == 0 ? sum - sum_former : sum_former; + AscendC::DataCopyPad(llUbTensor, lGmTensor[addrLOffset], + AscendC::DataCopyExtParams(splitNum, lseBlock * sizeof(float), srcStride * sizeof(float), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, lseBlockAlign - lseBlock, 0)); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ReduceMax + uint32_t reduceMaxShape[] = { splitNumAlign, lseBlockAlign }; + AscendC::ReduceMax(lmUbTensor, llUbTensor, tempReduceMax, reduceMaxShape, true); + AscendC::PipeBarrier(); + + // Broadcast Max + uint32_t dstShapeBroadcast[] = { splitNum, lseBlockAlign }; + uint32_t srcShapeBroadcast[] = { 1, lseBlockAlign }; + AscendC::BroadCast(broadCastTensor, lmUbTensor, dstShapeBroadcast, srcShapeBroadcast, tempReduceSum); + AscendC::PipeBarrier(); + + AscendC::Sub(tlUbTensor, llUbTensor, broadCastTensor, count); + AscendC::PipeBarrier(); + + // expf + AscendC::Exp(tlUbTensor, tlUbTensor, count); + AscendC::PipeBarrier(); + + // ReduceSum + uint32_t reduceSumShape[] = { splitNumAlign, lseBlockAlign }; + AscendC::ReduceSum(rsUbTensor, tlUbTensor, tempReduceSum, reduceSumShape, true); + AscendC::PipeBarrier(); + + // Ln + AscendC::Ln(rsUbTensor, rsUbTensor, lnCount); + AscendC::PipeBarrier(); + + // logf(lse_sum) + lse_max + AscendC::Add(tsUbTensor, rsUbTensor, lmUbTensor, lnCount); + AscendC::PipeBarrier(); + + // Broadcast scale + AscendC::BroadCast(broadCastScaleTensor, tsUbTensor, dstShapeBroadcast, srcShapeBroadcast, tempReduceSum); + AscendC::PipeBarrier(); + + + AscendC::Sub(glUbTensor, llUbTensor, broadCastScaleTensor, count); + AscendC::PipeBarrier(); + + AscendC::Exp(glUbTensor, glUbTensor, count); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID2); + for (uint32_t nIdx = 0; nIdx < splitNum; nIdx++) { + + AscendC::WaitFlag(EVENT_ID2); + AscendC::DataCopyPad(loFloatUbTensor, oCoreTmpGmTensor[addrOOffset + nIdx * oCount_vector], + AscendC::DataCopyExtParams(1, oCount * sizeof(float), 0, 0, 0), + AscendC::DataCopyPadExtParams(true, 0, 0, 0)); + + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + + uint32_t dstShapeO[2] = { lseBlockAlign, headSizeV }; + uint32_t srcShapeO[2] = { lseBlockAlign, 1 }; + AscendC::BroadCast(broadCastOTensor, glUbTensor[nIdx * lseBlockAlign], dstShapeO, srcShapeO, tempReduceSum); + AscendC::PipeBarrier(); + + AscendC::Mul(toUbTensor, loFloatUbTensor, broadCastOTensor, oCount); // toUbTensor和broadCastOTensor共用一块空间 + AscendC::PipeBarrier(); + + if (nIdx == 0) { + AscendC::Adds(goUbTensor, toUbTensor, 0.0f, oCount); // goUbTensor和loFloatUbTensor一块空间 + AscendC::PipeBarrier(); + } else { + AscendC::Add(goUbTensor, toUbTensor, goUbTensor, oCount); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID2); + } + AscendC::WaitFlag(EVENT_ID2); + + // Cast and move out + if (std::is_same::value) { + AscendC::Cast(go16UbTensor, goUbTensor, AscendC::RoundMode::CAST_RINT, oCount); + } else { + AscendC::Cast(go16UbTensor, goUbTensor, AscendC::RoundMode::CAST_NONE, oCount); + } + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + + if (q_len == 1) { + AscendC::DataCopyPad(oGmTensor[gmOScalar], go16UbTensor, AscendC::DataCopyExtParams(1, oCount * sizeof(ElementOutput) , 0, 0, 0)); + } else { + uint32_t q_half = q_len / 2; + if (vectorsubBlockID == 0) { + AscendC::DataCopyPad(oGmTensor[gmOScalar], go16UbTensor, + AscendC::DataCopyExtParams(q_half, (headEndIndx - headStartIndx) * headSizeV * sizeof(ElementOutput) , 0, (qHeads - (headEndIndx - headStartIndx)) * headSizeV * sizeof(ElementOutput), 0)); + } else { + AscendC::DataCopyPad(oGmTensor[gmOScalar], go16UbTensor, + AscendC::DataCopyExtParams(q_len - q_half, (headEndIndx - headStartIndx) * headSizeV * sizeof(ElementOutput) , 0, (qHeads - (headEndIndx - headStartIndx)) * headSizeV * sizeof(ElementOutput), 0)); + } + } + + + AscendC::SetFlag(EVENT_ID0); + } + AscendC::WaitFlag(EVENT_ID0); + } + +private: + uint32_t ptoLogicalBlockIdx = 0; + uint32_t ptoLogicalBlockNum = 1; + uint32_t ptoSubBlockIdx = 0; + uint32_t ptoLanesPerBlock = 1; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor tlUbTensor; + AscendC::LocalTensor rsUbTensor; + AscendC::LocalTensor tsUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor toUbTensor; + AscendC::LocalTensor goUbTensor; + AscendC::LocalTensor go16UbTensor; + AscendC::LocalTensor loFloatUbTensor; + + AscendC::LocalTensor tempReduceMax; + AscendC::LocalTensor tempReduceSum; + AscendC::LocalTensor broadCastTensor; + AscendC::LocalTensor broadCastScaleTensor; + AscendC::LocalTensor broadCastOTensor; +}; +} + +#endif // EPILOGUE_BLOCK_COMBINE_SCALE_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue.hpp new file mode 100644 index 0000000000..cad0338549 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue.hpp @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Block { + +template < + class DispatchPolicy, + class... Args +> +class BlockEpilogue { + static_assert(DEPENDENT_FALSE, "Could not find an epilogue specialization"); +}; + +} // namespace NpuArch::Epilogue::Block + +#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax.hpp" +#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp" +#include "../../../attn_infra/epilogue/block/block_epilogue_rescale_o.hpp" +#include "../../../attn_infra/epilogue/block/CombineScale.hpp" +#include "../../../attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp" +#include "../../../attn_infra/epilogue/block/block_epilogue_init_outputs.hpp" +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_init_outputs.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_init_outputs.hpp new file mode 100644 index 0000000000..cbbaafc63a --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_init_outputs.hpp @@ -0,0 +1,168 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_INIT_OUTPUTS_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_INIT_OUTPUTS_HPP + +#include +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/epilogue/dispatch_policy.hpp" +#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/matrix_coord.hpp" + +namespace NpuArch::Epilogue::Block { + +template < + class AttnOutType_, + class LseOutType_, + LseMode LSE_MODE_> +class BlockEpilogue< + EpilogueAtlasA2InitOutWhenZero, + AttnOutType_, + LseOutType_> +{ +public: + using DispatchPolicy = EpilogueAtlasA2InitOutWhenZero; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementAttnOut = typename AttnOutType_::Element; + using ElementLseOut = typename LseOutType_::Element; + + using LayoutAttnOut = typename AttnOutType_::Layout; + using LayoutLseOut = typename LseOutType_::Layout; + + static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE; + static constexpr float ATTN_OUT_INI = 0; + static constexpr float LSE_OUT_INI = std::numeric_limits::infinity(); + static constexpr uint32_t HALF_ELEM_NUM_PER_BLK = 16; + static constexpr uint32_t FLOAT_ELEM_NUM_PER_BLK = 8; + static constexpr uint32_t HALF_ELEM_NUM_PER_RPT = 128; + static constexpr uint32_t FLOAT_ELEM_NUM_PER_RPT = 64; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + + __aicore__ inline + BlockEpilogue() {} + + __aicore__ inline + void init(Arch::Resource &resource) + { + ptoSubBlockIdx = resource.ptoTopology.subBlockIdx; + ptoLanesPerBlock = resource.ptoTopology.lanesPerBlock; + // Allocate UB space + constexpr uint32_t ATTN_OUT_INIT_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LSE_OUT_INIT_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE; + + attnOutUbTensor = resource.ubBuf.template GetBufferByByte(ATTN_OUT_INIT_UB_TENSOR_OFFSET); + lseOutUbTensor = resource.ubBuf.template GetBufferByByte(LSE_OUT_INIT_UB_TENSOR_OFFSET); + } + + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gLse, + const LayoutAttnOut &layoutOutput, + const LayoutLseOut &layoutLse, + uint32_t qSThisSubBlock, uint32_t qNThisSubBlock) + { + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t qHeads = layoutLse.shape(1); + uint32_t embedV = oHiddenSize / qHeads; + uint32_t embedRoundV = NpuArch::Detail::Alignment::RoundUp(embedV, HALF_ELEM_NUM_PER_BLK); + AscendC::PipeBarrier(); + // init attnOut with 0 + AscendC::WaitFlag(EVENT_ID6); + AscendC::Duplicate(attnOutUbTensor, static_cast(ATTN_OUT_INI), embedRoundV * qSThisSubBlock); + AscendC::SetFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID6); + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + gOutput[qNIdx * embedV], + attnOutUbTensor, + AscendC::DataCopyExtParams( + qSThisSubBlock, embedV * sizeof(ElementAttnOut), + 0, (oHiddenSize - embedV) * sizeof(ElementAttnOut), 0)); + } + AscendC::SetFlag(EVENT_ID6); + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + // init lseOut with inf + AscendC::WaitFlag(EVENT_ID7); + AscendC::Duplicate(lseOutUbTensor, LSE_OUT_INI, qSThisSubBlock * FLOAT_ELEM_NUM_PER_BLK); + AscendC::SetFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID7); + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + gLse[qNIdx], + lseOutUbTensor, + AscendC::DataCopyExtParams( + qSThisSubBlock, sizeof(ElementLseOut), + 0, (qHeads - 1) * sizeof(ElementLseOut), 0)); + } + AscendC::SetFlag(EVENT_ID7); + } + AscendC::PipeBarrier(); + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gLse, + const LayoutAttnOut &layoutOutput, + const LayoutLseOut &layoutLse, + uint32_t qSBlockSize, uint32_t qNBlockSize) + { + uint32_t rowNum = qSBlockSize * qNBlockSize; + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t qHeads = layoutLse.shape(1); + uint32_t embedV = oHiddenSize / qHeads; + + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 1 + : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = + (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualSubBlock = (subBlockIdx == 1U) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetSubBlock = subBlockIdx * rowSplitSubBlock; + uint32_t outRowOffsetSubBlock = (qNBlockSize == 1U) ? rowOffsetSubBlock : 0; + uint32_t outColOffsetSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embedV; + uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? rowActualSubBlock : qSBlockSize; + int64_t outOffsetSubBlock = + layoutOutput.GetOffset(MatrixCoord(outRowOffsetSubBlock, outColOffsetSubBlock)); + auto gOutputSubBlock = gOutput[outOffsetSubBlock]; + auto layoutOutputSubBlock = layoutOutput; + + uint32_t outLseRowOffsetSubBlock = (qNBlockSize == 1U) ? + rowOffsetSubBlock : 0; + uint32_t outLseColOffsetSubBlock = (qNBlockSize == 1U) ? + 0 : subBlockIdx * qNSplitSubBlock; + int64_t lseOffsetSubBlock = + layoutLse.GetOffset(MatrixCoord(outLseRowOffsetSubBlock, outLseColOffsetSubBlock)); + auto gLseThisSubBlock = gLse[lseOffsetSubBlock]; + auto layoutLseThisSubBlock = layoutLse; + + if (rowActualSubBlock > 0U) { + SubCoreCompute( + gOutputSubBlock, gLseThisSubBlock, + layoutOutputSubBlock, layoutLseThisSubBlock, + qSThisSubBlock, qNThisSubBlock); + } + } +private: + uint32_t ptoSubBlockIdx = 0; + uint32_t ptoLanesPerBlock = 1; + AscendC::LocalTensor attnOutUbTensor; + AscendC::LocalTensor lseOutUbTensor; +}; +} +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_INIT_OUTPUTS_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax.hpp new file mode 100644 index 0000000000..9fe7f619c4 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax.hpp @@ -0,0 +1,2055 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP + +#include +#include +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/cross_core_sync.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/epilogue/dispatch_policy.hpp" +#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/matrix_coord.hpp" +#include "utils/std/algorithm.h" + +namespace NpuArch::Epilogue::Block { + +struct SinkLoopParam +{ + uint32_t rowOffsetIoGm; + uint32_t rowNumCurLoop; + uint32_t qSBlockSize; + uint32_t rowOffsetThisSubBlock; + + __aicore__ inline + SinkLoopParam( + uint32_t rowOffsetIoGm_, + uint32_t rowNumCurLoop_, + uint32_t qSBlockSize_, + uint32_t rowOffsetThisSubBlock_ + ) : + rowOffsetIoGm(rowOffsetIoGm_), + rowNumCurLoop(rowNumCurLoop_), + qSBlockSize(qSBlockSize_), + rowOffsetThisSubBlock(rowOffsetThisSubBlock_) + {} +}; + +template < + class OutputType_, + class InputType_, + class MaskType_, + class SinkType_, + class FullType_, + LseMode LSE_MODE_, + SinkMode SINK_MODE_, + MaskMode MASK_MODE_> +class BlockEpilogue< + EpilogueAtlasA2OnlineSoftmax, + OutputType_, + InputType_, + MaskType_, + SinkType_, + FullType_> +{ +public: + using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + using ElementSink = typename SinkType_::Element; + using ElementFull = typename FullType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + using LayoutFull = typename FullType_::Layout; + + static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE; + static constexpr SinkMode SINK_MODE = DispatchPolicy::SINK_MODE; + static constexpr MaskMode MASK_MODE = DispatchPolicy::MASK_MODE; + + static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32; + static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192; + + static constexpr uint32_t REDUCE_UB_SIZE = 1024; + static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; + static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256; + static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; + static constexpr uint32_t HEAD_NUM_2 = 2; + + static constexpr float NEG_INF = -std::numeric_limits::infinity(); + + __aicore__ inline + BlockEpilogue() {} + + __aicore__ inline + void init(Arch::Resource &resource, float scaleValue_) + { + ptoSubBlockIdx = resource.ptoTopology.subBlockIdx; + ptoLanesPerBlock = resource.ptoTopology.lanesPerBlock; + // Allocate UB space + constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; + constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t FULL32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t MASK_UB_PREMASK_TENSOR_OFFSET = 5 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t SEL_MASK_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET; + + constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t FULL16_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE; + scaleValue = scaleValue_; + lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); + lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + maskUbTensorUint8 = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + maskUbTensor16 = resource.ubBuf.template GetBufferByByte(MASK16_UB_TENSOR_OFFSET); + maskUbTensor32 = resource.ubBuf.template GetBufferByByte(MASK32_UB_TENSOR_OFFSET); + fullUbTensor16 = resource.ubBuf.template GetBufferByByte(FULL16_UB_TENSOR_OFFSET); + fullUbTensor32 = resource.ubBuf.template GetBufferByByte(FULL32_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + selMaskUbTensor = resource.ubBuf.template GetBufferByByte(SEL_MASK_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tempMaskTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_PREMASK_TENSOR_OFFSET); + } + + __aicore__ inline + ~BlockEpilogue() {} + + template + __aicore__ inline T Min(T a, T b) + { + return (a > b) ? b : a; + } + + __aicore__ inline + void SetVecMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = len % FLOAT_VECTOR_SIZE; + for (int64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE || len == 0) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + + __aicore__ inline + void SetBlockReduceMask(int32_t len) + { + if (len > 8 || len < 1) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + uint64_t subMask = ((uint64_t)1 << len) - 1; + uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) + + (subMask << 40) + (subMask << 24) + (subMask << 8); + AscendC::SetVectorMask(maskValue, maskValue); + } + + __aicore__ inline + void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum( + tvUbTensor, + srcUb, + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + rowsumUb, + tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void RowsumSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceSum( + tvUbTensor, + srcUb, + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + numRowsRound, + 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceSum( + rowsumUb, + tvUbTensor[REDUCE_UB_SIZE], + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + __aicore__ inline + void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum( + tvUbTensor, + srcUb, + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + rowsumUb, + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) { + AscendC::BlockReduceSum( + tvUbTensor, + srcUb[rowSumIdx * FLOAT_VECTOR_SIZE], + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add( + rowsumUb, + rowsumUb, + tvUbTensor[REDUCE_UB_SIZE], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceSum( + tvUbTensor, + srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask(NpuArch::Detail::Alignment::CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE)); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceSum( + rowsumUb, + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceSum( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Add( + rowsumUb, + rowsumUb, + tvUbTensor[REDUCE_UB_SIZE], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + __aicore__ inline + void RowmaxSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax( + tvUbTensor, + srcUb, + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + rowmaxUb, + tvUbTensor[REDUCE_UB_SIZE], + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void RowmaxSPECTILE256(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::BlockReduceMax( + tvUbTensor, + srcUb, + numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE, + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(ROW_OPS_SPEC_MASK_32); + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + numRowsRound, + 0, 1, 1, 4); + AscendC::PipeBarrier(); + SetBlockReduceMask(ROW_OPS_SPEC_MASK_4); + AscendC::BlockReduceMax( + rowmaxUb, + tvUbTensor[REDUCE_UB_SIZE], + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + __aicore__ inline + void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems >= FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax( + tvUbTensor, + srcUb, + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + rowmaxUb, + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) { + AscendC::BlockReduceMax( + tvUbTensor, + srcUb[rowmax_idx * FLOAT_VECTOR_SIZE], + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + AscendC::BlockReduceMax( + tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, + rowmaxUb, + tvUbTensor[REDUCE_UB_SIZE], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + if (numElems % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(numElems % FLOAT_VECTOR_SIZE); + AscendC::BlockReduceMax( + tvUbTensor, + srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + numRowsRound, + 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE); + AscendC::PipeBarrier(); + SetBlockReduceMask(NpuArch::Detail::Alignment::CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE)); + if (numElems < FLOAT_VECTOR_SIZE) { + AscendC::BlockReduceMax(rowmaxUb, + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + } else { + AscendC::BlockReduceMax(tvUbTensor[REDUCE_UB_SIZE], + tvUbTensor, + NpuArch::Detail::Alignment::CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE), + 0, 1, 1, 8); + AscendC::PipeBarrier(); + SetVecMask(numRowsRound); + AscendC::Max(rowmaxUb, + rowmaxUb, + tvUbTensor[REDUCE_UB_SIZE], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + __aicore__ inline + void CopySGmToUb( + AscendC::GlobalTensor gInput, + uint32_t sUbOffset, + uint32_t rowNumCurLoop, + uint32_t columnNumRound, + uint32_t columnNumPad) + { + AscendC::DataCopy( + lsUbTensor[sUbOffset], + gInput, + AscendC::DataCopyParams( + rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE, + (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE, 0)); + } + + __aicore__ inline void OperatePreMaskUb(uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + UpCastMask( + maskUbTensor16, + maskUbTensor, + rowNumCurLoop, + columnNumRound + ); + AscendC::CompareScalar( + maskUbTensorUint8, + maskUbTensor16, + static_cast(1.0), + AscendC::CMPMODE::NE, + REPEAT_SIZE_IN_BYTE / sizeof(half), + (rowNumCurLoop * columnNumRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8) + ); + AscendC::PipeBarrier(); + AscendC::Duplicate(tempMaskTensor, static_cast(1), rowNumCurLoop * columnNumRound); + AscendC::PipeBarrier(); + AscendC::Select(maskUbTensor16, maskUbTensorUint8, tempMaskTensor, static_cast(0), AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, rowNumCurLoop * columnNumRound); + AscendC::PipeBarrier(); + UpCastMask(maskUbTensor32, maskUbTensor16, rowNumCurLoop, columnNumRound); + } + + __aicore__ inline void OperateNextMaskUb(uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + UpCastMask( + maskUbTensor16, + maskUbTensor[MAX_UB_S_ELEM_NUM], + rowNumCurLoop, + columnNumRound + ); + UpCastMask( + maskUbTensor32, + maskUbTensor16, + rowNumCurLoop, + columnNumRound + ); + } + + + __aicore__ inline + void CopyMaskGmToUb( + AscendC::GlobalTensor gMask, + uint32_t columnNum, uint32_t columnNumRound, + uint32_t maskStride, uint32_t tokenNumPerHead, + uint32_t proTokenIdx, uint32_t proTokenNum, + uint32_t integralHeadNum, uint32_t epiTokenNum, bool isNextMask) + { + uint32_t innerUbRowOffset = isNextMask ? MAX_UB_S_ELEM_NUM : 0; + if (proTokenNum != 0) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride], + AscendC::DataCopyExtParams( + proTokenNum, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += proTokenNum * columnNumRound; + } + for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask, + AscendC::DataCopyExtParams( + tokenNumPerHead, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += tokenNumPerHead * columnNumRound; + } + if (epiTokenNum != 0) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], gMask, + AscendC::DataCopyExtParams( + epiTokenNum, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + } + } + + __aicore__ inline + void CalcGmFullShift(int64_t &offsetFull, const LayoutInput &layoutMask, + uint32_t rowOffer, uint32_t kvSStartIdx, uint32_t maskOffsetThisSubBlock) + { + uint32_t fullBlockStart = rowOffer; + uint32_t gmOffsetFullRow = fullBlockStart + maskOffsetThisSubBlock ; + uint32_t gmOffsetFullColumn = kvSStartIdx; + offsetFull = layoutMask.GetOffset(MatrixCoord(gmOffsetFullRow, gmOffsetFullColumn)); + } + + __aicore__ inline + void CopyFullGmToUb( + AscendC::GlobalTensor gFull, + uint32_t columnNum, uint32_t columnNumRound, uint32_t maskStride, uint32_t tokenNumPerHead, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t integralHeadNum, uint32_t epiTokenNum, + uint32_t &qNStartIdxVec, uint32_t qNThisSubBlock, uint32_t rowNumCurLoop, uint32_t BIdx, + uint32_t qHeads, int64_t offsetFull, int64_t pseQ, int64_t pseKv) + { + uint32_t innerUbRowOffset = 0; + int64_t gFullOffset = 0; + if (proTokenNum != 0) { + gFullOffset = BIdx * qHeads * pseQ * pseKv + qNStartIdxVec * pseQ * pseKv + offsetFull; + AscendC::DataCopyPad( + fullUbTensor16[innerUbRowOffset], gFull[gFullOffset + proTokenIdx * maskStride], + AscendC::DataCopyExtParams( + proTokenNum, columnNum * sizeof(ElementFull), + (maskStride - columnNum) * sizeof(ElementFull), + (columnNumRound - columnNum) * sizeof(ElementFull) / BLOCK_SIZE_IN_BYTE, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + UpCastMask(fullUbTensor32[innerUbRowOffset], fullUbTensor16[innerUbRowOffset], + rowNumCurLoop, columnNumRound); + innerUbRowOffset += proTokenNum * columnNumRound; + } + for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) { + if (qNThisSubBlock >= HEAD_NUM_2) { + if (proTokenNum > 0 && headIdx == 0) { + qNStartIdxVec++; + } + if (headIdx > 0) { + qNStartIdxVec++; + } + } + gFullOffset = BIdx * qHeads * pseQ * pseKv + qNStartIdxVec * pseQ * pseKv + offsetFull; + AscendC::DataCopyPad( + fullUbTensor16[innerUbRowOffset], gFull[gFullOffset], + AscendC::DataCopyExtParams( + tokenNumPerHead, columnNum * sizeof(ElementFull), + (maskStride - columnNum) * sizeof(ElementFull), + (columnNumRound - columnNum) * sizeof(ElementFull) / BLOCK_SIZE_IN_BYTE, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + UpCastMask(fullUbTensor32[innerUbRowOffset], fullUbTensor16[innerUbRowOffset], + rowNumCurLoop, columnNumRound); + innerUbRowOffset += tokenNumPerHead * columnNumRound; + } + if (epiTokenNum != 0) { + if (qNThisSubBlock >= HEAD_NUM_2) { + qNStartIdxVec++; + } + gFullOffset = BIdx * qHeads * pseQ * pseKv + qNStartIdxVec * pseQ * pseKv + offsetFull; + AscendC::DataCopyPad( + fullUbTensor16[innerUbRowOffset], gFull[gFullOffset], + AscendC::DataCopyExtParams( + epiTokenNum, columnNum * sizeof(ElementFull), + (maskStride - columnNum) * sizeof(ElementFull), + (columnNumRound - columnNum) * sizeof(ElementFull) / BLOCK_SIZE_IN_BYTE, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + UpCastMask(fullUbTensor32[innerUbRowOffset], fullUbTensor16[innerUbRowOffset], + rowNumCurLoop, columnNumRound); + } + } + + __aicore__ inline + void Applyfull(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + AscendC::Add( + lsUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + fullUbTensor32, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + AscendC::Muls( + lsUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + scaleValue, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + } + + template + __aicore__ inline + void UpCastMask( + const AscendC::LocalTensor &maskUbTensorDst, + const AscendC::LocalTensor &maskUbTensorSrc, + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + AscendC::Cast( + maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv( + rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))), + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void ApplyMask(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound, + uint32_t addMaskUbOffset) + { + AscendC::Muls( + maskUbTensor32, + maskUbTensor32, + (float)-3e38, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + if (maskColumnRound == columnNumRound) { + AscendC::Add( + lsUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + maskUbTensor32, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + } else { + uint32_t loop = maskColumnRound / FLOAT_VECTOR_SIZE; + for (uint32_t i = 0; i < loop; i++) { + AscendC::Add(lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE], + maskUbTensor32[i * FLOAT_VECTOR_SIZE], + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 1, + columnNumRound / FLOAT_BLOCK_SIZE, + columnNumRound / FLOAT_BLOCK_SIZE, + maskColumnRound / FLOAT_BLOCK_SIZE)); + } + if (maskColumnRound % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(maskColumnRound % FLOAT_VECTOR_SIZE); + AscendC::Add(lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE], + maskUbTensor32[loop * FLOAT_VECTOR_SIZE], + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 1, + columnNumRound / FLOAT_BLOCK_SIZE, + columnNumRound / FLOAT_BLOCK_SIZE, + maskColumnRound / FLOAT_BLOCK_SIZE)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcLocalRowMax(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound, + uint32_t rowOffset) + { + if (columnNum == 512) { + RowmaxSPECTILE512( + lsUbTensor[sUbOffset], + lmUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } else if (columnNum == 256) { + RowmaxSPECTILE256( + lsUbTensor[sUbOffset], + lmUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } else { + RowmaxTAILTILE( + lsUbTensor[sUbOffset], + lmUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } + } + + __aicore__ inline + void UpdateGlobalRowMax(uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, + uint32_t columnNumRound, uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + AscendC::DataCopy( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** hm = vmax(lm, gm) + AscendC::Max( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + gmUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** dm = gm - hm + AscendC::Sub( + dmUbTensor[dmUbOffsetCurCycle], + gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** dm = exp(dm) + AscendC::Exp( + dmUbTensor[dmUbOffsetCurCycle], + dmUbTensor[dmUbOffsetCurCycle], + (uint64_t)0, + 1, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + // *** gm = hm + AscendC::DataCopy( + gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void UpdateGlobalRowMax(AscendC::GlobalTensor gSink, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, + uint32_t columnNumRound, uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile, bool isLastStackTile, SinkLoopParam &curLoop) + { + if (isFirstStackTile) { + AscendC::DataCopy( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + // hm = Maxs(hm, sink) + if constexpr (SINK_MODE == SinkMode::ENABLE){ + if (isLastStackTile) { + UpdateRowMaxWithSink(gSink, rowOffset, dmUbOffsetCurCycle, curLoop); + } + } + } else { + SetVecMask(rowNumCurLoop); + // *** hm = vmax(lm, gm) + AscendC::Max( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + gmUbTensor[rowOffset], + (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // hm = Maxs(hm, sink) + if constexpr (SINK_MODE == SinkMode::ENABLE){ + if (isLastStackTile) { + UpdateRowMaxWithSink(gSink, rowOffset, dmUbOffsetCurCycle, curLoop); + SetVecMask(rowNumCurLoop); + } + } + + // *** dm = gm - hm + AscendC::Sub( + dmUbTensor[dmUbOffsetCurCycle], + gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** dm = exp(dm) + AscendC::Exp( + dmUbTensor[dmUbOffsetCurCycle], + dmUbTensor[dmUbOffsetCurCycle], + (uint64_t)0, 1, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + + // *** gm = hm + AscendC::DataCopy( + gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcExp(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, + uint32_t columnNumRound, uint32_t rowOffset) + { + // *** hm_block = expand_to_block(hm), 存放于 tv + AscendC::Brcb( + tvUbTensor.template ReinterpretCast(), + hmUbTensor[rowOffset].template ReinterpretCast(), + rowNumCurLoopRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** ls = ls - hm_block + for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) { + AscendC::Sub( + lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); + } + if (columnNum % FLOAT_VECTOR_SIZE > 0) { + SetVecMask(columnNum % FLOAT_VECTOR_SIZE); + AscendC::Sub( + lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** ls = exp(ls) + AscendC::Exp( + lsUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcLocalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound, + uint32_t rowOffset) + { + // *** ll = rowsum(ls32) + if (columnNum == 512) { + RowsumSPECTILE512( + lsUbTensor[sUbOffset], + llUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } else if (columnNum == 256) { + RowsumSPECTILE256( + lsUbTensor[sUbOffset], + llUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } else { + RowsumTAILTILE( + lsUbTensor[sUbOffset], + llUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } + } + + __aicore__ inline + void UpdateGlobalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, + uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + // *** gl = ll + AscendC::DataCopy( + glUbTensor[rowOffset], + llUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** gl = dm * gl + AscendC::Mul( + glUbTensor[rowOffset], + dmUbTensor[dmUbOffsetCurCycle], + glUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** gl = ll + gl + AscendC::Add( + glUbTensor[rowOffset], + glUbTensor[rowOffset], + llUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + __aicore__ inline + void UpdateGlobalRowSum(AscendC::GlobalTensor gSink, uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, + uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile, bool isLastStackTile, SinkLoopParam &curLoop) + { + if (isFirstStackTile) { + // *** gl = ll + AscendC::DataCopy( + glUbTensor[rowOffset], + llUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + + } else { + SetVecMask(rowNumCurLoop); + // *** gl = dm * gl + AscendC::Mul( + glUbTensor[rowOffset], + dmUbTensor[dmUbOffsetCurCycle], + glUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // *** gl = ll + gl + AscendC::Add( + glUbTensor[rowOffset], + glUbTensor[rowOffset], + llUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + // gl = gl + exp(sink-lm) + if constexpr (SINK_MODE == SinkMode::ENABLE) { + if (isLastStackTile) { + UpdateRowSumWithSink(rowOffset, curLoop.rowNumCurLoop); + } + } + } + + __aicore__ inline + void DownCastP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + // *** lp = castfp32to16(ls) + if (std::is_same::value) { + AscendC::Cast( + lpUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + AscendC::RoundMode::CAST_RINT, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + lpUbTensor[sUbOffset], + lsUbTensor[sUbOffset], + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + } + + __aicore__ inline + void CopyPUbToGm(AscendC::GlobalTensor gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop, + uint32_t columnNumRound, uint32_t columnNumPad) + { + AscendC::DataCopy( + gOutput, + lpUbTensor[sUbOffset], + AscendC::DataCopyParams( + rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE)); + } + + template + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, const LayoutOutput &layoutOutput, + uint32_t rowOffset, uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, + uint32_t isFirstRowLoop, uint32_t isLastRowLoop, + uint32_t columnNumRound, uint32_t pingpongFlag, + uint32_t curStackTileMod) + { + uint32_t rowNumCurLoop = layoutOutput.shape(0); + uint32_t rowNumCurLoopRound = NpuArch::Detail::Alignment::RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE); + uint32_t columnNum = layoutOutput.shape(1); + uint32_t columnNumPad = layoutOutput.stride(0); + uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM; + uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset; + + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + // In lse out-only mode, tv is used in the last stack tile to transport lse + if (isFirstStackTile && isFirstRowLoop) { + AscendC::WaitFlag(EVENT_ID4); + } + } + CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + UpdateGlobalRowMax( + rowNumCurLoop, rowNumCurLoopRound, + columnNum, columnNumRound, + dmUbOffsetCurCycle, + rowOffset, + isFirstStackTile); + + CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + if constexpr (!doTriUMask) { + AscendC::WaitFlag(pingpongFlag); + } + + DownCastP(sUbOffset, rowNumCurLoop, columnNumRound); + AscendC::SetFlag(pingpongFlag); + + CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + AscendC::SetFlag(pingpongFlag); + + AscendC::WaitFlag(pingpongFlag); + CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad); + if constexpr (!doTriUMask) { + AscendC::SetFlag(pingpongFlag); + if (isLastNoMaskStackTile && isLastRowLoop) { + AscendC::WaitFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + } + } else { + AscendC::SetFlag(EVENT_ID0); + } + UpdateGlobalRowSum( + sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); + } + + template + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gSink, const LayoutOutput &layoutOutput, + uint32_t rowOffset, uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, + uint32_t isFirstRowLoop, uint32_t isLastRowLoop, + uint32_t columnNumRound, uint32_t pingpongFlag, + uint32_t curStackTileMod, SinkLoopParam& sinkLoopParam, bool isLastStackTile, bool isSplitKV, bool startsWithMaskThenNomaskFlag) + { + uint32_t rowNumCurLoop = layoutOutput.shape(0); + uint32_t rowNumCurLoopRound = NpuArch::Detail::Alignment::RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE); + uint32_t columnNum = layoutOutput.shape(1); + uint32_t columnNumPad = layoutOutput.stride(0); + uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM; + uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset; + + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + // In lse out-only mode, tv is used in the last stack tile to transport lse + if (isFirstStackTile && isFirstRowLoop) { + AscendC::WaitFlag(EVENT_ID4); + } + } else { + if (isFirstStackTile && isFirstRowLoop && isSplitKV) { + AscendC::WaitFlag(EVENT_ID4); + } + } + CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + UpdateGlobalRowMax( + gSink, + rowNumCurLoop, rowNumCurLoopRound, + columnNum, columnNumRound, + dmUbOffsetCurCycle, + rowOffset, + isFirstStackTile, + isLastStackTile, + sinkLoopParam); + + CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + if constexpr (!doTriUMask) { + AscendC::WaitFlag(pingpongFlag); + } + + DownCastP(sUbOffset, rowNumCurLoop, columnNumRound); + AscendC::SetFlag(pingpongFlag); + + CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + AscendC::SetFlag(pingpongFlag); + + AscendC::WaitFlag(pingpongFlag); + CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad); + if constexpr (!doTriUMask) { + AscendC::SetFlag(pingpongFlag); + if (isLastNoMaskStackTile && isLastRowLoop) { + if(!startsWithMaskThenNomaskFlag) { + AscendC::WaitFlag(EVENT_ID0); + } + AscendC::SetFlag(EVENT_ID0); + } + } else { + AscendC::SetFlag(EVENT_ID0); + } + UpdateGlobalRowSum( + gSink, sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile, isLastStackTile, sinkLoopParam); + } + + __aicore__ inline + float ConvertElementSinkToFloat(const ElementSink& rawSinkVal) { + if constexpr (std::is_same_v) { + return AscendC::ToFloat(rawSinkVal); + } else if constexpr (std::is_same_v) { + return static_cast(rawSinkVal); + } + } + + __aicore__ inline + void SetSinkVecMask(uint32_t elemStart, uint32_t elemEnd) { + uint64_t mask = 0; + uint64_t one = 1; + if (elemStart < 0 || elemStart > elemEnd || elemEnd > FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)0, (uint64_t)0); + return; + } + + for (uint32_t elemIdx = elemStart; elemIdx <= elemEnd; elemIdx++) { + mask |= (one << elemIdx); + } + AscendC::SetVectorMask(0x0, mask); + } + + __aicore__ inline + void UpdateRowMaxWithSink(AscendC::GlobalTensor gSink, uint32_t rowOffset, uint32_t dmUbOffsetCurCycle, SinkLoopParam &curLoop) + { + const uint32_t loopStart = curLoop.rowOffsetIoGm; + const uint32_t loopEnd = curLoop.rowOffsetIoGm + curLoop.rowNumCurLoop - 1; + const uint32_t qSBlockSize = curLoop.qSBlockSize; + + const uint32_t firstHeadId = loopStart / qSBlockSize; + const uint32_t lastHeadId = loopEnd / qSBlockSize; + + SetVecMask(curLoop.rowNumCurLoop); + float zeroNum = 0.0f; + AscendC::Duplicate(lmUbTensor[rowOffset], zeroNum, (uint64_t)0, 1, 1, 8); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + + for (uint32_t headId = firstHeadId; headId <= lastHeadId; headId++) { + uint32_t curHeadQsBlockStartGm = headId * qSBlockSize; + uint32_t curHeadQsBlockEndGm = curHeadQsBlockStartGm + qSBlockSize - 1U; + + uint32_t headActualStartThisSubCore = AscendC::Std::max(curHeadQsBlockStartGm, loopStart) - curLoop.rowOffsetThisSubBlock - rowOffset; + uint32_t headActualEndThisSubCore = AscendC::Std::min(curHeadQsBlockEndGm, loopEnd) - curLoop.rowOffsetThisSubBlock - rowOffset; + + float sinkValue = ConvertElementSinkToFloat(gSink.GetValue(headId)); + uint32_t headRowNumThisSubCore = headActualEndThisSubCore - headActualStartThisSubCore + 1; + + SetSinkVecMask(headActualStartThisSubCore, headActualEndThisSubCore); + + AscendC::UnaryRepeatParams maxsRepeatParams(1, 1, 8, 8); + + // hm = Maxs(hm, sink) + AscendC::Maxs( + dmUbTensor[dmUbOffsetCurCycle], + hmUbTensor[rowOffset], + sinkValue, + (uint64_t)0, 1, + maxsRepeatParams + ); + + // sinkTensor + AscendC::Adds( + lmUbTensor[rowOffset], + lmUbTensor[rowOffset], + sinkValue, + (uint64_t)0, 1, + maxsRepeatParams + ); + + } + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + uint64_t mask = static_cast(curLoop.rowNumCurLoop); + + AscendC::CompareScalar(selMaskUbTensor, hmUbTensor[rowOffset], NEG_INF, AscendC::CMPMODE::EQ, + mask, 1, AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + + AscendC::Select(hmUbTensor[rowOffset], selMaskUbTensor, hmUbTensor[rowOffset], dmUbTensor[dmUbOffsetCurCycle], AscendC::SELMODE::VSEL_CMPMASK_SPR, + mask, 1, AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + } + + __aicore__ inline + void UpdateRowSumWithSink(uint32_t rowOffset, uint32_t rowNumCurLoop) + { + SetVecMask(rowNumCurLoop); + // sink = sink - hm + AscendC::Sub( + lmUbTensor[rowOffset], + lmUbTensor[rowOffset], + hmUbTensor[rowOffset], + (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // exp(sink-hm) + AscendC::Exp( + lmUbTensor[rowOffset], + lmUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + + // gl+exp(sink -m) + AscendC::Add( + glUbTensor[rowOffset], + glUbTensor[rowOffset], + lmUbTensor[rowOffset], + (uint64_t)0, 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutputBase, + AscendC::GlobalTensor gInputBase, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + GemmCoord actualBlockShape, + uint32_t isFirstStackTile, + uint32_t isLastNoMaskStackTile, + uint32_t qSBlockSize, + uint32_t qNBlockSize, + uint32_t curStackTileMod, + uint32_t kvNBlockSize, + uint64_t gmOffsetSBase, + uint64_t gmOffsetPBase, + Arch::CrossCoreFlag qkReady, + Arch::CrossCoreFlag softmaxReady) + { + Arch::CrossCoreWaitFlag(qkReady); + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t kvNSplitSubBlock = kvNBlockSize / subBlockNum; + uint32_t kvNThisSubBlock = (kvNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (kvNBlockSize - kvNSplitSubBlock) : kvNSplitSubBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) + : qNSplitSubBlock) + : (kvNThisSubBlock * qNBlockSize); + + uint32_t rowSplitSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock)) : + (qSBlockSize * qNBlockSize * kvNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInputBase[offsetInput]; + + AscendC::WaitFlag(pingpongFlag); + + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = + (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutputBase[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + AscendC::WaitFlag(pingpongFlag); + + // add sink + + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + SubCoreCompute( + gOutputCurLoop, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + isLastNoMaskStackTile, + (delayedRowLoopIdx == 0), + (delayedRowLoopIdx == rowLoopNum - 1), + columnNumRound, + pingpongFlag, + curStackTileMod + ); + } + } + } + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, AscendC::GlobalTensor gSink, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, + uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, uint32_t qSBlockSize, uint32_t qNBlockSize, + uint32_t curStackTileMod, bool isLastStackTile, bool isSplitKV = false, bool startsWithMaskTile = false, + bool startsWithMaskThenNomaskFlag = false) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;// 在整个块的起始偏移 + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + + AscendC::WaitFlag(pingpongFlag); + if (startsWithMaskTile && rowLoopIdx == 0) { + AscendC::WaitFlag(EVENT_ID0); + } + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = + (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + AscendC::WaitFlag(pingpongFlag); + + // add sink + SinkLoopParam curSinkLoop(rowOffsetIoGm, rowNumCurLoop, qSBlockSize, rowOffsetThisSubBlock); + + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + SubCoreCompute( + gOutputCurLoop, + gSink, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + isLastNoMaskStackTile, + (delayedRowLoopIdx == 0), + (delayedRowLoopIdx == rowLoopNum - 1), + columnNumRound, + pingpongFlag, + curStackTileMod, + curSinkLoop, + isLastStackTile, + isSplitKV, + startsWithMaskThenNomaskFlag); + } + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, AscendC::GlobalTensor gSink, + AscendC::GlobalTensor gMask, const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, + const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize, + uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown, + uint32_t kvSStartIdx, uint32_t kvSEndIdx, bool isLastStackTile, bool isSplitKV = false) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE_IN_BYTE); + uint32_t columnNumPad = layoutInput.stride(0); + uint32_t maskStride = layoutMask.stride(0); + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + + uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock); + uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? + rowOffsetThisSubBlock : 0; + + // calc mask shift in gm + uint32_t gmOffsetMaskRow; + uint32_t gmOffsetMaskColumn; + uint32_t maskColumn; + uint32_t addMaskUbOffset; + if (triUp >= kvSStartIdx) { + uint32_t triUpRoundDown = NpuArch::Detail::Alignment::RoundDown(triUp, BLOCK_SIZE_IN_BYTE); + gmOffsetMaskRow = triUp - triUpRoundDown; + gmOffsetMaskColumn = 0; + maskColumn = kvSEndIdx - triUpRoundDown; + addMaskUbOffset = triUpRoundDown - kvSStartIdx; + } else { + gmOffsetMaskRow = 0; + gmOffsetMaskColumn = kvSStartIdx - triUp; + maskColumn = columnNum; + addMaskUbOffset = 0; + } + uint32_t maskColumnRound = NpuArch::Detail::Alignment::RoundUp(maskColumn, BLOCK_SIZE_IN_BYTE); + + int64_t offsetMask = + layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn)); + auto gMaskThisSubBlock = gMask[offsetMask]; + auto layoutMaskThisSubBlock = layoutMask; + + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + if (rowActualThisSubBlock == 0) { + Arch::CrossCoreWaitFlag(qkReady); + return; + } + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // loop 0 mask load before cross core sync + if (rowLoopIdx == 0) { + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + CopyMaskGmToUb( + gMaskThisSubBlock, + maskColumn, maskColumnRound, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + AscendC::SetFlag(EVENT_ID2); + Arch::CrossCoreWaitFlag(qkReady); + } + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + AscendC::WaitFlag(pingpongFlag); + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + AscendC::WaitFlag(EVENT_ID2); + UpCastMask(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound); + UpCastMask(maskUbTensor32, maskUbTensor16, rowNumCurLoop, columnNumRound); + + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, columnNumRound, + maskColumnRound, addMaskUbOffset); + // next loop mask load + if (rowLoopIdx < rowLoopNum) { + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = + (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + CopyMaskGmToUb( + gMaskThisSubBlock, + maskColumn, maskColumnRound, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + AscendC::SetFlag(EVENT_ID2); + } + // online softmax vectorized compute + + // add sink + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + SinkLoopParam curSinkLoop(rowOffsetIoGm, rowNumCurLoop, qSBlockSize, rowOffsetThisSubBlock); + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + gSink, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + 0, + (delayedRowLoopIdx == 0), + (delayedRowLoopIdx == rowLoopNum - 1), + columnNumRound, + pingpongFlag, + curStackTileMod, + curSinkLoop, + isLastStackTile, + isSplitKV, + false); + } + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, AscendC::GlobalTensor gSink, + AscendC::GlobalTensor gMask, const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, + const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize, + uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, + int32_t kvSStartIdx, bool doTriUPreMask, bool doTriUNextMask, int32_t preTokenStartLen, + int32_t preTokenEndLen, int32_t nextTokenStartLen, int32_t nextTokenEndLen, bool isLastStackTile) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE_IN_BYTE); + uint32_t columnNumPad = layoutInput.stride(0); + uint32_t maskStride = layoutMask.stride(0); + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + + uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock); + uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? + rowOffsetThisSubBlock : 0; + + uint32_t addMaskUbOffset = 0; + + uint32_t gmOffsetMaskRowPre; + uint32_t gmOffsetMaskColumnPre; + uint32_t maskColumnPre; + + uint32_t gmOffsetMaskRowNext; + uint32_t gmOffsetMaskColumnNext; + uint32_t maskColumnNext; + if (doTriUPreMask) { + if (preTokenStartLen > kvSStartIdx) { + gmOffsetMaskRowPre = preTokenStartLen - kvSStartIdx - 1; + gmOffsetMaskColumnPre = 0; + maskColumnPre = columnNumRound; + } else { + gmOffsetMaskRowPre = 0; + gmOffsetMaskColumnPre = kvSStartIdx - preTokenStartLen + 1; + maskColumnPre = columnNumRound; + } + } + if (doTriUNextMask) { + if (nextTokenStartLen > kvSStartIdx) { + gmOffsetMaskRowNext = nextTokenStartLen - kvSStartIdx; + gmOffsetMaskColumnNext = 0; + maskColumnNext = columnNumRound; + } else { + gmOffsetMaskRowNext = 0; + gmOffsetMaskColumnNext = kvSStartIdx - nextTokenStartLen; + maskColumnNext = columnNumRound; + } + } + uint32_t columnNumRoundPre = NpuArch::Detail::Alignment::RoundUp(maskColumnPre, BLOCK_SIZE_IN_BYTE); + int64_t offsetMaskPre = + layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRowPre + maskOffsetThisSubBlock, gmOffsetMaskColumnPre)); + auto gMaskThisSubBlockPre = gMask[offsetMaskPre]; + + uint32_t columnNumRoundNext = NpuArch::Detail::Alignment::RoundUp(maskColumnNext, BLOCK_SIZE_IN_BYTE); + int64_t offsetMaskNext = + layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRowNext + maskOffsetThisSubBlock, gmOffsetMaskColumnNext)); + auto gMaskThisSubBlockNext = gMask[offsetMaskNext]; + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + if (rowActualThisSubBlock == 0) { + Arch::CrossCoreWaitFlag(qkReady); + return; + } + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // loop 0 mask load before cross core sync + if (rowLoopIdx == 0) { + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + if (doTriUPreMask && doTriUNextMask) { + CopyMaskGmToUb( + gMaskThisSubBlockPre, + maskColumnPre, columnNumRoundPre, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + } else if (doTriUPreMask){ + CopyMaskGmToUb( + gMaskThisSubBlockPre, + maskColumnPre, columnNumRoundPre, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + } else if (doTriUNextMask) { + CopyMaskGmToUb( + gMaskThisSubBlockNext, + maskColumnNext, columnNumRoundNext, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, true); + } + AscendC::SetFlag(EVENT_ID2); + Arch::CrossCoreWaitFlag(qkReady); + } + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + AscendC::WaitFlag(pingpongFlag); + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + AscendC::WaitFlag(EVENT_ID2); + if (doTriUPreMask && doTriUNextMask) { + OperatePreMaskUb(rowNumCurLoop, columnNumRound); + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, columnNumRound, + columnNumRoundPre, addMaskUbOffset); + AscendC::SetFlag(EVENT_ID6); + if (doTriUNextMask) { + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID6); + CopyMaskGmToUb( + gMaskThisSubBlockNext, + maskColumnNext, columnNumRoundNext, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, true); + AscendC::SetFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID6); + OperateNextMaskUb(rowNumCurLoop, columnNumRound); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, columnNumRound, + columnNumRoundNext, addMaskUbOffset); + } + } else if (doTriUPreMask) { + OperatePreMaskUb(rowNumCurLoop, columnNumRound); + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, columnNumRound, + columnNumRoundPre, addMaskUbOffset); + } else if (doTriUNextMask) { + OperateNextMaskUb(rowNumCurLoop, columnNumRound); + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, columnNumRound, + columnNumRoundNext, addMaskUbOffset); + } + // next loop mask load + if (rowLoopIdx < rowLoopNum) { + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = + (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // the token idx of the start token of the prologue part + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + uint32_t proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + if (doTriUPreMask && doTriUNextMask) { + CopyMaskGmToUb( + gMaskThisSubBlockPre, + maskColumnPre, columnNumRoundPre, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + } else if (doTriUPreMask){ + CopyMaskGmToUb( + gMaskThisSubBlockPre, + maskColumnPre, columnNumRoundPre, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, false); + } else if (doTriUNextMask) { + CopyMaskGmToUb( + gMaskThisSubBlockNext, + maskColumnNext, columnNumRoundNext, maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, true); + } + AscendC::SetFlag(EVENT_ID2); + } + // online softmax vectorized compute + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + // add sink + SinkLoopParam curSinkLoop(rowOffsetIoGm, rowNumCurLoop, qSBlockSize, rowOffsetThisSubBlock); + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + gSink, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + 0, + (delayedRowLoopIdx == 0), + (delayedRowLoopIdx == rowLoopNum - 1), + columnNumRound, + pingpongFlag, + curStackTileMod, + curSinkLoop, + isLastStackTile, + false, + false); + } + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gSink, AscendC::GlobalTensor gFull, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, + const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize, + uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t rowOffer, uint32_t kvSStartIdx, + uint32_t kvSEndIdx, uint32_t qNStartIdx, uint32_t BIdx, uint32_t qHeads, int64_t pseQ, int64_t pseKv, bool isLastStackTile) + { + uint32_t rowNum = actualBlockShape.m(); //当前fullmask块的总行数 + uint32_t columnNum = actualBlockShape.n(); //fullmask块的总列数 + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE_IN_BYTE); //列数向上对齐到32B + uint32_t columnNumPad = layoutInput.stride(0); //输入张量stride + uint32_t maskStride = layoutMask.stride(0); //掩码张量 + uint32_t subBlockIdx = ptoSubBlockIdx; // current AIV lane + uint32_t subBlockNum = ptoLanesPerBlock; + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; // 1 每核分到的头数 + uint32_t qNThisSubBlock = (qNBlockSize == 1) ? + 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; // 1 本核的头索引 + uint32_t rowSplitSubBlock = (qNBlockSize == 1) ? + (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock); //每核分到的行数 + uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? + (rowNum - rowSplitSubBlock) : rowSplitSubBlock; //本核实际要处理的行数 + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; //本核子块在fullmask的行偏移 + uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock); // + uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ? rowOffsetThisSubBlock : 0; // + // calc qNstartIdx per vec core + uint32_t qNStartIdxVec = (qNBlockSize == 1) ? qNStartIdx : (subBlockIdx == 1) ? (qNStartIdx + qNSplitSubBlock) : qNStartIdx; + // calc Full shift in gm + int64_t offsetFull = 0; + CalcGmFullShift(offsetFull, layoutMask, rowOffer, kvSStartIdx, maskOffsetThisSubBlock); + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + uint32_t preLoad = 1; + + if (rowActualThisSubBlock == 0) { + Arch::CrossCoreWaitFlag(qkReady); + return; + } + uint32_t proTokenIdx; + uint32_t proTokenNum; + uint32_t integralHeadNum; + uint32_t epiTokenNum; + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) { + if (rowLoopIdx < rowLoopNum) { + uint32_t pingpongFlag = rowLoopIdx % 2; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // loop 0 mask load before cross core sync + if (rowLoopIdx == 0) { + // the token idx of the start token of the prologue part + proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + // the token num of the epilogue part + integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + // the number of integral heads within a cycle + epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + CopyFullGmToUb( + gFull, + columnNum, columnNumRound, maskStride, + tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, + qNStartIdxVec, qNThisSubBlock, rowNumCurLoop, BIdx, qHeads, offsetFull, pseQ, pseKv); + Arch::CrossCoreWaitFlag(qkReady); + } + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + AscendC::WaitFlag(pingpongFlag); + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(pingpongFlag); + } + if (rowLoopIdx >= preLoad) { + uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad; + uint32_t pingpongFlag = delayedRowLoopIdx % 2; + uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ? + (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + AscendC::WaitFlag(pingpongFlag); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + Applyfull((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + // next loop mask load + // online softmax vectorized compute + + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + SinkLoopParam curSinkLoop(rowOffsetIoGm, rowNumCurLoop, qSBlockSize, rowOffsetThisSubBlock); + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + gSink, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + 0, + delayedRowLoopIdx == 0, + delayedRowLoopIdx == rowLoopNum - 1, + columnNumRound, + pingpongFlag, + curStackTileMod, + curSinkLoop, + isLastStackTile, + false, + false); + if (rowLoopIdx < rowLoopNum) { + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowNumCurLoop = + (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + // the token idx of the start token of the prologue part + proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + // the token num of the prologue part + proTokenNum = + Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock; + if ((qNThisSubBlock >= HEAD_NUM_2) && (proTokenIdx == 0) && proTokenNum > 0) { + qNStartIdxVec++; + } + // the number of integral heads within a cycle + integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + if ((qNThisSubBlock >= HEAD_NUM_2) && integralHeadNum > 0 && proTokenNum == 0) { + qNStartIdxVec++; + } + // the token num of the epilogue part + epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + AscendC::WaitFlag(EVENT_ID0); + + CopyFullGmToUb( + gFull, + columnNum, columnNumRound, maskStride, + tokenNumPerHeadThisSubBlock, proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum, + qNStartIdxVec, qNThisSubBlock, rowNumCurLoop, BIdx, qHeads, offsetFull, pseQ, pseKv); // 当前循环要搬的行数 + } + } + } + } + +private: + uint32_t ptoSubBlockIdx = 0; + uint32_t ptoLanesPerBlock = 1; + uint32_t kvheadIdx = 0; + float scaleValue; + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor lpUbTensor; + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor maskUbTensorUint8; + AscendC::LocalTensor maskUbTensor16; + AscendC::LocalTensor maskUbTensor32; + AscendC::LocalTensor fullUbTensor16; + AscendC::LocalTensor fullUbTensor32; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tempMaskTensor; + AscendC::LocalTensor selMaskUbTensor; +}; +} + +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp new file mode 100644 index 0000000000..9bd1d05e60 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp @@ -0,0 +1,868 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/cross_core_sync.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/epilogue/dispatch_policy.hpp" +#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/matrix_coord.hpp" +#include "utils/std/algorithm.h" + +namespace NpuArch::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class MaskType_, + class SinkType_, + class FullType_, + LseMode LSE_MODE_, + SinkMode SINK_MODE_, + MaskMode MASK_MODE_> +class BlockEpilogue< + EpilogueAtlasA2OnlineSoftmax, + OutputType_, + InputType_, + MaskType_, + SinkType_, + FullType_> +{ +public: + using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax; + using ArchTag = typename DispatchPolicy::ArchTag; + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementMask = typename MaskType_::Element; + using ElementSink = typename SinkType_::Element; + using ElementFull = typename FullType_::Element; + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutMask = typename MaskType_::Layout; + using LayoutFull = typename FullType_::Layout; + + static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE; + static constexpr SinkMode SINK_MODE = DispatchPolicy::SINK_MODE; + + static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32; + static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t MAX_UB_S_ELEM_NUM = 16384; + + static constexpr uint32_t REDUCE_UB_SIZE = 1024; + static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32; + static constexpr uint32_t ROW_OPS_SPEC_MASK_8 = 8; + static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4; + static constexpr uint32_t ROW_OPS_SPEC_MASK_2 = 2; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256; + static constexpr int64_t UB_FLOAT_LINE_SIZE = 64; + + static constexpr uint32_t SPLIT_COL_IDX_2 = 2; + static constexpr uint32_t SPLIT_COL_IDX_3 = 3; + __aicore__ inline + BlockEpilogue() {} + + __aicore__ inline + void init(Arch::Resource &resource, float scaleValue_) + { + // Allocate UB space + constexpr uint32_t LS_UB_TENSOR_OFFSET = 0; + constexpr uint32_t COMPUTE_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 0; + + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + + constexpr uint32_t MASK_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE; + + scaleValue = static_cast(scaleValue_); + lsUbTensor = resource.ubBuf.template GetBufferByByte(LS_UB_TENSOR_OFFSET); + computeUbTensor = resource.ubBuf.template GetBufferByByte(COMPUTE_UB_TENSOR_OFFSET); + lpUbTensor = resource.ubBuf.template GetBufferByByte(LP_UB_TENSOR_OFFSET); + maskUbTensor = resource.ubBuf.template GetBufferByByte(MASK_UB_TENSOR_OFFSET); + maskUbTensor16 = resource.ubBuf.template GetBufferByByte(MASK16_UB_TENSOR_OFFSET); + lmUbTensor = resource.ubBuf.template GetBufferByByte(LM_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + llUbTensor = resource.ubBuf.template GetBufferByByte(LL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + } + + __aicore__ inline + ~BlockEpilogue() {} + + __aicore__ inline + void SetVecMask(int32_t len) + { + const int32_t MAX_MASK_LEN = 128; + const int32_t HALF_MASK_LEN = 64; + if (len >= MAX_MASK_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; + int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; + if (len < HALF_MASK_LEN) { + AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); + } else { + AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); + } + } + + __aicore__ inline + void SetBlockReduceMask(int32_t len) + { + const int32_t MAX_LEN = 16; + if (len > MAX_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + uint64_t subMask = (static_cast(1) << len) - 1; + uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask; + AscendC::SetVectorMask(maskValue, maskValue); + } + + __aicore__ inline + void RowsumSPECTILE512(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + AscendC::Add( + srcUb, + srcUb, + srcUb[HALF_VECTOR_SIZE], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::Add( + srcUb[HALF_VECTOR_SIZE * SPLIT_COL_IDX_2], + srcUb[HALF_VECTOR_SIZE * SPLIT_COL_IDX_2], + srcUb[HALF_VECTOR_SIZE * SPLIT_COL_IDX_3], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + AscendC::Add( + srcUb, + srcUb, + srcUb[HALF_VECTOR_SIZE * SPLIT_COL_IDX_2], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + AscendC::WholeReduceSum( + rowsumUb, srcUb, (int32_t)0, numRowsRound, 1, 1, + numElemsAligned / BLOCK_SIZE); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void RowsumTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowsumUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems <= HALF_VECTOR_SIZE) { + SetVecMask(numElems); + AscendC::WholeReduceSum( + rowsumUb, srcUb, (int32_t)0, numRowsRound, 1, 1, + numElemsAligned / BLOCK_SIZE); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + for (uint32_t vmaxIdx = 1; vmaxIdx < numElems / HALF_VECTOR_SIZE; vmaxIdx++) { + AscendC::Add( + srcUb, + srcUb, + srcUb[vmaxIdx * HALF_VECTOR_SIZE], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + } + if (numElems % HALF_VECTOR_SIZE > 0) { + SetVecMask(numElems % HALF_VECTOR_SIZE); + AscendC::Add( + srcUb, + srcUb, + srcUb[numElems / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::WholeReduceSum( + rowsumUb, srcUb, (int32_t)0, numRowsRound, 1, 1, + numElemsAligned / BLOCK_SIZE); + } + AscendC::PipeBarrier(); + } + + __aicore__ inline + void RowmaxTAILTILE(const AscendC::LocalTensor &srcUb, const AscendC::LocalTensor &rowmaxUb, + const AscendC::LocalTensor &tvUbTensor, uint32_t numRowsRound, uint32_t numElems, + uint32_t numElemsAligned) + { + if (numElems <= HALF_VECTOR_SIZE) { + SetVecMask(numElems); + AscendC::WholeReduceMax( + rowmaxUb, srcUb, (int32_t)0, numRowsRound, 1, 1, + numElemsAligned / BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else { + AscendC::DataCopy( + lsUbTensor, + srcUb, + AscendC::DataCopyParams( + numRowsRound, + HALF_VECTOR_SIZE / BLOCK_SIZE, + (numElemsAligned - HALF_VECTOR_SIZE) / BLOCK_SIZE, + (numElemsAligned - HALF_VECTOR_SIZE) / BLOCK_SIZE)); + AscendC::PipeBarrier(); + for (uint32_t vmaxIdx = 1; vmaxIdx < numElems / HALF_VECTOR_SIZE; vmaxIdx++) { + AscendC::Max( + lsUbTensor, + lsUbTensor, + srcUb[vmaxIdx * HALF_VECTOR_SIZE], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + } + if (numElems % HALF_VECTOR_SIZE > 0) { + SetVecMask(numElems % HALF_VECTOR_SIZE); + AscendC::Max( + lsUbTensor, + lsUbTensor, + srcUb[numElems / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + (uint64_t)0, + numRowsRound, + AscendC::BinaryRepeatParams( + 1, 1, 1, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE, + numElemsAligned / BLOCK_SIZE)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::WholeReduceMax( + rowmaxUb, lsUbTensor, (int32_t)0, numRowsRound, 1, 1, + numElemsAligned / BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE); + } + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CopySGmToUb(AscendC::GlobalTensor gInput, uint32_t sUbOffset, uint32_t rowNumCurLoop, + uint32_t columnNumRound, uint32_t columnNumPad) + { + // input S + AscendC::DataCopy( + lsUbTensor, + gInput, + AscendC::DataCopyParams(rowNumCurLoop, + columnNumRound / BLOCK_SIZE, + (columnNumPad - columnNumRound) / BLOCK_SIZE, + 0)); + } + + __aicore__ inline + void CopyMaskGmToUb(AscendC::GlobalTensor gMask, uint32_t columnNum, uint32_t columnNumRound, + uint32_t maskStride, uint32_t tokenNumPerHead, uint32_t proTokenIdx, uint32_t proTokenNum, + uint32_t integralHeadNum, uint32_t epiTokenNum) + { + uint32_t innerUbRowOffset = 0; + if (proTokenNum != 0U) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], + gMask[proTokenIdx * maskStride], + AscendC::DataCopyExtParams( + proTokenNum, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += proTokenNum * columnNumRound; + } + for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], + gMask, + AscendC::DataCopyExtParams( + tokenNumPerHead, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + innerUbRowOffset += tokenNumPerHead * columnNumRound; + } + if (epiTokenNum != 0) { + AscendC::DataCopyPad( + maskUbTensor[innerUbRowOffset], + gMask, + AscendC::DataCopyExtParams( + epiTokenNum, columnNum * sizeof(ElementMask), + (maskStride - columnNum) * sizeof(ElementMask), 0, 0), + AscendC::DataCopyPadExtParams(false, 0, 0, 0)); + } + } + + __aicore__ inline + void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + // *** ls = scaleValue * ls + AscendC::Muls( + computeUbTensor, + lsUbTensor, + scaleValue, + (uint64_t)0, + (rowNumCurLoop * columnNumRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + } + + template + __aicore__ inline + void UpCastMask( + const AscendC::LocalTensor &maskUbTensorDst, + const AscendC::LocalTensor &maskUbTensorSrc, + uint32_t rowNumCurLoop, + uint32_t columnNumRound) + { + AscendC::Cast( + maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv( + rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))), + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void ApplyMask(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound, + uint32_t addMaskUbOffset) + { + AscendC::Muls( + maskUbTensor16, + maskUbTensor16, + (half)-6e4, // -65504 + (uint64_t)0, + (rowNumCurLoop * maskColumnRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + if (maskColumnRound == columnNumRound) { + AscendC::Add( + computeUbTensor, + computeUbTensor, + maskUbTensor16, + (uint64_t)0, + (rowNumCurLoop * maskColumnRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + } else { + uint32_t loop = maskColumnRound / HALF_VECTOR_SIZE; + for (uint32_t i = 0; i < loop; i++) { + AscendC::Add( + computeUbTensor[addMaskUbOffset + i * HALF_VECTOR_SIZE], + computeUbTensor[addMaskUbOffset + i * HALF_VECTOR_SIZE], + maskUbTensor16[i * HALF_VECTOR_SIZE], + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams(1, + 1, + 1, + columnNumRound / BLOCK_SIZE, + columnNumRound / BLOCK_SIZE, + maskColumnRound / BLOCK_SIZE)); + } + if (maskColumnRound % HALF_VECTOR_SIZE > 0) { + SetVecMask(maskColumnRound % HALF_VECTOR_SIZE); + AscendC::Add( + computeUbTensor[addMaskUbOffset + loop * HALF_VECTOR_SIZE], + computeUbTensor[addMaskUbOffset + loop * HALF_VECTOR_SIZE], + maskUbTensor16[loop * HALF_VECTOR_SIZE], + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams(1, + 1, + 1, + columnNumRound / BLOCK_SIZE, + columnNumRound / BLOCK_SIZE, + maskColumnRound / BLOCK_SIZE)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcLocalRowMax(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound, + uint32_t rowOffset) + { + RowmaxTAILTILE( + computeUbTensor, + lmUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } + + __aicore__ inline + void UpdateGlobalRowMax(uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, + uint32_t columnNumRound, uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + AscendC::DataCopy( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** hm = vmax(lm, gm) + AscendC::Max( + hmUbTensor[rowOffset], + lmUbTensor[rowOffset], + gmUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + + AscendC::PipeBarrier(); + // *** dm = gm - hm + AscendC::Sub( + dmUbTensor[dmUbOffsetCurCycle], + gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** dm = exp(dm) + AscendC::Exp(dmUbTensor[dmUbOffsetCurCycle], + dmUbTensor[dmUbOffsetCurCycle], + (uint64_t)0, + 1, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + } + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::PipeBarrier(); + // *** gm = hm + AscendC::DataCopy(gmUbTensor[rowOffset], + hmUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcExp(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum, + uint32_t columnNumRound, uint32_t rowOffset) + { + // *** hm_block = expand_to_block(hm), 存放于 tv + AscendC::Brcb( + tvUbTensor.template ReinterpretCast(), + hmUbTensor[rowOffset].template ReinterpretCast(), + rowNumCurLoopRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** ls = ls - hm_block + for (uint32_t subIdx = 0; subIdx < columnNum / HALF_VECTOR_SIZE; ++subIdx) { + AscendC::Sub( + computeUbTensor[subIdx * HALF_VECTOR_SIZE], + computeUbTensor[subIdx * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 0, columnNumRound / BLOCK_SIZE, columnNumRound / BLOCK_SIZE, 1)); + } + if (columnNum % HALF_VECTOR_SIZE > 0) { + SetVecMask(columnNum % HALF_VECTOR_SIZE); + AscendC::Sub( + computeUbTensor[columnNum / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + computeUbTensor[columnNum / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + rowNumCurLoop, + AscendC::BinaryRepeatParams( + 1, 1, 0, columnNumRound / BLOCK_SIZE, columnNumRound / BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** ls = exp(ls) + AscendC::Exp( + computeUbTensor, + computeUbTensor, + (uint64_t)0, + (rowNumCurLoop * columnNumRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CalcLocalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound, + uint32_t rowOffset) + { + // *** ll = rowsum(ls32) + if (columnNum == 512U) { + RowsumSPECTILE512(computeUbTensor, + llUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } else { + RowsumTAILTILE(computeUbTensor, + llUbTensor[rowOffset], + tvUbTensor, + rowNumCurLoopRound, + columnNum, + columnNumRound); + } + } + + __aicore__ inline + void UpdateGlobalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, + uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile) + { + if (isFirstStackTile) { + // *** gl = ll + AscendC::DataCopy( + glUbTensor[rowOffset], + llUbTensor[rowOffset], + AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0)); + AscendC::PipeBarrier(); + } else { + SetVecMask(rowNumCurLoop); + // *** gl = dm * gl + AscendC::Mul( + glUbTensor[rowOffset], + dmUbTensor[dmUbOffsetCurCycle], + glUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + // *** gl = ll + gl + AscendC::Add( + glUbTensor[rowOffset], + glUbTensor[rowOffset], + llUbTensor[rowOffset], + (uint64_t)0, + 1, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + } + + __aicore__ inline + void MoveP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound) + { + AscendC::DataCopyParams repeatParams; + repeatParams.blockCount = 1; + repeatParams.srcStride = 0; + repeatParams.blockLen = NpuArch::Detail::Alignment::CeilDiv(rowNumCurLoop * columnNumRound, BLOCK_SIZE); + AscendC::DataCopy(lpUbTensor, computeUbTensor, repeatParams); + AscendC::PipeBarrier(); + } + + __aicore__ inline + void CopyPUbToGm(AscendC::GlobalTensor gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop, + uint32_t columnNumRound, uint32_t columnNumPad) + { + AscendC::DataCopy(gOutput, + lpUbTensor, + AscendC::DataCopyParams( + rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE)); + } + + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, const LayoutOutput &layoutOutput, + uint32_t rowOffset, uint32_t isFirstStackTile, uint32_t isFirstRowLoop, + uint32_t columnNumRound, uint32_t pingpongFlag, + uint32_t curStackTileMod, bool isSplitKV) + { + uint32_t rowNumCurLoop = layoutOutput.shape(0); + uint32_t rowNumCurLoopRound = NpuArch::Detail::Alignment::RoundUp(rowNumCurLoop, BLOCK_SIZE); + uint32_t columnNum = layoutOutput.shape(1); + uint32_t columnNumPad = layoutOutput.stride(0); + uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM; + uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset; + + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + // In lse out-only mode, tv is used in the last stack tile to transport lse + if (isFirstStackTile && isFirstRowLoop) { + AscendC::WaitFlag(EVENT_ID4); + } + } else { + if (isFirstStackTile && isFirstRowLoop && isSplitKV) { + AscendC::WaitFlag(EVENT_ID4); + } + } + CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + AscendC::SetFlag(EVENT_ID0); + UpdateGlobalRowMax(rowNumCurLoop, + rowNumCurLoopRound, + columnNum, + columnNumRound, + dmUbOffsetCurCycle, + rowOffset, + isFirstStackTile); + CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + + AscendC::WaitFlag(EVENT_ID0); + MoveP(sUbOffset, rowNumCurLoop, columnNumRound); + AscendC::SetFlag(EVENT_ID0); + + CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset); + + AscendC::WaitFlag(EVENT_ID0); + CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(EVENT_ID0); + UpdateGlobalRowSum( + sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile); + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, AscendC::GlobalTensor gSink, + const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape, + uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile, + uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod, bool isLastStackTile, bool isSplitKV = false, + bool startsWithMaskTile = false, bool startsWithMaskThenNomaskFlag = false) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? + 0 : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1U) ? (qSBlockSize / 2U) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, HALF_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum; rowLoopIdx++) { + uint32_t pingpongFlag = rowLoopIdx % 2U; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = + (rowLoopIdx == rowLoopNum - 1U) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + + AscendC::WaitFlag(EVENT_ID0); + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + (rowLoopIdx == 0U), + columnNumRound, + pingpongFlag, + curStackTileMod, + isSplitKV); + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gOutput, AscendC::GlobalTensor gInput, AscendC::GlobalTensor gSink, + AscendC::GlobalTensor gMask, const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, + const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize, + uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown, + uint32_t kvSStartIdx, uint32_t kvSEndIdx, bool isLastStackTile, bool isSplitKV = false) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t columnNum = actualBlockShape.n(); + uint32_t columnNumRound = NpuArch::Detail::Alignment::RoundUp(columnNum, BLOCK_SIZE); + uint32_t columnNumPad = layoutInput.stride(0); + uint32_t maskStride = layoutMask.stride(0); + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? + 0 : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock; + uint32_t rowSplitSubBlock = (qNBlockSize == 1U) ? (qSBlockSize / 2U) : (qSBlockSize * qNSplitSubBlock); + uint32_t rowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock; + uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock; + + uint32_t tokenNumPerHeadThisSubBlock = AscendC::Std::min(qSBlockSize, rowActualThisSubBlock); + + uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1U) ? rowOffsetThisSubBlock : 0; + + uint32_t gmOffsetMaskRow; + uint32_t gmOffsetMaskColumn; + uint32_t maskColumn; + uint32_t addMaskUbOffset; + if (triUp >= kvSStartIdx) { + uint32_t triUpRoundDown = NpuArch::Detail::Alignment::RoundDown(triUp, BLOCK_SIZE); + gmOffsetMaskRow = triUp - triUpRoundDown; + gmOffsetMaskColumn = 0U; + maskColumn = kvSEndIdx - triUpRoundDown; + addMaskUbOffset = triUpRoundDown - kvSStartIdx; + } else { + gmOffsetMaskRow = 0U; + gmOffsetMaskColumn = kvSStartIdx - triUp; + maskColumn = columnNum; + addMaskUbOffset = 0U; + } + uint32_t maskColumnRound = NpuArch::Detail::Alignment::RoundUp(maskColumn, BLOCK_SIZE); + + int64_t offsetMask = + layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn)); + auto gMaskThisSubBlock = gMask[offsetMask]; + auto layoutMaskThisSubBlock = layoutMask; + + uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, BLOCK_SIZE); + rowNumTile = AscendC::Std::min(rowNumTile, HALF_VECTOR_SIZE); + uint32_t rowLoopNum = NpuArch::Detail::Alignment::CeilDiv(rowActualThisSubBlock, rowNumTile); + + if (rowActualThisSubBlock == 0U) { + Arch::CrossCoreWaitFlag(qkReady); + return; + } + Arch::CrossCoreWaitFlag(qkReady); + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum; rowLoopIdx++) { + uint32_t pingpongFlag = rowLoopIdx % 2U; + uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock; + uint32_t rowNumCurLoop = + (rowLoopIdx == rowLoopNum - 1U) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile; + + uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock; + uint32_t proTokenNum = AscendC::Std::min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % + tokenNumPerHeadThisSubBlock; + uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock; + uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock; + + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gInputCurLoop = gInput[offsetInput]; + AscendC::WaitFlag(EVENT_ID0); + CopySGmToUb( + gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound); + + AscendC::WaitFlag(EVENT_ID3); + CopyMaskGmToUb( + gMaskThisSubBlock, + maskColumn, + maskColumnRound, + maskStride, + tokenNumPerHeadThisSubBlock, + proTokenIdx, + proTokenNum, + integralHeadNum, + epiTokenNum); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + UpCastMask(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound); + AscendC::SetFlag(EVENT_ID3); + ApplyMask( + (pingpongFlag * MAX_UB_S_ELEM_NUM), + rowNumCurLoop, + columnNumRound, + maskColumnRound, + addMaskUbOffset); + + // online softmax vectorized compute + int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0)); + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum)); + SubCoreCompute( + gOutputCurLoop, + layoutOutputCurLoop, + rowOffsetCurLoop, + isFirstStackTile, + (rowLoopIdx == 0), + columnNumRound, + pingpongFlag, + curStackTileMod, + isSplitKV); + } + } + +private: + half scaleValue; + AscendC::LocalTensor lsUbTensor; + AscendC::LocalTensor computeUbTensor; + AscendC::LocalTensor lpUbTensor; + AscendC::LocalTensor maskUbTensor; + AscendC::LocalTensor maskUbTensor16; + AscendC::LocalTensor lmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor llUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor glUbTensor; +}; +} + +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o.hpp new file mode 100644 index 0000000000..7e9179e386 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o.hpp @@ -0,0 +1,1105 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP + +#include +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/epilogue/dispatch_policy.hpp" +#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/matrix_coord.hpp" + +namespace NpuArch::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class UpdateType_, + class LseType_, + LseMode LSE_MODE_> +class BlockEpilogue< + EpilogueAtlasA2RescaleO, + OutputType_, + InputType_, + UpdateType_, + LseType_> +{ +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2RescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementUpdate = typename UpdateType_::Element; + using ElementLse = typename LseType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutUpdate = typename UpdateType_::Layout; + using LayoutLse = typename LseType_::Layout; + + static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr float LSE_OUT_INI = std::numeric_limits::infinity(); + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t HALF_DM_UB_SIZE = 64; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t NUM4 = 4; + static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256; + static constexpr uint32_t SIZE_OF_16BIT = 2; + + struct SplitKVParams { + bool isSplitkv = false; + AscendC::GlobalTensor gCombineLse; + AscendC::GlobalTensor gCombineo; + const LayoutLse* layoutgmLse = nullptr; + const LayoutInput* layoutgmLo = nullptr; + }; + + __aicore__ inline + BlockEpilogue() {} + + __aicore__ inline + void init(Arch::Resource &resource) + { + ptoSubBlockIdx = resource.ptoTopology.subBlockIdx; + ptoLanesPerBlock = resource.ptoTopology.lanesPerBlock; + // Allocate UB space + constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LSE_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor16 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + goUbTensor32 = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte(LSE_UB_TENSOR_OFFSET); + } + + __aicore__ inline + ~BlockEpilogue() {} + + __aicore__ inline + void SetMask(int32_t len) + { + uint64_t mask = 0; + uint64_t one = 1; + uint64_t temp = static_cast(len) % static_cast(FLOAT_VECTOR_SIZE); + for (uint64_t i = 0; i < temp; i++) { + mask |= one << i; + } + + if (len == VECTOR_SIZE) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } else if (len >= FLOAT_VECTOR_SIZE) { + AscendC::SetVectorMask(mask, (uint64_t)-1); + } else { + AscendC::SetVectorMask(0x0, mask); + } + } + __aicore__ inline + void InvalidLineLSEProcess( + uint32_t qNThisSubBlock, int32_t delStartRow, uint32_t qSBlockIdx, uint32_t inRowOffsetThisSubBlock, + uint32_t totalRowNum, int32_t delEndRow, uint32_t qSeqlen, uint32_t qSThisSubBlock) + { + uint32_t qNSubBlockStartOffset = qNThisSubBlock == 0U ? qSBlockIdx * VECTOR_SIZE + inRowOffsetThisSubBlock : qSBlockIdx * VECTOR_SIZE; + uint32_t qNSubBlockEnbdOffset = totalRowNum + qNSubBlockStartOffset; + if (qNThisSubBlock == 0U && delStartRow != 0 && qNSubBlockEnbdOffset >= delStartRow) { + uint32_t start = qNSubBlockStartOffset > delStartRow ? 0 : (delStartRow - qNSubBlockStartOffset); + uint32_t end = totalRowNum; + AscendC::PipeBarrier(); + AscendC::Duplicate( + tvUbTensor[start * FLOAT_BLOCK_SIZE], + LSE_OUT_INI, + (end - start) * FLOAT_BLOCK_SIZE + ); + } + if (qNThisSubBlock == 0U && delEndRow != qSeqlen && qNSubBlockStartOffset < delEndRow) { + uint32_t rowStart = qNSubBlockStartOffset; + uint32_t start = 0; + uint32_t end = rowStart + totalRowNum >= delEndRow ? (delEndRow - rowStart) : totalRowNum; + AscendC::PipeBarrier(); + AscendC::Duplicate( + tvUbTensor[start * FLOAT_BLOCK_SIZE], + LSE_OUT_INI, + (end - start) * FLOAT_BLOCK_SIZE + ); + } + if (qNThisSubBlock != 0U && delStartRow != 0 && qNSubBlockEnbdOffset >= delStartRow) { + uint32_t start = delStartRow - qNSubBlockStartOffset; + uint32_t end = qSThisSubBlock; + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::PipeBarrier(); + AscendC::Duplicate( + tvUbTensor[(qNIdx * qSThisSubBlock + start) * FLOAT_BLOCK_SIZE], + LSE_OUT_INI, + (end - start) * FLOAT_BLOCK_SIZE + ); + } + } + if (qNThisSubBlock != 0U && delEndRow != qSeqlen && qNSubBlockStartOffset < delEndRow) { + uint32_t start = 0; + uint32_t end = qNSubBlockEnbdOffset >= delEndRow ? (delEndRow - qNSubBlockStartOffset) : totalRowNum; + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::PipeBarrier(); + AscendC::Duplicate( + tvUbTensor[(qNIdx * qSThisSubBlock + start) * FLOAT_BLOCK_SIZE], + LSE_OUT_INI, + (end - start) * FLOAT_BLOCK_SIZE + ); + } + } + } + __aicore__ inline + void CopyOToGm( + AscendC::GlobalTensor gOutput, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum, + uint32_t qSThisSubBlock, uint32_t embedV, uint32_t embedRoundV, uint32_t oHiddenSize) + { + uint32_t innerOGmOffset = 0; + uint32_t innerGOUbOffset = 0; + if (proTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], + goUbTensor16[innerGOUbOffset], + AscendC::DataCopyExtParams( + proTokenNum, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + innerOGmOffset += embedV; + innerGOUbOffset += proTokenNum * embedRoundV; + } + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor16[innerGOUbOffset], + AscendC::DataCopyExtParams( + qSThisSubBlock, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + innerOGmOffset += embedV; + innerGOUbOffset += qSThisSubBlock * embedRoundV; + } + if (epiTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor16[innerGOUbOffset], + AscendC::DataCopyExtParams( + epiTokenNum, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + } + } + + __aicore__ inline + void CopyOToGmFp32( + AscendC::GlobalTensor gOutput, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum, + uint32_t qSThisSubBlock, uint32_t embedV, uint32_t embedRoundV, uint32_t oHiddenSize, uint32_t oHiddenSize_gmlo) + { + uint32_t innerOGmOffset = 0; + uint32_t innerGOUbOffset = 0; + if (proTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], + goUbTensor32[innerGOUbOffset], + AscendC::DataCopyExtParams( + proTokenNum, embedV * sizeof(float), 0, (oHiddenSize_gmlo - embedV) * sizeof(float), 0)); + innerOGmOffset += embedV; + innerGOUbOffset += proTokenNum * embedRoundV; + } + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor32[innerGOUbOffset], + AscendC::DataCopyExtParams( + qSThisSubBlock, embedV * sizeof(float), 0, (oHiddenSize_gmlo - embedV) * sizeof(float), 0)); + innerOGmOffset += embedV; + innerGOUbOffset += qSThisSubBlock * embedRoundV; + } + if (epiTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor32[innerGOUbOffset], + AscendC::DataCopyExtParams( + epiTokenNum, embedV * sizeof(float), 0, (oHiddenSize_gmlo - embedV) * sizeof(float), 0)); + } + } + + + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embed = layoutInput.shape(1); + uint32_t embedRound = layoutInput.stride(0); + uint32_t curRowNumRound = NpuArch::Detail::Alignment::RoundUp(curRowNum, FLOAT_BLOCK_SIZE); + uint32_t qSBlockSize = layoutOutput.shape(0); + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t qHeads = layoutLse.shape(1); + uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop; + + if (!isFirstStackTile) { + AscendC::WaitFlag(EVENT_ID3); + AscendC::DataCopy( + loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + } + AscendC::WaitFlag(EVENT_ID6); + if (!isFirstStackTile) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast(), + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + if (needRowLoop) { + AscendC::DataCopy( + goUbTensor32, gUpdate, + AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) { + AscendC::Mul( + goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Mul( + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID0); + // *** go = lo + go + AscendC::Add( + goUbTensor32, + goUbTensor32, + loUbTensor, + (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID3); + } else { + // *** go = lo + AscendC::DataCopy( + goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + if (isLastStackTile) { + // *** gl_block = expand_to_block(gl), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast()[rowOffsetLoop], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + AscendC::Div( + goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + } + if (embed % FLOAT_VECTOR_SIZE > 0) { + SetMask(embed % FLOAT_VECTOR_SIZE); + AscendC::Div( + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + // *** go = castfp32to16(go) + if (std::is_same::value) { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_RINT, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_NONE, (uint64_t)0, + (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE, + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ***move O to GM + CopyOToGm( + gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed, embedRound, oHiddenSize); + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + if (isLastRowLoop) { + AscendC::PipeBarrier(); + AscendC::Ln( + lse32_ubuf_tensor, + glUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + AscendC::Add( + lse32_ubuf_tensor, + lse32_ubuf_tensor, + gmUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // *** lse_block = expand_to_block(lse), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + lse32_ubuf_tensor.ReinterpretCast(), + NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID4); + + if (qNThisSubBlock == 0U) { + AscendC::DataCopyPad( + gLse, tvUbTensor, + AscendC::DataCopyExtParams( + totalRowNum, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } else { + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + gLse[qNIdx], + tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE], + AscendC::DataCopyExtParams( + qSBlockSize, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } + } + AscendC::SetFlag(EVENT_ID4); + } + } + } else if (needRowLoop) { + AscendC::SetFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID5); + AscendC::DataCopy( + gUpdate, goUbTensor32, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0)); + } + AscendC::SetFlag(EVENT_ID6); + } + + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum, + uint32_t rowOffsetCurLoop, int32_t delStartRow, int32_t delEndRow, uint32_t qSeqlen, + uint32_t qSBlockIdx, uint32_t rowNum, uint32_t inRowOffsetThisSubBlock, + const SplitKVParams& splitParams, uint32_t curQNBlockTile) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embedV = layoutInput.shape(1); + uint32_t embedRoundV = layoutInput.stride(0); + uint32_t curRowNumRound = NpuArch::Detail::Alignment::RoundUp(curRowNum, FLOAT_BLOCK_SIZE); + uint32_t qSBlockSize = layoutOutput.shape(0); + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t qHeads = layoutLse.shape(1); + uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop; + + uint32_t oHiddenSize_gmlo = 0; + uint32_t qHeads_gmlse = 0; + if (splitParams.isSplitkv) { + oHiddenSize_gmlo = splitParams.layoutgmLo->shape(1); + qHeads_gmlse = splitParams.layoutgmLse->shape(1); + } + + if (!isFirstStackTile) { + AscendC::WaitFlag(EVENT_ID3); + AscendC::DataCopy( + loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRoundV / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + } + AscendC::WaitFlag(EVENT_ID6); + if (!isFirstStackTile) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb(tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast(), + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + if (needRowLoop) { + AscendC::DataCopy( + goUbTensor32, gUpdate, + AscendC::DataCopyParams(1, curRowNum * embedRoundV / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < embedV / FLOAT_VECTOR_SIZE; ++vmul_idx) { + AscendC::Mul( + goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / FLOAT_BLOCK_SIZE, embedRoundV / FLOAT_BLOCK_SIZE, 1)); + } + if (embedV % FLOAT_VECTOR_SIZE > 0) { + SetMask(embedV % FLOAT_VECTOR_SIZE); + AscendC::Mul( + goUbTensor32[embedV / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embedV / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / FLOAT_BLOCK_SIZE, embedRoundV / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID0); + // *** go = lo + go + AscendC::Add( + goUbTensor32, + goUbTensor32, + loUbTensor, + (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(curRowNum * embedRoundV, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID3); + } else { + // *** go = lo + AscendC::DataCopy( + goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRoundV / FLOAT_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + if (isLastStackTile) { + // *** gl_block = expand_to_block(gl), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast()[rowOffsetLoop], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < embedV / FLOAT_VECTOR_SIZE; ++vdiv_idx) { + AscendC::Div( + goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / FLOAT_BLOCK_SIZE, embedRoundV / FLOAT_BLOCK_SIZE, 1)); + } + if (embedV % FLOAT_VECTOR_SIZE > 0) { + SetMask(embedV % FLOAT_VECTOR_SIZE); + AscendC::Div( + goUbTensor32[embedV / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + goUbTensor32[embedV / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / FLOAT_BLOCK_SIZE, embedRoundV / FLOAT_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + + // *** go = castfp32to16(go) + if (!splitParams.isSplitkv) { + if (std::is_same::value) { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_RINT, (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(curRowNum * embedRoundV, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } else { + AscendC::Cast( + goUbTensor16, goUbTensor32, + AscendC::RoundMode::CAST_NONE, (uint64_t)0, + NpuArch::Detail::Alignment::CeilDiv(curRowNum * embedRoundV, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 4, 8)); + } + } + uint32_t rowStart = qSBlockIdx * VECTOR_SIZE + rowOffsetCurLoop ; + uint32_t innerGOUbOffset = 0; + uint32_t subBlockStart = (curQNBlockTile == 1U) ? rowStart : (rowStart >= qSeqlen ? rowStart - rowStart / qSeqlen * qSeqlen : rowStart); + if (delStartRow != 0) { + if (proTokenNum != 0U && subBlockStart + proTokenNum >= delStartRow) { + uint32_t start = subBlockStart >= delStartRow ? 0 : delStartRow - subBlockStart; + uint32_t end = proTokenNum; + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset + start * embedRoundV], + static_cast(0), + (end - start) * embedRoundV + ); + innerGOUbOffset += proTokenNum * embedRoundV; + } + if (subBlockStart + qSThisSubBlock >= delStartRow) { + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + uint32_t start = subBlockStart >= delStartRow ? 0 : delStartRow - subBlockStart; + uint32_t end = qSThisSubBlock; + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset + start * embedRoundV], + static_cast(0), + (end - start) * embedRoundV + ); + innerGOUbOffset += qSThisSubBlock * embedRoundV; + } + } + if (epiTokenNum != 0U && subBlockStart + epiTokenNum >= delStartRow) { + uint32_t start = subBlockStart >= delStartRow ? 0 : delStartRow - subBlockStart; + uint32_t end = epiTokenNum; + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset + start * embedRoundV], + static_cast(0), + (end - start) * embedRoundV + ); + } + } + if (delEndRow != qSeqlen) { + if (proTokenNum != 0U && subBlockStart < delEndRow) { + uint32_t start = curQNBlockTile == 1U ? rowStart : 0; + uint32_t end = (subBlockStart + proTokenNum >= delEndRow) ? + (curQNBlockTile == 1U ? delEndRow : delEndRow - subBlockStart) + : subBlockStart + proTokenNum; + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset], + static_cast(0), + (end - start) * embedRoundV + ); + innerGOUbOffset += proTokenNum * embedRoundV; + } + if (subBlockStart < delEndRow) { + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + uint32_t start = curQNBlockTile == 1U ? subBlockStart : proTokenNum; + uint32_t end = (subBlockStart + qSThisSubBlock >= delEndRow) ? + (curQNBlockTile == 1U ? delEndRow : delEndRow - subBlockStart) + : start + qSThisSubBlock; + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset], + static_cast(0), + (end - start) * embedRoundV + ); + innerGOUbOffset += qSThisSubBlock * embedRoundV; + } + } + if (epiTokenNum != 0U && subBlockStart < delEndRow) { + uint32_t start = curQNBlockTile == 1U ? subBlockStart : proTokenNum + integralHeadNum * qSThisSubBlock + subBlockStart; + uint32_t end = curQNBlockTile == 1U ? (subBlockStart + epiTokenNum >= delEndRow ? delEndRow : subBlockStart + epiTokenNum) : + (epiTokenNum >= delEndRow ? start + delEndRow: start + epiTokenNum); + AscendC::PipeBarrier(); + AscendC::Duplicate( + goUbTensor16[innerGOUbOffset], + static_cast(0), + (end - start) * embedRoundV + ); + } + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + if (splitParams.isSplitkv) { + CopyOToGmFp32( + splitParams.gCombineo, + proTokenIdx, + proTokenNum, + epiTokenNum, + integralHeadNum, + qSThisSubBlock, + embedV, + embedRoundV, + oHiddenSize, oHiddenSize_gmlo); + } else { + CopyOToGm( + gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, + qSThisSubBlock, embedV, embedRoundV, oHiddenSize); + } + + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + if (isLastRowLoop) { + AscendC::PipeBarrier(); + AscendC::Ln( + lse32_ubuf_tensor, + glUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + AscendC::Add( + lse32_ubuf_tensor, + lse32_ubuf_tensor, + gmUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // *** lse_block = expand_to_block(lse), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + lse32_ubuf_tensor.ReinterpretCast(), + NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), + AscendC::BrcbRepeatParams(1, 8)); + InvalidLineLSEProcess(qNThisSubBlock, delStartRow, qSBlockIdx, + inRowOffsetThisSubBlock, totalRowNum, delEndRow, qSeqlen, qSThisSubBlock); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID4); + + if (qNThisSubBlock == 0U) { + if (splitParams.isSplitkv) { + AscendC::DataCopyPad( + splitParams.gCombineLse, tvUbTensor, + AscendC::DataCopyExtParams( + totalRowNum, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0)); + } + AscendC::DataCopyPad( + gLse, tvUbTensor, + AscendC::DataCopyExtParams( + totalRowNum, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } else { + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + if (splitParams.isSplitkv) { + AscendC::DataCopyPad( + splitParams.gCombineLse[qNIdx], + tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE], + AscendC::DataCopyExtParams( + qSBlockSize, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0)); + } + AscendC::DataCopyPad( + gLse[qNIdx], + tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE], + AscendC::DataCopyExtParams( + qSBlockSize, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } + } + AscendC::SetFlag(EVENT_ID4); + } + } else { + if (splitParams.isSplitkv) { + if (isLastRowLoop) { + AscendC::PipeBarrier(); + AscendC::Ln( + lse32_ubuf_tensor, + glUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + + AscendC::PipeBarrier(); + AscendC::Add( + lse32_ubuf_tensor, + lse32_ubuf_tensor, + gmUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + + // *** lse_block = expand_to_block(lse), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + lse32_ubuf_tensor.ReinterpretCast(), + NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID4); + + if (qNThisSubBlock == 0U) { + AscendC::DataCopyPad( + splitParams.gCombineLse, tvUbTensor, + AscendC::DataCopyExtParams( + totalRowNum, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0)); + } else { + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + splitParams.gCombineLse[qNIdx], + tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE], + AscendC::DataCopyExtParams( + qSBlockSize, sizeof(float), 0, (qHeads_gmlse - 1) * sizeof(float), 0)); + } + } + AscendC::SetFlag(EVENT_ID4); + } + } + } + } else if (needRowLoop) { + AscendC::SetFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID5); + AscendC::DataCopy( + gUpdate, goUbTensor32, AscendC::DataCopyParams(1, curRowNum * embedRoundV / FLOAT_BLOCK_SIZE, 0, 0)); + } + AscendC::SetFlag(EVENT_ID6); + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + GemmCoord actualBlockShape, + uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t kvNBlockSize, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + uint32_t isNew) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t embed = actualBlockShape.n(); + uint32_t embedRoundV = (layoutInput.stride(0) == 0) ? BLOCK_SIZE : layoutInput.stride(0); + uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embedRoundV; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t kvNSplitSubBlock = kvNBlockSize / subBlockNum; + uint32_t kvNThisSubBlock = (kvNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (kvNBlockSize - kvNSplitSubBlock) + : kvNSplitSubBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) + : qNSplitSubBlock) + : (kvNThisSubBlock * qNBlockSize); + + uint32_t inRowSplitSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock)) : + (qSBlockSize * qNBlockSize * kvNSplitSubBlock); + + uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock; + uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock; + + uint32_t outRowOffsetThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0) : 0; + uint32_t outColOffsetThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? 0 : (subBlockIdx * qNSplitSubBlock * embed)) : + (subBlockIdx * kvNSplitSubBlock * qNBlockSize * embed); + + uint32_t qSThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize) : qSBlockSize; + + int64_t outOffsetSubBlock = + layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); + + uint32_t outLseRowOffsetThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0) : 0; + uint32_t outLseColOffsetThisSubBlock = (kvNBlockSize == 1U) ? + ((qNBlockSize == 1U) ? 0 : (subBlockIdx * qNSplitSubBlock)) : + (subBlockIdx * kvNSplitSubBlock * qNBlockSize); + + int64_t offsetLse = + layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock)); + + auto gLseThisSubBlock = gLse[offsetLse]; + auto layoutOutLseThisSubBlock = layoutLse; + + if (inRowActualThisSubBlock > 0U) { + uint32_t rowLoop = NpuArch::Detail::Alignment::CeilDiv(inRowActualThisSubBlock, rowNumTile); + uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0; + + uint32_t proTokenIdx = 0; + uint32_t proTokenIdxPre = 0; + uint32_t proTokenNum = 0; + uint32_t epiTokenNum = 0; + uint32_t integralHeadNum = 0; + uint32_t qSRemian = qSThisSubBlock; + + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { + uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop; + uint32_t rowActualCurLoop = + (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile; + + int64_t offsetOutput = + static_cast(rowLoopIdx * rowNumTile / qSThisSubBlock * embed) + outOffsetSubBlock; + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput; + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gInputCurLoop = gInput[offsetInput]; + + auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embed)); + int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gUpdateCurLoop = gUpdate[offsetUpdate]; + auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embed)); + + proTokenIdx = rowOffsetLoop % qSThisSubBlock; + proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock; + integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock; + epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock; + + SubCoreCompute( + gOutputCurLoop, + gInputCurLoop, + gUpdateCurLoop, + gLseThisSubBlock, + layoutOutputCurLoop, + layoutInputCurLoop, + layoutUpdateCurLoop, + layoutOutLseThisSubBlock, + qNThisSubBlock, + qSThisSubBlock, + inRowActualThisSubBlock, + isFirstStackTile, + isLastStackTile, + curStackTileMod, + needRowLoop, + (rowLoopIdx == rowLoop - 1U), + rowOffsetLoop, + proTokenIdx, + proTokenNum, + epiTokenNum, + integralHeadNum); + } + } + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + GemmCoord actualBlockShape, + uint32_t qSBlockSize, uint32_t qNBlockSize, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + int32_t delStartRow, int32_t delEndRow, uint32_t qSeqlen, uint32_t qSBlockIdx, uint32_t curQNBlockTile, + const SplitKVParams& splitParams = SplitKVParams()) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t embedV = actualBlockShape.n(); + uint32_t embedRoundV = (layoutInput.stride(0) == 0) ? BLOCK_SIZE : layoutInput.stride(0); + uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embedRoundV; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE); + + uint32_t subBlockIdx = ptoSubBlockIdx; + uint32_t subBlockNum = ptoLanesPerBlock; + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) + : qNSplitSubBlock; + uint32_t inRowSplitSubBlock = + (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock); + uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock; + uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock; + uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0; + uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embedV; + uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize; + int64_t outOffsetSubBlock = + layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); + + int64_t gmlooutOffsetSubBlock = 0; + if (splitParams.isSplitkv) { + gmlooutOffsetSubBlock = + splitParams.layoutgmLo->GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); + } + + uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1U) ? + inRowOffsetThisSubBlock : 0; + uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1U) ? + 0 : subBlockIdx * qNSplitSubBlock; + int64_t offsetLse = + layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock)); + auto gLseThisSubBlock = gLse[offsetLse]; + + auto layoutOutLseThisSubBlock = layoutLse; + + int64_t gmLseoffsetLse = 0; + if (splitParams.isSplitkv) { + gmLseoffsetLse = + splitParams.layoutgmLse->GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock)); + } + + // Prepare block params for SubCoreCompute + SplitKVParams blockParams = splitParams; + if (splitParams.isSplitkv) { + blockParams.gCombineLse = splitParams.gCombineLse[gmLseoffsetLse]; + } + + if (inRowActualThisSubBlock > 0U) { + uint32_t rowLoop = NpuArch::Detail::Alignment::CeilDiv(inRowActualThisSubBlock, rowNumTile); + uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0; + + // The rows of each cycle consist of multiple heads with several tokens. + // There are several integral heads, one prologue head, one epilogue head. + uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part + uint32_t proTokenIdxPre = 0; // the token idx of the start token of the pre prologue part + uint32_t proTokenNum = 0; // the token num of the prologue part + uint32_t epiTokenNum = 0; // the token num of the epilogue part + uint32_t integralHeadNum = 0; // the number of integral heads within a cycle + uint32_t qSRemian = qSThisSubBlock; + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { + uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop; + uint32_t rowActualCurLoop = + (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile; + + int64_t offsetOutput = + static_cast(rowLoopIdx * rowNumTile / qSThisSubBlock * embedV) + outOffsetSubBlock; + + int64_t gmloffset = 0; + if (splitParams.isSplitkv) { + gmloffset = + static_cast(rowLoopIdx * rowNumTile / qSThisSubBlock * embedV) + gmlooutOffsetSubBlock; + blockParams.gCombineo = splitParams.gCombineo[gmloffset]; + } + + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput; + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gInputCurLoop = gInput[offsetInput]; + auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embedV)); + + int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gUpdateCurLoop = gUpdate[offsetUpdate]; + auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embedV)); + + proTokenIdx = rowOffsetLoop % qSThisSubBlock; + proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock; + integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock; + epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock; + + SubCoreCompute( + gOutputCurLoop, + gInputCurLoop, + gUpdateCurLoop, + gLseThisSubBlock, + layoutOutputCurLoop, + layoutInputCurLoop, + layoutUpdateCurLoop, + layoutOutLseThisSubBlock, + qNThisSubBlock, + qSThisSubBlock, + inRowActualThisSubBlock, + isFirstStackTile, + isLastStackTile, + curStackTileMod, + needRowLoop, + (rowLoopIdx == rowLoop - 1U), + rowOffsetLoop, + proTokenIdx, + proTokenNum, + epiTokenNum, + integralHeadNum, + rowOffsetCurLoop, + delStartRow, + delEndRow, + qSeqlen, + qSBlockIdx, + rowNum, + inRowOffsetThisSubBlock, + blockParams, + curQNBlockTile); + } + } + } + + +private: + uint32_t ptoSubBlockIdx = 0; + uint32_t ptoLanesPerBlock = 1; + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor goUbTensor16; + AscendC::LocalTensor goUbTensor32; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor lse32_ubuf_tensor; +}; +} + +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp new file mode 100644 index 0000000000..e83c18ef74 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp @@ -0,0 +1,472 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP +#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/epilogue/dispatch_policy.hpp" +#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/matrix_coord.hpp" + +namespace NpuArch::Epilogue::Block { + +template < + class OutputType_, + class InputType_, + class UpdateType_, + class LseType_, + LseMode LSE_MODE_> +class BlockEpilogue< + EpilogueAtlasA2RescaleO, + OutputType_, + InputType_, + UpdateType_, + LseType_> +{ +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasA2RescaleO; + using ArchTag = typename DispatchPolicy::ArchTag; + + using ElementOutput = typename OutputType_::Element; + using ElementInput = typename InputType_::Element; + using ElementUpdate = typename UpdateType_::Element; + using ElementLse = typename LseType_::Element; + + using LayoutOutput = typename OutputType_::Layout; + using LayoutInput = typename InputType_::Layout; + using LayoutUpdate = typename UpdateType_::Layout; + using LayoutLse = typename LseType_::Layout; + + static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE; + + static constexpr uint32_t HALF_ELENUM_PER_BLK = 16; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128; + static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64; + static constexpr uint32_t HALF_ELENUM_PER_LINE = 256; + static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128; + static constexpr uint32_t MULTIPLIER = 2; + static constexpr uint32_t FLOAT_BLOCK_SIZE = 8; + static constexpr uint32_t HALF_BLOCK_SIZE = 16; + static constexpr uint32_t FLOAT_VECTOR_SIZE = 64; + static constexpr uint32_t HALF_VECTOR_SIZE = 128; + static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024; + static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384; + static constexpr uint32_t HALF_DM_UB_SIZE = 64; + static constexpr uint32_t HALF_LL_UB_SIZE = 256; + static constexpr uint32_t VECTOR_SIZE = 128; + static constexpr uint32_t NUM4 = 4; + static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192; + static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256; + static constexpr uint32_t SIZE_OF_16BIT = 2; + + __aicore__ inline + BlockEpilogue() {} + + __aicore__ inline + void init(Arch::Resource &resource) + { + // Allocate UB space + constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE; + constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE; + + constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LSE32_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t LSE16_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE; + constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE; + + loUbTensor = resource.ubBuf.template GetBufferByByte(LO_UB_TENSOR_OFFSET); + dmUbTensor = resource.ubBuf.template GetBufferByByte(DM_UB_TENSOR_OFFSET); + glUbTensor = resource.ubBuf.template GetBufferByByte(GL_UB_TENSOR_OFFSET); + tvUbTensor = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + tvUbTensor32 = resource.ubBuf.template GetBufferByByte(TV_UB_TENSOR_OFFSET); + goUbTensor = resource.ubBuf.template GetBufferByByte(GO_UB_TENSOR_OFFSET); + hmUbTensor = resource.ubBuf.template GetBufferByByte(HM_UB_TENSOR_OFFSET); + gmUbTensor = resource.ubBuf.template GetBufferByByte(GM_UB_TENSOR_OFFSET); + lse16_ubuf_tensor = resource.ubBuf.template GetBufferByByte(LSE16_UB_TENSOR_OFFSET); + lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte(LSE32_UB_TENSOR_OFFSET); + } + + __aicore__ inline + ~BlockEpilogue() {} + + __aicore__ inline + void SetMask(int32_t len) + { + const int32_t MAX_MASK_LEN = 128; + const int32_t HALF_MASK_LEN = 64; + if (len >= MAX_MASK_LEN) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + return; + } + int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0; + int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len; + if (len < HALF_MASK_LEN) { + AscendC::SetVectorMask(0x0, ((uint64_t)1 << lowMask) - 1); + } else { + AscendC::SetVectorMask(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff); + } + } + + __aicore__ inline + void CopyOToGm( + AscendC::GlobalTensor gOutput, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum, + uint32_t qSThisSubBlock, uint32_t embedV, uint32_t embedRoundV, uint32_t oHiddenSize) + { + uint32_t innerOGmOffset = 0; + uint32_t innerGOUbOffset = 0; + if (proTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset + proTokenIdx * oHiddenSize], + goUbTensor[innerGOUbOffset], + AscendC::DataCopyExtParams( + proTokenNum, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + innerOGmOffset += embedV; + innerGOUbOffset += proTokenNum * embedRoundV; + } + for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor[innerGOUbOffset], + AscendC::DataCopyExtParams( + qSThisSubBlock, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + innerOGmOffset += embedV; + innerGOUbOffset += qSThisSubBlock * embedRoundV; + } + if (epiTokenNum != 0U) { + AscendC::DataCopyPad( + gOutput[innerOGmOffset], + goUbTensor[innerGOUbOffset], + AscendC::DataCopyExtParams( + epiTokenNum, embedV * SIZE_OF_16BIT, 0, (oHiddenSize - embedV) * SIZE_OF_16BIT, 0)); + } + } + + __aicore__ inline + void SubCoreCompute( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop, + uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum) + { + uint32_t curRowNum = layoutInput.shape(0); + uint32_t embedV = layoutInput.shape(1); + uint32_t embedRoundV = layoutInput.stride(0); + uint32_t curRowNumRound = NpuArch::Detail::Alignment::RoundUp(curRowNum, HALF_BLOCK_SIZE); + uint32_t qSBlockSize = layoutOutput.shape(0); + uint32_t oHiddenSize = layoutOutput.shape(1); + uint32_t qHeads = layoutLse.shape(1); + uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop; + + if (!isFirstStackTile) { + AscendC::WaitFlag(EVENT_ID3); + AscendC::DataCopy( + loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRoundV / HALF_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + } + AscendC::WaitFlag(EVENT_ID6); + if (!isFirstStackTile) { + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast(), + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + if (needRowLoop) { + AscendC::DataCopy( + goUbTensor, gUpdate, + AscendC::DataCopyParams(1, curRowNum * embedRoundV / HALF_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } + // *** go = go * dm_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vmul_idx = 0; vmul_idx < embedV / HALF_VECTOR_SIZE; ++vmul_idx) { + AscendC::Mul( + goUbTensor[vmul_idx * HALF_VECTOR_SIZE], + goUbTensor[vmul_idx * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / HALF_BLOCK_SIZE, embedRoundV / HALF_BLOCK_SIZE, 1)); + } + if (embedV % HALF_VECTOR_SIZE > 0) { + SetMask(embedV % HALF_VECTOR_SIZE); + AscendC::Mul( + goUbTensor[embedV / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + goUbTensor[embedV / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / HALF_BLOCK_SIZE, embedRoundV / HALF_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID0); + // *** go = lo + go + AscendC::Add( + goUbTensor, + goUbTensor, + loUbTensor, + (uint64_t)0, + (curRowNum * embedRoundV + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE, + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID3); + } else { + // *** go = lo + AscendC::DataCopy( + goUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRoundV / HALF_BLOCK_SIZE, 0, 0)); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + if (isLastStackTile) { + // *** gl_block = expand_to_block(gl), 存放于 tv + AscendC::Brcb( + tvUbTensor.ReinterpretCast(), + glUbTensor.ReinterpretCast()[rowOffsetLoop], + curRowNumRound / FLOAT_BLOCK_SIZE, + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + // *** go = go / gl_block + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + for (uint32_t vdiv_idx = 0; vdiv_idx < embedV / HALF_VECTOR_SIZE; ++vdiv_idx) { + AscendC::Div( + goUbTensor[vdiv_idx * HALF_VECTOR_SIZE], + goUbTensor[vdiv_idx * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / HALF_BLOCK_SIZE, embedRoundV / HALF_BLOCK_SIZE, 1)); + } + if (embedV % HALF_VECTOR_SIZE > 0) { + SetMask(embedV % HALF_VECTOR_SIZE); + AscendC::Div( + goUbTensor[embedV / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + goUbTensor[embedV / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE], + tvUbTensor, + (uint64_t)0, + curRowNum, + AscendC::BinaryRepeatParams( + 1, 1, 0, embedRoundV / HALF_BLOCK_SIZE, embedRoundV / HALF_BLOCK_SIZE, 1)); + AscendC::SetVectorMask((uint64_t)-1, (uint64_t)-1); + } + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + // ***move O to GM + CopyOToGm( + gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, + qSThisSubBlock, embedV, embedRoundV, oHiddenSize); + if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) { + if (isLastRowLoop) { + AscendC::PipeBarrier(); + AscendC::Ln( + lse16_ubuf_tensor, + glUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, HALF_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Add( + lse16_ubuf_tensor, + lse16_ubuf_tensor, + gmUbTensor, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, HALF_VECTOR_SIZE), + AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)); + AscendC::PipeBarrier(); + AscendC::Cast( + lse32_ubuf_tensor, + lse16_ubuf_tensor, + AscendC::RoundMode::CAST_NONE, + (uint64_t)0, NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE), + AscendC::UnaryRepeatParams(1, 1, 8, 4)); + AscendC::PipeBarrier(); + + // *** lse_block = expand_to_block(lse), 存放于 tv + AscendC::Brcb( + tvUbTensor32.ReinterpretCast(), + lse32_ubuf_tensor.ReinterpretCast(), + NpuArch::Detail::Alignment::CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE), + AscendC::BrcbRepeatParams(1, 8)); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID4); + + if (qNThisSubBlock == 0U) { + AscendC::DataCopyPad( + gLse, tvUbTensor32, + AscendC::DataCopyExtParams( + totalRowNum, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } else { + for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) { + AscendC::DataCopyPad( + gLse[qNIdx], + tvUbTensor32[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE], + AscendC::DataCopyExtParams( + qSBlockSize, sizeof(float), 0, (qHeads - 1) * sizeof(float), 0)); + } + } + AscendC::SetFlag(EVENT_ID4); + } + } + } else if (needRowLoop) { + AscendC::SetFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID5); + AscendC::DataCopy( + gUpdate, goUbTensor, AscendC::DataCopyParams(1, curRowNum * embedRoundV / HALF_BLOCK_SIZE, 0, 0)); + } + AscendC::SetFlag(EVENT_ID6); + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gOutput, + AscendC::GlobalTensor gInput, + AscendC::GlobalTensor gUpdate, + AscendC::GlobalTensor gLse, + const LayoutOutput &layoutOutput, + const LayoutInput &layoutInput, + const LayoutUpdate &layoutUpdate, + const LayoutLse &layoutLse, + GemmCoord actualBlockShape, + uint32_t qSBlockSize, uint32_t qNBlockSize, + uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod, + int32_t delStartRow, int32_t delEndRow, uint32_t qSeqlen, uint32_t qSBlockIdx, uint32_t curQNBlockTile) + { + uint32_t rowNum = actualBlockShape.m(); + uint32_t embedV = actualBlockShape.n(); + uint32_t embedRoundV = (layoutInput.stride(0) == 0) ? BLOCK_SIZE : layoutInput.stride(0); + uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embedRoundV; + uint32_t rowNumTile = NpuArch::Detail::Alignment::RoundDown(maxRowNumPerLoop, HALF_BLOCK_SIZE); + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + + uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum; + uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 0 + : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) + : qNSplitSubBlock; + uint32_t inRowSplitSubBlock = + (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock); + uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock; + uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock; + uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0; + uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embedV; + uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize; + int64_t outOffsetSubBlock = + layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock)); + + uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1U) ? + inRowOffsetThisSubBlock : 0; + uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1U) ? + 0 : subBlockIdx * qNSplitSubBlock; + int64_t offsetLse = + layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock)); + auto gLseThisSubBlock = gLse[offsetLse]; + auto layoutOutLseThisSubBlock = layoutLse; + + if (inRowActualThisSubBlock > 0U) { + uint32_t rowLoop = NpuArch::Detail::Alignment::CeilDiv(inRowActualThisSubBlock, rowNumTile); + uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0; + + // The rows of each cycle consist of multiple heads with several tokens. + // There are several integral heads, one prologue head, one epilogue head. + uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part + uint32_t proTokenIdxPre = 0; // the token idx of the start token of the pre prologue part + uint32_t proTokenNum = 0; // the token num of the prologue part + uint32_t epiTokenNum = 0; // the token num of the epilogue part + uint32_t integralHeadNum = 0; // the number of integral heads within a cycle + uint32_t qSRemian = qSThisSubBlock; + for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) { + uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile; + uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop; + uint32_t rowActualCurLoop = + (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile; + + int64_t offsetOutput = + static_cast(rowLoopIdx * rowNumTile / qSThisSubBlock * embedV) + outOffsetSubBlock; + auto gOutputCurLoop = gOutput[offsetOutput]; + auto layoutOutputCurLoop = layoutOutput; + int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gInputCurLoop = gInput[offsetInput]; + auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embedV)); + + int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0)); + auto gUpdateCurLoop = gUpdate[offsetUpdate]; + auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embedV)); + + proTokenIdx = rowOffsetLoop % qSThisSubBlock; + proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock; + integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock; + epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock; + + SubCoreCompute( + gOutputCurLoop, + gInputCurLoop, + gUpdateCurLoop, + gLseThisSubBlock, + layoutOutputCurLoop, + layoutInputCurLoop, + layoutUpdateCurLoop, + layoutOutLseThisSubBlock, + qNThisSubBlock, + qSThisSubBlock, + inRowActualThisSubBlock, + isFirstStackTile, + isLastStackTile, + curStackTileMod, + needRowLoop, + (rowLoopIdx == rowLoop - 1U), + rowOffsetLoop, + proTokenIdx, + proTokenNum, + epiTokenNum, + integralHeadNum); + } + } + } + +private: + AscendC::LocalTensor loUbTensor; + AscendC::LocalTensor dmUbTensor; + AscendC::LocalTensor hmUbTensor; + AscendC::LocalTensor glUbTensor; + AscendC::LocalTensor tvUbTensor; + AscendC::LocalTensor tvUbTensor32; + AscendC::LocalTensor goUbTensor; + AscendC::LocalTensor gmUbTensor; + AscendC::LocalTensor lse16_ubuf_tensor; + AscendC::LocalTensor lse32_ubuf_tensor; +}; +} + +#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/dispatch_policy.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/dispatch_policy.hpp new file mode 100644 index 0000000000..7f763d93b0 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/dispatch_policy.hpp @@ -0,0 +1,56 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_DISPATCH_POLICY_HPP +#define EPILOGUE_DISPATCH_POLICY_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/arch/arch.hpp" + +namespace NpuArch::Epilogue +{ + +enum class LseMode {NONE = 0, OUT_ONLY = 1}; +enum class SinkMode {DISABLE = 0, ENABLE = 1}; +enum class MaskMode { + NO_MASK = 0, + MASK_CAUSAL = 1, + MASK_SPEC = 2, + MASK_SWA = 4 +}; +// For AtlasA2, FA Infer online Softmax +template +struct EpilogueAtlasA2OnlineSoftmax { + using ArchTag = Arch::AtlasA2; + using IntermPrec = SM_DTYPE_; + static constexpr LseMode LSE_MODE = LSE_MODE_; + static constexpr SinkMode SINK_MODE = SINK_MODE_; + static constexpr MaskMode MASK_MODE = MASK_MODE_; +}; + +// For AtlasA2, FA Infer RescaleO +template +struct EpilogueAtlasA2RescaleO { + using ArchTag = Arch::AtlasA2; + using IntermPrec = SM_DTYPE_; + static constexpr LseMode LSE_MODE = LSE_MODE_; +}; + +// For AtlasA2, FA Infer Deal kv-len=0 +template +struct EpilogueAtlasA2InitOutWhenZero { + using ArchTag = Arch::AtlasA2; + static constexpr LseMode LSE_MODE = LSE_MODE_; +}; + + +} // namespace NpuArch::Epilogue + +#endif // EPILOGUE_DISPATCH_POLICY_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp new file mode 100644 index 0000000000..c7db55445f --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP +#define EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + class ArchTag, + class GmType +> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = static_cast(BYTE_PER_BLK) / static_cast(sizeof(Element)); + + __aicore__ inline + CopyGm2Ub() = default; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutSrc.shape(0), + layoutSrc.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) * sizeof(Element), + (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK, + 0 + ); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = static_cast(BYTE_PER_BLK) / static_cast(sizeof(Element)); + + __aicore__ inline + CopyGm2Ub() = default; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0 + ); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +/// @brief This copy instruction used to copy per token scale from GM to UB. +/// Copy the scale of shape (m,1) on GM to the first column of shape (m,n) on UB, +/// and pad the first block of each row (i.e. pad to shape (m,8) when element type is float). +/// @tparam ArchTag: Architecture tag. +/// @tparam GmType: Type of data on GM. +template < + class ArchTag, + class GmType +> +struct CopyPerTokenScale2Ub { + static_assert(std::is_same_v, + "Unsupported layout for CopyPerTokenScale2Ub."); + + using Element = typename GmType::Element; + using LayoutSrc = typename GmType::Layout; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = static_cast(BYTE_PER_BLK) / static_cast(sizeof(Element)); + + __aicore__ inline + CopyPerTokenScale2Ub() = default; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, + LayoutSrc const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams; + AscendC::DataCopyPadExtParams padParams; + + dataCopyParams.blockCount = layoutSrc.shape(0); + dataCopyParams.blockLen = layoutSrc.shape(1) * sizeof(Element); // per token scale has only one column + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + // Pad the data to the complete block + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = 0; + + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + } +}; + +template < + class ArchTag, + class GmType +> +struct CopyGm2UbAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub aligned, can not find the specialization."); +}; + +template +struct CopyGm2UbAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = static_cast(BYTE_PER_BLK) / static_cast(sizeof(Element)); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + __aicore__ inline + CopyGm2UbAligned() = default; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = NpuArch::Detail::Alignment::CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp new file mode 100644 index 0000000000..4fc7b9a79c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp @@ -0,0 +1,144 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP +#define EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + class ArchTag, + class GmType +> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyUb2Gm() = default; + + __aicore__ inline + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0 + ); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + + +// new add vectorlayout version +template +struct CopyUb2Gm> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + __aicore__ inline + CopyUb2Gm() = default; + + __aicore__ inline + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutDst.shape(0) * sizeof(Element), + 0, + 0, + 0 + ); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + }; +}; + + +template < + class ArchTag, + class GmType +> +struct CopyUb2GmAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm aligned, can not find the specialization."); +}; + +template +struct CopyUb2GmAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + __aicore__ inline + CopyUb2GmAligned() = default; + + __aicore__ inline + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutDst.shape(0); + uint32_t cols = layoutDst.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = NpuArch::Detail::Alignment::CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_column.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_column.hpp new file mode 100644 index 0000000000..19dc7d13f9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_column.hpp @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_BROADCAST_INPLACE_BY_COLUMN_HPP +#define EPILOGUE_TILE_TILE_BROADCAST_INPLACE_BY_COLUMN_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class ComputeType_, + /// Length of the compute buffer + class TileShape_ +> +struct TileBroadcastInplaceByColumn { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileBroadcastInplaceByColumn() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubInOut + ) + { + constexpr uint32_t eleNumPerBlk = static_cast(BYTE_PER_BLK) / static_cast(sizeof(ElementCompute)); + constexpr uint32_t blkNumPerRow = TileShape::COLUMN / eleNumPerBlk; + + constexpr uint64_t defaultMask = BYTE_PER_VECTOR_FRACTAL / sizeof(ElementCompute); + constexpr uint64_t tailMask = (TileShape::ROW % BLK_NUM_PER_VECTOR_FRACTAL) * eleNumPerBlk; + + constexpr uint8_t repeatTimes = 1; + + AscendC::CopyRepeatParams repeatParams; + repeatParams.dstStride = blkNumPerRow; + repeatParams.srcStride = blkNumPerRow; + repeatParams.dstRepeatSize = 1; + repeatParams.srcRepeatSize = 1; + + for (uint32_t rowOffset = 0; rowOffset < TileShape::ROW; rowOffset += BLK_NUM_PER_VECTOR_FRACTAL) { + uint64_t mask = ((TileShape::ROW - rowOffset) >= BLK_NUM_PER_VECTOR_FRACTAL) ? defaultMask : tailMask; + for (uint32_t colOffset = eleNumPerBlk; colOffset < TileShape::COLUMN; colOffset += eleNumPerBlk) { + AscendC::Copy( + ubInOut[rowOffset * TileShape::COLUMN + colOffset], + ubInOut[rowOffset * TileShape::COLUMN], + mask, 1, repeatParams + ); + } + } + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_row.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_row.hpp new file mode 100644 index 0000000000..9aba9c7c39 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_inplace_by_row.hpp @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_BROADCAST_INPLACE_BY_ROW_HPP +#define EPILOGUE_TILE_TILE_BROADCAST_INPLACE_BY_ROW_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class ComputeType_, + /// Length of the compute buffer + class TileShape_ +> +struct TileBroadcastInplaceByRow { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileBroadcastInplaceByRow() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubInOut + ) + { + constexpr uint32_t eleNumPerVectorFractal = static_cast(BYTE_PER_VECTOR_FRACTAL) / static_cast(sizeof(ElementCompute)); + + constexpr uint64_t mask = eleNumPerVectorFractal; + constexpr uint8_t repeatTimes = TileShape::COLUMN / eleNumPerVectorFractal; + + AscendC::CopyRepeatParams repeatParams; + repeatParams.dstStride = 1; + repeatParams.srcStride = 1; + repeatParams.dstRepeatSize = BLK_NUM_PER_VECTOR_FRACTAL; + repeatParams.srcRepeatSize = BLK_NUM_PER_VECTOR_FRACTAL; + + for (uint32_t rowOffset = 1; rowOffset < TileShape::ROW; ++rowOffset) { + AscendC::Copy(ubInOut[rowOffset * TileShape::COLUMN], ubInOut, mask, repeatTimes, repeatParams); + } + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_mul.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_mul.hpp new file mode 100644 index 0000000000..cb988002ea --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_mul.hpp @@ -0,0 +1,140 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_BROADCAST_MUL_HPP +#define EPILOGUE_TILE_TILE_BROADCAST_MUL_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +/// BroadcastMul computes the elementwise multiplication of a tensor of shape (m, n) and a tensor +/// of shape (m, n) after broadcasting. There are two broadcast modes: row-broadcast and +/// column-broadcast. + +/// @brief Computes the elementwise multiplication of a tensor with shape (m, n) and a tensor with +/// original shape (1, n) broadcast to (m, n). +/// @tparam ArchTag_ is the architecture tag. +/// @tparam ComputeType_ includes the element type and layout information. +/// @tparam TileShape_ is the shape (m, n). +template < + class ArchTag_, + class ComputeType_, + class TileShape_ +> +struct TileRowBroadcastMul { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileRowBroadcastMul() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + constexpr uint32_t maxRepeatTimes = 255; + constexpr uint32_t eleNumPerBlk = static_cast(BYTE_PER_BLK) / static_cast(sizeof(ElementCompute)); + + constexpr uint32_t blkNumPerColumn = TileShape::COLUMN / eleNumPerBlk; + AscendC::BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.src0BlkStride = 1; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = blkNumPerColumn; + repeatParams.src0RepStride = blkNumPerColumn; + repeatParams.src1RepStride = 0; + + constexpr uint32_t rowNumPerCompute = maxRepeatTimes; + constexpr uint32_t colNumPerCompute = BYTE_PER_VECTOR_FRACTAL / sizeof(ElementCompute); + for (uint32_t rowOffset = 0; rowOffset < TileShape::ROW; rowOffset += rowNumPerCompute) { + uint32_t residueM = TileShape::ROW - rowOffset; + uint8_t temprepeatTimes = (residueM > rowNumPerCompute) ? rowNumPerCompute : residueM; + uint8_t repeatTimes = static_cast(temprepeatTimes); + for (uint32_t colOffset = 0; colOffset < TileShape::COLUMN; colOffset += colNumPerCompute) { + uint32_t residueN = TileShape::COLUMN - colOffset; + uint64_t mask = (residueN > colNumPerCompute) ? colNumPerCompute : residueN; + AscendC::Mul( + ubOut[rowOffset * TileShape::COLUMN + colOffset], + ubIn0[rowOffset * TileShape::COLUMN + colOffset], + ubIn1[colOffset], + mask, repeatTimes, repeatParams + ); + } + } + } +}; + +/// @brief Compute the elementwise multiplication of a tensor of shape (m, n) and a tensor of shape +/// (m, eleNumPerBlk), which is broadcast from a tensor of shape (m, 1), broadcast to (m, n). +/// @tparam ArchTag_ is the architecture tag. +/// @tparam ComputeType_ includes the element type and layout information. +/// @tparam TileShape_ is the shape (m, n). +template < + class ArchTag_, + class ComputeType_, + class TileShape_ +> +struct TileOneBlkColumnBroadcastMul { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileOneBlkColumnBroadcastMul() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + constexpr uint32_t maxRepeatNum = 255; + constexpr uint32_t eleNumPerBlk = static_cast(BYTE_PER_BLK) / static_cast(sizeof(ElementCompute)); + + constexpr uint32_t blkNumPerColumn = TileShape::COLUMN / eleNumPerBlk; + AscendC::BinaryRepeatParams repeatParams; + repeatParams.dstBlkStride = blkNumPerColumn; + repeatParams.src0BlkStride = blkNumPerColumn; + repeatParams.src1BlkStride = 1; + repeatParams.dstRepStride = 1; + repeatParams.src0RepStride = 1; + repeatParams.src1RepStride = 0; + + constexpr uint32_t rowNumPerCompute = BLK_NUM_PER_VECTOR_FRACTAL; + constexpr uint32_t colNumPerCompute = eleNumPerBlk * maxRepeatNum; + for (uint32_t rowOffset = 0; rowOffset < TileShape::ROW; rowOffset += rowNumPerCompute) { + uint32_t residueM = TileShape::ROW - rowOffset; + uint32_t currentRowNum = (residueM > rowNumPerCompute) ? rowNumPerCompute : residueM; + uint64_t mask = static_cast(currentRowNum) * static_cast(eleNumPerBlk); + for (uint32_t colOffset = 0; colOffset < TileShape::COLUMN; colOffset += colNumPerCompute) { + uint32_t residueN = TileShape::COLUMN - colOffset; + uint32_t currentColNum = (residueN > colNumPerCompute) ? colNumPerCompute : residueN; + uint8_t repeatTimes = static_cast(currentColNum / eleNumPerBlk); + AscendC::Mul( + ubOut[rowOffset * TileShape::COLUMN + colOffset], + ubIn0[rowOffset * TileShape::COLUMN + colOffset], + ubIn1[rowOffset * eleNumPerBlk], + mask, repeatTimes, repeatParams + ); + } + } + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_one_blk.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_one_blk.hpp new file mode 100644 index 0000000000..ba0142e4fb --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_broadcast_one_blk.hpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_BROADCAST_ONE_BLK_HPP +#define EPILOGUE_TILE_TILE_BROADCAST_ONE_BLK_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + class ArchTag_, + class ComputeType_, + uint32_t COMPUTE_LENGTH_ +> +struct TileBroadcastOneBlk { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; + + __aicore__ inline + TileBroadcastOneBlk() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn + ) + { + constexpr uint32_t maxRepeatNum = 255; + constexpr uint32_t eleNumPerBlk = static_cast(BYTE_PER_BLK) / static_cast(sizeof(ElementCompute)); + + AscendC::BrcbRepeatParams repeatParams; + repeatParams.dstBlkStride = 1; + repeatParams.dstRepStride = BLK_NUM_PER_VECTOR_FRACTAL; + + constexpr uint32_t eleNumPerCompute = + NpuArch::Detail::Alignment::RoundDown(maxRepeatNum * BLK_NUM_PER_VECTOR_FRACTAL); + for (uint32_t offset = 0; offset < COMPUTE_LENGTH; offset += eleNumPerCompute) { + uint32_t residueM = COMPUTE_LENGTH - offset; + uint32_t computeM = (residueM > eleNumPerCompute) ? eleNumPerCompute : residueM; + uint8_t repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(computeM)); + AscendC::Brcb( + ubOut[offset * eleNumPerBlk], ubIn[offset], + repeatTimes, repeatParams + ); + } + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_cast.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_cast.hpp new file mode 100644 index 0000000000..90b5295014 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_cast.hpp @@ -0,0 +1,49 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_CAST_HPP +#define EPILOGUE_TILE_TILE_CAST_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class DstType_, + class SrcType_, + /// Length of the compute buffer + class TileShape_ +> +struct TileCast { + using ArchTag = ArchTag_; + using ElementDst = typename DstType_::Element; + using ElementSrc = typename SrcType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileCast() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn + ) + { + AscendC::Cast(ubOut, ubIn, AscendC::RoundMode::CAST_RINT, TileShape::COUNT); + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_copy.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_copy.hpp new file mode 100644 index 0000000000..ca90ed0e7a --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_copy.hpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_COPY_HPP +#define EPILOGUE_TILE_TILE_COPY_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp" +#include "../../../attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag, + class... Args +> +struct TileCopy { + static_assert(DEPENDENT_FALSE, "Unsupported tile_common copy, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for X matrix operand + class XType, + /// GemmType for D matrix operand + class DType +> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType +> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementY = typename YType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyGmToUbY = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType +> +struct TileCopyBf16 { + using ElementC = typename CType::Element; + using ElementX = bfloat16_t; + using ElementY = bfloat16_t; + using ElementD = bfloat16_t; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub>; + using CopyGmToUbY = CopyGm2Ub>; + using CopyUbToGmD = CopyUb2Gm>; +}; + +template < + class ArchTag, + class CType, + class ScaleType, + class PerTokenScaleType, + class DType +> +struct TileCopyPerTokenDequant { + using ElementC = typename CType::Element; + using ElementScale = typename ScaleType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbScale = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyPerTokenScale2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; +} // namespace NpuArch::Epilogue::Tile + +#endif // EPILOGUE_TILE_TILE_COPY_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_add.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_add.hpp new file mode 100644 index 0000000000..1bbd011d16 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_add.hpp @@ -0,0 +1,50 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_ELEMWISE_ADD_HPP +#define EPILOGUE_TILE_TILE_ELEMWISE_ADD_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class ComputeType_, + /// Length of the compute buffer + uint32_t COMPUTE_LENGTH_ +> +struct TileElemWiseAdd { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + + static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; + + __aicore__ inline + TileElemWiseAdd() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + // Do the calculation + AscendC::Add(ubOut, ubIn0, ubIn1, COMPUTE_LENGTH); + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_mul.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_mul.hpp new file mode 100644 index 0000000000..1a9bb3ef32 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_mul.hpp @@ -0,0 +1,49 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_ELEMWISE_MUL_HPP +#define EPILOGUE_TILE_TILE_ELEMWISE_MUL_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +template < + /// Tag indicating architecture + class ArchTag_, + /// Compute data type + class ComputeType_, + /// Length of the compute buffer + class TileShape_ +> +struct TileElemwiseMul { + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + using TileShape = TileShape_; + + __aicore__ inline + TileElemwiseMul() {} + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &ubOut, + AscendC::LocalTensor const &ubIn0, + AscendC::LocalTensor const &ubIn1 + ) + { + // Do the calculation + AscendC::Mul(ubOut, ubIn0, ubIn1, TileShape::COUNT); + } +}; + +} // namespace NpuArch::Epilogue::Tile + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_muls.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_muls.hpp new file mode 100644 index 0000000000..eed9fac22c --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_elemwise_muls.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + + #ifndef EPILOGUE_TILE_TILE_ELEMWISE_MULS_HPP + #define EPILOGUE_TILE_TILE_ELEMWISE_MULS_HPP + + #include "../../../attn_infra/gemm/helper.hpp" + + namespace NpuArch::Epilogue::Tile + { + template< + class ArchTag_, + class ComputeType_, + uint32_t COMPUTE_LENGTH_ + > + struct TileElemWiseMuls{ + using ArchTag = ArchTag_; + using ElementCompute = typename ComputeType_::Element; + + static constexpr uint32_t COMPUTE_LENGTH = COMPUTE_LENGTH_; + + __aicore__ inline + TileElemWiseMuls(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstLocal, + AscendC::LocalTensor srcTensor, + ElementCompute scalar + ){ + AscendC::Muls(dstLocal, srcTensor, scalar, COMPUTE_LENGTH); + } + }; + } + + #endif // EPILOGUE_TILE_TILE_ELEMWISE_MULS_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_swizzle.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_swizzle.hpp new file mode 100644 index 0000000000..8d24f31a35 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_swizzle.hpp @@ -0,0 +1,93 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef EPILOGUE_TILE_TILE_SWIZZLE_HPP +#define EPILOGUE_TILE_TILE_SWIZZLE_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/detail/alignment.hpp" +#include "../../../attn_infra/matrix_coord.hpp" + +namespace NpuArch::Epilogue::Tile +{ + +struct EpilogueIdentityTileSwizzle { + MatrixCoord blockShape; + MatrixCoord tileShape; + MatrixCoord loopsMN; + + __aicore__ inline + EpilogueIdentityTileSwizzle() = default; + + __aicore__ inline + EpilogueIdentityTileSwizzle(MatrixCoord const &blockShape, MatrixCoord const &tileShape) : + blockShape(blockShape), + tileShape(tileShape) + { + loopsMN = NpuArch::Detail::Alignment::CeilDiv(blockShape, tileShape); + } + + __aicore__ inline + uint32_t GetLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + __aicore__ inline + MatrixCoord GetTileCoord(uint32_t loopIdx) const + { + return MatrixCoord{ loopIdx / loopsMN.column(), loopIdx % loopsMN.column() }; + } + + __aicore__ inline + MatrixCoord GetActualTileShape(MatrixCoord const &tileCoord) const + { + return MatrixCoord::Min(tileShape, blockShape - tileCoord * tileShape); + } +}; + +struct EpilogueHorizontalTileSwizzle { + MatrixCoord blockShape; + MatrixCoord tileShape; + MatrixCoord loopsMN; + + __aicore__ inline + EpilogueHorizontalTileSwizzle() = default; + + __aicore__ inline + EpilogueHorizontalTileSwizzle(MatrixCoord const &blockShape, MatrixCoord const &tileShape) : + blockShape(blockShape), + tileShape(tileShape) + { + loopsMN = NpuArch::Detail::Alignment::CeilDiv(blockShape, tileShape); + } + + __aicore__ inline + uint32_t GetLoops() const + { + return loopsMN.row() * loopsMN.column(); + } + + __aicore__ inline + MatrixCoord GetTileCoord(uint32_t loopIdx) const + { + return MatrixCoord{ loopIdx % loopsMN.row(), loopIdx / loopsMN.row() }; + } + + __aicore__ inline + MatrixCoord GetActualTileShape(MatrixCoord const &tileCoord) const + { + return MatrixCoord::Min(tileShape, blockShape - tileCoord * tileShape); + } +}; + +} + +#endif // EPILOGUE_TILE_TILE_SWIZZLE_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad.hpp new file mode 100644 index 0000000000..20e11a4d37 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad.hpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_BLOCK_BLOCK_MMAD_HPP +#define GEMM_BLOCK_BLOCK_MMAD_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp" + +namespace NpuArch::Gemm::Block { + +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class AType, + class BType, + class CType, + class BiasType = void, + class TileCopy = Gemm::Tile::TileCopy, + class TileMmad = Gemm::Tile::TileMmad +> +struct BlockMmad { + static_assert(DEPENDENT_FALSE, "BlockMmad is not implemented for this DispatchPolicy"); +}; + +} // namespace NpuArch::Gemm::Block + +#include "../../../attn_infra/gemm/block/block_mmad_qk.hpp" +#include "../../../attn_infra/gemm/block/block_mmad_qk_decode.hpp" +#include "../../../attn_infra/gemm/block/block_mmad_pv.hpp" +#include "../../../attn_infra/gemm/block/block_mmad_pv_decode.hpp" + +#endif // GEMM_BLOCK_BLOCK_MMAD_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv.hpp new file mode 100644 index 0000000000..07de666f7b --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv.hpp @@ -0,0 +1,325 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_BLOCK_MMAD_PV_HPP +#define GEMM_BLOCK_MMAD_PV_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/coord.hpp" +#include "../../../attn_infra/arch/cross_core_sync.hpp" +#include "../../../attn_infra/gemm/dispatch_policy.hpp" +#include "../../../attn_infra/gemm/helper.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp" +//////////////////////////////////////////////////////////////////// + +namespace NpuArch::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIPV, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIPV; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + static constexpr uint32_t COORD_DIM0 = 0; + static constexpr uint32_t COORD_DIM1 = 1; + static constexpr uint32_t COORD_DIM2 = 2; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + __aicore__ inline + BlockMmad() {} + + __aicore__ inline + void init(Arch::Resource &resource, uint32_t nDyn, uint32_t kDyn, uint32_t KVStackLen = 512, uint32_t l1BufAddrStart = 0) + { + maxKVStackLen = KVStackLen; + // Allocate L1 memory space + l1BTensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) * STAGES); + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + l1NDynamic = nDyn; + l1KDynamic = kDyn; + } + + /// Destructor + __aicore__ inline + ~BlockMmad() {} + + __aicore__ inline + void resetBlockStart(uint32_t kvStart, uint32_t pagedBlockSize) + { + blockStartOffset = kvStart * maxKVStackLen % pagedBlockSize; + } + + __aicore__ inline + void getBlockShape(GemmCoord &actualShape, uint32_t& nowLen) + { + actualShape[COORD_DIM2] = nowLen; + } + + __aicore__ inline + void getKVOffset(uint32_t &kOffset, uint32_t nIdx, uint32_t &strideKV) + { + kOffset = nIdx * maxKVStackLen * strideKV; + } + + __aicore__ inline + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t blockStartOffset, + uint32_t nowNIdx, uint32_t &strideKV, uint32_t &blockSize) + { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + blockStartOffset * strideKV; + } + + __aicore__ inline + void setBlockParam(uint32_t stackSeqTile, uint32_t &blockStart, uint32_t &blockEnd, uint32_t &curBlockTotalNum, uint32_t blockSize){ + if(stackSeqTile >= blockStart && blockSize != 0) { + blockEnd = ((stackSeqTile - blockStart) % blockSize == 0) ? blockSize : (stackSeqTile - blockStart) % blockSize; + curBlockTotalNum = (((stackSeqTile - blockStart) + blockSize - 1) / blockSize) + 1; + } else { + blockStart = stackSeqTile; + blockEnd = stackSeqTile + blockStartOffset; + curBlockTotalNum = 1; + } + } + + __aicore__ inline + void updateBlockOffset(uint32_t nowLen, uint32_t &curBlockIdx, uint32_t blockSize){ + if (blockStartOffset + nowLen == blockSize) { + blockStartOffset = 0; + } else { + blockStartOffset += nowLen; + } + curBlockIdx++; + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, + uint32_t blockStackNum, Arch::CrossCoreFlag softmaxFlag) + { + uint32_t rowNum = actualOriShape[COORD_DIM0]; + uint32_t embed = actualOriShape[COORD_DIM1]; + uint32_t stackSeqTile = actualOriShape[COORD_DIM2]; + GemmCoord actualShape{rowNum, embed, 0}; + uint32_t gBOffset = 0; + + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(stackSeqTile, embed); + AscendC::WaitFlag(EVENT_ID4); + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t curBlockIdx = 0; + uint32_t blockStart = blockSize - blockStartOffset; + uint32_t blockEnd = 0; + uint32_t curBlockTotalNum = 0; + setBlockParam(stackSeqTile, blockStart, blockEnd, curBlockTotalNum, blockSize); + while(curBlockIdx < curBlockTotalNum) { + uint32_t nowLen = (curBlockIdx < (curBlockTotalNum-1)) ? (blockSize - blockStartOffset) : (blockEnd - blockStartOffset); + uint32_t nowNIdx = nIdx * maxKVStackLen / blockSize + curBlockIdx; + getBlockShape(actualShape, nowLen); + getKVOffset(gBlockTable, gBOffset, blockStartOffset, nowNIdx, strideKV, blockSize); + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + uint32_t curBlockSize = (curBlockIdx > 0) ? ((curBlockIdx - 1) * blockSize + blockStart) : 0; + MatrixCoord l1BTileCoord{curBlockSize, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)]; + copyGmToL1B(l1BTile, gB[gBOffset], layoutBInL1, layoutBTile); + updateBlockOffset(nowLen, curBlockIdx, blockSize); + } + } else { + getBlockShape(actualShape, stackSeqTile); + getKVOffset(gBOffset, nIdx, strideKV); + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + copyGmToL1B(l1BTensor, gB[gBOffset], layoutBInL1, layoutBTile); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + Arch::CrossCoreWaitFlag(softmaxFlag); + + uint32_t mL1Loop = NpuArch::Detail::Alignment::CeilDiv(rowNum, L1TileShape::M); + uint32_t kL1Loop = NpuArch::Detail::Alignment::CeilDiv(stackSeqTile, l1KDynamic); + uint32_t nL1Loop = NpuArch::Detail::Alignment::CeilDiv(embed, L0TileShape::N); + + for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; nL1Idx++) { + uint32_t nL1Actual = (nL1Idx < nL1Loop - 1U) ? L0TileShape::N : (embed - nL1Idx * L0TileShape::N); + for (uint32_t mL1Idx = 0; mL1Idx < mL1Loop; mL1Idx++) { + uint32_t mL1Actual = (mL1Idx < mL1Loop - 1U) ? L1TileShape::M : (rowNum - mL1Idx * L1TileShape::M); + AscendC::WaitFlag(l0CPingPongFlag); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t kL1Actual = (kL1Idx < kL1Loop - 1U) ? l1KDynamic : (stackSeqTile - kL1Idx * l1KDynamic); + // load P + AscendC::WaitFlag(l1PPingPongFlag); + MatrixCoord gmATileCoord{mL1Idx * L1TileShape::M, kL1Idx * l1KDynamic}; + auto gmTileA = gA[layoutA.GetOffset(gmATileCoord)]; + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(mL1Actual, kL1Actual)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mL1Actual, kL1Actual); + copyGmToL1A(l1ATensor[l1PPingPongFlag], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1PPingPongFlag); + + uint32_t kL0Loop = NpuArch::Detail::Alignment::CeilDiv(kL1Actual, L0TileShape::K); + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = + (kL0Idx < kL0Loop - 1U) ? L0TileShape::K : (kL1Actual - kL0Idx * L0TileShape::K); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mL1Actual, kL0Actual); + MatrixCoord l1ATileCoord{0, kL0Idx * L0TileShape::K}; + auto l1ATile = l1ATensor[l1PPingPongFlag][layoutAInL1.GetOffset(l1ATileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag); + if (kL0Idx == 0U) { + AscendC::WaitFlag(l1PPingPongFlag); + } + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1); + if (kL0Idx == kL0Loop - 1U) { + AscendC::SetFlag(l1PPingPongFlag); + } + + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, nL1Actual); + MatrixCoord l1BTileCoord{kL1Idx * l1KDynamic + kL0Idx * L0TileShape::K, L0TileShape::N * nL1Idx}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag + 2U); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + bool initMmad = (kL1Idx == 0U) && (kL0Idx == 0U); + uint32_t mL0Align = (mL1Actual + BLOCK_SIZE - 1U) / BLOCK_SIZE * BLOCK_SIZE; + tileMmad(l0CTensor[l0CPingPongFlag], + l0ATensor[l0ABPingPongFlag], + l0BTensor[l0ABPingPongFlag], + mL0Align, + nL1Actual, + kL0Actual, + initMmad); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2U); + l0ABPingPongFlag = 1U - l0ABPingPongFlag; + } + l1PPingPongFlag = 1U - l1PPingPongFlag; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + MatrixCoord gmCTileCoord{mL1Idx * L0TileShape::M, L0TileShape::N * nL1Idx}; + LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL1Actual, nL1Actual)); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL1Actual, nL1Actual)); + copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C); + AscendC::SetFlag(l0CPingPongFlag); + l0CPingPongFlag = 1U - l0CPingPongFlag; + } + } + AscendC::SetFlag(EVENT_ID4); + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + uint32_t l1PPingPongFlag = 0; + uint32_t l0CPingPongFlag = 0; + uint32_t l0ABPingPongFlag = 0; + + uint32_t l1MDynamic = 0; + uint32_t l1NDynamic = 0; + uint32_t l1KDynamic = 0; + + uint32_t blockStartOffset = 0; + uint32_t maxKVStackLen = 0; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Block + +#endif // GEMM_BLOCK_MMAD_PV_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv_decode.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv_decode.hpp new file mode 100644 index 0000000000..d2918b1418 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv_decode.hpp @@ -0,0 +1,286 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_BLOCK_MMAD_PV_DECODE_HPP +#define GEMM_BLOCK_MMAD_PV_DECODE_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/coord.hpp" +#include "../../../attn_infra/arch/cross_core_sync.hpp" +#include "../../../attn_infra/gemm/dispatch_policy.hpp" +#include "../../../attn_infra/gemm/helper.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace NpuArch::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIPVDecode, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIPVDecode; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t LOAB_BLOCK = 1; + static constexpr uint32_t COORD_DIM0 = 0; + static constexpr uint32_t COORD_DIM1 = 1; + static constexpr uint32_t COORD_DIM2 = 2; + + // Check LayoutC + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + /// Construct + __aicore__ inline + BlockMmad() {} + + __aicore__ inline + void init(Arch::Resource &resource,uint32_t nDyn, uint32_t kDyn, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1BTensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) * STAGES); + for (uint32_t i = 0; i < STAGES; i++) { + l1ATensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + l1NDynamic = nDyn; + l1KDynamic = kDyn; + } + + /// Destructor + __aicore__ inline + ~BlockMmad() {} + + __aicore__ inline + void getBlockShape( + GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &nLoop, uint32_t &kvSeqlen, uint32_t &blockSize) + { + uint32_t nSplitSize = blockSize; + if (nowNIdx == nLoop - 1U) { + nSplitSize = kvSeqlen - nowNIdx * blockSize; + } + actualShape[COORD_DIM2] = nSplitSize; + } + + __aicore__ inline + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t &nowNIdx, + uint32_t &strideKV, uint32_t &blockSize) + { + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV; + } else { + kOffset = nowNIdx * blockSize * strideKV; + } + } + + __aicore__ inline + void operator()( + AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,GemmCoord actualOriShape, + uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV, + uint32_t blockStackNum, Arch::CrossCoreFlag softmaxFlag, uint32_t crossCoreSyncTrigger) + { + uint32_t rowNum = actualOriShape[COORD_DIM0]; + uint32_t embed = actualOriShape[COORD_DIM1]; + uint32_t stackSeqTile = actualOriShape[COORD_DIM2]; + GemmCoord actualShape{rowNum, embed, 0}; + uint32_t gBOffset = 0; + + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(stackSeqTile, embed); + AscendC::WaitFlag(EVENT_ID4); + for (uint32_t blockStackIdx = 0; (blockStackIdx < blockStackNum) && ((nIdx + blockStackIdx) < nLoop); + blockStackIdx++) { + uint32_t nowNIdx = nIdx + blockStackIdx; + getBlockShape(actualShape, nowNIdx, nLoop, kvSeqlen, blockSize); + getKVOffset(gBlockTable, gBOffset, nowNIdx, strideKV, blockSize); + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(actualShape.k(), actualShape.n())); + MatrixCoord l1BTileCoord{blockStackIdx * blockSize, 0}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)]; + copyGmToL1B(l1BTile, gB[gBOffset], layoutBInL1, layoutBTile); + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + if (crossCoreSyncTrigger) { + Arch::CrossCoreWaitFlag(softmaxFlag); + + } + + uint32_t mL1Loop = CeilDiv(rowNum, L1TileShape::M); + uint32_t kL1Loop = NpuArch::Detail::Alignment::CeilDiv(stackSeqTile, l1KDynamic); + uint32_t nL1Loop = CeilDiv(embed, L0TileShape::N); + + for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; nL1Idx++) { + uint32_t nL1Actual = (nL1Idx < nL1Loop - 1U) ? L0TileShape::N : (embed - nL1Idx * L0TileShape::N); + for (uint32_t mL1Idx = 0; mL1Idx < mL1Loop; mL1Idx++) { + uint32_t mL1Actual = (mL1Idx < mL1Loop - 1U) ? L1TileShape::M : (rowNum - mL1Idx * L1TileShape::M); + AscendC::WaitFlag(l0CPingPongFlag); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t kL1Actual = (kL1Idx < kL1Loop - 1U) ? l1KDynamic : (stackSeqTile - kL1Idx * l1KDynamic); + // load P + AscendC::WaitFlag(l1PPingPongFlag); + MatrixCoord gmATileCoord{mL1Idx * L1TileShape::M, kL1Idx * l1KDynamic}; + auto gmTileA = gA[layoutA.GetOffset(gmATileCoord)]; + auto layoutTileA = layoutA.GetTileLayout(MakeCoord(mL1Actual, kL1Actual)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(mL1Actual, kL1Actual); + copyGmToL1A(l1ATensor[l1PPingPongFlag], gmTileA, layoutAInL1, layoutTileA); + AscendC::SetFlag(l1PPingPongFlag); + + uint32_t kL0Loop = CeilDiv(kL1Actual, L0TileShape::K); + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = + (kL0Idx < kL0Loop - 1U) ? L0TileShape::K : (kL1Actual - kL0Idx * L0TileShape::K); + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mL1Actual, kL0Actual); + MatrixCoord l1ATileCoord{0, kL0Idx * L0TileShape::K}; + auto l1ATile = l1ATensor[l1PPingPongFlag][layoutAInL1.GetOffset(l1ATileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag); + if (kL0Idx == 0U) { + AscendC::WaitFlag(l1PPingPongFlag); + } + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1); + if (kL0Idx == kL0Loop - 1U) { + AscendC::SetFlag(l1PPingPongFlag); + } + + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, nL1Actual); + MatrixCoord l1BTileCoord{kL1Idx * l1KDynamic + kL0Idx * L0TileShape::K, L0TileShape::N * nL1Idx}; + auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag + 2U); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + bool initMmad = (kL1Idx == 0U) && (kL0Idx == 0U); + uint32_t mL0Align = (mL1Actual + BLOCK_SIZE - 1U) / BLOCK_SIZE * BLOCK_SIZE; + tileMmad(l0CTensor[l0CPingPongFlag], + l0ATensor[l0ABPingPongFlag], + l0BTensor[l0ABPingPongFlag], + mL0Align, + nL1Actual, + kL0Actual, + initMmad); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2U); + l0ABPingPongFlag = 1U - l0ABPingPongFlag; + } + l1PPingPongFlag = 1U - l1PPingPongFlag; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + MatrixCoord gmCTileCoord{mL1Idx * L0TileShape::M, L0TileShape::N * nL1Idx}; + LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL1Actual, nL1Actual)); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL1Actual, nL1Actual)); + copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C); + AscendC::SetFlag(l0CPingPongFlag); + l0CPingPongFlag = 1U - l0CPingPongFlag; + } + } + AscendC::SetFlag(EVENT_ID4); + } + +protected: + /// Data members + AscendC::LocalTensor l1ATensor[STAGES]; + AscendC::LocalTensor l1BTensor; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + uint32_t l1PPingPongFlag = 0; + uint32_t l0CPingPongFlag = 0; + uint32_t l0ABPingPongFlag = 0; + + uint32_t l1MDynamic = 0; + uint32_t l1NDynamic = 0; + uint32_t l1KDynamic = 0; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Block + +#endif // GEMM_BLOCK_MMAD_PV_DECODE_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk.hpp new file mode 100644 index 0000000000..6d172081ee --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk.hpp @@ -0,0 +1,348 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_BLOCK_MMAD_QK_HPP +#define GEMM_BLOCK_MMAD_QK_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/coord.hpp" +#include "../../../attn_infra/gemm/dispatch_policy.hpp" +#include "../../../attn_infra/gemm/helper.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp" +//////////////////////////////////////////////////////////////////// +namespace NpuArch::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIQK, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIQK; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t COORD_DIM0 = 0; + static constexpr uint32_t COORD_DIM1 = 1; + static constexpr uint32_t COORD_DIM2 = 2; + + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + __aicore__ inline + BlockMmad() {} + + __aicore__ inline + void init(Arch::Resource &resource, uint32_t nDyn, uint32_t kDyn, uint32_t KVStackLen = 512, uint32_t l1BufAddrStart = 0) + { + maxKVStackLen = KVStackLen; + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) + nDyn * kDyn * sizeof(ElementB) * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + l1NDynamic = nDyn; + l1KDynamic = kDyn; + } + + __aicore__ inline + ~BlockMmad() {} + + __aicore__ inline + void loadQGM( + AscendC::GlobalTensor gA, + LayoutA layoutA, + uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads) + { + uint32_t embed = layoutA.shape(1); + uint32_t rowNumRound = NpuArch::Detail::Alignment::RoundUp(rowNum, L1AAlignHelper::M_ALIGNED); + uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; + auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A( + l1ATensor, gA, + layoutAInL1, layoutSingleANd, + tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + } + + __aicore__ inline + void setBlockParam(uint32_t stackSeqTile, uint32_t &blockStart, uint32_t &blockEnd, uint32_t &curBlockTotalNum, uint32_t blockSize){ + if(stackSeqTile >= blockStart && blockSize != 0) { + blockEnd = ((stackSeqTile - blockStart) % blockSize == 0) ? blockSize : (stackSeqTile - blockStart) % blockSize; + curBlockTotalNum = (((stackSeqTile - blockStart) + blockSize - 1) / blockSize) + 1; + } else { + curBlockTotalNum = 1; + blockStart = stackSeqTile; + blockEnd = stackSeqTile + blockStartOffset; + } + } + + __aicore__ inline + void getBlockShape(GemmCoord &actualShape, uint32_t nL1Idx, uint32_t nL1Loop, uint32_t stackSeqTile) + { + uint32_t nSplitSize = l1NDynamic; + if (nL1Idx == nL1Loop - 1U) { + nSplitSize = stackSeqTile - nL1Idx * l1NDynamic; + } + actualShape[COORD_DIM1] = nSplitSize; + } + + __aicore__ inline + void getBlockShape(GemmCoord &actualShape, uint32_t& blockStartOffset, uint32_t& l1NResDynamic, uint32_t& kvL1Len, uint32_t& nowLen, uint32_t& blockSize) + { + nowLen = (blockSize - blockStartOffset < l1NResDynamic - kvL1Len) ? + blockSize - blockStartOffset : + l1NResDynamic - kvL1Len; + actualShape[COORD_DIM1] = nowLen; + } + + __aicore__ inline + void getKVOffset(uint32_t &kOffset, uint32_t nIdx, uint32_t nowNIdx, uint32_t strideKV) + { + kOffset = nIdx * maxKVStackLen * strideKV + nowNIdx * l1NDynamic * strideKV; + } + + __aicore__ inline + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t nowNIdx, + uint32_t startOffset, uint32_t strideKV, uint32_t blockSize) + { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + startOffset * strideKV; + } + + __aicore__ inline + void resetBlockStart(uint32_t kvStart, uint32_t pagedBlockSize) + { + blockStartOffset = kvStart * maxKVStackLen % pagedBlockSize; + } + + __aicore__ inline + void updateBlockOffset(uint32_t nowLen, uint32_t &curBlockIdx, uint32_t blockSize){ + if(blockStartOffset + nowLen == blockSize){ + blockStartOffset = 0; + curBlockIdx++; + } else{ + blockStartOffset += nowLen; + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualOriShape, + uint32_t nIdx, uint32_t nLoop, uint32_t blockSize, uint32_t strideKV) + { + uint32_t rowNum = actualOriShape[COORD_DIM0]; + uint32_t stackSeqTile = actualOriShape[COORD_DIM1]; + uint32_t embed = actualOriShape[COORD_DIM2]; + GemmCoord actualShape{rowNum, 0, embed}; + uint32_t gBOffset = 0; + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + uint32_t tileNNumPerBaseBlock = blockSize / l1NDynamic; + uint32_t nL1Loop = NpuArch::Detail::Alignment::CeilDiv(stackSeqTile, l1NDynamic); + uint32_t curBlockIdx = 0; + uint32_t blockStart = 0; + uint32_t blockEnd = 0; + uint32_t curBlockTotalNum = 0; + if constexpr (PAGED_CACHE_FLAG_){ + blockStart = blockSize - blockStartOffset; + setBlockParam(stackSeqTile, blockStart, blockEnd, curBlockTotalNum, blockSize); + } + for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; ++nL1Idx) { + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + if constexpr (PAGED_CACHE_FLAG_){ + uint32_t l1NResDynamic = (nL1Idx < (nL1Loop-1)) ? l1NDynamic : (stackSeqTile - nL1Idx * l1NDynamic); + layoutBInL1 = LayoutBInL1::template MakeLayout(embed, l1NResDynamic); + uint32_t kvL1Len = 0; + AscendC::WaitFlag(l1KvPingPongFlag); + while(kvL1Len < l1NResDynamic){ + uint32_t nowLen = 0; + uint32_t curBlockSize = (curBlockIdx < (curBlockTotalNum-1)) ? blockSize : blockEnd; + uint32_t nowNIdx = nIdx * maxKVStackLen / blockSize + curBlockIdx; + getBlockShape(actualShape, blockStartOffset, l1NResDynamic, kvL1Len, nowLen, curBlockSize); + getKVOffset(gBlockTable, gBOffset, nowNIdx, blockStartOffset, strideKV, blockSize); + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(embed, nowLen)); + MatrixCoord l1BTileCoord{0, kvL1Len}; + auto l1BTile = l1BTensor[l1KvPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)]; + copyGmToL1B(l1BTile, gB[gBOffset], layoutBInL1, layoutBTile); + kvL1Len += nowLen; + updateBlockOffset(nowLen, curBlockIdx, blockSize); + } + AscendC::SetFlag(l1KvPingPongFlag); + mActual = actualShape.m(); + kActual = actualShape.k(); + nActual = l1NResDynamic; + } else { + getBlockShape(actualShape, nL1Idx, nL1Loop, stackSeqTile); + getKVOffset(gBOffset, nIdx, nL1Idx, strideKV); + mActual = actualShape.m(); + kActual = actualShape.k(); + nActual = actualShape.n(); + layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB[gBOffset], layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag); + } + uint32_t mL0Loop = NpuArch::Detail::Alignment::CeilDiv(mActual, L0TileShape::M); + uint32_t kL0Loop = NpuArch::Detail::Alignment::CeilDiv(kActual, L0TileShape::K); + for (uint32_t mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1U) ? L0TileShape::M : (mActual - mL0Idx * L0TileShape::M); + AscendC::WaitFlag(l0CPingPongFlag); + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1U) ? L0TileShape::K : (kActual - kL0Idx * L0TileShape::K); + + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mL0Actual, kL0Actual); + MatrixCoord l1ATileCoord{mL0Idx * L0TileShape::M, kL0Idx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ATileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag); + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1); + + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, nActual); + MatrixCoord l1BTileCoord{kL0Idx * L0TileShape::K, 0}; + auto l1BTile = l1BTensor[l1KvPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)]; + if ((mL0Idx == 0U) && (kL0Idx == 0U)) { + AscendC::WaitFlag(l1KvPingPongFlag); + } + AscendC::WaitFlag(l0ABPingPongFlag + 2U); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1); + if ((mL0Idx == mL0Loop - 1U) && (kL0Idx == kL0Loop - 1U)) { + AscendC::SetFlag(l1KvPingPongFlag); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + bool initMmad = (kL0Idx == 0U); + uint32_t mL0Align = (mL0Actual + BLOCK_SIZE - 1U) / BLOCK_SIZE * BLOCK_SIZE; + tileMmad(l0CTensor[l0CPingPongFlag], + l0ATensor[l0ABPingPongFlag], + l0BTensor[l0ABPingPongFlag], + mL0Align, + nActual, + kL0Actual, + initMmad); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2U); + l0ABPingPongFlag = 1U - l0ABPingPongFlag; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + MatrixCoord gmCTileCoord{mL0Idx * L0TileShape::M, nL1Idx * l1NDynamic}; + LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL0Actual, nActual)); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL0Actual, nActual)); + copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C); + AscendC::SetFlag(l0CPingPongFlag); + l0CPingPongFlag = 1U - l0CPingPongFlag; + } + l1KvPingPongFlag = 1U - l1KvPingPongFlag; + } + } +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + uint32_t l1KvPingPongFlag = 0; + uint32_t l0CPingPongFlag = 0; + uint32_t l0ABPingPongFlag = 0; + + uint32_t l1MDynamic = 0; + uint32_t l1NDynamic = 0; + uint32_t l1KDynamic = 0; + + uint32_t blockStartOffset = 0; + uint32_t maxKVStackLen = 0; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Block + +#endif // GEMM_BLOCK_MMAD_QK_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk_decode.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk_decode.hpp new file mode 100644 index 0000000000..e6cb829de6 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk_decode.hpp @@ -0,0 +1,290 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_BLOCK_MMAD_QK_DECODE_HPP +#define GEMM_BLOCK_MMAD_QK_DECODE_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/resource.hpp" +#include "../../../attn_infra/coord.hpp" +#include "../../../attn_infra/gemm/dispatch_policy.hpp" +#include "../../../attn_infra/gemm/helper.hpp" +#include "../../../attn_infra/gemm_coord.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp" + +//////////////////////////////////////////////////////////////////// + +namespace NpuArch::Gemm::Block { +//////////////////////////////////////////////////////////////////// + +template < + bool PAGED_CACHE_FLAG_, + bool ENABLE_UNIT_FLAG_, + class L1TileShape_, + class L0TileShape_, + class AType_, + class BType_, + class CType_, + class BiasType_, + class TileCopy_, + class TileMmad_> +struct BlockMmad< + MmadAtlasA2FAIQKDecode, + L1TileShape_, + L0TileShape_, + AType_, + BType_, + CType_, + BiasType_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadAtlasA2FAIQKDecode; + using ArchTag = typename DispatchPolicy::ArchTag; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = typename AType_::Element; + using LayoutA = typename AType_::Layout; + using ElementB = typename BType_::Element; + using LayoutB = typename BType_::Layout; + using ElementC = typename CType_::Element; + using LayoutC = typename CType_::Layout; + using TileMmad = TileMmad_; + using CopyGmToL1A = typename TileCopy_::CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B; + using CopyL0CToGm = typename TileCopy_::CopyL0CToGm; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc; + using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc; + using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst; + using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst; + using LayoutCInL0 = layout::zN; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + static constexpr uint32_t STAGES = DispatchPolicy::STAGES; + static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA); + static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB); + static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE; + static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE; + static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE; + static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES; + static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES; + static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES; + static constexpr uint32_t BLOCK_SIZE = 16; + static constexpr uint32_t EMBED_SPLIT_SIZE = 128; + static constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4; + static constexpr uint32_t KV_BASE_BLOCK = 512; + static constexpr uint32_t KV_SPLIT_SIZE = 128; + static constexpr uint32_t COORD_DIM0 = 0; + static constexpr uint32_t COORD_DIM1 = 1; + static constexpr uint32_t COORD_DIM2 = 2; + + static_assert(std::is_same_v, "LayoutC only support RowMajor yet!"); + + __aicore__ inline + BlockMmad() {} + + __aicore__ inline + void init(Arch::Resource &resource, uint32_t nDyn, uint32_t kDyn, uint32_t l1BufAddrStart = 0) + { + // Allocate L1 memory space + l1ATensor = resource.l1Buf.template GetBufferByByte(l1BufAddrStart); + for (uint32_t i = 0; i < STAGES; i++) { + l1BTensor[i] = resource.l1Buf.template GetBufferByByte(l1BufAddrStart + + L1TileShape::M * kDyn * sizeof(ElementA) + nDyn * kDyn * sizeof(ElementB) * i); + l0ATensor[i] = resource.l0ABuf.template GetBufferByByte(L0A_PINGPONG_BUF_SIZE * i); + l0BTensor[i] = resource.l0BBuf.template GetBufferByByte(L0B_PINGPONG_BUF_SIZE * i); + l0CTensor[i] = resource.l0CBuf.template GetBufferByByte(L0C_PINGPONG_BUF_SIZE * i); + } + l1NDynamic = nDyn; + l1KDynamic = kDyn; + } + + __aicore__ inline + ~BlockMmad() {} + + __aicore__ inline + void loadQGM( + AscendC::GlobalTensor gA, + LayoutA layoutA, + uint32_t rowNum, uint32_t &singleGroupHeads, uint32_t &qHeads, + uint32_t kvNBlockSizeParam) + { + uint32_t embed = layoutA.shape(1); + uint32_t rowNumRound = RoundUp(rowNum, L1AAlignHelper::M_ALIGNED); + uint32_t tokenNumPerGroup = rowNum / singleGroupHeads; + auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed)); + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(rowNum, embed); + copyGmToL1A( + l1ATensor, gA, + layoutAInL1, layoutSingleANd, + tokenNumPerGroup, qHeads * embed, tokenNumPerGroup, BLOCK_SIZE, rowNumRound); + AscendC::SetFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3); + + l1ATotalRowNum = rowNum; + l1ARowNumPerKvHead = rowNum / kvNBlockSizeParam; + l1AEmbedDim = embed; + l1AKvNBlockSize = kvNBlockSizeParam; + } + + __aicore__ inline + void getBlockShape(GemmCoord &actualShape, uint32_t nL1Idx, uint32_t nL1Loop, uint32_t stackSeqTile) + { + uint32_t nSplitSize = l1NDynamic; + if (nL1Idx == nL1Loop - 1U) { + nSplitSize = stackSeqTile - nL1Idx * l1NDynamic; + } + actualShape[COORD_DIM1] = nSplitSize; + } + + __aicore__ inline + void getKVOffset(AscendC::GlobalTensor &gBlockTable, uint32_t &kOffset, uint32_t nowNIdx, uint32_t nL1Idx, + uint32_t strideKV, uint32_t blockSize) + { + if constexpr (PAGED_CACHE_FLAG_) { + uint32_t blockTableId = gBlockTable.GetValue(nowNIdx); + kOffset = blockTableId * blockSize * strideKV + nL1Idx * l1NDynamic * strideKV; + } else { + kOffset = nowNIdx * blockSize * strideKV + nL1Idx * l1NDynamic * strideKV; + } + } + + __aicore__ inline + void operator()(AscendC::GlobalTensor gA, + AscendC::GlobalTensor gB, + AscendC::GlobalTensor gC, + AscendC::GlobalTensor gBlockTable, + LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualOriShape, + uint32_t nIdx, uint32_t nLoop, uint32_t blockSize, uint32_t strideKV, + uint32_t kvNIncreIdx) + { + uint32_t rowNum = actualOriShape[COORD_DIM0]; + uint32_t stackSeqTile = actualOriShape[COORD_DIM1]; + uint32_t embed = actualOriShape[COORD_DIM2]; + + GemmCoord actualShape{rowNum, 0, embed}; + uint32_t gBOffset = 0; + + LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout(l1ATotalRowNum, l1AEmbedDim); + + uint32_t kvNRowOffset = kvNIncreIdx * l1ARowNumPerKvHead; + + uint32_t tileNNumPerBaseBlock = blockSize / l1NDynamic; + uint32_t nL1Loop = NpuArch::Detail::Alignment::CeilDiv(stackSeqTile, l1NDynamic); + for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; ++nL1Idx) { + uint32_t nowNIdx = nIdx + nL1Idx / tileNNumPerBaseBlock; + getBlockShape(actualShape, nL1Idx, nL1Loop, stackSeqTile); + getKVOffset(gBlockTable, gBOffset, nowNIdx, nL1Idx % tileNNumPerBaseBlock, strideKV, blockSize); + uint32_t mActual = actualShape.m(); + uint32_t kActual = actualShape.k(); + uint32_t nActual = actualShape.n(); + LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout(kActual, nActual); + + auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, nActual)); + AscendC::WaitFlag(l1KvPingPongFlag); + copyGmToL1B(l1BTensor[l1KvPingPongFlag], gB[gBOffset], layoutBInL1, layoutBTile); + AscendC::SetFlag(l1KvPingPongFlag); + + uint32_t mL0Loop = CeilDiv(mActual, L0TileShape::M); + uint32_t kL0Loop = CeilDiv(kActual, L0TileShape::K); + for (uint32_t mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1U) ? L0TileShape::M : (mActual - mL0Idx * L0TileShape::M); + AscendC::WaitFlag(l0CPingPongFlag); + for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1U) ? L0TileShape::K : (kActual - kL0Idx * L0TileShape::K); + + LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout(mL0Actual, kL0Actual); + MatrixCoord l1ATileCoord{mL0Idx * L0TileShape::M + kvNRowOffset, kL0Idx * L0TileShape::K}; + auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ATileCoord)]; + + AscendC::WaitFlag(l0ABPingPongFlag); + copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1); + + LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout(kL0Actual, nActual); + MatrixCoord l1BTileCoord{kL0Idx * L0TileShape::K, 0}; + auto l1BTile = l1BTensor[l1KvPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)]; + if ((mL0Idx == 0U) && (kL0Idx == 0U)) { + AscendC::WaitFlag(l1KvPingPongFlag); + } + AscendC::WaitFlag(l0ABPingPongFlag + 2U); + copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1); + if ((mL0Idx == mL0Loop - 1U) && (kL0Idx == kL0Loop - 1U)) { + AscendC::SetFlag(l1KvPingPongFlag); + } + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + bool initMmad = (kL0Idx == 0U); + uint32_t mL0Align = (mL0Actual + BLOCK_SIZE - 1U) / BLOCK_SIZE * BLOCK_SIZE; + tileMmad(l0CTensor[l0CPingPongFlag], + l0ATensor[l0ABPingPongFlag], + l0BTensor[l0ABPingPongFlag], + mL0Align, + nActual, + kL0Actual, + initMmad); + AscendC::SetFlag(l0ABPingPongFlag); + AscendC::SetFlag(l0ABPingPongFlag + 2U); + l0ABPingPongFlag = 1U - l0ABPingPongFlag; + } + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + MatrixCoord gmCTileCoord{mL0Idx * L0TileShape::M, nL1Idx * l1NDynamic}; + LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL0Actual, nActual)); + auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL0Actual, nActual)); + copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C); + AscendC::SetFlag(l0CPingPongFlag); + l0CPingPongFlag = 1U - l0CPingPongFlag; + } + l1KvPingPongFlag = 1U - l1KvPingPongFlag; + } + } +protected: + /// Data members + AscendC::LocalTensor l1ATensor; + AscendC::LocalTensor l1BTensor[STAGES]; + AscendC::LocalTensor l0ATensor[STAGES]; + AscendC::LocalTensor l0BTensor[STAGES]; + AscendC::LocalTensor l0CTensor[STAGES]; + + TileMmad tileMmad; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToGm copyL0CToGm; + + uint32_t l1KvPingPongFlag = 0; + uint32_t l0CPingPongFlag = 0; + uint32_t l0ABPingPongFlag = 0; + + uint32_t l1MDynamic = 0; + uint32_t l1NDynamic = 0; + uint32_t l1KDynamic = 0; + + uint32_t l1ATotalRowNum = 0; + uint32_t l1ARowNumPerKvHead = 0; + uint32_t l1AEmbedDim = 0; + uint32_t l1AKvNBlockSize = 1; +}; + +//////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Block + +#endif // GEMM_BLOCK_MMAD_QK_DECODE_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/dispatch_policy.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/dispatch_policy.hpp new file mode 100644 index 0000000000..d7ec7c6cc9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/dispatch_policy.hpp @@ -0,0 +1,76 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_DISPATCH_POLICY_HPP +#define GEMM_DISPATCH_POLICY_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/arch/arch.hpp" + +namespace NpuArch::Gemm +{ + +// Block Mmad Policies + +template +struct MmadAtlasA2Base { + using ArchTag = Arch::AtlasA2; + static constexpr uint32_t ASYNC = ASYNC_; +}; + +using MmadAtlasA2 = MmadAtlasA2Base; + +template +struct MmadAtlasA2FAIQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +// Dispatch policy for decoding scenario (pagedCacheFlag == true && qSeqlen == 1) +template +struct MmadAtlasA2FAIQKDecode : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAIPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +// Dispatch policy for PV decoding scenario (pagedCacheFlag == true && qSeqlen == 1) +template +struct MmadAtlasA2FAIPVDecode : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +} // namespace NpuArch::Gemm + +#endif // GEMM_DISPATCH_POLICY_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/gemm_type.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/gemm_type.hpp new file mode 100644 index 0000000000..95fd93c362 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/gemm_type.hpp @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_GEMM_TYPE_HPP +#define GEMM_GEMM_TYPE_HPP + +#include "../../attn_infra/base_defs.hpp" + +namespace NpuArch::Gemm +{ +template +struct GemmType +{ + using Element = Element_; + using Layout = Layout_; + static constexpr AscendC::TPosition POSITION = POSITION_; +}; + +} // namespace NpuArch::Gemm + +#endif // GEMM_GEMM_TYPE_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/helper.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/helper.hpp new file mode 100644 index 0000000000..0ca3a3c48e --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/helper.hpp @@ -0,0 +1,255 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_HELPER_HPP +#define GEMM_HELPER_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/layout/layout.hpp" +#include "../../attn_infra/gemm/gemm_type.hpp" + +namespace NpuArch::Gemm::helper +{ + +template +struct L1AlignHelper { + static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct ElementAccumulatorSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported element accumulator selector, can not find the specialization."); +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template +struct L1ATypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BiasTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BiasTypeSelector { + using GMBiasType = void; + using L1BiasType = void; + using L0BiasType = void; +}; + +template +struct L1BiasTypeSelector, ElementAccumulator> { + using GMBiasType = Gemm::GemmType; + using L1BiasType = Gemm::GemmType; + using L0BiasType = Gemm::GemmType; +}; + +/////////////////////////////////////// +// new add +template<> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1AndL0TypeSelectorGemm{ + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template<> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType>{ + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; +/////////////////////////////////////// +} // namespace NpuArch::Gemm::helper + +#endif // GEMM_HELPER_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_l1.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_l1.hpp new file mode 100644 index 0000000000..7281d87bb9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_l1.hpp @@ -0,0 +1,1067 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_GM_TO_L1_HPP +#define GEMM_TILE_COPY_GM_TO_L1_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp" + +namespace NpuArch::Gemm::Tile { + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void +> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void +> +struct CopyGmToL1IntervalDataCopy { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +//////////////////////////////////////// +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +/// Partial specialization for AtlasA2, half, RowMajor in and zN out. +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1IntervalDataCopy() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + NpuArch::Detail::Alignment::CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingRowMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1IntervalDataCopy() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, ColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1IntervalDataCopy() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + NpuArch::Detail::Alignment::CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template<> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1IntervalDataCopy() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0 + ); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(2)], dataCopyParams); + } + } +}; + +/// new add gemm +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(0); + uint32_t ndNum = layoutSrc.shape(0) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(0) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(1); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; //` + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(0) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(1)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(1) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(0); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; +//////////////////////////////////////// + +/////////////////////////////////////// +/// new add gemv, VectorLayout -> zN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + intriParams.nValue = 1; + intriParams.srcDValue = layoutSrc.shape(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + + + +/////////////////////////////////////// +/// new add gemv, ColumnMajor -> nN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Methods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +/// Partial specialization for zN in and zN out. +template < + class ArchTag, + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(1)); + uint32_t blockLen = NpuArch::Detail::Alignment::RoundUp(layoutSrc.orgShape(0)); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(3) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(3); + uint64_t srcOffset = i * layoutSrc.stride(3); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for nZ in and nZ out. +template < + class ArchTag, + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0)); + uint32_t blockLen = NpuArch::Detail::Alignment::RoundUp(layoutSrc.orgShape(1)); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(1) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(1); + uint64_t srcOffset = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element +> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = NpuArch::Detail::Alignment::CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride + ); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + } +}; + +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + // Mehtods + + __aicore__ inline + CopyGmToL1() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = layoutDst.shape(0) / ELE_NUM_PER_C0; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_GM_TO_L1_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_ub.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_ub.hpp new file mode 100644 index 0000000000..3a819a5309 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_gm_to_ub.hpp @@ -0,0 +1,22 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_GM_TO_UB_HPP +#define GEMM_TILE_COPY_GM_TO_UB_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp" +namespace NpuArch::Gemm::Tile { + + +} // NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_GM_TO_UB_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l0c_to_gm.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..660f00e5c6 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l0c_to_gm.hpp @@ -0,0 +1,209 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_L0C_TO_GM_HPP +#define GEMM_TILE_COPY_L0C_TO_GM_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" +namespace NpuArch::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT +> +struct CopyL0CToGmQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +// CopyL0CToGm cast fp32 to fp16 +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + float, half, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToGm cast fp32 to bf16 +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToGm output fp32 +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + float, float, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm output int32 +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + int32_t, int32_t, + ScaleGranularity::NO_QUANT +> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast int32_t to fp16 +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_TENSOR +> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + NpuArch::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_CHANNEL +> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false +> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_ +> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> +{ + using ArchTag = NpuArch::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = NpuArch::layout::zN; + using LayoutDst = NpuArch::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + __aicore__ inline + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_ +> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> +{ + using ArchTag = NpuArch::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = NpuArch::layout::zN; + using LayoutDst = NpuArch::layout::zN; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + __aicore__ inline + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void +> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_L0C_TO_GM_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_bt.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_bt.hpp new file mode 100644 index 0000000000..b3dbb5691f --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_bt.hpp @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_L1_TO_BT_HPP +#define GEMM_TILE_COPY_L1_TO_BT_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" + + +namespace NpuArch::Gemm::Tile { + + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_L1_TO_BT_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0a.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..b58e8acce6 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0a.hpp @@ -0,0 +1,356 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_L1_TO_L0A_HPP +#define GEMM_TILE_COPY_L1_TO_L0A_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp" + + +namespace NpuArch::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void +> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0A, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + __aicore__ inline + CopyL1ToL0A(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyL1ToL0A(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0)));; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for(uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0)); i++){ + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, NpuArch::Gemm::GemmType>{ + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyL1ToL0A(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0))); + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + for(uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(0)); i++){ + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1) * 2], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, NpuArch::Gemm::GemmType>{ + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyL1ToL0A(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; +////////////////////////////////////////// + +/// Partial specialization for zN in and zZ out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0A() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and zZ out. +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0A() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + __aicore__ inline + CopyL1ToL0A() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for int8_t, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0A() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +/// Partial specialization for float, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0A() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = NpuArch::Detail::Alignment::RoundUp(l1M); + uint16_t l1KAlign = NpuArch::Detail::Alignment::RoundUp(l1K); + uint16_t l0MAlign = NpuArch::Detail::Alignment::RoundUp(l0M); + uint16_t l0KAlign = NpuArch::Detail::Alignment::RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_L1_TO_L0A_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0b.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..1e3c1f0d29 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_l1_to_l0b.hpp @@ -0,0 +1,487 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_L1_TO_L0B_HPP +#define GEMM_TILE_COPY_L1_TO_L0B_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/layout/layout.hpp" +#include "../../../attn_infra/gemm/gemm_type.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp" + +namespace NpuArch::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void +> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyL1ToL0B(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for(uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++){ // K N + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + __aicore__ inline + CopyL1ToL0B(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))) - 1; + for(uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++){ // K N + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + __aicore__ inline + CopyL1ToL0B(){} + + __aicore__ inline + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc + ){ + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; +///////////////////////////////////////////// + +//////////////////////////////////////////// +/// new add gemv +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1)); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(3); i++) + { + AscendC::LoadData(dstTensor[i * layoutDst.stride(3)], srcTensor[i * layoutSrc.stride(3)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType> +{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = layoutDst.shape(1) * layoutDst.shape(3); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + }; +}; + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + using Element = float; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)) - 1; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv<2 * ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++) + { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3) * 2], + srcTensor[i * layoutSrc.stride(3)], + loadDataParams); + } + }; +}; + +template +struct CopyL1ToL0B, NpuArch::Gemm::GemmType>{ + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + using Element = int8_t; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1)); i++) + { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3)], + srcTensor[i * layoutSrc.stride(3) * 2], + loadDataParams); + } + } +}; +//////////////////////////////////////////// + +/// Partial specialization for int8_t, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1N = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0K = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0N = layoutDst.shape(2) * layoutDst.shape(3); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = NpuArch::Detail::Alignment::RoundUp(l1K); + uint16_t l1NAlign = NpuArch::Detail::Alignment::RoundUp(l1N); + uint16_t l0KAlign = NpuArch::Detail::Alignment::RoundUp(l0K); + uint16_t l0NAlign = NpuArch::Detail::Alignment::RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for zN in and nZ out. +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = + static_cast(NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < NpuArch::Detail::Alignment::CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + + // Methods + + __aicore__ inline + CopyL1ToL0B() {}; + + __aicore__ inline + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + if (layoutSrc.shape(3) == layoutDst.shape(3)) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1) * layoutDst.shape(3)); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } + + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_L1_TO_L0B_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_ub_to_gm.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_ub_to_gm.hpp new file mode 100644 index 0000000000..093b18cbf6 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/copy_ub_to_gm.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_COPY_UB_TO_GM_HPP +#define GEMM_TILE_COPY_UB_TO_GM_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/arch/arch.hpp" +#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp" +namespace NpuArch::Gemm::Tile { + +} // NpuArch::Gemm::Tile + +#endif // GEMM_TILE_COPY_UB_TO_GM_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy.hpp new file mode 100644 index 0000000000..7be077ecf9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy.hpp @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_TILE_COPY_HPP +#define GEMM_TILE_TILE_COPY_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_gm_to_l1.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_l0c_to_gm.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0a.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0b.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_l1_to_bt.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_gm_to_ub.hpp" +#include "../../../attn_infra/gemm/tile_common/copy_ub_to_gm.hpp" +#include "../../../attn_infra/gemm/helper.hpp" + + +namespace NpuArch::Gemm::Tile { + +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void +> +struct TileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A< + ArchTag, typename helper::L1ATypeSelector::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B< + ArchTag, typename helper::L1BTypeSelector::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t, + void, + Gemm::Tile::CopyGmToL1>; +}; + +////////////////////////////// +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_TILE_COPY_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy_tla.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy_tla.hpp new file mode 100644 index 0000000000..04256fb3a9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy_tla.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_TILE_COPY_TLA_HPP +#define GEMM_TILE_TILE_COPY_TLA_HPP + +#include "../../../attn_infra/base_defs.hpp" + +namespace NpuArch::Gemm::Tile { + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_TILE_COPY_TLA_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_mmad.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_mmad.hpp new file mode 100644 index 0000000000..61a71965a9 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_mmad.hpp @@ -0,0 +1,104 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef GEMM_TILE_TILE_MMAD_HPP +#define GEMM_TILE_TILE_MMAD_HPP + +#include "../../../attn_infra/base_defs.hpp" +#include "../../../attn_infra/gemm/helper.hpp" +namespace NpuArch::Gemm::Tile { + +/////////////////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag_, + /// GemmType for A matrix operand + class AType_, + /// GemmType type for B matrix operand + class BType_, + /// GemmType type for Bias operand + class BiasType_ +> +struct TileMmad { + using ElementA = typename AType_::Element; + using ElementB = typename BType_::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + // Methods + + __aicore__ inline + TileMmad() {} + + __aicore__ inline + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + __aicore__ inline + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + AscendC::LocalTensor const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + l0BiasTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace NpuArch::Gemm::Tile + +#endif // GEMM_TILE_TILE_MMAD_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm_coord.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm_coord.hpp new file mode 100644 index 0000000000..a2dc28a305 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm_coord.hpp @@ -0,0 +1,163 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! + * \file gemm_coord.hpp + * \brief + */ + +#ifndef GEMM_COORD_HPP +#define GEMM_COORD_HPP + +#include "../attn_infra/coord.hpp" + +namespace NpuArch { + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + uint32_t M_ = 1, + /// Columns of matrix product + uint32_t N_ = 1, + /// Inner dimension of matrix product + uint32_t K_ = 1 +> +struct GemmShape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + static constexpr int64_t MN = M * N; + static constexpr int64_t MK = M * K; + static constexpr int64_t KN = N * K; + static constexpr int64_t MNK = M * N * K; + + static constexpr int64_t COUNT = MNK; + + /// Returns a Coord object + HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } + + HOST_DEVICE + static Coord<2> ToCoordMN() + { + return MakeCoord(M, N); + } + + HOST_DEVICE + static Coord<2> ToCoordMK() + { + return MakeCoord(M, K); + } + + HOST_DEVICE + static Coord<2> ToCoordKN() + { + return MakeCoord(K, N); + } +}; + +/// GemmCoord is a structure derived from Coord<3> that specifies a location within the +/// coordinate space of a Gemm problem. +struct GemmCoord : public Coord<3, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// Gemm M dimension - rows of the output C matrix + static constexpr int M_INDEX = 0; + + /// Gemm N dimension - columns of the output C matrix + static constexpr int N_INDEX = 1; + + /// Gemm K dimension - inner dimension of the Gemm problem + static constexpr int K_INDEX = 2; + + /// Default ctor + HOST_DEVICE + GemmCoord() {} + + /// Constructs from Coord<3> and a batch + HOST_DEVICE + GemmCoord(Coord<3, Index> const &coord) : Base(coord) {} + + /// Helper to construct from a K, N, M, batch variables + HOST_DEVICE + GemmCoord(Index m, Index n, Index k) : Base(MakeCoord(m, n, k)) {} + + /// Returns the Gemm M coordinate + HOST_DEVICE + Index const &m() const + { + return this->At(M_INDEX); + } + + /// Returns reference to the Gemm M coordinate + HOST_DEVICE + Index &m() + { + return this->At(M_INDEX); + } + + /// Returns the Gemm N coordinate + HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + + /// Returns reference to the Gemm N coordinate + HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + /// Returns the Gemm K coordinate + HOST_DEVICE + Index const &k() const + { + return this->At(K_INDEX); + } + + /// Returns reference to the Gemm K coordinate + HOST_DEVICE + Index &k() + { + return this->At(K_INDEX); + } + + HOST_DEVICE + auto GetCoordMN() const + { + return this->GetCoordByAxis(); + } + + HOST_DEVICE + auto GetCoordMK() const + { + return this->GetCoordByAxis(); + } + + HOST_DEVICE + auto GetCoordKN() const + { + return this->GetCoordByAxis(); + } +}; + +} // namespace NpuArch + +#endif // GEMM_COORD_HPP diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/layout.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/layout.hpp new file mode 100644 index 0000000000..d43b668c52 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/layout.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef LAYOUT_LAYOUT_HPP +#define LAYOUT_LAYOUT_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/layout/matrix.hpp" +#include "../../attn_infra/layout/vector.hpp" + +#endif // LAYOUT_LAYOUT_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/matrix.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/matrix.hpp new file mode 100644 index 0000000000..9308ce5cbd --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/matrix.hpp @@ -0,0 +1,1208 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef LAYOUT_MATRIX_HPP +#define LAYOUT_MATRIX_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/coord.hpp" +#include "../../attn_infra/detail/alignment.hpp" +#include "../../attn_infra/matrix_coord.hpp" + +namespace NpuArch::layout +{ + +/// Mapping function for row-major matrices +struct RowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + HOST_DEVICE + RowMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(cols), LongIndex(1))) {} + + /// Constructor + HOST_DEVICE + RowMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(ldm, LongIndex(1))) {} + + /// Ctor + HOST_DEVICE + RowMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + template + HOST_DEVICE + static RowMajor MakeLayoutInUb(MatrixCoord const &shape) + { + return RowMajor(shape.row(), shape.column(), NpuArch::Detail::Alignment::RoundUp(shape.column())); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) * stride_[0] + LongIndex(coord.column()); + } + + /// Returns the layout of a tile_common. + HOST_DEVICE + RowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return RowMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for col-major matrices +struct ColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + HOST_DEVICE + ColumnMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), LongIndex(rows))) {} + + /// Constructor + HOST_DEVICE + ColumnMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), ldm)) {} + + /// Ctor + HOST_DEVICE + ColumnMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) + LongIndex(coord.column()) * stride_[1]; + } + + /// Returns the layout of a tile_common. + HOST_DEVICE + ColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return ColumnMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal +struct nZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + HOST_DEVICE constexpr + nZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + HOST_DEVICE constexpr + nZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + HOST_DEVICE constexpr + static nZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = NpuArch::Detail::Alignment::RoundUp(orgRows); + Index colsRound = NpuArch::Detail::Alignment::RoundUp(orgCols); + return nZ(orgRows, + orgCols, + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + 1, + colsRound * ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile_common. + HOST_DEVICE + nZ GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), NpuArch::Detail::Alignment::CeilDiv(tileOriShape.row(), shape(0)), + shape(2), NpuArch::Detail::Alignment::CeilDiv(tileOriShape.column(), shape(2)) + ); + return nZ(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and col-major between fractal +struct zN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + HOST_DEVICE constexpr + zN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + HOST_DEVICE constexpr + zN(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + HOST_DEVICE constexpr + static zN MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = NpuArch::Detail::Alignment::RoundUp(orgRows); + Index colsRound = NpuArch::Detail::Alignment::RoundUp(orgCols); + return zN(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * ELE_NUM_PER_C0); + } + + HOST_DEVICE + static zN MakeLayoutInL0C(MatrixCoord const &shape) + { + return zN(shape.row(), + shape.column(), + C0_NUM_PER_FRACTAL, + NpuArch::Detail::Alignment::CeilDiv(shape.row()), + C0_NUM_PER_FRACTAL, + NpuArch::Detail::Alignment::CeilDiv(shape.column()), + C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL * C0_NUM_PER_FRACTAL, + 1, + NpuArch::Detail::Alignment::RoundUp(shape.row()) * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile_common. + HOST_DEVICE + zN GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), NpuArch::Detail::Alignment::CeilDiv(tileOriShape.row(), shape(0)), + shape(2), NpuArch::Detail::Alignment::CeilDiv(tileOriShape.column(), shape(2)) + ); + return zN(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and row-major between fractal +struct zZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + HOST_DEVICE constexpr + zZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {} + + /// Ctor + HOST_DEVICE constexpr + zZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + HOST_DEVICE constexpr + static zZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = static_cast(BYTE_PER_C0) / static_cast(sizeof(Element)); + constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = NpuArch::Detail::Alignment::RoundUp(orgRows); + Index colsRound = NpuArch::Detail::Alignment::RoundUp(orgCols); + return zZ(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + colsRound * C0_NUM_PER_FRACTAL, + 1, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding rowmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is row-major within blocks and also row-major between blocks. +struct PaddingRowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + HOST_DEVICE + PaddingRowMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) : + orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, NpuArch::Detail::Alignment::CeilDiv(orgRows, blockRows), + blockCols, NpuArch::Detail::Alignment::CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)blockCols, + (LongIndex)blockRows * (LongIndex)NpuArch::Detail::Alignment::RoundUp(orgCols, blockCols), + (LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols)) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + + (LongIndex)coord.column() / blockCols * stride_[3] + + (LongIndex)coord.row() % blockRows * stride_[0] + + (LongIndex)coord.column() % blockCols; + } + + HOST_DEVICE + PaddingRowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingRowMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding columnmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is column-major within blocks and also column-major between blocks. +struct PaddingColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + HOST_DEVICE + PaddingColumnMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) : + orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, NpuArch::Detail::Alignment::CeilDiv(orgRows, blockRows), + blockCols, NpuArch::Detail::Alignment::CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols, (LongIndex)blockRows, + (LongIndex)NpuArch::Detail::Alignment::RoundUp(orgRows, blockRows) * (LongIndex)blockCols)) {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + + (LongIndex)coord.column() / blockCols * stride_[3] + + (LongIndex)coord.row() % blockRows + + (LongIndex)coord.column() % blockCols * stride_[2]; + } + + HOST_DEVICE + PaddingColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingColumnMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/////////////////////// +// new add layout nN +// nN layout +struct nN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + HOST_DEVICE + nN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) { + } + + /// Ctor + HOST_DEVICE + nN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), shape_(shape), stride_(stride) {} + + /// Make the layout of a coordinate (row, column) + template + HOST_DEVICE static nN MakeLayout(Index orgRows, Index orgCols) { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + Index rowsRound = NpuArch::Detail::Alignment::RoundUp(orgRows); + Index colsRound = NpuArch::Detail::Alignment::RoundUp(orgCols); + return nN(orgRows, + orgCols, + + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + + 1, + ELE_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + HOST_DEVICE + LongIndex GetOffset(MatrixCoord const& coord) const { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + HOST_DEVICE + typename OrgShape::Index& orgShape(int idx) { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape& shape() { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index& shape(int idx) { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride& stride() { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index& stride(int idx) { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; +} // namespace NpuArch::layout + +#endif // LAYOUT_MATRIX_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/vector.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/vector.hpp new file mode 100644 index 0000000000..408cccb0a4 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/vector.hpp @@ -0,0 +1,133 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +#ifndef LAYOUT_VECTOR_HPP +#define LAYOUT_VECTOR_HPP + +#include "../../attn_infra/base_defs.hpp" +#include "../../attn_infra/coord.hpp" + +namespace NpuArch::layout +{ + +struct VectorLayout { +public: + /// Logical rank of tensor + static constexpr int RANK = 1; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Shape vector + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + + /// Logical coordinate + using TensorCoord = Coord; + +public: + // Methods + + HOST_DEVICE + VectorLayout(Index size = 0) : shape_(MakeCoord(size)), stride_(MakeCoord(LongIndex(1))) {} + + HOST_DEVICE + VectorLayout(Shape shape, Stride stride) : shape_(shape), stride_(stride) {} + + template + HOST_DEVICE + static VectorLayout MakeLayoutInUb(TensorCoord const &tileShape) + { + return VectorLayout{NpuArch::Detail::Alignment::RoundUp(tileShape[0])}; + } + + HOST_DEVICE + LongIndex GetOffset(TensorCoord const &coord) const + { + return stride_[0] * coord[0]; + } + + /// Returns the layout of a tile_common. + HOST_DEVICE + VectorLayout GetTileLayout(TensorCoord const &tileShape) const + { + return VectorLayout(tileShape, stride()); + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Stride data member + Shape shape_; + Stride stride_; +}; + +} // namespace NpuArch::layout + +#endif // LAYOUT_VECTOR_HPP \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/matrix_coord.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/matrix_coord.hpp new file mode 100644 index 0000000000..e3c3953308 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/matrix_coord.hpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! + * \file matrix_coord.hpp + * \brief + */ + +#ifndef MATRIX_COORD_HPP +#define MATRIX_COORD_HPP + +#include "../attn_infra/coord.hpp" + +namespace NpuArch { + +template < + uint32_t ROW_ = 1, + uint32_t COLUMN_ = 1 +> +struct MatrixShape { + static constexpr uint32_t ROW = ROW_; + static constexpr uint32_t COLUMN = COLUMN_; + + static constexpr int64_t COUNT = ROW * COLUMN; + + HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(ROW, COLUMN); + } +}; + +/// MatrixCoord wraps Coord<2, uint32_t> to provide a helper for accessing named dimensions. Classes +/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord. +struct MatrixCoord : public Coord<2, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Rows dimension + static constexpr uint32_t ROW_INDEX = 0; + + /// Columns dimension + static constexpr uint32_t COLUMN_INDEX = 1; + + /// Default ctor + HOST_DEVICE + MatrixCoord() {} + + /// Constructs from Coord<2> + HOST_DEVICE + MatrixCoord(Coord<2, Index> const &coord) : Base(coord) {} + + /// Helper to construct from a row and column + HOST_DEVICE + MatrixCoord(Index row, Index column) : Base(MakeCoord(row, column)) {} + + /// Helper to construct from a row and column, which are LongIndex based + HOST_DEVICE + MatrixCoord(LongIndex row, LongIndex column) : Base(MakeCoord(Index(row), Index(column))) {} + + /// Returns the row of the coordinate + HOST_DEVICE + Index const &row() const { return this->At(ROW_INDEX); } + + /// Returns the row of the coordinate + HOST_DEVICE + Index &row() { return this->At(ROW_INDEX); } + + /// Returns the column of the coordinate + HOST_DEVICE + Index const &column() const { return this->At(COLUMN_INDEX); } + + /// Returns the column of the coordinate + HOST_DEVICE + Index &column() { return this->At(COLUMN_INDEX); } + + /// Element-wise addition + HOST_DEVICE + MatrixCoord operator+(Base const &b) const + { + return MatrixCoord(Base::operator+(b)); + } + + /// In-place addition + HOST_DEVICE + MatrixCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +} // namespace NpuArch + +#endif \ No newline at end of file diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/flash_attention_regular.h b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/flash_attention_regular.h new file mode 100644 index 0000000000..d77c556487 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/flash_attention_regular.h @@ -0,0 +1,1116 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! +* \file flash_attention_regular.h +* \brief +*/ +#ifndef FLASH_ATTENTION_REGULAR_H +#define FLASH_ATTENTION_REGULAR_H + +#include "kernel_common.hpp" + +using namespace NpuArch; +using namespace KernelCommon; + +namespace SplitFuse { + template < + class BlockMmadQK, + class BlockMmadPV, + class EpilogueOnlineSoftmax, + class EpilogueRescaleO, + class EpilogueInitOut, + bool PAGED_CACHE_FLAG, + FaiKernel::MaskType MASK_TYPE = FaiKernel::MaskType::NO_MASK, + FaiKernel::inputLayout INPUT_LAYOUT = FaiKernel::inputLayout::BSND, + class CombineScale = void, + bool IS_FD = false, + bool ENABLE_FD_COMBINE = true, + bool USE_EPOCH_FD_BARRIER = false> + class FAInferKernel { + public: + using ArchTag = typename BlockMmadQK::ArchTag; + using L1TileShape = typename BlockMmadQK::L1TileShape; + using ElementQ = typename BlockMmadQK::ElementA; + using LayoutQ = typename BlockMmadQK::LayoutA; + using ElementK = typename BlockMmadQK::ElementB; + using LayoutK = typename BlockMmadQK::LayoutB; + using ElementS = typename BlockMmadQK::ElementC; + using LayoutS = typename BlockMmadQK::LayoutC; + + using ElementP = typename BlockMmadPV::ElementA; + using LayoutP = typename BlockMmadPV::LayoutA; + using ElementV = typename BlockMmadPV::ElementB; + using LayoutV = typename BlockMmadPV::LayoutB; + + using ElementMask = typename EpilogueOnlineSoftmax::ElementMask; + using LayoutMask = typename EpilogueOnlineSoftmax::LayoutMask; + using ElementSink = typename EpilogueOnlineSoftmax::ElementSink; + + using ElementO = typename EpilogueRescaleO::ElementOutput; + using LayoutO = typename EpilogueRescaleO::LayoutOutput; + + using ElementOTmp = typename EpilogueRescaleO::ElementInput; + using LayoutOTmp = typename EpilogueRescaleO::LayoutInput; + + using ElementLse = typename EpilogueRescaleO::ElementLse; + using LayoutLse = typename EpilogueRescaleO::LayoutLse; + + using ElementUpdate = typename EpilogueRescaleO::ElementUpdate; + using LayoutUpdate = typename EpilogueRescaleO::LayoutUpdate; + + static constexpr Epilogue::LseMode LSE_MODE = EpilogueRescaleO::LSE_MODE; + static constexpr Epilogue::SinkMode SINK_MODE = EpilogueOnlineSoftmax::SINK_MODE; + + struct GlobalTensorBundle { + AscendC::GlobalTensor& gQ; + AscendC::GlobalTensor& gK; + AscendC::GlobalTensor& gV; + AscendC::GlobalTensor& gPseShift; + AscendC::GlobalTensor& gMask; + AscendC::GlobalTensor& gBlockTable; + AscendC::GlobalTensor& gActualQseqlen; + AscendC::GlobalTensor& gActualKvseqlen; + AscendC::GlobalTensor& gO; + AscendC::GlobalTensor& gLse; + AscendC::GlobalTensor& gLseFD; + AscendC::GlobalTensor& gOFD; + AscendC::GlobalTensor& gS; + AscendC::GlobalTensor& gP; + AscendC::GlobalTensor& gOTmp; + AscendC::GlobalTensor& gOUpdate; + AscendC::GlobalTensor& gSink; + }; + + __aicore__ inline + FAInferKernel() {} + + __aicore__ inline + void operator()( + FAIKernelParams const ¶ms, + Arch::PtoTopology const &ptoTopology, + __gm__ int32_t *barrierState = nullptr) + { + resource.ptoTopology = ptoTopology; + __gm__ FAInferTilingData *fATilingData = reinterpret_cast<__gm__ FAInferTilingData *>(params.tiling); + mm1OutSize = fATilingData->mm1OutSize; + smOnlineOutSize = fATilingData->smOnlineOutSize; + mm2OutSize = fATilingData->mm2OutSize; + batch = fATilingData->batch; + qHeads = fATilingData->numHeads; + kvHeads = fATilingData->kvHeads; + embed = fATilingData->embeddingSize; + embedV = fATilingData->embeddingSizeV; + pagedBlockSize = fATilingData->blockSize; + maxNumBlocksPerBatch = fATilingData->maxNumBlocksPerBatch; + firstBatchTaskNum = fATilingData->firstBatchTaskNum; + totalTaskNum = fATilingData->totalTaskNum; + blockSize = fATilingData->blockSize; + maskType = fATilingData->maskType; + scaleValue = fATilingData->scaleValue; + sparseMode = fATilingData->sparseMode; + preToken = fATilingData->preToken; + nextToken = fATilingData->nextToken; + pseQ = fATilingData->pseQ; + pseKv = fATilingData->pseKv; + uint64_t Lsesize = 0; + uint64_t Losize = 0; + if constexpr (IS_FD) { + Lsesize = fATilingData->splitLseTotalSize; + Losize = fATilingData->splitOTotalSize; + } + + AscendC::GlobalTensor gQ; + gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q); + // K/V point directly to vLLM-compatible paged-cache storage. + __gm__ uint8_t* currentKey = reinterpret_cast<__gm__ uint8_t*>(params.k); + __gm__ uint8_t* currentValue = reinterpret_cast<__gm__ uint8_t*>(params.v); + AscendC::GlobalTensor gK; + gK.SetGlobalBuffer((__gm__ ElementK *)currentKey); + AscendC::GlobalTensor gV; + gV.SetGlobalBuffer((__gm__ ElementK *)currentValue); + AscendC::GlobalTensor gPseShift; + gPseShift.SetGlobalBuffer((__gm__ ElementQ *)params.pseShift); + AscendC::GlobalTensor gMask; + gMask.SetGlobalBuffer((__gm__ ElementMask *)params.mask); + AscendC::GlobalTensor gBlockTable; + gBlockTable.SetGlobalBuffer((__gm__ int32_t *)(params.blockTables)); + AscendC::GlobalTensor gActualQseqlen; + gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen); + AscendC::GlobalTensor gActualKvseqlen; + gActualKvseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualKvseqlen); + AscendC::GlobalTensor gO; + gO.SetGlobalBuffer((__gm__ ElementO *)params.o); + AscendC::GlobalTensor gLse; + gLse.SetGlobalBuffer((__gm__ ElementLse *)params.lse); + AscendC::GlobalTensor gLseFD; + AscendC::GlobalTensor gOFD; + if constexpr (IS_FD) { + gLseFD.SetGlobalBuffer((__gm__ ElementLse *)(params.workSpace)); + gOFD.SetGlobalBuffer((__gm__ ElementLse *)(params.workSpace + Lsesize)); + } + AscendC::GlobalTensor gS; + gS.SetGlobalBuffer((__gm__ ElementS *)(params.workSpace + Lsesize + Losize)); + AscendC::GlobalTensor gP; + gP.SetGlobalBuffer((__gm__ ElementP *)(params.workSpace + Lsesize + Losize + mm1OutSize)); + AscendC::GlobalTensor gOTmp; + gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)(params.workSpace + Lsesize + Losize + mm1OutSize + smOnlineOutSize)); + AscendC::GlobalTensor gOUpdate; + gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)(params.workSpace + Lsesize + Losize + + mm1OutSize + smOnlineOutSize + mm2OutSize)); + AscendC::GlobalTensor gSink; + gSink.SetGlobalBuffer((__gm__ ElementSink *)(params.sink)); + + GlobalTensorBundle globalTensors{ + gQ, gK, gV, gPseShift, gMask, gBlockTable, + gActualQseqlen, gActualKvseqlen, + gO, gLse, gLseFD, gOFD, + gS, gP, gOTmp, gOUpdate, gSink + }; + + uint32_t coreIdx = ptoTopology.logicalBlockIdx; + uint32_t coreNum = ptoTopology.logicalBlockNum; +#ifdef __DAV_C220_CUBE__ + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + AscendC::SetFlag(EVENT_ID7); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + AscendC::SetFlag(EVENT_ID7); + + uint32_t kDynNum = NpuArch::Detail::Alignment::RoundUp(embed, NUM_128); + kDynNum = kDynNum < NUM_256 ? NUM_256 : kDynNum; + uint32_t maxQKPL1Size = L1_MAX_SIZE - embedV * MAX_KV_STACK_LEN * sizeof(ElementV); + uint32_t maxQL1Size = Q_TILE_CEIL * kDynNum * sizeof(ElementQ); + uint32_t maxNDynNum = + ((maxQKPL1Size - maxQL1Size) / kDynNum / sizeof(ElementV) / DOUBLE_BUFFER) / NUM_32 * NUM_32; + + uint32_t nDynNum = maxNDynNum < L1_MAX_N_NUM ? maxNDynNum : L1_MAX_N_NUM; + nDynNum = L1_MAX_N_NUM % nDynNum != 0 ? + NpuArch::Detail::Alignment::RoundDown((nDynNum - 1), NUM_32) : nDynNum; + + uint32_t L1_QK_SIZE = BlockMmadQK::L1TileShape::M * kDynNum * sizeof(ElementQ); + blockMmadQK.init(resource, nDynNum, kDynNum, MAX_KV_STACK_LEN); + uint32_t kPVDynNum = nDynNum * kDynNum / BlockMmadPV::L1TileShape::M; + blockMmadPV.init(resource, nDynNum, kPVDynNum, MAX_KV_STACK_LEN, L1_QK_SIZE); +#endif +#ifdef __DAV_C220_VEC__ + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID6); + AscendC::SetFlag(EVENT_ID7); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID4); + AscendC::SetFlag(EVENT_ID5); + AscendC::SetFlag(EVENT_ID6); + + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID3); + + epilogueOnlineSoftmax.init(resource, scaleValue); + epilogueRescaleO.init(resource); + epilogueInitOut.init(resource); + + coreIdx = ptoTopology.logicalBlockIdx; +#endif + strideQ = static_cast(qHeads * embed); + strideO = static_cast(qHeads * embedV); + strideK = static_cast(kvHeads * embed); + strideV = static_cast(kvHeads * embedV); + embedRound = NpuArch::Detail::Alignment::RoundUp(embed, FaiKernel::BLOCK_SIZE); + embedRoundV = NpuArch::Detail::Alignment::RoundUp(embedV, FaiKernel::BLOCK_SIZE); + groupSize = qHeads / kvHeads; + + totalQTokens = static_cast(gActualQseqlen.GetValue(batch - 1)); + + if constexpr (IS_FD) { + // Keep every launched block alive for the selected final rendezvous, + // but do not let untiled blocks consume uninitialized coreInfo. + if (coreIdx < fATilingData->needCoreNum) { + uint32_t startBIdx = fATilingData->coreInfo.startBIdx[coreIdx]; + uint32_t startN1Idx = fATilingData->coreInfo.startN1Idx[coreIdx]; + uint32_t startS1Idx = fATilingData->coreInfo.startS1Idx[coreIdx]; + uint32_t startS2Idx = fATilingData->coreInfo.startS2Idx[coreIdx]; + uint32_t endBIdx = fATilingData->coreInfo.endBIdx[coreIdx]; + uint32_t endN1Idx = fATilingData->coreInfo.endN1Idx[coreIdx]; + uint32_t endS1Idx = fATilingData->coreInfo.endS1Idx[coreIdx]; + uint32_t endS2Idx = fATilingData->coreInfo.endS2Idx[coreIdx]; + uint64_t gmOffsetLseFD = fATilingData->coreInfo.firstSplitKVTaskLseOffset[coreIdx]; + uint64_t gmOffsetOFD = fATilingData->coreInfo.firstSplitKVTaskOOffset[coreIdx]; + + for (uint32_t BIdx = startBIdx; BIdx <= endBIdx; BIdx++) { + uint32_t qSeqlenCur = static_cast(gActualQseqlen.GetValue(BIdx)); + uint32_t kvSeqlenCur = static_cast(gActualKvseqlen.GetValue(BIdx)); + if constexpr(INPUT_LAYOUT == FaiKernel::inputLayout::TND) { + uint32_t prevQSeqlenSum = (BIdx == 0) ? + 0 : static_cast(gActualQseqlen.GetValue(BIdx - 1)); + qSeqlenCur = qSeqlenCur - prevQSeqlenSum; + if constexpr (!PAGED_CACHE_FLAG) { + uint32_t prevKvSeqlenSum = (BIdx == 0) ? + 0 : static_cast(gActualKvseqlen.GetValue(BIdx - 1)); + kvSeqlenCur = kvSeqlenCur - prevKvSeqlenSum; + } + } + + uint32_t curQNBlockTileTmp = GetQNBlockTile(qSeqlenCur, groupSize); + uint32_t qNBlockNumPerGroupTmp = NpuArch::Detail::Alignment::CeilDiv(groupSize, curQNBlockTileTmp); + uint32_t curQNBlockNumTmp = qNBlockNumPerGroupTmp * kvHeads; + uint32_t curQSBlockTileTmp = GetQSBlockTile(kvSeqlenCur); + uint32_t curQSBlockNumTmp = NpuArch::Detail::Alignment::CeilDiv(qSeqlenCur, curQSBlockTileTmp); + uint32_t curKSBlockNumTmp = NpuArch::Detail::Alignment::CeilDiv(kvSeqlenCur, GetKSBlockTile(kvSeqlenCur)); + + int32_t stN1IdxNow = (BIdx == startBIdx) ? startN1Idx : 0; + int32_t enN1IdxNow = (BIdx == endBIdx) ? endN1Idx : curQNBlockNumTmp - 1; + + for (int32_t n1Idx = stN1IdxNow; n1Idx <= enN1IdxNow; n1Idx++) { + int32_t stS1IdxNow = (BIdx == startBIdx && n1Idx == stN1IdxNow) ? startS1Idx : 0; + int32_t enS1IdxNow = (BIdx == endBIdx && n1Idx == enN1IdxNow) ? endS1Idx : curQSBlockNumTmp - 1; + + for (int32_t s1Idx = stS1IdxNow; s1Idx <= enS1IdxNow; s1Idx++) { + int32_t stS2IdxNow = (BIdx == startBIdx && n1Idx == stN1IdxNow && s1Idx == stS1IdxNow) ? startS2Idx : 0; + int32_t enS2IdxNow = (BIdx == endBIdx && n1Idx == enN1IdxNow && s1Idx == enS1IdxNow) ? endS2Idx : curKSBlockNumTmp; + + bool isSplitKV = (enS2IdxNow - stS2IdxNow) > 0 && (enS2IdxNow - stS2IdxNow) < static_cast(curKSBlockNumTmp); + + runMainLoop( + coreIdx, BIdx, n1Idx, s1Idx, + isSplitKV, stS2IdxNow, enS2IdxNow, + gmOffsetLseFD, gmOffsetOFD, + globalTensors, pseQ, pseKv + ); + + if (isSplitKV) { + uint32_t qSBlockSizeTmp = (s1Idx == static_cast(curQSBlockNumTmp - 1U)) ? + (qSeqlenCur - s1Idx * curQSBlockTileTmp) : curQSBlockTileTmp; + uint32_t qNBlockIdxCurGroupTmp = n1Idx % qNBlockNumPerGroupTmp; + uint32_t qNBlockSizeTmp = (qNBlockIdxCurGroupTmp == (qNBlockNumPerGroupTmp - 1U)) ? + (groupSize - qNBlockIdxCurGroupTmp * curQNBlockTileTmp) : curQNBlockTileTmp; + gmOffsetLseFD += qSBlockSizeTmp * qNBlockSizeTmp; + gmOffsetOFD += qSBlockSizeTmp * qNBlockSizeTmp * embedV; + } + } + } + } + } + } + else { + for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum; taskIdx += uint32_t(coreNum)) { + uint32_t curBatchTmp = 0; + uint32_t preTotalTaskNumTmp = 0; + uint32_t curTotalTaskNumTmp = firstBatchTaskNum; + + while (taskIdx >= curTotalTaskNumTmp && curBatchTmp < batch - 1) { + ++curBatchTmp; + preTotalTaskNumTmp = curTotalTaskNumTmp; + + uint32_t qSeqlenTmp = static_cast(gActualQseqlen.GetValue(curBatchTmp)); + uint32_t kvSeqlenTmp = static_cast(gActualKvseqlen.GetValue(curBatchTmp)); + if constexpr(INPUT_LAYOUT == FaiKernel::inputLayout::TND) { + uint32_t prevQSeqlenSumTmp = (curBatchTmp == 0) ? + 0 : static_cast(gActualQseqlen.GetValue(curBatchTmp - 1)); + qSeqlenTmp = qSeqlenTmp - prevQSeqlenSumTmp; + if constexpr (!PAGED_CACHE_FLAG) { + uint32_t prevKvSeqlenSumTmp = (curBatchTmp == 0) ? + 0 : static_cast(gActualKvseqlen.GetValue(curBatchTmp - 1)); + kvSeqlenTmp = kvSeqlenTmp - prevKvSeqlenSumTmp; + } + } + + uint32_t curQNBlockTileTmp = GetQNBlockTile(qSeqlenTmp, groupSize); + uint32_t qNBlockNumPerGroupTmp = NpuArch::Detail::Alignment::CeilDiv(groupSize, curQNBlockTileTmp); + uint32_t curQNBlockNumTmp = qNBlockNumPerGroupTmp * kvHeads; + uint32_t curQSBlockTileTmp = GetQSBlockTile(kvSeqlenTmp); + uint32_t curQSBlockNumTmp = NpuArch::Detail::Alignment::CeilDiv(qSeqlenTmp, curQSBlockTileTmp); + curTotalTaskNumTmp += curQNBlockNumTmp * curQSBlockNumTmp; + } + + uint32_t qSeqlenCur = static_cast(gActualQseqlen.GetValue(curBatchTmp)); + uint32_t kvSeqlenCur = static_cast(gActualKvseqlen.GetValue(curBatchTmp)); + if constexpr(INPUT_LAYOUT == FaiKernel::inputLayout::TND) { + uint32_t prevQSeqlenSumCur = (curBatchTmp == 0) ? + 0 : static_cast(gActualQseqlen.GetValue(curBatchTmp - 1)); + qSeqlenCur = qSeqlenCur - prevQSeqlenSumCur; + if constexpr (!PAGED_CACHE_FLAG) { + uint32_t prevKvSeqlenSumCur = (curBatchTmp == 0) ? + 0 : static_cast(gActualKvseqlen.GetValue(curBatchTmp - 1)); + kvSeqlenCur = kvSeqlenCur - prevKvSeqlenSumCur; + } + } + + uint32_t curQNBlockTileCur = GetQNBlockTile(qSeqlenCur, groupSize); + uint32_t qNBlockNumPerGroupCur = NpuArch::Detail::Alignment::CeilDiv(groupSize, curQNBlockTileCur); + uint32_t curQNBlockNumCur = qNBlockNumPerGroupCur * kvHeads; + + uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNumTmp; + uint32_t qSBlockIdxCur = taskIdxCurBatch / curQNBlockNumCur; + uint32_t qNBlockIdxCur = taskIdxCurBatch - qSBlockIdxCur * curQNBlockNumCur; + + runMainLoop( + coreIdx, curBatchTmp, qNBlockIdxCur, qSBlockIdxCur, + false, 0, 0, + 0, 0, + globalTensors, pseQ, pseKv + ); + } + } + +#ifdef __DAV_C220_CUBE__ + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID7); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID7); +#endif +#ifdef __DAV_C220_VEC__ + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID7); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID4); + AscendC::WaitFlag(EVENT_ID5); + AscendC::WaitFlag(EVENT_ID6); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID3); +#endif + AscendC::PipeBarrier(); + + if constexpr (IS_FD && ENABLE_FD_COMBINE) { + if constexpr (USE_EPOCH_FD_BARRIER) { +#ifdef __DAV_C220_VEC__ + epochBarrierArriveAndWait(barrierState, ptoTopology); +#endif + } else { + AscendC::SyncAll(); + } +#ifdef __DAV_C220_VEC__ + CombineScale combineScale; + combineScale.init(resource); + combineScale( + qHeads, + fATilingData->totalSplitNodeNum, + embedV, + &fATilingData->splitInfo, + gLseFD, + gOFD, + gO, + gActualQseqlen, + true + ); +#endif + } + } + + __aicore__ inline void epochBarrierArriveAndWait( + __gm__ int32_t *state, + Arch::PtoTopology const &topology) + { + constexpr uint32_t SLOT_STRIDE_WORDS = 512 / sizeof(int32_t); + uint32_t physicalCount = topology.logicalBlockNum * topology.lanesPerBlock; + uint32_t physicalId = + topology.logicalBlockIdx * topology.lanesPerBlock + topology.subBlockIdx; + __gm__ int32_t *mySlotRaw = state + physicalId * SLOT_STRIDE_WORDS; + volatile __gm__ int32_t *mySlot = + reinterpret_cast(mySlotRaw); + + // Split outputs reach GM through MTE3 and the PIPE_ALL above has + // already completed those stores. Only the scalar epoch slot uses + // the data cache, so flushing the entire cache here is unnecessary. + dcci(mySlotRaw, SINGLE_CACHE_LINE); + dsb(DSB_DDR); + uint32_t localEpoch = static_cast(*mySlot) + 1; + *mySlot = static_cast(localEpoch); + dcci(mySlotRaw, SINGLE_CACHE_LINE, CACHELINE_OUT); + dsb(DSB_DDR); + + uint64_t expectedSum = static_cast(physicalCount) * localEpoch; + while (true) { + for (uint32_t i = 0; i < physicalCount; ++i) { + dcci(state + i * SLOT_STRIDE_WORDS, SINGLE_CACHE_LINE); + } + dsb(DSB_DDR); + + uint64_t epochSum = 0; + for (uint32_t i = 0; i < physicalCount; ++i) { + volatile __gm__ int32_t *slot = + reinterpret_cast( + state + i * SLOT_STRIDE_WORDS); + epochSum += static_cast(*slot); + } + if (epochSum >= expectedSum) { + break; + } + } + pipe_barrier(PIPE_ALL); + } + + __aicore__ inline void runMainLoop( + uint32_t coreIdx, + uint32_t BIdx, + uint32_t qNBlockIdx, + uint32_t qSBlockIdx, + bool isSplitKV, + int32_t stS2IdxNow, + int32_t enS2IdxNow, + uint64_t gmOffsetLseFD, + uint64_t gmOffsetOFD, + GlobalTensorBundle& globalTensors, + int64_t pseQ, + int64_t pseKv + ) { + auto& gQ = globalTensors.gQ; + auto& gK = globalTensors.gK; + auto& gV = globalTensors.gV; + auto& gPseShift = globalTensors.gPseShift; + auto& gMask = globalTensors.gMask; + auto& gBlockTable = globalTensors.gBlockTable; + auto& gActualQseqlen = globalTensors.gActualQseqlen; + auto& gActualKvseqlen = globalTensors.gActualKvseqlen; + auto& gO = globalTensors.gO; + auto& gLse = globalTensors.gLse; + auto& gLseFD = globalTensors.gLseFD; + auto& gOFD = globalTensors.gOFD; + auto& gS = globalTensors.gS; + auto& gP = globalTensors.gP; + auto& gOTmp = globalTensors.gOTmp; + auto& gOUpdate = globalTensors.gOUpdate; + auto& gSink = globalTensors.gSink; + + uint32_t qSeqlen = static_cast(gActualQseqlen.GetValue(BIdx)); + uint32_t kvSeqlen = static_cast(gActualKvseqlen.GetValue(BIdx)); + uint32_t prevQSeqlenSum = 0; + uint32_t prevKvSeqlenSum = 0; + + if constexpr(INPUT_LAYOUT == FaiKernel::inputLayout::TND) { + prevQSeqlenSum = (BIdx == 0) ? + 0 : static_cast(gActualQseqlen.GetValue(BIdx - 1)); + qSeqlen = qSeqlen - prevQSeqlenSum; + if constexpr (!PAGED_CACHE_FLAG) { + prevKvSeqlenSum = (BIdx == 0) ? + 0 : static_cast(gActualKvseqlen.GetValue(BIdx - 1)); + kvSeqlen = kvSeqlen - prevKvSeqlenSum; + } + } + + uint64_t qBOffset = static_cast(prevQSeqlenSum) * strideQ; + uint64_t kBOffset = 0; + uint64_t vBOffset = 0; + uint64_t blockBOffset = 0; + if constexpr (!PAGED_CACHE_FLAG) { + kBOffset = static_cast(prevKvSeqlenSum) * strideK; + vBOffset = static_cast(prevKvSeqlenSum) * strideV; + } else { + blockBOffset = BIdx * static_cast(maxNumBlocksPerBatch); + } + uint64_t oBOffset = static_cast(prevQSeqlenSum) * strideO; + uint64_t lseBOffset = static_cast(prevQSeqlenSum) * qHeads; + + uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize); + uint32_t qNBlockNumPerGroup = NpuArch::Detail::Alignment::CeilDiv(groupSize, curQNBlockTile); + uint32_t curQSBlockTile = GetQSBlockTile(kvSeqlen); + uint32_t curQSBlockNum = NpuArch::Detail::Alignment::CeilDiv(qSeqlen, curQSBlockTile); + uint32_t curKSBlockNum = NpuArch::Detail::Alignment::CeilDiv(kvSeqlen, GetKSBlockTile(kvSeqlen)); + + uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup; + uint32_t kvNIdx = qNBlockIdx / qNBlockNumPerGroup; + uint32_t qNStartIdx = kvNIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile; + uint32_t lseTokenOffset = qSBlockIdx * curQSBlockTile * qHeads; + + uint64_t gmOffsetQ = qBOffset + + static_cast(qSBlockIdx * curQSBlockTile) * strideQ + + static_cast(qNStartIdx * embed); + uint64_t gmOffsetK = kBOffset + static_cast(kvNIdx * embed); + uint64_t gmOffsetV = vBOffset + static_cast(kvNIdx * embedV); + uint64_t gmOffsetO = oBOffset + + static_cast(qSBlockIdx * curQSBlockTile) * strideO + + static_cast(qNStartIdx * embedV); + uint64_t gmOffsetLse = lseBOffset + + static_cast(lseTokenOffset + qNStartIdx); + uint64_t gmOffsetSink = qNStartIdx; + + uint32_t qSBlockSize = (qSBlockIdx == (curQSBlockNum - 1U)) ? + (qSeqlen - qSBlockIdx * curQSBlockTile) : curQSBlockTile; + uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1U)) ? + (groupSize - qNBlockIdxCurGroup * curQNBlockTile) : curQNBlockTile; + + int64_t noSkipKvS = static_cast(kvSeqlen); + uint32_t kvSLoopNumTotal = 0; + uint32_t startIdx = 0; + int32_t preTokenStartLen = 0; + int32_t preTokenEndLen = 0; + int32_t nextTokenStartLen = 0; + int32_t nextTokenEndLen = 0; + bool notPreMask = true; + bool notNextMask = true; + int32_t delStartRow = 0; + int32_t delEndRow = qSeqlen; + bool startsWithMaskTile = false; + bool startsWithMaskThenNomaskFlag = false; + if constexpr (IS_FD) { + noSkipKvS = kvSeqlen; + if (maskType != 0U) { + int64_t diffS = kvSeqlen - qSeqlen; + diffS = (diffS < 0) ? 0 : diffS; + noSkipKvS = (qSBlockIdx + 1U) * curQSBlockTile + diffS; + noSkipKvS = AscendC::Std::min(static_cast(kvSeqlen), noSkipKvS); + } + kvSLoopNumTotal = NpuArch::Detail::Alignment::CeilDiv(static_cast(noSkipKvS), MAX_KV_STACK_LEN); + } else { + if (maskType != 0U && sparseMode != 4U && maskType != 4U) { + int64_t diffS = kvSeqlen - qSeqlen; + diffS = (diffS < 0) ? 0 : diffS; + noSkipKvS = (qSBlockIdx + 1U) * curQSBlockTile + diffS; + noSkipKvS = AscendC::Std::min(static_cast(kvSeqlen), noSkipKvS); + kvSLoopNumTotal = NpuArch::Detail::Alignment::CeilDiv(noSkipKvS, MAX_KV_STACK_LEN); + } else if (maskType != 0U && sparseMode == 4U) { + int32_t leftPointPreToken = kvSeqlen; + int32_t leftPointNextToken = 0; + if (preToken < 0 && preToken * (-1) >= qSeqlen) { + startIdx = kvSeqlen / MAX_KV_STACK_LEN + 1; + } else if (preToken != SPARSE_MODE_INT_MAX) { + leftPointPreToken = kvSeqlen - qSeqlen - preToken; + preTokenStartLen = qSBlockIdx * curQSBlockTile + leftPointPreToken; + preTokenEndLen = qSBlockIdx * curQSBlockTile + qSBlockSize + leftPointPreToken; + startIdx = AscendC::Std::max(static_cast(0), preTokenStartLen) / static_cast(MAX_KV_STACK_LEN); + notPreMask = false; + } else { + startIdx = 0; + } + if (nextToken < 0 && nextToken * (-1) >= kvSeqlen) { + kvSLoopNumTotal = 0; + } else if (nextToken != SPARSE_MODE_INT_MAX) { + leftPointNextToken = kvSeqlen - qSeqlen + nextToken; + nextTokenStartLen = qSBlockIdx * curQSBlockTile + leftPointNextToken; + nextTokenEndLen = qSBlockIdx * curQSBlockTile + qSBlockSize + leftPointNextToken; + noSkipKvS = AscendC::Std::min(static_cast(kvSeqlen), NpuArch::Detail::Alignment::RoundUp(nextTokenEndLen, static_cast(MAX_KV_STACK_LEN))); + noSkipKvS = noSkipKvS <= 0 ? kvSeqlen : noSkipKvS; + kvSLoopNumTotal = NpuArch::Detail::Alignment::CeilDiv(static_cast(noSkipKvS), MAX_KV_STACK_LEN); + notNextMask = false; + } else { + noSkipKvS = kvSeqlen; + kvSLoopNumTotal = NpuArch::Detail::Alignment::CeilDiv(static_cast(noSkipKvS), MAX_KV_STACK_LEN); + } + if (preTokenEndLen > static_cast(kvSeqlen) && preToken != SPARSE_MODE_INT_MAX) { + delStartRow = kvSeqlen - leftPointPreToken; + } else if (nextTokenStartLen < 0 && nextToken != SPARSE_MODE_INT_MAX) { + delEndRow = -leftPointNextToken; + } + } else { + kvSLoopNumTotal = NpuArch::Detail::Alignment::CeilDiv(noSkipKvS, MAX_KV_STACK_LEN); + } + } + + uint32_t kvStart = 0; + uint32_t kvEnd = kvSLoopNumTotal; + if constexpr (IS_FD) { + kvStart = stS2IdxNow; + kvEnd = (enS2IdxNow == static_cast(curKSBlockNum)) ? kvSLoopNumTotal : enS2IdxNow; + } else { + kvStart = startIdx; + kvEnd = kvSLoopNumTotal; + } + + uint32_t rowNum = qSBlockSize * qNBlockSize; + int32_t stackSeqCount = 0; + uint32_t preKVNum = PRE_LAUNCH; + uint32_t blockStackNum = (MAX_KV_STACK_LEN - 1 + pagedBlockSize) / pagedBlockSize; + uint32_t stackSeqTile = MAX_KV_STACK_LEN; + uint32_t stackSeqTilePad = MAX_KV_STACK_LEN; + bool isLastStackTile = false; + + +#ifdef __DAV_C220_VEC__ + if (kvSLoopNumTotal <= 0 || startIdx >= kvSLoopNumTotal) { + LayoutO layoutO(qSeqlen, embed * qHeads); + LayoutLse layoutLse(totalQTokens, qHeads); + epilogueInitOut(gO[gmOffsetO], gLse[gmOffsetLse], layoutO, layoutLse, qSBlockSize, qNBlockSize); + } +#endif +#ifdef __DAV_C220_CUBE__ + LayoutQ layoutQTemp(rowNum, embed); + LayoutK layoutKTemp(strideK, stackSeqTile); + LayoutV layoutVTemp(stackSeqTile, strideV); + blockMmadQK.resetBlockStart(kvStart, pagedBlockSize); + blockMmadPV.resetBlockStart(kvStart, pagedBlockSize); + blockMmadQK.loadQGM(gQ[gmOffsetQ], layoutQTemp, rowNum, qNBlockSize, qHeads); +#endif + for (uint32_t kvSIdx = kvStart; kvSIdx < kvEnd + preKVNum; kvSIdx++) { + if (kvSIdx < kvEnd) { + if (kvSIdx + 1 > kvSLoopNumTotal - 1U) { + stackSeqTile = noSkipKvS - kvSIdx * MAX_KV_STACK_LEN; + } else { + stackSeqTile = MAX_KV_STACK_LEN; + } + isLastStackTile = (kvSIdx + 1) >= kvSLoopNumTotal; + uint32_t curStackTileMod = stackSeqCount % (PRE_LAUNCH + 1U); + uint64_t gmOffsetS = + static_cast(coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1U) + + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB); + GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed}; + LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad); +#ifdef __DAV_C220_CUBE__ + if constexpr (PAGED_CACHE_FLAG) { + blockMmadQK( + gQ[gmOffsetQ], + gK[gmOffsetK], + gS[gmOffsetS], + gBlockTable[blockBOffset], + layoutQTemp, + layoutKTemp, + layOutS, + actualBlockShapeQK, + kvSIdx, + kvSLoopNumTotal, + pagedBlockSize, + strideK); + } else { + blockMmadQK( + gQ[gmOffsetQ], + gK[gmOffsetK], + gS[gmOffsetS], + gBlockTable, + layoutQTemp, + layoutKTemp, + layOutS, + actualBlockShapeQK, + kvSIdx, + kvSLoopNumTotal, + pagedBlockSize, + strideK); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady); +#endif +#ifdef __DAV_C220_VEC__ + LayoutP layOutP(rowNum, stackSeqTile, stackSeqTilePad); + LayoutMask layOutMask(COMP_TRIU_MASK_DIM_LEN, COMP_TRIU_MASK_DIM_LEN); + LayoutQ layOutFullMask(pseQ, pseKv); + uint64_t gmOffsetP = gmOffsetS; + uint32_t kvSStartIdx = kvSIdx * MAX_KV_STACK_LEN; + uint32_t kvSEndIdx = kvSStartIdx + stackSeqTile; + if constexpr (MASK_TYPE == FaiKernel::MaskType::MASK_CAUSAL) { + uint32_t triUp = noSkipKvS - qSBlockSize; + uint32_t triDown = noSkipKvS; + bool doTriUMask = triUp < kvSEndIdx - 1; + if (doTriUMask) { + if constexpr (IS_FD) { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + gMask, + layOutP, + layOutS, + layOutMask, + actualBlockShapeQK, + (stackSeqCount == 0), + qSBlockSize, + qNBlockSize, + curStackTileMod, + qkReady, + triUp, + triDown, + kvSStartIdx, + kvSEndIdx, + isLastStackTile, + isSplitKV); + } else { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + gMask, + layOutP, + layOutS, + layOutMask, + actualBlockShapeQK, + (stackSeqCount == 0), + qSBlockSize, + qNBlockSize, + curStackTileMod, + qkReady, + triUp, + triDown, + kvSStartIdx, + kvSEndIdx, + isLastStackTile, + false); + } + } else { + uint32_t noMaskStackSeqNum = (triUp + 1) / MAX_KV_STACK_LEN; + Arch::CrossCoreWaitFlag(qkReady); + if constexpr (IS_FD) { + int32_t localLastNoMaskStackId = (int32_t)(noMaskStackSeqNum - 1) - kvStart; + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + layOutP, + layOutS, + actualBlockShapeQK, + (stackSeqCount == 0), + (stackSeqCount == localLastNoMaskStackId), + qSBlockSize, + qNBlockSize, + curStackTileMod, + isLastStackTile, + isSplitKV, + false, + startsWithMaskThenNomaskFlag); + } else { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + layOutP, + layOutS, + actualBlockShapeQK, + (stackSeqCount == 0), + (stackSeqCount == noMaskStackSeqNum - 1), + qSBlockSize, + qNBlockSize, + curStackTileMod, + isLastStackTile, + false, + false, + startsWithMaskThenNomaskFlag); + } + } + } else if constexpr (MASK_TYPE == FaiKernel::MaskType::MASK_SWA) { + if constexpr (!IS_FD) { + bool doTriUPreMask = (sparseMode != 4 || notPreMask) ? false : + (preTokenStartLen >= kvSStartIdx && preTokenStartLen < kvSEndIdx) || + (preTokenEndLen > kvSStartIdx && preTokenEndLen <= kvSEndIdx) || + (preTokenStartLen <= kvSStartIdx && preTokenEndLen >= kvSEndIdx); + bool doTriUNextMask = (sparseMode != 4 || notNextMask) ? false : + (nextTokenStartLen >= kvSStartIdx && nextTokenStartLen < kvSEndIdx) || + (nextTokenEndLen > kvSStartIdx && nextTokenEndLen <= kvSEndIdx) || + (nextTokenStartLen <= kvSStartIdx && nextTokenEndLen >= kvSEndIdx); + bool doTriUMask = (doTriUPreMask || doTriUNextMask); + if (doTriUMask) { + startsWithMaskTile = true; + startsWithMaskThenNomaskFlag = true; + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + gMask, + layOutP, + layOutS, + layOutMask, + actualBlockShapeQK, + (stackSeqCount == 0), + qSBlockSize, + qNBlockSize, + curStackTileMod, + qkReady, + kvSStartIdx, + doTriUPreMask, + doTriUNextMask, + preTokenStartLen, + preTokenEndLen, + nextTokenStartLen, + nextTokenEndLen, + isLastStackTile); + } else { + bool isLastNoMaskStackTile = (nextTokenStartLen >= kvSeqlen) || (nextTokenStartLen < 0); + uint32_t kvSeqlenLimit = isLastNoMaskStackTile ? kvSeqlen : nextTokenStartLen; + uint32_t alignedKvSeqlenLimit = isLastNoMaskStackTile ? + NpuArch::Detail::Alignment::RoundUp(kvSeqlenLimit, MAX_KV_STACK_LEN) : + NpuArch::Detail::Alignment::RoundDown(kvSeqlenLimit, MAX_KV_STACK_LEN); + uint32_t noMaskStackSeqNum = (alignedKvSeqlenLimit - kvStart * MAX_KV_STACK_LEN) / MAX_KV_STACK_LEN; + Arch::CrossCoreWaitFlag(qkReady); + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + layOutP, + layOutS, + actualBlockShapeQK, + (stackSeqCount == 0), + (stackSeqCount == noMaskStackSeqNum - 1), + qSBlockSize, + qNBlockSize, + curStackTileMod, + isLastStackTile, + false, + startsWithMaskTile, + startsWithMaskThenNomaskFlag); + startsWithMaskTile = false; + } + } + } else if constexpr (MASK_TYPE == FaiKernel::MaskType::FULL_MASK) { + uint32_t rowOffer = qSBlockIdx * curQSBlockTile; + if constexpr (!IS_FD) { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + gPseShift, + layOutP, + layOutS, + layOutFullMask, // pseShift + actualBlockShapeQK, + (stackSeqCount == 0), + qSBlockSize, + qNBlockSize, + curStackTileMod, + qkReady, + rowOffer, + kvSStartIdx, + kvSEndIdx, + qNStartIdx, + BIdx, + qHeads, + pseQ, + pseKv, + isLastStackTile); + } + } else { + Arch::CrossCoreWaitFlag(qkReady); + if constexpr (IS_FD) { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + layOutP, + layOutS, + actualBlockShapeQK, + (stackSeqCount == 0), + 0, + qSBlockSize, + qNBlockSize, + curStackTileMod, + isLastStackTile, + isSplitKV, + false); + } else { + epilogueOnlineSoftmax( + gP[gmOffsetP], + gS[gmOffsetS], + gSink[gmOffsetSink], + layOutP, + layOutS, + actualBlockShapeQK, + (stackSeqCount == 0), + 0, + qSBlockSize, + qNBlockSize, + curStackTileMod, + isLastStackTile, + false, + false, + startsWithMaskThenNomaskFlag); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxReady); +#endif + } + if (kvSIdx >= kvStart + preKVNum) { + uint32_t nowkvSIdx = kvSIdx - preKVNum; + if (nowkvSIdx + 1 > kvSLoopNumTotal - 1U) { + stackSeqTile = noSkipKvS - nowkvSIdx * MAX_KV_STACK_LEN; + } else { + stackSeqTile = MAX_KV_STACK_LEN; + } + uint32_t curStackTileMod = (stackSeqCount - PRE_LAUNCH) % (PRE_LAUNCH + 1U); + uint64_t gmOffsetOTmp = + static_cast(coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1U) + + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB); + GemmCoord actualBlockShapePV{rowNum, embedV, stackSeqTile}; + LayoutOTmp layoutOTmp(rowNum, embedV, embedRoundV); +#ifdef __DAV_C220_CUBE__ + LayoutP layoutPTemp(rowNum, stackSeqTile, stackSeqTilePad); + uint64_t gmOffsetP = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) + curStackTileMod * WORKSPACE_BLOCK_SIZE_DB; + if constexpr (PAGED_CACHE_FLAG) { + blockMmadPV( + gP[gmOffsetP], + gV[gmOffsetV], + gOTmp[gmOffsetOTmp], + gBlockTable[blockBOffset], + layoutPTemp, + layoutVTemp, + layoutOTmp, + actualBlockShapePV, + nowkvSIdx, + kvSLoopNumTotal, + pagedBlockSize, + noSkipKvS, + strideV, + blockStackNum, + softmaxReady); + } else { + blockMmadPV( + gP[gmOffsetP], + gV[gmOffsetV], + gOTmp[gmOffsetOTmp], + gBlockTable, + layoutPTemp, + layoutVTemp, + layoutOTmp, + actualBlockShapePV, + nowkvSIdx, + kvSLoopNumTotal, + pagedBlockSize, + noSkipKvS, + strideV, + blockStackNum, + softmaxReady); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady); +#endif +#ifdef __DAV_C220_VEC__ + LayoutO layoutO(qSeqlen, embed * qHeads); + LayoutUpdate layoutUpdate(rowNum, embed, embedRound); + LayoutLse layoutLse(totalQTokens, qHeads); + uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * WORKSPACE_BLOCK_SIZE_DB); + Arch::CrossCoreWaitFlag(pvReady); + + if constexpr (IS_FD) { + LayoutLse layoutgmLse(qSBlockSize, qNBlockSize); + LayoutLse layoutgmLo(qSBlockSize, embed * qNBlockSize); + typename EpilogueRescaleO::SplitKVParams splitParams; + splitParams.isSplitkv = isSplitKV; + splitParams.gCombineLse = gLseFD[gmOffsetLseFD]; + splitParams.gCombineo = gOFD[gmOffsetOFD]; + splitParams.layoutgmLse = &layoutgmLse; + splitParams.layoutgmLo = &layoutgmLo; + + epilogueRescaleO( + gO[gmOffsetO], + gOTmp[gmOffsetOTmp], + gOUpdate[gmOffsetUpdate], + gLse[gmOffsetLse], + layoutO, + layoutOTmp, + layoutUpdate, + layoutLse, + actualBlockShapePV, + qSBlockSize, + qNBlockSize, + (stackSeqCount - PRE_LAUNCH == 0), + nowkvSIdx + 1 >= kvEnd, + curStackTileMod, + delStartRow, + delEndRow, + qSeqlen, + qSBlockIdx, + curQNBlockTile, + splitParams); + } else { + epilogueRescaleO( + gO[gmOffsetO], + gOTmp[gmOffsetOTmp], + gOUpdate[gmOffsetUpdate], + gLse[gmOffsetLse], + layoutO, + layoutOTmp, + layoutUpdate, + layoutLse, + actualBlockShapePV, + qSBlockSize, + qNBlockSize, + (stackSeqCount - PRE_LAUNCH == 0), + nowkvSIdx + 1 >= kvSLoopNumTotal, + curStackTileMod, + delStartRow, + delEndRow, + qSeqlen, + qSBlockIdx, + curQNBlockTile); + } +#endif + } + stackSeqCount++; + } + } + + private: + uint64_t mm1OutSize; + uint64_t smOnlineOutSize; + uint64_t mm2OutSize; + uint32_t batch; + uint32_t qHeads; + uint32_t kvHeads; + uint32_t embed; + uint32_t embedV; + uint32_t pagedBlockSize; + uint32_t maxNumBlocksPerBatch; + uint32_t firstBatchTaskNum; + uint32_t totalTaskNum; + uint32_t blockSize; + uint32_t maskType; + float scaleValue; + uint32_t sparseMode; + int64_t preToken; + int64_t nextToken; + int64_t pseQ; + int64_t pseKv; + uint32_t totalQTokens; + + uint64_t strideQ; + uint64_t strideO; + uint64_t strideK; + uint64_t strideV; + uint32_t embedRound; + uint32_t embedRoundV; + uint32_t groupSize; + + Arch::Resource resource; + Arch::CrossCoreFlag qkReady{QK_READY_ID}; + Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID}; + Arch::CrossCoreFlag pvReady{PV_READY_ID}; + + BlockMmadQK blockMmadQK; + BlockMmadPV blockMmadPV; + EpilogueOnlineSoftmax epilogueOnlineSoftmax; + EpilogueRescaleO epilogueRescaleO; + EpilogueInitOut epilogueInitOut; + }; +} +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/kernel_common.hpp b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/kernel_common.hpp new file mode 100644 index 0000000000..0e2c21cb0b --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/kernel_common.hpp @@ -0,0 +1,155 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * 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. + */ + +/*! +* \file kernel_common.hpp +* \brief +*/ + +#ifndef KERNEL_COMMON +#define KERNEL_COMMON + +#include "attn_infra/base_defs.hpp" +#include "attn_infra/arch/arch.hpp" +#include "attn_infra/layout/layout.hpp" + +#include "attn_infra/gemm/block/block_mmad.hpp" +#include "attn_infra/gemm/dispatch_policy.hpp" +#include "attn_infra/gemm/gemm_type.hpp" + +#include "attn_infra/arch/cross_core_sync.hpp" +#include "attn_infra/arch/resource.hpp" +#include "attn_infra/epilogue/block/block_epilogue.hpp" +#include "attn_infra/epilogue/dispatch_policy.hpp" +#if ASC_DEVKIT_MAJOR >= 9 +#include "kernel_vec_intf.h" +#include "kernel_cube_intf.h" +#else +#include "kernel_operator.h" +#endif +#include "kernel_operator_list_tensor_intf.h" +#include "../../generated/kernel_tiling/kernel_tiling.h" + +namespace KernelCommon { + constexpr uint32_t QK_READY_ID = 1; + constexpr uint32_t SOFTMAX_READY_ID = 2; + constexpr uint32_t PV_READY_ID = 3; + constexpr uint32_t PRE_LAUNCH = 2; + constexpr uint32_t N_SPLIT_HELPER = 2; + constexpr uint32_t MAX_KV_STACK_LEN = 512; + constexpr uint32_t Q_TILE_CEIL = 128; + constexpr uint32_t WORKSPACE_BLOCK_SIZE_DB = Q_TILE_CEIL * MAX_KV_STACK_LEN; + constexpr uint32_t L1_MAX_SIZE = 524288; + constexpr uint32_t L1_MAX_N_NUM = 128; + constexpr uint32_t DOUBLE_BUFFER = 2; + constexpr uint32_t COMP_TRIU_MASK_DIM_LEN = 2048; + constexpr uint32_t NUM_32 = 32; + constexpr uint32_t NUM_128 = 128; + constexpr uint32_t NUM_256 = 256; + constexpr uint32_t FLOAT_SIZE = 4; + constexpr int64_t SPARSE_MODE_INT_MAX = 2147483647; + + template + __aicore__ inline + T AlignUp(T a, T b) + { + return (b == 0) ? 0 : (a + b - 1) / b * b; + } + + template + __aicore__ inline + T Max(T a, T b) + { + return (a > b) ? a : b; + } + + namespace FaiKernel { + constexpr uint32_t BLOCK_SIZE = 16; + + enum class cvPipeLineType : uint32_t { + FAI_COMMON_NORMAL = 0, + FAI_COMMON_CHUNK_MASK = 1, + }; + + enum class MaskType : uint32_t { + NO_MASK = 0, + MASK_CAUSAL = 1, + MASK_SPEC = 2, + MASK_SWA = 4, + FULL_MASK = 5 + }; + + enum class inputLayout : uint32_t { + BSND = 0, + TND = 1 + }; + }; + + struct FAIKernelParams { + // Data members + GM_ADDR q; + GM_ADDR k; + GM_ADDR v; + GM_ADDR pseShift; + GM_ADDR mask; + GM_ADDR blockTables; + GM_ADDR actualQseqlen; + GM_ADDR actualKvseqlen; + GM_ADDR o; + GM_ADDR lse; + GM_ADDR workSpace; + GM_ADDR tiling; + GM_ADDR sink; + + // Methods + __aicore__ inline FAIKernelParams() {} + + __aicore__ inline FAIKernelParams(GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR pseShift_, GM_ADDR mask_, GM_ADDR blockTables_, + GM_ADDR actualQseqlen_, GM_ADDR actualKvseqlen_, GM_ADDR o_, GM_ADDR lse_, GM_ADDR workSpace_, GM_ADDR tiling_, GM_ADDR sink_) + : q(q_), k(k_), v(v_), pseShift(pseShift_), mask(mask_), blockTables(blockTables_), actualQseqlen(actualQseqlen_), + actualKvseqlen(actualKvseqlen_), o(o_), lse(lse_), workSpace(workSpace_), tiling(tiling_), sink(sink_) {} + }; + + __aicore__ inline uint32_t GetQNBlockTile(uint32_t qSeqlen, uint32_t groupSize) + { + uint32_t qNBlockTile = (qSeqlen != 0) ? + (Q_TILE_CEIL / qSeqlen) / N_SPLIT_HELPER * N_SPLIT_HELPER : Q_TILE_CEIL; + qNBlockTile = qNBlockTile < groupSize ? qNBlockTile : groupSize; + qNBlockTile = qNBlockTile < 1 ? 1 : qNBlockTile; + return qNBlockTile; + } + + __aicore__ inline uint32_t GetKvNBlockTile(uint32_t rowNumPerQSGTile, uint32_t kvHead) + { + uint32_t rowNumCeilPerQSGKvNTile = Q_TILE_CEIL; + uint32_t kvNBlockTile = rowNumCeilPerQSGKvNTile / rowNumPerQSGTile; + kvNBlockTile = kvNBlockTile < kvHead ? kvNBlockTile : kvHead; + kvNBlockTile = kvNBlockTile < 1 ? 1 : kvNBlockTile; + return kvNBlockTile; + } + + __aicore__ inline uint32_t GetQSBlockTile(uint32_t kvSeqlen) + { + uint32_t qSBlockTile = Q_TILE_CEIL; + return qSBlockTile; + } + + __aicore__ inline uint32_t GetQSBlockTileDecode(uint32_t qSeqlen) + { + uint32_t qSBlockTile = Q_TILE_CEIL < qSeqlen ? Q_TILE_CEIL : qSeqlen; + return qSBlockTile; + } + __aicore__ inline uint32_t GetKSBlockTile(uint32_t kvSeqlen) + { + uint32_t kSBlockTile = MAX_KV_STACK_LEN; + return kSBlockTile; + } +} +#endif diff --git a/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py b/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py new file mode 100644 index 0000000000..bdd92b4b45 --- /dev/null +++ b/examples/a2a3/host_build_graph/qwen3_14b_decode/test_qwen3_14b_decode.py @@ -0,0 +1,456 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Qwen3-14B 40-layer decode (CANN fused-attention) — SceneTestCase. + +Self-contained port of pypto-lib ``models/qwen3/14b/decode_fwd.py`` entry +``decode_fwd_layers`` with ``_CHUNK_NLAYERS == 40``: the whole Qwen3-14B decode +stack as ONE fused dispatch (hidden -> hidden, no LM head), carrying the +inter-layer residual in FP32. The layer loop is a real loop in the generated +orchestration, so a 40-layer chunk costs one extra literal over a 2-layer one, +not 20x the kernels. + +The C++ under ``kernels/`` is harvested pypto codegen (orchestration + 18 AIC + +16 AIV) plus the hand-written CANN attention extern under +``kernels/paged_attention_cce/``; ``simpler_setup/goldens/qwen3_14b_decode.py`` +is the matching torch reference. See README.md for provenance and how to +regenerate. + +Parameter regime matches ``stress_profile.py`` (vLLM serving stress): BATCH=16, +MAX_SEQ=5500 (= max_model_len), fixed decode seq_len=3500. Weights and the paged +KV pool are stacked x40 (one slice per layer); every layer reuses layer-0 +weights, per the lib const-layer-0 stacked-fwd reference, while each layer reads +and writes its own KV pool. +""" + +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, scene_test +from simpler_setup.goldens.qwen3_14b_decode import ( + compute_golden as _decode_golden, +) +from simpler_setup.goldens.qwen3_14b_decode import ( + generate_inputs as _decode_generate_inputs, +) + +# CANN devkit headers for the attention extern, which builds on AscendC and the +# vendored FusedInferAttentionScore under kernels/paged_attention_cce/vendor/. +# `vendor/.../attn_infra/base_defs.hpp` selects its AscendC entry header under +# `#if ASC_DEVKIT_MAJOR >= 9`, which ccec predefines from the installed devkit, +# so a CANN 9 box must be able to resolve `basic_api/kernel_basic_intf.h` from +# one of these. +# +# `$ASCEND_HOME_PATH` keeps the paths machine-independent and is expanded at +# compile time, not import time — this file is collected on sim and macOS runners +# that have no CANN at all. The devkit's arch subdirectory is named differently +# across installs, so both layouts are listed; missing ones are dropped, and the +# scene-test resolver raises if *every* entry is missing. +_CANN_SUBDIRS = ( + "include", + "asc", + "asc/impl/adv_api", + "asc/impl/basic_api", + "asc/impl/basic_api/reg_compute", + "asc/impl/c_api", + "asc/impl/simt_api", + "asc/impl/utils", + "asc/include", + "asc/include/adv_api", + "asc/include/aicpu_api", + "asc/include/basic_api", + "asc/include/basic_api/reg_compute", + "asc/include/c_api", + "asc/include/interface", + "asc/include/simt_api", + "asc/include/utils", + "tikcpp/tikcfw", + "tikcpp/tikcfw/impl", + "tikcpp/tikcfw/interface", +) + +_CANN_INCLUDE_DIRS = [f"$ASCEND_HOME_PATH/{prefix}{sub}" for prefix in ("aarch64-linux/", "") for sub in _CANN_SUBDIRS] + + +# Validates the full 40-layer fused decode against a torch reference. +@scene_test(level=2, runtime="host_build_graph") +class TestQwen314BDecodeHostBuildGraph(SceneTestCase): + """Qwen3-14B decode, all 40 layers in one dispatch, against a torch reference.""" + + RTOL = 5e-2 + ATOL = 1e-1 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/decode_fwd_layers.cpp", + "function_name": "aicpu_orchestration_entry", + # decode_fwd_layers takes k_cache / v_cache as plain inputs, but the + # attention extern writes the current token's KV into them. Declaring + # them INOUT here has simpler copy the pools back, so the golden can + # check all 40 layers' KV writes and not just the hidden output. + "signature": [ + D.IN, # 0 hidden_states + D.IN, # 1 input_rms_weight + D.IN, # 2 wq + D.IN, # 3 wk + D.IN, # 4 wv + D.IN, # 5 q_norm_weight + D.IN, # 6 k_norm_weight + D.IN, # 7 seq_lens + D.IN, # 8 block_table + D.IN, # 9 slot_mapping + D.IN, # 10 rope_cos + D.IN, # 11 rope_sin + D.INOUT, # 12 k_cache + D.INOUT, # 13 v_cache + D.IN, # 14 wo + D.IN, # 15 w_gate + D.IN, # 16 w_up + D.IN, # 17 w_down + D.IN, # 18 post_rms_weight + D.OUT, # 19 out + ], + }, + # 37 incores (func_id 0..36), transcribed from the pypto codegen + # kernel_config.py for decode_fwd_layers (N=40). func_id 0/11/12 are the + # CANN attention externs; 11 and 12 are the same source dispatched as the + # AIC and AIV halves of one mixed task. + "incores": [ + { + "func_id": 0, + "name": "paged_attention_tiling_cce", + "source": "kernels/vendor/paged_attention_cce/tiling/entry.cpp", + "core_type": "aiv", + "extra_include_dirs": _CANN_INCLUDE_DIRS, + "signature": [D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "copy_hidden", + "source": "kernels/aiv/copy_hidden.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.IN], + }, + { + "func_id": 2, + "name": "x_gamma0", + "source": "kernels/aiv/x_gamma0.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.IN, D.IN], + }, + { + "func_id": 3, + "name": "attn_out_seed", + "source": "kernels/aiv/attn_out_seed.cpp", + "core_type": "aiv", + "signature": [D.IN], + }, + { + "func_id": 4, + "name": "rms_recip", + "source": "kernels/aiv/rms_recip.cpp", + "core_type": "aiv", + "signature": [D.IN, D.INOUT], + }, + { + "func_id": 5, + "name": "q_seed", + "source": "kernels/aiv/q_seed.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 6, + "name": "q_proj", + "source": "kernels/aic/q_proj.cpp", + "core_type": "aic", + "signature": [D.INOUT, D.IN, D.IN], + }, + { + "func_id": 7, + "name": "kv_seed", + "source": "kernels/aiv/kv_seed.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT], + }, + { + "func_id": 8, + "name": "mlp_out_seed", + "source": "kernels/aiv/mlp_out_seed.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.INOUT, D.INOUT], + }, + { + "func_id": 9, + "name": "k_proj", + "source": "kernels/aic/k_proj.cpp", + "core_type": "aic", + "signature": [D.INOUT, D.IN, D.IN], + }, + { + "func_id": 10, + "name": "v_proj", + "source": "kernels/aic/v_proj.cpp", + "core_type": "aic", + "signature": [D.INOUT, D.IN, D.IN], + }, + { + "func_id": 11, + "name": "paged_attention_rope_cce_aic", + "source": "kernels/vendor/paged_attention_cce/attention_rope/entry.cpp", + "core_type": "aic", + "extra_include_dirs": _CANN_INCLUDE_DIRS, + "signature": [ + D.INOUT, + D.INOUT, + D.INOUT, + D.INOUT, + D.IN, + D.INOUT, + D.INOUT, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + ], + }, + { + "func_id": 12, + "name": "paged_attention_rope_cce_aiv", + "source": "kernels/vendor/paged_attention_cce/attention_rope/entry.cpp", + "core_type": "aiv", + "extra_include_dirs": _CANN_INCLUDE_DIRS, + "signature": [ + D.INOUT, + D.INOUT, + D.INOUT, + D.INOUT, + D.IN, + D.INOUT, + D.INOUT, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + D.IN, + ], + }, + { + "func_id": 13, + "name": "out_proj", + "source": "kernels/aic/out_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 14, + "name": "out_proj_0", + "source": "kernels/aic/out_proj_0.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 15, + "name": "residual_rms_cast", + "source": "kernels/aiv/residual_rms_cast.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + }, + { + "func_id": 16, + "name": "residual_rms_cast_0", + "source": "kernels/aiv/residual_rms_cast_0.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + }, + { + "func_id": 17, + "name": "residual_rms_cast_1", + "source": "kernels/aiv/residual_rms_cast_1.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + }, + { + "func_id": 18, + "name": "residual_rms_cast_2", + "source": "kernels/aiv/residual_rms_cast_2.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + }, + { + "func_id": 19, + "name": "residual_rms_cast_3", + "source": "kernels/aiv/residual_rms_cast_3.cpp", + "core_type": "aiv", + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + }, + { + "func_id": 20, + "name": "post_rms_reduce", + "source": "kernels/aiv/post_rms_reduce.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 21, + "name": "gate_proj", + "source": "kernels/aic/gate_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 22, + "name": "up_proj", + "source": "kernels/aic/up_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 23, + "name": "gate_proj_0", + "source": "kernels/aic/gate_proj_0.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 24, + "name": "up_proj_0", + "source": "kernels/aic/up_proj_0.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 25, + "name": "gate_proj_1", + "source": "kernels/aic/gate_proj_1.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 26, + "name": "up_proj_1", + "source": "kernels/aic/up_proj_1.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 27, + "name": "gate_proj_2", + "source": "kernels/aic/gate_proj_2.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 28, + "name": "up_proj_2", + "source": "kernels/aic/up_proj_2.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 29, + "name": "gate_proj_3", + "source": "kernels/aic/gate_proj_3.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 30, + "name": "up_proj_3", + "source": "kernels/aic/up_proj_3.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 31, + "name": "gate_proj_4", + "source": "kernels/aic/gate_proj_4.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 32, + "name": "up_proj_4", + "source": "kernels/aic/up_proj_4.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 33, + "name": "silu", + "source": "kernels/aiv/silu.cpp", + "core_type": "aiv", + "signature": [D.IN, D.INOUT, D.IN, D.IN], + }, + { + "func_id": 34, + "name": "down_proj", + "source": "kernels/aic/down_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], + }, + { + "func_id": 35, + "name": "dcr_xgamma", + "source": "kernels/aiv/dcr_xgamma.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.INOUT, D.IN, D.INOUT], + }, + { + "func_id": 36, + "name": "copy_out", + "source": "kernels/aiv/copy_out.cpp", + "core_type": "aiv", + "signature": [D.OUT, D.IN], + }, + ], + } + + CASES = [ + { + "name": "StressBatch16Seq3500", + "platforms": ["a2a3"], + # A run takes the whole device, matching the lib default. + # + # The heap must hold all 40 layers' intermediates simultaneously: + # host_build_graph builds the whole graph before the device schedules + # anything, so no task has completed while the graph is being built + # and the heap tail never advances off 0 — nothing is reclaimed + # mid-orchestration. The 256 MiB default runs out in the last layers + # (`Task Allocator Deadlock - Heap Exhausted`, tail=0, ~254 MiB of + # 256 used); 512 MiB carries it. The task window is not the + # constraint — the graph is ~10.6K tasks, inside the 16384 default. + # + # The tensormap_and_ringbuffer variant needs no sizing at all, + # because each layer's scope frees its intermediates as + # orchestration proceeds and the live set stays flat in layer count. + "config": {"aicpu_thread_num": 4, "runtime_env": {"ring_heap": 536870912}}, + "params": {"seed": 1234, "seq_len": 3500}, + }, + ] + + def generate_args(self, params): + return _decode_generate_inputs(params.get("seed", 1234), params.get("seq_len", 3500)) + + def compute_golden(self, args, params): + _decode_golden(args) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__)