diff --git a/docs/repro/pipelining/.gitignore b/docs/repro/pipelining/.gitignore new file mode 100644 index 0000000000..4478ddd95b --- /dev/null +++ b/docs/repro/pipelining/.gitignore @@ -0,0 +1,12 @@ +build/ +sim_outputs/ +device/build/ + +*.so +*.o +log_ca/ +instr.bin +__pycache__/ +*.pyc +.pytest_cache/ +*.egg-info/ diff --git a/docs/repro/pipelining/README.md b/docs/repro/pipelining/README.md new file mode 100644 index 0000000000..dc7c4f800d --- /dev/null +++ b/docs/repro/pipelining/README.md @@ -0,0 +1,96 @@ +# pipelining — VMI Persistent stages vs CCE ping-pong + +Minimal, **kernel-agnostic** compile (+ optional on-device CCE) repro for the +PTOAS gap: long-K Persistent loops want `stages=2` and `block_k > 512` so MTE +and vector stay overlapped. Ascend CCE can express and **run** stages=2 +ping-pong; VMI on tag **vmi-v0.1.3** cannot reach the same wide-tile schedule +— layout assignment rejects 1024-wide tiles, and the hand-written MI path +crashes bisheng even after ptoas emits LLVM IR. + +Product-shaped MHC / quant patterns and ASC↔CCE **bandwidth** live in the +companion study under `vmi-demo-pipeline/pipeline_study` (separate repo). + +## Pin + +| Component | Ref | +|-----------|-----| +| PTOAS | tag **vmi-v0.1.3** (`ptoas` from that build; not PyPI-only) | +| Branch | `zjw/pipelining_issue` | +| CANN | 9.0 / `bisheng` | +| `PTOAS_ROOT` | repo root (`../../..` from this directory) | + +If `install/bin/ptoas` crashes on missing MLIR libs, point tools at a build tree: + +```bash +export PTOAS_TOOLS_ROOT=/path/to/PTOAS-built-at-vmi-v0.1.3 +``` + +## What breaks + +| Path | stages | block_k | Fault mode | +|------|-------:|--------:|------------| +| CCE (`fixtures/reference_asc_cce.asc`, `device/scale_stages2.asc`) | 2 | 512–1024 | **OK** — compile; device host numerical OK | +| VMI stages=1 / 512 (`current_slow_vmi.*`) | 1 | 512 | **emit-vpto OK** | +| VMI stages=2 / 512 (`isolate_stages2_blockk512_vmi.*`) | 2 | 512 | **emit-vpto OK** (dual-buffer alone is fine) | +| VMI stages=1 / 1024 (`isolate_stages1_blockk1024_vmi.*`) | 1 | 1024 | **Layout reject** on 1024-wide vmul | +| VMI stages=2 / 1024 (`desired_vmi.*`) | 2 | 1024 | **Layout reject** (stages+wide) | +| VMI target MI (`target_mi.pto`) | 2 | 1024 | **Bisheng crash** after LLVM IR emit | + +The blocker for the desired schedule is **wide `block_k=1024` layout**, not +dual-buffer stages alone. Fixtures use a trivial scale body so the failure +stays schedule/layout — not a product kernel. + +## Check results (vmi-v0.1.3 tools, 2026-07-28) + +```bash +source scripts/env.sh +# optional: export PTOAS_TOOLS_ROOT=... +./scripts/check_stages.sh +``` + +| Step | Result | +|------|--------| +| `reference_asc_cce.asc` | **PASS** bisheng `--cce-aicore-only` | +| `device/scale_stages2.asc` aicore object | **PASS** | +| `current_slow_vmi.pto` `--emit-vpto` | **PASS** | +| `isolate_stages2_blockk512_vmi.pto` | **PASS** | +| `isolate_stages1_blockk1024_vmi.pto` | **FAIL** layout | +| `desired_vmi.pto` | **FAIL** layout | +| `lowered_vpto.pto` → LLVM IR | **PASS** | +| `target_mi.pto` → emit + LLVM IR | **PASS**; bisheng `.o` **FAIL** | + +Full log: `sim_outputs/check_stages/compile_results.txt` + +## On-device CCE (correctness only) + +```bash +cd device +./run.sh # build + launch stages=2 f32 scale, check max_err +MSOPPROF=1 ./run.sh # optional msopprof wrap if installed +``` + +Measured 2026-07-28: + +``` +scale_stages2: N=2048 mism=0 max_err=0 OK +``` + +CCE stages=2 runs with dual-buffer HardEvent overlap. VMI stages=2 / +`block_k=1024` does not compile through layout — that is the PTOAS gap. + +For e2e bandwidth (stages=1 vs stages=2, MHC / quant-shaped bodies), see +`pipeline_study` in the vmi-demo-pipeline repo. + +## Layout + +``` +pipelining/ + fixtures/ CCE shell, VMI stages×block_k isolate, target_mi + device/ AscendC host for stages=2 CCE numerical check + scripts/ env.sh, check_stages.sh +``` + +## Done when + +VMI can compile and run stages=2 / `block_k=1024` Persistent tiles with +dual-buffer ping-pong — same overlap pattern as CCE and `target_mi.pto`. diff --git a/docs/repro/pipelining/device/build.sh b/docs/repro/pipelining/device/build.sh new file mode 100755 index 0000000000..bc61cc0646 --- /dev/null +++ b/docs/repro/pipelining/device/build.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +# Build launchable CCE stages=2 scale host (relative CANN / PTOAS paths). +set -euo pipefail +DEVICE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +# shellcheck disable=SC1091 +source "${DEVICE_DIR}/../scripts/env.sh" + +AICORE_ARCH="${PTO_AICORE_ARCH:-dav-c310}" +OUT="${DEVICE_DIR}/build" +mkdir -p "${OUT}" + +echo "bisheng=${BISHENG}" +echo "ASCEND=${ASCEND}" + +"${BISHENG}" -c -fPIC -O2 -std=c++17 -xcce \ + -Xhost-start -Xhost-end \ + -mllvm -cce-aicore-stack-size=0x8000 \ + -mllvm -cce-aicore-function-stack-size=0x8000 \ + -mllvm -cce-aicore-record-overflow=true \ + -mllvm -cce-aicore-addr-transform \ + -mllvm -cce-aicore-dcci-insert-for-scalar=false \ + --cce-aicore-arch="${AICORE_ARCH}" \ + -I"${ASCEND}/include" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/impl" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/interface" \ + "${DEVICE_DIR}/scale_stages2.asc" \ + -o "${OUT}/scale_stages2.o" + +"${BISHENG}" -fPIC -shared --cce-fatobj-link \ + -o "${OUT}/libscale_stages2.so" \ + "${OUT}/scale_stages2.o" \ + -Wl,--no-as-needed -L"${ASCEND}/lib64" -lruntime + +"${BISHENG}" -O2 -std=c++17 -xc++ \ + -I"${ASCEND}/include" \ + "${DEVICE_DIR}/main.cpp" \ + -o "${OUT}/scale_stages2_host" \ + -L"${OUT}" -L"${ASCEND}/lib64" \ + -lscale_stages2 -lascendcl -lruntime -ltiling_api -lplatform -lc_sec -ldl -lnnopbase \ + -Wl,-rpath,"${OUT}" -Wl,-rpath,"${ASCEND}/lib64" + +echo "OK: ${OUT}/scale_stages2_host" diff --git a/docs/repro/pipelining/device/main.cpp b/docs/repro/pipelining/device/main.cpp new file mode 100644 index 0000000000..793e8bba6f --- /dev/null +++ b/docs/repro/pipelining/device/main.cpp @@ -0,0 +1,65 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// Minimal host: fill f32 input, launch CCE stages=2 scale, check vs host ref. +#include "acl/acl.h" +#include +#include +#include +#include + +void LaunchScaleStages2(void* x, void* scale, void* out, void* stream); + +int main() { + constexpr int N = 2048; + const size_t nbytes = N * sizeof(float); + const float scale_v = 1.5f; + + aclInit(nullptr); + aclrtSetDevice(0); + aclrtStream stream; + aclrtCreateStream(&stream); + + float *x_h = nullptr, *y_h = nullptr, *scale_h = nullptr; + void *x_d = nullptr, *y_d = nullptr, *scale_d = nullptr; + aclrtMallocHost((void**)&x_h, nbytes); + aclrtMallocHost((void**)&y_h, nbytes); + aclrtMallocHost((void**)&scale_h, sizeof(float)); + aclrtMalloc(&x_d, nbytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&y_d, nbytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&scale_d, sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST); + + scale_h[0] = scale_v; + for (int i = 0; i < N; ++i) { + x_h[i] = 0.25f * static_cast(i % 17); + y_h[i] = 0.f; + } + + aclrtMemcpy(x_d, nbytes, x_h, nbytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(scale_d, sizeof(float), scale_h, sizeof(float), ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(y_d, nbytes, y_h, nbytes, ACL_MEMCPY_HOST_TO_DEVICE); + + LaunchScaleStages2(x_d, scale_d, y_d, stream); + aclrtSynchronizeStream(stream); + aclrtMemcpy(y_h, nbytes, y_d, nbytes, ACL_MEMCPY_DEVICE_TO_HOST); + + int mism = 0; + float max_err = 0.f; + for (int i = 0; i < N; ++i) { + float ref = x_h[i] * scale_v; + float err = std::fabs(ref - y_h[i]); + if (err > max_err) max_err = err; + if (err > 1e-5f) ++mism; + } + std::printf("scale_stages2: N=%d mism=%d max_err=%.4g %s\n", N, mism, max_err, + mism == 0 ? "OK" : "FAIL"); + + aclrtFree(x_d); + aclrtFree(y_d); + aclrtFree(scale_d); + aclrtFreeHost(x_h); + aclrtFreeHost(y_h); + aclrtFreeHost(scale_h); + aclrtDestroyStream(stream); + aclrtResetDevice(0); + aclFinalize(); + return mism == 0 ? 0 : 1; +} diff --git a/docs/repro/pipelining/device/run.sh b/docs/repro/pipelining/device/run.sh new file mode 100755 index 0000000000..9fd7f074fb --- /dev/null +++ b/docs/repro/pipelining/device/run.sh @@ -0,0 +1,12 @@ +#!/usr/bin/env bash +# Build and run CCE stages=2 on NPU. Optional: MSOPPROF=1 for msopprof wrap. +set -euo pipefail +DEVICE_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +bash "${DEVICE_DIR}/build.sh" +# shellcheck disable=SC1091 +source "${DEVICE_DIR}/../scripts/env.sh" +HOST="${DEVICE_DIR}/build/scale_stages2_host" +if [ "${MSOPPROF:-0}" = "1" ] && command -v msopprof >/dev/null 2>&1; then + msopprof --application="${HOST}" --output="${DEVICE_DIR}/build/msopprof" || true +fi +exec "${HOST}" diff --git a/docs/repro/pipelining/device/scale_stages2.asc b/docs/repro/pipelining/device/scale_stages2.asc new file mode 100644 index 0000000000..49487f548c --- /dev/null +++ b/docs/repro/pipelining/device/scale_stages2.asc @@ -0,0 +1,68 @@ +// Launchable CCE stages=2 scale: GM f32 * scale -> GM f32 (real MTE + VF). +// Proves dual-buffer HardEvent ping-pong runs on device; VMI stages=2 is the gap. +#include "kernel_operator.h" + +__simd_vf__ inline void vf_scale_f32(__ubuf__ float* x_ub, __ubuf__ float* scale_ub, + __ubuf__ float* out_ub) { + vector_bool v32 = pset_b32(PAT_ALL); + for (int i = 0; i < 512; i += 64) { + vector_f32 xv, sv, yv; + vlds(xv, x_ub + i, 0, NORM); + vlds(sv, scale_ub, 0, BRC_B32); + vmul(yv, xv, sv, v32, MODE_ZEROING); + vsts(yv, out_ub + i, 0, NORM_B32, v32); + } +} + +// N = 2048 f32, block_k = 512, stages = 2, 4 tiles. +extern "C" __global__ __vector__ void scale_stages2(__gm__ float* x, __gm__ float* scale, + __gm__ float* out) { + AscendC::InitSocState(); + __ubuf__ uint8_t* ub = (__ubuf__ uint8_t*)0; + + constexpr int kBlockK = 512; + constexpr int kTiles = 4; + constexpr int kBytes = kBlockK * 4; + + __ubuf__ float* sc_ub = (__ubuf__ float*)(ub + 4 * kBytes); + copy_gm_to_ubuf_align_v2((__ubuf__ uint8_t*)sc_ub, (__gm__ uint8_t*)scale, 0, 1, 4, 0, 0, 0, 0, + 32, 32); + AscendC::SetFlag(2); + AscendC::WaitFlag(2); + + AscendC::SetFlag(0); + AscendC::SetFlag(1); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + + for (int w = 0; w < kTiles; ++w) { + int slot = w & 1; + __ubuf__ float* x_ub = (__ubuf__ float*)(ub + slot * kBytes); + __ubuf__ float* o_ub = (__ubuf__ float*)(ub + (2 + slot) * kBytes); + + AscendC::WaitFlag(slot); + copy_gm_to_ubuf_align_v2((__ubuf__ uint8_t*)x_ub, (__gm__ uint8_t*)(x + w * kBlockK), 0, 1, + kBytes, 0, 0, 0, 0, kBytes, kBytes); + AscendC::SetFlag(slot); + AscendC::WaitFlag(slot); + AscendC::WaitFlag(slot); + + vf_scale_f32(x_ub, sc_ub, o_ub); + + AscendC::SetFlag(slot); + AscendC::SetFlag(slot); + AscendC::WaitFlag(slot); + copy_ubuf_to_gm_align_v2((__gm__ void*)(out + w * kBlockK), (__ubuf__ void*)o_ub, 0, 1, kBytes, + 4, kBytes, kBytes); + AscendC::SetFlag(slot); + } + + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); +} + +void LaunchScaleStages2(void* x, void* scale, void* out, void* stream) { + scale_stages2<<<1, nullptr, stream>>>((__gm__ float*)x, (__gm__ float*)scale, (__gm__ float*)out); +} diff --git a/docs/repro/pipelining/fixtures/current_slow_vmi.pto b/docs/repro/pipelining/fixtures/current_slow_vmi.pto new file mode 100644 index 0000000000..4494a4cce1 --- /dev/null +++ b/docs/repro/pipelining/fixtures/current_slow_vmi.pto @@ -0,0 +1,47 @@ +// 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. + +// Current slow path: Persistent vector tile: stages=1, block_k=512. + +module attributes { + pto.target_arch = "a5", + pto.kernel_kind = #pto.kernel_kind + // stages = 1 + // block_k = 512 +} { + func.func @persistent_quant_tile_k512( + %x_ub: !pto.ptr, + %y_ub: !pto.ptr, + %scale_slot: !pto.vmi.vreg<1xf32>, + %block_count: index) + attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + + pto.vecscope { + %scale = pto.vmi.vbrc %scale_slot + : !pto.vmi.vreg<1xf32> -> !pto.vmi.vreg<512xf32> + scf.for %i = %c0 to %block_count step %c1 { + %off = arith.muli %i, %c512 : index + %x_bf16 = pto.vmi.vload %x_ub[%off] + : !pto.ptr -> !pto.vmi.vreg<512xbf16> + %x_f32 = pto.vmi.vcvt %x_bf16 + : !pto.vmi.vreg<512xbf16> -> !pto.vmi.vreg<512xf32> + %scaled = pto.vmi.vmul %x_f32, %scale + : !pto.vmi.vreg<512xf32>, !pto.vmi.vreg<512xf32> + -> !pto.vmi.vreg<512xf32> + %y = pto.vmi.vcvt %scaled {saturate = "SAT"} + : !pto.vmi.vreg<512xf32> -> !pto.vmi.vreg<512xf8E4M3FN> + pto.vmi.vstore %y, %y_ub[%off] + : !pto.vmi.vreg<512xf8E4M3FN>, !pto.ptr + } + } + return + } +} diff --git a/docs/repro/pipelining/fixtures/current_slow_vmi.py b/docs/repro/pipelining/fixtures/current_slow_vmi.py new file mode 100644 index 0000000000..2c0f4e96b2 --- /dev/null +++ b/docs/repro/pipelining/fixtures/current_slow_vmi.py @@ -0,0 +1,36 @@ +# 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. + +"""Python DSL form of ``current_slow_vmi.pto`` — stages=1, block_k=512 Persistent tile. + +Canonical IR: ``current_slow_vmi.pto``. +""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def persistent_quant_tile_k512(*, BLOCK_COUNT: pto.const_expr = 4): + # stages = 1, block_k = 512 + x_ub = pto.alloc_tile(shape=[1, 512], dtype=pto.bf16) + y_ub = pto.alloc_tile(shape=[1, 512], dtype=pto.f8e4m3) + scale_ub = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off0 = pto.const(0, dtype=pto.index) + scale_slot = pto.vmi.vload(scale_ub.as_ptr(), off0, size=1) + scale = pto.vmi.vbrc(scale_slot, size=512) + for i in range(BLOCK_COUNT): + off = i * 512 + x_bf16 = pto.vmi.vload(x_ub.as_ptr(), off, size=512) + x_f32 = pto.vmi.vcvt(x_bf16, pto.f32) + scaled = pto.vmi.vmul(x_f32, scale) + y = pto.vmi.vcvt(scaled, pto.f8e4m3, saturate="SAT") + pto.vmi.vstore(y, y_ub.as_ptr(), off) + + +if __name__ == "__main__": + print(persistent_quant_tile_k512.compile().mlir_text()) diff --git a/docs/repro/pipelining/fixtures/desired_vmi.pto b/docs/repro/pipelining/fixtures/desired_vmi.pto new file mode 100644 index 0000000000..da24d2cdb3 --- /dev/null +++ b/docs/repro/pipelining/fixtures/desired_vmi.pto @@ -0,0 +1,62 @@ +// 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. + +// Idiomatic Persistent VMI: stages=2 and block_k=1024 (does not lower today). + +module attributes { + pto.target_arch = "a5", + pto.kernel_kind = #pto.kernel_kind, + pto.pipeline.stages = 2 : i32, + pto.block_k = 1024 : i32 +} { + func.func @persistent_quant_tile_k1024( + %x_ub0: !pto.ptr, + %x_ub1: !pto.ptr, + %y_ub0: !pto.ptr, + %y_ub1: !pto.ptr, + %scale_slot: !pto.vmi.vreg<1xf32>, + %block_count: index) + attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c1024 = arith.constant 1024 : index + + pto.vecscope { + %scale = pto.vmi.vbrc %scale_slot + : !pto.vmi.vreg<1xf32> -> !pto.vmi.vreg<1024xf32> + scf.for %i = %c0 to %block_count step %c1 { + %off = arith.muli %i, %c1024 : index + %stage = arith.remui %i, %c2 : index + %use1 = arith.cmpi eq, %stage, %c1 : index + %x_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %x_ub1 : !pto.ptr + } else { + scf.yield %x_ub0 : !pto.ptr + } + %y_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %y_ub1 : !pto.ptr + } else { + scf.yield %y_ub0 : !pto.ptr + } + %x_bf16 = pto.vmi.vload %x_ub[%off] + : !pto.ptr -> !pto.vmi.vreg<1024xbf16> + %x_f32 = pto.vmi.vcvt %x_bf16 + : !pto.vmi.vreg<1024xbf16> -> !pto.vmi.vreg<1024xf32> + %scaled = pto.vmi.vmul %x_f32, %scale + : !pto.vmi.vreg<1024xf32>, !pto.vmi.vreg<1024xf32> + -> !pto.vmi.vreg<1024xf32> + %y = pto.vmi.vcvt %scaled {saturate = "SAT"} + : !pto.vmi.vreg<1024xf32> -> !pto.vmi.vreg<1024xf8E4M3FN> + pto.vmi.vstore %y, %y_ub[%off] + : !pto.vmi.vreg<1024xf8E4M3FN>, !pto.ptr + } + } + return + } +} diff --git a/docs/repro/pipelining/fixtures/desired_vmi.py b/docs/repro/pipelining/fixtures/desired_vmi.py new file mode 100644 index 0000000000..cdee3d6176 --- /dev/null +++ b/docs/repro/pipelining/fixtures/desired_vmi.py @@ -0,0 +1,48 @@ +# 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. + +"""Python DSL form of ``desired_vmi.pto`` — stages=2, block_k=1024 dual-buffer tile. + +Canonical IR: ``desired_vmi.pto`` (module attrs ``pto.pipeline.stages`` / +``pto.block_k``). Layout assignment rejects the 1024-wide multiply today. +""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def persistent_quant_tile_k1024(*, BLOCK_COUNT: pto.const_expr = 4): + # stages = 2, block_k = 1024 — dual UB ping-pong + x_ub0 = pto.alloc_tile(shape=[1, 1024], dtype=pto.bf16) + x_ub1 = pto.alloc_tile(shape=[1, 1024], dtype=pto.bf16) + y_ub0 = pto.alloc_tile(shape=[1, 1024], dtype=pto.f8e4m3) + y_ub1 = pto.alloc_tile(shape=[1, 1024], dtype=pto.f8e4m3) + scale_ub = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off0 = pto.const(0, dtype=pto.index) + scale_slot = pto.vmi.vload(scale_ub.as_ptr(), off0, size=1) + scale = pto.vmi.vbrc(scale_slot, size=1024) + for i in range(BLOCK_COUNT): + off = i * 1024 + if (i % 2) == 1: + x_ptr = x_ub1.as_ptr() + y_ptr = y_ub1.as_ptr() + else: + x_ptr = x_ub0.as_ptr() + y_ptr = y_ub0.as_ptr() + x_bf16 = pto.vmi.vload(x_ptr, off, size=1024) + x_f32 = pto.vmi.vcvt(x_bf16, pto.f32) + scaled = pto.vmi.vmul(x_f32, scale) + y = pto.vmi.vcvt(scaled, pto.f8e4m3, saturate="SAT") + pto.vmi.vstore(y, y_ptr, off) + + +if __name__ == "__main__": + try: + print(persistent_quant_tile_k1024.compile().mlir_text()) + except Exception as exc: # noqa: BLE001 — document the gap + print(f"desired emit/lower may fail (layout / stages): {exc}") diff --git a/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.pto b/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.pto new file mode 100644 index 0000000000..f7a7d21e96 --- /dev/null +++ b/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.pto @@ -0,0 +1,47 @@ +// 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. + +// Isolate: stages=1, block_k=1024 — wide tile without dual-buffer stages. + +module attributes { + pto.target_arch = "a5", + pto.kernel_kind = #pto.kernel_kind, + pto.pipeline.stages = 1 : i32, + pto.block_k = 1024 : i32 +} { + func.func @persistent_quant_stages1_k1024( + %x_ub: !pto.ptr, + %y_ub: !pto.ptr, + %scale_slot: !pto.vmi.vreg<1xf32>, + %block_count: index) + attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c1024 = arith.constant 1024 : index + + pto.vecscope { + %scale = pto.vmi.vbrc %scale_slot + : !pto.vmi.vreg<1xf32> -> !pto.vmi.vreg<1024xf32> + scf.for %i = %c0 to %block_count step %c1 { + %off = arith.muli %i, %c1024 : index + %x_bf16 = pto.vmi.vload %x_ub[%off] + : !pto.ptr -> !pto.vmi.vreg<1024xbf16> + %x_f32 = pto.vmi.vcvt %x_bf16 + : !pto.vmi.vreg<1024xbf16> -> !pto.vmi.vreg<1024xf32> + %scaled = pto.vmi.vmul %x_f32, %scale + : !pto.vmi.vreg<1024xf32>, !pto.vmi.vreg<1024xf32> + -> !pto.vmi.vreg<1024xf32> + %y = pto.vmi.vcvt %scaled {saturate = "SAT"} + : !pto.vmi.vreg<1024xf32> -> !pto.vmi.vreg<1024xf8E4M3FN> + pto.vmi.vstore %y, %y_ub[%off] + : !pto.vmi.vreg<1024xf8E4M3FN>, !pto.ptr + } + } + return + } +} diff --git a/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.py b/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.py new file mode 100644 index 0000000000..61d3a341d0 --- /dev/null +++ b/docs/repro/pipelining/fixtures/isolate_stages1_blockk1024_vmi.py @@ -0,0 +1,35 @@ +# 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. + +"""Isolate cell: stages=1, block_k=1024 — wide tile without dual-buffer stages.""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def persistent_quant_stages1_k1024(*, BLOCK_COUNT: pto.const_expr = 4): + x_ub = pto.alloc_tile(shape=[1, 1024], dtype=pto.bf16) + y_ub = pto.alloc_tile(shape=[1, 1024], dtype=pto.f8e4m3) + scale_ub = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off0 = pto.const(0, dtype=pto.index) + scale_slot = pto.vmi.vload(scale_ub.as_ptr(), off0, size=1) + scale = pto.vmi.vbrc(scale_slot, size=1024) + for i in range(BLOCK_COUNT): + off = i * 1024 + x_bf16 = pto.vmi.vload(x_ub.as_ptr(), off, size=1024) + x_f32 = pto.vmi.vcvt(x_bf16, pto.f32) + scaled = pto.vmi.vmul(x_f32, scale) + y = pto.vmi.vcvt(scaled, pto.f8e4m3, saturate="SAT") + pto.vmi.vstore(y, y_ub.as_ptr(), off) + + +if __name__ == "__main__": + try: + print(persistent_quant_stages1_k1024.compile().mlir_text()) + except Exception as exc: # noqa: BLE001 + print(f"emit/lower may fail (layout): {exc}") diff --git a/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.pto b/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.pto new file mode 100644 index 0000000000..0f6c918a9b --- /dev/null +++ b/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.pto @@ -0,0 +1,62 @@ +// 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. + +// Isolate: stages=2, block_k=512 — dual-buffer without 1024-wide layout. + +module attributes { + pto.target_arch = "a5", + pto.kernel_kind = #pto.kernel_kind, + pto.pipeline.stages = 2 : i32, + pto.block_k = 512 : i32 +} { + func.func @persistent_quant_stages2_k512( + %x_ub0: !pto.ptr, + %x_ub1: !pto.ptr, + %y_ub0: !pto.ptr, + %y_ub1: !pto.ptr, + %scale_slot: !pto.vmi.vreg<1xf32>, + %block_count: index) + attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c512 = arith.constant 512 : index + + pto.vecscope { + %scale = pto.vmi.vbrc %scale_slot + : !pto.vmi.vreg<1xf32> -> !pto.vmi.vreg<512xf32> + scf.for %i = %c0 to %block_count step %c1 { + %off = arith.muli %i, %c512 : index + %stage = arith.remui %i, %c2 : index + %use1 = arith.cmpi eq, %stage, %c1 : index + %x_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %x_ub1 : !pto.ptr + } else { + scf.yield %x_ub0 : !pto.ptr + } + %y_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %y_ub1 : !pto.ptr + } else { + scf.yield %y_ub0 : !pto.ptr + } + %x_bf16 = pto.vmi.vload %x_ub[%off] + : !pto.ptr -> !pto.vmi.vreg<512xbf16> + %x_f32 = pto.vmi.vcvt %x_bf16 + : !pto.vmi.vreg<512xbf16> -> !pto.vmi.vreg<512xf32> + %scaled = pto.vmi.vmul %x_f32, %scale + : !pto.vmi.vreg<512xf32>, !pto.vmi.vreg<512xf32> + -> !pto.vmi.vreg<512xf32> + %y = pto.vmi.vcvt %scaled {saturate = "SAT"} + : !pto.vmi.vreg<512xf32> -> !pto.vmi.vreg<512xf8E4M3FN> + pto.vmi.vstore %y, %y_ub[%off] + : !pto.vmi.vreg<512xf8E4M3FN>, !pto.ptr + } + } + return + } +} diff --git a/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.py b/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.py new file mode 100644 index 0000000000..ca7e01e7d3 --- /dev/null +++ b/docs/repro/pipelining/fixtures/isolate_stages2_blockk512_vmi.py @@ -0,0 +1,38 @@ +# 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. + +"""Isolate cell: stages=2, block_k=512 — dual-buffer without wide-tile layout.""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def persistent_quant_stages2_k512(*, BLOCK_COUNT: pto.const_expr = 4): + x_ub0 = pto.alloc_tile(shape=[1, 512], dtype=pto.bf16) + x_ub1 = pto.alloc_tile(shape=[1, 512], dtype=pto.bf16) + y_ub0 = pto.alloc_tile(shape=[1, 512], dtype=pto.f8e4m3) + y_ub1 = pto.alloc_tile(shape=[1, 512], dtype=pto.f8e4m3) + scale_ub = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off0 = pto.const(0, dtype=pto.index) + scale_slot = pto.vmi.vload(scale_ub.as_ptr(), off0, size=1) + scale = pto.vmi.vbrc(scale_slot, size=512) + for i in range(BLOCK_COUNT): + off = i * 512 + if (i % 2) == 1: + x_ptr, y_ptr = x_ub1.as_ptr(), y_ub1.as_ptr() + else: + x_ptr, y_ptr = x_ub0.as_ptr(), y_ub0.as_ptr() + x_bf16 = pto.vmi.vload(x_ptr, off, size=512) + x_f32 = pto.vmi.vcvt(x_bf16, pto.f32) + scaled = pto.vmi.vmul(x_f32, scale) + y = pto.vmi.vcvt(scaled, pto.f8e4m3, saturate="SAT") + pto.vmi.vstore(y, y_ptr, off) + + +if __name__ == "__main__": + print(persistent_quant_stages2_k512.compile().mlir_text()) diff --git a/docs/repro/pipelining/fixtures/lowered_vpto.pto b/docs/repro/pipelining/fixtures/lowered_vpto.pto new file mode 100644 index 0000000000..7bcff4e15e --- /dev/null +++ b/docs/repro/pipelining/fixtures/lowered_vpto.pto @@ -0,0 +1,100 @@ +module attributes {pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @persistent_quant_tile_k512(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.vreg<64xf32>, %arg3: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c512 = arith.constant 512 : index + pto.vecscope { + %0 = pto.pset_b32 "PAT_ALL" : !pto.mask + %1 = pto.vdup %arg2, %0 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %2 = pto.pset_b32 "PAT_ALL" : !pto.mask + %3 = pto.vdup %arg2, %2 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %4 = pto.pset_b32 "PAT_ALL" : !pto.mask + %5 = pto.vdup %arg2, %4 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %6 = pto.pset_b32 "PAT_ALL" : !pto.mask + %7 = pto.vdup %arg2, %6 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %9 = pto.vdup %arg2, %8 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %11 = pto.vdup %arg2, %10 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %13 = pto.vdup %arg2, %12 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %15 = pto.vdup %arg2, %14 {position = "LOWEST"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + scf.for %arg4 = %c0 to %arg3 step %c1 { + %16 = arith.muli %arg4, %c512 : index + %result = pto.vlds %arg0[%16] : !pto.ptr -> !pto.vreg<128xbf16> + %c128 = arith.constant 128 : index + %17 = arith.addi %16, %c128 : index + %result_0 = pto.vlds %arg0[%17] : !pto.ptr -> !pto.vreg<128xbf16> + %c256 = arith.constant 256 : index + %18 = arith.addi %16, %c256 : index + %result_1 = pto.vlds %arg0[%18] : !pto.ptr -> !pto.vreg<128xbf16> + %c384 = arith.constant 384 : index + %19 = arith.addi %16, %c384 : index + %result_2 = pto.vlds %arg0[%19] : !pto.ptr -> !pto.vreg<128xbf16> + %20 = pto.pset_b16 "PAT_ALL" : !pto.mask + %21 = pto.vcvt %result, %20 {part = "EVEN"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %22 = pto.vcvt %result_0, %20 {part = "EVEN"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vcvt %result_1, %20 {part = "EVEN"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %24 = pto.vcvt %result_2, %20 {part = "EVEN"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %25 = pto.vcvt %result, %20 {part = "ODD"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vcvt %result_0, %20 {part = "ODD"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %27 = pto.vcvt %result_1, %20 {part = "ODD"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %28 = pto.vcvt %result_2, %20 {part = "ODD"} : !pto.vreg<128xbf16>, !pto.mask -> !pto.vreg<64xf32> + %low, %high = pto.vintlv %21, %25 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_3, %high_4 = pto.vintlv %22, %26 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_5, %high_6 = pto.vintlv %23, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_7, %high_8 = pto.vintlv %24, %28 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_9, %high_10 = pto.vdintlv %low, %high : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_11, %high_12 = pto.vdintlv %low_3, %high_4 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_13, %high_14 = pto.vdintlv %low_9, %low_11 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_15, %high_16 = pto.vdintlv %high_10, %high_12 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_17, %high_18 = pto.vdintlv %low_5, %high_6 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_19, %high_20 = pto.vdintlv %low_7, %high_8 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_21, %high_22 = pto.vdintlv %low_17, %low_19 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %low_23, %high_24 = pto.vdintlv %high_18, %high_20 : !pto.vreg<64xf32>, !pto.vreg<64xf32> -> !pto.vreg<64xf32>, !pto.vreg<64xf32> + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %30 = pto.vmul %low_13, %1, %29 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %31 = pto.pset_b32 "PAT_ALL" : !pto.mask + %32 = pto.vmul %low_21, %3, %31 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %34 = pto.vmul %low_15, %5, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.pset_b32 "PAT_ALL" : !pto.mask + %36 = pto.vmul %low_23, %7, %35 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %37 = pto.pset_b32 "PAT_ALL" : !pto.mask + %38 = pto.vmul %high_14, %9, %37 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.pset_b32 "PAT_ALL" : !pto.mask + %40 = pto.vmul %high_22, %11, %39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %41 = pto.pset_b32 "PAT_ALL" : !pto.mask + %42 = pto.vmul %high_16, %13, %41 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.pset_b32 "PAT_ALL" : !pto.mask + %44 = pto.vmul %high_24, %15, %43 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %45 = pto.pset_b32 "PAT_ALL" : !pto.mask + %46 = pto.pset_b8 "PAT_ALL" : !pto.mask + %47 = pto.vcvt %30, %45 {part = "P0", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %48 = pto.vcvt %34, %45 {part = "P1", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %49 = pto.vcvt %38, %45 {part = "P2", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %50 = pto.vcvt %42, %45 {part = "P3", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %51 = pto.vor %47, %48, %46 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %52 = pto.vor %51, %49, %46 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %53 = pto.vor %52, %50, %46 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %54 = pto.pset_b8 "PAT_ALL" : !pto.mask + %55 = pto.vcvt %32, %45 {part = "P0", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %56 = pto.vcvt %36, %45 {part = "P1", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %57 = pto.vcvt %40, %45 {part = "P2", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %58 = pto.vcvt %44, %45 {part = "P3", rnd = "R", sat = "SAT"} : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %59 = pto.vor %55, %56, %54 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %60 = pto.vor %59, %57, %54 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %61 = pto.vor %60, %58, %54 : !pto.vreg<256xf8E4M3FN>, !pto.vreg<256xf8E4M3FN>, !pto.mask -> !pto.vreg<256xf8E4M3FN> + %62 = pto.pset_b8 "PAT_ALL" : !pto.mask + pto.vsts %53, %arg1[%16], %62 : !pto.vreg<256xf8E4M3FN>, !pto.ptr, !pto.mask + %63 = pto.pset_b8 "PAT_ALL" : !pto.mask + %c256_25 = arith.constant 256 : index + %64 = arith.addi %16, %c256_25 : index + pto.vsts %61, %arg1[%64], %63 : !pto.vreg<256xf8E4M3FN>, !pto.ptr, !pto.mask + } + } + return + } +} + diff --git a/docs/repro/pipelining/fixtures/reference_asc_cce.asc b/docs/repro/pipelining/fixtures/reference_asc_cce.asc new file mode 100644 index 0000000000..a7989be4b7 --- /dev/null +++ b/docs/repro/pipelining/fixtures/reference_asc_cce.asc @@ -0,0 +1,58 @@ +// ASC/CCE working baseline for gap 05: stages=2 ping-pong + UNPK/vmul/PK body. +// Compiles with bisheng --cce-aicore-only (see README). +#include "kernel_operator.h" + +__simd_vf__ inline void vf_scale_mul(__ubuf__ bfloat16_t* x_ub, __ubuf__ float* scale_ub, + __ubuf__ bfloat16_t* out_ub) { + vector_bool v16 = pset_b16(PAT_ALL); + vector_bool v32 = pset_b32(PAT_ALL); + vector_bf16 xb; + vlds(xb, x_ub, 0, UNPK_B16); + vector_f32 xf; + vcvt(xf, xb, v16, PART_EVEN, MODE_ZEROING); + vector_f32 sc; + vlds(sc, scale_ub, 0, BRC_B32); + vector_f32 yf; + vmul(yf, xf, sc, v32, MODE_ZEROING); + vector_bf16 yb; + vcvt(yb, yf, v16, ROUND_R, RS_ENABLE, PART_EVEN, MODE_ZEROING); + vsts(yb, out_ub, 0, PK_B32, v16); +} + +extern "C" __global__ __vector__ void ref_gap05_stages2(__gm__ bfloat16_t* /*x*/, + __gm__ float* /*sf*/, + __gm__ bfloat16_t* /*out*/) { + AscendC::InitSocState(); + __ubuf__ uint8_t* ub = (__ubuf__ uint8_t*)0; + + AscendC::SetFlag(0); + AscendC::SetFlag(1); + AscendC::SetFlag(0); + AscendC::SetFlag(1); + + constexpr int kBlockK = 1024; + constexpr int kStages = 2; + for (int w = 0; w < 4; ++w) { + int slot = w & 1; + AscendC::WaitFlag(slot); + // Producer would MTE2 into the ping-pong UB slot for this stage. + AscendC::SetFlag(slot); + AscendC::WaitFlag(slot); + AscendC::WaitFlag(slot); + + __ubuf__ bfloat16_t* x_ub = (__ubuf__ bfloat16_t*)(ub + slot * kBlockK * 2); + __ubuf__ float* sc_ub = (__ubuf__ float*)(ub + kStages * kBlockK * 2); + __ubuf__ bfloat16_t* o_ub = (__ubuf__ bfloat16_t*)(ub + slot * kBlockK * 2 + 4096); + vf_scale_mul(x_ub, sc_ub, o_ub); + + AscendC::SetFlag(slot); + AscendC::SetFlag(slot); + AscendC::WaitFlag(slot); + AscendC::SetFlag(slot); + } + + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); + AscendC::WaitFlag(0); + AscendC::WaitFlag(1); +} diff --git a/docs/repro/pipelining/fixtures/target_mi.pto b/docs/repro/pipelining/fixtures/target_mi.pto new file mode 100644 index 0000000000..39407b79f1 --- /dev/null +++ b/docs/repro/pipelining/fixtures/target_mi.pto @@ -0,0 +1,84 @@ +// 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. + +// Desired MI: stages=2 / block_k=1024 Persistent schedule — dual UB ping-pong, +// pipe-event sync on slot = i % 2, and UNPK / mul / PK4 VF body per tile. +// Mirrors the ASC/CCE reference_asc_cce.asc HardEvent shell. + +module attributes { + pto.target_arch = "a5", + pto.kernel_kind = #pto.kernel_kind, + pto.pipeline.stages = 2 : i32, + pto.block_k = 1024 : i32 +} { + func.func @persistent_quant_tile_k1024_mi( + %x_ub0: !pto.ptr, + %x_ub1: !pto.ptr, + %y_ub0: !pto.ptr, + %y_ub1: !pto.ptr, + %scale: !pto.vreg<64xf32>, + %block_count: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c1024 = arith.constant 1024 : index + + // Prime dual-slot events (ASC: SetFlag MTE3_V / V_MTE2 on slots 0 and 1). + pto.set_flag["PIPE_MTE3", "PIPE_V", "EVENT_ID0"] + pto.set_flag["PIPE_MTE3", "PIPE_V", "EVENT_ID1"] + pto.set_flag["PIPE_V", "PIPE_MTE2", "EVENT_ID0"] + pto.set_flag["PIPE_V", "PIPE_MTE2", "EVENT_ID1"] + + scf.for %i = %c0 to %block_count step %c1 { + %slot = arith.remui %i, %c2 : index + %use1 = arith.cmpi eq, %slot, %c1 : index + %off = arith.muli %i, %c1024 : index + + // Wait producer-ready / consumer-done for this ping-pong slot. + pto.wait_flag_dyn [#pto.pipe, #pto.pipe, %slot] + // (Host MTE2 would fill the selected UB slot here.) + pto.set_flag_dyn [#pto.pipe, #pto.pipe, %slot] + pto.wait_flag_dyn [#pto.pipe, #pto.pipe, %slot] + pto.wait_flag_dyn [#pto.pipe, #pto.pipe, %slot] + + %x_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %x_ub1 : !pto.ptr + } else { + scf.yield %x_ub0 : !pto.ptr + } + %y_ub = scf.if %use1 -> (!pto.ptr) { + scf.yield %y_ub1 : !pto.ptr + } else { + scf.yield %y_ub0 : !pto.ptr + } + + pto.vecscope { + %m32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %x0 = pto.vlds %x_ub[%off] {dist = "UNPK_B16"} + : !pto.ptr -> !pto.vreg<64xf32> + %y0 = pto.vmul %x0, %scale, %m32 + : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask + -> !pto.vreg<64xf32> + pto.vsts %y0, %y_ub[%off], %m32 {dist = "PK4_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + } + + pto.set_flag_dyn [#pto.pipe, #pto.pipe, %slot] + pto.set_flag_dyn [#pto.pipe, #pto.pipe, %slot] + pto.wait_flag_dyn [#pto.pipe, #pto.pipe, %slot] + pto.set_flag_dyn [#pto.pipe, #pto.pipe, %slot] + } + + // Drain both slots. + pto.wait_flag["PIPE_MTE3", "PIPE_V", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE3", "PIPE_V", "EVENT_ID1"] + pto.wait_flag["PIPE_V", "PIPE_MTE2", "EVENT_ID0"] + pto.wait_flag["PIPE_V", "PIPE_MTE2", "EVENT_ID1"] + return + } +} diff --git a/docs/repro/pipelining/scripts/check_stages.sh b/docs/repro/pipelining/scripts/check_stages.sh new file mode 100755 index 0000000000..33ec55def5 --- /dev/null +++ b/docs/repro/pipelining/scripts/check_stages.sh @@ -0,0 +1,186 @@ +#!/usr/bin/env bash +# Compile-check CCE stages=2 vs VMI stages×block_k isolate matrix (kernel-agnostic). +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +KERNEL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +FIXTURES="${KERNEL_ROOT}/fixtures" +OUT="${KERNEL_ROOT}/sim_outputs/check_stages" +mkdir -p "${OUT}" +LOG="${OUT}/compile_results.txt" + +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/env.sh" + +NPU_ARCH="${ASCEND_NPU_ARCH:-dav-3510}" + +: > "${LOG}" +echo "pipelining check_stages — $(date -u +%Y-%m-%dT%H:%M:%SZ)" | tee -a "${LOG}" +echo "PTOAS_ROOT=${PTOAS_ROOT}" | tee -a "${LOG}" +echo "BISHENG=${BISHENG}" | tee -a "${LOG}" +echo | tee -a "${LOG}" + +compile_asc() { + local name="$1" + local src="${FIXTURES}/${name}" + local obj="${OUT}/${name%.asc}.o" + local logf="${OUT}/${name%.asc}.log" + echo "=== ${name} (bisheng --cce-aicore-only) ===" | tee -a "${LOG}" + set +e + "${BISHENG}" -O2 -fPIC -std=c++17 --npu-arch="${NPU_ARCH}" --cce-aicore-only -c \ + "${src}" -o "${obj}" \ + -I"${ASCEND}/include" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/impl" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/interface" \ + > "${logf}" 2>&1 + local rc=$? + set -e + if [ "${rc}" -eq 0 ] && [ -s "${obj}" ]; then + echo "PASS: ${name} -> non-empty .o" | tee -a "${LOG}" + else + echo "FAIL: ${name} compile exit ${rc}" | tee -a "${LOG}" + tail -20 "${logf}" | tee -a "${LOG}" + fi + echo | tee -a "${LOG}" +} + +emit_vpto() { + local name="$1" + local expect="$2" # pass|fail + local src="${FIXTURES}/${name}" + local outf="${OUT}/${name%.pto}.vpto" + local logf="${OUT}/${name%.pto}_emit.log" + echo "=== ${name} (ptoas --emit-vpto) expect=${expect} ===" | tee -a "${LOG}" + set +e + if command -v ptoas >/dev/null 2>&1; then + ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto \ + --emit-vpto "${src}" -o "${outf}" > "${logf}" 2>&1 + local rc=$? + else + echo "SKIP: ptoas not on PATH" | tee -a "${LOG}" + rc=127 + fi + set -e + if [ "${rc}" -eq 0 ] && [ -s "${outf}" ]; then + if [ "${expect}" = "pass" ]; then + echo "PASS: ${name} --emit-vpto" | tee -a "${LOG}" + else + echo "PASS: ${name} --emit-vpto (unexpected — gap may be closed)" | tee -a "${LOG}" + fi + else + if [ "${expect}" = "fail" ]; then + echo "FAIL: ${name} --emit-vpto exit ${rc} (expected layout/lower reject)" | tee -a "${LOG}" + else + echo "FAIL: ${name} --emit-vpto exit ${rc}" | tee -a "${LOG}" + fi + tail -12 "${logf}" | tee -a "${LOG}" + fi + echo | tee -a "${LOG}" +} + +ptodsl_py() { + local name="$1" + local outf="${OUT}/${name%.py}.mlir" + local errf="${OUT}/${name%.py}.err" + echo "=== ${name} (ptodsl frontend) ===" | tee -a "${LOG}" + set +e + python3 "${FIXTURES}/${name}" > "${outf}" 2> "${errf}" + local rc=$? + set -e + if [ "${rc}" -eq 0 ]; then + echo "PASS: ${name} emitted MLIR (frontend only)" | tee -a "${LOG}" + else + echo "FAIL: ${name} exit ${rc}" | tee -a "${LOG}" + tail -20 "${errf}" | tee -a "${LOG}" + fi + echo | tee -a "${LOG}" +} + +# --- CCE baseline --- +compile_asc "reference_asc_cce.asc" + +echo "=== device/scale_stages2.asc (bisheng --cce-aicore-only) ===" | tee -a "${LOG}" +set +e +"${BISHENG}" -O2 -fPIC -std=c++17 --npu-arch="${NPU_ARCH}" --cce-aicore-only -c \ + "${KERNEL_ROOT}/device/scale_stages2.asc" -o "${OUT}/scale_stages2_aicore.o" \ + -I"${ASCEND}/include" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/impl" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/interface" \ + > "${OUT}/scale_stages2_aicore.log" 2>&1 +rc=$? +set -e +if [ "${rc}" -eq 0 ] && [ -s "${OUT}/scale_stages2_aicore.o" ]; then + echo "PASS: device/scale_stages2.asc -> non-empty .o" | tee -a "${LOG}" +else + echo "FAIL: device/scale_stages2.asc compile exit ${rc}" | tee -a "${LOG}" + tail -20 "${OUT}/scale_stages2_aicore.log" | tee -a "${LOG}" +fi +echo | tee -a "${LOG}" + +# --- VMI matrix: stages × block_k --- +ptodsl_py "current_slow_vmi.py" +emit_vpto "current_slow_vmi.pto" "pass" + +echo "=== lowered_vpto.pto (stages=1 chunked MI -> LLVM IR) ===" | tee -a "${LOG}" +set +e +ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto \ + --emit-vpto-llvm-ir "${FIXTURES}/lowered_vpto.pto" -o "${OUT}/lowered_vpto.ll" \ + > "${OUT}/lowered_vpto_ptoas.log" 2>&1 +lowered_rc=$? +set -e +if [ "${lowered_rc}" -eq 0 ] && [ -s "${OUT}/lowered_vpto.ll" ]; then + echo "PASS: lowered_vpto.pto -> LLVM IR" | tee -a "${LOG}" +else + echo "FAIL: lowered_vpto.pto exit ${lowered_rc}" | tee -a "${LOG}" + tail -15 "${OUT}/lowered_vpto_ptoas.log" | tee -a "${LOG}" +fi +echo | tee -a "${LOG}" + +ptodsl_py "isolate_stages2_blockk512_vmi.py" +emit_vpto "isolate_stages2_blockk512_vmi.pto" "pass" + +ptodsl_py "isolate_stages1_blockk1024_vmi.py" +emit_vpto "isolate_stages1_blockk1024_vmi.pto" "fail" + +ptodsl_py "desired_vmi.py" +emit_vpto "desired_vmi.pto" "fail" + +# target_mi: emit OK, bisheng object often crashes +echo "=== target_mi.pto (emit-vpto + LLVM IR + bisheng .o) ===" | tee -a "${LOG}" +set +e +ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto \ + --emit-vpto "${FIXTURES}/target_mi.pto" -o "${OUT}/target_mi.vpto" \ + > "${OUT}/target_mi_emit_vpto.log" 2>&1 +emit_rc=$? +set -e +if [ "${emit_rc}" -eq 0 ] && [ -s "${OUT}/target_mi.vpto" ]; then + echo "PASS: target_mi.pto --emit-vpto" | tee -a "${LOG}" +else + echo "FAIL: target_mi.pto --emit-vpto exit ${emit_rc}" | tee -a "${LOG}" +fi +set +e +ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto \ + --emit-vpto-llvm-ir "${FIXTURES}/target_mi.pto" -o "${OUT}/target_mi.ll" \ + > "${OUT}/target_mi_ptoas.log" 2>&1 +ll_rc=$? +set -e +if [ "${ll_rc}" -eq 0 ]; then + echo "PASS: target_mi.pto --emit-vpto-llvm-ir" | tee -a "${LOG}" + set +e + "${BISHENG}" -O2 -fPIC -std=c++17 --npu-arch="${NPU_ARCH}" -c -x ir \ + "${OUT}/target_mi.ll" -o "${OUT}/target_mi.o" > "${OUT}/target_mi_bisheng.log" 2>&1 + b_rc=$? + set -e + if [ "${b_rc}" -eq 0 ] && [ -s "${OUT}/target_mi.o" ]; then + echo "PASS: bisheng object (unexpected — gap may be closed)" | tee -a "${LOG}" + else + echo "FAIL: bisheng object exit ${b_rc} (expected on vmi-v0.1.3)" | tee -a "${LOG}" + tail -8 "${OUT}/target_mi_bisheng.log" | tee -a "${LOG}" + fi +else + echo "FAIL: target_mi llvm-ir exit ${ll_rc}" | tee -a "${LOG}" +fi + +echo | tee -a "${LOG}" +echo "Full log: ${LOG}" | tee -a "${LOG}" diff --git a/docs/repro/pipelining/scripts/env.sh b/docs/repro/pipelining/scripts/env.sh new file mode 100755 index 0000000000..e0b682a9d9 --- /dev/null +++ b/docs/repro/pipelining/scripts/env.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# CANN + PTOAS env for pipelining stage compile checks. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +KERNEL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +ASCEND_HOME_PATH="${ASCEND_HOME_PATH:-${ASCEND_TOOLKIT_HOME:-/usr/local/Ascend/cann-9.0.0}}" + +# pipelining -> repro -> docs -> repo (three levels up from repro root). +PTOAS_ROOT="${PTOAS_ROOT:-$(cd "${KERNEL_ROOT}/../../.." && pwd)}" + +# shellcheck disable=SC1091 +source "${ASCEND_HOME_PATH}/bin/setenv.bash" +if [ -n "${PTOAS_ROOT}" ] && [ -f "${PTOAS_ROOT}/scripts/ptoas_env.sh" ]; then + export PTOAS_ENV_SKIP_SMOKE_TEST="${PTOAS_ENV_SKIP_SMOKE_TEST:-1}" + # shellcheck disable=SC1091 + source "${PTOAS_ROOT}/scripts/ptoas_env.sh" + # Prefer build/tools wrappers over install/bin (install wrappers often miss MLIR libs). + [ -d "${PTOAS_ROOT}/install/bin" ] && export PATH="${PTOAS_ROOT}/install/bin:${PATH}" + [ -d "${PTOAS_ROOT}/build/tools/ptoas" ] && export PATH="${PTOAS_ROOT}/build/tools/ptoas:${PATH}" +fi + +# Optional: point at a known-good ptoas tree (e.g. built at tag vmi-v0.1.3). +if [ -n "${PTOAS_TOOLS_ROOT:-}" ]; then + [ -d "${PTOAS_TOOLS_ROOT}/build/tools/ptoas" ] && export PATH="${PTOAS_TOOLS_ROOT}/build/tools/ptoas:${PATH}" +fi + +# Prefer in-tree ptodsl when present (else site-packages). +if [ -d "${PTOAS_ROOT}/ptodsl" ]; then + export PYTHONPATH="${PTOAS_ROOT}/ptodsl${PYTHONPATH:+:${PYTHONPATH}}" +fi + +export ASCEND_NPU_ARCH="${ASCEND_NPU_ARCH:-dav-3510}" +export BISHENG="${BISHENG:-${ASCEND_HOME_PATH}/tools/bisheng_compiler/bin/bisheng}" +export ASCEND="${ASCEND_HOME_PATH}"