From 370d1086b8bd6a4a6519b66629c9c28aced16aeb Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Sun, 26 Jul 2026 19:37:50 -0700 Subject: [PATCH] Update: re-harvest qwen3_14b_decode and scale it to all 40 layers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #1484 The example carried pypto codegen last refreshed on 2026-06-25, cited a provenance file upstream had deleted (`models/qwen3/14b/decode_layer.py`), and covered 2 of Qwen3-14B's 40 layers. Re-harvested from `decode_fwd.py`'s `decode_fwd_layers` at pypto-lib `45be52c` with `_CHUNK_NLAYERS = 40`, and the README now records that commit plus the ptoas and pto-isa pins so the next drift is detectable rather than archaeological. Scaling to the full model is nearly free. The generated orchestration loops over layers (`for (int64_t i = 0; i < 40; i += 1)`) instead of unrolling them, so going from 2 to 40 moves a trip count and leaves the kernel set the same size. The lib chunks its dispatches because a 40-layer one used to exceed a 2000 ms stream-sync timeout; simpler raised that to 50 s in #1175 and a 40-layer step is ~40 ms, so the constraint is gone. What did change is attention. Since pypto-lib #765 and #796 it is a CANN FusedInferAttentionScore extern rather than generated kernels, folding per-head QK-norm, RoPE, the paged KV write, the flash inner loop and the online softmax into one mixed AIC+2xAIV task behind an FFTS barrier. Seven kernels go away (`rope_qkv`, `qk_gamma`, `qk_recip`, `fa_work_build`, `fa_fused_aic`/`_aiv`, `online_softmax`, `attn_fence`) and the harvested extern arrives under `kernels/vendor/paged_attention_cce/`. That extern needs CANN devkit headers, which nothing could express: every incore compiled against the same fixed include set, even though `KernelCompiler.compile_incore` has always taken `extra_include_dirs`. The scene-test `CALLABLE` now carries a per-incore `extra_include_dirs`, resolved when the kernel is compiled rather than when its module is imported — sim and macOS runners collect a case declaring a CANN dependency without having CANN. Entries may use `$VAR` and may name directories that do not exist here, so one candidate list spans SDK layouts; an unset variable or a set where nothing survived raises, which beats a "file not found" from four includes deep in vendored code. The harvested tree is copied verbatim and must stay that way: `fai_body.hpp` reaches its dependencies by relative include, and reformatting it would make the next refresh diff against our copy instead of against upstream — the exact drift this example exists to expose. It therefore lives under `kernels/vendor/`, and the repo's header, formatting and language lint skip anything below a `vendor/` directory. Keying the carve-out on the directory rather than on this operator's name means the next harvested extern needs no lint change at all. The extern also moves the paged KV layout from NSND to vLLM's BSND (`[page, token, kv_head, dim]`, so `slot_mapping` is directly the row index), which the golden now models for both the attention read and the KV write-back. The previous harvest needed one hand-edit, replacing a `[[block_local]] static` in `fa_fused_aiv` that the AICore loader rejected. That kernel no longer exists and the current codegen emits no such construct, so this harvest is unmodified codegen throughout. At 40 layers the case takes ~5 min against a 38.05 GiB fixture (24.61 weights + 13.44 paged KV, bf16) — over half the onboard sweep's 600 s session budget. It runs as its own CI step on a dedicated device instead, keeping that sweep's short timeout meaningful as a hang detector. Verification: passes on device with output and all 40 layers' KV caches matching the torch reference at RTOL=5e-2 / ATOL=1e-1, with no ring-sizing overrides — per-layer intermediates live inside that iteration's scope, so the live set does not grow with layer count. The l0-swimlane docs used this example as their "real SPMD workload", shrinking the generated fa_fused via `--set-arg 0=96` on its work-item count. The extern takes its work distribution from runtime tiling metadata instead, so no `--set-arg` shrinks it and those references are removed rather than repointed. The dfx buffer-capacity row measured at 2 layers is dropped for the same reason: it would now be a stale number rather than a wrong one. Co-Authored-By: Claude Opus 5 (1M context) --- .claude/skills/l0-swimlane/SKILL.md | 1 - .claude/skills/multi-repo-qwen-setup/SKILL.md | 16 +- .github/workflows/ci.yml | 24 +- .pre-commit-config.yaml | 8 +- docs/dfx/dfx-buffer-capacity-audit.md | 7 +- docs/dfx/l0-swimlane-profiling.md | 14 +- docs/testing.md | 11 + .../qwen3_14b_decode/README.md | 210 +- .../kernels/aic/down_proj.cpp | 519 +-- .../kernels/aic/fa_fused_aic.cpp | 767 ----- .../kernels/aic/gate_proj.cpp | 508 +-- .../kernels/aic/gate_proj_0.cpp | 621 ++++ .../kernels/aic/gate_proj_1.cpp | 621 ++++ .../kernels/aic/gate_proj_2.cpp | 621 ++++ .../kernels/aic/gate_proj_3.cpp | 621 ++++ .../kernels/aic/gate_proj_4.cpp | 622 ++++ .../qwen3_14b_decode/kernels/aic/k_proj.cpp | 718 +++-- .../qwen3_14b_decode/kernels/aic/out_proj.cpp | 413 +-- .../kernels/aic/out_proj_0.cpp | 579 ++++ .../qwen3_14b_decode/kernels/aic/q_proj.cpp | 665 ++-- .../qwen3_14b_decode/kernels/aic/up_proj.cpp | 508 +-- .../kernels/aic/up_proj_0.cpp | 621 ++++ .../kernels/aic/up_proj_1.cpp | 621 ++++ .../kernels/aic/up_proj_2.cpp | 621 ++++ .../kernels/aic/up_proj_3.cpp | 621 ++++ .../kernels/aic/up_proj_4.cpp | 622 ++++ .../qwen3_14b_decode/kernels/aic/v_proj.cpp | 718 +++-- .../kernels/aiv/attn_fence.cpp | 143 - .../aiv/{k_seed.cpp => attn_out_seed.cpp} | 40 +- .../kernels/aiv/copy_hidden.cpp | 81 +- .../qwen3_14b_decode/kernels/aiv/copy_out.cpp | 76 +- .../kernels/aiv/dcr_xgamma.cpp | 226 ++ .../kernels/aiv/down_cast_residual.cpp | 180 -- .../kernels/aiv/down_seed.cpp | 156 - .../kernels/aiv/fa_fused_aiv.cpp | 779 ----- .../kernels/aiv/fa_work_build.cpp | 268 -- .../kernels/aiv/gate_seed.cpp | 158 - .../kernels/aiv/{v_seed.cpp => kv_seed.cpp} | 72 +- .../kernels/aiv/mlp_out_seed.cpp | 391 +++ .../kernels/aiv/online_softmax.cpp | 1185 ------- .../kernels/aiv/out_consolidate.cpp | 129 - .../qwen3_14b_decode/kernels/aiv/out_seed.cpp | 134 - .../kernels/aiv/post_rms_reduce.cpp | 363 +-- .../qwen3_14b_decode/kernels/aiv/q_seed.cpp | 50 +- .../qwen3_14b_decode/kernels/aiv/qk_gamma.cpp | 294 -- .../qwen3_14b_decode/kernels/aiv/qk_recip.cpp | 442 --- .../kernels/aiv/residual_rms_cast.cpp | 363 +-- .../kernels/aiv/residual_rms_cast_0.cpp | 363 +-- .../kernels/aiv/residual_rms_cast_1.cpp | 363 +-- .../kernels/aiv/residual_rms_cast_2.cpp | 363 +-- .../kernels/aiv/residual_rms_cast_3.cpp | 363 +-- .../kernels/aiv/rms_recip.cpp | 456 ++- .../qwen3_14b_decode/kernels/aiv/rope_qkv.cpp | 1051 ------ .../qwen3_14b_decode/kernels/aiv/silu.cpp | 326 +- .../qwen3_14b_decode/kernels/aiv/up_seed.cpp | 158 - .../qwen3_14b_decode/kernels/aiv/x_gamma.cpp | 383 --- .../qwen3_14b_decode/kernels/aiv/x_gamma0.cpp | 244 ++ .../orchestration/decode_fwd_layers.cpp | 2820 +++++++++++------ .../paged_attention_cce/attention/entry.cpp | 44 + .../attention_rope/entry.cpp | 50 + .../generated/kernel_tiling/kernel_tiling.h | 94 + .../paged_attention_cce/kernel/fai_body.hpp | 240 ++ .../kernel/metadata_layout.h | 40 + .../kernel/rope_qkv_generated.hpp | 1246 ++++++++ .../paged_attention_cce/tiling/entry.cpp | 116 + .../tiling/qwen_fai_runtime_tiler.hpp | 346 ++ .../attn_infra/arch/arch.hpp | 55 + .../attn_infra/arch/cross_core_sync.hpp | 120 + .../attn_infra/arch/local_tensor_buffer.hpp | 234 ++ .../attn_infra/arch/resource.hpp | 51 + .../attn_infra/base_defs.hpp | 46 + .../attn_infra/coord.hpp | 442 +++ .../attn_infra/detail/alignment.hpp | 75 + .../attn_infra/detail/dependent_false.hpp | 20 + .../attn_infra/detail/macros.hpp | 16 + .../epilogue/block/CombineScale.hpp | 290 ++ .../epilogue/block/block_epilogue.hpp | 34 + .../block/block_epilogue_init_outputs.hpp | 168 + .../block/block_epilogue_online_softmax.hpp | 2055 ++++++++++++ ...block_epilogue_online_softmax_low_prec.hpp | 868 +++++ .../block/block_epilogue_rescale_o.hpp | 1105 +++++++ .../block_epilogue_rescale_o_low_prec.hpp | 472 +++ .../attn_infra/epilogue/dispatch_policy.hpp | 56 + .../epilogue/tile_common/copy_gm_to_ub.hpp | 188 ++ .../epilogue/tile_common/copy_ub_to_gm.hpp | 144 + .../tile_broadcast_inplace_by_column.hpp | 69 + .../tile_broadcast_inplace_by_row.hpp | 59 + .../tile_common/tile_broadcast_mul.hpp | 140 + .../tile_common/tile_broadcast_one_blk.hpp | 62 + .../epilogue/tile_common/tile_cast.hpp | 49 + .../epilogue/tile_common/tile_copy.hpp | 108 + .../tile_common/tile_elemwise_add.hpp | 50 + .../tile_common/tile_elemwise_mul.hpp | 49 + .../tile_common/tile_elemwise_muls.hpp | 43 + .../epilogue/tile_common/tile_swizzle.hpp | 93 + .../attn_infra/gemm/block/block_mmad.hpp | 42 + .../attn_infra/gemm/block/block_mmad_pv.hpp | 325 ++ .../gemm/block/block_mmad_pv_decode.hpp | 286 ++ .../attn_infra/gemm/block/block_mmad_qk.hpp | 348 ++ .../gemm/block/block_mmad_qk_decode.hpp | 290 ++ .../attn_infra/gemm/dispatch_policy.hpp | 76 + .../attn_infra/gemm/gemm_type.hpp | 28 + .../attn_infra/gemm/helper.hpp | 255 ++ .../gemm/tile_common/copy_gm_to_l1.hpp | 1067 +++++++ .../gemm/tile_common/copy_gm_to_ub.hpp | 22 + .../gemm/tile_common/copy_l0c_to_gm.hpp | 209 ++ .../gemm/tile_common/copy_l1_to_bt.hpp | 27 + .../gemm/tile_common/copy_l1_to_l0a.hpp | 356 +++ .../gemm/tile_common/copy_l1_to_l0b.hpp | 487 +++ .../gemm/tile_common/copy_ub_to_gm.hpp | 21 + .../attn_infra/gemm/tile_common/tile_copy.hpp | 63 + .../gemm/tile_common/tile_copy_tla.hpp | 20 + .../attn_infra/gemm/tile_common/tile_mmad.hpp | 104 + .../attn_infra/gemm_coord.hpp | 163 + .../attn_infra/layout/layout.hpp | 18 + .../attn_infra/layout/matrix.hpp | 1208 +++++++ .../attn_infra/layout/vector.hpp | 133 + .../attn_infra/matrix_coord.hpp | 108 + .../flash_attention_regular.h | 1116 +++++++ .../kernel_common.hpp | 155 + .../qwen3_14b_decode/test_qwen3_14b_decode.py | 351 +- simpler_setup/goldens/qwen3_14b_decode.py | 60 +- simpler_setup/scene_test.py | 48 +- tests/lint/check_headers.py | 11 +- 124 files changed, 29999 insertions(+), 11053 deletions(-) delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/fa_fused_aic.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_0.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_1.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_2.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_3.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/gate_proj_4.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/out_proj_0.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_0.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_1.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_2.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_3.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aic/up_proj_4.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/attn_fence.cpp rename examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/{k_seed.cpp => attn_out_seed.cpp} (53%) create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/dcr_xgamma.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_cast_residual.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/down_seed.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_fused_aiv.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/fa_work_build.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/gate_seed.cpp rename examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/{v_seed.cpp => kv_seed.cpp} (51%) create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/mlp_out_seed.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/online_softmax.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_consolidate.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/out_seed.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_gamma.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/qk_recip.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/rope_qkv.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/up_seed.cpp delete mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/aiv/x_gamma0.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention/entry.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/attention_rope/entry.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/generated/kernel_tiling/kernel_tiling.h create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/fai_body.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/metadata_layout.h create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/kernel/rope_qkv_generated.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/entry.cpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/tiling/qwen_fai_runtime_tiler.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/arch.hpp create mode 100644 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 create mode 100644 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 create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/arch/resource.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/base_defs.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/coord.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/alignment.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/dependent_false.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/detail/macros.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/block/CombineScale.hpp create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/epilogue/dispatch_policy.hpp create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/dispatch_policy.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/gemm_type.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm/helper.hpp create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 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 create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/gemm_coord.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/layout.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/matrix.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/layout/vector.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/attn_infra/matrix_coord.hpp create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/flash_attention_regular.h create mode 100644 examples/a2a3/tensormap_and_ringbuffer/qwen3_14b_decode/kernels/vendor/paged_attention_cce/vendor/fused_infer_attention_score/kernel_common.hpp 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