diff --git a/.claude/skills/l0-swimlane/SKILL.md b/.claude/skills/l0-swimlane/SKILL.md index 7afdbbf971..5e18e1c18e 100644 --- a/.claude/skills/l0-swimlane/SKILL.md +++ b/.claude/skills/l0-swimlane/SKILL.md @@ -81,7 +81,6 @@ Verified examples (slots read from source): | Test | Loop bound | Flag | | ---- | ---------- | ---- | | `paged_attention_unroll` | `aic_qk_matmul.cpp` `args[4] = n_blocks` (scalar) | `--set-arg 4=4` | -| `qwen3_14b_decode` (fa_fused) | `fa_fused_aic/aiv.cpp` outer loop `for(i=block_idx; i()` 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 -The 36 sources under `kernels/` (orchestration + 35 incores: 8 AIC + 27 AIV) and -the golden `simpler_setup/goldens/qwen3_14b_decode.py` are **harvested pypto -codegen** for `decode_fwd_layers` (`_CHUNK_NLAYERS=2`, `PTO2_MANUAL_MAX_SEQ=5500`): -the orchestration + kernel `.cpp` (license header prepended, otherwise verbatim) -and the `CALLABLE` transcribed from that run's `kernel_config.py`. The golden -ports the per-layer `golden_decode_layer` math (RoPE θ=1e4, controlled scales, -FP32 residual, bf16 cast points) composed over the two layers with FP32 carry + -per-layer KV pools. - -The kernels are used essentially as generated, with **one hand-edit** to -`fa_fused_aiv`: the codegen emitted the AIV sub-block id as a `[[block_local]] -static` (`pypto_runtime_subblock_id`) whose **non-branch** `.text` relocation -simpler's strict `.text`-only loader rejects. The fix swaps it for -`setEntryOffset(get_sub_block_id(args) * …)` (legacy workaround). New MIX -kernels should prefer direct `TPUSH`/`TPOP` per -[docs/tpush-tpop.md](../../../docs/tpush-tpop.md). See -Status below. - -To regenerate: in pypto-lib, set `_CHUNK_NLAYERS=2`, `PTO2_MANUAL_MAX_SEQ=5500`, -build the ×2-stacked inputs, and `decode_fwd_layers.compile_for_test(...)`; then -harvest `orchestration/` + `kernels/` + `kernel_config.py`. No post-edit needed. +| 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 @@ -65,42 +150,29 @@ pytest examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode \ python examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/test_qwen3_14b_decode.py -p a2a3 -d ${DEVICE} ``` -DFX (the `--enable-l2-swimlane` the lib command implies) is opt-in via the -existing flags — no kernel changes needed: +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 -The case passes deterministically on device — output and both layers' KV-cache -match the torch reference (finite `max_abs_diff ≈ 0.03`). - -### Sub-block-id handling (the one hand-edit) - -The fused-attention tile-pipe (`TPUSH/TPOP<…, TILE_UP_DOWN>`) computes its per-AIV -FIFO offset from the **no-arg** ISA `get_subblockid()`, which under this MIX -dispatch returns **0 for both** AIV0 and AIV1 (confirmed by args dump). Left -uncorrected, both lanes collide on the same FIFO half and attention is poisoned → -NaN. The correct lane id (0/1) is only available from simpler's -`get_sub_block_id(args)`. - -The codegen's original workaround cached that value once (in `kernel_entry`) into -a per-core `[[block_local]]` static and `#define`d `get_subblockid()` to read it. -That static's `.text` relocation is rejected by `elf_parser` (which rejects -**any** `.rela.text` entry — the guard against unapplied *branch* relocations, -issues #900 / #830 / #831), so the verbatim kernel would not load. - -The fix follows the `spmd_paged_attention` pattern: drop the `[[block_local]]` -static and macro, and apply the per-lane split explicitly on the tile-pipe in -`kernel_entry`/`fa_fused_aiv` — -`v36.cons/prod.setEntryOffset(get_sub_block_id(args) * rows * cols * sizeof(T))`. -`get_subblockid()` then stays at its native 0 (the library auto-split contributes -nothing) and the explicit offset carries the lane separation. No per-core static, -no relocation — loads under the strict loader and computes correctly. - -The cleaner long-term fix is direct `TPUSH`/`TPOP` with platform-correct -`get_subblockid()` (see [docs/tpush-tpop.md](../../../docs/tpush-tpop.md) -and [`spmd_paged_attention`](../../../../tests/st/a2a3/tensormap_and_ringbuffer/spmd_paged_attention/kernels/mix/paged_attention_parallel.cpp)). +Passes on device: output **and all 40 layers' KV caches** match the torch +reference at `RTOL=5e-2 / ATOL=1e-1`. + +In CI it runs as its own `st-onboard-a2a3` step on a dedicated device, not in +the general scene-test sweep: at ~5 min it would eat over half that sweep's 600 s +session budget, and keeping the budget short there is what makes it a hang +detector. See `.github/workflows/ci.yml`. + +Measured on an idle a2a3 die, one device: **~5 min wall** for the whole case, +of which ~11 s is fixture generation and ~49 s the torch golden (both host-side, +single-threaded torch over 40 layers × 16 batch × 8 KV heads); the rest is +kernel compilation, the 38 GiB host→device upload, and the device run itself. diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/down_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/down_proj.cpp index ffb86fdbed..774948fce5 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/down_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/down_proj.cpp @@ -56,145 +56,146 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ const int64_t v8 = 2; const int64_t v9 = 15; const int64_t v10 = 32; - const int64_t v11 = 1024; - const int64_t v12 = 64; - const int64_t v13 = 5120; - const int64_t v14 = 1; - const int64_t v15 = 17408; - const int64_t v16 = 16; + 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 = 512; - const int64_t v20 = 0; - const int64_t v21 = 135168; - const int64_t v22 = 4096; + 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 v23 = (size_t)v12; - size_t v24 = (size_t)v20; - size_t v25 = (size_t)v10; + 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> - v26 = Tile< + v27 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v12); - uint64_t v27 = (uint64_t)v22; - TASSIGN(v26, v27); - pto::Shape<1, 1, 1, 16, 64> v28 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v29 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v30 = GlobalTensor< + v31 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + v4 * v14), v28, v29 + 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(v26, v30); + TLOAD(v27, v31); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v31 = Tile< + v32 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v32 = (uint64_t)v21; - TASSIGN(v31, v32); - int64_t v33 = (int64_t)((uint64_t)v5 + (uint64_t)v4); - pto::Shape<1, 1, 1, 64, 1024> v34 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v35 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v36 = GlobalTensor< + v37 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( - v2 + (v20 + v33 * v13 + v6 * v14), v34, v35 + v2 + (v21 + v34 * v12 + v6 * v13), v35, v36 ); - TLOAD(v31, 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 v37 = v24; v37 < v23; v37 += v25) { - int64_t v38 = (int64_t)v37; + 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> - v39 = Tile< + v40 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v16); - uint64_t v40 = (uint64_t)v20; - TASSIGN(v39, v40); + 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(v39, v26, v20, v37); + TEXTRACT(v40, v27, v21, v38); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v41 = Tile< + v42 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v42 = (uint64_t)v20; - TASSIGN(v41, v42); - TEXTRACT(v41, v31, v37, v20); + 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> - v43 = Tile< + v44 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v16); - uint64_t v44 = (uint64_t)v19; - TASSIGN(v43, v44); - int64_t v45 = (int64_t)((uint64_t)v38 + (uint64_t)v16); + 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(v43, v26, v20, v45); + TEXTRACT(v44, v27, v21, v46); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v46 = Tile< + v47 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v47 = (uint64_t)v18; - TASSIGN(v46, v47); - TEXTRACT(v46, v31, v45, v20); + 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 (v38 == v20) { + if (v39 == v21) { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v48 = Tile< + v49 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v49 = (uint64_t)v20; - TASSIGN(v48, v49); + CompactMode::Null>(v15, v20); + uint64_t v50 = (uint64_t)v21; + TASSIGN(v49, v50); pipe_barrier(PIPE_M); - TMATMUL(v48, v39, v41); + TMATMUL(v49, v40, v42); } else { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v50 = Tile< + v51 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v51 = (uint64_t)v20; - TASSIGN(v50, v51); + CompactMode::Null>(v15, v20); + uint64_t v52 = (uint64_t)v21; + TASSIGN(v51, v52); pipe_barrier(PIPE_M); - TMATMUL_ACC(v50, v50, v39, v41); + 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> - v52 = Tile< + v53 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v53 = (uint64_t)v20; - TASSIGN(v52, v53); + 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(v52, v52, v43, v46); + TMATMUL_ACC(v53, v53, v44, v47); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -204,147 +205,147 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ 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 v54 = (size_t)v14; v54 < ((size_t)v9); v54 += (size_t)v8) { - int64_t v55 = (int64_t)((uint64_t)((int64_t)v54) * (uint64_t)v12); - int64_t v56 = (int64_t)((uint64_t)v55 + (uint64_t)v12); + 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> - v57 = Tile< + v58 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v12); - uint64_t v58 = (uint64_t)v20; - TASSIGN(v57, v58); - pto::Shape<1, 1, 1, 16, 64> v59 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v60 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v61 = GlobalTensor< + v62 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, - pto::Layout::ND>(v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v55) * v14), v59, v60); + 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(v57, v61); + TLOAD(v58, v62); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v62 = Tile< + v63 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v63 = (uint64_t)v21; - TASSIGN(v62, v63); - pto::Shape<1, 1, 1, 64, 1024> v64 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v65 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v66 = GlobalTensor< + v67 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v55) * v13 + v6 * v14), v64, v65); - TLOAD(v62, v66); + 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> - v67 = Tile< + v68 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v12); - uint64_t v68 = (uint64_t)v17; - TASSIGN(v67, v68); - pto::Shape<1, 1, 1, 16, 64> v69 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v70 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v71 = GlobalTensor< + v72 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, - pto::Layout::ND>(v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v14), v69, v70); + 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(v67, v71); + TLOAD(v68, v72); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v72 = Tile< + v73 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v73 = (uint64_t)v22; - TASSIGN(v72, v73); - pto::Shape<1, 1, 1, 64, 1024> v74 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v75 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v76 = GlobalTensor< + v77 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v56) * v13 + v6 * v14), v74, v75); - TLOAD(v72, v76); + 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_M, PIPE_MTE1, EVENT_ID3); 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 v77 = v24; v77 < v23; v77 += v25) { + 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> - v78 = Tile< + v79 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v79 = (uint64_t)v20; - TASSIGN(v78, v79); + 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(v78, v57, v20, v77); + TEXTRACT(v79, v58, v21, v78); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v80 = Tile< + v81 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v81 = (uint64_t)v20; - TASSIGN(v80, v81); - TEXTRACT(v80, v62, v77, v20); + 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> - v82 = Tile< + v83 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v83 = (uint64_t)v19; - TASSIGN(v82, v83); - int64_t v84 = (int64_t)((uint64_t)((int64_t)v77) + (uint64_t)v16); + 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(v82, v57, v20, v84); + TEXTRACT(v83, v58, v21, v85); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v85 = Tile< + v86 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v86 = (uint64_t)v18; - TASSIGN(v85, v86); - TEXTRACT(v85, v62, v84, v20); + 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> - v87 = Tile< + v88 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v88 = (uint64_t)v20; - TASSIGN(v87, v88); + 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(v87, v87, v78, v80); + 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> - v89 = Tile< + v90 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v90 = (uint64_t)v20; - TASSIGN(v89, v90); + 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(v89, v89, v82, v85); + TMATMUL_ACC(v90, v90, v83, v86); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); @@ -355,71 +356,71 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ 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 v91 = v24; v91 < v23; v91 += v25) { + 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> - v92 = Tile< + v93 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v93 = (uint64_t)v20; - TASSIGN(v92, v93); + PadValue::Null, CompactMode::Null>(v15, v15); + uint64_t v94 = (uint64_t)v21; + TASSIGN(v93, v94); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - TEXTRACT(v92, v67, v20, v91); + TEXTRACT(v93, v68, v21, v92); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v94 = Tile< + v95 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v95 = (uint64_t)v20; - TASSIGN(v94, v95); - TEXTRACT(v94, v72, v91, v20); + 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> - v96 = Tile< + v97 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v97 = (uint64_t)v19; - TASSIGN(v96, v97); - int64_t v98 = (int64_t)((uint64_t)((int64_t)v91) + (uint64_t)v16); + 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(v96, v67, v20, v98); + TEXTRACT(v97, v68, v21, v99); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v99 = Tile< + v100 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v100 = (uint64_t)v18; - TASSIGN(v99, v100); - TEXTRACT(v99, v72, v98, v20); + 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> - v101 = Tile< + v102 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v102 = (uint64_t)v20; - TASSIGN(v101, v102); + 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(v101, v101, v92, v94); + 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> - v103 = Tile< + v104 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v104 = (uint64_t)v20; - TASSIGN(v103, v104); + 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(v103, v103, v96, v99); + TMATMUL_ACC(v104, v104, v97, v100); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); @@ -433,108 +434,108 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v105 = Tile< + v106 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v12); - uint64_t v106 = (uint64_t)v22; - TASSIGN(v105, v106); - pto::Shape<1, 1, 1, 16, 64> v107 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v108 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v109 = GlobalTensor< + v110 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v14), v107, v108 + v1 + (v21 + v21 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v108, v109 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v105, v109); + TLOAD(v106, v110); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v110 = Tile< + v111 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v111 = (uint64_t)v21; - TASSIGN(v110, v111); - pto::Shape<1, 1, 1, 64, 1024> v112 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v113 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v114 = GlobalTensor< + v115 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( - v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v7) * v13 + v6 * v14), v112, v113 + v2 + (v21 + (int64_t)((uint64_t)v34 + (uint64_t)v7) * v12 + v6 * v13), v113, v114 ); - TLOAD(v110, v114); + TLOAD(v111, v115); set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); 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 v115 = v24; v115 < v23; v115 += v25) { + 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> - v116 = Tile< + v117 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v16); - uint64_t v117 = (uint64_t)v20; - TASSIGN(v116, v117); + CompactMode::Null>(v15, v15); + uint64_t v118 = (uint64_t)v20; + TASSIGN(v117, v118); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - TEXTRACT(v116, v105, v20, v115); + TEXTRACT(v117, v106, v21, v116); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v118 = Tile< + v119 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v119 = (uint64_t)v20; - TASSIGN(v118, v119); - TEXTRACT(v118, v110, v115, v20); + 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> - v120 = Tile< + v121 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v16, v16); - uint64_t v121 = (uint64_t)v19; - TASSIGN(v120, v121); - int64_t v122 = (int64_t)((uint64_t)((int64_t)v115) + (uint64_t)v16); + 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(v120, v105, v20, v122); + TEXTRACT(v121, v106, v21, v123); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v123 = Tile< + v124 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v124 = (uint64_t)v18; - TASSIGN(v123, v124); - TEXTRACT(v123, v110, v122, v20); + 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> - v125 = Tile< + v126 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v126 = (uint64_t)v20; - TASSIGN(v125, v126); + 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(v125, v125, v116, v118); + 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> - v127 = Tile< + v128 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v128 = (uint64_t)v20; - TASSIGN(v127, v128); + 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(v127, v127, v120, v123); + TMATMUL_ACC(v128, v128, v121, v124); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -543,17 +544,17 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v129 = Tile< + v130 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v130 = (uint64_t)v20; - TASSIGN(v129, v130); - pto::Shape<1, 1, 1, 16, 1024> v131 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v132 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v133 = GlobalTensor< + v134 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v20 + v20 * v13 + v6 * v14), v131, v132 + v3 + (v21 + v21 * v12 + v6 * v13), v132, v133 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); TSTORE< @@ -561,7 +562,7 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ 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>(v133, v129); + AtomicType::AtomicAdd>(v134, v130); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); #endif // __DAV_CUBE__ @@ -571,50 +572,50 @@ down_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_tile_inline238__rv_v2 - __gm__ Tensor *mlp_tile_inline238__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_tile_inline238__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_tile_inline238__rv_v2_tensor->buffer.addr) + - mlp_tile_inline238__rv_v2_tensor->start_offset; + // 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_inline251__iter_v6 - __gm__ Tensor *down_acc_all_inline251__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *down_acc_all_inline251__iter_v6 = - reinterpret_cast<__gm__ float *>(down_acc_all_inline251__iter_v6_tensor->buffer.addr) + - down_acc_all_inline251__iter_v6_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_inline202__ssa_v11 + // Unpack scalar: k0_inline113__ssa_v8 union { uint64_t u64; int64_t val; - } k0_inline202__ssa_v11_conv; - k0_inline202__ssa_v11_conv.u64 = args[3]; - int64_t k0_inline202__ssa_v11 = k0_inline202__ssa_v11_conv.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_inline231__ssa_v0 + // Unpack scalar: layer_inter_base_inline107__ssa_v0 union { uint64_t u64; int64_t val; - } layer_inter_base_inline231__ssa_v0_conv; - layer_inter_base_inline231__ssa_v0_conv.u64 = args[4]; - int64_t layer_inter_base_inline231__ssa_v0 = layer_inter_base_inline231__ssa_v0_conv.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_inline142__ssa_v8 + // Unpack scalar: n0_inline122__ssa_v8 union { uint64_t u64; int64_t val; - } n0_inline142__ssa_v8_conv; - n0_inline142__ssa_v8_conv.u64 = args[5]; - int64_t n0_inline142__ssa_v8 = n0_inline142__ssa_v8_conv.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_inline238__rv_v2, w_down__ssa_v0, down_acc_all_inline251__iter_v6, k0_inline202__ssa_v11, - layer_inter_base_inline231__ssa_v0, n0_inline142__ssa_v8 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/fa_fused_aic.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/fa_fused_aic.cpp deleted file mode 100644 index f5d9ba49b7..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/fa_fused_aic.cpp +++ /dev/null @@ -1,767 +0,0 @@ -/* - * 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: fa_fused_aic -// 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 fa_fused_aic( - __gm__ int32_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ int32_t *v5, __gm__ int32_t *v6, - __gm__ int32_t *v7, __gm__ bfloat16_t *v8, __gm__ bfloat16_t *v9, __gm__ bfloat16_t *v10, __gm__ float *v11, - int64_t v12, int64_t v13, int32_t v14, int32_t v15 -) { - const int64_t v16 = 5; - const int64_t v17 = 8; - const int64_t v18 = 43; - const int64_t v19 = 24; - const int64_t v20 = 16; - const int64_t v21 = 688; - const int64_t v22 = 128; - const int64_t v23 = 1; - const int64_t v24 = 0; - const int64_t v25 = 36864; - const int64_t v26 = 32768; - const int64_t v27 = 7; - const int32_t v28 = 0; - using T = float; - -#if defined(__DAV_CUBE__) - size_t v29 = (size_t)v23; - auto v30 = TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>(v11, v28, v28); - int32_t v31 = v1[v24]; - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - for (size_t v32 = (size_t)((int64_t)v14); v32 < ((size_t)((int64_t)v31)); v32 += (size_t)v19) { - int32_t v33 = v5[v32]; - int64_t v34 = (int64_t)v33; - int64_t v35 = v34 / v18; - int64_t v36 = v34 % v18; - int32_t v37 = v7[(int64_t)((uint64_t)((int64_t)(uint64_t)v35 * (uint64_t)v12) + (uint64_t)v36)]; - int64_t v38 = (int64_t)((uint64_t)v35 * (uint64_t)v22); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v39 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v40 = (uint64_t)v26; - TASSIGN(v39, v40); - pto::Shape<1, 1, 1, 16, 128> v41 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<2048, 2048, 2048, 128, 1> v42 = pto::Stride<2048, 2048, 2048, 128, 1>(); - GlobalTensor, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> - v43 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>( - v8 + (v24 + v38 * v22 + v24 * v23), v41, v42 - ); - TLOAD(v39, v43); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - int64_t v44 = (int64_t)((uint64_t)((int64_t)v37) * (uint64_t)v17); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v45 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null>(v22, v22); - uint64_t v46 = (uint64_t)v25; - TASSIGN(v45, v46); - pto::Shape<1, 1, 1, 128, 128> v47 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<128, 128, 128, 1, 128> v48 = pto::Stride<128, 128, 128, 1, 128>(); - GlobalTensor, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> - v49 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>( - v9 + (v24 + v24 * v23 + - (int64_t)((uint64_t)v13 + (uint64_t)((int64_t)(uint64_t)v44 * (uint64_t)v22)) * v22), - v47, v48 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - TLOAD(v45, v49); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v50 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v51 = (uint64_t)v24; - TASSIGN(v50, v51); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - TMOV(v50, v39); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v52 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v53 = (uint64_t)v24; - TASSIGN(v52, v53); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - TMOV(v52, v45); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v54 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v55 = (uint64_t)v24; - TASSIGN(v54, v55); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - TMATMUL(v54, v50, v52); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v54); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID1); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID1); - for (size_t v56 = v29; v56 < ((size_t)v17); v56 += v29) { - int64_t v57 = (int64_t)v56; - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v58 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v20, v22); - uint64_t v59 = (uint64_t)v26; - TASSIGN(v58, v59); - pto::Shape<1, 1, 1, 16, 128> v60 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<2048, 2048, 2048, 128, 1> v61 = pto::Stride<2048, 2048, 2048, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> - v62 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>( - v8 + (v24 + (int64_t)((uint64_t)v38 + (uint64_t)((int64_t)(uint64_t)v57 * (uint64_t)v20)) * v22 + - v24 * v23), - v60, v61 - ); - TLOAD(v58, v62); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v63 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v64 = (uint64_t)v25; - TASSIGN(v63, v64); - pto::Shape<1, 1, 1, 128, 128> v65 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<128, 128, 128, 1, 128> v66 = pto::Stride<128, 128, 128, 1, 128>(); - GlobalTensor, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> - v67 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>( - v9 + (v24 + v24 * v23 + - (int64_t)((uint64_t)v13 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v57) * - (uint64_t)v22)) * - v22), - v65, v66 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v63, v67); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v68 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v20, v22); - uint64_t v69 = (uint64_t)v24; - TASSIGN(v68, v69); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TMOV(v68, v58); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null> - v70 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v71 = (uint64_t)v24; - TASSIGN(v70, v71); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - TMOV(v70, v63); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v72 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v73 = (uint64_t)v24; - TASSIGN(v72, v73); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - TMATMUL(v72, v68, v70); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID1); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v72); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID3); - int64_t v74 = (int64_t)((uint64_t)v57 - (uint64_t)v23); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v75 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v76 = (uint64_t)v25; - TASSIGN(v75, v76); - pto::Shape<1, 1, 1, 128, 128> v77 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<16384, 16384, 16384, 128, 1> v78 = pto::Stride<16384, 16384, 16384, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND> - v79 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, - pto::Layout::ND>( - v10 + (v24 + - (int64_t)((uint64_t)v13 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v74) * - (uint64_t)v22)) * - v22 + - v24 * v23), - v77, v78 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v75, v79); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v80 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>(v20, v22); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v80); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v81 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>(v20, v22); - uint64_t v82 = (uint64_t)v24; - TASSIGN(v81, v82); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); - TMOV(v81, v80); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v30); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null> - v83 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v84 = (uint64_t)v24; - TASSIGN(v83, v84); - TMOV(v83, v75); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v85 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v86 = (uint64_t)v24; - TASSIGN(v85, v86); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID3); - TMATMUL(v85, v81, v83); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID2); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - v85.SetValidShape(v16, v22); - pto::Shape<1, 1, 1, 5, 128> v87 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v88 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v89 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v4 + (v24 + - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v35 * - (uint64_t)v17) + - (uint64_t)v74) * - (uint64_t)v21) + - (uint64_t)((int64_t)(uint64_t)v36 * (uint64_t)v20)) * - v22 + - v24 * v23), - v87, v88 - ); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID2); - TSTORE(v89, v85); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - }; - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID4); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v90 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v22, v22); - uint64_t v91 = (uint64_t)v25; - TASSIGN(v90, v91); - pto::Shape<1, 1, 1, 128, 128> v92 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<16384, 16384, 16384, 128, 1> v93 = pto::Stride<16384, 16384, 16384, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND> - v94 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND>( - v10 + - (v24 + - (int64_t)((uint64_t)v13 + (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v27) * - (uint64_t)v22)) * - v22 + - v24 * v23), - v92, v93 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); - TLOAD(v90, v94); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v95 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>(v20, v22); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v95); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v96 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>(v20, v22); - uint64_t v97 = (uint64_t)v24; - TASSIGN(v96, v97); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - TMOV(v96, v95); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v30); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v98 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v99 = (uint64_t)v24; - TASSIGN(v98, v99); - TMOV(v98, v90); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v100 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v101 = (uint64_t)v24; - TASSIGN(v100, v101); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID4); - TMATMUL(v100, v96, v98); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID3); - v100.SetValidShape(v16, v22); - pto::Shape<1, 1, 1, 5, 128> v102 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v103 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v104 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v4 + (v24 + - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v35 * - (uint64_t)v17) + - (uint64_t)v27) * - (uint64_t)v21) + - (uint64_t)((int64_t)(uint64_t)v36 * (uint64_t)v20)) * - v22 + - v24 * v23), - v102, v103 - ); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID3); - TSTORE(v104, v100); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - } - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2); -#endif // __DAV_CUBE__ - - ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); - return; -} - -static __aicore__ void fa_fused_aiv( - __gm__ int32_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ int32_t *v5, __gm__ int32_t *v6, - __gm__ int32_t *v7, __gm__ bfloat16_t *v8, __gm__ bfloat16_t *v9, __gm__ bfloat16_t *v10, __gm__ float *v11, - int64_t v12, int64_t v13, int32_t v14, int32_t v15 -) { - SaturationMode v16 = SaturationMode::OFF; - RoundMode v17 = RoundMode::CAST_ROUND; - const int64_t v18 = 5; - const float v19 = 0.0883883461f; - const int64_t v20 = 8; - const int64_t v21 = 43; - const int64_t v22 = 24; - const int64_t v23 = 0; - const int64_t v24 = 16; - const int64_t v25 = 688; - const int64_t v26 = 128; - const int64_t v27 = 88064; - const int64_t v28 = 1; - const int64_t v29 = 43040; - const int64_t v30 = 40992; - const int64_t v31 = 40960; - const int64_t v32 = 32768; - const int64_t v33 = 36864; - const int32_t v34 = 0; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - int64_t v35 = get_subblockid(); - auto v36 = TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>(v11, v34, v34); - int32_t v37 = v1[v23]; - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - for (size_t v38 = (size_t)((int64_t)v14); v38 < ((size_t)((int64_t)v37)); v38 += (size_t)v22) { - int32_t v39 = v5[v38]; - int64_t v40 = (int64_t)v39; - int64_t v41 = v40 / v21; - int64_t v42 = v40 % v21; - int32_t v43 = v6[v41]; - int64_t v44 = (int64_t)((uint64_t)((int64_t)v43) - (uint64_t)((int64_t)(uint64_t)v42 * (uint64_t)v26)); - int64_t v45 = v44 < v26 ? v44 : v26; - for (size_t v46 = (size_t)v23; v46 < ((size_t)v20); v46 += (size_t)v28) { - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v47 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v36, v47); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v48 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v49 = (uint64_t)v33; - TASSIGN(v48, v49); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - pipe_barrier(PIPE_V); - TMULS(v48, v47, v19); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v36); - v48.SetValidShape(v18, v45); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v50 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v51 = (uint64_t)v32; - TASSIGN(v50, v51); - pipe_barrier(PIPE_V); - TFILLPAD(v50, v48); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v52 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v53 = (uint64_t)v33; - TASSIGN(v52, v53); - Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v54 = Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v28); - uint64_t v55 = (uint64_t)v31; - TASSIGN(v54, v55); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TROWMAX(v54, v50, v52); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v56 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v57 = (uint64_t)v32; - TASSIGN(v56, v57); - pipe_barrier(PIPE_V); - TROWEXPANDSUB(v56, v50, v54); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v58 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v59 = (uint64_t)v32; - TASSIGN(v58, v59); - pipe_barrier(PIPE_V); - TEXP(v58, v56); - Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v60 = Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v61 = (uint64_t)v30; - TASSIGN(v60, v61); - pipe_barrier(PIPE_V); - TCVT(v60, v58, v17, v16); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v62 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v63 = (uint64_t)v32; - TASSIGN(v62, v63); - pipe_barrier(PIPE_V); - TCVT(v62, v60, v17, v16); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v64 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v65 = (uint64_t)v33; - TASSIGN(v64, v65); - Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v66 = Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v28); - uint64_t v67 = (uint64_t)v29; - TASSIGN(v66, v67); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TROWSUM(v66, v62, v64); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v36, v60); - int64_t v68 = - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t) - v41 * - (uint64_t)v20) + - (uint64_t)((int64_t)v46)) * - (uint64_t)v25) + - (uint64_t)((int64_t)(uint64_t)v42 * (uint64_t)v24)) + - (uint64_t)((int64_t)(uint64_t)((int64_t)v35) * (uint64_t)v20)); - pto::Shape<1, 1, 1, 8, 1> v69 = pto::Shape<1, 1, 1, 8, 1>(); - pto::Stride<8, 8, 8, 1, 88064> v70 = pto::Stride<8, 8, 8, 1, 88064>(); - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN> v71 = - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN>( - v3 + (v23 + v68 * v28 + v23 * v27), v69, v70 - ); - TSTORE(v71, v54); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pto::Shape<1, 1, 1, 8, 1> v72 = pto::Shape<1, 1, 1, 8, 1>(); - pto::Stride<8, 8, 8, 1, 88064> v73 = pto::Stride<8, 8, 8, 1, 88064>(); - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN> v74 = - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN>( - v2 + (v23 + v68 * v28 + v23 * v27), v72, v73 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v74, v66); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - }; - } - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - 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) { - // 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: fa_total_inline183__ssa_v0 - __gm__ Tensor *fa_total_inline183__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ int32_t *fa_total_inline183__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_total_inline183__ssa_v0_tensor->buffer.addr) + - fa_total_inline183__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_cur_li_inline190__ssa_v0 - __gm__ Tensor *all_cur_li_inline190__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *all_cur_li_inline190__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_cur_li_inline190__ssa_v0_tensor->buffer.addr) + - all_cur_li_inline190__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_cur_mi_inline174__ssa_v0 - __gm__ Tensor *all_cur_mi_inline174__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *all_cur_mi_inline174__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_cur_mi_inline174__ssa_v0_tensor->buffer.addr) + - all_cur_mi_inline174__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_oi_tmp_inline145__ssa_v0 - __gm__ Tensor *all_oi_tmp_inline145__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *all_oi_tmp_inline145__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_oi_tmp_inline145__ssa_v0_tensor->buffer.addr) + - all_oi_tmp_inline145__ssa_v0_tensor->start_offset; - - // Unpack tensor: fa_work_table_inline126__ssa_v0 - __gm__ Tensor *fa_work_table_inline126__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ int32_t *fa_work_table_inline126__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_work_table_inline126__ssa_v0_tensor->buffer.addr) + - fa_work_table_inline126__ssa_v0_tensor->start_offset; - - // Unpack tensor: seq_lens__ssa_v0 - __gm__ Tensor *seq_lens__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[5]); - __gm__ int32_t *seq_lens__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(seq_lens__ssa_v0_tensor->buffer.addr) + - seq_lens__ssa_v0_tensor->start_offset; - - // Unpack tensor: block_table__ssa_v0 - __gm__ Tensor *block_table__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[6]); - __gm__ int32_t *block_table__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(block_table__ssa_v0_tensor->buffer.addr) + - block_table__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_q_padded_inline168__rv_v2 - __gm__ Tensor *all_q_padded_inline168__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[7]); - __gm__ bfloat16_t *all_q_padded_inline168__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(all_q_padded_inline168__rv_v2_tensor->buffer.addr) + - all_q_padded_inline168__rv_v2_tensor->start_offset; - - // Unpack tensor: k_cache__rv_v4 - __gm__ Tensor *k_cache__rv_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[8]); - __gm__ bfloat16_t *k_cache__rv_v4 = - reinterpret_cast<__gm__ bfloat16_t *>(k_cache__rv_v4_tensor->buffer.addr) + k_cache__rv_v4_tensor->start_offset; - - // Unpack tensor: v_cache__rv_v4 - __gm__ Tensor *v_cache__rv_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[9]); - __gm__ bfloat16_t *v_cache__rv_v4 = - reinterpret_cast<__gm__ bfloat16_t *>(v_cache__rv_v4_tensor->buffer.addr) + v_cache__rv_v4_tensor->start_offset; - - // Unpack tensor: __gm_pipe_buffer - __gm__ Tensor *__gm_pipe_buffer_tensor = reinterpret_cast<__gm__ Tensor *>(args[10]); - // SPMD: shard GM pipe workspace by logical block_idx to avoid overlap. - int64_t __pypto_gm_block_num = static_cast(__pypto_spmd_block_num); - if (__pypto_gm_block_num <= 0) __pypto_gm_block_num = 1; - int64_t __pypto_gm_total_elems = static_cast(__gm_pipe_buffer_tensor->shapes[0]); - int64_t __pypto_gm_elems_per_block = __pypto_gm_total_elems / __pypto_gm_block_num; - int64_t __pypto_gm_block_offset = static_cast(__pypto_spmd_block_idx) * __pypto_gm_elems_per_block; - __gm__ float *__gm_pipe_buffer = reinterpret_cast<__gm__ float *>(__gm_pipe_buffer_tensor->buffer.addr) + - __gm_pipe_buffer_tensor->start_offset + __pypto_gm_block_offset; - - // Unpack scalar: max_blocks_per_seq_inline165__ssa_v0 - union { - uint64_t u64; - int64_t val; - } max_blocks_per_seq_inline165__ssa_v0_conv; - max_blocks_per_seq_inline165__ssa_v0_conv.u64 = args[11]; - int64_t max_blocks_per_seq_inline165__ssa_v0 = max_blocks_per_seq_inline165__ssa_v0_conv.val; - - // Unpack scalar: layer_cache_base_inline137__ssa_v0 - union { - uint64_t u64; - int64_t val; - } layer_cache_base_inline137__ssa_v0_conv; - layer_cache_base_inline137__ssa_v0_conv.u64 = args[12]; - int64_t layer_cache_base_inline137__ssa_v0 = layer_cache_base_inline137__ssa_v0_conv.val; - - // Forward to ptoas-generated function - fa_fused_aic( - fa_total_inline183__ssa_v0, all_cur_li_inline190__ssa_v0, all_cur_mi_inline174__ssa_v0, - all_oi_tmp_inline145__ssa_v0, fa_work_table_inline126__ssa_v0, seq_lens__ssa_v0, block_table__ssa_v0, - all_q_padded_inline168__rv_v2, k_cache__rv_v4, v_cache__rv_v4, __gm_pipe_buffer, - max_blocks_per_seq_inline165__ssa_v0, layer_cache_base_inline137__ssa_v0, __pypto_spmd_block_idx, - __pypto_spmd_block_num - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj.cpp index 53dc55c329..de4335aeac 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj.cpp @@ -27,6 +27,7 @@ #include #include "tensor.h" +#include "intrinsic.h" using namespace pto; @@ -50,149 +51,152 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void -gate_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 = 1024; +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 = 2048; - const int64_t v18 = 32768; - const int64_t v19 = 512; - const int64_t v20 = 0; - const int64_t v21 = 135168; - const int64_t v22 = 4096; + 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 v23 = (size_t)v12; - size_t v24 = (size_t)v20; - size_t v25 = (size_t)v10; + 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> - v26 = Tile< + v29 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v27 = (uint64_t)v22; - TASSIGN(v26, v27); - pto::Shape<1, 1, 1, 16, 64> v28 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v30 = GlobalTensor< + v33 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + v4 * v14), v28, v29 + 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(v26, v30); + TLOAD(v29, v33); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v31 = Tile< + v34 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v32 = (uint64_t)v21; - TASSIGN(v31, v32); - int64_t v33 = (int64_t)((uint64_t)v5 + (uint64_t)v4); - pto::Shape<1, 1, 1, 64, 1024> v34 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v35 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v36 = GlobalTensor< + v39 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + v33 * v13 + v6 * v14), v34, v35); - TLOAD(v31, v36); + 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 v37 = v24; v37 < v23; v37 += v25) { - int64_t v38 = (int64_t)v37; + 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> - v39 = Tile< + v42 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v40 = (uint64_t)v20; - TASSIGN(v39, v40); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); pipe_barrier(PIPE_MTE1); - TEXTRACT(v39, v26, v20, v37); + TEXTRACT(v42, v29, v22, v40); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v41 = Tile< + v44 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v42 = (uint64_t)v20; - TASSIGN(v41, v42); - TEXTRACT(v41, v31, v37, v20); + 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> - v43 = Tile< + v46 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v44 = (uint64_t)v19; - TASSIGN(v43, v44); - int64_t v45 = (int64_t)((uint64_t)v38 + (uint64_t)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(v43, v26, v20, v45); + TEXTRACT(v46, v29, v22, v48); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v46 = Tile< + v49 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v47 = (uint64_t)v18; - TASSIGN(v46, v47); - TEXTRACT(v46, v31, v45, v20); + 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 (v38 == v20) { + if (v41 == v22) { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v48 = Tile< + v51 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v49 = (uint64_t)v20; - TASSIGN(v48, v49); + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); pipe_barrier(PIPE_M); - TMATMUL(v48, v39, v41); + TMATMUL(v51, v42, v44); } else { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v50 = Tile< + v53 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v51 = (uint64_t)v20; - TASSIGN(v50, v51); + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); pipe_barrier(PIPE_M); - TMATMUL_ACC(v50, v50, v39, v41); + 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> - v52 = Tile< + v55 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v53 = (uint64_t)v20; - TASSIGN(v52, v53); + 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(v52, v52, v43, v46); + TMATMUL_ACC(v55, v55, v46, v49); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -202,151 +206,151 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ 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 v54 = (size_t)v14; v54 < ((size_t)v9); v54 += (size_t)v8) { - int64_t v55 = (int64_t)((uint64_t)((int64_t)v54) * (uint64_t)v12); - int64_t v56 = (int64_t)((uint64_t)v55 + (uint64_t)v12); + 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> - v57 = Tile< + v60 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v58 = (uint64_t)v20; - TASSIGN(v57, v58); - pto::Shape<1, 1, 1, 16, 64> v59 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v60 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v61 = GlobalTensor< + v64 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v55) * v14), v59, v60 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - TLOAD(v57, v61); + TLOAD(v60, v64); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v62 = Tile< + v65 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v63 = (uint64_t)v21; - TASSIGN(v62, v63); - pto::Shape<1, 1, 1, 64, 1024> v64 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v65 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v66 = GlobalTensor< + v69 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v55) * v13 + v6 * v14), v64, v65); - TLOAD(v62, v66); + 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> - v67 = Tile< + v70 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v68 = (uint64_t)v17; - TASSIGN(v67, v68); - pto::Shape<1, 1, 1, 16, 64> v69 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v70 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v71 = GlobalTensor< + v74 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v14), v69, v70 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v67, v71); + TLOAD(v70, v74); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v72 = Tile< + v75 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v73 = (uint64_t)v22; - TASSIGN(v72, v73); - pto::Shape<1, 1, 1, 64, 1024> v74 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v75 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v76 = GlobalTensor< + v79 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v56) * v13 + v6 * v14), v74, v75); - TLOAD(v72, v76); + 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_M, PIPE_MTE1, EVENT_ID3); 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 v77 = v24; v77 < v23; v77 += v25) { + 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> - v78 = Tile< + v81 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v79 = (uint64_t)v20; - TASSIGN(v78, v79); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); pipe_barrier(PIPE_MTE1); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - TEXTRACT(v78, v57, v20, v77); + TEXTRACT(v81, v60, v22, v80); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v80 = Tile< + v83 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v81 = (uint64_t)v20; - TASSIGN(v80, v81); - TEXTRACT(v80, v62, v77, v20); + 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> - v82 = Tile< + v85 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v83 = (uint64_t)v19; - TASSIGN(v82, v83); - int64_t v84 = (int64_t)((uint64_t)((int64_t)v77) + (uint64_t)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(v82, v57, v20, v84); + TEXTRACT(v85, v60, v22, v87); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v85 = Tile< + v88 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v86 = (uint64_t)v18; - TASSIGN(v85, v86); - TEXTRACT(v85, v62, v84, v20); + 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> - v87 = Tile< + v90 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v88 = (uint64_t)v20; - TASSIGN(v87, v88); + 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(v87, v87, v78, v80); + 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> - v89 = Tile< + v92 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v90 = (uint64_t)v20; - TASSIGN(v89, v90); + 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(v89, v89, v82, v85); + TMATMUL_ACC(v92, v92, v85, v88); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); @@ -357,71 +361,71 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ 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 v91 = v24; v91 < v23; v91 += v25) { + 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> - v92 = Tile< + v95 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v93 = (uint64_t)v20; - TASSIGN(v92, v93); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - TEXTRACT(v92, v67, v20, v91); + TEXTRACT(v95, v70, v22, v94); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v94 = Tile< + v97 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v95 = (uint64_t)v20; - TASSIGN(v94, v95); - TEXTRACT(v94, v72, v91, v20); + 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> - v96 = Tile< + v99 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v97 = (uint64_t)v19; - TASSIGN(v96, v97); - int64_t v98 = (int64_t)((uint64_t)((int64_t)v91) + (uint64_t)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(v96, v67, v20, v98); + TEXTRACT(v99, v70, v22, v101); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v99 = Tile< + v102 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v100 = (uint64_t)v18; - TASSIGN(v99, v100); - TEXTRACT(v99, v72, v98, v20); + 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> - v101 = Tile< + v104 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v102 = (uint64_t)v20; - TASSIGN(v101, v102); + 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(v101, v101, v92, v94); + 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> - v103 = Tile< + v106 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v104 = (uint64_t)v20; - TASSIGN(v103, v104); + 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(v103, v103, v96, v99); + TMATMUL_ACC(v106, v106, v99, v102); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); @@ -435,106 +439,106 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v105 = Tile< + v108 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v106 = (uint64_t)v22; - TASSIGN(v105, v106); - pto::Shape<1, 1, 1, 16, 64> v107 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v108 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v109 = GlobalTensor< + v112 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v14), v107, v108 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v105, v109); + TLOAD(v108, v112); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v110 = Tile< + v113 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v111 = (uint64_t)v21; - TASSIGN(v110, v111); - pto::Shape<1, 1, 1, 64, 1024> v112 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v113 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v114 = GlobalTensor< + v117 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v7) * v13 + v6 * v14), v112, v113); - TLOAD(v110, v114); + 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_M, PIPE_MTE1, EVENT_ID0); 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 v115 = v24; v115 < v23; v115 += v25) { + 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> - v116 = Tile< + v119 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v117 = (uint64_t)v20; - TASSIGN(v116, v117); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - TEXTRACT(v116, v105, v20, v115); + TEXTRACT(v119, v108, v22, v118); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v118 = Tile< + v121 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v119 = (uint64_t)v20; - TASSIGN(v118, v119); - TEXTRACT(v118, v110, v115, v20); + 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> - v120 = Tile< + v123 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v121 = (uint64_t)v19; - TASSIGN(v120, v121); - int64_t v122 = (int64_t)((uint64_t)((int64_t)v115) + (uint64_t)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(v120, v105, v20, v122); + TEXTRACT(v123, v108, v22, v125); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v123 = Tile< + v126 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v124 = (uint64_t)v18; - TASSIGN(v123, v124); - TEXTRACT(v123, v110, v122, v20); + 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> - v125 = Tile< + v128 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v126 = (uint64_t)v20; - TASSIGN(v125, v126); + 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(v125, v125, v116, v118); + 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> - v127 = Tile< + v130 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v128 = (uint64_t)v20; - TASSIGN(v127, v128); + 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(v127, v127, v120, v123); + TMATMUL_ACC(v130, v130, v123, v126); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -543,17 +547,17 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v129 = Tile< + v132 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v130 = (uint64_t)v20; - TASSIGN(v129, v130); - pto::Shape<1, 1, 1, 16, 1024> v131 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v132 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v133 = GlobalTensor< + v136 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v3 + (v20 + v20 * v13 + v6 * v14), v131, v132 + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); TSTORE< @@ -562,7 +566,7 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ CompactMode::Null>, GlobalTensor< float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, - AtomicType::AtomicAdd>(v133, v129); + AtomicType::AtomicAdd>(v136, v132); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); #endif // __DAV_CUBE__ @@ -572,50 +576,46 @@ gate_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_ } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v14 - __gm__ Tensor *mlp_norm_in_inline250__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v14 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v14_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v14_tensor->start_offset; + // 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_inline170__iter_v6 - __gm__ Tensor *gate_acc_all_inline170__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *gate_acc_all_inline170__iter_v6 = - reinterpret_cast<__gm__ float *>(gate_acc_all_inline170__iter_v6_tensor->buffer.addr) + - gate_acc_all_inline170__iter_v6_tensor->start_offset; - - // Unpack scalar: k0_inline202__ssa_v10 - union { - uint64_t u64; - int64_t val; - } k0_inline202__ssa_v10_conv; - k0_inline202__ssa_v10_conv.u64 = args[3]; - int64_t k0_inline202__ssa_v10 = k0_inline202__ssa_v10_conv.val; + // 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: layer_hidden_base_inline135__ssa_v0 + // Unpack scalar: gu_k0_inline131__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[4]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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: n0_inline142__ssa_v6 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } n0_inline142__ssa_v6_conv; - n0_inline142__ssa_v6_conv.u64 = args[5]; - int64_t n0_inline142__ssa_v6 = n0_inline142__ssa_v6_conv.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_inline250__rv_v14, w_gate__ssa_v0, gate_acc_all_inline170__iter_v6, k0_inline202__ssa_v10, - layer_hidden_base_inline135__ssa_v0, n0_inline142__ssa_v6 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp new file mode 100644 index 0000000000..e05dc60546 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp new file mode 100644 index 0000000000..a5688a2059 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp new file mode 100644 index 0000000000..1b7ecba2cc --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp new file mode 100644 index 0000000000..4012cd6415 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp new file mode 100644 index 0000000000..caf3aa32c0 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/k_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/k_proj.cpp index 35abc291e0..79666ceefb 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/k_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/k_proj.cpp @@ -53,415 +53,512 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT 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 = 5; + 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 = 1280; - const int64_t v13 = 512; - const int64_t v14 = 4; - const int64_t v15 = 5120; - const int64_t v16 = 1; - const int64_t v17 = 1024; - const int64_t v18 = 16; + 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 = 2048; - const int64_t v22 = 0; - const int64_t v23 = 147456; - const int64_t v24 = 16384; + 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 v25 = (size_t)v16; - size_t v26 = (size_t)v22; - size_t v27 = (size_t)v11; - size_t v28 = (size_t)v10; - size_t v29 = (size_t)v9; - int64_t v30 = (int64_t)v5; - int64_t v31 = (int64_t)((uint64_t)(v30 % v14) * (uint64_t)v12); + 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_MTE1, PIPE_MTE2, EVENT_ID2); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - for (size_t v32 = v26; v32 < v27; v32 += v25) { - int64_t v33 = (int64_t)((uint64_t)((int64_t)(uint64_t)(v30 / v14) * (uint64_t)v13) + - (uint64_t)((int64_t)(uint64_t)((int64_t)v32) * (uint64_t)v10)); + 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> - v34 = Tile< + v33 = Tile< TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v35 = (uint64_t)v24; - 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>(); + 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> - v38 = GlobalTensor< + v37 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v22 + v22 * v15 + v31 * v16), v36, v37 + v2 + (v23 + v23 * v14 + v30 * v15), v35, v36 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - TLOAD(v34, v38); + TLOAD(v33, v37); Tile< TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v39 = Tile< + v38 = Tile< TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v40 = (uint64_t)v23; - TASSIGN(v39, v40); - pto::Shape<1, 1, 1, 256, 256> v41 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v42 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + 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> - v43 = GlobalTensor< + v42 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v31) * v17 + v33 * v16), v41, v42); - TLOAD(v39, v43); + 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); - Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v44 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v45 = (uint64_t)v22; - TASSIGN(v44, v45); 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 v46 = v26; v46 < v28; v46 += v29) { - int64_t v47 = (int64_t)v46; + 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> - v48 = Tile< + v45 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v49 = (uint64_t)v22; - TASSIGN(v48, v49); + 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(v48, v34, v22, v46); + TEXTRACT(v45, v33, v23, v43); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v50 = Tile< + v47 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v51 = (uint64_t)v22; - TASSIGN(v50, v51); - TEXTRACT(v50, v39, v46, v22); + 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> - v52 = Tile< + v49 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v53 = (uint64_t)v21; - TASSIGN(v52, v53); - int64_t v54 = (int64_t)((uint64_t)v47 + (uint64_t)v8); + 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(v52, v34, v22, v54); + TEXTRACT(v49, v33, v23, v51); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v55 = Tile< + v52 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v56 = (uint64_t)v20; - TASSIGN(v55, v56); - TEXTRACT(v55, v39, v54, v22); + 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 (v47 == v22) { + if (v44 == v23) { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v57 = Tile< + v54 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v58 = (uint64_t)v22; - TASSIGN(v57, v58); + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); pipe_barrier(PIPE_M); - TMATMUL(v57, v48, v50); + TMATMUL(v54, v45, v47); } else { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v59 = Tile< + v56 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); pipe_barrier(PIPE_M); - TMATMUL_ACC(v59, v59, v48, v50); + 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> - v61 = Tile< + v58 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v62 = (uint64_t)v22; - TASSIGN(v61, v62); + 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(v61, v61, v52, v55); + 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 v63 = v25; v63 < ((size_t)v7); v63 += v27) { - int64_t v64 = (int64_t)((uint64_t)((int64_t)v63) * (uint64_t)v10); - int64_t v65 = (int64_t)((uint64_t)v31 + (uint64_t)v64); - int64_t v66 = (int64_t)((uint64_t)v31 + (uint64_t)((int64_t)(uint64_t)v64 + (uint64_t)v10)); + for (size_t v82 = v26; v82 < v27; v82 += v28) { Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v67 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pto::Shape<1, 1, 1, 16, 256> v69 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v70 = 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> - v71 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v22 + v22 * v15 + v65 * v16), v69, v70); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v67, v71); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v72 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v73 = (uint64_t)v23; - TASSIGN(v72, v73); - pto::Shape<1, 1, 1, 256, 256> v74 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v75 = 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> - v76 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v65) * v17 + v33 * v16), v74, v75); - TLOAD(v72, v76); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + 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::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v77 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v78 = (uint64_t)v19; - TASSIGN(v77, v78); - pto::Shape<1, 1, 1, 16, 256> v79 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v80 = 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> - v81 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v22 + v22 * v15 + v66 * v16), v79, v80); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v77, v81); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v82 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v83 = (uint64_t)v24; - TASSIGN(v82, v83); - pto::Shape<1, 1, 1, 256, 256> v84 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v85 = 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> - v86 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v66) * v17 + v33 * v16), v84, v85); - TLOAD(v82, v86); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - for (size_t v87 = v26; v87 < v28; v87 += v29) { - Tile< + }; + 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> - v88 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v89 = (uint64_t)v22; - TASSIGN(v88, v89); - pipe_barrier(PIPE_MTE1); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - TEXTRACT(v88, v67, v22, v87); - Tile< + 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> - 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)v22; - TASSIGN(v90, v91); - TEXTRACT(v90, v72, v87, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - Tile< + 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> - v92 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v93 = (uint64_t)v21; - TASSIGN(v92, v93); - int64_t v94 = (int64_t)((uint64_t)((int64_t)v87) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - TEXTRACT(v92, v67, v22, v94); - Tile< + 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> - v95 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v96 = (uint64_t)v20; - TASSIGN(v95, v96); - TEXTRACT(v95, v72, v94, v22); - 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> - v97 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v98 = (uint64_t)v22; - TASSIGN(v97, v98); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v97, v97, v88, v90); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - Tile< + 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> - v99 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v100 = (uint64_t)v22; - TASSIGN(v99, v100); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - TMATMUL_ACC(v99, v99, v92, v95); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - }; - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + 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); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - for (size_t v101 = v26; v101 < v28; v101 += v29) { - Tile< + }; + 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> - v102 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v103 = (uint64_t)v22; - TASSIGN(v102, v103); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - TEXTRACT(v102, v77, v22, v101); - Tile< + 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> - 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)v22; - TASSIGN(v104, v105); - TEXTRACT(v104, v82, v101, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - Tile< + 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> - v106 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v107 = (uint64_t)v21; - TASSIGN(v106, v107); - int64_t v108 = (int64_t)((uint64_t)((int64_t)v101) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TEXTRACT(v106, v77, v22, v108); - Tile< + 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> - v109 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v110 = (uint64_t)v20; - TASSIGN(v109, v110); - TEXTRACT(v109, v82, v108, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - Tile< + 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> - v111 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v112 = (uint64_t)v22; - TASSIGN(v111, v112); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v111, v111, v102, v104); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - Tile< + 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> - v113 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v114 = (uint64_t)v22; - TASSIGN(v113, v114); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - TMATMUL_ACC(v113, v113, v106, v109); - 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_M, PIPE_MTE1, EVENT_ID4); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + 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_ID4); + 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); - pto::Shape<1, 1, 1, 16, 256> v115 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v116 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + 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> - v117 = GlobalTensor< + v139 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v1 + (v22 + v22 * v17 + v33 * v16), v115, v116 + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); pipe_barrier(PIPE_FIX); @@ -471,14 +568,12 @@ k_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v CompactMode::Null>, GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>, - AtomicType::AtomicAdd>(v117, v44); + 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); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); #endif // __DAV_CUBE__ ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); @@ -490,34 +585,33 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in int32_t __pypto_spmd_block_idx = get_block_idx(args); int32_t __pypto_spmd_block_num = get_block_num(args); - // Unpack tensor: k_proj_inline128__ssa_v1 - __gm__ Tensor *k_proj_inline128__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *k_proj_inline128__ssa_v1 = - reinterpret_cast<__gm__ float *>(k_proj_inline128__ssa_v1_tensor->buffer.addr) + - k_proj_inline128__ssa_v1_tensor->start_offset; + // 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_states_inline146__rv_v2 - __gm__ Tensor *normed_states_inline146__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *normed_states_inline146__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(normed_states_inline146__rv_v2_tensor->buffer.addr) + - normed_states_inline146__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: 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_inline135__ssa_v0 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[3]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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_inline128__ssa_v1, normed_states_inline146__rv_v2, wk__ssa_v0, layer_hidden_base_inline135__ssa_v0, + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj.cpp index 0c9dce0b5d..1d35867e0e 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj.cpp @@ -61,29 +61,30 @@ out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t const int64_t v13 = 1; const int64_t v14 = 5120; const int64_t v15 = 16; - const int64_t v16 = 2048; + const int64_t v16 = 1024; const int64_t v17 = 32768; - const int64_t v18 = 1024; - const int64_t v19 = 0; - const int64_t v20 = 69632; - const int64_t v21 = 4096; + 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> - v22 = Tile< + v23 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v12); - uint64_t v23 = (uint64_t)v21; - TASSIGN(v22, v23); - pto::Shape<1, 1, 1, 16, 64> v24 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v25 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v26 = GlobalTensor< + v27 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v19 + v19 * v14 + v4 * v13), v24, v25 + v1 + (v20 + v20 * v14 + v4 * v13), v25, v26 ); set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); @@ -91,304 +92,304 @@ out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t 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(v22, v26); + 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> - v27 = Tile< + v28 = Tile< TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v12, v11); - uint64_t v28 = (uint64_t)v20; - TASSIGN(v27, v28); - int64_t v29 = (int64_t)((uint64_t)v5 + (uint64_t)v4); - pto::Shape<1, 1, 1, 64, 512> v30 = pto::Shape<1, 1, 1, 64, 512>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v31 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v32 = GlobalTensor< + v33 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( - v2 + (v19 + v29 * v14 + v6 * v13), v30, v31 + v2 + (v20 + v30 * v14 + v6 * v13), v31, v32 ); - TLOAD(v27, 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> - v33 = Tile< + v34 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v34 = (uint64_t)v19; - TASSIGN(v33, v34); + uint64_t v35 = (uint64_t)v19; + TASSIGN(v34, v35); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - TEXTRACT(v33, v22, v19, v19); + TEXTRACT(v34, v23, v20, v20); Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v35 = Tile< + v36 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v36 = (uint64_t)v19; - TASSIGN(v35, v36); + uint64_t v37 = (uint64_t)v20; + TASSIGN(v36, v37); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - TEXTRACT(v35, v27, v19, v19); + 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> - v37 = Tile< + v38 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v38 = (uint64_t)v18; - TASSIGN(v37, v38); - TEXTRACT(v37, v22, v19, 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> - v39 = Tile< + v40 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v40 = (uint64_t)v17; - TASSIGN(v39, v40); - TEXTRACT(v39, v27, v10, v19); + 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> - v41 = Tile< + v42 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v42 = (uint64_t)v19; - TASSIGN(v41, v42); + uint64_t v43 = (uint64_t)v20; + TASSIGN(v42, v43); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - TMATMUL(v41, v33, v35); + TMATMUL(v42, v34, v36); Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v43 = Tile< + v44 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v44 = (uint64_t)v19; - TASSIGN(v43, v44); + uint64_t v45 = (uint64_t)v20; + TASSIGN(v44, v45); pipe_barrier(PIPE_M); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - TMATMUL_ACC(v43, v43, v37, v39); + 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 v45 = (size_t)v13; v45 < ((size_t)v9); v45 += (size_t)v8) { - int64_t v46 = (int64_t)((uint64_t)((int64_t)v45) * (uint64_t)v12); - int64_t v47 = (int64_t)((uint64_t)v46 + (uint64_t)v12); + 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> - v48 = Tile< + v49 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v12); - uint64_t v49 = (uint64_t)v19; - TASSIGN(v48, v49); - pto::Shape<1, 1, 1, 16, 64> v50 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v51 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v52 = GlobalTensor< + v53 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v19 + v19 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v46) * v13), v50, v51 + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v47) * v13), v51, v52 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - TLOAD(v48, v52); + 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> - v53 = Tile< + v54 = Tile< TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v12, v11); - uint64_t v54 = (uint64_t)v20; - TASSIGN(v53, v54); - pto::Shape<1, 1, 1, 64, 512> v55 = pto::Shape<1, 1, 1, 64, 512>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v56 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v57 = GlobalTensor< + v58 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, - pto::Layout::ND>(v2 + (v19 + (int64_t)((uint64_t)v29 + (uint64_t)v46) * v14 + v6 * v13), v55, v56); + 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(v53, v57); + 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> - v58 = Tile< + v59 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v12); - uint64_t v59 = (uint64_t)v16; - 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>(); + 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> - v62 = GlobalTensor< + v63 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v19 + v19 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v47) * v13), v60, v61 + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v48) * v13), v61, v62 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v58, v62); + 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> - v63 = Tile< + v64 = Tile< TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v12, v11); - uint64_t v64 = (uint64_t)v21; - TASSIGN(v63, v64); - pto::Shape<1, 1, 1, 64, 512> v65 = pto::Shape<1, 1, 1, 64, 512>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v66 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v67 = GlobalTensor< + v68 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, - pto::Layout::ND>(v2 + (v19 + (int64_t)((uint64_t)v29 + (uint64_t)v47) * v14 + v6 * v13), v65, v66); + 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(v63, v67); + 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> - v68 = Tile< + v69 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v69 = (uint64_t)v19; - TASSIGN(v68, v69); + uint64_t v70 = (uint64_t)v19; + TASSIGN(v69, v70); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - TEXTRACT(v68, v48, v19, v19); + TEXTRACT(v69, v49, v20, v20); Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v70 = Tile< + v71 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v71 = (uint64_t)v19; - TASSIGN(v70, v71); + uint64_t v72 = (uint64_t)v20; + TASSIGN(v71, v72); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - TEXTRACT(v70, v53, v19, v19); + 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> - v72 = Tile< + v73 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v73 = (uint64_t)v18; - TASSIGN(v72, v73); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TEXTRACT(v72, v48, v19, 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> - v74 = Tile< + v75 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v75 = (uint64_t)v17; - TASSIGN(v74, v75); - TEXTRACT(v74, v53, v10, v19); + 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> - v76 = Tile< + v77 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v77 = (uint64_t)v19; - TASSIGN(v76, v77); + uint64_t v78 = (uint64_t)v20; + TASSIGN(v77, v78); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); pipe_barrier(PIPE_M); - TMATMUL_ACC(v76, v76, v68, v70); + 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> - v78 = Tile< + v79 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v79 = (uint64_t)v19; - TASSIGN(v78, v79); + uint64_t v80 = (uint64_t)v20; + TASSIGN(v79, v80); pipe_barrier(PIPE_M); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - TMATMUL_ACC(v78, v78, v72, v74); + 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> - v80 = Tile< + v81 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v81 = (uint64_t)v19; - TASSIGN(v80, v81); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); + uint64_t v82 = (uint64_t)v20; + TASSIGN(v81, v82); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); - TEXTRACT(v80, v58, v19, v19); + TEXTRACT(v81, v59, v20, v20); Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v82 = Tile< + v83 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v83 = (uint64_t)v19; - TASSIGN(v82, v83); + 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(v82, v63, v19, v19); + 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> - v84 = Tile< + v85 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v85 = (uint64_t)v18; - TASSIGN(v84, v85); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - TEXTRACT(v84, v58, v19, 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> - v86 = Tile< + v87 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v87 = (uint64_t)v17; - TASSIGN(v86, v87); - TEXTRACT(v86, v63, v10, v19); + 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> - v88 = Tile< + v89 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v89 = (uint64_t)v19; - TASSIGN(v88, v89); + uint64_t v90 = (uint64_t)v20; + TASSIGN(v89, v90); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); pipe_barrier(PIPE_M); - TMATMUL_ACC(v88, v88, v80, v82); + 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> - v90 = Tile< + v91 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v91 = (uint64_t)v19; - TASSIGN(v90, v91); + uint64_t v92 = (uint64_t)v20; + TASSIGN(v91, v92); pipe_barrier(PIPE_M); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - TMATMUL_ACC(v90, v90, v84, v86); + TMATMUL_ACC(v91, v91, v85, v87); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); } set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); @@ -396,116 +397,116 @@ out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v92 = Tile< + v93 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v12); - uint64_t v93 = (uint64_t)v21; - TASSIGN(v92, v93); - pto::Shape<1, 1, 1, 16, 64> v94 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v95 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v96 = GlobalTensor< + v97 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v19 + v19 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v94, v95 + v1 + (v20 + v20 * v14 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v13), v95, v96 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID5); - TLOAD(v92, v96); + 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> - v97 = Tile< + v98 = Tile< TileType::Mat, bfloat16_t, 64, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v12, v11); - uint64_t v98 = (uint64_t)v20; - TASSIGN(v97, v98); - pto::Shape<1, 1, 1, 64, 512> v99 = pto::Shape<1, 1, 1, 64, 512>(); - pto::Stride<327680, 327680, 327680, 5120, 1> v100 = pto::Stride<327680, 327680, 327680, 5120, 1>(); + 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> - v101 = GlobalTensor< + v102 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 512>, pto::Stride<327680, 327680, 327680, 5120, 1>, pto::Layout::ND>( - v2 + (v19 + (int64_t)((uint64_t)v29 + (uint64_t)v7) * v14 + v6 * v13), v99, v100 + v2 + (v20 + (int64_t)((uint64_t)v30 + (uint64_t)v7) * v14 + v6 * v13), v100, v101 ); - TLOAD(v97, 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> - v102 = Tile< + v103 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v103 = (uint64_t)v19; - TASSIGN(v102, v103); + uint64_t v104 = (uint64_t)v19; + TASSIGN(v103, v104); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID6); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - TEXTRACT(v102, v92, v19, v19); + TEXTRACT(v103, v93, v20, v20); Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v104 = Tile< + v105 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v105 = (uint64_t)v19; - TASSIGN(v104, v105); + uint64_t v106 = (uint64_t)v20; + TASSIGN(v105, v106); wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID7); - TEXTRACT(v104, v97, v19, v19); + 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> - v106 = Tile< + v107 = Tile< TileType::Left, bfloat16_t, 16, 32, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v107 = (uint64_t)v18; - TASSIGN(v106, v107); - TEXTRACT(v106, v92, v19, 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> - v108 = Tile< + v109 = Tile< TileType::Right, bfloat16_t, 32, 512, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v10, v11); - uint64_t v109 = (uint64_t)v17; - TASSIGN(v108, v109); - TEXTRACT(v108, v97, v10, v19); + 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> - v110 = Tile< + v111 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v111 = (uint64_t)v19; - TASSIGN(v110, v111); + uint64_t v112 = (uint64_t)v20; + TASSIGN(v111, v112); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID6); pipe_barrier(PIPE_M); - TMATMUL_ACC(v110, v110, v102, v104); + TMATMUL_ACC(v111, v111, v103, v105); Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v112 = Tile< + v113 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v113 = (uint64_t)v19; - TASSIGN(v112, v113); + uint64_t v114 = (uint64_t)v20; + TASSIGN(v113, v114); pipe_barrier(PIPE_M); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID7); - TMATMUL_ACC(v112, v112, v106, v108); + 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> - v114 = Tile< + v115 = Tile< TileType::Acc, float, 16, 512, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v15, v11); - uint64_t v115 = (uint64_t)v19; - TASSIGN(v114, v115); - pto::Shape<1, 1, 1, 16, 512> v116 = pto::Shape<1, 1, 1, 16, 512>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v117 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> v118 = + 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 + (v19 + v19 * v14 + v6 * v13), v116, v117 + v3 + (v20 + v20 * v14 + v6 * v13), v117, v118 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); TSTORE< @@ -513,7 +514,7 @@ out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t 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>(v118, v114); + 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); @@ -527,50 +528,50 @@ out_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: attn_out_inline157__rv_v2 - __gm__ Tensor *attn_out_inline157__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *attn_out_inline157__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline157__rv_v2_tensor->buffer.addr) + - attn_out_inline157__rv_v2_tensor->start_offset; + // 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_inline248__iter_v6 - __gm__ Tensor *attn_proj_fp32_inline248__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__iter_v6 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__iter_v6_tensor->buffer.addr) + - attn_proj_fp32_inline248__iter_v6_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_inline151__ssa_v0 + // Unpack scalar: k_op_inline266__ssa_v0 union { uint64_t u64; int64_t val; - } k_op_inline151__ssa_v0_conv; - k_op_inline151__ssa_v0_conv.u64 = args[3]; - int64_t k_op_inline151__ssa_v0 = k_op_inline151__ssa_v0_conv.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_inline135__ssa_v0 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[4]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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_inline211__ssa_v0 + // Unpack scalar: n_op_inline64__ssa_v0 union { uint64_t u64; int64_t val; - } n_op_inline211__ssa_v0_conv; - n_op_inline211__ssa_v0_conv.u64 = args[5]; - int64_t n_op_inline211__ssa_v0 = n_op_inline211__ssa_v0_conv.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_inline157__rv_v2, wo__ssa_v0, attn_proj_fp32_inline248__iter_v6, k_op_inline151__ssa_v0, - layer_hidden_base_inline135__ssa_v0, n_op_inline211__ssa_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj_0.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj_0.cpp new file mode 100644 index 0000000000..50b5cc3582 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/q_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/q_proj.cpp index 49d22c8041..cc978d8d64 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/q_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/q_proj.cpp @@ -53,39 +53,37 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT 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 = 5; + 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 = 1280; + const int64_t v12 = 1024; const int64_t v13 = 512; - const int64_t v14 = 4; + const int64_t v14 = 5; const int64_t v15 = 1; const int64_t v16 = 5120; const int64_t v17 = 16; - const int64_t v18 = 8192; - const int64_t v19 = 32768; - const int64_t v20 = 2048; - const int64_t v21 = 0; - const int64_t v22 = 147456; - const int64_t v23 = 16384; + 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 v24 = (size_t)v15; - size_t v25 = (size_t)v21; - size_t v26 = (size_t)v11; + 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_MTE1, PIPE_MTE2, EVENT_ID2); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - for (size_t v31 = v25; v31 < v26; v31 += v24) { + 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< @@ -94,7 +92,7 @@ q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v 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)v23; + 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>(); @@ -102,7 +100,7 @@ q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v 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 + (v21 + v21 * v16 + v30 * v15), v35, v36 + v2 + (v23 + v23 * v16 + v30 * v15), v35, v36 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); TLOAD(v33, v37); @@ -112,7 +110,7 @@ q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v 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)v22; + 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>(); @@ -120,347 +118,447 @@ q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v 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 + (v21 + (int64_t)((uint64_t)v4 + (uint64_t)v30) * v16 + v32 * v15), v40, v41); + 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); - Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v43 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v17, v10); - uint64_t v44 = (uint64_t)v21; - TASSIGN(v43, v44); 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 v45 = v25; v45 < v27; v45 += v28) { - int64_t v46 = (int64_t)v45; + 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> - v47 = Tile< + v45 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v17, v8); - uint64_t v48 = (uint64_t)v21; - TASSIGN(v47, v48); + uint64_t v46 = (uint64_t)v22; + TASSIGN(v45, v46); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); pipe_barrier(PIPE_MTE1); - TEXTRACT(v47, v33, v21, v45); + TEXTRACT(v45, v33, v23, v43); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v49 = Tile< + v47 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v50 = (uint64_t)v21; - TASSIGN(v49, v50); - TEXTRACT(v49, v38, v45, v21); + 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> - v51 = Tile< + v49 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v17, v8); - uint64_t v52 = (uint64_t)v20; - TASSIGN(v51, v52); - int64_t v53 = (int64_t)((uint64_t)v46 + (uint64_t)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(v51, v33, v21, v53); + TEXTRACT(v49, v33, v23, v51); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v54 = Tile< + v52 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v55 = (uint64_t)v19; - TASSIGN(v54, v55); - TEXTRACT(v54, v38, v53, v21); + 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 (v46 == v21) { + if (v44 == v23) { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v56 = Tile< + v54 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v57 = (uint64_t)v21; - TASSIGN(v56, v57); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); pipe_barrier(PIPE_M); - TMATMUL(v56, v47, v49); + TMATMUL(v54, v45, v47); } else { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v58 = Tile< + v56 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v59 = (uint64_t)v21; - TASSIGN(v58, v59); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); pipe_barrier(PIPE_M); - TMATMUL_ACC(v58, v58, v47, v49); + 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> - v60 = Tile< + v58 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v61 = (uint64_t)v21; - TASSIGN(v60, v61); + uint64_t v59 = (uint64_t)v23; + TASSIGN(v58, v59); pipe_barrier(PIPE_M); wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - TMATMUL_ACC(v60, v60, v51, v54); + 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 v62 = v24; v62 < ((size_t)v7); v62 += v26) { - int64_t v63 = (int64_t)((uint64_t)((int64_t)v62) * (uint64_t)v10); - int64_t v64 = (int64_t)((uint64_t)v30 + (uint64_t)v63); - int64_t v65 = (int64_t)((uint64_t)v30 + (uint64_t)((int64_t)(uint64_t)v63 + (uint64_t)v10)); + for (size_t v82 = v26; v82 < v27; v82 += v28) { Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v66 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v67 = (uint64_t)v21; - TASSIGN(v66, v67); - pto::Shape<1, 1, 1, 16, 256> v68 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v69 = 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> - v70 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v21 + v21 * v16 + v64 * v15), v68, v69); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v66, v70); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v71 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v72 = (uint64_t)v22; - TASSIGN(v71, v72); - pto::Shape<1, 1, 1, 256, 256> v73 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<1310720, 1310720, 1310720, 5120, 1> v74 = 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> - v75 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, - pto::Layout::ND>(v3 + (v21 + (int64_t)((uint64_t)v4 + (uint64_t)v64) * v16 + v32 * v15), v73, v74); - TLOAD(v71, v75); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + 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::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v76 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v77 = (uint64_t)v18; - TASSIGN(v76, v77); - pto::Shape<1, 1, 1, 16, 256> v78 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v79 = 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> - v80 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v21 + v21 * v16 + v65 * v15), v78, v79); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v76, v80); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v81 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v82 = (uint64_t)v23; - TASSIGN(v81, v82); - pto::Shape<1, 1, 1, 256, 256> v83 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<1310720, 1310720, 1310720, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<1310720, 1310720, 1310720, 5120, 1>, - pto::Layout::ND>(v3 + (v21 + (int64_t)((uint64_t)v4 + (uint64_t)v65) * v16 + v32 * v15), v83, v84); - TLOAD(v81, v85); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - for (size_t v86 = v25; v86 < v27; v86 += v28) { - Tile< + }; + 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> - 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); - pipe_barrier(PIPE_MTE1); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - TEXTRACT(v87, v66, v21, v86); - Tile< + 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> - v89 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v90 = (uint64_t)v21; - TASSIGN(v89, v90); - TEXTRACT(v89, v71, v86, v21); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - Tile< + 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> - v91 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v17, v8); - uint64_t v92 = (uint64_t)v20; - TASSIGN(v91, v92); - int64_t v93 = (int64_t)((uint64_t)((int64_t)v86) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - TEXTRACT(v91, v66, v21, v93); - Tile< + 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> - v94 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v95 = (uint64_t)v19; - TASSIGN(v94, v95); - TEXTRACT(v94, v71, v93, v21); - 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> - v96 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v97 = (uint64_t)v21; - TASSIGN(v96, v97); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v96, v96, v87, v89); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - Tile< + 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> - v98 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v99 = (uint64_t)v21; - TASSIGN(v98, v99); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - TMATMUL_ACC(v98, v98, v91, v94); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - }; - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + 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); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - for (size_t v100 = v25; v100 < v27; v100 += v28) { - Tile< + }; + 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> - 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)v21; - TASSIGN(v101, v102); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - TEXTRACT(v101, v76, v21, v100); - Tile< + 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> - v103 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v104 = (uint64_t)v21; - TASSIGN(v103, v104); - TEXTRACT(v103, v81, v100, v21); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - Tile< + 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> - v105 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v17, v8); - uint64_t v106 = (uint64_t)v20; - TASSIGN(v105, v106); - int64_t v107 = (int64_t)((uint64_t)((int64_t)v100) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TEXTRACT(v105, v76, v21, v107); - Tile< + 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> - v108 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v109 = (uint64_t)v19; - TASSIGN(v108, v109); - TEXTRACT(v108, v81, v107, v21); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - Tile< + 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> - v110 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v111 = (uint64_t)v21; - TASSIGN(v110, v111); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v110, v110, v101, v103); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - Tile< + 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> - v112 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v17, v10); - uint64_t v113 = (uint64_t)v21; - TASSIGN(v112, v113); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - TMATMUL_ACC(v112, v112, v105, v108); - 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_M, PIPE_MTE1, EVENT_ID4); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + 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_ID4); + 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); - pto::Shape<1, 1, 1, 16, 256> v114 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v115 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v116 = GlobalTensor< + v139 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v21 + v21 * v16 + v32 * v15), v114, v115 + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); pipe_barrier(PIPE_FIX); @@ -470,14 +568,12 @@ q_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v CompactMode::Null>, GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>, - AtomicType::AtomicAdd>(v116, v43); + 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); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); #endif // __DAV_CUBE__ ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); @@ -489,34 +585,33 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in int32_t __pypto_spmd_block_idx = get_block_idx(args); int32_t __pypto_spmd_block_num = get_block_num(args); - // Unpack tensor: q_proj_inline129__rv_v2 - __gm__ Tensor *q_proj_inline129__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *q_proj_inline129__rv_v2 = - reinterpret_cast<__gm__ float *>(q_proj_inline129__rv_v2_tensor->buffer.addr) + - q_proj_inline129__rv_v2_tensor->start_offset; + // 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_states_inline146__rv_v2 - __gm__ Tensor *normed_states_inline146__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *normed_states_inline146__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(normed_states_inline146__rv_v2_tensor->buffer.addr) + - normed_states_inline146__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_inline135__ssa_v0 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[3]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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_inline129__rv_v2, normed_states_inline146__rv_v2, wq__ssa_v0, layer_hidden_base_inline135__ssa_v0, + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj.cpp index 917df5c0b3..09f48d3335 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj.cpp @@ -27,6 +27,7 @@ #include #include "tensor.h" +#include "intrinsic.h" using namespace pto; @@ -50,149 +51,152 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void -up_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 = 1024; +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 = 2048; - const int64_t v18 = 32768; - const int64_t v19 = 512; - const int64_t v20 = 0; - const int64_t v21 = 135168; - const int64_t v22 = 4096; + 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 v23 = (size_t)v12; - size_t v24 = (size_t)v20; - size_t v25 = (size_t)v10; + 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> - v26 = Tile< + v29 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v27 = (uint64_t)v22; - TASSIGN(v26, v27); - pto::Shape<1, 1, 1, 16, 64> v28 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v29 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v30 = GlobalTensor< + v33 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + v4 * v14), v28, v29 + 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(v26, v30); + TLOAD(v29, v33); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v31 = Tile< + v34 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v32 = (uint64_t)v21; - TASSIGN(v31, v32); - int64_t v33 = (int64_t)((uint64_t)v5 + (uint64_t)v4); - pto::Shape<1, 1, 1, 64, 1024> v34 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v35 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v36 = GlobalTensor< + v39 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + v33 * v13 + v6 * v14), v34, v35); - TLOAD(v31, v36); + 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 v37 = v24; v37 < v23; v37 += v25) { - int64_t v38 = (int64_t)v37; + 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> - v39 = Tile< + v42 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v40 = (uint64_t)v20; - TASSIGN(v39, v40); + uint64_t v43 = (uint64_t)v21; + TASSIGN(v42, v43); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); pipe_barrier(PIPE_MTE1); - TEXTRACT(v39, v26, v20, v37); + TEXTRACT(v42, v29, v22, v40); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v41 = Tile< + v44 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v42 = (uint64_t)v20; - TASSIGN(v41, v42); - TEXTRACT(v41, v31, v37, v20); + 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> - v43 = Tile< + v46 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v44 = (uint64_t)v19; - TASSIGN(v43, v44); - int64_t v45 = (int64_t)((uint64_t)v38 + (uint64_t)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(v43, v26, v20, v45); + TEXTRACT(v46, v29, v22, v48); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v46 = Tile< + v49 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v47 = (uint64_t)v18; - TASSIGN(v46, v47); - TEXTRACT(v46, v31, v45, v20); + 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 (v38 == v20) { + if (v41 == v22) { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v48 = Tile< + v51 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v49 = (uint64_t)v20; - TASSIGN(v48, v49); + CompactMode::Null>(v16, v21); + uint64_t v52 = (uint64_t)v22; + TASSIGN(v51, v52); pipe_barrier(PIPE_M); - TMATMUL(v48, v39, v41); + TMATMUL(v51, v42, v44); } else { Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v50 = Tile< + v53 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v51 = (uint64_t)v20; - TASSIGN(v50, v51); + CompactMode::Null>(v16, v21); + uint64_t v54 = (uint64_t)v22; + TASSIGN(v53, v54); pipe_barrier(PIPE_M); - TMATMUL_ACC(v50, v50, v39, v41); + 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> - v52 = Tile< + v55 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v53 = (uint64_t)v20; - TASSIGN(v52, v53); + 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(v52, v52, v43, v46); + TMATMUL_ACC(v55, v55, v46, v49); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -202,151 +206,151 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t 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 v54 = (size_t)v14; v54 < ((size_t)v9); v54 += (size_t)v8) { - int64_t v55 = (int64_t)((uint64_t)((int64_t)v54) * (uint64_t)v12); - int64_t v56 = (int64_t)((uint64_t)v55 + (uint64_t)v12); + 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> - v57 = Tile< + v60 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v58 = (uint64_t)v20; - TASSIGN(v57, v58); - pto::Shape<1, 1, 1, 16, 64> v59 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v60 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v61 = GlobalTensor< + v64 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v55) * v14), v59, v60 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v58) * v14), v62, v63 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - TLOAD(v57, v61); + TLOAD(v60, v64); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v62 = Tile< + v65 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v63 = (uint64_t)v21; - TASSIGN(v62, v63); - pto::Shape<1, 1, 1, 64, 1024> v64 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v65 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v66 = GlobalTensor< + v69 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v55) * v13 + v6 * v14), v64, v65); - TLOAD(v62, v66); + 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> - v67 = Tile< + v70 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v68 = (uint64_t)v17; - TASSIGN(v67, v68); - pto::Shape<1, 1, 1, 16, 64> v69 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v70 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v71 = GlobalTensor< + v74 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v56) * v14), v69, v70 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v59) * v14), v72, v73 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v67, v71); + TLOAD(v70, v74); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v72 = Tile< + v75 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v73 = (uint64_t)v22; - TASSIGN(v72, v73); - pto::Shape<1, 1, 1, 64, 1024> v74 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v75 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v76 = GlobalTensor< + v79 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v56) * v13 + v6 * v14), v74, v75); - TLOAD(v72, v76); + 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_M, PIPE_MTE1, EVENT_ID3); 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 v77 = v24; v77 < v23; v77 += v25) { + 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> - v78 = Tile< + v81 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v79 = (uint64_t)v20; - TASSIGN(v78, v79); + uint64_t v82 = (uint64_t)v21; + TASSIGN(v81, v82); pipe_barrier(PIPE_MTE1); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - TEXTRACT(v78, v57, v20, v77); + TEXTRACT(v81, v60, v22, v80); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v80 = Tile< + v83 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v81 = (uint64_t)v20; - TASSIGN(v80, v81); - TEXTRACT(v80, v62, v77, v20); + 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> - v82 = Tile< + v85 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v83 = (uint64_t)v19; - TASSIGN(v82, v83); - int64_t v84 = (int64_t)((uint64_t)((int64_t)v77) + (uint64_t)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(v82, v57, v20, v84); + TEXTRACT(v85, v60, v22, v87); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v85 = Tile< + v88 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v86 = (uint64_t)v18; - TASSIGN(v85, v86); - TEXTRACT(v85, v62, v84, v20); + 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> - v87 = Tile< + v90 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v88 = (uint64_t)v20; - TASSIGN(v87, v88); + 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(v87, v87, v78, v80); + 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> - v89 = Tile< + v92 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v90 = (uint64_t)v20; - TASSIGN(v89, v90); + 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(v89, v89, v82, v85); + TMATMUL_ACC(v92, v92, v85, v88); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); @@ -357,71 +361,71 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t 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 v91 = v24; v91 < v23; v91 += v25) { + 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> - v92 = Tile< + v95 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v93 = (uint64_t)v20; - TASSIGN(v92, v93); + uint64_t v96 = (uint64_t)v22; + TASSIGN(v95, v96); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - TEXTRACT(v92, v67, v20, v91); + TEXTRACT(v95, v70, v22, v94); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v94 = Tile< + v97 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v95 = (uint64_t)v20; - TASSIGN(v94, v95); - TEXTRACT(v94, v72, v91, v20); + 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> - v96 = Tile< + v99 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v97 = (uint64_t)v19; - TASSIGN(v96, v97); - int64_t v98 = (int64_t)((uint64_t)((int64_t)v91) + (uint64_t)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(v96, v67, v20, v98); + TEXTRACT(v99, v70, v22, v101); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v99 = Tile< + v102 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v100 = (uint64_t)v18; - TASSIGN(v99, v100); - TEXTRACT(v99, v72, v98, v20); + 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> - v101 = Tile< + v104 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v102 = (uint64_t)v20; - TASSIGN(v101, v102); + 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(v101, v101, v92, v94); + 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> - v103 = Tile< + v106 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v104 = (uint64_t)v20; - TASSIGN(v103, v104); + 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(v103, v103, v96, v99); + TMATMUL_ACC(v106, v106, v99, v102); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); }; wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); @@ -435,106 +439,106 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v105 = Tile< + v108 = Tile< TileType::Mat, bfloat16_t, 16, 64, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v12); - uint64_t v106 = (uint64_t)v22; - TASSIGN(v105, v106); - pto::Shape<1, 1, 1, 16, 64> v107 = pto::Shape<1, 1, 1, 16, 64>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v108 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v109 = GlobalTensor< + v112 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 64>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v7) * v14), v107, v108 + v1 + (v22 + v22 * v15 + (int64_t)((uint64_t)v4 + (uint64_t)v8) * v14), v110, v111 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v105, v109); + TLOAD(v108, v112); Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v110 = Tile< + v113 = Tile< TileType::Mat, bfloat16_t, 64, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v12, v11); - uint64_t v111 = (uint64_t)v21; - TASSIGN(v110, v111); - pto::Shape<1, 1, 1, 64, 1024> v112 = pto::Shape<1, 1, 1, 64, 1024>(); - pto::Stride<1114112, 1114112, 1114112, 17408, 1> v113 = pto::Stride<1114112, 1114112, 1114112, 17408, 1>(); + 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> - v114 = GlobalTensor< + v117 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 64, 1024>, pto::Stride<1114112, 1114112, 1114112, 17408, 1>, - pto::Layout::ND>(v2 + (v20 + (int64_t)((uint64_t)v33 + (uint64_t)v7) * v13 + v6 * v14), v112, v113); - TLOAD(v110, v114); + 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_M, PIPE_MTE1, EVENT_ID0); 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 v115 = v24; v115 < v23; v115 += v25) { + 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> - v116 = Tile< + v119 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v117 = (uint64_t)v20; - TASSIGN(v116, v117); + uint64_t v120 = (uint64_t)v21; + TASSIGN(v119, v120); wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - TEXTRACT(v116, v105, v20, v115); + TEXTRACT(v119, v108, v22, v118); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v118 = Tile< + v121 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v119 = (uint64_t)v20; - TASSIGN(v118, v119); - TEXTRACT(v118, v110, v115, v20); + 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> - v120 = Tile< + v123 = Tile< TileType::Left, bfloat16_t, 16, 16, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v16, v16); - uint64_t v121 = (uint64_t)v19; - TASSIGN(v120, v121); - int64_t v122 = (int64_t)((uint64_t)((int64_t)v115) + (uint64_t)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(v120, v105, v20, v122); + TEXTRACT(v123, v108, v22, v125); Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v123 = Tile< + v126 = Tile< TileType::Right, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v16, v11); - uint64_t v124 = (uint64_t)v18; - TASSIGN(v123, v124); - TEXTRACT(v123, v110, v122, v20); + 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> - v125 = Tile< + v128 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v126 = (uint64_t)v20; - TASSIGN(v125, v126); + 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(v125, v125, v116, v118); + 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> - v127 = Tile< + v130 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v128 = (uint64_t)v20; - TASSIGN(v127, v128); + 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(v127, v127, v120, v123); + TMATMUL_ACC(v130, v130, v123, v126); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); } wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); @@ -543,17 +547,17 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v129 = Tile< + v132 = Tile< TileType::Acc, float, 16, 1024, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v130 = (uint64_t)v20; - TASSIGN(v129, v130); - pto::Shape<1, 1, 1, 16, 1024> v131 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v132 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v133 = GlobalTensor< + v136 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v3 + (v20 + v20 * v13 + v6 * v14), v131, v132 + v3 + (v22 + v22 * v13 + v28 * v14), v134, v135 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); TSTORE< @@ -562,7 +566,7 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t CompactMode::Null>, GlobalTensor< float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>, - AtomicType::AtomicAdd>(v133, v129); + AtomicType::AtomicAdd>(v136, v132); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); #endif // __DAV_CUBE__ @@ -572,50 +576,46 @@ up_proj(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v14 - __gm__ Tensor *mlp_norm_in_inline250__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v14 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v14_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v14_tensor->start_offset; + // 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_inline247__iter_v6 - __gm__ Tensor *up_acc_all_inline247__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *up_acc_all_inline247__iter_v6 = - reinterpret_cast<__gm__ float *>(up_acc_all_inline247__iter_v6_tensor->buffer.addr) + - up_acc_all_inline247__iter_v6_tensor->start_offset; - - // Unpack scalar: k0_inline202__ssa_v10 - union { - uint64_t u64; - int64_t val; - } k0_inline202__ssa_v10_conv; - k0_inline202__ssa_v10_conv.u64 = args[3]; - int64_t k0_inline202__ssa_v10 = k0_inline202__ssa_v10_conv.val; + // 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: layer_hidden_base_inline135__ssa_v0 + // Unpack scalar: gu_k0_inline131__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[4]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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: n0_inline142__ssa_v6 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } n0_inline142__ssa_v6_conv; - n0_inline142__ssa_v6_conv.u64 = args[5]; - int64_t n0_inline142__ssa_v6 = n0_inline142__ssa_v6_conv.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_inline250__rv_v14, w_up__ssa_v0, up_acc_all_inline247__iter_v6, k0_inline202__ssa_v10, - layer_hidden_base_inline135__ssa_v0, n0_inline142__ssa_v6 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_0.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_0.cpp new file mode 100644 index 0000000000..01a120e8df --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_1.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_1.cpp new file mode 100644 index 0000000000..5238de9792 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_2.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_2.cpp new file mode 100644 index 0000000000..e7e043d882 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_3.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_3.cpp new file mode 100644 index 0000000000..8f85c9a488 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_4.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_4.cpp new file mode 100644 index 0000000000..136ab89c18 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/v_proj.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/v_proj.cpp index 64df754bb1..6b98eb96c3 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/v_proj.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/v_proj.cpp @@ -53,415 +53,512 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT 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 = 5; + 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 = 1280; - const int64_t v13 = 512; - const int64_t v14 = 4; - const int64_t v15 = 5120; - const int64_t v16 = 1; - const int64_t v17 = 1024; - const int64_t v18 = 16; + 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 = 2048; - const int64_t v22 = 0; - const int64_t v23 = 147456; - const int64_t v24 = 16384; + 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 v25 = (size_t)v16; - size_t v26 = (size_t)v22; - size_t v27 = (size_t)v11; - size_t v28 = (size_t)v10; - size_t v29 = (size_t)v9; - int64_t v30 = (int64_t)v5; - int64_t v31 = (int64_t)((uint64_t)(v30 % v14) * (uint64_t)v12); + 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_MTE1, PIPE_MTE2, EVENT_ID2); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - for (size_t v32 = v26; v32 < v27; v32 += v25) { - int64_t v33 = (int64_t)((uint64_t)((int64_t)(uint64_t)(v30 / v14) * (uint64_t)v13) + - (uint64_t)((int64_t)(uint64_t)((int64_t)v32) * (uint64_t)v10)); + 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> - v34 = Tile< + v33 = Tile< TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v35 = (uint64_t)v24; - 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>(); + 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> - v38 = GlobalTensor< + v37 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v22 + v22 * v15 + v31 * v16), v36, v37 + v2 + (v23 + v23 * v14 + v30 * v15), v35, v36 ); wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - TLOAD(v34, v38); + TLOAD(v33, v37); Tile< TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v39 = Tile< + v38 = Tile< TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v40 = (uint64_t)v23; - TASSIGN(v39, v40); - pto::Shape<1, 1, 1, 256, 256> v41 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v42 = pto::Stride<262144, 262144, 262144, 1024, 1>(); + 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> - v43 = GlobalTensor< + v42 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v31) * v17 + v33 * v16), v41, v42); - TLOAD(v39, v43); + 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); - Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v44 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v45 = (uint64_t)v22; - TASSIGN(v44, v45); 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 v46 = v26; v46 < v28; v46 += v29) { - int64_t v47 = (int64_t)v46; + 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> - v48 = Tile< + v45 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v49 = (uint64_t)v22; - TASSIGN(v48, v49); + 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(v48, v34, v22, v46); + TEXTRACT(v45, v33, v23, v43); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v50 = Tile< + v47 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v51 = (uint64_t)v22; - TASSIGN(v50, v51); - TEXTRACT(v50, v39, v46, v22); + 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> - v52 = Tile< + v49 = Tile< TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v53 = (uint64_t)v21; - TASSIGN(v52, v53); - int64_t v54 = (int64_t)((uint64_t)v47 + (uint64_t)v8); + 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(v52, v34, v22, v54); + TEXTRACT(v49, v33, v23, v51); Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v55 = Tile< + v52 = Tile< TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v56 = (uint64_t)v20; - TASSIGN(v55, v56); - TEXTRACT(v55, v39, v54, v22); + 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 (v47 == v22) { + if (v44 == v23) { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v57 = Tile< + v54 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v58 = (uint64_t)v22; - TASSIGN(v57, v58); + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v55 = (uint64_t)v23; + TASSIGN(v54, v55); pipe_barrier(PIPE_M); - TMATMUL(v57, v48, v50); + TMATMUL(v54, v45, v47); } else { Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, CompactMode::Null> - v59 = Tile< + v56 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); + PadValue::Null, CompactMode::Null>(v17, v10); + uint64_t v57 = (uint64_t)v23; + TASSIGN(v56, v57); pipe_barrier(PIPE_M); - TMATMUL_ACC(v59, v59, v48, v50); + 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> - v61 = Tile< + v58 = Tile< TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v18, v10); - uint64_t v62 = (uint64_t)v22; - TASSIGN(v61, v62); + 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(v61, v61, v52, v55); + 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 v63 = v25; v63 < ((size_t)v7); v63 += v27) { - int64_t v64 = (int64_t)((uint64_t)((int64_t)v63) * (uint64_t)v10); - int64_t v65 = (int64_t)((uint64_t)v31 + (uint64_t)v64); - int64_t v66 = (int64_t)((uint64_t)v31 + (uint64_t)((int64_t)(uint64_t)v64 + (uint64_t)v10)); + for (size_t v82 = v26; v82 < v27; v82 += v28) { Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v67 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pto::Shape<1, 1, 1, 16, 256> v69 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v70 = 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> - v71 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v22 + v22 * v15 + v65 * v16), v69, v70); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v67, v71); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v72 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v73 = (uint64_t)v23; - TASSIGN(v72, v73); - pto::Shape<1, 1, 1, 256, 256> v74 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v75 = 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> - v76 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v65) * v17 + v33 * v16), v74, v75); - TLOAD(v72, v76); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); + 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::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, CompactMode::Null> - v77 = Tile< - TileType::Mat, bfloat16_t, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v78 = (uint64_t)v19; - TASSIGN(v77, v78); - pto::Shape<1, 1, 1, 16, 256> v79 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v80 = 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> - v81 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, - pto::Layout::ND>(v2 + (v22 + v22 * v15 + v66 * v16), v79, v80); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v77, v81); + 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::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, + TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, CompactMode::Null> - v82 = Tile< - TileType::Mat, bfloat16_t, 256, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v10, v10); - uint64_t v83 = (uint64_t)v24; - TASSIGN(v82, v83); - pto::Shape<1, 1, 1, 256, 256> v84 = pto::Shape<1, 1, 1, 256, 256>(); - pto::Stride<262144, 262144, 262144, 1024, 1> v85 = 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> - v86 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 256, 256>, pto::Stride<262144, 262144, 262144, 1024, 1>, - pto::Layout::ND>(v3 + (v22 + (int64_t)((uint64_t)v4 + (uint64_t)v66) * v17 + v33 * v16), v84, v85); - TLOAD(v82, v86); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - for (size_t v87 = v26; v87 < v28; v87 += v29) { - Tile< + }; + 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> - v88 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v89 = (uint64_t)v22; - TASSIGN(v88, v89); - pipe_barrier(PIPE_MTE1); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - TEXTRACT(v88, v67, v22, v87); - Tile< + 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> - 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)v22; - TASSIGN(v90, v91); - TEXTRACT(v90, v72, v87, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - Tile< + 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> - v92 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v93 = (uint64_t)v21; - TASSIGN(v92, v93); - int64_t v94 = (int64_t)((uint64_t)((int64_t)v87) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - TEXTRACT(v92, v67, v22, v94); - Tile< + 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> - v95 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v96 = (uint64_t)v20; - TASSIGN(v95, v96); - TEXTRACT(v95, v72, v94, v22); - 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> - v97 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v98 = (uint64_t)v22; - TASSIGN(v97, v98); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v97, v97, v88, v90); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - Tile< + 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> - v99 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v100 = (uint64_t)v22; - TASSIGN(v99, v100); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - TMATMUL_ACC(v99, v99, v92, v95); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); - }; - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID5); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID6); + 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); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID7); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); + 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); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - for (size_t v101 = v26; v101 < v28; v101 += v29) { - Tile< + }; + 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> - v102 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v103 = (uint64_t)v22; - TASSIGN(v102, v103); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - TEXTRACT(v102, v77, v22, v101); - Tile< + 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> - 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)v22; - TASSIGN(v104, v105); - TEXTRACT(v104, v82, v101, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - Tile< + 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> - v106 = Tile< - TileType::Left, bfloat16_t, 16, 64, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v18, v8); - uint64_t v107 = (uint64_t)v21; - TASSIGN(v106, v107); - int64_t v108 = (int64_t)((uint64_t)((int64_t)v101) + (uint64_t)v8); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TEXTRACT(v106, v77, v22, v108); - Tile< + 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> - v109 = Tile< - TileType::Right, bfloat16_t, 64, 256, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v8, v10); - uint64_t v110 = (uint64_t)v20; - TASSIGN(v109, v110); - TEXTRACT(v109, v82, v108, v22); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - Tile< + 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> - v111 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v112 = (uint64_t)v22; - TASSIGN(v111, v112); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID4); - pipe_barrier(PIPE_M); - TMATMUL_ACC(v111, v111, v102, v104); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - Tile< + 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> - v113 = Tile< - TileType::Acc, float, 16, 256, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, - PadValue::Null, CompactMode::Null>(v18, v10); - uint64_t v114 = (uint64_t)v22; - TASSIGN(v113, v114); - pipe_barrier(PIPE_M); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID5); - TMATMUL_ACC(v113, v113, v106, v109); - 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_M, PIPE_MTE1, EVENT_ID4); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); + 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_ID4); + 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); - pto::Shape<1, 1, 1, 16, 256> v115 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v116 = pto::Stride<16384, 16384, 16384, 1024, 1>(); + 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> - v117 = GlobalTensor< + v139 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v1 + (v22 + v22 * v17 + v33 * v16), v115, v116 + v1 + (v23 + v23 * v16 + v32 * v15), v137, v138 ); wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); pipe_barrier(PIPE_FIX); @@ -471,14 +568,12 @@ v_proj(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ bfloat16_t *v3, int64_t v CompactMode::Null>, GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>, - AtomicType::AtomicAdd>(v117, v44); + 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); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); #endif // __DAV_CUBE__ ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); @@ -490,34 +585,33 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in int32_t __pypto_spmd_block_idx = get_block_idx(args); int32_t __pypto_spmd_block_num = get_block_num(args); - // Unpack tensor: v_proj_inline152__ssa_v1 - __gm__ Tensor *v_proj_inline152__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *v_proj_inline152__ssa_v1 = - reinterpret_cast<__gm__ float *>(v_proj_inline152__ssa_v1_tensor->buffer.addr) + - v_proj_inline152__ssa_v1_tensor->start_offset; + // 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_states_inline146__rv_v2 - __gm__ Tensor *normed_states_inline146__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *normed_states_inline146__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(normed_states_inline146__rv_v2_tensor->buffer.addr) + - normed_states_inline146__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: 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_inline135__ssa_v0 + // Unpack scalar: layer_hidden_base_inline151__ssa_v0 union { uint64_t u64; int64_t val; - } layer_hidden_base_inline135__ssa_v0_conv; - layer_hidden_base_inline135__ssa_v0_conv.u64 = args[3]; - int64_t layer_hidden_base_inline135__ssa_v0 = layer_hidden_base_inline135__ssa_v0_conv.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_inline152__ssa_v1, normed_states_inline146__rv_v2, wv__ssa_v0, layer_hidden_base_inline135__ssa_v0, + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_fence.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_fence.cpp deleted file mode 100644 index c8ec8ea090..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_fence.cpp +++ /dev/null @@ -1,143 +0,0 @@ -/* - * 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_fence -// 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_fence(__gm__ bfloat16_t *v1, __gm__ float *v2) { - SaturationMode v3 = SaturationMode::OFF; - RoundMode v4 = RoundMode::CAST_ROUND; - const int64_t v5 = 256; - 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; - const int64_t v11 = 16384; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v12 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v5); - uint64_t v13 = (uint64_t)v11; - TASSIGN(v12, v13); - pto::Shape<1, 1, 1, 16, 256> v14 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v15 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v16 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v10 * v6), v14, v15 - ); - TLOAD(v12, v16); - 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> - v17 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v5); - uint64_t v18 = (uint64_t)v10; - TASSIGN(v17, v18); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v17, v12, v4, v3); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v19 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v5); - uint64_t v20 = (uint64_t)v11; - TASSIGN(v19, v20); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v21 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v6); - uint64_t v22 = (uint64_t)v9; - TASSIGN(v21, v22); - pipe_barrier(PIPE_V); - TROWSUM(v21, v17, v19); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1> v23 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 16> v24 = pto::Stride<16, 16, 16, 1, 16>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v25 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN>( - v2 + (v10 + v10 * v6 + v10 * v8), v23, v24 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v25, v21); -#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_out_inline157__rv_v2 - __gm__ Tensor *attn_out_inline157__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *attn_out_inline157__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline157__rv_v2_tensor->buffer.addr) + - attn_out_inline157__rv_v2_tensor->start_offset; - - // Unpack tensor: attn_fence_dummy_inline258__ssa_v0 - __gm__ Tensor *attn_fence_dummy_inline258__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *attn_fence_dummy_inline258__ssa_v0 = - reinterpret_cast<__gm__ float *>(attn_fence_dummy_inline258__ssa_v0_tensor->buffer.addr) + - attn_fence_dummy_inline258__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - attn_fence(attn_out_inline157__rv_v2, attn_fence_dummy_inline258__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/k_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp similarity index 53% rename from examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/k_seed.cpp rename to examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp index 730e1aa6e8..4fc5a9778d 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/k_seed.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_out_seed.cpp @@ -8,7 +8,7 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -// Kernel Function: k_seed +// Kernel Function: attn_out_seed // Generated by PyPTO IR Compiler (PTO backend) #include @@ -50,48 +50,24 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void k_seed(__gm__ float *v1) { - const float v2 = 0.0f; - const int64_t v3 = 1; - const int64_t v4 = 1024; - const int64_t v5 = 16; - const int64_t v6 = 0; +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); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v7 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v5, v4); - uint64_t v8 = (uint64_t)v6; - TASSIGN(v7, v8); - TEXPANDS(v7, v2); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1024> v9 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v10 = pto::Stride<16384, 16384, 16384, 1024, 1>(); - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> v11 = - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v1 + (v6 + v6 * v4 + v6 * v3), v9, v10 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v11, v7); #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_inline128__ssa_v0 - __gm__ Tensor *k_proj_inline128__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *k_proj_inline128__ssa_v0 = - reinterpret_cast<__gm__ float *>(k_proj_inline128__ssa_v0_tensor->buffer.addr) + - k_proj_inline128__ssa_v0_tensor->start_offset; + // 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 - k_seed(k_proj_inline128__ssa_v0); + attn_out_seed(attn_out_inline282__ssa_v0); } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp index 42ebf81015..b743ef0c9e 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_hidden.cpp @@ -50,53 +50,70 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void copy_hidden(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, int64_t v3) { - const int64_t v4 = 256; - const int64_t v5 = 20; - const int64_t v6 = 1; - const int64_t v7 = 5120; - const int64_t v8 = 16; - const int64_t v9 = 0; +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_MTE3, PIPE_MTE2, EVENT_ID0); - for (size_t v10 = (size_t)v9; v10 < ((size_t)v5); v10 += (size_t)v6) { - int64_t v11 = (int64_t)((uint64_t)((int64_t)v10) * (uint64_t)v4); + 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> - v12 = Tile< + v15 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v13 = (uint64_t)v9; - TASSIGN(v12, v13); - pto::Shape<1, 1, 1, 16, 256> v14 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v15 = 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> - v16 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v9 + v3 * v7 + v11 * v6), v14, v15 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v12, v16); - set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); + 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>( - v1 + (v9 + v3 * v7 + v11 * v6), v17, v18 + 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_MTE2, PIPE_MTE3, EVENT_ID0); - TSTORE(v19, v12); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TSTORE(v24, v20); + set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); } - wait_flag(PIPE_MTE3, PIPE_MTE2, 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); @@ -106,8 +123,8 @@ static __aicore__ void copy_hidden(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, 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__ bfloat16_t *cur__iter_v1 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v1_tensor->buffer.addr) + cur__iter_v1_tensor->start_offset; + __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]); diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_out.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_out.cpp index 091274a479..16110bc219 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_out.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/copy_out.cpp @@ -50,50 +50,62 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void copy_out(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, int64_t v3) { - const int64_t v4 = 256; - const int64_t v5 = 20; - const int64_t v6 = 1; - const int64_t v7 = 5120; - const int64_t v8 = 16; - const int64_t v9 = 0; +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 v10 = (size_t)v9; v10 < ((size_t)v5); v10 += (size_t)v6) { - int64_t v11 = (int64_t)((uint64_t)((int64_t)v10) * (uint64_t)v4); + 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, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v12 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v13 = (uint64_t)v9; - TASSIGN(v12, v13); - pto::Shape<1, 1, 1, 16, 256> v14 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v15 = 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> - v16 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v9 + v3 * v7 + v11 * v6), v14, v15 + 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(v12, v16); - set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - 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>(); + 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> - v19 = GlobalTensor< + v23 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v9 + v3 * v7 + v11 * v6), v17, v18 + v1 + (v11 + v3 * v9 + v13 * v8), v21, v22 ); - wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - TSTORE(v19, v12); + 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); @@ -111,8 +123,8 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Unpack tensor: cur__rv_v7 __gm__ Tensor *cur__rv_v7_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *cur__rv_v7 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__rv_v7_tensor->buffer.addr) + cur__rv_v7_tensor->start_offset; + __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 { diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp new file mode 100644 index 0000000000..7af717198c --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_cast_residual.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_cast_residual.cpp deleted file mode 100644 index d5fbc120c3..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_cast_residual.cpp +++ /dev/null @@ -1,180 +0,0 @@ -/* - * 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_cast_residual -// 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_cast_residual(__gm__ bfloat16_t *v1, __gm__ float *v2, __gm__ bfloat16_t *v3, int64_t v4) { - SaturationMode v5 = SaturationMode::OFF; - RoundMode v6 = RoundMode::CAST_ROUND; - const int64_t v7 = 1024; - const int64_t v8 = 1; - const int64_t v9 = 5120; - const int64_t v10 = 16; - const int64_t v11 = 0; - const int64_t v12 = 65536; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - Tile< - TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v13 = Tile< - TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v10, v7); - uint64_t v14 = (uint64_t)v12; - TASSIGN(v13, v14); - pto::Shape<1, 1, 1, 16, 1024> v15 = pto::Shape<1, 1, 1, 16, 1024>(); - 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< - bfloat16_t, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v11 + v11 * v9 + v4 * v8), v15, v16 - ); - TLOAD(v13, v17); - 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> - v18 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v10, v7); - uint64_t v19 = (uint64_t)v11; - TASSIGN(v18, v19); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v18, v13, v6, v5); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v20 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v10, v7); - uint64_t v21 = (uint64_t)v12; - TASSIGN(v20, v21); - pto::Shape<1, 1, 1, 16, 1024> v22 = pto::Shape<1, 1, 1, 16, 1024>(); - 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, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v11 + v11 * v9 + v4 * v8), v22, v23 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TLOAD(v20, v24); - 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> - v25 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v10, v7); - uint64_t v26 = (uint64_t)v11; - TASSIGN(v25, v26); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(v25, v20, v18); - Tile< - TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v27 = Tile< - TileType::Vec, bfloat16_t, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v10, v7); - uint64_t v28 = (uint64_t)v11; - TASSIGN(v27, v28); - pipe_barrier(PIPE_V); - TCVT(v27, v25, v6, v5); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1024> v29 = pto::Shape<1, 1, 1, 16, 1024>(); - 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, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v11 + v11 * v9 + v4 * v8), v29, v30 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v31, v27); -#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: post_norm_partial_inline130__rv_v14 - __gm__ Tensor *post_norm_partial_inline130__rv_v14_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *post_norm_partial_inline130__rv_v14 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__rv_v14_tensor->buffer.addr) + - post_norm_partial_inline130__rv_v14_tensor->start_offset; - - // Unpack tensor: down_acc_all_inline251__rv_v5 - __gm__ Tensor *down_acc_all_inline251__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *down_acc_all_inline251__rv_v5 = - reinterpret_cast<__gm__ float *>(down_acc_all_inline251__rv_v5_tensor->buffer.addr) + - down_acc_all_inline251__rv_v5_tensor->start_offset; - - // Unpack tensor: out_partial_inline312__iter_v1 - __gm__ Tensor *out_partial_inline312__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ bfloat16_t *out_partial_inline312__iter_v1 = - reinterpret_cast<__gm__ bfloat16_t *>(out_partial_inline312__iter_v1_tensor->buffer.addr) + - out_partial_inline312__iter_v1_tensor->start_offset; - - // Unpack scalar: n0_inline142__ssa_v9 - union { - uint64_t u64; - int64_t val; - } n0_inline142__ssa_v9_conv; - n0_inline142__ssa_v9_conv.u64 = args[3]; - int64_t n0_inline142__ssa_v9 = n0_inline142__ssa_v9_conv.val; - - // Forward to ptoas-generated function - down_cast_residual( - post_norm_partial_inline130__rv_v14, down_acc_all_inline251__rv_v5, out_partial_inline312__iter_v1, - n0_inline142__ssa_v9 - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_seed.cpp deleted file mode 100644 index 24855e3146..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_seed.cpp +++ /dev/null @@ -1,156 +0,0 @@ -/* - * 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_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 down_seed(__gm__ float *v1) { - const int64_t v2 = 4096; - const float v3 = 0.0f; - const int64_t v4 = 1024; - const int64_t v5 = 2; - const int64_t v6 = 4; - const int64_t v7 = 1; - const int64_t v8 = 5120; - const int64_t v9 = 16; - const int64_t v10 = 65536; - 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_V, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - for (size_t v12 = (size_t)v11; v12 < ((size_t)v6); v12 += (size_t)v5) { - int64_t v13 = (int64_t)((uint64_t)((int64_t)v12) * (uint64_t)v4); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v14 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v9, v4); - uint64_t v15 = (uint64_t)v11; - TASSIGN(v14, v15); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TEXPANDS(v14, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v16 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v9, v4); - uint64_t v17 = (uint64_t)v10; - TASSIGN(v16, v17); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TEXPANDS(v16, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 1024> v18 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v19 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v20 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v11 + v11 * v8 + v13 * v7), v18, v19 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pipe_barrier(PIPE_MTE3); - TSTORE(v20, v14); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1024> v21 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v22 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v23 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v11 + v11 * v8 + (int64_t)((uint64_t)v13 + (uint64_t)v4) * v7), v21, v22 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v23, v16); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - } - 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> - v24 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v9, v4); - uint64_t v25 = (uint64_t)v11; - TASSIGN(v24, v25); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); - TEXPANDS(v24, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - 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, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v11 + v11 * v8 + v2 * v7), v26, v27 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pipe_barrier(PIPE_MTE3); - TSTORE(v28, v24); - 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: down_acc_all_inline251__ssa_v0 - __gm__ Tensor *down_acc_all_inline251__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *down_acc_all_inline251__ssa_v0 = - reinterpret_cast<__gm__ float *>(down_acc_all_inline251__ssa_v0_tensor->buffer.addr) + - down_acc_all_inline251__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - down_seed(down_acc_all_inline251__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_fused_aiv.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_fused_aiv.cpp deleted file mode 100644 index 5be25921f3..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_fused_aiv.cpp +++ /dev/null @@ -1,779 +0,0 @@ -/* - * 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: fa_fused_aiv -// 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 fa_fused_aic( - __gm__ int32_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ int32_t *v5, __gm__ int32_t *v6, - __gm__ int32_t *v7, __gm__ bfloat16_t *v8, __gm__ bfloat16_t *v9, __gm__ bfloat16_t *v10, __gm__ float *v11, - int64_t v12, int64_t v13, int32_t v14, int32_t v15 -) { - const int64_t v16 = 5; - const int64_t v17 = 8; - const int64_t v18 = 43; - const int64_t v19 = 24; - const int64_t v20 = 16; - const int64_t v21 = 688; - const int64_t v22 = 128; - const int64_t v23 = 1; - const int64_t v24 = 0; - const int64_t v25 = 36864; - const int64_t v26 = 32768; - const int64_t v27 = 7; - const int32_t v28 = 0; - using T = float; - -#if defined(__DAV_CUBE__) - size_t v29 = (size_t)v23; - auto v30 = TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>(v11, v28, v28); - int32_t v31 = v1[v24]; - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - for (size_t v32 = (size_t)((int64_t)v14); v32 < ((size_t)((int64_t)v31)); v32 += (size_t)v19) { - int32_t v33 = v5[v32]; - int64_t v34 = (int64_t)v33; - int64_t v35 = v34 / v18; - int64_t v36 = v34 % v18; - int32_t v37 = v7[(int64_t)((uint64_t)((int64_t)(uint64_t)v35 * (uint64_t)v12) + (uint64_t)v36)]; - int64_t v38 = (int64_t)((uint64_t)v35 * (uint64_t)v22); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v39 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v40 = (uint64_t)v26; - TASSIGN(v39, v40); - pto::Shape<1, 1, 1, 16, 128> v41 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<2048, 2048, 2048, 128, 1> v42 = pto::Stride<2048, 2048, 2048, 128, 1>(); - GlobalTensor, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> - v43 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>( - v8 + (v24 + v38 * v22 + v24 * v23), v41, v42 - ); - TLOAD(v39, v43); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - int64_t v44 = (int64_t)((uint64_t)((int64_t)v37) * (uint64_t)v17); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v45 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null>(v22, v22); - uint64_t v46 = (uint64_t)v25; - TASSIGN(v45, v46); - pto::Shape<1, 1, 1, 128, 128> v47 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<128, 128, 128, 1, 128> v48 = pto::Stride<128, 128, 128, 1, 128>(); - GlobalTensor, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> - v49 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>( - v9 + (v24 + v24 * v23 + - (int64_t)((uint64_t)v13 + (uint64_t)((int64_t)(uint64_t)v44 * (uint64_t)v22)) * v22), - v47, v48 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - TLOAD(v45, v49); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v50 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v51 = (uint64_t)v24; - TASSIGN(v50, v51); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - TMOV(v50, v39); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v52 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v53 = (uint64_t)v24; - TASSIGN(v52, v53); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID1); - TMOV(v52, v45); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v54 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v55 = (uint64_t)v24; - TASSIGN(v54, v55); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - TMATMUL(v54, v50, v52); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v54); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID1); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID1); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID1); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID1); - for (size_t v56 = v29; v56 < ((size_t)v17); v56 += v29) { - int64_t v57 = (int64_t)v56; - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v58 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v20, v22); - uint64_t v59 = (uint64_t)v26; - TASSIGN(v58, v59); - pto::Shape<1, 1, 1, 16, 128> v60 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<2048, 2048, 2048, 128, 1> v61 = pto::Stride<2048, 2048, 2048, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> - v62 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>( - v8 + (v24 + (int64_t)((uint64_t)v38 + (uint64_t)((int64_t)(uint64_t)v57 * (uint64_t)v20)) * v22 + - v24 * v23), - v60, v61 - ); - TLOAD(v58, v62); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v63 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v64 = (uint64_t)v25; - TASSIGN(v63, v64); - pto::Shape<1, 1, 1, 128, 128> v65 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<128, 128, 128, 1, 128> v66 = pto::Stride<128, 128, 128, 1, 128>(); - GlobalTensor, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN> - v67 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<128, 128, 128, 1, 128>, pto::Layout::DN>( - v9 + (v24 + v24 * v23 + - (int64_t)((uint64_t)v13 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v57) * - (uint64_t)v22)) * - v22), - v65, v66 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - TLOAD(v63, v67); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v68 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v20, v22); - uint64_t v69 = (uint64_t)v24; - TASSIGN(v68, v69); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - TMOV(v68, v58); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null> - v70 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v71 = (uint64_t)v24; - TASSIGN(v70, v71); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID3); - TMOV(v70, v63); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v72 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v73 = (uint64_t)v24; - TASSIGN(v72, v73); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID1); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - TMATMUL(v72, v68, v70); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID1); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID1); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v72); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID3); - int64_t v74 = (int64_t)((uint64_t)v57 - (uint64_t)v23); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v75 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v76 = (uint64_t)v25; - TASSIGN(v75, v76); - pto::Shape<1, 1, 1, 128, 128> v77 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<16384, 16384, 16384, 128, 1> v78 = pto::Stride<16384, 16384, 16384, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND> - v79 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, - pto::Layout::ND>( - v10 + (v24 + - (int64_t)((uint64_t)v13 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v74) * - (uint64_t)v22)) * - v22 + - v24 * v23), - v77, v78 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID3); - TLOAD(v75, v79); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v80 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>(v20, v22); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v80); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v81 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, - PadValue::Min, CompactMode::Null>(v20, v22); - uint64_t v82 = (uint64_t)v24; - TASSIGN(v81, v82); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID4); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID3); - TMOV(v81, v80); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v30); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null> - v83 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v84 = (uint64_t)v24; - TASSIGN(v83, v84); - TMOV(v83, v75); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v85 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v86 = (uint64_t)v24; - TASSIGN(v85, v86); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID2); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID3); - TMATMUL(v85, v81, v83); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID2); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - v85.SetValidShape(v16, v22); - pto::Shape<1, 1, 1, 5, 128> v87 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v88 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v89 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v4 + (v24 + - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v35 * - (uint64_t)v17) + - (uint64_t)v74) * - (uint64_t)v21) + - (uint64_t)((int64_t)(uint64_t)v36 * (uint64_t)v20)) * - v22 + - v24 * v23), - v87, v88 - ); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID2); - TSTORE(v89, v85); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID2); - }; - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID4); - Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null> - v90 = Tile< - TileType::Mat, bfloat16_t, 128, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Null, - CompactMode::Null>(v22, v22); - uint64_t v91 = (uint64_t)v25; - TASSIGN(v90, v91); - pto::Shape<1, 1, 1, 128, 128> v92 = pto::Shape<1, 1, 1, 128, 128>(); - pto::Stride<16384, 16384, 16384, 128, 1> v93 = pto::Stride<16384, 16384, 16384, 128, 1>(); - GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND> - v94 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 128, 128>, pto::Stride<16384, 16384, 16384, 128, 1>, pto::Layout::ND>( - v10 + - (v24 + - (int64_t)((uint64_t)v13 + (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v44 + (uint64_t)v27) * - (uint64_t)v22)) * - v22 + - v24 * v23), - v92, v93 - ); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID4); - TLOAD(v90, v94); - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v95 = Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>(v20, v22); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Mat, bfloat16_t, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v30, v95); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); - Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null> - v96 = Tile< - TileType::Left, bfloat16_t, 16, 128, BLayout::RowMajor, -1, -1, SLayout::RowMajor, 512, PadValue::Min, - CompactMode::Null>(v20, v22); - uint64_t v97 = (uint64_t)v24; - TASSIGN(v96, v97); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID5); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID4); - TMOV(v96, v95); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v30); - Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, PadValue::Null, - CompactMode::Null> - v98 = Tile< - TileType::Right, bfloat16_t, 128, 128, BLayout::RowMajor, -1, -1, SLayout::ColMajor, 512, - PadValue::Null, CompactMode::Null>(v22, v22); - uint64_t v99 = (uint64_t)v24; - TASSIGN(v98, v99); - TMOV(v98, v90); - set_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null> - v100 = Tile< - TileType::Acc, float, 16, 128, BLayout::ColMajor, -1, -1, SLayout::RowMajor, 1024, PadValue::Null, - CompactMode::Null>(v20, v22); - uint64_t v101 = (uint64_t)v24; - TASSIGN(v100, v101); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID3); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID4); - TMATMUL(v100, v96, v98); - set_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID3); - v100.SetValidShape(v16, v22); - pto::Shape<1, 1, 1, 5, 128> v102 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v103 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v104 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v4 + (v24 + - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v35 * - (uint64_t)v17) + - (uint64_t)v27) * - (uint64_t)v21) + - (uint64_t)((int64_t)(uint64_t)v36 * (uint64_t)v20)) * - v22 + - v24 * v23), - v102, v103 - ); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID3); - TSTORE(v104, v100); - set_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - } - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_M, PIPE_MTE1, EVENT_ID2); - wait_flag(PIPE_FIX, PIPE_M, EVENT_ID2); -#endif // __DAV_CUBE__ - - ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); - return; -} - -static __aicore__ void fa_fused_aiv( - __gm__ int32_t *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ int32_t *v5, __gm__ int32_t *v6, - __gm__ int32_t *v7, __gm__ bfloat16_t *v8, __gm__ bfloat16_t *v9, __gm__ bfloat16_t *v10, __gm__ float *v11, - int64_t v12, int64_t v13, int32_t v14, int32_t v15, int32_t sub_block_id -) { - SaturationMode v16 = SaturationMode::OFF; - RoundMode v17 = RoundMode::CAST_ROUND; - const int64_t v18 = 5; - const float v19 = 0.0883883461f; - const int64_t v20 = 8; - const int64_t v21 = 43; - const int64_t v22 = 24; - const int64_t v23 = 0; - const int64_t v24 = 16; - const int64_t v25 = 688; - const int64_t v26 = 128; - const int64_t v27 = 88064; - const int64_t v28 = 1; - const int64_t v29 = 43040; - const int64_t v30 = 40992; - const int64_t v31 = 40960; - const int64_t v32 = 32768; - const int64_t v33 = 36864; - const int32_t v34 = 0; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - int64_t v35 = sub_block_id; - auto v36 = TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>(v11, v34, v34); - // simpler's runtime leaves the AICore sub-block register at 0 for both AIV - // lanes, so the library's TILE_UP_DOWN auto-split (get_subblockid() * bytes) - // contributes nothing. Supply the per-lane offset explicitly from the - // runtime lane id, as spmd_paged_attention does. Tiles are 8x128: - // cons = C2V float pop, prod = V2C bf16 push. - v36.cons.setEntryOffset(static_cast(sub_block_id) * 8 * 128 * static_cast(sizeof(float))); - v36.prod.setEntryOffset(static_cast(sub_block_id) * 8 * 128 * static_cast(sizeof(bfloat16_t))); - int32_t v37 = v1[v23]; - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - for (size_t v38 = (size_t)((int64_t)v14); v38 < ((size_t)((int64_t)v37)); v38 += (size_t)v22) { - int32_t v39 = v5[v38]; - int64_t v40 = (int64_t)v39; - int64_t v41 = v40 / v21; - int64_t v42 = v40 % v21; - int32_t v43 = v6[v41]; - int64_t v44 = (int64_t)((uint64_t)((int64_t)v43) - (uint64_t)((int64_t)(uint64_t)v42 * (uint64_t)v26)); - int64_t v45 = v44 < v26 ? v44 : v26; - for (size_t v46 = (size_t)v23; v46 < ((size_t)v20); v46 += (size_t)v28) { - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v47 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TPOP< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v36, v47); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v48 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v49 = (uint64_t)v33; - TASSIGN(v48, v49); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - pipe_barrier(PIPE_V); - TMULS(v48, v47, v19); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TFREE, TileSplitAxis::TILE_UP_DOWN>(v36); - v48.SetValidShape(v18, v45); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v50 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v51 = (uint64_t)v32; - TASSIGN(v50, v51); - pipe_barrier(PIPE_V); - TFILLPAD(v50, v48); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v52 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v53 = (uint64_t)v33; - TASSIGN(v52, v53); - Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v54 = Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v28); - uint64_t v55 = (uint64_t)v31; - TASSIGN(v54, v55); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TROWMAX(v54, v50, v52); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v56 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v57 = (uint64_t)v32; - TASSIGN(v56, v57); - pipe_barrier(PIPE_V); - TROWEXPANDSUB(v56, v50, v54); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v58 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v59 = (uint64_t)v32; - TASSIGN(v58, v59); - pipe_barrier(PIPE_V); - TEXP(v58, v56); - Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v60 = Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v61 = (uint64_t)v30; - TASSIGN(v60, v61); - pipe_barrier(PIPE_V); - TCVT(v60, v58, v17, v16); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null> - v62 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>(v20, v26); - uint64_t v63 = (uint64_t)v32; - TASSIGN(v62, v63); - pipe_barrier(PIPE_V); - TCVT(v62, v60, v17, v16); - Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v64 = Tile< - TileType::Vec, float, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v26); - uint64_t v65 = (uint64_t)v33; - TASSIGN(v64, v65); - Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v66 = Tile< - TileType::Vec, float, 8, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v28); - uint64_t v67 = (uint64_t)v29; - TASSIGN(v66, v67); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TROWSUM(v66, v62, v64); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TPUSH< - TPipe<0, Direction::DIR_BOTH, 8192, 4, 4, false>, - Tile< - TileType::Vec, bfloat16_t, 8, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Min, - CompactMode::Null>, - TileSplitAxis::TILE_UP_DOWN>(v36, v60); - int64_t v68 = - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t) - v41 * - (uint64_t)v20) + - (uint64_t)((int64_t)v46)) * - (uint64_t)v25) + - (uint64_t)((int64_t)(uint64_t)v42 * (uint64_t)v24)) + - (uint64_t)((int64_t)(uint64_t)((int64_t)v35) * (uint64_t)v20)); - pto::Shape<1, 1, 1, 8, 1> v69 = pto::Shape<1, 1, 1, 8, 1>(); - pto::Stride<8, 8, 8, 1, 88064> v70 = pto::Stride<8, 8, 8, 1, 88064>(); - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN> v71 = - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN>( - v3 + (v23 + v68 * v28 + v23 * v27), v69, v70 - ); - TSTORE(v71, v54); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pto::Shape<1, 1, 1, 8, 1> v72 = pto::Shape<1, 1, 1, 8, 1>(); - pto::Stride<8, 8, 8, 1, 88064> v73 = pto::Stride<8, 8, 8, 1, 88064>(); - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN> v74 = - GlobalTensor, pto::Stride<8, 8, 8, 1, 88064>, pto::Layout::DN>( - v2 + (v23 + v68 * v28 + v23 * v27), v72, v73 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v74, v66); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - }; - } - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - 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) { - // Read A2A3 mixed-task subblock id from runtime dispatch context. - // Must use the simpler runtime accessor (not the CCE get_subblockid() - // builtin, which is stale 0 for both AIV lanes under MIX dispatch). - int32_t __pypto_sub_block_id = get_sub_block_id(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: fa_total_inline183__ssa_v0 - __gm__ Tensor *fa_total_inline183__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ int32_t *fa_total_inline183__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_total_inline183__ssa_v0_tensor->buffer.addr) + - fa_total_inline183__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_cur_li_inline190__ssa_v0 - __gm__ Tensor *all_cur_li_inline190__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *all_cur_li_inline190__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_cur_li_inline190__ssa_v0_tensor->buffer.addr) + - all_cur_li_inline190__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_cur_mi_inline174__ssa_v0 - __gm__ Tensor *all_cur_mi_inline174__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *all_cur_mi_inline174__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_cur_mi_inline174__ssa_v0_tensor->buffer.addr) + - all_cur_mi_inline174__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_oi_tmp_inline145__ssa_v0 - __gm__ Tensor *all_oi_tmp_inline145__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *all_oi_tmp_inline145__ssa_v0 = - reinterpret_cast<__gm__ float *>(all_oi_tmp_inline145__ssa_v0_tensor->buffer.addr) + - all_oi_tmp_inline145__ssa_v0_tensor->start_offset; - - // Unpack tensor: fa_work_table_inline126__ssa_v0 - __gm__ Tensor *fa_work_table_inline126__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ int32_t *fa_work_table_inline126__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_work_table_inline126__ssa_v0_tensor->buffer.addr) + - fa_work_table_inline126__ssa_v0_tensor->start_offset; - - // Unpack tensor: seq_lens__ssa_v0 - __gm__ Tensor *seq_lens__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[5]); - __gm__ int32_t *seq_lens__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(seq_lens__ssa_v0_tensor->buffer.addr) + - seq_lens__ssa_v0_tensor->start_offset; - - // Unpack tensor: block_table__ssa_v0 - __gm__ Tensor *block_table__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[6]); - __gm__ int32_t *block_table__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(block_table__ssa_v0_tensor->buffer.addr) + - block_table__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_q_padded_inline168__rv_v2 - __gm__ Tensor *all_q_padded_inline168__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[7]); - __gm__ bfloat16_t *all_q_padded_inline168__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(all_q_padded_inline168__rv_v2_tensor->buffer.addr) + - all_q_padded_inline168__rv_v2_tensor->start_offset; - - // Unpack tensor: k_cache__rv_v4 - __gm__ Tensor *k_cache__rv_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[8]); - __gm__ bfloat16_t *k_cache__rv_v4 = - reinterpret_cast<__gm__ bfloat16_t *>(k_cache__rv_v4_tensor->buffer.addr) + k_cache__rv_v4_tensor->start_offset; - - // Unpack tensor: v_cache__rv_v4 - __gm__ Tensor *v_cache__rv_v4_tensor = reinterpret_cast<__gm__ Tensor *>(args[9]); - __gm__ bfloat16_t *v_cache__rv_v4 = - reinterpret_cast<__gm__ bfloat16_t *>(v_cache__rv_v4_tensor->buffer.addr) + v_cache__rv_v4_tensor->start_offset; - - // Unpack tensor: __gm_pipe_buffer - __gm__ Tensor *__gm_pipe_buffer_tensor = reinterpret_cast<__gm__ Tensor *>(args[10]); - // SPMD: shard GM pipe workspace by logical block_idx to avoid overlap. - int64_t __pypto_gm_block_num = static_cast(__pypto_spmd_block_num); - if (__pypto_gm_block_num <= 0) __pypto_gm_block_num = 1; - int64_t __pypto_gm_total_elems = static_cast(__gm_pipe_buffer_tensor->shapes[0]); - int64_t __pypto_gm_elems_per_block = __pypto_gm_total_elems / __pypto_gm_block_num; - int64_t __pypto_gm_block_offset = static_cast(__pypto_spmd_block_idx) * __pypto_gm_elems_per_block; - __gm__ float *__gm_pipe_buffer = reinterpret_cast<__gm__ float *>(__gm_pipe_buffer_tensor->buffer.addr) + - __gm_pipe_buffer_tensor->start_offset + __pypto_gm_block_offset; - - // Unpack scalar: max_blocks_per_seq_inline165__ssa_v0 - union { - uint64_t u64; - int64_t val; - } max_blocks_per_seq_inline165__ssa_v0_conv; - max_blocks_per_seq_inline165__ssa_v0_conv.u64 = args[11]; - int64_t max_blocks_per_seq_inline165__ssa_v0 = max_blocks_per_seq_inline165__ssa_v0_conv.val; - - // Unpack scalar: layer_cache_base_inline137__ssa_v0 - union { - uint64_t u64; - int64_t val; - } layer_cache_base_inline137__ssa_v0_conv; - layer_cache_base_inline137__ssa_v0_conv.u64 = args[12]; - int64_t layer_cache_base_inline137__ssa_v0 = layer_cache_base_inline137__ssa_v0_conv.val; - - // Forward to ptoas-generated function - fa_fused_aiv( - fa_total_inline183__ssa_v0, all_cur_li_inline190__ssa_v0, all_cur_mi_inline174__ssa_v0, - all_oi_tmp_inline145__ssa_v0, fa_work_table_inline126__ssa_v0, seq_lens__ssa_v0, block_table__ssa_v0, - all_q_padded_inline168__rv_v2, k_cache__rv_v4, v_cache__rv_v4, __gm_pipe_buffer, - max_blocks_per_seq_inline165__ssa_v0, layer_cache_base_inline137__ssa_v0, __pypto_spmd_block_idx, - __pypto_spmd_block_num, __pypto_sub_block_id - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_work_build.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_work_build.cpp deleted file mode 100644 index aafce55b3b..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_work_build.cpp +++ /dev/null @@ -1,268 +0,0 @@ -/* - * 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: fa_work_build -// 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 fa_work_build(__gm__ int32_t *v1, __gm__ int32_t *v2, __gm__ int32_t *v3) { - const int64_t v4 = 645; - const int64_t v5 = 15; - const int64_t v6 = 602; - const int64_t v7 = 14; - const int64_t v8 = 559; - const int64_t v9 = 13; - const int64_t v10 = 516; - const int64_t v11 = 12; - const int64_t v12 = 473; - const int64_t v13 = 11; - const int64_t v14 = 430; - const int64_t v15 = 10; - const int64_t v16 = 387; - const int64_t v17 = 9; - const int64_t v18 = 344; - const int64_t v19 = 8; - const int64_t v20 = 301; - const int64_t v21 = 7; - const int64_t v22 = 258; - const int64_t v23 = 6; - const int64_t v24 = 215; - const int64_t v25 = 5; - const int64_t v26 = 172; - const int64_t v27 = 4; - const int64_t v28 = 129; - const int64_t v29 = 3; - const int64_t v30 = 86; - const int64_t v31 = 2; - const int64_t v32 = 43; - const int64_t v33 = 128; - const int64_t v34 = 127; - const int64_t v35 = 0; - const int64_t v36 = 1; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - size_t v37 = (size_t)v36; - size_t v38 = (size_t)v35; - int32_t v39 = v1[v35]; - int64_t v40 = (int64_t)((uint64_t)((int64_t)v39) + (uint64_t)v34) / v33; - for (size_t v41 = v38; v41 < ((size_t)v40); v41 += v37) { - int32_t v42 = (int32_t)((int64_t)v41); - v2[v41] = v42; - } - int32_t v43 = v1[v36]; - int64_t v44 = (int64_t)((uint64_t)((int64_t)v43) + (uint64_t)v34) / v33; - for (size_t v45 = v38; v45 < ((size_t)v44); v45 += v37) { - int64_t v46 = (int64_t)v45; - int32_t v47 = (int32_t)((int64_t)(uint64_t)v46 + (uint64_t)v32); - int64_t v48 = (int64_t)((uint64_t)v40 + (uint64_t)v46); - v2[v48] = v47; - } - int64_t v49 = (int64_t)((uint64_t)v40 + (uint64_t)v44); - int32_t v50 = v1[v31]; - int64_t v51 = (int64_t)((uint64_t)((int64_t)v50) + (uint64_t)v34) / v33; - for (size_t v52 = v38; v52 < ((size_t)v51); v52 += v37) { - int64_t v53 = (int64_t)v52; - int32_t v54 = (int32_t)((int64_t)(uint64_t)v53 + (uint64_t)v30); - int64_t v55 = (int64_t)((uint64_t)v49 + (uint64_t)v53); - v2[v55] = v54; - } - int64_t v56 = (int64_t)((uint64_t)v49 + (uint64_t)v51); - int32_t v57 = v1[v29]; - int64_t v58 = (int64_t)((uint64_t)((int64_t)v57) + (uint64_t)v34) / v33; - for (size_t v59 = v38; v59 < ((size_t)v58); v59 += v37) { - int64_t v60 = (int64_t)v59; - int32_t v61 = (int32_t)((int64_t)(uint64_t)v60 + (uint64_t)v28); - int64_t v62 = (int64_t)((uint64_t)v56 + (uint64_t)v60); - v2[v62] = v61; - } - int64_t v63 = (int64_t)((uint64_t)v56 + (uint64_t)v58); - int32_t v64 = v1[v27]; - int64_t v65 = (int64_t)((uint64_t)((int64_t)v64) + (uint64_t)v34) / v33; - for (size_t v66 = v38; v66 < ((size_t)v65); v66 += v37) { - int64_t v67 = (int64_t)v66; - int32_t v68 = (int32_t)((int64_t)(uint64_t)v67 + (uint64_t)v26); - int64_t v69 = (int64_t)((uint64_t)v63 + (uint64_t)v67); - v2[v69] = v68; - } - int64_t v70 = (int64_t)((uint64_t)v63 + (uint64_t)v65); - int32_t v71 = v1[v25]; - int64_t v72 = (int64_t)((uint64_t)((int64_t)v71) + (uint64_t)v34) / v33; - for (size_t v73 = v38; v73 < ((size_t)v72); v73 += v37) { - int64_t v74 = (int64_t)v73; - int32_t v75 = (int32_t)((int64_t)(uint64_t)v74 + (uint64_t)v24); - int64_t v76 = (int64_t)((uint64_t)v70 + (uint64_t)v74); - v2[v76] = v75; - } - int64_t v77 = (int64_t)((uint64_t)v70 + (uint64_t)v72); - int32_t v78 = v1[v23]; - int64_t v79 = (int64_t)((uint64_t)((int64_t)v78) + (uint64_t)v34) / v33; - for (size_t v80 = v38; v80 < ((size_t)v79); v80 += v37) { - int64_t v81 = (int64_t)v80; - int32_t v82 = (int32_t)((int64_t)(uint64_t)v81 + (uint64_t)v22); - int64_t v83 = (int64_t)((uint64_t)v77 + (uint64_t)v81); - v2[v83] = v82; - } - int64_t v84 = (int64_t)((uint64_t)v77 + (uint64_t)v79); - int32_t v85 = v1[v21]; - int64_t v86 = (int64_t)((uint64_t)((int64_t)v85) + (uint64_t)v34) / v33; - for (size_t v87 = v38; v87 < ((size_t)v86); v87 += v37) { - int64_t v88 = (int64_t)v87; - int32_t v89 = (int32_t)((int64_t)(uint64_t)v88 + (uint64_t)v20); - int64_t v90 = (int64_t)((uint64_t)v84 + (uint64_t)v88); - v2[v90] = v89; - } - int64_t v91 = (int64_t)((uint64_t)v84 + (uint64_t)v86); - int32_t v92 = v1[v19]; - int64_t v93 = (int64_t)((uint64_t)((int64_t)v92) + (uint64_t)v34) / v33; - for (size_t v94 = v38; v94 < ((size_t)v93); v94 += v37) { - int64_t v95 = (int64_t)v94; - int32_t v96 = (int32_t)((int64_t)(uint64_t)v95 + (uint64_t)v18); - int64_t v97 = (int64_t)((uint64_t)v91 + (uint64_t)v95); - v2[v97] = v96; - } - int64_t v98 = (int64_t)((uint64_t)v91 + (uint64_t)v93); - int32_t v99 = v1[v17]; - int64_t v100 = (int64_t)((uint64_t)((int64_t)v99) + (uint64_t)v34) / v33; - for (size_t v101 = v38; v101 < ((size_t)v100); v101 += v37) { - int64_t v102 = (int64_t)v101; - int32_t v103 = (int32_t)((int64_t)(uint64_t)v102 + (uint64_t)v16); - int64_t v104 = (int64_t)((uint64_t)v98 + (uint64_t)v102); - v2[v104] = v103; - } - int64_t v105 = (int64_t)((uint64_t)v98 + (uint64_t)v100); - int32_t v106 = v1[v15]; - int64_t v107 = (int64_t)((uint64_t)((int64_t)v106) + (uint64_t)v34) / v33; - for (size_t v108 = v38; v108 < ((size_t)v107); v108 += v37) { - int64_t v109 = (int64_t)v108; - int32_t v110 = (int32_t)((int64_t)(uint64_t)v109 + (uint64_t)v14); - int64_t v111 = (int64_t)((uint64_t)v105 + (uint64_t)v109); - v2[v111] = v110; - } - int64_t v112 = (int64_t)((uint64_t)v105 + (uint64_t)v107); - int32_t v113 = v1[v13]; - int64_t v114 = (int64_t)((uint64_t)((int64_t)v113) + (uint64_t)v34) / v33; - for (size_t v115 = v38; v115 < ((size_t)v114); v115 += v37) { - int64_t v116 = (int64_t)v115; - int32_t v117 = (int32_t)((int64_t)(uint64_t)v116 + (uint64_t)v12); - int64_t v118 = (int64_t)((uint64_t)v112 + (uint64_t)v116); - v2[v118] = v117; - } - int64_t v119 = (int64_t)((uint64_t)v112 + (uint64_t)v114); - int32_t v120 = v1[v11]; - int64_t v121 = (int64_t)((uint64_t)((int64_t)v120) + (uint64_t)v34) / v33; - for (size_t v122 = v38; v122 < ((size_t)v121); v122 += v37) { - int64_t v123 = (int64_t)v122; - int32_t v124 = (int32_t)((int64_t)(uint64_t)v123 + (uint64_t)v10); - int64_t v125 = (int64_t)((uint64_t)v119 + (uint64_t)v123); - v2[v125] = v124; - } - int64_t v126 = (int64_t)((uint64_t)v119 + (uint64_t)v121); - int32_t v127 = v1[v9]; - int64_t v128 = (int64_t)((uint64_t)((int64_t)v127) + (uint64_t)v34) / v33; - for (size_t v129 = v38; v129 < ((size_t)v128); v129 += v37) { - int64_t v130 = (int64_t)v129; - int32_t v131 = (int32_t)((int64_t)(uint64_t)v130 + (uint64_t)v8); - int64_t v132 = (int64_t)((uint64_t)v126 + (uint64_t)v130); - v2[v132] = v131; - } - int64_t v133 = (int64_t)((uint64_t)v126 + (uint64_t)v128); - int32_t v134 = v1[v7]; - int64_t v135 = (int64_t)((uint64_t)((int64_t)v134) + (uint64_t)v34) / v33; - for (size_t v136 = v38; v136 < ((size_t)v135); v136 += v37) { - int64_t v137 = (int64_t)v136; - int32_t v138 = (int32_t)((int64_t)(uint64_t)v137 + (uint64_t)v6); - int64_t v139 = (int64_t)((uint64_t)v133 + (uint64_t)v137); - v2[v139] = v138; - } - int64_t v140 = (int64_t)((uint64_t)v133 + (uint64_t)v135); - int32_t v141 = v1[v5]; - int64_t v142 = (int64_t)((uint64_t)((int64_t)v141) + (uint64_t)v34) / v33; - for (size_t v143 = v38; v143 < ((size_t)v142); v143 += v37) { - int64_t v144 = (int64_t)v143; - int32_t v145 = (int32_t)((int64_t)(uint64_t)v144 + (uint64_t)v4); - int64_t v146 = (int64_t)((uint64_t)v140 + (uint64_t)v144); - v2[v146] = v145; - } - int32_t v147 = (int32_t)((int64_t)(uint64_t)v140 + (uint64_t)v142); - v3[v35] = v147; - pipe_barrier(PIPE_ALL); - dcci((__gm__ void *)0, ENTIRE_DATA_CACHE, CACHELINE_OUT); - dsb((mem_dsb_t)0); -#endif // __DAV_VEC__ - - pipe_barrier(PIPE_ALL); - dcci((__gm__ void *)0, ENTIRE_DATA_CACHE, CACHELINE_OUT); - dsb((mem_dsb_t)0); - 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: seq_lens__ssa_v0 - __gm__ Tensor *seq_lens__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ int32_t *seq_lens__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(seq_lens__ssa_v0_tensor->buffer.addr) + - seq_lens__ssa_v0_tensor->start_offset; - - // Unpack tensor: fa_work_table_inline126__ssa_v0 - __gm__ Tensor *fa_work_table_inline126__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ int32_t *fa_work_table_inline126__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_work_table_inline126__ssa_v0_tensor->buffer.addr) + - fa_work_table_inline126__ssa_v0_tensor->start_offset; - - // Unpack tensor: fa_total_inline183__ssa_v0 - __gm__ Tensor *fa_total_inline183__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ int32_t *fa_total_inline183__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(fa_total_inline183__ssa_v0_tensor->buffer.addr) + - fa_total_inline183__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - fa_work_build(seq_lens__ssa_v0, fa_work_table_inline126__ssa_v0, fa_total_inline183__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/gate_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/gate_seed.cpp deleted file mode 100644 index 54d96c9658..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/gate_seed.cpp +++ /dev/null @@ -1,158 +0,0 @@ -/* - * 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_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 gate_seed(__gm__ float *v1) { - const int64_t v2 = 16384; - const float v3 = 0.0f; - const int64_t v4 = 1024; - const int64_t v5 = 2; - const int64_t v6 = 1; - const int64_t v7 = 17408; - const int64_t v8 = 16; - const int64_t v9 = 65536; - 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)v8); v11 += (size_t)v5) { - int64_t v12 = (int64_t)((uint64_t)((int64_t)v11) * (uint64_t)v4); - 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>(v8, v4); - uint64_t v14 = (uint64_t)v10; - TASSIGN(v13, v14); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TEXPANDS(v13, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v15 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v16 = (uint64_t)v9; - TASSIGN(v15, v16); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TEXPANDS(v15, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 1024> v17 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v18 = 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> - v19 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v12 * v6), v17, v18 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pipe_barrier(PIPE_MTE3); - TSTORE(v19, v13); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1024> v20 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v21 = 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> - v22 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + (int64_t)((uint64_t)v12 + (uint64_t)v4) * v6), v20, v21 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v22, v15); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - } - 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> - v23 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v24 = (uint64_t)v10; - TASSIGN(v23, v24); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); - TEXPANDS(v23, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pto::Shape<1, 1, 1, 16, 1024> v25 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v26 = pto::Stride<278528, 278528, 278528, 17408, 1>(); - GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> - v27 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v2 * v6), v25, v26 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pipe_barrier(PIPE_MTE3); - TSTORE(v27, v23); - 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: gate_acc_all_inline170__ssa_v0 - __gm__ Tensor *gate_acc_all_inline170__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *gate_acc_all_inline170__ssa_v0 = - reinterpret_cast<__gm__ float *>(gate_acc_all_inline170__ssa_v0_tensor->buffer.addr) + - gate_acc_all_inline170__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - gate_seed(gate_acc_all_inline170__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/v_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/kv_seed.cpp similarity index 51% rename from examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/v_seed.cpp rename to examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/kv_seed.cpp index b54c1a5f92..632e4bb663 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/v_seed.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/kv_seed.cpp @@ -8,7 +8,7 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -// Kernel Function: v_seed +// Kernel Function: kv_seed // Generated by PyPTO IR Compiler (PTO backend) #include @@ -50,12 +50,12 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void v_seed(__gm__ float *v1) { - const float v2 = 0.0f; - const int64_t v3 = 1; - const int64_t v4 = 1024; - const int64_t v5 = 16; - const int64_t v6 = 0; +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__) @@ -64,21 +64,41 @@ static __aicore__ void v_seed(__gm__ float *v1) { Tile< TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v7 = Tile< + v8 = Tile< TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v5, v4); - uint64_t v8 = (uint64_t)v6; - TASSIGN(v7, v8); - TEXPANDS(v7, v2); + 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> v9 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v10 = pto::Stride<16384, 16384, 16384, 1024, 1>(); - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> v11 = + 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 + (v6 + v6 * v4 + v6 * v3), v9, v10 + v1 + (v7 + v7 * v5 + v7 * v4), v10, v11 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v11, v7); + 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); @@ -86,12 +106,18 @@ static __aicore__ void v_seed(__gm__ float *v1) { } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: v_proj_inline152__ssa_v0 - __gm__ Tensor *v_proj_inline152__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *v_proj_inline152__ssa_v0 = - reinterpret_cast<__gm__ float *>(v_proj_inline152__ssa_v0_tensor->buffer.addr) + - v_proj_inline152__ssa_v0_tensor->start_offset; + // 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 - v_seed(v_proj_inline152__ssa_v0); + kv_seed(k_proj_inline135__ssa_v0, v_proj_inline255__ssa_v0); } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp new file mode 100644 index 0000000000..d528272138 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/online_softmax.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/online_softmax.cpp deleted file mode 100644 index 36dc59acc1..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/online_softmax.cpp +++ /dev/null @@ -1,1185 +0,0 @@ -/* - * 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: online_softmax -// 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 online_softmax( - __gm__ bfloat16_t *v1, __gm__ int32_t *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, int32_t v6, - int32_t v7 -) { - SaturationMode v8 = SaturationMode::OFF; - RoundMode v9 = RoundMode::CAST_ROUND; - const int64_t v10 = 640; - const int64_t v11 = 2; - const int64_t v12 = 688; - const int64_t v13 = 5; - const int64_t v14 = 127; - const int64_t v15 = 8; - const int64_t v16 = 48; - const int64_t v17 = 128; - const int64_t v18 = 88064; - const int64_t v19 = 1; - const int64_t v20 = 5120; - const int64_t v21 = 16; - const int64_t v22 = 8512; - const int64_t v23 = 8448; - const int64_t v24 = 8384; - const int64_t v25 = 8320; - const int64_t v26 = 8256; - const int64_t v27 = 64; - const int64_t v28 = 0; - const int64_t v29 = 8576; - const int64_t v30 = 25088; - const int64_t v31 = 25024; - const int64_t v32 = 24960; - const int64_t v33 = 16768; - 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> - v34 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v35 = (uint64_t)v29; - TASSIGN(v34, v35); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v36 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v37 = (uint64_t)v30; - TASSIGN(v36, v37); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - for (size_t v38 = (size_t)((int64_t)v6); v38 < ((size_t)v17); v38 += (size_t)v16) { - int64_t v39 = (int64_t)v38; - int64_t v40 = v39 / v15; - int64_t v41 = v39 % v15; - int32_t v42 = v2[v40]; - int64_t v43 = (int64_t)((uint64_t)((int64_t)v42) + (uint64_t)v14) / v17; - int64_t v44 = - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v40 * (uint64_t)v15) + (uint64_t)v41) * - (uint64_t)v12); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v45 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v46 = (uint64_t)v33; - TASSIGN(v45, v46); - pto::Shape<1, 1, 1, 16, 128> v47 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<2048, 2048, 2048, 128, 1> v48 = pto::Stride<2048, 2048, 2048, 128, 1>(); - GlobalTensor, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND> v49 = - GlobalTensor, pto::Stride<2048, 2048, 2048, 128, 1>, pto::Layout::ND>( - v3 + (v28 + v44 * v17 + v28 * v19), v47, v48 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v45, v49); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v50 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v51 = (uint64_t)v32; - TASSIGN(v50, v51); - pto::Shape<1, 1, 1, 16, 1> v52 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 88064> v53 = pto::Stride<16, 16, 16, 1, 88064>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 88064>, pto::Layout::DN> v54 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 88064>, pto::Layout::DN>( - v4 + (v28 + v44 * v19 + v28 * v18), v52, v53 - ); - TLOAD(v50, 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>(v21, v19); - uint64_t v56 = (uint64_t)v31; - TASSIGN(v55, v56); - pto::Shape<1, 1, 1, 16, 1> v57 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 88064> v58 = pto::Stride<16, 16, 16, 1, 88064>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 88064>, pto::Layout::DN> v59 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 88064>, pto::Layout::DN>( - v5 + (v28 + v44 * v19 + v28 * v18), v57, v58 - ); - TLOAD(v55, v59); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - int64_t v60 = - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)((uint64_t)v43 - (uint64_t)v19) / v11) * (uint64_t)v11) + - (uint64_t)v19); - for (size_t v61 = (size_t)v19; v61 < ((size_t)v60); v61 += (size_t)v11) { - int64_t v62 = (int64_t)((uint64_t)((int64_t)v61) * (uint64_t)v21); - int64_t v63 = (int64_t)((uint64_t)v44 + (uint64_t)v62); - int64_t v64 = (int64_t)((uint64_t)v44 + (uint64_t)((int64_t)(uint64_t)v62 + (uint64_t)v21)); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v65 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v66 = (uint64_t)v30; - TASSIGN(v65, v66); - pto::Shape<1, 1, 1, 5, 128> v67 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v68 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v69 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v3 + (v28 + v63 * v17 + v28 * v19), v67, v68 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TLOAD(v65, v69); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v70 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v71 = (uint64_t)v29; - TASSIGN(v70, v71); - pto::Shape<1, 1, 1, 5, 1> v72 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v73 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v74 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v4 + (v28 + v63 * v19 + v28 * v18), v72, v73 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - TLOAD(v70, v74); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - 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, v19); - uint64_t v76 = (uint64_t)v28; - TASSIGN(v75, v76); - pto::Shape<1, 1, 1, 5, 1> v77 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v78 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v79 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v5 + (v28 + v63 * v19 + v28 * v18), v77, v78 - ); - TLOAD(v75, v79); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v80 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v81 = (uint64_t)v27; - TASSIGN(v80, v81); - pto::Shape<1, 1, 1, 5, 128> v82 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v83 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v84 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v3 + (v28 + v64 * v17 + v28 * v19), v82, v83 - ); - TLOAD(v80, 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, v19); - uint64_t v86 = (uint64_t)v26; - TASSIGN(v85, v86); - pto::Shape<1, 1, 1, 5, 1> v87 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v88 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v89 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v4 + (v28 + v64 * v19 + v28 * v18), v87, v88 - ); - TLOAD(v85, v89); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v90 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v91 = (uint64_t)v25; - TASSIGN(v90, v91); - pto::Shape<1, 1, 1, 5, 1> v92 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v93 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v94 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v5 + (v28 + v64 * v19 + v28 * v18), v92, v93 - ); - TLOAD(v90, v94); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - 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>(v19, v21); - uint64_t v96 = (uint64_t)v32; - 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>(v19, v21); - uint64_t v98 = (uint64_t)v29; - TASSIGN(v97, v98); - 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>(v19, v21); - uint64_t v100 = (uint64_t)v24; - TASSIGN(v99, v100); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TMAX(v99, v95, v97); - 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>(v19, v21); - uint64_t v102 = (uint64_t)v32; - TASSIGN(v101, v102); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v103 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v104 = (uint64_t)v24; - TASSIGN(v103, v104); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v105 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v106 = (uint64_t)v23; - TASSIGN(v105, v106); - pipe_barrier(PIPE_V); - TSUB(v105, v101, v103); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v107 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v108 = (uint64_t)v23; - TASSIGN(v107, v108); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v109 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v110 = (uint64_t)v23; - TASSIGN(v109, v110); - pipe_barrier(PIPE_V); - TEXP(v109, v107); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v111 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v112 = (uint64_t)v23; - TASSIGN(v111, v112); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v113 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v114 = (uint64_t)v29; - TASSIGN(v113, v114); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v115 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v116 = (uint64_t)v24; - TASSIGN(v115, v116); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v117 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v118 = (uint64_t)v29; - TASSIGN(v117, v118); - TSUB(v117, v113, v115); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v119 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v120 = (uint64_t)v29; - TASSIGN(v119, v120); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v121 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v122 = (uint64_t)v22; - TASSIGN(v121, v122); - pipe_barrier(PIPE_V); - TEXP(v121, v119); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v123 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v124 = (uint64_t)v22; - TASSIGN(v123, v124); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v125 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v126 = (uint64_t)v23; - TASSIGN(v125, v126); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v127 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v128 = (uint64_t)v31; - TASSIGN(v127, v128); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v129 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v130 = (uint64_t)v29; - TASSIGN(v129, v130); - pipe_barrier(PIPE_V); - TMUL(v129, v125, v127); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v131 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v132 = (uint64_t)v22; - TASSIGN(v131, v132); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v133 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v134 = (uint64_t)v28; - TASSIGN(v133, v134); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v135 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v136 = (uint64_t)v28; - TASSIGN(v135, v136); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TMUL(v135, v131, v133); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v137 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v138 = (uint64_t)v29; - TASSIGN(v137, v138); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v139 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v140 = (uint64_t)v28; - TASSIGN(v139, v140); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v141 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v142 = (uint64_t)v28; - TASSIGN(v141, v142); - pipe_barrier(PIPE_V); - TADD(v141, v137, v139); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v143 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v144 = (uint64_t)v29; - TASSIGN(v143, v144); - pipe_barrier(PIPE_V); - TROWEXPANDMUL(v143, v45, v111); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v145 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v146 = (uint64_t)v30; - TASSIGN(v145, v146); - TROWEXPANDMUL(v145, v65, v123); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v147 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v148 = (uint64_t)v30; - TASSIGN(v147, v148); - pipe_barrier(PIPE_V); - TADD(v147, v143, v145); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v149 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v150 = (uint64_t)v24; - TASSIGN(v149, v150); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v151 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v152 = (uint64_t)v26; - TASSIGN(v151, v152); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v153 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v154 = (uint64_t)v29; - TASSIGN(v153, v154); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TMAX(v153, v149, v151); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v155 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v156 = (uint64_t)v24; - TASSIGN(v155, v156); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v157 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v158 = (uint64_t)v29; - TASSIGN(v157, v158); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v159 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v160 = (uint64_t)v24; - TASSIGN(v159, v160); - pipe_barrier(PIPE_V); - TSUB(v159, v155, v157); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v161 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v162 = (uint64_t)v24; - TASSIGN(v161, v162); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v163 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v164 = (uint64_t)v24; - TASSIGN(v163, v164); - pipe_barrier(PIPE_V); - TEXP(v163, v161); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v165 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v166 = (uint64_t)v24; - TASSIGN(v165, v166); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v167 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v168 = (uint64_t)v26; - TASSIGN(v167, v168); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v169 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v170 = (uint64_t)v29; - TASSIGN(v169, v170); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v171 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v172 = (uint64_t)v26; - TASSIGN(v171, v172); - TSUB(v171, v167, v169); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v173 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v174 = (uint64_t)v26; - TASSIGN(v173, v174); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v175 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v176 = (uint64_t)v26; - TASSIGN(v175, v176); - pipe_barrier(PIPE_V); - TEXP(v175, v173); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v177 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v178 = (uint64_t)v26; - TASSIGN(v177, v178); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v179 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v180 = (uint64_t)v24; - TASSIGN(v179, v180); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v181 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v182 = (uint64_t)v28; - TASSIGN(v181, v182); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v183 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v184 = (uint64_t)v28; - TASSIGN(v183, v184); - TMUL(v183, v179, v181); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v185 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v186 = (uint64_t)v26; - TASSIGN(v185, v186); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v187 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v188 = (uint64_t)v25; - TASSIGN(v187, v188); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v189 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v190 = (uint64_t)v25; - TASSIGN(v189, v190); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TMUL(v189, v185, v187); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v191 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v192 = (uint64_t)v28; - TASSIGN(v191, v192); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v193 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v194 = (uint64_t)v25; - TASSIGN(v193, v194); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v195 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v196 = (uint64_t)v28; - TASSIGN(v195, v196); - pipe_barrier(PIPE_V); - TADD(v195, v191, v193); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v197 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v198 = (uint64_t)v28; - TASSIGN(v197, v198); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v199 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v200 = (uint64_t)v30; - TASSIGN(v199, v200); - TROWEXPANDMUL(v199, v147, v165); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v201 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v202 = (uint64_t)v27; - TASSIGN(v201, v202); - TROWEXPANDMUL(v201, v80, v177); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v203 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v204 = (uint64_t)v33; - TASSIGN(v203, v204); - pipe_barrier(PIPE_V); - TADD(v203, v199, v201); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v205 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v206 = (uint64_t)v29; - TASSIGN(v205, v206); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v207 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v208 = (uint64_t)v31; - TASSIGN(v207, v208); - TMOV(v207, v197); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v209 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v210 = (uint64_t)v32; - TASSIGN(v209, v210); - TMOV(v209, v205); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - }; - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - if ((int64_t)((uint64_t)v43 - (uint64_t)v60) == v19) { - int64_t v211 = (int64_t)((uint64_t)v44 + (uint64_t)((int64_t)(uint64_t)v60 * (uint64_t)v21)); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v212 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v213 = (uint64_t)v30; - TASSIGN(v212, v213); - pto::Shape<1, 1, 1, 5, 128> v214 = pto::Shape<1, 1, 1, 5, 128>(); - pto::Stride<640, 640, 640, 128, 1> v215 = pto::Stride<640, 640, 640, 128, 1>(); - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND> v216 = - GlobalTensor, pto::Stride<640, 640, 640, 128, 1>, pto::Layout::ND>( - v3 + (v28 + v211 * v17 + v28 * v19), v214, v215 - ); - TLOAD(v212, v216); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v217 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v218 = (uint64_t)v27; - TASSIGN(v217, v218); - pto::Shape<1, 1, 1, 5, 1> v219 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v220 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v221 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v4 + (v28 + v211 * v19 + v28 * v18), v219, v220 - ); - TLOAD(v217, v221); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v222 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v223 = (uint64_t)v29; - TASSIGN(v222, v223); - pto::Shape<1, 1, 1, 5, 1> v224 = pto::Shape<1, 1, 1, 5, 1>(); - pto::Stride<5, 5, 5, 1, 88064> v225 = pto::Stride<5, 5, 5, 1, 88064>(); - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN> v226 = - GlobalTensor, pto::Stride<5, 5, 5, 1, 88064>, pto::Layout::DN>( - v5 + (v28 + v211 * v19 + v28 * v18), v224, v225 - ); - TLOAD(v222, v226); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID6); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v227 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v228 = (uint64_t)v32; - TASSIGN(v227, v228); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v229 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v230 = (uint64_t)v27; - TASSIGN(v229, v230); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v231 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v232 = (uint64_t)v28; - TASSIGN(v231, v232); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - pipe_barrier(PIPE_V); - TMAX(v231, v227, v229); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v233 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v234 = (uint64_t)v32; - TASSIGN(v233, v234); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v235 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v236 = (uint64_t)v28; - TASSIGN(v235, v236); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v237 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v238 = (uint64_t)v26; - TASSIGN(v237, v238); - pipe_barrier(PIPE_V); - TSUB(v237, v233, v235); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v239 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v240 = (uint64_t)v26; - TASSIGN(v239, v240); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v241 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v242 = (uint64_t)v26; - TASSIGN(v241, v242); - pipe_barrier(PIPE_V); - TEXP(v241, v239); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v243 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v244 = (uint64_t)v26; - TASSIGN(v243, v244); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v245 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v246 = (uint64_t)v27; - TASSIGN(v245, v246); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v247 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v248 = (uint64_t)v28; - TASSIGN(v247, v248); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v249 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v250 = (uint64_t)v27; - TASSIGN(v249, v250); - TSUB(v249, v245, v247); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v251 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v252 = (uint64_t)v27; - TASSIGN(v251, v252); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v253 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v254 = (uint64_t)v25; - TASSIGN(v253, v254); - pipe_barrier(PIPE_V); - TEXP(v253, v251); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v255 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v19); - uint64_t v256 = (uint64_t)v25; - TASSIGN(v255, v256); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v257 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v258 = (uint64_t)v26; - TASSIGN(v257, v258); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v259 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v260 = (uint64_t)v31; - TASSIGN(v259, v260); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v261 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v262 = (uint64_t)v27; - TASSIGN(v261, v262); - pipe_barrier(PIPE_V); - TMUL(v261, v257, v259); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v263 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v264 = (uint64_t)v25; - TASSIGN(v263, v264); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v265 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v266 = (uint64_t)v29; - TASSIGN(v265, v266); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v267 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v268 = (uint64_t)v29; - TASSIGN(v267, v268); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID6); - TMUL(v267, v263, v265); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v269 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v270 = (uint64_t)v27; - TASSIGN(v269, v270); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v271 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v272 = (uint64_t)v29; - TASSIGN(v271, v272); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v273 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v21); - uint64_t v274 = (uint64_t)v29; - TASSIGN(v273, v274); - pipe_barrier(PIPE_V); - TADD(v273, v269, v271); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v275 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v276 = (uint64_t)v27; - TASSIGN(v275, v276); - pipe_barrier(PIPE_V); - TROWEXPANDMUL(v275, v45, v243); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v277 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v17); - uint64_t v278 = (uint64_t)v30; - TASSIGN(v277, v278); - TROWEXPANDMUL(v277, v212, v255); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v279 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v280 = (uint64_t)v30; - TASSIGN(v279, v280); - pipe_barrier(PIPE_V); - TADD(v279, v275, v277); - } else { - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v281 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v282 = (uint64_t)v29; - TASSIGN(v281, v282); - pipe_barrier(PIPE_V); - TMOV(v281, v55); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v283 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v19); - uint64_t v284 = (uint64_t)v28; - TASSIGN(v283, v284); - TMOV(v283, v50); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v285 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v286 = (uint64_t)v30; - TASSIGN(v285, v286); - TMOV(v285, v45); - }; - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v287 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v21, v17); - uint64_t v288 = (uint64_t)v33; - TASSIGN(v287, v288); - pipe_barrier(PIPE_V); - TROWEXPANDDIV(v287, v36, v34); - Tile< - TileType::Vec, float, 1, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v289 = Tile< - TileType::Vec, float, 1, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v10); - uint64_t v290 = (uint64_t)v33; - TASSIGN(v289, v290); - Tile< - TileType::Vec, bfloat16_t, 1, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v291 = Tile< - TileType::Vec, bfloat16_t, 1, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v19, v10); - uint64_t v292 = (uint64_t)v33; - TASSIGN(v291, v292); - pipe_barrier(PIPE_V); - TCVT(v291, v289, v9, v8); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pto::Shape<1, 1, 1, 1, 640> v293 = pto::Shape<1, 1, 1, 1, 640>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v294 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> - v295 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 640>, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( - v1 + (v28 + v40 * v20 + - (int64_t)((uint64_t)((int64_t)(uint64_t)v41 * (uint64_t)v13) * (uint64_t)v17) * v19), - v293, v294 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v295, v291); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - } - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, 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: attn_out_inline157__ssa_v0 - __gm__ Tensor *attn_out_inline157__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *attn_out_inline157__ssa_v0 = - reinterpret_cast<__gm__ bfloat16_t *>(attn_out_inline157__ssa_v0_tensor->buffer.addr) + - attn_out_inline157__ssa_v0_tensor->start_offset; - - // Unpack tensor: seq_lens__ssa_v0 - __gm__ Tensor *seq_lens__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ int32_t *seq_lens__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(seq_lens__ssa_v0_tensor->buffer.addr) + - seq_lens__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_oi_tmp_inline145__rv_v2 - __gm__ Tensor *all_oi_tmp_inline145__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *all_oi_tmp_inline145__rv_v2 = - reinterpret_cast<__gm__ float *>(all_oi_tmp_inline145__rv_v2_tensor->buffer.addr) + - all_oi_tmp_inline145__rv_v2_tensor->start_offset; - - // Unpack tensor: all_cur_mi_inline174__rv_v2 - __gm__ Tensor *all_cur_mi_inline174__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *all_cur_mi_inline174__rv_v2 = - reinterpret_cast<__gm__ float *>(all_cur_mi_inline174__rv_v2_tensor->buffer.addr) + - all_cur_mi_inline174__rv_v2_tensor->start_offset; - - // Unpack tensor: all_cur_li_inline190__rv_v2 - __gm__ Tensor *all_cur_li_inline190__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ float *all_cur_li_inline190__rv_v2 = - reinterpret_cast<__gm__ float *>(all_cur_li_inline190__rv_v2_tensor->buffer.addr) + - all_cur_li_inline190__rv_v2_tensor->start_offset; - - // Forward to ptoas-generated function - online_softmax( - attn_out_inline157__ssa_v0, seq_lens__ssa_v0, all_oi_tmp_inline145__rv_v2, all_cur_mi_inline174__rv_v2, - all_cur_li_inline190__rv_v2, __pypto_spmd_block_idx, __pypto_spmd_block_num - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_consolidate.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_consolidate.cpp deleted file mode 100644 index b12518db91..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_consolidate.cpp +++ /dev/null @@ -1,129 +0,0 @@ -/* - * 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_consolidate -// 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_consolidate(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, int64_t v3) { - const int64_t v4 = 256; - const int64_t v5 = 20; - const int64_t v6 = 1; - const int64_t v7 = 5120; - const int64_t v8 = 16; - const int64_t v9 = 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 v10 = (size_t)v9; v10 < ((size_t)v5); v10 += (size_t)v6) { - int64_t v11 = (int64_t)((uint64_t)((int64_t)v10) * (uint64_t)v4); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v12 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v13 = (uint64_t)v9; - TASSIGN(v12, v13); - pto::Shape<1, 1, 1, 16, 256> v14 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v15 = 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> - v16 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v9 + v3 * v7 + v11 * v6), v14, v15 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v12, v16); - set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - 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>( - v1 + (v9 + v3 * v7 + v11 * v6), v17, v18 - ); - wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); - TSTORE(v19, v12); - 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: next_hidden__iter_v1 - __gm__ Tensor *next_hidden__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *next_hidden__iter_v1 = - reinterpret_cast<__gm__ bfloat16_t *>(next_hidden__iter_v1_tensor->buffer.addr) + - next_hidden__iter_v1_tensor->start_offset; - - // Unpack tensor: out_partial_inline312__rv_v2 - __gm__ Tensor *out_partial_inline312__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *out_partial_inline312__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(out_partial_inline312__rv_v2_tensor->buffer.addr) + - out_partial_inline312__rv_v2_tensor->start_offset; - - // Unpack scalar: oc_b_inline93__idx_v0 - union { - uint64_t u64; - int64_t val; - } oc_b_inline93__idx_v0_conv; - oc_b_inline93__idx_v0_conv.u64 = args[2]; - int64_t oc_b_inline93__idx_v0 = oc_b_inline93__idx_v0_conv.val; - - // Forward to ptoas-generated function - out_consolidate(next_hidden__iter_v1, out_partial_inline312__rv_v2, oc_b_inline93__idx_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_seed.cpp deleted file mode 100644 index 736b11eb12..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_seed.cpp +++ /dev/null @@ -1,134 +0,0 @@ -/* - * 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_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 out_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) { - int64_t v12 = (int64_t)((uint64_t)((int64_t)v11) * (uint64_t)v3); - Tile< - TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v13 = Tile< - TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v3); - uint64_t v14 = (uint64_t)v10; - TASSIGN(v13, v14); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TEXPANDS(v13, v2); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v15 = Tile< - TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v3); - uint64_t v16 = (uint64_t)v9; - TASSIGN(v15, v16); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TEXPANDS(v15, v2); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 512> v17 = pto::Shape<1, 1, 1, 16, 512>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v18 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v19 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v12 * v6), v17, v18 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pipe_barrier(PIPE_MTE3); - TSTORE(v19, v13); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - 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)v12 + (uint64_t)v3) * v6), v20, v21 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v22, v15); - 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: attn_proj_fp32_inline248__ssa_v0 - __gm__ Tensor *attn_proj_fp32_inline248__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *attn_proj_fp32_inline248__ssa_v0 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__ssa_v0_tensor->buffer.addr) + - attn_proj_fp32_inline248__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - out_seed(attn_proj_fp32_inline248__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp index 8586d4ff42..0c236440d2 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/post_rms_reduce.cpp @@ -50,24 +50,23 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void post_rms_reduce(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ float *v3) { - SaturationMode v4 = SaturationMode::OFF; - RoundMode v5 = RoundMode::CAST_ROUND; - const float v6 = 9.99999997E-7f; - const float v7 = 1.95312503E-4f; - const int64_t v8 = 256; - const int64_t v9 = 2; - const int64_t v10 = 20; - const float v11 = 0.0f; - const int64_t v12 = 1; - const int64_t v13 = 5120; - const int64_t v14 = 16; - const int64_t v15 = 49216; - const int64_t v16 = 41024; - const int64_t v17 = 24640; - const int64_t v18 = 16448; - const int64_t v19 = 64; - const int64_t v20 = 0; +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__) @@ -76,281 +75,255 @@ static __aicore__ void post_rms_reduce(__gm__ float *v1, __gm__ bfloat16_t *v2, Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v21 = Tile< + v20 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v22 = (uint64_t)v20; - TASSIGN(v21, v22); + CompactMode::Null>(v10, v12); + uint64_t v21 = (uint64_t)v19; + TASSIGN(v20, v21); set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - TEXPANDS(v21, v11); - for (size_t v23 = (size_t)v20; v23 < ((size_t)v10); v23 += (size_t)v9) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v8); - int64_t v25 = (int64_t)((uint64_t)v24 + (uint64_t)v8); + 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> - v26 = Tile< + v25 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v27 = (uint64_t)v19; - 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>(); + 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> - v30 = GlobalTensor< + v29 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v13 + v24 * v12), v28, v29 + v1 + (v14 + v14 * v11 + v23 * v10), v27, v28 ); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TLOAD(v26, v30); + TLOAD(v25, v29); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v31 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v32 = (uint64_t)v18; - 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< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v35 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v20 + v20 * v13 + v24 * v12), v33, v34 + 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 ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - TLOAD(v31, v35); + 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> - v36 = Tile< + v35 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v37 = (uint64_t)v17; - 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>(); + 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> - v40 = GlobalTensor< + v39 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v20 + v20 * v13 + v25 * v12), v38, v39 + v1 + (v14 + v14 * v11 + v24 * v10), v37, v38 ); - TLOAD(v36, v40); + TLOAD(v35, v39); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v41 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v42 = (uint64_t)v16; - 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< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v45 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v20 + v20 * v13 + v25 * v12), v43, v44 + 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(v41, v45); + 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> - v46 = Tile< + v45 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v47 = (uint64_t)v15; - TASSIGN(v46, v47); + CompactMode::Null>(v12, v6); + uint64_t v46 = (uint64_t)v18; + TASSIGN(v45, v46); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v46, v31, v5, v4); - 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, v8); - uint64_t v49 = (uint64_t)v19; - TASSIGN(v48, v49); - pipe_barrier(PIPE_V); - TADD(v48, v26, v46); + TADD(v45, v25, v30); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v50 = Tile< + v47 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v51 = (uint64_t)v19; - TASSIGN(v50, v51); + CompactMode::Null>(v12, v6); + uint64_t v48 = (uint64_t)v18; + TASSIGN(v47, v48); pipe_barrier(PIPE_V); - TMUL(v50, v48, v48); + TMUL(v47, v45, v45); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v52 = Tile< + v49 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v53 = (uint64_t)v15; - TASSIGN(v52, v53); + 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> - v54 = Tile< + v51 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v55 = (uint64_t)v18; - TASSIGN(v54, v55); + CompactMode::Null>(v12, v10); + uint64_t v52 = (uint64_t)v14; + TASSIGN(v51, v52); pipe_barrier(PIPE_V); - TROWSUM(v54, v50, v52); + TROWSUM(v51, v47, v49); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v56 = Tile< + v53 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v57 = (uint64_t)v18; - TASSIGN(v56, v57); + 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> - v58 = Tile< + v55 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v59 = (uint64_t)v15; - TASSIGN(v58, v59); + CompactMode::Null>(v10, v12); + uint64_t v56 = (uint64_t)v18; + TASSIGN(v55, v56); pipe_barrier(PIPE_V); - TADD(v58, v21, v56); + TADD(v55, v20, v53); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v60 = Tile< + v57 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v61 = (uint64_t)v19; - TASSIGN(v60, v61); + CompactMode::Null>(v12, v6); + uint64_t v58 = (uint64_t)v16; + TASSIGN(v57, v58); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCVT(v60, v41, v5, v4); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v62 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v63 = (uint64_t)v19; - TASSIGN(v62, v63); - pipe_barrier(PIPE_V); - TADD(v62, v36, v60); + TADD(v57, v35, v40); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v64 = Tile< + v59 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v65 = (uint64_t)v19; - TASSIGN(v64, v65); + CompactMode::Null>(v12, v6); + uint64_t v60 = (uint64_t)v16; + TASSIGN(v59, v60); pipe_barrier(PIPE_V); - TMUL(v64, v62, v62); + TMUL(v59, v57, v57); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v66 = Tile< + v61 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v67 = (uint64_t)v17; - TASSIGN(v66, v67); + 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> - v68 = Tile< + v63 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v69 = (uint64_t)v18; - TASSIGN(v68, v69); + CompactMode::Null>(v12, v10); + uint64_t v64 = (uint64_t)v13; + TASSIGN(v63, v64); pipe_barrier(PIPE_V); - TROWSUM(v68, v64, v66); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TROWSUM(v63, v59, v61); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v70 = Tile< + v65 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v71 = (uint64_t)v18; - TASSIGN(v70, v71); + 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> - v72 = Tile< + v67 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v73 = (uint64_t)v20; - TASSIGN(v72, v73); + CompactMode::Null>(v10, v12); + uint64_t v68 = (uint64_t)v19; + TASSIGN(v67, v68); pipe_barrier(PIPE_V); - TADD(v72, v58, v70); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + 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> - v74 = Tile< + v69 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v75 = (uint64_t)v19; - TASSIGN(v74, v75); + CompactMode::Null>(v10, v12); + uint64_t v70 = (uint64_t)v18; + TASSIGN(v69, v70); pipe_barrier(PIPE_V); - TMULS(v74, v21, v7); + TMULS(v69, v20, v5); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v76 = Tile< + v71 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v77 = (uint64_t)v19; - TASSIGN(v76, v77); + CompactMode::Null>(v10, v12); + uint64_t v72 = (uint64_t)v18; + TASSIGN(v71, v72); pipe_barrier(PIPE_V); - TADDS(v76, v74, v6); + TADDS(v71, v69, v4); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v78 = Tile< + v73 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v79 = (uint64_t)v19; - TASSIGN(v78, v79); + CompactMode::Null>(v10, v12); + uint64_t v74 = (uint64_t)v18; + TASSIGN(v73, v74); pipe_barrier(PIPE_V); - TSQRT(v78, v76); + TSQRT(v73, v71); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v80 = Tile< + v75 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v81 = (uint64_t)v17; - TASSIGN(v80, v81); + CompactMode::Null>(v10, v12); + uint64_t v76 = (uint64_t)v17; + TASSIGN(v75, v76); pipe_barrier(PIPE_V); - TRECIP(v80, v78); + 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> - v82 = Tile< + v77 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v83 = (uint64_t)v17; - TASSIGN(v82, v83); - pto::Shape<1, 1, 1, 16, 1> v84 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 16> v85 = pto::Stride<16, 16, 16, 1, 16>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v86 = + 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 + (v20 + v20 * v12 + v20 * v14), v84, v85 + v3 + (v14 + v14 * v10 + v14 * v12), v79, v80 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v86, v82); + TSTORE(v81, v77); wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); #endif // __DAV_VEC__ ptoas_auto_sync_tail(PTOAutoSyncTailMode::kBarrierAll); @@ -358,23 +331,23 @@ static __aicore__ void post_rms_reduce(__gm__ float *v1, __gm__ bfloat16_t *v2, } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: attn_proj_fp32_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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[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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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_inline255__ssa_v0 - __gm__ Tensor *inv_rms_tile_inline255__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *inv_rms_tile_inline255__ssa_v0 = - reinterpret_cast<__gm__ float *>(inv_rms_tile_inline255__ssa_v0_tensor->buffer.addr) + - inv_rms_tile_inline255__ssa_v0_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_inline248__rv_v5, cur__iter_v6, inv_rms_tile_inline255__ssa_v0); + post_rms_reduce(attn_proj_fp32_inline220__ssa_v7, cur__iter_v6, inv_rms_tile_inline126__ssa_v0); } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/q_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/q_seed.cpp index 95eb3db6e9..f0495a0ddc 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/q_seed.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/q_seed.cpp @@ -79,39 +79,39 @@ static __aicore__ void q_seed(__gm__ float *v1) { 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> - v14 = Tile< + v18 = Tile< TileType::Vec, float, 16, 512, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v8, v3); - uint64_t v15 = (uint64_t)v9; - TASSIGN(v14, v15); + uint64_t v19 = (uint64_t)v9; + TASSIGN(v18, v19); wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TEXPANDS(v14, v2); + TEXPANDS(v18, v2); set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - int64_t v16 = (int64_t)((uint64_t)((int64_t)v11) * (uint64_t)v3); - pto::Shape<1, 1, 1, 16, 512> v17 = pto::Shape<1, 1, 1, 16, 512>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v18 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v19 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 512>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v16 * v6), v17, v18 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pipe_barrier(PIPE_MTE3); - TSTORE(v19, v12); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); 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)v16 + (uint64_t)v3) * v6), v20, v21 + v1 + (v10 + v10 * v7 + (int64_t)((uint64_t)v14 + (uint64_t)v3) * v6), v20, v21 ); - pipe_barrier(PIPE_MTE3); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v22, v14); + pipe_barrier(PIPE_MTE3); + TSTORE(v22, v18); set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); } wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); @@ -123,12 +123,12 @@ static __aicore__ void q_seed(__gm__ float *v1) { } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: q_proj_inline129__ssa_v0 - __gm__ Tensor *q_proj_inline129__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *q_proj_inline129__ssa_v0 = - reinterpret_cast<__gm__ float *>(q_proj_inline129__ssa_v0_tensor->buffer.addr) + - q_proj_inline129__ssa_v0_tensor->start_offset; + // 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_inline129__ssa_v0); + q_seed(q_proj_inline139__ssa_v0); } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_gamma.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_gamma.cpp deleted file mode 100644 index 1272e2f9aa..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_gamma.cpp +++ /dev/null @@ -1,294 +0,0 @@ -/* - * 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: qk_gamma -// 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 qk_gamma( - __gm__ float *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5, __gm__ float *v6, - __gm__ float *v7 -) { - const int64_t v8 = 80; - const int64_t v9 = 640; - const int64_t v10 = 8; - const int64_t v11 = 128; - const int64_t v12 = 5120; - const int64_t v13 = 1; - const int64_t v14 = 1024; - const int64_t v15 = 16; - const int64_t v16 = 0; - const int64_t v17 = 41536; - const int64_t v18 = 41024; - const int64_t v19 = 40960; - 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> - v20 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v13); - uint64_t v21 = (uint64_t)v19; - TASSIGN(v20, v21); - pto::Shape<1, 1, 1, 16, 1> v22 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 1> v23 = pto::Stride<16, 16, 16, 1, 1>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 1>, pto::Layout::DN> v24 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 1>, pto::Layout::DN>( - v4 + (v16 + v16 * v13 + v16 * v13), v22, v23 - ); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v20, v24); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v25 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v11); - uint64_t v26 = (uint64_t)v18; - TASSIGN(v25, v26); - pto::Shape<1, 1, 1, 1, 128> v27 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<128, 128, 128, 128, 1> v28 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v29 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v5 + (v16 + v16 * v11 + v16 * v13), v27, v28 - ); - TLOAD(v25, v29); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v30 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v11); - uint64_t v31 = (uint64_t)v17; - TASSIGN(v30, v31); - pto::Shape<1, 1, 1, 1, 128> v32 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<128, 128, 128, 128, 1> v33 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v34 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v7 + (v16 + v16 * v11 + v16 * v13), v32, v33 - ); - TLOAD(v30, v34); - for (size_t v35 = (size_t)v16; v35 < ((size_t)v10); v35 += (size_t)v13) { - int64_t v36 = (int64_t)v35; - int64_t v37 = (int64_t)((uint64_t)v36 * (uint64_t)v9); - Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v38 = Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v39 = (uint64_t)v16; - TASSIGN(v38, v39); - pto::Shape<1, 1, 1, 16, 640> v40 = pto::Shape<1, 1, 1, 16, 640>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v41 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v42 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 640>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v16 + v16 * v12 + v37 * v13), v40, v41 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v38, v42); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v43 = Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v44 = (uint64_t)v16; - TASSIGN(v43, v44); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TROWEXPANDMUL(v43, v38, v20); - Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v45 = Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v11); - uint64_t v46 = (uint64_t)v16; - TASSIGN(v45, v46); - Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v47 = Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v11); - uint64_t v48 = (uint64_t)v16; - TASSIGN(v47, v48); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v47, v45, v25); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v49 = Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v50 = (uint64_t)v16; - TASSIGN(v49, v50); - pto::Shape<1, 1, 1, 16, 640> v51 = pto::Shape<1, 1, 1, 16, 640>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v52 = pto::Stride<81920, 81920, 81920, 5120, 1>(); - GlobalTensor, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v53 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 640>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v16 + v16 * v12 + v37 * v13), v51, v52 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v53, v49); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - int64_t v54 = (int64_t)((uint64_t)v36 * (uint64_t)v11); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v55 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v11); - uint64_t v56 = (uint64_t)v16; - TASSIGN(v55, v56); - pto::Shape<1, 1, 1, 16, 128> v57 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v58 = pto::Stride<16384, 16384, 16384, 1024, 1>(); - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> - v59 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v6 + (v16 + v16 * v14 + v54 * v13), v57, v58 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - TLOAD(v55, v59); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v60 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v11); - uint64_t v61 = (uint64_t)v16; - TASSIGN(v60, v61); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TROWEXPANDMUL(v60, v55, v20); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v62 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v11); - uint64_t v63 = (uint64_t)v16; - TASSIGN(v62, v63); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v62, v60, v30); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 128> v64 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v65 = pto::Stride<16384, 16384, 16384, 1024, 1>(); - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> - v66 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v1 + (v16 + v16 * v14 + v54 * v13), v64, v65 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v66, v62); - 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: k_proj_norm_inline155__ssa_v0 - __gm__ Tensor *k_proj_norm_inline155__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *k_proj_norm_inline155__ssa_v0 = - reinterpret_cast<__gm__ float *>(k_proj_norm_inline155__ssa_v0_tensor->buffer.addr) + - k_proj_norm_inline155__ssa_v0_tensor->start_offset; - - // Unpack tensor: q_proj_norm_inline197__ssa_v0 - __gm__ Tensor *q_proj_norm_inline197__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *q_proj_norm_inline197__ssa_v0 = - reinterpret_cast<__gm__ float *>(q_proj_norm_inline197__ssa_v0_tensor->buffer.addr) + - q_proj_norm_inline197__ssa_v0_tensor->start_offset; - - // Unpack tensor: q_proj_inline129__rv_v5 - __gm__ Tensor *q_proj_inline129__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *q_proj_inline129__rv_v5 = - reinterpret_cast<__gm__ float *>(q_proj_inline129__rv_v5_tensor->buffer.addr) + - q_proj_inline129__rv_v5_tensor->start_offset; - - // Unpack tensor: inv_rms_col_inline89__ssa_v0 - __gm__ Tensor *inv_rms_col_inline89__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *inv_rms_col_inline89__ssa_v0 = - reinterpret_cast<__gm__ float *>(inv_rms_col_inline89__ssa_v0_tensor->buffer.addr) + - inv_rms_col_inline89__ssa_v0_tensor->start_offset; - - // Unpack tensor: q_norm_w_inline242__ssa_v0 - __gm__ Tensor *q_norm_w_inline242__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ float *q_norm_w_inline242__ssa_v0 = - reinterpret_cast<__gm__ float *>(q_norm_w_inline242__ssa_v0_tensor->buffer.addr) + - q_norm_w_inline242__ssa_v0_tensor->start_offset; - - // Unpack tensor: k_proj_inline128__rv_v3 - __gm__ Tensor *k_proj_inline128__rv_v3_tensor = reinterpret_cast<__gm__ Tensor *>(args[5]); - __gm__ float *k_proj_inline128__rv_v3 = - reinterpret_cast<__gm__ float *>(k_proj_inline128__rv_v3_tensor->buffer.addr) + - k_proj_inline128__rv_v3_tensor->start_offset; - - // Unpack tensor: k_norm_w_inline138__ssa_v0 - __gm__ Tensor *k_norm_w_inline138__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[6]); - __gm__ float *k_norm_w_inline138__ssa_v0 = - reinterpret_cast<__gm__ float *>(k_norm_w_inline138__ssa_v0_tensor->buffer.addr) + - k_norm_w_inline138__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - qk_gamma( - k_proj_norm_inline155__ssa_v0, q_proj_norm_inline197__ssa_v0, q_proj_inline129__rv_v5, - inv_rms_col_inline89__ssa_v0, q_norm_w_inline242__ssa_v0, k_proj_inline128__rv_v3, k_norm_w_inline138__ssa_v0 - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_recip.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_recip.cpp deleted file mode 100644 index 807571479a..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_recip.cpp +++ /dev/null @@ -1,442 +0,0 @@ -/* - * 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: qk_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 -qk_recip(__gm__ float *v1, __gm__ float *v2, __gm__ float *v3, __gm__ float *v4, __gm__ float *v5) { - const float v6 = 9.99999997E-7f; - const float v7 = 0.0078125f; - const int64_t v8 = 80; - const int64_t v9 = 8; - const int64_t v10 = 1024; - const int64_t v11 = 5120; - const int64_t v12 = 16; - const int64_t v13 = 640; - const int64_t v14 = 1; - const int64_t v15 = 128; - const int64_t v16 = 0; - const int64_t v17 = 41344; - const int64_t v18 = 384; - const int64_t v19 = 320; - 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> - v20 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v21 = (uint64_t)v19; - TASSIGN(v20, v21); - pto::Shape<1, 1, 1, 16, 1> v22 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 1> v23 = pto::Stride<16, 16, 16, 1, 1>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 1>, pto::Layout::DN> v24 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 1>, pto::Layout::DN>( - v4 + (v16 + v16 * v14 + v16 * v14), v22, v23 - ); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TLOAD(v20, v24); - for (size_t v25 = (size_t)v16; v25 < ((size_t)v9); v25 += (size_t)v14) { - int64_t v26 = (int64_t)v25; - Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v27 = Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v13); - uint64_t v28 = (uint64_t)v18; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 640> v29 = pto::Shape<1, 1, 1, 16, 640>(); - 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< - float, pto::Shape<1, 1, 1, 16, 640>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v16 + v16 * v11 + (int64_t)((uint64_t)v26 * (uint64_t)v13) * v14), v29, v30 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v32 = Tile< - TileType::Vec, float, 16, 640, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v13); - uint64_t v33 = (uint64_t)v18; - TASSIGN(v32, v33); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TROWEXPANDMUL(v32, v27, v20); - Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v34 = Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v15); - uint64_t v35 = (uint64_t)v18; - TASSIGN(v34, v35); - Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v36 = Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v15); - uint64_t v37 = (uint64_t)v18; - TASSIGN(v36, v37); - pipe_barrier(PIPE_V); - TMUL(v36, v34, v34); - Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v38 = Tile< - TileType::Vec, float, 80, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v15); - uint64_t v39 = (uint64_t)v17; - TASSIGN(v38, v39); - Tile< - TileType::Vec, float, 80, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v40 = Tile< - TileType::Vec, float, 80, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v14); - uint64_t v41 = (uint64_t)v16; - TASSIGN(v40, v41); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TROWSUM(v40, v36, v38); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v42 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v43 = (uint64_t)v16; - TASSIGN(v42, v43); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v44 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v45 = (uint64_t)v18; - TASSIGN(v44, v45); - pipe_barrier(PIPE_V); - TMULS(v44, v42, v7); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v46 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v47 = (uint64_t)v18; - TASSIGN(v46, v47); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v48 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v49 = (uint64_t)v18; - TASSIGN(v48, v49); - pipe_barrier(PIPE_V); - TADDS(v48, v46, v6); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v50 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v51 = (uint64_t)v18; - TASSIGN(v50, v51); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v52 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v53 = (uint64_t)v18; - TASSIGN(v52, v53); - pipe_barrier(PIPE_V); - TSQRT(v52, v50); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v54 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v55 = (uint64_t)v18; - TASSIGN(v54, v55); - Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v56 = Tile< - TileType::Vec, float, 1, 80, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v8); - uint64_t v57 = (uint64_t)v17; - TASSIGN(v56, v57); - pipe_barrier(PIPE_V); - TRECIP(v56, v54); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 80, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v58 = Tile< - TileType::Vec, float, 80, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v14); - uint64_t v59 = (uint64_t)v17; - TASSIGN(v58, v59); - pto::Shape<1, 1, 1, 80, 1> v60 = pto::Shape<1, 1, 1, 80, 1>(); - pto::Stride<80, 80, 80, 1, 640> v61 = pto::Stride<80, 80, 80, 1, 640>(); - GlobalTensor, pto::Stride<80, 80, 80, 1, 640>, pto::Layout::DN> v62 = - GlobalTensor, pto::Stride<80, 80, 80, 1, 640>, pto::Layout::DN>( - v2 + (v16 + (int64_t)((uint64_t)v26 * (uint64_t)v8) * v14 + v16 * v13), v60, v61 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v62, v58); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v63 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v15); - uint64_t v64 = (uint64_t)v18; - TASSIGN(v63, v64); - pto::Shape<1, 1, 1, 16, 128> v65 = pto::Shape<1, 1, 1, 16, 128>(); - pto::Stride<16384, 16384, 16384, 1024, 1> v66 = pto::Stride<16384, 16384, 16384, 1024, 1>(); - GlobalTensor, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND> - v67 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 128>, pto::Stride<16384, 16384, 16384, 1024, 1>, pto::Layout::ND>( - v5 + (v16 + v16 * v10 + (int64_t)((uint64_t)v26 * (uint64_t)v15) * v14), v65, v66 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - TLOAD(v63, v67); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v68 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v15); - uint64_t v69 = (uint64_t)v18; - TASSIGN(v68, v69); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TROWEXPANDMUL(v68, v63, v20); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v70 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v15); - uint64_t v71 = (uint64_t)v18; - TASSIGN(v70, v71); - pipe_barrier(PIPE_V); - TMUL(v70, v68, v68); - Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v72 = Tile< - TileType::Vec, float, 16, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v15); - uint64_t v73 = (uint64_t)v17; - TASSIGN(v72, v73); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v74 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v75 = (uint64_t)v16; - TASSIGN(v74, v75); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - pipe_barrier(PIPE_V); - TROWSUM(v74, v70, v72); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v76 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v77 = (uint64_t)v16; - TASSIGN(v76, v77); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v78 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v79 = (uint64_t)v18; - TASSIGN(v78, v79); - pipe_barrier(PIPE_V); - TMULS(v78, v76, v7); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v80 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v81 = (uint64_t)v18; - TASSIGN(v80, v81); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v82 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v83 = (uint64_t)v18; - TASSIGN(v82, v83); - pipe_barrier(PIPE_V); - TADDS(v82, v80, v6); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v84 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v85 = (uint64_t)v18; - TASSIGN(v84, v85); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v86 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v87 = (uint64_t)v18; - TASSIGN(v86, v87); - pipe_barrier(PIPE_V); - TSQRT(v86, v84); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v88 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v89 = (uint64_t)v18; - TASSIGN(v88, v89); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v90 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v12); - uint64_t v91 = (uint64_t)v17; - TASSIGN(v90, v91); - pipe_barrier(PIPE_V); - TRECIP(v90, v88); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v92 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v12, v14); - uint64_t v93 = (uint64_t)v17; - TASSIGN(v92, v93); - pto::Shape<1, 1, 1, 16, 1> v94 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 128> v95 = pto::Stride<16, 16, 16, 1, 128>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 128>, pto::Layout::DN> v96 = - GlobalTensor, pto::Stride<16, 16, 16, 1, 128>, pto::Layout::DN>( - v1 + (v16 + (int64_t)((uint64_t)v26 * (uint64_t)v12) * v14 + v16 * v15), v94, v95 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v96, v92); - 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: k_inv_states_inline90__ssa_v0 - __gm__ Tensor *k_inv_states_inline90__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *k_inv_states_inline90__ssa_v0 = - reinterpret_cast<__gm__ float *>(k_inv_states_inline90__ssa_v0_tensor->buffer.addr) + - k_inv_states_inline90__ssa_v0_tensor->start_offset; - - // Unpack tensor: q_inv_states_inline99__ssa_v0 - __gm__ Tensor *q_inv_states_inline99__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *q_inv_states_inline99__ssa_v0 = - reinterpret_cast<__gm__ float *>(q_inv_states_inline99__ssa_v0_tensor->buffer.addr) + - q_inv_states_inline99__ssa_v0_tensor->start_offset; - - // Unpack tensor: q_proj_inline129__rv_v5 - __gm__ Tensor *q_proj_inline129__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *q_proj_inline129__rv_v5 = - reinterpret_cast<__gm__ float *>(q_proj_inline129__rv_v5_tensor->buffer.addr) + - q_proj_inline129__rv_v5_tensor->start_offset; - - // Unpack tensor: inv_rms_col_inline89__ssa_v0 - __gm__ Tensor *inv_rms_col_inline89__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *inv_rms_col_inline89__ssa_v0 = - reinterpret_cast<__gm__ float *>(inv_rms_col_inline89__ssa_v0_tensor->buffer.addr) + - inv_rms_col_inline89__ssa_v0_tensor->start_offset; - - // Unpack tensor: k_proj_inline128__rv_v3 - __gm__ Tensor *k_proj_inline128__rv_v3_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ float *k_proj_inline128__rv_v3 = - reinterpret_cast<__gm__ float *>(k_proj_inline128__rv_v3_tensor->buffer.addr) + - k_proj_inline128__rv_v3_tensor->start_offset; - - // Forward to ptoas-generated function - qk_recip( - k_inv_states_inline90__ssa_v0, q_inv_states_inline99__ssa_v0, q_proj_inline129__rv_v5, - inv_rms_col_inline89__ssa_v0, k_proj_inline128__rv_v3 - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp index 0ea90da0b8..5d44c597bd 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast.cpp @@ -51,286 +51,261 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } static __aicore__ void residual_rms_cast( - __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ bfloat16_t *v4, __gm__ float *v5, int64_t v6, + __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 = 4; - const int64_t v12 = 2; + 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 = 1024; - const int64_t v17 = 0; - const int64_t v18 = 59392; - const int64_t v19 = 43008; - const int64_t v20 = 41984; - const int64_t v21 = 33792; - const int64_t v22 = 17408; + 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); - for (size_t v23 = (size_t)v17; v23 < ((size_t)v11); v23 += (size_t)v12) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v10); - int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)v24); - int64_t v26 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v24 + (uint64_t)v10)); + 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> - v27 = Tile< + v26 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v28 = (uint64_t)v22; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v31 = GlobalTensor< + v30 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v25 * v13), v29, v30 + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v33 = (uint64_t)v21; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v25 * v13), v34, v35 + 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 ); - TLOAD(v32, v36); + 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> - v37 = Tile< + v36 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v38 = (uint64_t)v20; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 1, 256> v39 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v40 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v41 = + 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 + (v17 + v7 * v14 + v25 * v13), v39, v40 + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 ); - TLOAD(v37, v41); + 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> - v42 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v43 = (uint64_t)v19; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v46 = GlobalTensor< + v45 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v26 * v13), v44, v45 + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 ); - TLOAD(v42, v46); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v47 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v48 = (uint64_t)v18; - TASSIGN(v47, v48); - pto::Shape<1, 1, 1, 16, 256> v49 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v50 = 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> - v51 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v26 * v13), v49, v50 + 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 ); - TLOAD(v47, v51); + 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> - v52 = Tile< + v51 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v53 = (uint64_t)v17; - TASSIGN(v52, v53); - pto::Shape<1, 1, 1, 1, 256> v54 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v55 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v56 = + 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 + (v17 + v7 * v14 + v26 * v13), v54, v55 + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 ); - TLOAD(v52, v56); + 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> - v57 = Tile< + v56 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v58 = (uint64_t)v16; - TASSIGN(v57, v58); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v57, v32, v9, v8); - 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>(v15, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - TADD(v59, v27, v57); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v62 = (uint64_t)v16; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v9, v8); + 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> - v63 = Tile< + v58 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v64 = (uint64_t)v22; - TASSIGN(v63, v64); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v63, v59, v37); + 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> - v65 = Tile< + v60 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v66 = (uint64_t)v21; - TASSIGN(v65, v66); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); pipe_barrier(PIPE_V); - TCVT(v65, v63, v9, v8); + 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> - v67 = Tile< + v68 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pipe_barrier(PIPE_V); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v67, v47, v9, v8); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v70 = (uint64_t)v22; - TASSIGN(v69, v70); - pipe_barrier(PIPE_V); - TADD(v69, v42, v67); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); - pipe_barrier(PIPE_V); - TCVT(v71, v69, v9, v8); + 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> - v73 = Tile< + v70 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v74 = (uint64_t)v22; - TASSIGN(v73, v74); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v73, v69, v52); + 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> - v75 = Tile< + v72 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v76 = (uint64_t)v22; - TASSIGN(v75, v76); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); pipe_barrier(PIPE_V); - TCVT(v75, v73, v9, v8); + 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>( - v2 + (v17 + v17 * v14 + v25 * v13), v77, v78 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v79, v61); - pto::Shape<1, 1, 1, 16, 256> v80 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v81 = 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> - v82 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v25 * v13), v80, v81 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v82, v65); - pto::Shape<1, 1, 1, 16, 256> v83 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v17 + v17 * v14 + v26 * v13), v83, v84 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v85, v71); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v26 * v13), v86, v87 + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v88, v75); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -338,28 +313,28 @@ static __aicore__ void residual_rms_cast( } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__ssa_v0 - __gm__ Tensor *mlp_norm_in_inline250__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__ssa_v0 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__ssa_v0_tensor->buffer.addr) + - mlp_norm_in_inline250__ssa_v0_tensor->start_offset; + // 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_inline130__ssa_v0 - __gm__ Tensor *post_norm_partial_inline130__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *post_norm_partial_inline130__ssa_v0 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__ssa_v0_tensor->buffer.addr) + - post_norm_partial_inline130__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_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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]); @@ -367,13 +342,13 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + post_rms_weight__ssa_v0_tensor->start_offset; - // Unpack scalar: k_base_inline177__ssa_v0 + // Unpack scalar: k_base_inline111__ssa_v0 union { uint64_t u64; int64_t val; - } k_base_inline177__ssa_v0_conv; - k_base_inline177__ssa_v0_conv.u64 = args[5]; - int64_t k_base_inline177__ssa_v0 = k_base_inline177__ssa_v0_conv.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 { @@ -385,7 +360,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Forward to ptoas-generated function residual_rms_cast( - mlp_norm_in_inline250__ssa_v0, post_norm_partial_inline130__ssa_v0, attn_proj_fp32_inline248__rv_v5, - cur__iter_v6, post_rms_weight__ssa_v0, k_base_inline177__ssa_v0, i__idx_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp index 2dda71a02f..d3dc242f7c 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_0.cpp @@ -51,286 +51,261 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } static __aicore__ void residual_rms_cast_0( - __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ bfloat16_t *v4, __gm__ float *v5, int64_t v6, + __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 = 4; - const int64_t v12 = 2; + 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 = 1024; - const int64_t v17 = 0; - const int64_t v18 = 59392; - const int64_t v19 = 43008; - const int64_t v20 = 41984; - const int64_t v21 = 33792; - const int64_t v22 = 17408; + 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); - for (size_t v23 = (size_t)v17; v23 < ((size_t)v11); v23 += (size_t)v12) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v10); - int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)v24); - int64_t v26 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v24 + (uint64_t)v10)); + 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> - v27 = Tile< + v26 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v28 = (uint64_t)v22; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v31 = GlobalTensor< + v30 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v25 * v13), v29, v30 + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v33 = (uint64_t)v21; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v25 * v13), v34, v35 + 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 ); - TLOAD(v32, v36); + 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> - v37 = Tile< + v36 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v38 = (uint64_t)v20; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 1, 256> v39 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v40 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v41 = + 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 + (v17 + v7 * v14 + v25 * v13), v39, v40 + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 ); - TLOAD(v37, v41); + 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> - v42 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v43 = (uint64_t)v19; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v46 = GlobalTensor< + v45 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v26 * v13), v44, v45 + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 ); - TLOAD(v42, v46); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v47 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v48 = (uint64_t)v18; - TASSIGN(v47, v48); - pto::Shape<1, 1, 1, 16, 256> v49 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v50 = 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> - v51 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v26 * v13), v49, v50 + 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 ); - TLOAD(v47, v51); + 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> - v52 = Tile< + v51 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v53 = (uint64_t)v17; - TASSIGN(v52, v53); - pto::Shape<1, 1, 1, 1, 256> v54 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v55 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v56 = + 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 + (v17 + v7 * v14 + v26 * v13), v54, v55 + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 ); - TLOAD(v52, v56); + 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> - v57 = Tile< + v56 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v58 = (uint64_t)v16; - TASSIGN(v57, v58); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v57, v32, v9, v8); - 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>(v15, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - TADD(v59, v27, v57); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v62 = (uint64_t)v16; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v9, v8); + 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> - v63 = Tile< + v58 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v64 = (uint64_t)v22; - TASSIGN(v63, v64); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v63, v59, v37); + 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> - v65 = Tile< + v60 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v66 = (uint64_t)v21; - TASSIGN(v65, v66); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); pipe_barrier(PIPE_V); - TCVT(v65, v63, v9, v8); + 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> - v67 = Tile< + v68 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pipe_barrier(PIPE_V); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v67, v47, v9, v8); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v70 = (uint64_t)v22; - TASSIGN(v69, v70); - pipe_barrier(PIPE_V); - TADD(v69, v42, v67); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); - pipe_barrier(PIPE_V); - TCVT(v71, v69, v9, v8); + 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> - v73 = Tile< + v70 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v74 = (uint64_t)v22; - TASSIGN(v73, v74); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v73, v69, v52); + 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> - v75 = Tile< + v72 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v76 = (uint64_t)v22; - TASSIGN(v75, v76); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); pipe_barrier(PIPE_V); - TCVT(v75, v73, v9, v8); + 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>( - v2 + (v17 + v17 * v14 + v25 * v13), v77, v78 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v79, v61); - pto::Shape<1, 1, 1, 16, 256> v80 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v81 = 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> - v82 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v25 * v13), v80, v81 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v82, v65); - pto::Shape<1, 1, 1, 16, 256> v83 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v17 + v17 * v14 + v26 * v13), v83, v84 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v85, v71); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v26 * v13), v86, v87 + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v88, v75); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -338,28 +313,28 @@ static __aicore__ void residual_rms_cast_0( } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v2 - __gm__ Tensor *mlp_norm_in_inline250__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v2_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v2_tensor->start_offset; + // 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_inline130__rv_v2 - __gm__ Tensor *post_norm_partial_inline130__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *post_norm_partial_inline130__rv_v2 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__rv_v2_tensor->buffer.addr) + - post_norm_partial_inline130__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_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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]); @@ -367,13 +342,13 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + post_rms_weight__ssa_v0_tensor->start_offset; - // Unpack scalar: k_base_inline177__ssa_v1 + // Unpack scalar: k_base_inline111__ssa_v1 union { uint64_t u64; int64_t val; - } k_base_inline177__ssa_v1_conv; - k_base_inline177__ssa_v1_conv.u64 = args[5]; - int64_t k_base_inline177__ssa_v1 = k_base_inline177__ssa_v1_conv.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 { @@ -385,7 +360,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Forward to ptoas-generated function residual_rms_cast_0( - mlp_norm_in_inline250__rv_v2, post_norm_partial_inline130__rv_v2, attn_proj_fp32_inline248__rv_v5, cur__iter_v6, - post_rms_weight__ssa_v0, k_base_inline177__ssa_v1, i__idx_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp index 5f5fe87c7e..0a54a73863 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_1.cpp @@ -51,286 +51,261 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } static __aicore__ void residual_rms_cast_1( - __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ bfloat16_t *v4, __gm__ float *v5, int64_t v6, + __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 = 4; - const int64_t v12 = 2; + 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 = 1024; - const int64_t v17 = 0; - const int64_t v18 = 59392; - const int64_t v19 = 43008; - const int64_t v20 = 41984; - const int64_t v21 = 33792; - const int64_t v22 = 17408; + 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); - for (size_t v23 = (size_t)v17; v23 < ((size_t)v11); v23 += (size_t)v12) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v10); - int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)v24); - int64_t v26 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v24 + (uint64_t)v10)); + 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> - v27 = Tile< + v26 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v28 = (uint64_t)v22; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v31 = GlobalTensor< + v30 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v25 * v13), v29, v30 + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v33 = (uint64_t)v21; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v25 * v13), v34, v35 + 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 ); - TLOAD(v32, v36); + 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> - v37 = Tile< + v36 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v38 = (uint64_t)v20; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 1, 256> v39 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v40 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v41 = + 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 + (v17 + v7 * v14 + v25 * v13), v39, v40 + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 ); - TLOAD(v37, v41); + 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> - v42 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v43 = (uint64_t)v19; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v46 = GlobalTensor< + v45 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v26 * v13), v44, v45 + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 ); - TLOAD(v42, v46); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v47 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v48 = (uint64_t)v18; - TASSIGN(v47, v48); - pto::Shape<1, 1, 1, 16, 256> v49 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v50 = 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> - v51 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v26 * v13), v49, v50 + 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 ); - TLOAD(v47, v51); + 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> - v52 = Tile< + v51 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v53 = (uint64_t)v17; - TASSIGN(v52, v53); - pto::Shape<1, 1, 1, 1, 256> v54 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v55 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v56 = + 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 + (v17 + v7 * v14 + v26 * v13), v54, v55 + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 ); - TLOAD(v52, v56); + 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> - v57 = Tile< + v56 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v58 = (uint64_t)v16; - TASSIGN(v57, v58); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v57, v32, v9, v8); - 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>(v15, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - TADD(v59, v27, v57); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v62 = (uint64_t)v16; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v9, v8); + 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> - v63 = Tile< + v58 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v64 = (uint64_t)v22; - TASSIGN(v63, v64); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v63, v59, v37); + 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> - v65 = Tile< + v60 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v66 = (uint64_t)v21; - TASSIGN(v65, v66); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); pipe_barrier(PIPE_V); - TCVT(v65, v63, v9, v8); + 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> - v67 = Tile< + v68 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pipe_barrier(PIPE_V); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v67, v47, v9, v8); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v70 = (uint64_t)v22; - TASSIGN(v69, v70); - pipe_barrier(PIPE_V); - TADD(v69, v42, v67); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); - pipe_barrier(PIPE_V); - TCVT(v71, v69, v9, v8); + 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> - v73 = Tile< + v70 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v74 = (uint64_t)v22; - TASSIGN(v73, v74); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v73, v69, v52); + 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> - v75 = Tile< + v72 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v76 = (uint64_t)v22; - TASSIGN(v75, v76); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); pipe_barrier(PIPE_V); - TCVT(v75, v73, v9, v8); + 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>( - v2 + (v17 + v17 * v14 + v25 * v13), v77, v78 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v79, v61); - pto::Shape<1, 1, 1, 16, 256> v80 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v81 = 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> - v82 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v25 * v13), v80, v81 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v82, v65); - pto::Shape<1, 1, 1, 16, 256> v83 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v17 + v17 * v14 + v26 * v13), v83, v84 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v85, v71); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v26 * v13), v86, v87 + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v88, v75); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -338,28 +313,28 @@ static __aicore__ void residual_rms_cast_1( } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v5 - __gm__ Tensor *mlp_norm_in_inline250__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v5 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v5_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v5_tensor->start_offset; + // 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_inline130__rv_v5 - __gm__ Tensor *post_norm_partial_inline130__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *post_norm_partial_inline130__rv_v5 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__rv_v5_tensor->buffer.addr) + - post_norm_partial_inline130__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_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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]); @@ -367,13 +342,13 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + post_rms_weight__ssa_v0_tensor->start_offset; - // Unpack scalar: k_base_inline177__ssa_v2 + // Unpack scalar: k_base_inline111__ssa_v2 union { uint64_t u64; int64_t val; - } k_base_inline177__ssa_v2_conv; - k_base_inline177__ssa_v2_conv.u64 = args[5]; - int64_t k_base_inline177__ssa_v2 = k_base_inline177__ssa_v2_conv.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 { @@ -385,7 +360,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Forward to ptoas-generated function residual_rms_cast_1( - mlp_norm_in_inline250__rv_v5, post_norm_partial_inline130__rv_v5, attn_proj_fp32_inline248__rv_v5, cur__iter_v6, - post_rms_weight__ssa_v0, k_base_inline177__ssa_v2, i__idx_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp index cf16156875..3b85e537b8 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_2.cpp @@ -51,286 +51,261 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } static __aicore__ void residual_rms_cast_2( - __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ bfloat16_t *v4, __gm__ float *v5, int64_t v6, + __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 = 4; - const int64_t v12 = 2; + 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 = 1024; - const int64_t v17 = 0; - const int64_t v18 = 59392; - const int64_t v19 = 43008; - const int64_t v20 = 41984; - const int64_t v21 = 33792; - const int64_t v22 = 17408; + 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); - for (size_t v23 = (size_t)v17; v23 < ((size_t)v11); v23 += (size_t)v12) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v10); - int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)v24); - int64_t v26 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v24 + (uint64_t)v10)); + 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> - v27 = Tile< + v26 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v28 = (uint64_t)v22; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v31 = GlobalTensor< + v30 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v25 * v13), v29, v30 + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v33 = (uint64_t)v21; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v25 * v13), v34, v35 + 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 ); - TLOAD(v32, v36); + 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> - v37 = Tile< + v36 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v38 = (uint64_t)v20; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 1, 256> v39 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v40 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v41 = + 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 + (v17 + v7 * v14 + v25 * v13), v39, v40 + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 ); - TLOAD(v37, v41); + 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> - v42 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v43 = (uint64_t)v19; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v46 = GlobalTensor< + v45 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v26 * v13), v44, v45 + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 ); - TLOAD(v42, v46); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v47 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v48 = (uint64_t)v18; - TASSIGN(v47, v48); - pto::Shape<1, 1, 1, 16, 256> v49 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v50 = 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> - v51 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v26 * v13), v49, v50 + 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 ); - TLOAD(v47, v51); + 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> - v52 = Tile< + v51 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v53 = (uint64_t)v17; - TASSIGN(v52, v53); - pto::Shape<1, 1, 1, 1, 256> v54 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v55 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v56 = + 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 + (v17 + v7 * v14 + v26 * v13), v54, v55 + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 ); - TLOAD(v52, v56); + 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> - v57 = Tile< + v56 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v58 = (uint64_t)v16; - TASSIGN(v57, v58); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v57, v32, v9, v8); - 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>(v15, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - TADD(v59, v27, v57); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v62 = (uint64_t)v16; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v9, v8); + 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> - v63 = Tile< + v58 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v64 = (uint64_t)v22; - TASSIGN(v63, v64); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v63, v59, v37); + 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> - v65 = Tile< + v60 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v66 = (uint64_t)v21; - TASSIGN(v65, v66); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); pipe_barrier(PIPE_V); - TCVT(v65, v63, v9, v8); + 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> - v67 = Tile< + v68 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pipe_barrier(PIPE_V); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v67, v47, v9, v8); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v70 = (uint64_t)v22; - TASSIGN(v69, v70); - pipe_barrier(PIPE_V); - TADD(v69, v42, v67); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); - pipe_barrier(PIPE_V); - TCVT(v71, v69, v9, v8); + 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> - v73 = Tile< + v70 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v74 = (uint64_t)v22; - TASSIGN(v73, v74); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v73, v69, v52); + 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> - v75 = Tile< + v72 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v76 = (uint64_t)v22; - TASSIGN(v75, v76); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); pipe_barrier(PIPE_V); - TCVT(v75, v73, v9, v8); + 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>( - v2 + (v17 + v17 * v14 + v25 * v13), v77, v78 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v79, v61); - pto::Shape<1, 1, 1, 16, 256> v80 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v81 = 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> - v82 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v25 * v13), v80, v81 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v82, v65); - pto::Shape<1, 1, 1, 16, 256> v83 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v17 + v17 * v14 + v26 * v13), v83, v84 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v85, v71); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v26 * v13), v86, v87 + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v88, v75); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -338,28 +313,28 @@ static __aicore__ void residual_rms_cast_2( } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v8 - __gm__ Tensor *mlp_norm_in_inline250__rv_v8_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v8 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v8_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v8_tensor->start_offset; + // 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_inline130__rv_v8 - __gm__ Tensor *post_norm_partial_inline130__rv_v8_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *post_norm_partial_inline130__rv_v8 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__rv_v8_tensor->buffer.addr) + - post_norm_partial_inline130__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_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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]); @@ -367,13 +342,13 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + post_rms_weight__ssa_v0_tensor->start_offset; - // Unpack scalar: k_base_inline177__ssa_v3 + // Unpack scalar: k_base_inline111__ssa_v3 union { uint64_t u64; int64_t val; - } k_base_inline177__ssa_v3_conv; - k_base_inline177__ssa_v3_conv.u64 = args[5]; - int64_t k_base_inline177__ssa_v3 = k_base_inline177__ssa_v3_conv.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 { @@ -385,7 +360,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Forward to ptoas-generated function residual_rms_cast_2( - mlp_norm_in_inline250__rv_v8, post_norm_partial_inline130__rv_v8, attn_proj_fp32_inline248__rv_v5, cur__iter_v6, - post_rms_weight__ssa_v0, k_base_inline177__ssa_v3, i__idx_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp index 01ebc4c377..9532317302 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/residual_rms_cast_3.cpp @@ -51,286 +51,261 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } static __aicore__ void residual_rms_cast_3( - __gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, __gm__ bfloat16_t *v4, __gm__ float *v5, int64_t v6, + __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 = 4; - const int64_t v12 = 2; + 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 = 1024; - const int64_t v17 = 0; - const int64_t v18 = 59392; - const int64_t v19 = 43008; - const int64_t v20 = 41984; - const int64_t v21 = 33792; - const int64_t v22 = 17408; + 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); - for (size_t v23 = (size_t)v17; v23 < ((size_t)v11); v23 += (size_t)v12) { - int64_t v24 = (int64_t)((uint64_t)((int64_t)v23) * (uint64_t)v10); - int64_t v25 = (int64_t)((uint64_t)v6 + (uint64_t)v24); - int64_t v26 = (int64_t)((uint64_t)v6 + (uint64_t)((int64_t)(uint64_t)v24 + (uint64_t)v10)); + 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> - v27 = Tile< + v26 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v28 = (uint64_t)v22; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v31 = GlobalTensor< + v30 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v25 * v13), v29, v30 + v3 + (v16 + v16 * v14 + v24 * v13), v28, v29 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v27, v31); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TLOAD(v26, v30); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v31 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v33 = (uint64_t)v21; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v25 * v13), v34, v35 + 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 ); - TLOAD(v32, v36); + 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> - v37 = Tile< + v36 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v38 = (uint64_t)v20; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 1, 256> v39 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v40 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v41 = + 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 + (v17 + v7 * v14 + v25 * v13), v39, v40 + v5 + (v16 + v7 * v14 + v24 * v13), v38, v39 ); - TLOAD(v37, v41); + 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> - v42 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v43 = (uint64_t)v19; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = pto::Stride<81920, 81920, 81920, 5120, 1>(); + 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> - v46 = GlobalTensor< + v45 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v3 + (v17 + v17 * v14 + v26 * v13), v44, v45 + v3 + (v16 + v16 * v14 + v25 * v13), v43, v44 ); - TLOAD(v42, v46); + wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); + wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + TLOAD(v41, v45); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v47 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + v46 = Tile< + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v48 = (uint64_t)v18; - TASSIGN(v47, v48); - pto::Shape<1, 1, 1, 16, 256> v49 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v50 = 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> - v51 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v4 + (v17 + v17 * v14 + v26 * v13), v49, v50 + 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 ); - TLOAD(v47, v51); + 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> - v52 = Tile< + v51 = Tile< TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v13, v10); - uint64_t v53 = (uint64_t)v17; - TASSIGN(v52, v53); - pto::Shape<1, 1, 1, 1, 256> v54 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v55 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v56 = + 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 + (v17 + v7 * v14 + v26 * v13), v54, v55 + v5 + (v16 + v7 * v14 + v25 * v13), v53, v54 ); - TLOAD(v52, v56); + 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> - v57 = Tile< + v56 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v58 = (uint64_t)v16; - TASSIGN(v57, v58); + uint64_t v57 = (uint64_t)v21; + TASSIGN(v56, v57); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v57, v32, v9, v8); - 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>(v15, v10); - uint64_t v60 = (uint64_t)v22; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - TADD(v59, v27, v57); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v62 = (uint64_t)v16; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v9, v8); + 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> - v63 = Tile< + v58 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v64 = (uint64_t)v22; - TASSIGN(v63, v64); + uint64_t v59 = (uint64_t)v20; + TASSIGN(v58, v59); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v63, v59, v37); + 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> - v65 = Tile< + v60 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v66 = (uint64_t)v21; - TASSIGN(v65, v66); + uint64_t v61 = (uint64_t)v20; + TASSIGN(v60, v61); pipe_barrier(PIPE_V); - TCVT(v65, v63, v9, v8); + 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> - v67 = Tile< + v68 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v68 = (uint64_t)v22; - TASSIGN(v67, v68); - pipe_barrier(PIPE_V); + uint64_t v69 = (uint64_t)v18; + TASSIGN(v68, v69); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v67, v47, v9, v8); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v70 = (uint64_t)v22; - TASSIGN(v69, v70); - pipe_barrier(PIPE_V); - TADD(v69, v42, v67); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v10); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); - pipe_barrier(PIPE_V); - TCVT(v71, v69, v9, v8); + 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> - v73 = Tile< + v70 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v74 = (uint64_t)v22; - TASSIGN(v73, v74); + uint64_t v71 = (uint64_t)v17; + TASSIGN(v70, v71); pipe_barrier(PIPE_V); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v73, v69, v52); + 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> - v75 = Tile< + v72 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v15, v10); - uint64_t v76 = (uint64_t)v22; - TASSIGN(v75, v76); + uint64_t v73 = (uint64_t)v17; + TASSIGN(v72, v73); pipe_barrier(PIPE_V); - TCVT(v75, v73, v9, v8); + 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>( - v2 + (v17 + v17 * v14 + v25 * v13), v77, v78 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v79, v61); - pto::Shape<1, 1, 1, 16, 256> v80 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v81 = 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> - v82 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v25 * v13), v80, v81 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v82, v65); - pto::Shape<1, 1, 1, 16, 256> v83 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v84 = 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> - v85 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v17 + v17 * v14 + v26 * v13), v83, v84 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v85, v71); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v17 + v17 * v14 + v26 * v13), v86, v87 + v1 + (v16 + v16 * v14 + v25 * v13), v77, v78 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v88, v75); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -338,28 +313,28 @@ static __aicore__ void residual_rms_cast_3( } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: mlp_norm_in_inline250__rv_v11 - __gm__ Tensor *mlp_norm_in_inline250__rv_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *mlp_norm_in_inline250__rv_v11 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_norm_in_inline250__rv_v11_tensor->buffer.addr) + - mlp_norm_in_inline250__rv_v11_tensor->start_offset; + // 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_inline130__rv_v11 - __gm__ Tensor *post_norm_partial_inline130__rv_v11_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *post_norm_partial_inline130__rv_v11 = - reinterpret_cast<__gm__ bfloat16_t *>(post_norm_partial_inline130__rv_v11_tensor->buffer.addr) + - post_norm_partial_inline130__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_inline248__rv_v5 - __gm__ Tensor *attn_proj_fp32_inline248__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *attn_proj_fp32_inline248__rv_v5 = - reinterpret_cast<__gm__ float *>(attn_proj_fp32_inline248__rv_v5_tensor->buffer.addr) + - attn_proj_fp32_inline248__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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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]); @@ -367,13 +342,13 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in reinterpret_cast<__gm__ float *>(post_rms_weight__ssa_v0_tensor->buffer.addr) + post_rms_weight__ssa_v0_tensor->start_offset; - // Unpack scalar: k_base_inline177__ssa_v4 + // Unpack scalar: k_base_inline111__ssa_v4 union { uint64_t u64; int64_t val; - } k_base_inline177__ssa_v4_conv; - k_base_inline177__ssa_v4_conv.u64 = args[5]; - int64_t k_base_inline177__ssa_v4 = k_base_inline177__ssa_v4_conv.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 { @@ -385,7 +360,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // Forward to ptoas-generated function residual_rms_cast_3( - mlp_norm_in_inline250__rv_v11, post_norm_partial_inline130__rv_v11, attn_proj_fp32_inline248__rv_v5, - cur__iter_v6, post_rms_weight__ssa_v0, k_base_inline177__ssa_v4, i__idx_v0 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rms_recip.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rms_recip.cpp index 92e91d0462..a4a9666bcb 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rms_recip.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rms_recip.cpp @@ -50,27 +50,31 @@ static __aicore__ inline void ptoas_auto_sync_tail(PTOAutoSyncTailMode mode = PT } } -static __aicore__ void rms_recip(__gm__ bfloat16_t *v1, __gm__ float *v2) { - SaturationMode v3 = SaturationMode::OFF; - RoundMode v4 = RoundMode::CAST_ROUND; - const float v5 = 9.99999997E-7f; - const float v6 = 1.95312503E-4f; - const int64_t v7 = 768; - const int64_t v8 = 512; - const int64_t v9 = 256; - const int64_t v10 = 4; - const int64_t v11 = 20; - const float v12 = 0.0f; - const int64_t v13 = 1; - const int64_t v14 = 5120; - const int64_t v15 = 16; - const int64_t v16 = 0; - const int64_t v17 = 24704; - const int64_t v18 = 16512; - const int64_t v19 = 8320; - const int64_t v20 = 128; - const int64_t v21 = 41088; - const int64_t v22 = 64; +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__) @@ -79,181 +83,159 @@ static __aicore__ void rms_recip(__gm__ bfloat16_t *v1, __gm__ float *v2) { Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v23 = Tile< + v27 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v24 = (uint64_t)v22; - TASSIGN(v23, v24); + 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); - TEXPANDS(v23, v12); - for (size_t v25 = (size_t)v16; v25 < ((size_t)v11); v25 += (size_t)v10) { - int64_t v26 = (int64_t)((uint64_t)((int64_t)v25) * (uint64_t)v9); + 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, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v27 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v28 = (uint64_t)v21; - TASSIGN(v27, v28); - pto::Shape<1, 1, 1, 16, 256> v29 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v30 = 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> - v31 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v16 + v16 * v14 + v26 * v13), v29, v30 + 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(v27, v31); + TLOAD(v31, v35); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v32 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v33 = (uint64_t)v20; - TASSIGN(v32, v33); - pto::Shape<1, 1, 1, 16, 256> v34 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v35 = 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> - v36 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v16 + v16 * v14 + (int64_t)((uint64_t)v26 + (uint64_t)v9) * v13), v34, v35 + 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(v32, v36); + TLOAD(v36, v40); set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v37 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v38 = (uint64_t)v19; - TASSIGN(v37, v38); - pto::Shape<1, 1, 1, 16, 256> v39 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v40 = 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> - v41 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v16 + v16 * v14 + (int64_t)((uint64_t)v26 + (uint64_t)v8) * v13), v39, v40 - ); - TLOAD(v37, v41); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v42 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v43 = (uint64_t)v18; - TASSIGN(v42, v43); - pto::Shape<1, 1, 1, 16, 256> v44 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v45 = 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> - v46 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v16 + v16 * v14 + (int64_t)((uint64_t)v26 + (uint64_t)v7) * v13), v44, v45 - ); - TLOAD(v42, v46); - 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> - v47 = Tile< + v41 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v48 = (uint64_t)v17; - TASSIGN(v47, v48); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TCVT(v47, v27, v4, v3); + 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> - v49 = Tile< + v46 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v50 = (uint64_t)v17; - TASSIGN(v49, v50); - pipe_barrier(PIPE_V); - TMUL(v49, v47, v47); + 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>(v15, v9); - uint64_t v52 = (uint64_t)v21; + 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, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v53 = Tile< - TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v13); - uint64_t v54 = (uint64_t)v16; + 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); - pipe_barrier(PIPE_V); - TROWSUM(v53, v49, v51); Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 1, BLayout::ColMajor, -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>(v13, v15); - uint64_t v56 = (uint64_t)v16; + 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>(v13, v15); - uint64_t v58 = (uint64_t)v16; + CompactMode::Null>(v11, v13); + uint64_t v58 = (uint64_t)v20; TASSIGN(v57, v58); - pipe_barrier(PIPE_V); - TADD(v57, v23, v55); Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 1, 16, 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>(v15, v9); - uint64_t v60 = (uint64_t)v17; + 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); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCVT(v59, v32, v4, v3); + 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>(v15, v9); - uint64_t v62 = (uint64_t)v17; + CompactMode::Null>(v13, v7); + uint64_t v62 = (uint64_t)v24; TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TMUL(v61, v59, v59); + 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>(v15, v9); - uint64_t v64 = (uint64_t)v21; + 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>(v15, v13); - uint64_t v66 = (uint64_t)v20; + CompactMode::Null>(v13, v11); + uint64_t v66 = (uint64_t)v18; TASSIGN(v65, v66); pipe_barrier(PIPE_V); TROWSUM(v65, v61, v63); @@ -262,209 +244,191 @@ static __aicore__ void rms_recip(__gm__ bfloat16_t *v1, __gm__ float *v2) { CompactMode::Null> v67 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v68 = (uint64_t)v20; + 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>(v13, v15); - uint64_t v70 = (uint64_t)v20; + CompactMode::Null>(v11, v13); + uint64_t v70 = (uint64_t)v24; TASSIGN(v69, v70); pipe_barrier(PIPE_V); - TADD(v69, v57, v67); + 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>(v15, v9); - uint64_t v72 = (uint64_t)v17; + CompactMode::Null>(v13, v7); + uint64_t v72 = (uint64_t)v23; TASSIGN(v71, v72); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v71, v37, v4, v3); + 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>(v15, v9); + CompactMode::Null>(v13, v7); uint64_t v74 = (uint64_t)v17; TASSIGN(v73, v74); - pipe_barrier(PIPE_V); - TMUL(v73, v71, v71); Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v75 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v76 = (uint64_t)v21; + 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, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 1, 16, BLayout::RowMajor, -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>(v15, v13); - uint64_t v78 = (uint64_t)v19; + 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); - pipe_barrier(PIPE_V); - TROWSUM(v77, v73, v75); 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>(v13, v15); - uint64_t v80 = (uint64_t)v19; + CompactMode::Null>(v11, v13); + uint64_t v80 = (uint64_t)v23; TASSIGN(v79, v80); - Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v81 = Tile< - TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v82 = (uint64_t)v20; - TASSIGN(v81, v82); pipe_barrier(PIPE_V); - TADD(v81, v69, v79); + 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> - v83 = Tile< + v81 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v84 = (uint64_t)v17; - TASSIGN(v83, v84); + CompactMode::Null>(v13, v7); + uint64_t v82 = (uint64_t)v22; + TASSIGN(v81, v82); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCVT(v83, v42, v4, v3); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v85 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v86 = (uint64_t)v17; - TASSIGN(v85, v86); - pipe_barrier(PIPE_V); - TMUL(v85, v83, v83); + TMUL(v81, v46, v46); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v87 = Tile< + v83 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v9); - uint64_t v88 = (uint64_t)v21; - TASSIGN(v87, v88); + 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> - v89 = Tile< + v85 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v13); - uint64_t v90 = (uint64_t)v19; - TASSIGN(v89, v90); + CompactMode::Null>(v13, v11); + uint64_t v86 = (uint64_t)v14; + TASSIGN(v85, v86); pipe_barrier(PIPE_V); - TROWSUM(v89, v85, v87); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); + TROWSUM(v85, v81, v83); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v91 = Tile< + v87 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v92 = (uint64_t)v19; - TASSIGN(v91, v92); + 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> - v93 = Tile< + v89 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v94 = (uint64_t)v22; - TASSIGN(v93, v94); + CompactMode::Null>(v11, v13); + uint64_t v90 = (uint64_t)v26; + TASSIGN(v89, v90); pipe_barrier(PIPE_V); - TADD(v93, v81, v91); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); + 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> - v95 = Tile< + v91 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v96 = (uint64_t)v17; - TASSIGN(v95, v96); + CompactMode::Null>(v11, v13); + uint64_t v92 = (uint64_t)v25; + TASSIGN(v91, v92); pipe_barrier(PIPE_V); - TMULS(v95, v23, v6); + TMULS(v91, v27, v4); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v97 = Tile< + v93 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v98 = (uint64_t)v17; - TASSIGN(v97, v98); + CompactMode::Null>(v11, v13); + uint64_t v94 = (uint64_t)v25; + TASSIGN(v93, v94); pipe_barrier(PIPE_V); - TADDS(v97, v95, v5); + TADDS(v93, v91, v3); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v99 = Tile< + v95 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v100 = (uint64_t)v17; - TASSIGN(v99, v100); + 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> - v101 = Tile< + v97 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v102 = (uint64_t)v17; - TASSIGN(v101, v102); + CompactMode::Null>(v11, v13); + uint64_t v98 = (uint64_t)v25; + TASSIGN(v97, v98); pipe_barrier(PIPE_V); - TSQRT(v101, v99); + TSQRT(v97, v95); Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v103 = Tile< + v99 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v104 = (uint64_t)v17; - TASSIGN(v103, v104); + 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> - v105 = Tile< + v101 = Tile< TileType::Vec, float, 1, 16, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v13, v15); - uint64_t v106 = (uint64_t)v21; - TASSIGN(v105, v106); + CompactMode::Null>(v11, v13); + uint64_t v102 = (uint64_t)v24; + TASSIGN(v101, v102); pipe_barrier(PIPE_V); - TRECIP(v105, v103); + 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> - v107 = Tile< + v103 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v15, v13); - uint64_t v108 = (uint64_t)v21; - TASSIGN(v107, v108); - pto::Shape<1, 1, 1, 16, 1> v109 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 16> v110 = pto::Stride<16, 16, 16, 1, 16>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v111 = + 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 + (v16 + v16 * v13 + v16 * v15), v109, v110 + v2 + (v19 + v19 * v11 + v19 * v13), v105, v106 ); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v111, v107); + 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); @@ -474,15 +438,15 @@ static __aicore__ void rms_recip(__gm__ bfloat16_t *v1, __gm__ float *v2) { 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__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_tensor->start_offset; + __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_inline160__ssa_v0 - __gm__ Tensor *inv_rms_states_inline160__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *inv_rms_states_inline160__ssa_v0 = - reinterpret_cast<__gm__ float *>(inv_rms_states_inline160__ssa_v0_tensor->buffer.addr) + - inv_rms_states_inline160__ssa_v0_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_inline160__ssa_v0); + rms_recip(cur__iter_v6, inv_rms_states_inline176__ssa_v0); } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rope_qkv.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rope_qkv.cpp deleted file mode 100644 index a67bf80200..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rope_qkv.cpp +++ /dev/null @@ -1,1051 +0,0 @@ -/* - * 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: rope_qkv -// 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 rope_qkv( - __gm__ int32_t *v1, __gm__ float *v2, __gm__ int32_t *v3, __gm__ float *v4, __gm__ float *v5, __gm__ bfloat16_t *v6, - __gm__ bfloat16_t *v7, __gm__ bfloat16_t *v8, __gm__ float *v9, __gm__ float *v10, __gm__ float *v11, - __gm__ float *v12, __gm__ float *v13, int64_t v14, int32_t v15, int32_t v16 -) { - SaturationMode v17 = SaturationMode::OFF; - RoundMode v18 = RoundMode::CAST_ROUND; - const float v19 = 0.0f; - const int64_t v20 = 11; - const int64_t v21 = 80; - const int64_t v22 = 5; - const int64_t v23 = 2; - const int64_t v24 = 8; - const int64_t v25 = 64; - const int64_t v26 = 5120; - const int64_t v27 = 128; - const int64_t v28 = 1; - const int64_t v29 = 16; - const int64_t v30 = 8704; - const int64_t v31 = 11264; - const int64_t v32 = 2688; - const int64_t v33 = 2560; - const int64_t v34 = 2304; - const int64_t v35 = 2048; - const int64_t v36 = 3072; - const int64_t v37 = 1536; - const int64_t v38 = 1024; - const int64_t v39 = 8448; - const int64_t v40 = 2816; - const int64_t v41 = 768; - const int64_t v42 = 512; - const int64_t v43 = 256; - const int64_t v44 = 0; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - size_t v45 = (size_t)v28; - size_t v46 = (size_t)v44; - size_t v47 = (size_t)v22; - int64_t v48 = (int64_t)v15; - int32_t v49 = v1[v48]; - float v50 = v2[v48]; - int64_t v51 = (int64_t)((uint64_t)((int64_t)v49) - (uint64_t)v28); - int32_t v52 = v3[v48]; - int64_t v53 = (int64_t)v52; - int64_t v54 = v53 / v27; - int64_t v55 = (int64_t)((uint64_t)v53 - (uint64_t)((int64_t)(uint64_t)v54 * (uint64_t)v27)); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v56 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v57 = (uint64_t)v44; - TASSIGN(v56, v57); - pto::Shape<1, 1, 1, 1, 64> v58 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v59 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v60 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v4 + (v44 + v51 * v27 + v44 * v28), v58, v59 - ); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); - TLOAD(v56, v60); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v61 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v62 = (uint64_t)v43; - TASSIGN(v61, v62); - pto::Shape<1, 1, 1, 1, 64> v63 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v64 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v65 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v4 + (v44 + v51 * v27 + v25 * v28), v63, v64 - ); - TLOAD(v61, v65); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v66 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v67 = (uint64_t)v42; - TASSIGN(v66, v67); - pto::Shape<1, 1, 1, 1, 64> v68 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v69 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v70 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v5 + (v44 + v51 * v27 + v44 * v28), v68, v69 - ); - TLOAD(v66, v70); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v71 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v72 = (uint64_t)v41; - TASSIGN(v71, v72); - pto::Shape<1, 1, 1, 1, 64> v73 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v74 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v75 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v5 + (v44 + v51 * v27 + v25 * v28), v73, v74 - ); - TLOAD(v71, v75); - for (size_t v76 = v46; v76 < ((size_t)v24); v76 += (size_t)v23) { - int64_t v77 = (int64_t)v76; - int64_t v78 = (int64_t)((uint64_t)v77 * (uint64_t)v27); - int64_t v79 = (int64_t)((uint64_t)v77 * (uint64_t)v29); - int64_t v80 = (int64_t)((uint64_t)v79 + (uint64_t)v48); - float v81 = v9[v80]; - int64_t v82 = (int64_t)((uint64_t)v54 * (uint64_t)v24); - int64_t v83 = (int64_t)((uint64_t)((int64_t)(uint64_t)v14 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v82 + (uint64_t)v77) * - (uint64_t)v27)) + - (uint64_t)v55); - int64_t v84 = (int64_t)((uint64_t)v77 * (uint64_t)v22); - int64_t v85 = (int64_t)((uint64_t)v48 * (uint64_t)v27); - int64_t v86 = (int64_t)((uint64_t)v85 + (uint64_t)v79); - int64_t v87 = (int64_t)((uint64_t)((int64_t)(uint64_t)v77 * (uint64_t)v21) + - (uint64_t)((int64_t)(uint64_t)v48 * (uint64_t)v22)); - int64_t v88 = (int64_t)((uint64_t)v78 + (uint64_t)v27); - float v89 = v9[(int64_t)((uint64_t)v80 + (uint64_t)v29)]; - int64_t v90 = - (int64_t)((uint64_t)((int64_t)(uint64_t)v14 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v82 + - (uint64_t)((int64_t)(uint64_t)v77 + (uint64_t)v28)) * - (uint64_t)v27)) + - (uint64_t)v55); - int64_t v91 = (int64_t)((uint64_t)v85 + (uint64_t)((int64_t)(uint64_t)v79 + (uint64_t)v29)); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v92 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v93 = (uint64_t)v40; - TASSIGN(v92, v93); - pto::Shape<1, 1, 1, 1, 128> v94 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<1024, 1024, 1024, 1024, 1> v95 = pto::Stride<1024, 1024, 1024, 1024, 1>(); - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v96 = - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>( - v10 + (v44 + v48 * v38 + v78 * v28), v94, v95 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v92, v96); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v97 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v98 = (uint64_t)v39; - TASSIGN(v97, v98); - pto::Shape<1, 1, 1, 1, 128> v99 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<1024, 1024, 1024, 1024, 1> v100 = pto::Stride<1024, 1024, 1024, 1024, 1>(); - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v101 = - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>( - v11 + (v44 + v48 * v38 + v78 * v28), v99, v100 - ); - TLOAD(v97, v101); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v102 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v103 = (uint64_t)v38; - TASSIGN(v102, v103); - pto::Shape<1, 1, 1, 1, 128> v104 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<1024, 1024, 1024, 1024, 1> v105 = pto::Stride<1024, 1024, 1024, 1024, 1>(); - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v106 = - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>( - v10 + (v44 + v48 * v38 + v88 * v28), v104, v105 - ); - TLOAD(v102, v106); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v107 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v108 = (uint64_t)v37; - TASSIGN(v107, v108); - pto::Shape<1, 1, 1, 1, 128> v109 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<1024, 1024, 1024, 1024, 1> v110 = pto::Stride<1024, 1024, 1024, 1024, 1>(); - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND> v111 = - GlobalTensor, pto::Stride<1024, 1024, 1024, 1024, 1>, pto::Layout::ND>( - v11 + (v44 + v48 * v38 + v88 * v28), v109, v110 - ); - TLOAD(v107, v111); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v112 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v113 = (uint64_t)v40; - TASSIGN(v112, v113); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TMULS(v112, v92, v81); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v114 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v115 = (uint64_t)v40; - TASSIGN(v114, v115); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v116 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v117 = (uint64_t)v36; - TASSIGN(v116, v117); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v118 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v119 = (uint64_t)v35; - TASSIGN(v118, v119); - pipe_barrier(PIPE_V); - TEXTRACT(v114, v112, v44, v44); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v118, v114, v56); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v120 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v121 = (uint64_t)v34; - TASSIGN(v120, v121); - TEXTRACT(v116, v112, v44, v25); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v120, v116, v66); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v122 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v123 = (uint64_t)v35; - TASSIGN(v122, v123); - pipe_barrier(PIPE_V); - TSUB(v122, v118, v120); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v124 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v125 = (uint64_t)v34; - TASSIGN(v124, v125); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v124, v116, v61); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v126 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v127 = (uint64_t)v40; - TASSIGN(v126, v127); - TCOLEXPANDMUL(v126, v114, v71); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v128 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v129 = (uint64_t)v40; - TASSIGN(v128, v129); - pipe_barrier(PIPE_V); - TADD(v128, v124, v126); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v130 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v131 = (uint64_t)v33; - TASSIGN(v130, v131); - TCVT(v130, v122, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v132 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v133 = (uint64_t)v32; - TASSIGN(v132, v133); - pipe_barrier(PIPE_V); - TCVT(v132, v128, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v134 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v135 = (uint64_t)v40; - TASSIGN(v134, v135); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TMULS(v134, v97, v50); - Tile< - TileType::Vec, bfloat16_t, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v136 = Tile< - TileType::Vec, bfloat16_t, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v137 = (uint64_t)v35; - TASSIGN(v136, v137); - pipe_barrier(PIPE_V); - TCVT(v136, v134, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - for (size_t v138 = v46; v138 < v47; v138 += v45) { - int64_t v139 = (int64_t)v138; - float v140 = v12[(int64_t)((uint64_t)v87 + (uint64_t)v139)]; - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v141 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v142 = (uint64_t)v40; - TASSIGN(v141, v142); - pto::Shape<1, 1, 1, 1, 128> v143 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v144 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> - v145 = GlobalTensor< - float, pto::Shape<1, 1, 1, 1, 128>, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( - v13 + (v44 + v48 * v26 + - (int64_t)((uint64_t)((int64_t)(uint64_t)v84 + (uint64_t)v139) * (uint64_t)v27) * v28), - v143, v144 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - TLOAD(v141, v145); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v146 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v147 = (uint64_t)v40; - TASSIGN(v146, v147); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TMULS(v146, v141, v140); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v148 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v149 = (uint64_t)v40; - TASSIGN(v148, v149); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v150 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v151 = (uint64_t)v36; - TASSIGN(v150, v151); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v152 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v153 = (uint64_t)v39; - TASSIGN(v152, v153); - pipe_barrier(PIPE_V); - TEXTRACT(v148, v146, v44, v44); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v152, v148, v56); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v154 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v155 = (uint64_t)v34; - TASSIGN(v154, v155); - TEXTRACT(v150, v146, v44, v25); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v154, v150, v66); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v156 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v157 = (uint64_t)v39; - TASSIGN(v156, v157); - pipe_barrier(PIPE_V); - TSUB(v156, v152, v154); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v158 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v159 = (uint64_t)v34; - TASSIGN(v158, v159); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v158, v150, v61); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v160 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v161 = (uint64_t)v40; - TASSIGN(v160, v161); - TCOLEXPANDMUL(v160, v148, v71); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v162 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v163 = (uint64_t)v40; - TASSIGN(v162, v163); - pipe_barrier(PIPE_V); - TADD(v162, v158, v160); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v164 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v165 = (uint64_t)v39; - TASSIGN(v164, v165); - TCVT(v164, v156, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v166 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v167 = (uint64_t)v40; - TASSIGN(v166, v167); - pipe_barrier(PIPE_V); - TCVT(v166, v162, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - int64_t v168 = (int64_t)((uint64_t)v86 + (uint64_t)v139); - pto::Shape<1, 1, 1, 1, 64> v169 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v170 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v171 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 64>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v6 + (v44 + v168 * v27 + v44 * v28), v169, v170 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v171, v164); - pto::Shape<1, 1, 1, 1, 64> v172 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v173 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v174 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 64>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v6 + (v44 + v168 * v27 + v25 * v28), v172, v173 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v174, v166); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - }; - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - Tile< - TileType::Vec, float, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v175 = Tile< - TileType::Vec, float, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v27); - uint64_t v176 = (uint64_t)v40; - TASSIGN(v175, v176); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pipe_barrier(PIPE_V); - TEXPANDS(v175, v19); - Tile< - TileType::Vec, bfloat16_t, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v177 = Tile< - TileType::Vec, bfloat16_t, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v27); - uint64_t v178 = (uint64_t)v39; - TASSIGN(v177, v178); - pipe_barrier(PIPE_V); - TCVT(v177, v175, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v179 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v180 = (uint64_t)v40; - TASSIGN(v179, v180); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TMULS(v179, v102, v89); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v181 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v182 = (uint64_t)v40; - TASSIGN(v181, v182); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v183 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v184 = (uint64_t)v36; - TASSIGN(v183, v184); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v185 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v186 = (uint64_t)v38; - TASSIGN(v185, v186); - pipe_barrier(PIPE_V); - TEXTRACT(v181, v179, v44, v44); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v185, v181, v56); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v187 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v188 = (uint64_t)v34; - TASSIGN(v187, v188); - TEXTRACT(v183, v179, v44, v25); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v187, v183, v66); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v189 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v190 = (uint64_t)v38; - TASSIGN(v189, v190); - pipe_barrier(PIPE_V); - TSUB(v189, v185, v187); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v191 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v192 = (uint64_t)v34; - TASSIGN(v191, v192); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v191, v183, v61); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v193 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v194 = (uint64_t)v40; - TASSIGN(v193, v194); - TCOLEXPANDMUL(v193, v181, v71); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v195 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v196 = (uint64_t)v40; - TASSIGN(v195, v196); - pipe_barrier(PIPE_V); - TADD(v195, v191, v193); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v197 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v198 = (uint64_t)v34; - TASSIGN(v197, v198); - pipe_barrier(PIPE_V); - TCVT(v197, v189, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v199 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v200 = (uint64_t)v31; - TASSIGN(v199, v200); - TCVT(v199, v195, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v201 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v202 = (uint64_t)v40; - TASSIGN(v201, v202); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - TMULS(v201, v107, v50); - Tile< - TileType::Vec, bfloat16_t, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v203 = Tile< - TileType::Vec, bfloat16_t, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v204 = (uint64_t)v38; - TASSIGN(v203, v204); - pipe_barrier(PIPE_V); - TCVT(v203, v201, v18, v17); - Tile< - TileType::Vec, float, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v205 = Tile< - TileType::Vec, float, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v27); - uint64_t v206 = (uint64_t)v40; - TASSIGN(v205, v206); - pipe_barrier(PIPE_V); - TEXPANDS(v205, v19); - Tile< - TileType::Vec, bfloat16_t, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v207 = Tile< - TileType::Vec, bfloat16_t, 11, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v20, v27); - uint64_t v208 = (uint64_t)v40; - TASSIGN(v207, v208); - pipe_barrier(PIPE_V); - TCVT(v207, v205, v18, v17); - pto::Shape<1, 1, 1, 1, 64> v209 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v210 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v211 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v7 + (v44 + v83 * v27 + v44 * v28), v209, v210 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v211, v130); - pto::Shape<1, 1, 1, 1, 64> v212 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v213 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v214 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v7 + (v44 + v83 * v27 + v25 * v28), v212, v213 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID3); - TSTORE(v214, v132); - pto::Shape<1, 1, 1, 1, 128> v215 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<128, 128, 128, 128, 1> v216 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v217 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 128>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v8 + (v44 + v83 * v27 + v44 * v28), v215, v216 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID4); - TSTORE(v217, v136); - pto::Shape<1, 1, 1, 11, 128> v218 = pto::Shape<1, 1, 1, 11, 128>(); - pto::Stride<1408, 1408, 1408, 128, 1> v219 = pto::Stride<1408, 1408, 1408, 128, 1>(); - GlobalTensor, pto::Stride<1408, 1408, 1408, 128, 1>, pto::Layout::ND> - v220 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 11, 128>, pto::Stride<1408, 1408, 1408, 128, 1>, pto::Layout::ND>( - v6 + (v44 + (int64_t)((uint64_t)v86 + (uint64_t)v22) * v27 + v44 * v28), v218, v219 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID5); - TSTORE(v220, v177); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID2); - for (size_t v221 = v46; v221 < v47; v221 += v45) { - int64_t v222 = (int64_t)v221; - float v223 = v12[(int64_t)((uint64_t)((int64_t)(uint64_t)v87 + (uint64_t)v21) + (uint64_t)v222)]; - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v224 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v225 = (uint64_t)v39; - TASSIGN(v224, v225); - pto::Shape<1, 1, 1, 1, 128> v226 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v227 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> - v228 = GlobalTensor< - float, pto::Shape<1, 1, 1, 1, 128>, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( - v13 + (v44 + v48 * v26 + - (int64_t)((uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v84 + (uint64_t)v22) + - (uint64_t)v222) * - (uint64_t)v27) * - v28), - v226, v227 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); - TLOAD(v224, v228); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v229 = Tile< - TileType::Vec, float, 1, 128, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v27); - uint64_t v230 = (uint64_t)v39; - TASSIGN(v229, v230); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - TMULS(v229, v224, v223); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v231 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v232 = (uint64_t)v39; - TASSIGN(v231, v232); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v233 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v234 = (uint64_t)v30; - TASSIGN(v233, v234); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v235 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v236 = (uint64_t)v37; - TASSIGN(v235, v236); - pipe_barrier(PIPE_V); - TEXTRACT(v231, v229, v44, v44); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v235, v231, v56); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v237 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v238 = (uint64_t)v35; - TASSIGN(v237, v238); - TEXTRACT(v233, v229, v44, v25); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v237, v233, v66); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v239 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v240 = (uint64_t)v37; - TASSIGN(v239, v240); - pipe_barrier(PIPE_V); - TSUB(v239, v235, v237); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v241 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v242 = (uint64_t)v35; - TASSIGN(v241, v242); - pipe_barrier(PIPE_V); - TCOLEXPANDMUL(v241, v233, v61); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v243 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v244 = (uint64_t)v39; - TASSIGN(v243, v244); - TCOLEXPANDMUL(v243, v231, v71); - Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v245 = Tile< - TileType::Vec, float, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v246 = (uint64_t)v39; - TASSIGN(v245, v246); - pipe_barrier(PIPE_V); - TADD(v245, v241, v243); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v247 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v248 = (uint64_t)v37; - TASSIGN(v247, v248); - TCVT(v247, v239, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID6); - Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v249 = Tile< - TileType::Vec, bfloat16_t, 1, 64, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v28, v25); - uint64_t v250 = (uint64_t)v39; - TASSIGN(v249, v250); - pipe_barrier(PIPE_V); - TCVT(v249, v245, v18, v17); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID7); - int64_t v251 = (int64_t)((uint64_t)v91 + (uint64_t)v222); - pto::Shape<1, 1, 1, 1, 64> v252 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v253 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v254 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 64>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v6 + (v44 + v251 * v27 + v44 * v28), v252, v253 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID6); - TSTORE(v254, v247); - pto::Shape<1, 1, 1, 1, 64> v255 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v256 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v257 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 64>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v6 + (v44 + v251 * v27 + v25 * v28), v255, v256 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID7); - TSTORE(v257, v249); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID3); - }; - pto::Shape<1, 1, 1, 1, 64> v258 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v259 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v260 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v7 + (v44 + v90 * v27 + v44 * v28), v258, v259 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v260, v197); - pto::Shape<1, 1, 1, 1, 64> v261 = pto::Shape<1, 1, 1, 1, 64>(); - pto::Stride<128, 128, 128, 128, 1> v262 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> v263 = - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v7 + (v44 + v90 * v27 + v25 * v28), v261, v262 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v263, v199); - pto::Shape<1, 1, 1, 1, 128> v264 = pto::Shape<1, 1, 1, 1, 128>(); - pto::Stride<128, 128, 128, 128, 1> v265 = pto::Stride<128, 128, 128, 128, 1>(); - GlobalTensor, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND> - v266 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 1, 128>, pto::Stride<128, 128, 128, 128, 1>, pto::Layout::ND>( - v8 + (v44 + v90 * v27 + v44 * v28), v264, v265 - ); - pipe_barrier(PIPE_ALL); - TSTORE(v266, v203); - pto::Shape<1, 1, 1, 11, 128> v267 = pto::Shape<1, 1, 1, 11, 128>(); - pto::Stride<1408, 1408, 1408, 128, 1> v268 = pto::Stride<1408, 1408, 1408, 128, 1>(); - GlobalTensor, pto::Stride<1408, 1408, 1408, 128, 1>, pto::Layout::ND> - v269 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 11, 128>, pto::Stride<1408, 1408, 1408, 128, 1>, pto::Layout::ND>( - v6 + (v44 + (int64_t)((uint64_t)v91 + (uint64_t)v22) * v27 + v44 * v28), v267, v268 - ); - pipe_barrier(PIPE_ALL); - TSTORE(v269, v207); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - } - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_MTE3, 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) { - // 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: seq_lens__ssa_v0 - __gm__ Tensor *seq_lens__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ int32_t *seq_lens__ssa_v0 = reinterpret_cast<__gm__ int32_t *>(seq_lens__ssa_v0_tensor->buffer.addr) + - seq_lens__ssa_v0_tensor->start_offset; - - // Unpack tensor: inv_rms_states_inline160__ssa_v1 - __gm__ Tensor *inv_rms_states_inline160__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ float *inv_rms_states_inline160__ssa_v1 = - reinterpret_cast<__gm__ float *>(inv_rms_states_inline160__ssa_v1_tensor->buffer.addr) + - inv_rms_states_inline160__ssa_v1_tensor->start_offset; - - // Unpack tensor: slot_mapping__ssa_v0 - __gm__ Tensor *slot_mapping__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ int32_t *slot_mapping__ssa_v0 = - reinterpret_cast<__gm__ int32_t *>(slot_mapping__ssa_v0_tensor->buffer.addr) + - slot_mapping__ssa_v0_tensor->start_offset; - - // Unpack tensor: rope_cos__ssa_v0 - __gm__ Tensor *rope_cos__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *rope_cos__ssa_v0 = - reinterpret_cast<__gm__ float *>(rope_cos__ssa_v0_tensor->buffer.addr) + rope_cos__ssa_v0_tensor->start_offset; - - // Unpack tensor: rope_sin__ssa_v0 - __gm__ Tensor *rope_sin__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[4]); - __gm__ float *rope_sin__ssa_v0 = - reinterpret_cast<__gm__ float *>(rope_sin__ssa_v0_tensor->buffer.addr) + rope_sin__ssa_v0_tensor->start_offset; - - // Unpack tensor: all_q_padded_inline168__ssa_v0 - __gm__ Tensor *all_q_padded_inline168__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[5]); - __gm__ bfloat16_t *all_q_padded_inline168__ssa_v0 = - reinterpret_cast<__gm__ bfloat16_t *>(all_q_padded_inline168__ssa_v0_tensor->buffer.addr) + - all_q_padded_inline168__ssa_v0_tensor->start_offset; - - // Unpack tensor: k_cache__iter_v1 - __gm__ Tensor *k_cache__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[6]); - __gm__ bfloat16_t *k_cache__iter_v1 = reinterpret_cast<__gm__ bfloat16_t *>(k_cache__iter_v1_tensor->buffer.addr) + - k_cache__iter_v1_tensor->start_offset; - - // Unpack tensor: v_cache__iter_v1 - __gm__ Tensor *v_cache__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[7]); - __gm__ bfloat16_t *v_cache__iter_v1 = reinterpret_cast<__gm__ bfloat16_t *>(v_cache__iter_v1_tensor->buffer.addr) + - v_cache__iter_v1_tensor->start_offset; - - // Unpack tensor: k_inv_states_inline90__rv_v2 - __gm__ Tensor *k_inv_states_inline90__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[8]); - __gm__ float *k_inv_states_inline90__rv_v2 = - reinterpret_cast<__gm__ float *>(k_inv_states_inline90__rv_v2_tensor->buffer.addr) + - k_inv_states_inline90__rv_v2_tensor->start_offset; - - // Unpack tensor: k_proj_norm_inline155__rv_v2 - __gm__ Tensor *k_proj_norm_inline155__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[9]); - __gm__ float *k_proj_norm_inline155__rv_v2 = - reinterpret_cast<__gm__ float *>(k_proj_norm_inline155__rv_v2_tensor->buffer.addr) + - k_proj_norm_inline155__rv_v2_tensor->start_offset; - - // Unpack tensor: v_proj_inline152__rv_v3 - __gm__ Tensor *v_proj_inline152__rv_v3_tensor = reinterpret_cast<__gm__ Tensor *>(args[10]); - __gm__ float *v_proj_inline152__rv_v3 = - reinterpret_cast<__gm__ float *>(v_proj_inline152__rv_v3_tensor->buffer.addr) + - v_proj_inline152__rv_v3_tensor->start_offset; - - // Unpack tensor: q_inv_states_inline99__rv_v2 - __gm__ Tensor *q_inv_states_inline99__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[11]); - __gm__ float *q_inv_states_inline99__rv_v2 = - reinterpret_cast<__gm__ float *>(q_inv_states_inline99__rv_v2_tensor->buffer.addr) + - q_inv_states_inline99__rv_v2_tensor->start_offset; - - // Unpack tensor: q_proj_norm_inline197__rv_v2 - __gm__ Tensor *q_proj_norm_inline197__rv_v2_tensor = reinterpret_cast<__gm__ Tensor *>(args[12]); - __gm__ float *q_proj_norm_inline197__rv_v2 = - reinterpret_cast<__gm__ float *>(q_proj_norm_inline197__rv_v2_tensor->buffer.addr) + - q_proj_norm_inline197__rv_v2_tensor->start_offset; - - // Unpack scalar: layer_cache_base_inline137__ssa_v0 - union { - uint64_t u64; - int64_t val; - } layer_cache_base_inline137__ssa_v0_conv; - layer_cache_base_inline137__ssa_v0_conv.u64 = args[13]; - int64_t layer_cache_base_inline137__ssa_v0 = layer_cache_base_inline137__ssa_v0_conv.val; - - // Forward to ptoas-generated function - rope_qkv( - seq_lens__ssa_v0, inv_rms_states_inline160__ssa_v1, slot_mapping__ssa_v0, rope_cos__ssa_v0, rope_sin__ssa_v0, - all_q_padded_inline168__ssa_v0, k_cache__iter_v1, v_cache__iter_v1, k_inv_states_inline90__rv_v2, - k_proj_norm_inline155__rv_v2, v_proj_inline152__rv_v3, q_inv_states_inline99__rv_v2, - q_proj_norm_inline197__rv_v2, layer_cache_base_inline137__ssa_v0, __pypto_spmd_block_idx, __pypto_spmd_block_num - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/silu.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/silu.cpp index 8b847bee0a..327ce0009c 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/silu.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/silu.cpp @@ -60,13 +60,15 @@ static __aicore__ void silu(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ floa const int64_t v12 = 17408; const int64_t v13 = 1; const int64_t v14 = 16; - const int64_t v15 = 16384; - const int64_t v16 = 0; - const int64_t v17 = 81984; - const int64_t v18 = 65600; - const int64_t v19 = 49216; - const int64_t v20 = 32832; - const int64_t v21 = 32768; + 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__) @@ -75,301 +77,305 @@ static __aicore__ void silu(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ floa Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v22 = Tile< + v24 = Tile< TileType::Vec, float, 16, 1, BLayout::ColMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v13); - uint64_t v23 = (uint64_t)v21; - TASSIGN(v22, v23); - pto::Shape<1, 1, 1, 16, 1> v24 = pto::Shape<1, 1, 1, 16, 1>(); - pto::Stride<16, 16, 16, 1, 16> v25 = pto::Stride<16, 16, 16, 1, 16>(); - GlobalTensor, pto::Stride<16, 16, 16, 1, 16>, pto::Layout::DN> v26 = + 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 + (v16 + v16 * v13 + v16 * v14), v24, v25 + v1 + (v18 + v18 * v13 + v18 * v14), v26, v27 ); set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v22, v26); - for (size_t v27 = (size_t)v16; v27 < ((size_t)v11); v27 += (size_t)v10) { - int64_t v28 = (int64_t)((uint64_t)((int64_t)v27) * (uint64_t)v9); - int64_t v29 = (int64_t)((uint64_t)v5 + (uint64_t)v28); - int64_t v30 = (int64_t)((uint64_t)v5 + (uint64_t)((int64_t)(uint64_t)v28 + (uint64_t)v9)); + 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> - v31 = Tile< + v33 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - 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<278528, 278528, 278528, 17408, 1> v34 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v35 = GlobalTensor< + v37 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v3 + (v16 + v16 * v12 + v29 * v13), v33, v34 + v3 + (v18 + v18 * v12 + v31 * v13), v35, v36 ); wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - TLOAD(v31, v35); + 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> - v36 = Tile< + v38 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v37 = (uint64_t)v19; - TASSIGN(v36, v37); - pto::Shape<1, 1, 1, 16, 256> v38 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v39 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v40 = GlobalTensor< + v42 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v4 + (v16 + v16 * v12 + v29 * v13), v38, v39 + v4 + (v18 + v18 * v12 + v31 * v13), v40, v41 ); - TLOAD(v36, v40); + 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> - v41 = Tile< + v43 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - 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<278528, 278528, 278528, 17408, 1> v44 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v45 = GlobalTensor< + v47 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v3 + (v16 + v16 * v12 + v30 * v13), v43, v44 + v3 + (v18 + v18 * v12 + v32 * v13), v45, v46 ); - TLOAD(v41, v45); + 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> - v46 = Tile< + v48 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - 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<278528, 278528, 278528, 17408, 1> v49 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v50 = GlobalTensor< + v52 = GlobalTensor< float, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v4 + (v16 + v16 * v12 + v30 * v13), v48, v49 + v4 + (v18 + v18 * v12 + v32 * v13), v50, v51 ); - TLOAD(v46, v50); + 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> - v51 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v14, v9); - uint64_t v52 = (uint64_t)v20; - TASSIGN(v51, v52); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - TROWEXPANDMUL(v51, v31, v22); 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)v19; + uint64_t v54 = (uint64_t)v22; TASSIGN(v53, v54); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TROWEXPANDMUL(v53, v36, v22); + 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)v16; + uint64_t v56 = (uint64_t)v21; TASSIGN(v55, v56); - pipe_barrier(PIPE_V); - TNEG(v55, v51); + 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)v16; + uint64_t v58 = (uint64_t)v18; TASSIGN(v57, v58); pipe_barrier(PIPE_V); - TEXP(v57, v55); + 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)v16; + uint64_t v60 = (uint64_t)v18; TASSIGN(v59, v60); pipe_barrier(PIPE_V); - TADDS(v59, v57, v8); + 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)v15; + uint64_t v62 = (uint64_t)v18; TASSIGN(v61, v62); pipe_barrier(PIPE_V); - TRECIP(v61, v59); + 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)v20; + uint64_t v64 = (uint64_t)v17; TASSIGN(v63, v64); pipe_barrier(PIPE_V); - TMUL(v63, v51, v61); + 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)v20; + uint64_t v66 = (uint64_t)v22; TASSIGN(v65, v66); pipe_barrier(PIPE_V); - TMUL(v65, v63, v53); + TMUL(v65, v53, v63); Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v67 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v68 = (uint64_t)v16; + uint64_t v68 = (uint64_t)v22; TASSIGN(v67, v68); pipe_barrier(PIPE_V); - TCVT(v67, v65, v7, v6); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); + TMUL(v67, v65, v55); Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> v69 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, + TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v70 = (uint64_t)v20; + 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(v69, v41, v22); + TROWEXPANDMUL(v74, v43, v24); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v71 = Tile< + v76 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v72 = (uint64_t)v19; - TASSIGN(v71, v72); + uint64_t v77 = (uint64_t)v19; + TASSIGN(v76, v77); wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TROWEXPANDMUL(v71, v46, v22); + TROWEXPANDMUL(v76, v48, v24); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v73 = Tile< + v78 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v74 = (uint64_t)v18; - TASSIGN(v73, v74); + uint64_t v79 = (uint64_t)v16; + TASSIGN(v78, v79); pipe_barrier(PIPE_V); - TNEG(v73, v69); + TNEG(v78, v74); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v75 = Tile< + v80 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v76 = (uint64_t)v18; - TASSIGN(v75, v76); + uint64_t v81 = (uint64_t)v16; + TASSIGN(v80, v81); pipe_barrier(PIPE_V); - TEXP(v75, v73); + TEXP(v80, v78); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v77 = Tile< + v82 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v78 = (uint64_t)v18; - TASSIGN(v77, v78); + uint64_t v83 = (uint64_t)v16; + TASSIGN(v82, v83); pipe_barrier(PIPE_V); - TADDS(v77, v75, v8); + TADDS(v82, v80, v8); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v79 = Tile< + v84 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v80 = (uint64_t)v17; - TASSIGN(v79, v80); + uint64_t v85 = (uint64_t)v15; + TASSIGN(v84, v85); pipe_barrier(PIPE_V); - TRECIP(v79, v77); + TRECIP(v84, v82); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v81 = Tile< + v86 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v82 = (uint64_t)v20; - TASSIGN(v81, v82); + uint64_t v87 = (uint64_t)v20; + TASSIGN(v86, v87); pipe_barrier(PIPE_V); - TMUL(v81, v69, v79); + TMUL(v86, v74, v84); Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v83 = Tile< + v88 = Tile< TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v84 = (uint64_t)v20; - TASSIGN(v83, v84); + uint64_t v89 = (uint64_t)v20; + TASSIGN(v88, v89); pipe_barrier(PIPE_V); - TMUL(v83, v81, v71); + TMUL(v88, v86, v76); Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null> - v85 = Tile< + v90 = Tile< TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null>(v14, v9); - uint64_t v86 = (uint64_t)v20; - TASSIGN(v85, v86); + uint64_t v91 = (uint64_t)v20; + TASSIGN(v90, v91); pipe_barrier(PIPE_V); - TCVT(v85, v83, v7, v6); + TCVT(v90, v88, v7, v6); set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 256> v87 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v88 = pto::Stride<278528, 278528, 278528, 17408, 1>(); + 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> - v89 = GlobalTensor< + v94 = GlobalTensor< bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, - pto::Layout::ND>(v2 + (v16 + v16 * v12 + v29 * v13), v87, v88); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v89, v67); - pto::Shape<1, 1, 1, 16, 256> v90 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v91 = 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> - v92 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<278528, 278528, 278528, 17408, 1>, - pto::Layout::ND>(v2 + (v16 + v16 * v12 + v30 * v13), v90, v91); - pipe_barrier(PIPE_MTE3); + pto::Layout::ND>(v2 + (v18 + v18 * v12 + v32 * v13), v92, v93); wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v92, v85); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); + 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); @@ -377,41 +383,41 @@ static __aicore__ void silu(__gm__ float *v1, __gm__ bfloat16_t *v2, __gm__ floa } // --- Kernel entry point --- extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { - // Unpack tensor: inv_rms_tile_inline255__ssa_v1 - __gm__ Tensor *inv_rms_tile_inline255__ssa_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *inv_rms_tile_inline255__ssa_v1 = - reinterpret_cast<__gm__ float *>(inv_rms_tile_inline255__ssa_v1_tensor->buffer.addr) + - inv_rms_tile_inline255__ssa_v1_tensor->start_offset; + // 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_inline238__iter_v1 - __gm__ Tensor *mlp_tile_inline238__iter_v1_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *mlp_tile_inline238__iter_v1 = - reinterpret_cast<__gm__ bfloat16_t *>(mlp_tile_inline238__iter_v1_tensor->buffer.addr) + - mlp_tile_inline238__iter_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_inline170__rv_v5 - __gm__ Tensor *gate_acc_all_inline170__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ float *gate_acc_all_inline170__rv_v5 = - reinterpret_cast<__gm__ float *>(gate_acc_all_inline170__rv_v5_tensor->buffer.addr) + - gate_acc_all_inline170__rv_v5_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_inline247__rv_v5 - __gm__ Tensor *up_acc_all_inline247__rv_v5_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - __gm__ float *up_acc_all_inline247__rv_v5 = - reinterpret_cast<__gm__ float *>(up_acc_all_inline247__rv_v5_tensor->buffer.addr) + - up_acc_all_inline247__rv_v5_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_inline142__ssa_v7 + // Unpack scalar: n0_inline122__ssa_v7 union { uint64_t u64; int64_t val; - } n0_inline142__ssa_v7_conv; - n0_inline142__ssa_v7_conv.u64 = args[4]; - int64_t n0_inline142__ssa_v7 = n0_inline142__ssa_v7_conv.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_inline255__ssa_v1, mlp_tile_inline238__iter_v1, gate_acc_all_inline170__rv_v5, - up_acc_all_inline247__rv_v5, n0_inline142__ssa_v7 + 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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/up_seed.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/up_seed.cpp deleted file mode 100644 index 972a442819..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/up_seed.cpp +++ /dev/null @@ -1,158 +0,0 @@ -/* - * 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_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 up_seed(__gm__ float *v1) { - const int64_t v2 = 16384; - const float v3 = 0.0f; - const int64_t v4 = 1024; - const int64_t v5 = 2; - const int64_t v6 = 1; - const int64_t v7 = 17408; - const int64_t v8 = 16; - const int64_t v9 = 65536; - 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)v8); v11 += (size_t)v5) { - int64_t v12 = (int64_t)((uint64_t)((int64_t)v11) * (uint64_t)v4); - 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>(v8, v4); - uint64_t v14 = (uint64_t)v10; - TASSIGN(v13, v14); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TEXPANDS(v13, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v15 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v16 = (uint64_t)v9; - TASSIGN(v15, v16); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - TEXPANDS(v15, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 1024> v17 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v18 = 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> - v19 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v12 * v6), v17, v18 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - pipe_barrier(PIPE_MTE3); - TSTORE(v19, v13); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 1024> v20 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v21 = 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> - v22 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + (int64_t)((uint64_t)v12 + (uint64_t)v4) * v6), v20, v21 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v22, v15); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - } - 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> - v23 = Tile< - TileType::Vec, float, 16, 1024, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v8, v4); - uint64_t v24 = (uint64_t)v10; - TASSIGN(v23, v24); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID2); - TEXPANDS(v23, v3); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pto::Shape<1, 1, 1, 16, 1024> v25 = pto::Shape<1, 1, 1, 16, 1024>(); - pto::Stride<278528, 278528, 278528, 17408, 1> v26 = pto::Stride<278528, 278528, 278528, 17408, 1>(); - GlobalTensor, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND> - v27 = GlobalTensor< - float, pto::Shape<1, 1, 1, 16, 1024>, pto::Stride<278528, 278528, 278528, 17408, 1>, pto::Layout::ND>( - v1 + (v10 + v10 * v7 + v2 * v6), v25, v26 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pipe_barrier(PIPE_MTE3); - TSTORE(v27, v23); - 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: up_acc_all_inline247__ssa_v0 - __gm__ Tensor *up_acc_all_inline247__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ float *up_acc_all_inline247__ssa_v0 = - reinterpret_cast<__gm__ float *>(up_acc_all_inline247__ssa_v0_tensor->buffer.addr) + - up_acc_all_inline247__ssa_v0_tensor->start_offset; - - // Forward to ptoas-generated function - up_seed(up_acc_all_inline247__ssa_v0); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma.cpp deleted file mode 100644 index 974f103cd9..0000000000 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma.cpp +++ /dev/null @@ -1,383 +0,0 @@ -/* - * 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_gamma -// 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_gamma(__gm__ bfloat16_t *v1, __gm__ bfloat16_t *v2, __gm__ float *v3, int64_t v4, int32_t v5, int32_t v6) { - SaturationMode v7 = SaturationMode::OFF; - RoundMode v8 = RoundMode::CAST_ROUND; - const int64_t v9 = 5; - const int64_t v10 = 1280; - const int64_t v11 = 256; - const int64_t v12 = 20; - const int64_t v13 = 10; - const int64_t v14 = 4; - const int64_t v15 = 2; - const int64_t v16 = 1; - const int64_t v17 = 5120; - const int64_t v18 = 16; - const int64_t v19 = 18432; - const int64_t v20 = 17408; - const int64_t v21 = 9216; - const int64_t v22 = 8192; - const int64_t v23 = 0; - using T = float; - -#if defined(__DAV_VEC__) - set_mask_norm(); - set_vector_mask(-1, -1); - int64_t v24 = (int64_t)v5; - int64_t v25 = (int64_t)((uint64_t)v14 - (uint64_t)v24); - int64_t v26 = v25 / v13; - int64_t v27 = (int64_t)((uint64_t)v24 + - (uint64_t)((int64_t)(uint64_t)((int64_t)(uint64_t)v26 * (uint64_t)v13) + (uint64_t)v12)); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - for (size_t v28 = (size_t)v24; v28 < ((size_t)v27); v28 += (size_t)v13) { - int64_t v29 = (int64_t)((uint64_t)((int64_t)v28) * (uint64_t)v11); - int64_t v30 = (int64_t)((uint64_t)v29 + (uint64_t)v10); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v31 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v32 = (uint64_t)v23; - 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< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v35 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v23 + v23 * v17 + v29 * v16), v33, v34 - ); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - 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>(v16, v11); - uint64_t v37 = (uint64_t)v22; - 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>( - v3 + (v23 + v4 * v17 + v29 * v16), v38, v39 - ); - TLOAD(v36, v40); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v41 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v42 = (uint64_t)v21; - 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< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND> - v45 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v23 + v23 * v17 + v30 * v16), v43, v44 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - TLOAD(v41, v45); - 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> - v46 = Tile< - TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v47 = (uint64_t)v20; - TASSIGN(v46, v47); - pto::Shape<1, 1, 1, 1, 256> v48 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v49 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v50 = - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( - v3 + (v23 + v4 * v17 + v30 * v16), v48, v49 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1); - 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>(v18, v11); - uint64_t v52 = (uint64_t)v19; - TASSIGN(v51, v52); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - TCVT(v51, v31, v8, v7); - 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>(v18, v11); - uint64_t v54 = (uint64_t)v19; - TASSIGN(v53, v54); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TCOLEXPANDMUL(v53, v51, v36); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v55 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v56 = (uint64_t)v23; - TASSIGN(v55, v56); - pipe_barrier(PIPE_V); - TCVT(v55, v53, v8, v7); - 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> - v57 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v58 = (uint64_t)v19; - TASSIGN(v57, v58); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID2); - TCVT(v57, v41, v8, v7); - 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> - v59 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v60 = (uint64_t)v19; - TASSIGN(v59, v60); - pipe_barrier(PIPE_V); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID3); - TCOLEXPANDMUL(v59, v57, v46); - 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> - v61 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v62 = (uint64_t)v19; - TASSIGN(v61, v62); - pipe_barrier(PIPE_V); - TCVT(v61, v59, v8, v7); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - pto::Shape<1, 1, 1, 16, 256> v63 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v64 = 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> - v65 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v23 + v23 * v17 + v29 * v16), v63, v64 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(v65, v55); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - pto::Shape<1, 1, 1, 16, 256> v66 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v67 = 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> - v68 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v23 + v23 * v17 + v30 * v16), v66, v67 - ); - pipe_barrier(PIPE_MTE3); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID1); - TSTORE(v68, v61); - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0); - } - set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID1); - wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1); - if ((int64_t)((uint64_t)(v25 / v9) - (uint64_t)((int64_t)(uint64_t)v26 * (uint64_t)v15)) == v16) { - int64_t v69 = (int64_t)((uint64_t)v27 * (uint64_t)v11); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v70 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v71 = (uint64_t)v23; - TASSIGN(v70, v71); - pto::Shape<1, 1, 1, 16, 256> v72 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v73 = 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> - v74 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v2 + (v23 + v23 * v17 + v69 * v16), v72, v73 - ); - TLOAD(v70, v74); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v75 = Tile< - TileType::Vec, float, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v76 = (uint64_t)v19; - TASSIGN(v75, v76); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID4); - TCVT(v75, v70, v8, v7); - set_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); - Tile< - TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v77 = Tile< - TileType::Vec, float, 1, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v16, v11); - uint64_t v78 = (uint64_t)v23; - TASSIGN(v77, v78); - pto::Shape<1, 1, 1, 1, 256> v79 = pto::Shape<1, 1, 1, 1, 256>(); - pto::Stride<5120, 5120, 5120, 5120, 1> v80 = pto::Stride<5120, 5120, 5120, 5120, 1>(); - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND> v81 = - GlobalTensor, pto::Stride<5120, 5120, 5120, 5120, 1>, pto::Layout::ND>( - v3 + (v23 + v4 * v17 + v69 * v16), v79, v80 - ); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID2); - TLOAD(v77, v81); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - 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>(v18, v11); - uint64_t v83 = (uint64_t)v19; - TASSIGN(v82, v83); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID5); - TCOLEXPANDMUL(v82, v75, v77); - Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null> - v84 = Tile< - TileType::Vec, bfloat16_t, 16, 256, BLayout::RowMajor, -1, -1, SLayout::NoneBox, 512, PadValue::Null, - CompactMode::Null>(v18, v11); - uint64_t v85 = (uint64_t)v19; - TASSIGN(v84, v85); - pipe_barrier(PIPE_V); - TCVT(v84, v82, v8, v7); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - pto::Shape<1, 1, 1, 16, 256> v86 = pto::Shape<1, 1, 1, 16, 256>(); - pto::Stride<81920, 81920, 81920, 5120, 1> v87 = 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> - v88 = GlobalTensor< - bfloat16_t, pto::Shape<1, 1, 1, 16, 256>, pto::Stride<81920, 81920, 81920, 5120, 1>, pto::Layout::ND>( - v1 + (v23 + v23 * v17 + v69 * v16), v86, v87 - ); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID2); - TSTORE(v88, v84); - } else { - } - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE2, 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) { - // 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_states_inline146__ssa_v0 - __gm__ Tensor *normed_states_inline146__ssa_v0_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ bfloat16_t *normed_states_inline146__ssa_v0 = - reinterpret_cast<__gm__ bfloat16_t *>(normed_states_inline146__ssa_v0_tensor->buffer.addr) + - normed_states_inline146__ssa_v0_tensor->start_offset; - - // Unpack tensor: cur__iter_v6 - __gm__ Tensor *cur__iter_v6_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ bfloat16_t *cur__iter_v6 = - reinterpret_cast<__gm__ bfloat16_t *>(cur__iter_v6_tensor->buffer.addr) + cur__iter_v6_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; - - // Unpack scalar: i__idx_v0 - union { - uint64_t u64; - int64_t val; - } i__idx_v0_conv; - i__idx_v0_conv.u64 = args[3]; - int64_t i__idx_v0 = i__idx_v0_conv.val; - - // Forward to ptoas-generated function - x_gamma( - normed_states_inline146__ssa_v0, cur__iter_v6, input_rms_weight__ssa_v0, i__idx_v0, __pypto_spmd_block_idx, - __pypto_spmd_block_num - ); -} diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp new file mode 100644 index 0000000000..268481d684 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp index 60acdde710..db08d899aa 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/orchestration/decode_fwd_layers.cpp @@ -52,1082 +52,1926 @@ __attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2Ta 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::BFLOAT16); - TaskOutputTensors alloc_0 = alloc_tensors(cur_ci); - const Tensor &cur = alloc_0.get_ref(0); + 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 0: copy_hidden - L0TaskArgs params_t0; - params_t0.add_output(cur); - params_t0.add_input(ext_hidden_states); - params_t0.add_scalar(cb0); - rt_submit_aiv_task(0, params_t0); + // 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; - for (int64_t i = 0; i < 2; i += 1) { + 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::BFLOAT16); - uint32_t normed_states_inline146_ci_shapes[2] = {16, 5120}; - TensorCreateInfo normed_states_inline146_ci(normed_states_inline146_ci_shapes, 2, DataType::BFLOAT16); - uint32_t out_partial_inline312_ci_shapes[2] = {16, 5120}; - TensorCreateInfo out_partial_inline312_ci(out_partial_inline312_ci_shapes, 2, DataType::BFLOAT16); - uint32_t inv_rms_states_inline160_ci_shapes[2] = {16, 1}; - TensorCreateInfo inv_rms_states_inline160_ci(inv_rms_states_inline160_ci_shapes, 2, DataType::FLOAT32); - uint32_t q_proj_inline129_ci_shapes[2] = {16, 5120}; - TensorCreateInfo q_proj_inline129_ci(q_proj_inline129_ci_shapes, 2, DataType::FLOAT32); - uint32_t k_proj_inline128_ci_shapes[2] = {16, 1024}; - TensorCreateInfo k_proj_inline128_ci(k_proj_inline128_ci_shapes, 2, DataType::FLOAT32); - uint32_t v_proj_inline152_ci_shapes[2] = {16, 1024}; - TensorCreateInfo v_proj_inline152_ci(v_proj_inline152_ci_shapes, 2, DataType::FLOAT32); - uint32_t all_q_padded_inline168_ci_shapes[2] = {2048, 128}; - TensorCreateInfo all_q_padded_inline168_ci(all_q_padded_inline168_ci_shapes, 2, DataType::BFLOAT16); - uint32_t attn_out_inline157_ci_shapes[2] = {16, 5120}; - TensorCreateInfo attn_out_inline157_ci(attn_out_inline157_ci_shapes, 2, DataType::BFLOAT16); - uint32_t all_oi_tmp_inline145_ci_shapes[2] = {88064, 128}; - TensorCreateInfo all_oi_tmp_inline145_ci(all_oi_tmp_inline145_ci_shapes, 2, DataType::FLOAT32); - uint32_t all_cur_mi_inline174_ci_shapes[2] = {88064, 1}; - TensorCreateInfo all_cur_mi_inline174_ci(all_cur_mi_inline174_ci_shapes, 2, DataType::FLOAT32); - uint32_t all_cur_li_inline190_ci_shapes[2] = {88064, 1}; - TensorCreateInfo all_cur_li_inline190_ci(all_cur_li_inline190_ci_shapes, 2, DataType::FLOAT32); - uint32_t fa_work_table_inline126_ci_shapes[2] = {688, 1}; - TensorCreateInfo fa_work_table_inline126_ci(fa_work_table_inline126_ci_shapes, 2, DataType::INT32); - uint32_t fa_total_inline183_ci_shapes[2] = {1, 1}; - TensorCreateInfo fa_total_inline183_ci(fa_total_inline183_ci_shapes, 2, DataType::INT32); - uint32_t q_proj_norm_inline197_ci_shapes[2] = {16, 5120}; - TensorCreateInfo q_proj_norm_inline197_ci(q_proj_norm_inline197_ci_shapes, 2, DataType::FLOAT32); - uint32_t k_proj_norm_inline155_ci_shapes[2] = {16, 1024}; - TensorCreateInfo k_proj_norm_inline155_ci(k_proj_norm_inline155_ci_shapes, 2, DataType::FLOAT32); + 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, normed_states_inline146_ci, out_partial_inline312_ci, inv_rms_states_inline160_ci, - q_proj_inline129_ci, k_proj_inline128_ci, v_proj_inline152_ci, all_q_padded_inline168_ci, - attn_out_inline157_ci, all_oi_tmp_inline145_ci, all_cur_mi_inline174_ci, all_cur_li_inline190_ci, - fa_work_table_inline126_ci, fa_total_inline183_ci, q_proj_norm_inline197_ci, - k_proj_norm_inline155_ci + 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 &normed_states_inline146 = alloc_1.get_ref(1); - const Tensor &out_partial_inline312 = alloc_1.get_ref(2); - const Tensor &inv_rms_states_inline160 = alloc_1.get_ref(3); - const Tensor &q_proj_inline129 = alloc_1.get_ref(4); - const Tensor &k_proj_inline128 = alloc_1.get_ref(5); - const Tensor &v_proj_inline152 = alloc_1.get_ref(6); - const Tensor &all_q_padded_inline168 = alloc_1.get_ref(7); - const Tensor &attn_out_inline157 = alloc_1.get_ref(8); - const Tensor &all_oi_tmp_inline145 = alloc_1.get_ref(9); - const Tensor &all_cur_mi_inline174 = alloc_1.get_ref(10); - const Tensor &all_cur_li_inline190 = alloc_1.get_ref(11); - const Tensor &fa_work_table_inline126 = alloc_1.get_ref(12); - const Tensor &fa_total_inline183 = alloc_1.get_ref(13); - const Tensor &q_proj_norm_inline197 = alloc_1.get_ref(14); - const Tensor &k_proj_norm_inline155 = alloc_1.get_ref(15); - uint32_t q_inv_states_inline99_ci_shapes[2] = {640, 1}; - TensorCreateInfo q_inv_states_inline99_ci(q_inv_states_inline99_ci_shapes, 2, DataType::FLOAT32); - uint32_t k_inv_states_inline90_ci_shapes[2] = {128, 1}; - TensorCreateInfo k_inv_states_inline90_ci(k_inv_states_inline90_ci_shapes, 2, DataType::FLOAT32); - uint32_t down_acc_all_inline251_ci_shapes[2] = {16, 5120}; - TensorCreateInfo down_acc_all_inline251_ci(down_acc_all_inline251_ci_shapes, 2, DataType::FLOAT32); - uint32_t gate_acc_all_inline170_ci_shapes[2] = {16, 17408}; - TensorCreateInfo gate_acc_all_inline170_ci(gate_acc_all_inline170_ci_shapes, 2, DataType::FLOAT32); - uint32_t up_acc_all_inline247_ci_shapes[2] = {16, 17408}; - TensorCreateInfo up_acc_all_inline247_ci(up_acc_all_inline247_ci_shapes, 2, DataType::FLOAT32); - uint32_t gm_pipe_buffer_0_ci_shapes[1] = {static_cast((8192) * (24))}; - TensorCreateInfo gm_pipe_buffer_0_ci( - gm_pipe_buffer_0_ci_shapes, 1, DataType::FLOAT32, /*manual_dep=*/true - ); - uint32_t attn_proj_fp32_inline248_ci_shapes[2] = {16, 5120}; - TensorCreateInfo attn_proj_fp32_inline248_ci(attn_proj_fp32_inline248_ci_shapes, 2, DataType::FLOAT32); - uint32_t post_norm_partial_inline130_ci_shapes[2] = {16, 5120}; - TensorCreateInfo post_norm_partial_inline130_ci( - post_norm_partial_inline130_ci_shapes, 2, DataType::BFLOAT16 - ); - uint32_t mlp_norm_in_inline250_ci_shapes[2] = {16, 5120}; - TensorCreateInfo mlp_norm_in_inline250_ci(mlp_norm_in_inline250_ci_shapes, 2, DataType::BFLOAT16); - uint32_t inv_rms_tile_inline255_ci_shapes[2] = {16, 1}; - TensorCreateInfo inv_rms_tile_inline255_ci(inv_rms_tile_inline255_ci_shapes, 2, DataType::FLOAT32); - uint32_t mlp_tile_inline238_ci_shapes[2] = {16, 17408}; - TensorCreateInfo mlp_tile_inline238_ci(mlp_tile_inline238_ci_shapes, 2, DataType::BFLOAT16); - uint32_t attn_fence_dummy_inline258_ci_shapes[2] = {16, 1}; - TensorCreateInfo attn_fence_dummy_inline258_ci( - attn_fence_dummy_inline258_ci_shapes, 2, DataType::FLOAT32 - ); - TaskOutputTensors alloc_2 = alloc_tensors( - q_inv_states_inline99_ci, k_inv_states_inline90_ci, down_acc_all_inline251_ci, - gate_acc_all_inline170_ci, up_acc_all_inline247_ci, gm_pipe_buffer_0_ci, - attn_proj_fp32_inline248_ci, post_norm_partial_inline130_ci, mlp_norm_in_inline250_ci, - inv_rms_tile_inline255_ci, mlp_tile_inline238_ci, attn_fence_dummy_inline258_ci - ); - const Tensor &q_inv_states_inline99 = alloc_2.get_ref(0); - const Tensor &k_inv_states_inline90 = alloc_2.get_ref(1); - const Tensor &down_acc_all_inline251 = alloc_2.get_ref(2); - const Tensor &gate_acc_all_inline170 = alloc_2.get_ref(3); - const Tensor &up_acc_all_inline247 = alloc_2.get_ref(4); - const Tensor &gm_pipe_buffer_0 = alloc_2.get_ref(5); - const Tensor &attn_proj_fp32_inline248 = alloc_2.get_ref(6); - const Tensor &post_norm_partial_inline130 = alloc_2.get_ref(7); - const Tensor &mlp_norm_in_inline250 = alloc_2.get_ref(8); - const Tensor &inv_rms_tile_inline255 = alloc_2.get_ref(9); - const Tensor &mlp_tile_inline238 = alloc_2.get_ref(10); - const Tensor &attn_fence_dummy_inline258 = alloc_2.get_ref(11); - int64_t layer_hidden_base_inline135 = (static_cast(i) * 5120); - int64_t layer_inter_base_inline231 = (static_cast(i) * 17408); - int64_t num_layers_actual_inline140 = 2; - int64_t t = 1409024; - int64_t layer_cache_rows_inline153 = (t / num_layers_actual_inline140); - int64_t layer_cache_base_inline137 = (static_cast(i) * layer_cache_rows_inline153); - int64_t user_batch_inline123 = 16; - int64_t t__tmp_v2 = 688; - int64_t max_blocks_per_seq_inline165 = (t__tmp_v2 / user_batch_inline123); - uint32_t q_norm_w_inline242_offsets[2] = {static_cast(i), 0}; - uint32_t q_norm_w_inline242_shapes[2] = { - (q_norm_w_inline242_offsets[0] >= ext_q_norm_weight.shapes[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_inline242_offsets[0])), - (q_norm_w_inline242_offsets[1] >= ext_q_norm_weight.shapes[1] ? + 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_inline242_offsets[1])) + std::min(128, ext_q_norm_weight.shapes[1] - q_norm_w_inline124_offsets[1])) }; - Tensor q_norm_w_inline242 = - ext_q_norm_weight.view(q_norm_w_inline242_shapes, q_norm_w_inline242_offsets); - uint32_t k_norm_w_inline138_offsets[2] = {static_cast(i), 0}; - uint32_t k_norm_w_inline138_shapes[2] = { - (k_norm_w_inline138_offsets[0] >= ext_k_norm_weight.shapes[0] ? + 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_inline138_offsets[0])), - (k_norm_w_inline138_offsets[1] >= ext_k_norm_weight.shapes[1] ? + 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_inline138_offsets[1])) + std::min(128, ext_k_norm_weight.shapes[1] - k_norm_w_inline114_offsets[1])) }; - Tensor k_norm_w_inline138 = - ext_k_norm_weight.view(k_norm_w_inline138_shapes, k_norm_w_inline138_offsets); - PTO2TaskId dcr_tids_inline154[5]; - for (int64_t __init_i = 0; __init_i < 5; ++__init_i) - dcr_tids_inline154[__init_i] = PTO2TaskId::invalid(); + 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; - // Spmd x_gamma_spmd: x_gamma - L0TaskArgs params_t1; - params_t1.add_inout(normed_states_inline146); - params_t1.add_input(cur__rv_v7); - params_t1.add_input(ext_input_rms_weight); - params_t1.add_scalar(i); - params_t1.launch_spec.set_block_num(5); - rt_submit_aiv_task(1, params_t1); + // 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 2: rms_recip - L0TaskArgs params_t2; - params_t2.add_input(cur__rv_v7); - params_t2.add_inout(inv_rms_states_inline160); - rt_submit_aiv_task(2, params_t2); + // 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 3: q_seed - L0TaskArgs params_t3; - params_t3.add_inout(q_proj_inline129); - rt_submit_aiv_task(3, params_t3); + // 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_t4; - params_t4.add_inout(q_proj_inline129); - params_t4.add_input(normed_states_inline146); - params_t4.add_input(ext_wq); - params_t4.add_scalar(layer_hidden_base_inline135); - params_t4.launch_spec.set_block_num(40); - rt_submit_aic_task(4, params_t4); + // 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 5: k_seed - L0TaskArgs params_t5; - params_t5.add_inout(k_proj_inline128); - rt_submit_aiv_task(5, params_t5); + // 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; - // Spmd k_proj_spmd: k_proj - L0TaskArgs params_t6; - params_t6.add_inout(k_proj_inline128); - params_t6.add_input(normed_states_inline146); - params_t6.add_input(ext_wk); - params_t6.add_scalar(layer_hidden_base_inline135); - params_t6.launch_spec.set_block_num(8); - rt_submit_aic_task(6, params_t6); + // 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(); - // Task 7: v_seed - L0TaskArgs params_t7; - params_t7.add_inout(v_proj_inline152); - rt_submit_aiv_task(7, params_t7); + // 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_t8; - params_t8.add_inout(v_proj_inline152); - params_t8.add_input(normed_states_inline146); - params_t8.add_input(ext_wv); - params_t8.add_scalar(layer_hidden_base_inline135); - params_t8.launch_spec.set_block_num(8); - rt_submit_aic_task(8, params_t8); + // 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; - // Task 9: fa_work_build - L0TaskArgs params_t9; - params_t9.add_input(ext_seq_lens); - params_t9.add_inout(fa_work_table_inline126); - params_t9.add_inout(fa_total_inline183); - rt_submit_aiv_task(9, params_t9); - uint32_t inv_rms_col_inline89_offsets[2] = {0, 0}; - uint32_t inv_rms_col_inline89_shapes[2] = { - (inv_rms_col_inline89_offsets[0] >= inv_rms_states_inline160.shapes[0] ? - 0u : - std::min(16, inv_rms_states_inline160.shapes[0] - inv_rms_col_inline89_offsets[0])), - (inv_rms_col_inline89_offsets[1] >= inv_rms_states_inline160.shapes[1] ? - 0u : - std::min(1, inv_rms_states_inline160.shapes[1] - inv_rms_col_inline89_offsets[1])) - }; - Tensor inv_rms_col_inline89 = - inv_rms_states_inline160.view(inv_rms_col_inline89_shapes, inv_rms_col_inline89_offsets); + // 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(); - // Task 10: qk_gamma - L0TaskArgs params_t10; - params_t10.add_inout(k_proj_norm_inline155); - params_t10.add_inout(q_proj_norm_inline197); - params_t10.add_input(q_proj_inline129); - params_t10.add_input(inv_rms_col_inline89); - params_t10.add_input(q_norm_w_inline242); - params_t10.add_input(k_proj_inline128); - params_t10.add_input(k_norm_w_inline138); - rt_submit_aiv_task(10, params_t10); + // 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 11: qk_recip - L0TaskArgs params_t11; - params_t11.add_inout(k_inv_states_inline90); - params_t11.add_inout(q_inv_states_inline99); - params_t11.add_input(q_proj_inline129); - params_t11.add_input(inv_rms_col_inline89); - params_t11.add_input(k_proj_inline128); - rt_submit_aiv_task(11, params_t11); + // 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 rope_qkv_spmd: rope_qkv - L0TaskArgs params_t12; - params_t12.add_input(ext_seq_lens); - params_t12.add_input(inv_rms_states_inline160); - params_t12.add_input(ext_slot_mapping); - params_t12.add_input(ext_rope_cos); - params_t12.add_input(ext_rope_sin); - params_t12.add_inout(all_q_padded_inline168); - params_t12.add_inout(ext_k_cache); - params_t12.add_inout(ext_v_cache); - params_t12.add_input(k_inv_states_inline90); - params_t12.add_input(k_proj_norm_inline155); - params_t12.add_input(v_proj_inline152); - params_t12.add_input(q_inv_states_inline99); - params_t12.add_input(q_proj_norm_inline197); - params_t12.add_scalar(layer_cache_base_inline137); - params_t12.launch_spec.set_block_num(16); - rt_submit_aiv_task(12, params_t12); + // 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 13: down_seed - L0TaskArgs params_t13; - params_t13.add_inout(down_acc_all_inline251); - TaskOutputTensors task_13_outs = rt_submit_aiv_task(13, params_t13); - PTO2TaskId seed_tid_inline132 = task_13_outs.task_id(); + // 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 14: gate_seed - L0TaskArgs params_t14; - params_t14.add_inout(gate_acc_all_inline170); - TaskOutputTensors task_14_outs = rt_submit_aiv_task(14, params_t14); - PTO2TaskId gate_seed_tid_inline304 = task_14_outs.task_id(); + // 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 15: up_seed - L0TaskArgs params_t15; - params_t15.add_inout(up_acc_all_inline247); - TaskOutputTensors task_15_outs = rt_submit_aiv_task(15, params_t15); - PTO2TaskId up_seed_tid_inline49 = task_15_outs.task_id(); + // 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; - // Group fa_fused: MixedKernels (AIC + AIV lanes) - L0TaskArgs params_t16; - params_t16.add_input(fa_total_inline183); - params_t16.add_inout(all_cur_li_inline190); - params_t16.add_inout(all_cur_mi_inline174); - params_t16.add_inout(all_oi_tmp_inline145); - params_t16.add_input(fa_work_table_inline126); - params_t16.add_input(ext_seq_lens); - params_t16.add_input(ext_block_table); - params_t16.add_input(all_q_padded_inline168); - params_t16.add_input(ext_k_cache); - params_t16.add_input(ext_v_cache); - params_t16.add_inout(gm_pipe_buffer_0); - params_t16.add_scalar(max_blocks_per_seq_inline165); - params_t16.add_scalar(layer_cache_base_inline137); - MixedKernels mixed_16 = {16, 17, 17}; - params_t16.launch_spec.set_block_num(24); - rt_submit_task(mixed_16, params_t16); + // 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; - // Spmd online_softmax_spmd: online_softmax - L0TaskArgs params_t17; - params_t17.add_inout(attn_out_inline157); - params_t17.add_input(ext_seq_lens); - params_t17.add_input(all_oi_tmp_inline145); - params_t17.add_input(all_cur_mi_inline174); - params_t17.add_input(all_cur_li_inline190); - params_t17.launch_spec.set_block_num(48); - rt_submit_aiv_task(18, params_t17); + // 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 18: out_seed - L0TaskArgs params_t18; - params_t18.add_inout(attn_proj_fp32_inline248); - TaskOutputTensors task_18_outs = rt_submit_aiv_task(19, params_t18); - PTO2TaskId out_seed_tid_inline253 = task_18_outs.task_id(); + // 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; - // Task 19: attn_fence - L0TaskArgs params_t19; - params_t19.add_input(attn_out_inline157); - params_t19.add_inout(attn_fence_dummy_inline258); - TaskOutputTensors task_19_outs = rt_submit_aiv_task(20, params_t19); - PTO2TaskId attn_done_tid_inline181 = task_19_outs.task_id(); - PTO2_SCOPE(PTO2ScopeMode::MANUAL) { - PTO2TaskId silu_tids_inline169[17]; - for (int64_t __init_i = 0; __init_i < 17; ++__init_i) - silu_tids_inline169[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId down_tids_inline261[85]; - for (int64_t __init_i = 0; __init_i < 85; ++__init_i) - down_tids_inline261[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId gate_tids_inline263[85]; - for (int64_t __init_i = 0; __init_i < 85; ++__init_i) - gate_tids_inline263[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId up_tids_inline265[85]; - for (int64_t __init_i = 0; __init_i < 85; ++__init_i) - up_tids_inline265[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId cast_tids_inline267[5]; - for (int64_t __init_i = 0; __init_i < 5; ++__init_i) - cast_tids_inline267[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId out_tids_inline268[50]; - for (int64_t __init_i = 0; __init_i < 50; ++__init_i) - out_tids_inline268[__init_i] = PTO2TaskId::invalid(); - for (int64_t n_out_proj_inline71 = 0; n_out_proj_inline71 < 10; n_out_proj_inline71 += 1) { - int64_t n_op_inline211 = (n_out_proj_inline71 * 512); - for (int64_t k_split_out_inline193 = 0; k_split_out_inline193 < 5; k_split_out_inline193 += 1) { - int64_t k_op_inline151 = (k_split_out_inline193 * 1024); + // 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; - // Task 20: out_proj - L0TaskArgs params_t20; - params_t20.add_input(attn_out_inline157); - params_t20.add_input(ext_wo); - params_t20.add_inout(attn_proj_fp32_inline248); - params_t20.add_scalar(k_op_inline151); - params_t20.add_scalar(layer_hidden_base_inline135); - params_t20.add_scalar(n_op_inline211); - PTO2TaskId params_t20_deps[2]; - uint32_t params_t20_deps_count = 0; - if (out_seed_tid_inline253.is_valid()) - params_t20_deps[params_t20_deps_count++] = out_seed_tid_inline253; - if (attn_done_tid_inline181.is_valid()) - params_t20_deps[params_t20_deps_count++] = attn_done_tid_inline181; - params_t20.set_dependencies(params_t20_deps, params_t20_deps_count); - TaskOutputTensors task_20_outs = rt_submit_aic_task(21, params_t20); - PTO2TaskId out_tid_inline269 = task_20_outs.task_id(); - out_tids_inline268[((n_out_proj_inline71 * 5) + k_split_out_inline193)] = out_tid_inline269; - } + // 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(); } - int64_t k_base_inline177 = 0; - int64_t n_split_base_inline47 = 0; - PTO2TaskId _submit_deps_buf_inline75[10]; - for (int64_t __init_i = 0; __init_i < 10; ++__init_i) - _submit_deps_buf_inline75[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v86 = out_tids_inline268[(n_split_base_inline47 * 5)]; - _submit_deps_buf_inline75[0] = t__tmp_v86; - PTO2TaskId t__tmp_v87 = out_tids_inline268[((n_split_base_inline47 * 5) + 1)]; - _submit_deps_buf_inline75[1] = t__tmp_v87; - PTO2TaskId t__tmp_v88 = out_tids_inline268[((n_split_base_inline47 * 5) + 2)]; - _submit_deps_buf_inline75[2] = t__tmp_v88; - PTO2TaskId t__tmp_v89 = out_tids_inline268[((n_split_base_inline47 * 5) + 3)]; - _submit_deps_buf_inline75[3] = t__tmp_v89; - PTO2TaskId t__tmp_v90 = out_tids_inline268[((n_split_base_inline47 * 5) + 4)]; - _submit_deps_buf_inline75[4] = t__tmp_v90; - PTO2TaskId t__tmp_v91 = out_tids_inline268[((n_split_base_inline47 * 5) + 5)]; - _submit_deps_buf_inline75[5] = t__tmp_v91; - PTO2TaskId t__tmp_v92 = out_tids_inline268[((n_split_base_inline47 * 5) + 6)]; - _submit_deps_buf_inline75[6] = t__tmp_v92; - PTO2TaskId t__tmp_v93 = out_tids_inline268[((n_split_base_inline47 * 5) + 7)]; - _submit_deps_buf_inline75[7] = t__tmp_v93; - PTO2TaskId t__tmp_v94 = out_tids_inline268[((n_split_base_inline47 * 5) + 8)]; - _submit_deps_buf_inline75[8] = t__tmp_v94; - PTO2TaskId t__tmp_v95 = out_tids_inline268[((n_split_base_inline47 * 5) + 9)]; - _submit_deps_buf_inline75[9] = t__tmp_v95; + 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; - // Task 21: residual_rms_cast + // Spmd up_proj_spmd: up_proj L0TaskArgs params_t21; - params_t21.add_inout(mlp_norm_in_inline250); - params_t21.add_inout(post_norm_partial_inline130); - params_t21.add_input(attn_proj_fp32_inline248); - params_t21.add_input(cur__rv_v7); - params_t21.add_input(ext_post_rms_weight); - params_t21.add_scalar(k_base_inline177); - params_t21.add_scalar(i); - PTO2TaskId params_t21_deps[10]; + 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_inline75[0].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[0]; - if (_submit_deps_buf_inline75[1].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[1]; - if (_submit_deps_buf_inline75[2].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[2]; - if (_submit_deps_buf_inline75[3].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[3]; - if (_submit_deps_buf_inline75[4].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[4]; - if (_submit_deps_buf_inline75[5].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[5]; - if (_submit_deps_buf_inline75[6].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[6]; - if (_submit_deps_buf_inline75[7].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[7]; - if (_submit_deps_buf_inline75[8].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[8]; - if (_submit_deps_buf_inline75[9].is_valid()) - params_t21_deps[params_t21_deps_count++] = _submit_deps_buf_inline75[9]; + 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_aiv_task(22, params_t21); - PTO2TaskId cast_tid_k_inline171 = task_21_outs.task_id(); - cast_tids_inline267[0] = cast_tid_k_inline171; - int64_t k_base_inline177__ssa_v1 = 1024; - int64_t n_split_base_inline47__ssa_v1 = 2; - PTO2TaskId _submit_deps_buf_inline75__ssa_v1[10]; - for (int64_t __init_i = 0; __init_i < 10; ++__init_i) - _submit_deps_buf_inline75__ssa_v1[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v100 = out_tids_inline268[(n_split_base_inline47__ssa_v1 * 5)]; - _submit_deps_buf_inline75__ssa_v1[0] = t__tmp_v100; - PTO2TaskId t__tmp_v101 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 1)]; - _submit_deps_buf_inline75__ssa_v1[1] = t__tmp_v101; - PTO2TaskId t__tmp_v102 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 2)]; - _submit_deps_buf_inline75__ssa_v1[2] = t__tmp_v102; - PTO2TaskId t__tmp_v103 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 3)]; - _submit_deps_buf_inline75__ssa_v1[3] = t__tmp_v103; - PTO2TaskId t__tmp_v104 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 4)]; - _submit_deps_buf_inline75__ssa_v1[4] = t__tmp_v104; - PTO2TaskId t__tmp_v105 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 5)]; - _submit_deps_buf_inline75__ssa_v1[5] = t__tmp_v105; - PTO2TaskId t__tmp_v106 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 6)]; - _submit_deps_buf_inline75__ssa_v1[6] = t__tmp_v106; - PTO2TaskId t__tmp_v107 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 7)]; - _submit_deps_buf_inline75__ssa_v1[7] = t__tmp_v107; - PTO2TaskId t__tmp_v108 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 8)]; - _submit_deps_buf_inline75__ssa_v1[8] = t__tmp_v108; - PTO2TaskId t__tmp_v109 = out_tids_inline268[((n_split_base_inline47__ssa_v1 * 5) + 9)]; - _submit_deps_buf_inline75__ssa_v1[9] = t__tmp_v109; + 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; - // Task 22: residual_rms_cast_0 + // Spmd gate_proj_spmd_0: gate_proj_0 L0TaskArgs params_t22; - params_t22.add_inout(mlp_norm_in_inline250); - params_t22.add_inout(post_norm_partial_inline130); - params_t22.add_input(attn_proj_fp32_inline248); - params_t22.add_input(cur__rv_v7); - params_t22.add_input(ext_post_rms_weight); - params_t22.add_scalar(k_base_inline177__ssa_v1); - params_t22.add_scalar(i); - PTO2TaskId params_t22_deps[10]; + 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_inline75__ssa_v1[0].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[0]; - if (_submit_deps_buf_inline75__ssa_v1[1].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[1]; - if (_submit_deps_buf_inline75__ssa_v1[2].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[2]; - if (_submit_deps_buf_inline75__ssa_v1[3].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[3]; - if (_submit_deps_buf_inline75__ssa_v1[4].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[4]; - if (_submit_deps_buf_inline75__ssa_v1[5].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[5]; - if (_submit_deps_buf_inline75__ssa_v1[6].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[6]; - if (_submit_deps_buf_inline75__ssa_v1[7].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[7]; - if (_submit_deps_buf_inline75__ssa_v1[8].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[8]; - if (_submit_deps_buf_inline75__ssa_v1[9].is_valid()) - params_t22_deps[params_t22_deps_count++] = _submit_deps_buf_inline75__ssa_v1[9]; + 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_aiv_task(23, params_t22); - PTO2TaskId cast_tid_k_inline171__ssa_v1 = task_22_outs.task_id(); - cast_tids_inline267[1] = cast_tid_k_inline171__ssa_v1; - int64_t k_base_inline177__ssa_v2 = 2048; - int64_t n_split_base_inline47__ssa_v2 = 4; - PTO2TaskId _submit_deps_buf_inline75__ssa_v2[10]; - for (int64_t __init_i = 0; __init_i < 10; ++__init_i) - _submit_deps_buf_inline75__ssa_v2[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v114 = out_tids_inline268[(n_split_base_inline47__ssa_v2 * 5)]; - _submit_deps_buf_inline75__ssa_v2[0] = t__tmp_v114; - PTO2TaskId t__tmp_v115 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 1)]; - _submit_deps_buf_inline75__ssa_v2[1] = t__tmp_v115; - PTO2TaskId t__tmp_v116 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 2)]; - _submit_deps_buf_inline75__ssa_v2[2] = t__tmp_v116; - PTO2TaskId t__tmp_v117 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 3)]; - _submit_deps_buf_inline75__ssa_v2[3] = t__tmp_v117; - PTO2TaskId t__tmp_v118 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 4)]; - _submit_deps_buf_inline75__ssa_v2[4] = t__tmp_v118; - PTO2TaskId t__tmp_v119 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 5)]; - _submit_deps_buf_inline75__ssa_v2[5] = t__tmp_v119; - PTO2TaskId t__tmp_v120 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 6)]; - _submit_deps_buf_inline75__ssa_v2[6] = t__tmp_v120; - PTO2TaskId t__tmp_v121 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 7)]; - _submit_deps_buf_inline75__ssa_v2[7] = t__tmp_v121; - PTO2TaskId t__tmp_v122 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 8)]; - _submit_deps_buf_inline75__ssa_v2[8] = t__tmp_v122; - PTO2TaskId t__tmp_v123 = out_tids_inline268[((n_split_base_inline47__ssa_v2 * 5) + 9)]; - _submit_deps_buf_inline75__ssa_v2[9] = t__tmp_v123; + 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; - // Task 23: residual_rms_cast_1 + // Spmd up_proj_spmd_0: up_proj_0 L0TaskArgs params_t23; - params_t23.add_inout(mlp_norm_in_inline250); - params_t23.add_inout(post_norm_partial_inline130); - params_t23.add_input(attn_proj_fp32_inline248); - params_t23.add_input(cur__rv_v7); - params_t23.add_input(ext_post_rms_weight); - params_t23.add_scalar(k_base_inline177__ssa_v2); - params_t23.add_scalar(i); - PTO2TaskId params_t23_deps[10]; + 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_inline75__ssa_v2[0].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[0]; - if (_submit_deps_buf_inline75__ssa_v2[1].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[1]; - if (_submit_deps_buf_inline75__ssa_v2[2].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[2]; - if (_submit_deps_buf_inline75__ssa_v2[3].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[3]; - if (_submit_deps_buf_inline75__ssa_v2[4].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[4]; - if (_submit_deps_buf_inline75__ssa_v2[5].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[5]; - if (_submit_deps_buf_inline75__ssa_v2[6].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[6]; - if (_submit_deps_buf_inline75__ssa_v2[7].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[7]; - if (_submit_deps_buf_inline75__ssa_v2[8].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[8]; - if (_submit_deps_buf_inline75__ssa_v2[9].is_valid()) - params_t23_deps[params_t23_deps_count++] = _submit_deps_buf_inline75__ssa_v2[9]; + 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_aiv_task(24, params_t23); - PTO2TaskId cast_tid_k_inline171__ssa_v2 = task_23_outs.task_id(); - cast_tids_inline267[2] = cast_tid_k_inline171__ssa_v2; - int64_t k_base_inline177__ssa_v3 = 3072; - int64_t n_split_base_inline47__ssa_v3 = 6; - PTO2TaskId _submit_deps_buf_inline75__ssa_v3[10]; - for (int64_t __init_i = 0; __init_i < 10; ++__init_i) - _submit_deps_buf_inline75__ssa_v3[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v128 = out_tids_inline268[(n_split_base_inline47__ssa_v3 * 5)]; - _submit_deps_buf_inline75__ssa_v3[0] = t__tmp_v128; - PTO2TaskId t__tmp_v129 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 1)]; - _submit_deps_buf_inline75__ssa_v3[1] = t__tmp_v129; - PTO2TaskId t__tmp_v130 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 2)]; - _submit_deps_buf_inline75__ssa_v3[2] = t__tmp_v130; - PTO2TaskId t__tmp_v131 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 3)]; - _submit_deps_buf_inline75__ssa_v3[3] = t__tmp_v131; - PTO2TaskId t__tmp_v132 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 4)]; - _submit_deps_buf_inline75__ssa_v3[4] = t__tmp_v132; - PTO2TaskId t__tmp_v133 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 5)]; - _submit_deps_buf_inline75__ssa_v3[5] = t__tmp_v133; - PTO2TaskId t__tmp_v134 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 6)]; - _submit_deps_buf_inline75__ssa_v3[6] = t__tmp_v134; - PTO2TaskId t__tmp_v135 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 7)]; - _submit_deps_buf_inline75__ssa_v3[7] = t__tmp_v135; - PTO2TaskId t__tmp_v136 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 8)]; - _submit_deps_buf_inline75__ssa_v3[8] = t__tmp_v136; - PTO2TaskId t__tmp_v137 = out_tids_inline268[((n_split_base_inline47__ssa_v3 * 5) + 9)]; - _submit_deps_buf_inline75__ssa_v3[9] = t__tmp_v137; + 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; - // Task 24: residual_rms_cast_2 + // Spmd gate_proj_spmd_1: gate_proj_1 L0TaskArgs params_t24; - params_t24.add_inout(mlp_norm_in_inline250); - params_t24.add_inout(post_norm_partial_inline130); - params_t24.add_input(attn_proj_fp32_inline248); - params_t24.add_input(cur__rv_v7); - params_t24.add_input(ext_post_rms_weight); - params_t24.add_scalar(k_base_inline177__ssa_v3); - params_t24.add_scalar(i); - PTO2TaskId params_t24_deps[10]; + 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_inline75__ssa_v3[0].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[0]; - if (_submit_deps_buf_inline75__ssa_v3[1].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[1]; - if (_submit_deps_buf_inline75__ssa_v3[2].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[2]; - if (_submit_deps_buf_inline75__ssa_v3[3].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[3]; - if (_submit_deps_buf_inline75__ssa_v3[4].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[4]; - if (_submit_deps_buf_inline75__ssa_v3[5].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[5]; - if (_submit_deps_buf_inline75__ssa_v3[6].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[6]; - if (_submit_deps_buf_inline75__ssa_v3[7].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[7]; - if (_submit_deps_buf_inline75__ssa_v3[8].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[8]; - if (_submit_deps_buf_inline75__ssa_v3[9].is_valid()) - params_t24_deps[params_t24_deps_count++] = _submit_deps_buf_inline75__ssa_v3[9]; + 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_aiv_task(25, params_t24); - PTO2TaskId cast_tid_k_inline171__ssa_v3 = task_24_outs.task_id(); - cast_tids_inline267[3] = cast_tid_k_inline171__ssa_v3; - int64_t k_base_inline177__ssa_v4 = 4096; - int64_t n_split_base_inline47__ssa_v4 = 8; - PTO2TaskId _submit_deps_buf_inline75__ssa_v4[10]; - for (int64_t __init_i = 0; __init_i < 10; ++__init_i) - _submit_deps_buf_inline75__ssa_v4[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v142 = out_tids_inline268[(n_split_base_inline47__ssa_v4 * 5)]; - _submit_deps_buf_inline75__ssa_v4[0] = t__tmp_v142; - PTO2TaskId t__tmp_v143 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 1)]; - _submit_deps_buf_inline75__ssa_v4[1] = t__tmp_v143; - PTO2TaskId t__tmp_v144 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 2)]; - _submit_deps_buf_inline75__ssa_v4[2] = t__tmp_v144; - PTO2TaskId t__tmp_v145 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 3)]; - _submit_deps_buf_inline75__ssa_v4[3] = t__tmp_v145; - PTO2TaskId t__tmp_v146 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 4)]; - _submit_deps_buf_inline75__ssa_v4[4] = t__tmp_v146; - PTO2TaskId t__tmp_v147 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 5)]; - _submit_deps_buf_inline75__ssa_v4[5] = t__tmp_v147; - PTO2TaskId t__tmp_v148 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 6)]; - _submit_deps_buf_inline75__ssa_v4[6] = t__tmp_v148; - PTO2TaskId t__tmp_v149 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 7)]; - _submit_deps_buf_inline75__ssa_v4[7] = t__tmp_v149; - PTO2TaskId t__tmp_v150 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 8)]; - _submit_deps_buf_inline75__ssa_v4[8] = t__tmp_v150; - PTO2TaskId t__tmp_v151 = out_tids_inline268[((n_split_base_inline47__ssa_v4 * 5) + 9)]; - _submit_deps_buf_inline75__ssa_v4[9] = t__tmp_v151; + 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; - // Task 25: residual_rms_cast_3 + // Spmd up_proj_spmd_1: up_proj_1 L0TaskArgs params_t25; - params_t25.add_inout(mlp_norm_in_inline250); - params_t25.add_inout(post_norm_partial_inline130); - params_t25.add_input(attn_proj_fp32_inline248); - params_t25.add_input(cur__rv_v7); - params_t25.add_input(ext_post_rms_weight); - params_t25.add_scalar(k_base_inline177__ssa_v4); - params_t25.add_scalar(i); - PTO2TaskId params_t25_deps[10]; + 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_inline75__ssa_v4[0].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[0]; - if (_submit_deps_buf_inline75__ssa_v4[1].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[1]; - if (_submit_deps_buf_inline75__ssa_v4[2].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[2]; - if (_submit_deps_buf_inline75__ssa_v4[3].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[3]; - if (_submit_deps_buf_inline75__ssa_v4[4].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[4]; - if (_submit_deps_buf_inline75__ssa_v4[5].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[5]; - if (_submit_deps_buf_inline75__ssa_v4[6].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[6]; - if (_submit_deps_buf_inline75__ssa_v4[7].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[7]; - if (_submit_deps_buf_inline75__ssa_v4[8].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[8]; - if (_submit_deps_buf_inline75__ssa_v4[9].is_valid()) - params_t25_deps[params_t25_deps_count++] = _submit_deps_buf_inline75__ssa_v4[9]; + 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_aiv_task(26, params_t25); - PTO2TaskId cast_tid_k_inline171__ssa_v4 = task_25_outs.task_id(); - cast_tids_inline267[4] = cast_tid_k_inline171__ssa_v4; + 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; - // Task 26: post_rms_reduce + // Spmd gate_proj_spmd_2: gate_proj_2 L0TaskArgs params_t26; - params_t26.add_input(attn_proj_fp32_inline248); - params_t26.add_input(cur__rv_v7); - params_t26.add_inout(inv_rms_tile_inline255); - PTO2TaskId params_t26_deps[50]; + 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 (out_tids_inline268[0].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[0]; - if (out_tids_inline268[1].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[1]; - if (out_tids_inline268[2].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[2]; - if (out_tids_inline268[3].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[3]; - if (out_tids_inline268[4].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[4]; - if (out_tids_inline268[5].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[5]; - if (out_tids_inline268[6].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[6]; - if (out_tids_inline268[7].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[7]; - if (out_tids_inline268[8].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[8]; - if (out_tids_inline268[9].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[9]; - if (out_tids_inline268[10].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[10]; - if (out_tids_inline268[11].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[11]; - if (out_tids_inline268[12].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[12]; - if (out_tids_inline268[13].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[13]; - if (out_tids_inline268[14].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[14]; - if (out_tids_inline268[15].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[15]; - if (out_tids_inline268[16].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[16]; - if (out_tids_inline268[17].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[17]; - if (out_tids_inline268[18].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[18]; - if (out_tids_inline268[19].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[19]; - if (out_tids_inline268[20].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[20]; - if (out_tids_inline268[21].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[21]; - if (out_tids_inline268[22].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[22]; - if (out_tids_inline268[23].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[23]; - if (out_tids_inline268[24].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[24]; - if (out_tids_inline268[25].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[25]; - if (out_tids_inline268[26].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[26]; - if (out_tids_inline268[27].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[27]; - if (out_tids_inline268[28].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[28]; - if (out_tids_inline268[29].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[29]; - if (out_tids_inline268[30].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[30]; - if (out_tids_inline268[31].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[31]; - if (out_tids_inline268[32].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[32]; - if (out_tids_inline268[33].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[33]; - if (out_tids_inline268[34].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[34]; - if (out_tids_inline268[35].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[35]; - if (out_tids_inline268[36].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[36]; - if (out_tids_inline268[37].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[37]; - if (out_tids_inline268[38].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[38]; - if (out_tids_inline268[39].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[39]; - if (out_tids_inline268[40].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[40]; - if (out_tids_inline268[41].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[41]; - if (out_tids_inline268[42].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[42]; - if (out_tids_inline268[43].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[43]; - if (out_tids_inline268[44].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[44]; - if (out_tids_inline268[45].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[45]; - if (out_tids_inline268[46].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[46]; - if (out_tids_inline268[47].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[47]; - if (out_tids_inline268[48].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[48]; - if (out_tids_inline268[49].is_valid()) - params_t26_deps[params_t26_deps_count++] = out_tids_inline268[49]; + 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_aiv_task(27, params_t26); - PTO2TaskId reduce_tid_inline127 = task_26_outs.task_id(); - for (int64_t n_out_inline279 = 0; n_out_inline279 < 17; n_out_inline279 += 1) { - int64_t n0_inline142 = (n_out_inline279 * 1024); - for (int64_t k_split_inline303 = 0; k_split_inline303 < 5; k_split_inline303 += 1) { - int64_t k0_inline202 = (k_split_inline303 * 1024); - PTO2TaskId _submit_deps_buf_inline305[2]; - for (int64_t __init_i = 0; __init_i < 2; ++__init_i) - _submit_deps_buf_inline305[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v163 = cast_tids_inline267[k_split_inline303]; - _submit_deps_buf_inline305[0] = t__tmp_v163; - _submit_deps_buf_inline305[1] = gate_seed_tid_inline304; + 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; - // Task 27: gate_proj - L0TaskArgs params_t27; - params_t27.add_input(mlp_norm_in_inline250); - params_t27.add_input(ext_w_gate); - params_t27.add_inout(gate_acc_all_inline170); - params_t27.add_scalar(k0_inline202); - params_t27.add_scalar(layer_hidden_base_inline135); - params_t27.add_scalar(n0_inline142); - PTO2TaskId params_t27_deps[2]; - uint32_t params_t27_deps_count = 0; - if (_submit_deps_buf_inline305[0].is_valid()) - params_t27_deps[params_t27_deps_count++] = _submit_deps_buf_inline305[0]; - if (_submit_deps_buf_inline305[1].is_valid()) - params_t27_deps[params_t27_deps_count++] = _submit_deps_buf_inline305[1]; - params_t27.set_dependencies(params_t27_deps, params_t27_deps_count); - TaskOutputTensors task_27_outs = rt_submit_aic_task(28, params_t27); - PTO2TaskId gate_tid_inline121 = task_27_outs.task_id(); - gate_tids_inline263[((n_out_inline279 * 5) + k_split_inline303)] = gate_tid_inline121; - PTO2TaskId _submit_deps_buf_inline178[2]; - for (int64_t __init_i = 0; __init_i < 2; ++__init_i) - _submit_deps_buf_inline178[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v164 = cast_tids_inline267[k_split_inline303]; - _submit_deps_buf_inline178[0] = t__tmp_v164; - _submit_deps_buf_inline178[1] = up_seed_tid_inline49; + // 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; - // Task 28: up_proj - L0TaskArgs params_t28; - params_t28.add_input(mlp_norm_in_inline250); - params_t28.add_input(ext_w_up); - params_t28.add_inout(up_acc_all_inline247); - params_t28.add_scalar(k0_inline202); - params_t28.add_scalar(layer_hidden_base_inline135); - params_t28.add_scalar(n0_inline142); - PTO2TaskId params_t28_deps[2]; - uint32_t params_t28_deps_count = 0; - if (_submit_deps_buf_inline178[0].is_valid()) - params_t28_deps[params_t28_deps_count++] = _submit_deps_buf_inline178[0]; - if (_submit_deps_buf_inline178[1].is_valid()) - params_t28_deps[params_t28_deps_count++] = _submit_deps_buf_inline178[1]; - params_t28.set_dependencies(params_t28_deps, params_t28_deps_count); - TaskOutputTensors task_28_outs = rt_submit_aic_task(29, params_t28); - PTO2TaskId up_tid_inline266 = task_28_outs.task_id(); - up_tids_inline265[((n_out_inline279 * 5) + k_split_inline303)] = up_tid_inline266; - } + // 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(); } - for (int64_t n_out_inline297 = 0; n_out_inline297 < 17; n_out_inline297 += 1) { - int64_t n0_inline142__ssa_v7 = (n_out_inline297 * 1024); - PTO2TaskId _submit_deps_buf_inline131[11]; - for (int64_t __init_i = 0; __init_i < 11; ++__init_i) - _submit_deps_buf_inline131[__init_i] = PTO2TaskId::invalid(); - _submit_deps_buf_inline131[0] = reduce_tid_inline127; - PTO2TaskId t__tmp_v165 = gate_tids_inline263[(n_out_inline297 * 5)]; - _submit_deps_buf_inline131[1] = t__tmp_v165; - PTO2TaskId t__tmp_v166 = gate_tids_inline263[((n_out_inline297 * 5) + 1)]; - _submit_deps_buf_inline131[2] = t__tmp_v166; - PTO2TaskId t__tmp_v167 = gate_tids_inline263[((n_out_inline297 * 5) + 2)]; - _submit_deps_buf_inline131[3] = t__tmp_v167; - PTO2TaskId t__tmp_v168 = gate_tids_inline263[((n_out_inline297 * 5) + 3)]; - _submit_deps_buf_inline131[4] = t__tmp_v168; - PTO2TaskId t__tmp_v169 = gate_tids_inline263[((n_out_inline297 * 5) + 4)]; - _submit_deps_buf_inline131[5] = t__tmp_v169; - PTO2TaskId t__tmp_v170 = up_tids_inline265[(n_out_inline297 * 5)]; - _submit_deps_buf_inline131[6] = t__tmp_v170; - PTO2TaskId t__tmp_v171 = up_tids_inline265[((n_out_inline297 * 5) + 1)]; - _submit_deps_buf_inline131[7] = t__tmp_v171; - PTO2TaskId t__tmp_v172 = up_tids_inline265[((n_out_inline297 * 5) + 2)]; - _submit_deps_buf_inline131[8] = t__tmp_v172; - PTO2TaskId t__tmp_v173 = up_tids_inline265[((n_out_inline297 * 5) + 3)]; - _submit_deps_buf_inline131[9] = t__tmp_v173; - PTO2TaskId t__tmp_v174 = up_tids_inline265[((n_out_inline297 * 5) + 4)]; - _submit_deps_buf_inline131[10] = t__tmp_v174; + 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; - // Task 29: silu - L0TaskArgs params_t29; - params_t29.add_input(inv_rms_tile_inline255); - params_t29.add_inout(mlp_tile_inline238); - params_t29.add_input(gate_acc_all_inline170); - params_t29.add_input(up_acc_all_inline247); - params_t29.add_scalar(n0_inline142__ssa_v7); - PTO2TaskId params_t29_deps[11]; - uint32_t params_t29_deps_count = 0; - if (_submit_deps_buf_inline131[0].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[0]; - if (_submit_deps_buf_inline131[1].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[1]; - if (_submit_deps_buf_inline131[2].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[2]; - if (_submit_deps_buf_inline131[3].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[3]; - if (_submit_deps_buf_inline131[4].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[4]; - if (_submit_deps_buf_inline131[5].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[5]; - if (_submit_deps_buf_inline131[6].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[6]; - if (_submit_deps_buf_inline131[7].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[7]; - if (_submit_deps_buf_inline131[8].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[8]; - if (_submit_deps_buf_inline131[9].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[9]; - if (_submit_deps_buf_inline131[10].is_valid()) - params_t29_deps[params_t29_deps_count++] = _submit_deps_buf_inline131[10]; - params_t29.set_dependencies(params_t29_deps, params_t29_deps_count); - TaskOutputTensors task_29_outs = rt_submit_aiv_task(30, params_t29); - PTO2TaskId silu_tid_inline88 = task_29_outs.task_id(); - silu_tids_inline169[n_out_inline297] = silu_tid_inline88; + // 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(); } - for (int64_t n_out_inline25 = 0; n_out_inline25 < 5; n_out_inline25 += 1) { - int64_t n0_inline142__ssa_v8 = (n_out_inline25 * 1024); - for (int64_t k_split_inline122 = 0; k_split_inline122 < 17; k_split_inline122 += 1) { - int64_t k0_inline202__ssa_v11 = (k_split_inline122 * 1024); - PTO2TaskId _submit_deps_buf_inline24[2]; - for (int64_t __init_i = 0; __init_i < 2; ++__init_i) - _submit_deps_buf_inline24[__init_i] = PTO2TaskId::invalid(); - _submit_deps_buf_inline24[0] = seed_tid_inline132; - PTO2TaskId t__tmp_v180 = silu_tids_inline169[k_split_inline122]; - _submit_deps_buf_inline24[1] = t__tmp_v180; + 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: down_proj + // Task 30: gate_proj_4 L0TaskArgs params_t30; - params_t30.add_input(mlp_tile_inline238); - params_t30.add_input(ext_w_down); - params_t30.add_inout(down_acc_all_inline251); - params_t30.add_scalar(k0_inline202__ssa_v11); - params_t30.add_scalar(layer_inter_base_inline231); - params_t30.add_scalar(n0_inline142__ssa_v8); - PTO2TaskId params_t30_deps[2]; + 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_inline24[0].is_valid()) - params_t30_deps[params_t30_deps_count++] = _submit_deps_buf_inline24[0]; - if (_submit_deps_buf_inline24[1].is_valid()) - params_t30_deps[params_t30_deps_count++] = _submit_deps_buf_inline24[1]; + 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 down_tid_inline22 = task_30_outs.task_id(); - down_tids_inline261[((n_out_inline25 * 17) + k_split_inline122)] = down_tid_inline22; - } - } - for (int64_t n_out_inline159 = 0; n_out_inline159 < 5; n_out_inline159 += 1) { - int64_t n0_inline142__ssa_v9 = (n_out_inline159 * 1024); - PTO2TaskId _submit_deps_buf_inline18[17]; - for (int64_t __init_i = 0; __init_i < 17; ++__init_i) - _submit_deps_buf_inline18[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v181 = down_tids_inline261[(n_out_inline159 * 17)]; - _submit_deps_buf_inline18[0] = t__tmp_v181; - PTO2TaskId t__tmp_v182 = down_tids_inline261[((n_out_inline159 * 17) + 1)]; - _submit_deps_buf_inline18[1] = t__tmp_v182; - PTO2TaskId t__tmp_v183 = down_tids_inline261[((n_out_inline159 * 17) + 2)]; - _submit_deps_buf_inline18[2] = t__tmp_v183; - PTO2TaskId t__tmp_v184 = down_tids_inline261[((n_out_inline159 * 17) + 3)]; - _submit_deps_buf_inline18[3] = t__tmp_v184; - PTO2TaskId t__tmp_v185 = down_tids_inline261[((n_out_inline159 * 17) + 4)]; - _submit_deps_buf_inline18[4] = t__tmp_v185; - PTO2TaskId t__tmp_v186 = down_tids_inline261[((n_out_inline159 * 17) + 5)]; - _submit_deps_buf_inline18[5] = t__tmp_v186; - PTO2TaskId t__tmp_v187 = down_tids_inline261[((n_out_inline159 * 17) + 6)]; - _submit_deps_buf_inline18[6] = t__tmp_v187; - PTO2TaskId t__tmp_v188 = down_tids_inline261[((n_out_inline159 * 17) + 7)]; - _submit_deps_buf_inline18[7] = t__tmp_v188; - PTO2TaskId t__tmp_v189 = down_tids_inline261[((n_out_inline159 * 17) + 8)]; - _submit_deps_buf_inline18[8] = t__tmp_v189; - PTO2TaskId t__tmp_v190 = down_tids_inline261[((n_out_inline159 * 17) + 9)]; - _submit_deps_buf_inline18[9] = t__tmp_v190; - PTO2TaskId t__tmp_v191 = down_tids_inline261[((n_out_inline159 * 17) + 10)]; - _submit_deps_buf_inline18[10] = t__tmp_v191; - PTO2TaskId t__tmp_v192 = down_tids_inline261[((n_out_inline159 * 17) + 11)]; - _submit_deps_buf_inline18[11] = t__tmp_v192; - PTO2TaskId t__tmp_v193 = down_tids_inline261[((n_out_inline159 * 17) + 12)]; - _submit_deps_buf_inline18[12] = t__tmp_v193; - PTO2TaskId t__tmp_v194 = down_tids_inline261[((n_out_inline159 * 17) + 13)]; - _submit_deps_buf_inline18[13] = t__tmp_v194; - PTO2TaskId t__tmp_v195 = down_tids_inline261[((n_out_inline159 * 17) + 14)]; - _submit_deps_buf_inline18[14] = t__tmp_v195; - PTO2TaskId t__tmp_v196 = down_tids_inline261[((n_out_inline159 * 17) + 15)]; - _submit_deps_buf_inline18[15] = t__tmp_v196; - PTO2TaskId t__tmp_v197 = down_tids_inline261[((n_out_inline159 * 17) + 16)]; - _submit_deps_buf_inline18[16] = t__tmp_v197; + 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: down_cast_residual - L0TaskArgs params_t31; - params_t31.add_input(post_norm_partial_inline130); - params_t31.add_input(down_acc_all_inline251); - params_t31.add_inout(out_partial_inline312); - params_t31.add_scalar(n0_inline142__ssa_v9); - PTO2TaskId params_t31_deps[17]; - uint32_t params_t31_deps_count = 0; - if (_submit_deps_buf_inline18[0].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[0]; - if (_submit_deps_buf_inline18[1].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[1]; - if (_submit_deps_buf_inline18[2].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[2]; - if (_submit_deps_buf_inline18[3].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[3]; - if (_submit_deps_buf_inline18[4].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[4]; - if (_submit_deps_buf_inline18[5].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[5]; - if (_submit_deps_buf_inline18[6].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[6]; - if (_submit_deps_buf_inline18[7].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[7]; - if (_submit_deps_buf_inline18[8].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[8]; - if (_submit_deps_buf_inline18[9].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[9]; - if (_submit_deps_buf_inline18[10].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[10]; - if (_submit_deps_buf_inline18[11].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[11]; - if (_submit_deps_buf_inline18[12].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[12]; - if (_submit_deps_buf_inline18[13].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[13]; - if (_submit_deps_buf_inline18[14].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[14]; - if (_submit_deps_buf_inline18[15].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[15]; - if (_submit_deps_buf_inline18[16].is_valid()) - params_t31_deps[params_t31_deps_count++] = _submit_deps_buf_inline18[16]; - params_t31.set_dependencies(params_t31_deps, params_t31_deps_count); - TaskOutputTensors task_31_outs = rt_submit_aiv_task(32, params_t31); - PTO2TaskId dcr_tid_inline8 = task_31_outs.task_id(); - dcr_tids_inline154[n_out_inline159] = dcr_tid_inline8; + // 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 oc_b_inline93 = 0; oc_b_inline93 < 16; oc_b_inline93 += 16) { - PTO2_SCOPE() { - PTO2TaskId _submit_deps_buf_inline4[5]; - for (int64_t __init_i = 0; __init_i < 5; ++__init_i) - _submit_deps_buf_inline4[__init_i] = PTO2TaskId::invalid(); - PTO2TaskId t__tmp_v199 = dcr_tids_inline154[0]; - _submit_deps_buf_inline4[0] = t__tmp_v199; - PTO2TaskId t__tmp_v200 = dcr_tids_inline154[1]; - _submit_deps_buf_inline4[1] = t__tmp_v200; - PTO2TaskId t__tmp_v201 = dcr_tids_inline154[2]; - _submit_deps_buf_inline4[2] = t__tmp_v201; - PTO2TaskId t__tmp_v202 = dcr_tids_inline154[3]; - _submit_deps_buf_inline4[3] = t__tmp_v202; - PTO2TaskId t__tmp_v203 = dcr_tids_inline154[4]; - _submit_deps_buf_inline4[4] = t__tmp_v203; + 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: out_consolidate + // Task 32: silu L0TaskArgs params_t32; - params_t32.add_inout(next_hidden); - params_t32.add_input(out_partial_inline312); - params_t32.add_scalar(oc_b_inline93); - PTO2TaskId params_t32_deps[5]; + 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_inline4[0].is_valid()) - params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline4[0]; - if (_submit_deps_buf_inline4[1].is_valid()) - params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline4[1]; - if (_submit_deps_buf_inline4[2].is_valid()) - params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline4[2]; - if (_submit_deps_buf_inline4[3].is_valid()) - params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline4[3]; - if (_submit_deps_buf_inline4[4].is_valid()) - params_t32_deps[params_t32_deps_count++] = _submit_deps_buf_inline4[4]; + 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 33: copy_out - L0TaskArgs params_t33; - params_t33.add_output(ext_out); - params_t33.add_input(cur__rv_v7); - params_t33.add_scalar(ob0); - rt_submit_aiv_task(34, params_t33); + // 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); } } } diff --git a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp new file mode 100644 index 0000000000..cadbef80ef --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp new file mode 100644 index 0000000000..c087c9f179 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/generated/kernel_tiling/kernel_tiling.h b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp new file mode 100644 index 0000000000..4381fe0660 --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h new file mode 100644 index 0000000000..b21f1991ca --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/rope_qkv_generated.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp new file mode 100644 index 0000000000..3c48e349fa --- /dev/null +++ b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/qwen_fai_runtime_tiler.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/arch.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/cross_core_sync.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/local_tensor_buffer.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/resource.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/base_defs.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/coord.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/alignment.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/dependent_false.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/macros.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/CombineScale.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/block_epilogue.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/dispatch_policy.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_cast.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_copy.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/tile_common/tile_swizzle.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_pv.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/block/block_mmad_qk.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/dispatch_policy.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/gemm_type.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/helper.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_copy.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/tile_common/tile_mmad.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm_coord.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/layout.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/matrix.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/vector.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/matrix_coord.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/flash_attention_regular.h b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/kernel_common.hpp b/examples/a2a3/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/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/tensormap_and_ringbuffer/qwen3_14b_decode/test_qwen3_14b_decode.py b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/test_qwen3_14b_decode.py index 6c6aaa2bd3..a90cf9068f 100644 --- a/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/test_qwen3_14b_decode.py +++ b/examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/test_qwen3_14b_decode.py @@ -7,20 +7,26 @@ # 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 2-layer decode (load-balanced fused attention) — SceneTestCase. +"""Qwen3-14B 40-layer decode (CANN fused-attention) — SceneTestCase. -Self-contained port of pypto-lib ``models/qwen3/14b/decode_layer.py`` entry -``decode_fwd_layers`` with ``_CHUNK_NLAYERS == 2``: a fused chunk of two -consecutive Qwen3-14B decode layers (hidden -> hidden, no LM head), with FP32 -inter-layer residual carry. The 36 C++ sources under ``kernels/`` (orchestration -+ 35 incores: 8 AIC + 27 AIV) and ``simpler_setup/goldens/qwen3_14b_decode.py`` -are the pypto codegen for that entry, harvested verbatim so simpler developers -run it directly — no descent through pypto-lib / the JIT. +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 (the ~3500-token -prompt). Weights + paged KV pool are stacked x2 (one slice per layer; the two -layers reuse layer-0 weights, per the lib const-layer-0 stacked-fwd reference). +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 @@ -33,11 +39,48 @@ 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 fused 2-layer decode against a torch reference. + +# Validates the full 40-layer fused decode against a torch reference. @scene_test(level=2, runtime="tensormap_and_ringbuffer") class TestQwen314BDecode(SceneTestCase): - """Two-layer Qwen3-14B decode (decode_fwd_layers N=2) against a torch reference.""" + """Qwen3-14B decode, all 40 layers in one dispatch, against a torch reference.""" RTOL = 5e-2 ATOL = 1e-1 @@ -46,6 +89,10 @@ class TestQwen314BDecode(SceneTestCase): "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 @@ -69,250 +116,304 @@ class TestQwen314BDecode(SceneTestCase): D.OUT, # 19 out ], }, - # 35 incores (func_id 0..34), transcribed from the pypto codegen - # kernel_config.py for decode_fwd_layers (N=2). fa_fused is the - # codegen-split mixed kernel (fa_fused_aic + fa_fused_aiv). + # 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": 1, - "name": "x_gamma", - "source": "kernels/aiv/x_gamma.cpp", + "func_id": 2, + "name": "x_gamma0", + "source": "kernels/aiv/x_gamma0.cpp", "core_type": "aiv", - "signature": [D.INOUT, D.IN, D.IN], + "signature": [D.OUT, D.IN, D.IN], }, { - "func_id": 2, + "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": 3, + "func_id": 5, "name": "q_seed", "source": "kernels/aiv/q_seed.cpp", "core_type": "aiv", "signature": [D.INOUT], }, { - "func_id": 4, + "func_id": 6, "name": "q_proj", "source": "kernels/aic/q_proj.cpp", "core_type": "aic", "signature": [D.INOUT, D.IN, D.IN], }, { - "func_id": 5, - "name": "k_seed", - "source": "kernels/aiv/k_seed.cpp", + "func_id": 7, + "name": "kv_seed", + "source": "kernels/aiv/kv_seed.cpp", "core_type": "aiv", - "signature": [D.INOUT], + "signature": [D.INOUT, D.INOUT], }, { - "func_id": 6, + "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": 7, - "name": "v_seed", - "source": "kernels/aiv/v_seed.cpp", - "core_type": "aiv", - "signature": [D.INOUT], - }, - { - "func_id": 8, + "func_id": 10, "name": "v_proj", "source": "kernels/aic/v_proj.cpp", "core_type": "aic", "signature": [D.INOUT, D.IN, D.IN], }, - { - "func_id": 9, - "name": "fa_work_build", - "source": "kernels/aiv/fa_work_build.cpp", - "core_type": "aiv", - "signature": [D.IN, D.INOUT, D.INOUT], - }, - { - "func_id": 10, - "name": "qk_gamma", - "source": "kernels/aiv/qk_gamma.cpp", - "core_type": "aiv", - "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN, D.IN, D.IN], - }, { "func_id": 11, - "name": "qk_recip", - "source": "kernels/aiv/qk_recip.cpp", - "core_type": "aiv", - "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], + "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": "rope_qkv", - "source": "kernels/aiv/rope_qkv.cpp", + "name": "paged_attention_rope_cce_aiv", + "source": "kernels/vendor/paged_attention_cce/attention_rope/entry.cpp", "core_type": "aiv", - "signature": [D.IN, D.IN, D.IN, D.IN, D.IN, D.INOUT, D.INOUT, D.INOUT, D.IN, D.IN, D.IN, D.IN, D.IN], + "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": "down_seed", - "source": "kernels/aiv/down_seed.cpp", - "core_type": "aiv", - "signature": [D.INOUT], + "name": "out_proj", + "source": "kernels/aic/out_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], }, { "func_id": 14, - "name": "gate_seed", - "source": "kernels/aiv/gate_seed.cpp", - "core_type": "aiv", - "signature": [D.INOUT], + "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": "up_seed", - "source": "kernels/aiv/up_seed.cpp", + "name": "residual_rms_cast", + "source": "kernels/aiv/residual_rms_cast.cpp", "core_type": "aiv", - "signature": [D.INOUT], + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], }, { "func_id": 16, - "name": "fa_fused_aic", - "source": "kernels/aic/fa_fused_aic.cpp", - "core_type": "aic", - "signature": [D.IN, D.INOUT, D.INOUT, D.INOUT, D.IN, D.IN, D.IN, D.IN, D.IN, D.IN, D.INOUT], + "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": "fa_fused_aiv", - "source": "kernels/aiv/fa_fused_aiv.cpp", + "name": "residual_rms_cast_1", + "source": "kernels/aiv/residual_rms_cast_1.cpp", "core_type": "aiv", - "signature": [D.IN, D.INOUT, D.INOUT, D.INOUT, D.IN, D.IN, D.IN, D.IN, D.IN, D.IN, D.INOUT], + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], }, { "func_id": 18, - "name": "online_softmax", - "source": "kernels/aiv/online_softmax.cpp", + "name": "residual_rms_cast_2", + "source": "kernels/aiv/residual_rms_cast_2.cpp", "core_type": "aiv", - "signature": [D.INOUT, D.IN, D.IN, D.IN, D.IN], + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], }, { "func_id": 19, - "name": "out_seed", - "source": "kernels/aiv/out_seed.cpp", + "name": "residual_rms_cast_3", + "source": "kernels/aiv/residual_rms_cast_3.cpp", "core_type": "aiv", - "signature": [D.INOUT], + "signature": [D.INOUT, D.INOUT, D.IN, D.IN, D.IN], }, { "func_id": 20, - "name": "attn_fence", - "source": "kernels/aiv/attn_fence.cpp", + "name": "post_rms_reduce", + "source": "kernels/aiv/post_rms_reduce.cpp", "core_type": "aiv", - "signature": [D.IN, D.INOUT], + "signature": [D.IN, D.IN, D.INOUT], }, { "func_id": 21, - "name": "out_proj", - "source": "kernels/aic/out_proj.cpp", + "name": "gate_proj", + "source": "kernels/aic/gate_proj.cpp", "core_type": "aic", "signature": [D.IN, D.IN, D.INOUT], }, { "func_id": 22, - "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], + "name": "up_proj", + "source": "kernels/aic/up_proj.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.INOUT], }, { "func_id": 23, - "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], + "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": "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], + "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": "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], + "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": "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], + "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": "post_rms_reduce", - "source": "kernels/aiv/post_rms_reduce.cpp", - "core_type": "aiv", + "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": "gate_proj", - "source": "kernels/aic/gate_proj.cpp", + "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": "up_proj", - "source": "kernels/aic/up_proj.cpp", + "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": "silu", - "source": "kernels/aiv/silu.cpp", - "core_type": "aiv", - "signature": [D.IN, D.INOUT, D.IN, D.IN], + "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": "down_proj", - "source": "kernels/aic/down_proj.cpp", + "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": "down_cast_residual", - "source": "kernels/aiv/down_cast_residual.cpp", - "core_type": "aiv", + "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": "out_consolidate", - "source": "kernels/aiv/out_consolidate.cpp", + "name": "silu", + "source": "kernels/aiv/silu.cpp", "core_type": "aiv", - "signature": [D.INOUT, D.IN], + "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", diff --git a/simpler_setup/goldens/qwen3_14b_decode.py b/simpler_setup/goldens/qwen3_14b_decode.py index 08b9cd18d2..b56c94a57a 100644 --- a/simpler_setup/goldens/qwen3_14b_decode.py +++ b/simpler_setup/goldens/qwen3_14b_decode.py @@ -6,24 +6,28 @@ # 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. # ----------------------------------------------------------------------------------------------------------- -"""Golden reference + fixture for the Qwen3-14B 2-layer decode SceneTestCase. +"""Golden reference + fixture for the Qwen3-14B 40-layer decode SceneTestCase. -Ported from pypto-lib ``models/qwen3/14b/decode_layer.py`` (entry -``decode_fwd_layers`` with ``_CHUNK_NLAYERS == 2``): a fused chunk of two -consecutive Qwen3-14B decode layers, hidden -> hidden, no LM head. Weights and -the paged KV pool are STACKED along dim 0 (one slice per layer); the two layers -use the same (replicated) weights — matching the lib's ``--validate-fwd`` -const-layer-0 stack — but each reads/writes its own KV pool. The inter-layer -hidden is carried in FP32 (no per-layer bf16 round); the only chunk-boundary -casts are ``copy_hidden`` (bf16 input embed) and ``copy_out`` (FP32->bf16). +Ported from pypto-lib ``models/qwen3/14b/decode_fwd.py`` (entry +``decode_fwd_layers`` with ``_CHUNK_NLAYERS == 40``): the full Qwen3-14B decode +stack as one fused dispatch, hidden -> hidden, no LM head. Weights and the paged +KV pool are STACKED along dim 0 (one slice per layer); every layer uses the same +(replicated) weights — matching the lib's ``--validate-fwd`` const-layer-0 stack +— but each reads/writes its own KV pool. The inter-layer hidden is carried in +FP32 (no per-layer bf16 round); the only chunk-boundary casts are +``copy_hidden`` (bf16 input embed) and ``copy_out`` (FP32->bf16). Parameter regime matches ``stress_profile.py`` (the vLLM serving stress run): BATCH=16 (CONCURRENCY, aligned with decode kernel BATCH=16), MAX_SEQ=5500 (= max_model_len), and a fixed decode sequence length of 3500 (the ~3500-token prompt). ``MAX_SEQ`` MUST match the harvested kernels' codegen-time value. -KV-cache layout per layer pool: row = (phys_page * NUM_KV_HEADS + kv_head) * -BLOCK_SIZE + pos_in_block; layer L's pool starts at L * CACHE_ROWS. +KV-cache layout is vLLM's paged **BSND**: each physical page holds +``[BLOCK_SIZE, KV_HIDDEN]`` laid out as ``[page, token, kv_head, dim]``, so a +layer pool reads as ``[NUM_PAGES * BLOCK_SIZE, KV_HIDDEN]`` — row +``phys_page * BLOCK_SIZE + pos_in_block``, columns +``kv_head * HEAD_DIM``. ``slot_mapping[b]`` is already that row index. Layer L's +pool starts at ``L * CACHE_ROWS`` in the flat ``[rows, HEAD_DIM]`` storage. Weight matrices are ``[in_features, out_features]`` -> ``y = x @ w``. """ @@ -49,7 +53,7 @@ ROPE_THETA = 1.0e4 # ── Chunk / paging (must match the harvested kernels) ── -N_LAYERS = 2 +N_LAYERS = 40 # Qwen3-14B num_layers; the whole model in one fused dispatch MAX_SEQ = 5500 # = stress_profile max_model_len; codegen-time KV-pool / RoPE sizing SEQ_TILE = 128 BLOCK_SIZE = SEQ_TILE @@ -112,12 +116,15 @@ def _paged_block_table_slot_mapping(seq_lens: torch.Tensor) -> tuple[torch.Tenso def generate_inputs(seed: int = 1234, seq_len: int = DEFAULT_SEQ_LEN) -> TaskArgsBuilder: - """Deterministic fixture for decode_fwd_layers (N=2), stacked x2 along dim 0. + """Deterministic fixture for decode_fwd_layers (N=40), stacked x40 along dim 0. Every lane uses sequence length ``seq_len`` (default 3500, the stress prompt). - Per-layer weights are replicated (stack0) so layer 1 reuses layer 0's weights, - matching the lib's const-layer-0 stacked-fwd reference; each layer still has - its own KV pool. + Per-layer weights are replicated (stack0) so every layer reuses layer 0's + weights, matching the lib's const-layer-0 stacked-fwd reference; each layer + still has its own KV pool. + + The stacks are the bulk of the fixture: ~24.6 GiB of weights plus ~13.4 GiB + of paged KV at this regime. """ if not (1 <= seq_len <= MAX_SEQ): raise ValueError(f"seq_len must be in [1, {MAX_SEQ}], got {seq_len}") @@ -187,8 +194,9 @@ def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor: slot_mapping = args.slot_mapping rope_cos = args.rope_cos.float() rope_sin = args.rope_sin.float() - k_cache = args.k_cache[cb : cb + CACHE_ROWS].float() # this layer's pool - v_cache = args.v_cache[cb : cb + CACHE_ROWS].float() + # This layer's pool, read in its native BSND shape: [page * BLOCK_SIZE + pos, kv_head * HEAD_DIM]. + k_cache = args.k_cache[cb : cb + CACHE_ROWS].float().view(NUM_PAGES * BLOCK_SIZE, KV_HIDDEN) + v_cache = args.v_cache[cb : cb + CACHE_ROWS].float().view(NUM_PAGES * BLOCK_SIZE, KV_HIDDEN) inv_rms = _rmsnorm_inv(x) normed = _bf16(x * irw) @@ -216,11 +224,12 @@ def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor: v_lane = torch.empty(slen, HEAD_DIM) for sb in range(n_blocks): pbid = int(block_table[b * MAX_BLOCKS_PER_SEQ + sb].item()) - row = (pbid * NUM_KV_HEADS + kvh) * BLOCK_SIZE + row = pbid * BLOCK_SIZE + col = kvh * HEAD_DIM lo = sb * BLOCK_SIZE blk = min(BLOCK_SIZE, slen - lo) - k_lane[lo : lo + blk] = k_cache[row : row + blk] - v_lane[lo : lo + blk] = v_cache[row : row + blk] + k_lane[lo : lo + blk] = k_cache[row : row + blk, col : col + HEAD_DIM] + v_lane[lo : lo + blk] = v_cache[row : row + blk, col : col + HEAD_DIM] k_lane[p] = k_cur[kvh] v_lane[p] = v_cur[kvh] cur_k[(b, kvh)] = k_cur[kvh] @@ -243,12 +252,13 @@ def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor: mlp = _bf16(sg * torch.sigmoid(sg) * su) down = mlp @ wd - # Write the current token into THIS layer's pool (INOUT compare). + # Write the current token into THIS layer's pool (INOUT compare). slot_mapping + # is already the BSND row (phys_page * BLOCK_SIZE + pos_in_block); within that + # row each kv_head owns HEAD_DIM columns, i.e. one row of the flat storage. for b in range(BATCH): slot = int(slot_mapping[b].item()) - sblk, soff = slot // BLOCK_SIZE, slot % BLOCK_SIZE for kvh in range(NUM_KV_HEADS): - row = cb + (sblk * NUM_KV_HEADS + kvh) * BLOCK_SIZE + soff + row = cb + slot * NUM_KV_HEADS + kvh args.k_cache[row] = cur_k[(b, kvh)].to(torch.bfloat16) args.v_cache[row] = cur_v[(b, kvh)].to(torch.bfloat16) @@ -256,7 +266,7 @@ def _one_layer(args, layer: int, x: torch.Tensor) -> torch.Tensor: def compute_golden(args: TaskArgsBuilder) -> None: - """Fill ``args.out`` (and INOUT k_cache/v_cache) for the 2-layer decode chunk.""" + """Fill ``args.out`` (and INOUT k_cache/v_cache) for the 40-layer decode chunk.""" cur = args.hidden_states.float() # copy_hidden: bf16 input embedded as FP32 for layer in range(N_LAYERS): cur = _one_layer(args, layer, cur) diff --git a/simpler_setup/scene_test.py b/simpler_setup/scene_test.py index 686d569210..a51c0b7f0d 100644 --- a/simpler_setup/scene_test.py +++ b/simpler_setup/scene_test.py @@ -510,10 +510,52 @@ def _resolve_chip_entry_paths(entry, cls_dir): k = dict(k) if "source" in k and not os.path.isabs(k["source"]): k["source"] = str(cls_dir / k["source"]) + if k.get("extra_include_dirs"): + k["extra_include_dirs"] = [ + d if "$" in str(d) or os.path.isabs(str(d)) else str((cls_dir / str(d)).resolve()) + for d in k["extra_include_dirs"] + ] resolved.append(k) entry["incores"] = resolved +def _resolve_incore_include_dirs(dirs, incore): + """Resolve an incore's ``extra_include_dirs`` down to usable ``-I`` paths. + + Entries may use ``$VAR`` (typically ``$ASCEND_HOME_PATH``, so a CANN-linked + kernel does not hardcode a machine-specific path); relative ones were made + absolute against the test file's directory when the class was declared. + + This runs when the kernel is compiled, not when its module is imported — + a case declaring a CANN dependency is still collected on sim and macOS + runners that have no CANN at all, and must not fail there. + + A vendored toolkit ships one over-broad candidate list across SDK versions, + so paths that do not exist here are dropped. An unset variable raises, and so + does a declared set where nothing survived: that means the SDK is absent, and + saying so beats a "file not found" from deep inside vendored code. + """ + name = incore.get("name", incore.get("source", "?")) + resolved = [] + for raw in dirs: + expanded = os.path.expandvars(str(raw)) + if "$" in expanded: + raise ValueError( + f"incore {name}: extra_include_dirs entry {raw!r} references an environment " + f"variable that is not set. Export it (e.g. ASCEND_HOME_PATH for CANN headers) " + f"before running." + ) + if Path(expanded).is_dir(): + resolved.append(expanded) + if not resolved: + raise ValueError( + f"incore {name}: none of its {len(dirs)} extra_include_dirs exist on this machine, " + f"so the SDK they come from is not installed where the test expects. Tried:\n " + + "\n ".join(os.path.expandvars(str(d)) for d in dirs) + ) + return resolved + + def _extract_name_map(callable_spec: dict) -> dict: """Extract name mapping from a CALLABLE spec. @@ -921,8 +963,12 @@ def _compile_chip_callable_from_spec(spec, platform, runtime, cache_key): kernel_binaries = [] for k in incores: signature = k.get("signature", []) + extra = _resolve_incore_include_dirs(k["extra_include_dirs"], k) if k.get("extra_include_dirs") else [] incore = kc.compile_incore( - k["source"], core_type=k["core_type"], pto_isa_root=pto_isa_root, extra_include_dirs=inc_dirs + k["source"], + core_type=k["core_type"], + pto_isa_root=pto_isa_root, + extra_include_dirs=inc_dirs + extra, ) if not is_sim: incore = extract_text_section(incore) diff --git a/tests/lint/check_headers.py b/tests/lint/check_headers.py index d086b411b6..429840989e 100644 --- a/tests/lint/check_headers.py +++ b/tests/lint/check_headers.py @@ -69,6 +69,13 @@ ".clang-tidy": PY_HEADER, } +# Paths carrying third-party or harvested-verbatim sources, which keep their +# upstream copyright notice. Rewriting those headers would misattribute the code, +# and the next refresh would copy the originals back over the edit anyway. Drop a +# harvested tree under a `vendor/` directory and it is covered here and by the +# matching excludes in .pre-commit-config.yaml, with no config change per tree. +EXCLUDED_PATH_PARTS = ["reference", "3rdparty", "vendor/"] + def get_git_tracked_files(root_dir: Path) -> list[Path]: """Get list of files tracked by git.""" @@ -179,7 +186,7 @@ def _collect_files(files: list[str], excluded_patterns: list[str]) -> Optional[l return None all_files = get_git_tracked_files(root_path) - filtered = [f for f in all_files if not any(str(f.relative_to(root_path)).startswith(p) for p in excluded_patterns)] + filtered = [f for f in all_files if not any(p in str(f.relative_to(root_path)) for p in excluded_patterns)] return [f for f in filtered if _is_checkable(f)] @@ -195,7 +202,7 @@ def main(): args = parser.parse_args() - files_to_check = _collect_files(args.files, ["reference"]) + files_to_check = _collect_files(args.files, EXCLUDED_PATH_PARTS) if files_to_check is None: return 1