diff --git a/.gitignore b/.gitignore index a2cfb9d562..43a15d52a6 100644 --- a/.gitignore +++ b/.gitignore @@ -99,3 +99,12 @@ test/samples/**/npu_validation/* .planning/ .work/ .local/ + +# local cannsim / build junk under docs/repro +docs/repro/**/build/ +docs/repro/**/sim_outputs/ +docs/repro/**/*.so +docs/repro/**/traces/ +docs/repro/**/*.zip +docs/repro/**/route_table.txt +docs/repro/**/rtb_debug.txt diff --git a/docs/repro/store_pad8/.gitignore b/docs/repro/store_pad8/.gitignore new file mode 100644 index 0000000000..999aae588e --- /dev/null +++ b/docs/repro/store_pad8/.gitignore @@ -0,0 +1,10 @@ +build/ +sim_outputs/ +*.so +*.asc +log_ca/ +instr.bin +__pycache__/ +*.pyc +.pytest_cache/ +*.egg-info/ diff --git a/docs/repro/store_pad8/README.md b/docs/repro/store_pad8/README.md new file mode 100644 index 0000000000..a68535e93a --- /dev/null +++ b/docs/repro/store_pad8/README.md @@ -0,0 +1,113 @@ +# store_pad8 — VMI pad-8 reduce store repro + +Self-contained cannsim microbench for PTOAS gap 03: after a cross-lane +`vcadd`, CCE stores one f32 scalar per group with `vsts ... ONEPT_B32`. +Product VMI kernels often broadcast the reduced value to 8 lanes (`vbrc`) +and store under an 8-wide mask. + +This repro isolates that epilogue and compares **kernel formulations** on a +built PTOAS pin (branch `zjw/store_pad_issue`, tag lineage **vmi-v0.1.3**). + +## What to compare + +| Path | Store shape | Notes | +|------|-------------|-------| +| CCE (`cce/`) | ONEPT / 1 scalar | Baseline — see `fixtures/reference_asc_cce.asc` | +| VMI pad-8 (`vmi/store_pad8_vmi.py`) | pad-8 + mask8 | Product workaround | +| VMI mask1 (`vmi/store_mask1_vmi.py`) | compact `[N]` + mask1 | Masked 1-lane store | +| VMI **group1** (`vmi/store_group1_vmi.py`) | `vcadd(group=1)` + `vstore(group=1)` | Compact score → point store / `1PT_B32` | +| VMI desired (`fixtures/desired_vmi.*`) | mask-1 after `vcadd` | Compile-only single-tile fixture | +| VMI target MI (`fixtures/target_mi.pto`) | `dist = "1PT_B32"` | Hand-written MI; bisheng may still crash | + +**Best VMI path:** `group1`. It lowers to `vsts ... {dist = "1PT_B32"}` and +matches CCE’s `vcadd` + scalar store shape in cannsim (see table below). + +## Layout + +``` +store_pad8/ + cce/ CCE ONEPT kernel + ctypes launcher + vmi/ VMI pad-8 / mask1 / group1 kernels (ptodsl) + common/ torch runtime, golden ref, build helper, top launcher + fixtures/ IR dumps and reference ASC from gap 03 + test/ correctness test + scripts/ env, cannsim runners, check_desired, run_three_way +``` + +## Prerequisites + +- CANN 9.0 simulator (`cannsim`, `bisheng`) +- Built PTOAS from this repo (`./quick_install.sh` with + `-DPTOAS_ENABLE_WERROR=OFF` if needed for older gcc). LLVM/MLIR comes from + the docker image. Install the produced wheel so `ptoas` on `PATH` is the + wheel launcher (not only `build/tools/ptoas`). +- `torch_npu` for cannsim host launch + +Push remote hint: https://github.com/learning-chip/PTOAS.git + +## Run + +From this directory (`docs/repro/store_pad8/`): + +```bash +source scripts/env.sh + +# Four-way: CCE + VMI pad8 / mask1 / group1 +STORE_PAD8_CASE=large ./scripts/run_three_way.sh +STORE_PAD8_CASE=small ./scripts/run_three_way.sh + +# Or one backend +STORE_PAD8_CASE=large ./scripts/run_cannsim.sh cce +STORE_PAD8_VARIANT=group1 STORE_PAD8_CASE=large \ + ./scripts/run_sim.sh test/test_store_pad8.py sim_outputs/store_pad8_large_vmi_group1 + +./scripts/check_desired.sh +``` + +Correctness only (no cannsim): + +```bash +source scripts/env.sh +TLVF_VMI_BACKEND=cce STORE_PAD8_CASE=large python3 test/test_store_pad8.py +TLVF_VMI_BACKEND=vmi STORE_PAD8_VARIANT=group1 STORE_PAD8_CASE=large python3 test/test_store_pad8.py +``` + +## Measured RVEC (Ascend950 cannsim, 2026-07-29) + +Built PTOAS from this branch (wheel `ptoas 0.52`). All paths PASS +(`maxDiff≈1e-6`). + +### N_ACC=20 (LARGE) + +| Backend | RVEC span | vs CCE | Top ops (approx.) | +|---------|----------:|-------:|-------------------| +| CCE ONEPT | **63** | 1.00× | `RV_VLDI=20, RV_VCADD=20, RV_VSTI=20, RV_PSET=1` | +| VMI **group1** | **59** | **0.94×** | `RV_VCADD=20, RV_VSTI=20, RV_VLDI=16, RV_VLDS=4, RV_PSET=3` | +| VMI mask1 | **66** | 1.05× | `RV_VCADD=20, RV_VADD=20, RV_VSTS=17, RV_VLDI=16, RV_VLDS=4` | +| VMI pad8 | **99** | 1.57× | `RV_VCADD=20, RV_VADD=20, RV_VDUP=20, RV_VSTI=20, RV_VLDI=16` | + +### N_ACC=4 (SMALL) + +| Backend | RVEC span | vs CCE | Top ops (approx.) | +|---------|----------:|-------:|-------------------| +| CCE ONEPT | **38** | 1.00× | `RV_VLDI=4, RV_VCADD=4, RV_VSTI=4, RV_PSET=1` | +| VMI **group1** | **43** | 1.13× | `RV_VLDI=4, RV_VCADD=4, RV_VSTI=4, RV_PSET=3, RV_VDUP=1` | +| VMI mask1 | **43** | 1.13× | `RV_VLDI=4, RV_VCADD=4, RV_VADD=4, RV_PSET=3, RV_VSTS=3` | +| VMI pad8 | **53** | 1.39× | `RV_PSET=4, RV_VLDI=4, RV_VCADD=4, RV_VADD=4, RV_VDUP=4, RV_VSTI=4` | + +**Read:** `vcadd(group=1)` + `vstore(group=1, stride=1)` is the VMI form that +gets ONEPT/`1PT_B32` lowering. At N_ACC=20 it matches (slightly beats) CCE +RVEC. Pad-8 and mask1 remain useful as slower baselines. Older PyPI +`ptoas 0.51` without a built tree overstated the pad-8 tax (~281 RVEC). + +## Done when + +Product kernels use the group1 / compact-score store (or equivalent ONEPT +lowering) instead of pad-8. This repro already shows the fast path on the +built pin above. + +## Pin + +- PTOAS: this branch `zjw/store_pad_issue` (from tag **vmi-v0.1.3**), built + + wheel-installed (`ptoas 0.52` in the measured run) +- `PTOAS_ROOT` defaults to the repo root (`../../..` from this repro directory) diff --git a/docs/repro/store_pad8/cce/csrc/inc/common.h b/docs/repro/store_pad8/cce/csrc/inc/common.h new file mode 100644 index 0000000000..0bfad8022b --- /dev/null +++ b/docs/repro/store_pad8/cce/csrc/inc/common.h @@ -0,0 +1,22 @@ +#pragma once + +#include "acl/acl.h" +#include "kernel_operator.h" +#include "simt_api/asc_bf16.h" +#include "simt_api/asc_fp16.h" + +#include "numeric_limits.h" +#include "simd_inst.h" + +// Reduce N_ACC live f32 accumulator vectors (VL=64 lanes each) to N_ACC +// scalars on one core. N_ACC=20 (LARGE) matches a typical residual-mix +// epilogue. N_ACC=4 (SMALL) is a quick smoke case. +constexpr int32_t kVL = 64; +constexpr int32_t kMaxNAcc = 20; + +// UB layout (bytes): +// acc[N_ACC * kVL] f32 @0 +// reduced[N_ACC] f32 @kMaxNAcc*kVL*4 (compact ONEPT store target) +constexpr int32_t kAccUb = 0; +constexpr int32_t kReducedUb = kMaxNAcc * kVL * 4; +constexpr int32_t kUbBytes = kReducedUb + kMaxNAcc * 4; diff --git a/docs/repro/store_pad8/cce/csrc/inc/simd_inst.h b/docs/repro/store_pad8/cce/csrc/inc/simd_inst.h new file mode 100644 index 0000000000..6d7bf73654 --- /dev/null +++ b/docs/repro/store_pad8/cce/csrc/inc/simd_inst.h @@ -0,0 +1,39 @@ +#pragma once + +// Minimal CCE SIMD helpers for store_pad8_vf_sim_kernel.cpp. +// Only float vlds / vcadd / vsts — not a full TileLang simd_inst dump. + +#include "kernel_operator.h" + +namespace simd_inst { + +template struct vec {}; + +template <> struct vec { + using type = vector_f32; +}; + +template using vec_t = typename vec::type; + +template +__simd_callee__ inline vec_t vlds(__ubuf__ T *src, Offset offset, + Dist dist) { + vec_t dst; + ::vlds(dst, src, offset, dist); + return dst; +} + +template +__simd_callee__ inline SrcVec vcadd(SrcVec src, vector_bool mask, Mode mode) { + SrcVec dst; + ::vcadd(dst, src, mask, mode); + return dst; +} + +template +__simd_callee__ inline void vsts(vec_t data, __ubuf__ T *base, + int32_t offset, Dist dist, vector_bool mask) { + ::vsts(data, base, offset, dist, mask); +} + +} // namespace simd_inst diff --git a/docs/repro/store_pad8/cce/csrc/store_pad8_vf_sim_kernel.cpp b/docs/repro/store_pad8/cce/csrc/store_pad8_vf_sim_kernel.cpp new file mode 100644 index 0000000000..6c6e6b32af --- /dev/null +++ b/docs/repro/store_pad8/cce/csrc/store_pad8_vf_sim_kernel.cpp @@ -0,0 +1,61 @@ +/*=========================================================================== + * store_pad8 microbench — CCE ONEPT path + * + * Per accumulator: vcadd then a single scalar vsts ... ONEPT_B32 store. + * No padding, no broadcast — the CCE baseline the VMI pad-8 port + * (../vmi/store_pad8_vmi.py) is compared against. + *===========================================================================*/ + +#include "inc/common.h" + +namespace { + +template +__simd_vf__ inline void simd_vf_store_pad8_cce(__ubuf__ uint8_t *ub) { + auto *acc = (__ubuf__ float *)(ub + kAccUb); + auto *reduced = (__ubuf__ float *)(ub + kReducedUb); +#pragma unroll + for (int32_t i = 0; i < N_ACC; ++i) { + vector_bool v = pset_b32(PAT_ALL); + auto acc_v = simd_inst::vlds(acc + i * kVL, 0, NORM); + auto red = simd_inst::vcadd(acc_v, v, MODE_ZEROING); + simd_inst::vsts(red, reduced + i, 0, ONEPT_B32, v); + } +} + +} // namespace + +template +__global__ __vector__ void store_pad8_cce_kernel(__gm__ float *acc_g, + __gm__ float *reduced_g) { + AscendC::InitSocState(); + __ubuf__ uint8_t *ub = (__ubuf__ uint8_t *)0x0; + + copy_gm_to_ubuf_align_v2(ub + kAccUb, (__gm__ uint8_t *)acc_g, 0, 1, + N_ACC * kVL * 4, 0, 0, 0, 0, N_ACC * kVL * 4, + N_ACC * kVL * 4); + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + + simd_vf_store_pad8_cce(ub); + + AscendC::SetFlag(0); + AscendC::WaitFlag(0); + copy_ubuf_to_gm_align_v2((__gm__ void *)reduced_g, + (__ubuf__ void *)(ub + kReducedUb), 0, 1, N_ACC * 4, + 4, N_ACC * 4, N_ACC * 4); +} + +extern "C" { + +void call_store_pad8_cce_large(void *stream, void *acc, void *reduced) { + store_pad8_cce_kernel<20><<<1, kUbBytes, stream>>>( + (__gm__ float *)acc, (__gm__ float *)reduced); +} + +void call_store_pad8_cce_small(void *stream, void *acc, void *reduced) { + store_pad8_cce_kernel<4><<<1, kUbBytes, stream>>>( + (__gm__ float *)acc, (__gm__ float *)reduced); +} + +} // extern "C" diff --git a/docs/repro/store_pad8/cce/launcher.py b/docs/repro/store_pad8/cce/launcher.py new file mode 100644 index 0000000000..4e4422d9dd --- /dev/null +++ b/docs/repro/store_pad8/cce/launcher.py @@ -0,0 +1,50 @@ +"""ctypes launcher for the store_pad8 CCE VF sim kernel.""" + +from __future__ import annotations + +import ctypes +from pathlib import Path +import sys + +import torch + +_REPRO = Path(__file__).resolve().parent.parent +if str(_REPRO) not in sys.path: + sys.path.insert(0, str(_REPRO)) + +from common.torch_runtime import empty_npu, stream_ptr, sync # noqa: E402 + +_VL = 64 + +_LIB: ctypes.CDLL | None = None + + +def _lib() -> ctypes.CDLL: + global _LIB + if _LIB is None: + from common.cce_vf_build import build_cce_root + + path = build_cce_root(Path(__file__).resolve().parent) + lib = ctypes.CDLL(str(path)) + for name in ("call_store_pad8_cce_large", "call_store_pad8_cce_small"): + fn = getattr(lib, name) + fn.argtypes = [ctypes.c_void_p] * 3 + fn.restype = None + _LIB = lib + return _LIB + + +def _vp(t: torch.Tensor) -> ctypes.c_void_p: + return ctypes.c_void_p(t.data_ptr()) + + +def launch(acc: torch.Tensor, n_acc: int) -> torch.Tensor: + """acc: [n_acc, VL] f32 on npu -> [n_acc] f32 (compact ONEPT reduce).""" + lib = _lib() + name = "call_store_pad8_cce_large" if n_acc == 20 else "call_store_pad8_cce_small" + if n_acc not in (4, 20): + raise ValueError(f"unsupported n_acc={n_acc} (build only has large=20 / small=4)") + reduced = empty_npu((n_acc,), torch.float32) + getattr(lib, name)(stream_ptr(), _vp(acc), _vp(reduced)) + sync() + return reduced diff --git a/docs/repro/store_pad8/common/__init__.py b/docs/repro/store_pad8/common/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/docs/repro/store_pad8/common/cce_vf_build.py b/docs/repro/store_pad8/common/cce_vf_build.py new file mode 100644 index 0000000000..41b5ed3026 --- /dev/null +++ b/docs/repro/store_pad8/common/cce_vf_build.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""bisheng --shared build for store_pad8 *_vf_sim_kernel.cpp shells.""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import sys +from pathlib import Path + + +def _ascend_home() -> Path: + home = os.environ.get("ASCEND_HOME_PATH") or os.environ.get("ASCEND_TOOLKIT_HOME") + if not home: + raise EnvironmentError("ASCEND_HOME_PATH is not set.") + return Path(home) + + +def _bisheng() -> str: + candidate = _ascend_home() / "bin" / "bisheng" + if candidate.is_file(): + return str(candidate) + found = shutil.which("bisheng") + if found: + return found + raise FileNotFoundError("bisheng compiler not found") + + +def _npu_arch() -> str: + return os.environ.get("ASCEND_NPU_ARCH", "dav-3510") + + +def _repro_root(cce_root: Path) -> Path: + # /cce -> parent == store_pad8 repro root + return cce_root.resolve().parent + + +def _so_path(cce_root: Path, sources: list[Path]) -> Path: + build_dir = cce_root / "build" + if len(sources) == 1: + stem = sources[0].stem + if stem.endswith("_kernel"): + stem = stem[: -len("_kernel")] + return build_dir / f"lib{stem}.so" + return build_dir / "libstore_pad8_vf_sim.so" + + +def build_cce_root(cce_root: Path, force: bool = False) -> Path: + """Build the VF sim .so from all *_vf_sim_kernel.cpp under cce_root/csrc.""" + cce_root = cce_root.resolve() + csrc = cce_root / "csrc" + build_dir = cce_root / "build" + shared_inc = _repro_root(cce_root) / "common" / "inc" + sources = sorted(csrc.glob("*_vf_sim_kernel.cpp")) + if not sources: + raise FileNotFoundError(f"no *_vf_sim_kernel.cpp under {csrc}") + + so_path = _so_path(cce_root, sources) + build_dir.mkdir(parents=True, exist_ok=True) + if so_path.is_file() and not force: + return so_path + + asc_files: list[str] = [] + for src in sources: + asc = build_dir / f"{src.stem}.asc" + if asc.is_symlink() or asc.exists(): + asc.unlink() + asc.symlink_to(src.resolve()) + asc_files.append(str(asc)) + + cmd = [ + _bisheng(), + "-fPIC", + "--shared", + f"--npu-arch={_npu_arch()}", + "-std=c++17", + f"-I{shared_inc}", + f"-I{csrc}", + f"-I{csrc / 'inc'}", + "-O2", + "-Wno-ignored-attributes", + "-Wno-unknown-attributes", + "-Wno-macro-redefined", + *asc_files, + "-o", + str(so_path), + ] + print("==>", " ".join(cmd), flush=True) + subprocess.run(cmd, cwd=build_dir, check=True) + print(f"Built {so_path}", flush=True) + return so_path + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "cce_root", + nargs="?", + type=Path, + default=Path.cwd(), + help="Path to cce/ (default: cwd)", + ) + parser.add_argument("--force", action="store_true") + args = parser.parse_args(argv) + try: + build_cce_root(args.cce_root.resolve(), force=args.force) + except (EnvironmentError, FileNotFoundError, subprocess.CalledProcessError) as exc: + print(f"build failed: {exc}", file=sys.stderr) + raise SystemExit(1) from exc + + +if __name__ == "__main__": + main() diff --git a/docs/repro/store_pad8/common/golden.py b/docs/repro/store_pad8/common/golden.py new file mode 100644 index 0000000000..022634be1c --- /dev/null +++ b/docs/repro/store_pad8/common/golden.py @@ -0,0 +1,52 @@ +"""Torch reference for the store_pad8 microbench. + +Each of N_ACC live f32 accumulator vectors (length VL) is reduced to one +scalar. CCE stores compactly (vcadd + vsts ONEPT_B32). VMI pad-8 broadcasts +to 8 lanes under an 8-wide mask; VMI mask1 stores compactly under a 1-lane +mask (still not CCE ONEPT lowering on vmi-v0.1.3). + +N_ACC=20 (LARGE) matches a typical residual-mix epilogue row count. +N_ACC=4 (SMALL) is a quick smoke case. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +import torch + +VL = 64 +REDUCE_PAD = 8 +SEED = 123 + + +@dataclass(frozen=True) +class StorePad8Case: + name: str + n_acc: int + + +LARGE = StorePad8Case(name="large", n_acc=20) +SMALL = StorePad8Case(name="small", n_acc=4) + + +def reduce_ref(acc: torch.Tensor) -> torch.Tensor: + """acc: [N_ACC, VL] f32 -> [N_ACC] f32 (sum over lanes).""" + return acc.sum(dim=-1) + + +def generate_case(case: StorePad8Case) -> dict: + torch.manual_seed(SEED + case.n_acc) + acc = torch.randn((case.n_acc, VL), dtype=torch.float32) + reduced = reduce_ref(acc) + return { + "case": case, + "acc": acc.numpy(), + "reduced": reduced.numpy(), + } + + +def extract_padded(padded: np.ndarray, n_acc: int, pad: int = REDUCE_PAD) -> np.ndarray: + """VMI stores each scalar broadcast across pad lanes; take lane 0.""" + return np.asarray(padded, dtype=np.float32).reshape(n_acc, pad)[:, 0] diff --git a/docs/repro/store_pad8/common/inc/numeric_limits.h b/docs/repro/store_pad8/common/inc/numeric_limits.h new file mode 100644 index 0000000000..5f25590951 --- /dev/null +++ b/docs/repro/store_pad8/common/inc/numeric_limits.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +namespace fa { +namespace limits { + +// Use compiler builtin so device compiles can stay on C++17 (avoids +// clang-15 + libstdc++-12 breakage under -std=c++20 on some hosts). + +inline constexpr float kFloatInf = __builtin_bit_cast(float, 0x7F800000u); +inline constexpr float kFloatNInf = __builtin_bit_cast(float, 0xFF800000u); +inline constexpr float kFloatNaN = __builtin_bit_cast(float, 0x7FC00000u); + +} // namespace limits +} // namespace fa diff --git a/docs/repro/store_pad8/common/launcher.py b/docs/repro/store_pad8/common/launcher.py new file mode 100644 index 0000000000..ed7b5de4eb --- /dev/null +++ b/docs/repro/store_pad8/common/launcher.py @@ -0,0 +1,86 @@ +"""Dispatch the store_pad8 microbench to CCE or VMI.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parent.parent +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from common.torch_runtime import empty_npu, stream_ptr, sync # noqa: E402 + +_COMPILED: dict[str, object] = {} + +_COMPACT_VARIANTS = frozenset({"mask1", "group1"}) + + +def launch_cce(acc: torch.Tensor, n_acc: int) -> torch.Tensor: + import importlib.util + + path = ROOT / "cce" / "launcher.py" + spec = importlib.util.spec_from_file_location("store_pad8_cce_launcher", path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load CCE launcher from {path}") + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod.launch(acc, n_acc) + + +def launch_vmi(acc: torch.Tensor, n_acc: int, variant: str | None = None) -> torch.Tensor: + """VMI backends. + + variant: + pad8 — product workaround; output [N,8], host gathers lane 0 + mask1 — vcadd + 1-lane masked vstore; compact [N] out + group1 — vcadd(group=1) + vstore(group=1); compact [N] out (ONEPT hint) + """ + variant = (variant or os.environ.get("STORE_PAD8_VARIANT", "pad8")).lower() + if variant not in ("pad8", "mask1", "group1"): + raise ValueError(f"unsupported STORE_PAD8_VARIANT={variant!r}") + if n_acc not in (4, 20): + raise ValueError(f"unsupported n_acc={n_acc}") + + key = f"vmi:{variant}:{n_acc}" + if key not in _COMPILED: + sys.path.insert(0, str(ROOT / "vmi")) + if variant == "pad8": + from store_pad8_vmi import store_pad8_vmi_large, store_pad8_vmi_small + + kn = store_pad8_vmi_large if n_acc == 20 else store_pad8_vmi_small + elif variant == "mask1": + from store_mask1_vmi import store_mask1_vmi_large, store_mask1_vmi_small + + kn = store_mask1_vmi_large if n_acc == 20 else store_mask1_vmi_small + else: + from store_group1_vmi import store_group1_vmi_large, store_group1_vmi_small + + kn = store_group1_vmi_large if n_acc == 20 else store_group1_vmi_small + _COMPILED[key] = kn.compile() + + compiled = _COMPILED[key] + if variant == "pad8": + out = empty_npu((n_acc * 8,), torch.float32) + else: + out = empty_npu((n_acc,), torch.float32) + compiled[1, stream_ptr()](acc.data_ptr(), out.data_ptr()) + sync() + return out + + +def launch( + acc: torch.Tensor, + n_acc: int, + backend: str | None = None, + variant: str | None = None, +) -> torch.Tensor: + backend = (backend or os.environ.get("TLVF_VMI_BACKEND", "vmi")).lower() + if backend == "cce": + return launch_cce(acc, n_acc) + if backend == "vmi": + return launch_vmi(acc, n_acc, variant=variant) + raise ValueError(f"unsupported backend: {backend}") diff --git a/docs/repro/store_pad8/common/torch_runtime.py b/docs/repro/store_pad8/common/torch_runtime.py new file mode 100644 index 0000000000..99633a4191 --- /dev/null +++ b/docs/repro/store_pad8/common/torch_runtime.py @@ -0,0 +1,35 @@ +"""Shared torch_npu helpers for store_pad8 cannsim runs.""" + +from __future__ import annotations + +import os + +import torch +import torch_npu + +_DEVICE = f"npu:{os.environ.get('NPU_DEVICE', '0')}" + + +def device_str() -> str: + return _DEVICE + + +def init_torch_npu(device: str | None = None) -> None: + global _DEVICE + if device is not None: + _DEVICE = device + torch.npu.config.allow_internal_format = False + torch_npu.npu.set_compile_mode(jit_compile=False) + torch.npu.set_device(_DEVICE) + + +def empty_npu(shape, dtype: torch.dtype) -> torch.Tensor: + return torch.empty(shape, dtype=dtype, device=_DEVICE) + + +def stream_ptr() -> int: + return torch.npu.current_stream()._as_parameter_ + + +def sync() -> None: + torch.npu.synchronize() diff --git a/docs/repro/store_pad8/fixtures/current_slow_vmi.pto b/docs/repro/store_pad8/fixtures/current_slow_vmi.pto new file mode 100644 index 0000000000..782e54a785 --- /dev/null +++ b/docs/repro/store_pad8/fixtures/current_slow_vmi.pto @@ -0,0 +1,25 @@ +// 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: group-reduce into an 8-slot pad, then group-store. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @padded_reduce_store( + %tile: !pto.vmi.vreg<64xf32>, + %mask: !pto.vmi.mask<64xpred>, + %dst: !pto.ptr, + %off: index) { + %c1 = arith.constant 1 : index + %sum8 = pto.vmi.vcadd %tile, %mask {group = 8, reassoc} + : !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<8xf32> + pto.vmi.vstore %sum8, %dst[%off], %c1 {group = 8} + : !pto.vmi.vreg<8xf32>, !pto.ptr + return + } +} diff --git a/docs/repro/store_pad8/fixtures/current_slow_vmi.py b/docs/repro/store_pad8/fixtures/current_slow_vmi.py new file mode 100644 index 0000000000..9eea0b1c67 --- /dev/null +++ b/docs/repro/store_pad8/fixtures/current_slow_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. + +"""Python DSL form of ``current_slow_vmi.pto`` — group-reduce into 8-slot pad + store. + +Canonical IR: ``current_slow_vmi.pto``. +""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def padded_reduce_store(): + tile_ub = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + dst = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off = pto.const(0, dtype=pto.index) + mask = pto.vmi.create_mask(pto.const(64, dtype=pto.index), size=64) + tile = pto.vmi.vload(tile_ub.as_ptr(), off, size=64) + sum8 = pto.vmi.vcadd(tile, mask, group=8, reassoc=True) + pto.vmi.vstore( + sum8, + dst.as_ptr(), + off, + group=8, + stride=pto.const(1, dtype=pto.index), + ) + + +if __name__ == "__main__": + print(padded_reduce_store.compile().mlir_text()) diff --git a/docs/repro/store_pad8/fixtures/desired_vmi.pto b/docs/repro/store_pad8/fixtures/desired_vmi.pto new file mode 100644 index 0000000000..c341694d91 --- /dev/null +++ b/docs/repro/store_pad8/fixtures/desired_vmi.pto @@ -0,0 +1,28 @@ +// 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 VMI: reduce to one scalar and store with a 1-lane mask +// (mirrors test/lit/vmi_new/vmi_to_vpto_reduce_addf_store.pto). + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @direct_scalar_reduce_store( + %tile: !pto.vmi.vreg<64xf32>, + %dst: !pto.ptr, + %off: index) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %mask64 = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> + %mask1 = pto.vmi.create_mask %c1 : index -> !pto.vmi.mask<1xpred> + %sum = pto.vmi.vcadd %tile, %mask64 {reassoc} + : !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<1xf32> + pto.vmi.vstore %sum, %dst[%off], %mask1 + : !pto.vmi.vreg<1xf32>, !pto.ptr, !pto.vmi.mask<1xpred> + return + } +} diff --git a/docs/repro/store_pad8/fixtures/desired_vmi.py b/docs/repro/store_pad8/fixtures/desired_vmi.py new file mode 100644 index 0000000000..df48713a05 --- /dev/null +++ b/docs/repro/store_pad8/fixtures/desired_vmi.py @@ -0,0 +1,31 @@ +# 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`` — scalar reduce + 1-lane store. + +Canonical IR: ``desired_vmi.pto``. +""" + +from ptodsl import pto + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def direct_scalar_reduce_store(): + tile_ub = pto.alloc_tile(shape=[1, 64], dtype=pto.f32) + # 32B-aligned tile; store uses a 1-lane mask (ONEPT / 1PT shape). + dst = pto.alloc_tile(shape=[1, 8], dtype=pto.f32) + off = pto.const(0, dtype=pto.index) + mask64 = pto.vmi.create_mask(pto.const(64, dtype=pto.index), size=64) + mask1 = pto.vmi.create_mask(pto.const(1, dtype=pto.index), size=1) + tile = pto.vmi.vload(tile_ub.as_ptr(), off, size=64) + total = pto.vmi.vcadd(tile, mask64, reassoc=True) + pto.vmi.vstore(total, dst.as_ptr(), off, mask1) + + +if __name__ == "__main__": + print(direct_scalar_reduce_store.compile().mlir_text()) diff --git a/docs/repro/store_pad8/fixtures/lowered_vpto.pto b/docs/repro/store_pad8/fixtures/lowered_vpto.pto new file mode 100644 index 0000000000..8cf259476b --- /dev/null +++ b/docs/repro/store_pad8/fixtures/lowered_vpto.pto @@ -0,0 +1,10 @@ +module attributes {pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5"} { + func.func @padded_reduce_store(%arg0: !pto.vreg<64xf32>, %arg1: !pto.mask, %arg2: !pto.ptr, %arg3: index) { + %c1 = arith.constant 1 : index + %0 = pto.vcgadd %arg0, %arg1 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %1 = pto.pge_b32 "PAT_VL8" : !pto.mask + pto.vsts %0, %arg2[%arg3], %1 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + return + } +} + diff --git a/docs/repro/store_pad8/fixtures/target_mi.pto b/docs/repro/store_pad8/fixtures/target_mi.pto new file mode 100644 index 0000000000..6281dc02e1 --- /dev/null +++ b/docs/repro/store_pad8/fixtures/target_mi.pto @@ -0,0 +1,26 @@ +// 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: vcadd + 1PT_B32 store (ASC ONEPT-style), no pad-to-8. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @direct_scalar_reduce_store_mi( + %tile: !pto.vreg<64xf32>, + %mask: !pto.mask, + %dst: !pto.ptr, + %off: index) { + pto.vecscope { + %sum = pto.vcadd %tile, %mask {reassoc} + : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %store_m = pto.pge_b32 "PAT_VL1" : !pto.mask + pto.vsts %sum, %dst[%off], %store_m {dist = "1PT_B32"} + : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + } + return + } +} diff --git a/docs/repro/store_pad8/scripts/cannsim_metrics.py b/docs/repro/store_pad8/scripts/cannsim_metrics.py new file mode 100644 index 0000000000..d87b61b593 --- /dev/null +++ b/docs/repro/store_pad8/scripts/cannsim_metrics.py @@ -0,0 +1,457 @@ +""" +Parse cannsim performance artifacts for per_block_cast VF sim. + +Primary VF metrics (per ops-simulator guidance): +- RVEC span from trace_core*.json (vector-pipe VF body cycles) +- PUSHQ VF dispatch duration (fallback) +- MaxDur / span from log_ca/*.instr_log.dump (AI-core kernel window) + +SoC cycles from cannsim.log are kept for reference but are too coarse (~420) +to compare VF variants or tile sizes. +""" + +from __future__ import annotations + +import glob +import json +import os +import re +import sys +from collections import Counter +from dataclasses import dataclass, field +from typing import Iterable + +INSTR_LOG_RE = re.compile( + r"start:\s*(\d+),\s*tick:\s*(\d+).*?blkDim:\s*(\d+)", + re.DOTALL, +) +SOC_CYCLE_RE = re.compile( + r"\[Hardware\]\s+parallel simulation finish\.\s+sim time:\s*" + r"(?:SoC sub \d+ )?([\d.]+)s,\s*cycle:\s*(\d+)", +) + + +@dataclass(frozen=True) +class LaunchRecord: + start: int + tick: int + blk_dim: int + + @property + def duration(self) -> int: + return self.tick - self.start + + +@dataclass(frozen=True) +class InstrLogMetrics: + launches: tuple[LaunchRecord, ...] + span: int + max_dur: int + max_blk_dim: int + + @classmethod + def empty(cls) -> InstrLogMetrics: + return cls(launches=(), span=0, max_dur=0, max_blk_dim=0) + + +@dataclass(frozen=True) +class TraceMetrics: + rvec_span: int | None = None + pushq_vf_dur: int | None = None + mte2_span: int | None = None + mte3_span: int | None = None + vector_span: int | None = None + rvec_op_counts: dict[str, int] = field(default_factory=dict) + rvec_event_count: int = 0 + trace_path: str | None = None + + @property + def vf_cycles(self) -> int | None: + """Best available VF latency estimate.""" + if self.rvec_span is not None and self.rvec_span > 0: + return self.rvec_span + if self.pushq_vf_dur is not None and self.pushq_vf_dur > 0: + return self.pushq_vf_dur + return None + + +@dataclass(frozen=True) +class SocCycleRecord: + sim_wall_s: float + soc_cycles: int + + +@dataclass(frozen=True) +class RunMetrics: + out_dir: str + cannsim_run_dir: str + instr: InstrLogMetrics + trace: TraceMetrics + soc_cycles: tuple[SocCycleRecord, ...] + steady_soc_cycles: int | None + measured_kernel_cycles: int | None + + @property + def primary_vf_cycles(self) -> int | None: + if self.trace.vf_cycles is not None: + return self.trace.vf_cycles + if self.instr.max_dur > 0: + return self.instr.max_dur + return self.measured_kernel_cycles + + +def find_cannsim_run_dir(out_dir: str) -> str: + pattern = os.path.join(out_dir, "cannsim_*") + candidates = [p for p in glob.glob(pattern) if os.path.isdir(p)] + if not candidates: + raise FileNotFoundError(f"No cannsim_* directory under {out_dir}") + return max(candidates, key=os.path.getmtime) + + +def _parse_instr_log_file(path: str) -> list[LaunchRecord]: + with open(path, encoding="utf-8", errors="replace") as fh: + text = fh.read() + records: list[LaunchRecord] = [] + for match in INSTR_LOG_RE.finditer(text): + start, tick, blk_dim = (int(match.group(i)) for i in range(1, 4)) + records.append(LaunchRecord(start=start, tick=tick, blk_dim=blk_dim)) + return records + + +def _group_launches( + launches: list[LaunchRecord], gap_threshold: int = 3000 +) -> list[list[LaunchRecord]]: + if not launches: + return [] + sorted_launches = sorted(launches, key=lambda r: r.start) + groups: list[list[LaunchRecord]] = [[sorted_launches[0]]] + for record in sorted_launches[1:]: + prev_end = max(r.tick for r in groups[-1]) + if record.start - prev_end > gap_threshold: + groups.append([record]) + else: + groups[-1].append(record) + return groups + + +def metrics_from_launches(launches: Iterable[LaunchRecord]) -> InstrLogMetrics: + items = list(launches) + if not items: + return InstrLogMetrics.empty() + starts = [r.start for r in items] + ticks = [r.tick for r in items] + durations = [r.duration for r in items] + blk_dims = [r.blk_dim for r in items] + return InstrLogMetrics( + launches=tuple(items), + span=max(ticks) - min(starts), + max_dur=max(durations), + max_blk_dim=max(blk_dims), + ) + + +def parse_instr_log_dir(log_ca_dir: str, measured_only: bool = True) -> InstrLogMetrics: + pattern = os.path.join(log_ca_dir, "core*.veccore*.instr_log.dump") + paths = sorted(glob.glob(pattern)) + if not paths: + return InstrLogMetrics.empty() + + all_launches: list[LaunchRecord] = [] + for path in paths: + all_launches.extend(_parse_instr_log_file(path)) + + if not all_launches: + return InstrLogMetrics.empty() + + if measured_only: + max_blk = max(r.blk_dim for r in all_launches) + kernel_launches = [r for r in all_launches if r.blk_dim == max_blk] + if kernel_launches: + groups = _group_launches(kernel_launches, gap_threshold=2000) + all_launches = groups[-1] if groups else kernel_launches + + return metrics_from_launches(all_launches) + + +def parse_soc_cycles(cannsim_log: str) -> list[SocCycleRecord]: + with open(cannsim_log, encoding="utf-8", errors="replace") as fh: + text = fh.read() + + records: list[SocCycleRecord] = [] + for match in SOC_CYCLE_RE.finditer(text): + sim_wall_s = float(match.group(1)) + soc_cycles = int(match.group(2)) + records.append(SocCycleRecord(sim_wall_s=sim_wall_s, soc_cycles=soc_cycles)) + return records + + +def steady_state_soc_cycles(records: Iterable[SocCycleRecord]) -> int | None: + items = list(records) + if not items: + return None + if len(items) == 1: + return items[0].soc_cycles + return items[-1].soc_cycles + + +def find_trace_json(run_dir: str) -> str | None: + patterns = [ + os.path.join(run_dir, "report", "trace_core0.json"), + os.path.join(run_dir, "**", "trace_core0.json"), + ] + for pattern in patterns: + matches = glob.glob(pattern, recursive=True) + if matches: + return max(matches, key=os.path.getmtime) + return None + + +def _load_trace_events(trace_path: str) -> list[dict]: + with open(trace_path, encoding="utf-8") as fh: + payload = json.load(fh) + if isinstance(payload, list): + return payload + return payload.get("traceEvents", []) + + +def _pipe_span(events: list[dict], processes: dict[int, str], *needles: str) -> int | None: + pids = {pid for pid, name in processes.items() if any(n in name.upper() for n in needles)} + exec_ev = [e for e in events if e.get("ph") == "X" and "dur" in e and e.get("pid") in pids] + if not exec_ev: + return None + start = min(e["ts"] for e in exec_ev) + end = max(e["ts"] + e["dur"] for e in exec_ev) + span = int(end - start) + return span if span > 0 else None + + +def parse_trace_metrics(trace_path: str) -> TraceMetrics: + if not os.path.isfile(trace_path): + return TraceMetrics() + + events = _load_trace_events(trace_path) + processes = { + e["pid"]: e["args"]["name"] + for e in events + if e.get("ph") == "M" and e.get("name") == "process_name" + } + exec_ev = [e for e in events if e.get("ph") == "X" and "dur" in e] + + pushq_pids = {pid for pid, name in processes.items() if "PUSHQ" in name.upper()} + vf_dispatch = [ + e for e in exec_ev if e.get("pid") in pushq_pids and "VF" in e.get("name", "") + ] + pushq_vf_dur = max((e.get("dur", 0) for e in vf_dispatch), default=0) or None + + rvec_pids = {pid for pid, name in processes.items() if "RVEC" in name.upper()} + rvec_events = [e for e in exec_ev if e.get("pid") in rvec_pids] + rvec_span = None + rvec_op_counts: dict[str, int] = {} + rvec_event_count = 0 + if rvec_events: + start = min(e["ts"] for e in rvec_events) + end = max(e["ts"] + e["dur"] for e in rvec_events) + rvec_span = int(end - start) + rvec_op_counts = dict(Counter(e.get("name", "?") for e in rvec_events)) + rvec_event_count = len(rvec_events) + + return TraceMetrics( + rvec_span=rvec_span, + pushq_vf_dur=pushq_vf_dur, + mte2_span=_pipe_span(events, processes, "MTE2"), + mte3_span=_pipe_span(events, processes, "MTE3"), + vector_span=_pipe_span(events, processes, "RVEC", "VECTOR", "VEC"), + rvec_op_counts=rvec_op_counts, + rvec_event_count=rvec_event_count, + trace_path=trace_path, + ) + + +def parse_marker_soc_cycles(cannsim_log: str) -> tuple[list[SocCycleRecord], list[SocCycleRecord]]: + with open(cannsim_log, encoding="utf-8", errors="replace") as fh: + lines = fh.readlines() + + in_window = False + all_records: list[SocCycleRecord] = [] + window_records: list[SocCycleRecord] = [] + + for line in lines: + if "CYCLE_MARKER" in line: + in_window = True + window_records.clear() + continue + if "CYCLE_DONE" in line: + in_window = False + continue + match = SOC_CYCLE_RE.search(line) + if not match: + continue + rec = SocCycleRecord(sim_wall_s=float(match.group(1)), soc_cycles=int(match.group(2))) + all_records.append(rec) + if in_window: + window_records.append(rec) + + return all_records, window_records + + +def measured_kernel_soc_cycles(cannsim_log: str, warmup_cycle: int = 248) -> int | None: + _, window = parse_marker_soc_cycles(cannsim_log) + measured = [r.soc_cycles for r in window if r.soc_cycles != warmup_cycle] + if not measured: + return None + return measured[-1] + + +def parse_run_metrics(out_dir: str) -> RunMetrics: + run_dir = find_cannsim_run_dir(out_dir) + log_ca = os.path.join(run_dir, "log_ca") + cannsim_log = os.path.join(run_dir, "cannsim.log") + + instr = parse_instr_log_dir(log_ca) if os.path.isdir(log_ca) else InstrLogMetrics.empty() + soc = parse_soc_cycles(cannsim_log) if os.path.isfile(cannsim_log) else [] + measured = measured_kernel_soc_cycles(cannsim_log) if os.path.isfile(cannsim_log) else None + + trace_path = find_trace_json(run_dir) + trace = parse_trace_metrics(trace_path) if trace_path else TraceMetrics() + + return RunMetrics( + out_dir=out_dir, + cannsim_run_dir=run_dir, + instr=instr, + trace=trace, + soc_cycles=tuple(soc), + steady_soc_cycles=steady_state_soc_cycles(soc), + measured_kernel_cycles=measured, + ) + + +def format_run_summary(m: RunMetrics, label: str | None = None) -> str: + """Human-readable one-run summary.""" + title = label or os.path.basename(m.out_dir.rstrip("/")) + lines = [f"=== {title} ({m.out_dir}) ===", f"cannsim run dir: {m.cannsim_run_dir}"] + + primary = m.primary_vf_cycles + if primary is not None: + if m.trace.rvec_span and m.trace.rvec_span > 0: + source = "RVEC span" + elif m.trace.pushq_vf_dur and m.trace.pushq_vf_dur > 0: + source = "PUSHQ VF dur" + elif m.instr.max_dur > 0: + source = "instr MaxDur" + else: + source = "SoC (fallback)" + lines.append(f"primary VF cycles: {primary} ({source})") + + if m.trace.rvec_span: + lines.append(f"RVEC span: {m.trace.rvec_span}") + if m.trace.pushq_vf_dur: + lines.append(f"PUSHQ VF dur: {m.trace.pushq_vf_dur}") + if m.trace.mte2_span: + lines.append(f"MTE2 span: {m.trace.mte2_span}") + if m.trace.mte3_span: + lines.append(f"MTE3 span: {m.trace.mte3_span}") + if m.trace.rvec_op_counts: + top_ops = sorted(m.trace.rvec_op_counts.items(), key=lambda kv: -kv[1])[:6] + lines.append("top RVEC ops: " + ", ".join(f"{k}={v}" for k, v in top_ops)) + + if m.instr.max_dur: + lines.append( + f"instr MaxDur / span: {m.instr.max_dur} / {m.instr.span} " + f"(blkDim={m.instr.max_blk_dim})" + ) + elif not os.path.isdir(os.path.join(m.cannsim_run_dir, "log_ca")): + lines.append("instr MaxDur / span: (log_ca unavailable)") + + if m.measured_kernel_cycles is not None: + lines.append( + f"SoC cycles (measured): {m.measured_kernel_cycles} " + "(coarse; ~420 is normal)" + ) + elif m.steady_soc_cycles is not None: + lines.append(f"SoC cycles (steady): {m.steady_soc_cycles} (coarse)") + + if m.trace.trace_path: + lines.append(f"trace: {m.trace.trace_path}") + elif not find_trace_json(m.cannsim_run_dir): + lines.append("trace: (none — need instr.bin + cannsim report)") + + return "\n".join(lines) + + +def format_table(rows: list[tuple[str, RunMetrics]]) -> str: + header = ( + f"{'case':<22} {'primary':>8} {'RVEC':>8} {'MaxDur':>8} " + f"{'Span':>8} {'SoC':>6}" + ) + lines = [header, "-" * len(header)] + for label, m in rows: + primary = m.primary_vf_cycles + lines.append( + f"{label:<22} " + f"{primary if primary is not None else '-':>8} " + f"{m.trace.rvec_span if m.trace.rvec_span else '-':>8} " + f"{m.instr.max_dur if m.instr.max_dur else '-':>8} " + f"{m.instr.span if m.instr.span else '-':>8} " + f"{m.measured_kernel_cycles if m.measured_kernel_cycles else '-':>6}" + ) + return "\n".join(lines) + + +def default_cycle_out_dirs(sim_root: str | None = None) -> list[str]: + """Return sim_outputs/cycle_{dtype}_{mode} dirs in canonical variant order.""" + root = sim_root or os.path.join(os.path.dirname(__file__), "..", "sim_outputs") + try: + from ref.tile_config import DTYPES, MODES + + ordered = [ + os.path.join(root, f"cycle_{dtype}_{mode}") + for dtype in DTYPES + for mode in MODES + if os.path.isdir(os.path.join(root, f"cycle_{dtype}_{mode}")) + ] + if ordered: + return ordered + except ImportError: + pass + return sorted( + p for p in glob.glob(os.path.join(root, "cycle_*")) if os.path.isdir(p) + ) + + +def main(argv: list[str] | None = None) -> int: + import argparse + + parser = argparse.ArgumentParser(description="Parse rope VF cannsim cycle metrics") + parser.add_argument( + "out_dirs", + nargs="*", + help="sim_outputs subdirs (default: all cycle_{dtype}_{mode})", + ) + parser.add_argument("--table", action="store_true", help="Print compact table") + args = parser.parse_args(argv) + + dirs = args.out_dirs if args.out_dirs else default_cycle_out_dirs() + if not dirs: + print("No sim_outputs/cycle_* dirs found. Run scripts/run_cycle_bench.sh first.", file=sys.stderr) + return 1 + + rows: list[tuple[str, RunMetrics]] = [] + for path in dirs: + try: + m = parse_run_metrics(path) + except FileNotFoundError as exc: + print(str(exc), file=sys.stderr) + continue + label = os.path.basename(path.rstrip("/")) + rows.append((label, m)) + if not args.table: + print(format_run_summary(m, label)) + print() + + if args.table and rows: + print(format_table(rows)) + return 0 if rows else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/repro/store_pad8/scripts/check_desired.sh b/docs/repro/store_pad8/scripts/check_desired.sh new file mode 100755 index 0000000000..85f8ab2dfc --- /dev/null +++ b/docs/repro/store_pad8/scripts/check_desired.sh @@ -0,0 +1,95 @@ +#!/usr/bin/env bash +# Try compiling the desired (1-lane) and target (1PT_B32) VMI forms; record results. +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_desired" +mkdir -p "${OUT}" +LOG="${OUT}/compile_results.txt" + +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/env.sh" + +BISHENG="${BISHENG:-${ASCEND_HOME_PATH}/bin/bisheng}" +ASCEND="${ASCEND_HOME_PATH}" +NPU_ARCH="${ASCEND_NPU_ARCH:-dav-3510}" + +: > "${LOG}" +echo "store_pad8 check_desired — $(date -u +%Y-%m-%dT%H:%M:%SZ)" | tee -a "${LOG}" +echo "PTOAS_ROOT=${PTOAS_ROOT}" | tee -a "${LOG}" +echo | tee -a "${LOG}" + +# --- desired_vmi.py: mask-1 store after vcadd (should compile via ptodsl) --- +echo "=== desired_vmi.py (ptodsl compile) ===" | tee -a "${LOG}" +set +e +python3 "${FIXTURES}/desired_vmi.py" > "${OUT}/desired_vmi.mlir" 2> "${OUT}/desired_vmi.err" +rc=$? +set -e +if [ "${rc}" -eq 0 ]; then + echo "PASS: desired_vmi.py compiled (MLIR written to sim_outputs/check_desired/desired_vmi.mlir)" | tee -a "${LOG}" +else + echo "FAIL: desired_vmi.py compile exit ${rc}" | tee -a "${LOG}" + cat "${OUT}/desired_vmi.err" | tee -a "${LOG}" +fi +echo | tee -a "${LOG}" + +# --- target_mi.pto: ptoas emit-vpto-llvm-ir + bisheng -c -x ir (expected to crash) --- +echo "=== target_mi.pto (ptoas -> LLVM IR -> bisheng object) ===" | tee -a "${LOG}" +TARGET_MI="${FIXTURES}/target_mi.pto" +LLVM_IR="${OUT}/target_mi.ll" +OBJ="${OUT}/target_mi.o" + +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-llvm-ir "${TARGET_MI}" -o "${LLVM_IR}" > "${OUT}/target_mi_ptoas.log" 2>&1 + pto_rc=$? +else + echo "SKIP: ptoas not on PATH" | tee -a "${LOG}" + pto_rc=127 +fi +set -e + +if [ "${pto_rc}" -eq 0 ]; then + echo "PASS: ptoas --emit-vpto-llvm-ir target_mi.pto" | tee -a "${LOG}" + set +e + "${BISHENG}" -O2 -fPIC -std=c++17 --npu-arch="${NPU_ARCH}" -c -x ir \ + "${LLVM_IR}" -o "${OBJ}" > "${OUT}/target_mi_bisheng.log" 2>&1 + b_rc=$? + set -e + if [ "${b_rc}" -eq 0 ] && [ -s "${OBJ}" ]; then + echo "PASS: bisheng object compile (unexpected — gap may be closed)" | tee -a "${LOG}" + else + echo "FAIL: bisheng object compile exit ${b_rc} (expected on vmi-v0.1.3)" | tee -a "${LOG}" + tail -20 "${OUT}/target_mi_bisheng.log" | tee -a "${LOG}" + fi +else + echo "FAIL: ptoas emit-vpto-llvm-ir exit ${pto_rc}" | tee -a "${LOG}" + tail -20 "${OUT}/target_mi_ptoas.log" | tee -a "${LOG}" +fi +echo | tee -a "${LOG}" + +# --- reference_asc_cce.asc: hand-written CCE ONEPT baseline (should compile) --- +echo "=== reference_asc_cce.asc (bisheng --cce-aicore-only) ===" | tee -a "${LOG}" +REF="${FIXTURES}/reference_asc_cce.asc" +REF_OBJ="${OUT}/reference_asc_cce.o" +set +e +"${BISHENG}" -O2 -fPIC -std=c++17 --npu-arch="${NPU_ARCH}" --cce-aicore-only -c \ + "${REF}" -o "${REF_OBJ}" \ + -I"${ASCEND}/include" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/impl" \ + -I"${ASCEND}/compiler/tikcpp/tikcfw/interface" \ + > "${OUT}/reference_asc_cce.log" 2>&1 +ref_rc=$? +set -e +if [ "${ref_rc}" -eq 0 ] && [ -s "${REF_OBJ}" ]; then + echo "PASS: reference_asc_cce.asc -> non-empty .o (CCE ONEPT path is legal)" | tee -a "${LOG}" +else + echo "FAIL: reference_asc_cce.asc compile exit ${ref_rc}" | tee -a "${LOG}" + tail -20 "${OUT}/reference_asc_cce.log" | tee -a "${LOG}" +fi + +echo | tee -a "${LOG}" +echo "Full log: ${LOG}" | tee -a "${LOG}" diff --git a/docs/repro/store_pad8/scripts/env.sh b/docs/repro/store_pad8/scripts/env.sh new file mode 100755 index 0000000000..0f9af3a407 --- /dev/null +++ b/docs/repro/store_pad8/scripts/env.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +# CANN + PTOAS env for store_pad8 cannsim runs. +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}}" + +# PTOAS repo root: store_pad8 -> 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" + # ptoas_env prepends build/tools/ptoas; that stub can clash with the pip/wheel + # launcher (LLVM CL option double-register). Prefer the wheel on PATH. + PATH="$(echo "${PATH}" | tr ':' '\n' | grep -v "/build/tools/ptoas$" | paste -sd:)" + export PATH +fi + +export ASCEND_NPU_ARCH="${ASCEND_NPU_ARCH:-dav-3510}" +ARCH=$(uname -m)-linux +CANN="${ASCEND_HOME_PATH}" +export LD_LIBRARY_PATH="${CANN}/${ARCH}/simulator/dav_3510/camodel:${CANN}/${ARCH}/simulator/dav_3510/lib:${CANN}/${ARCH}/lib64:${LD_LIBRARY_PATH:-}" +export TORCH_DEVICE_BACKEND_AUTOLOAD=0 +export PTOAS_HOST_TARGET_CPU="${PTOAS_HOST_TARGET_CPU:-znver3}" +unset CCC_OVERRIDE_OPTIONS || true diff --git a/docs/repro/store_pad8/scripts/run_cannsim.sh b/docs/repro/store_pad8/scripts/run_cannsim.sh new file mode 100755 index 0000000000..22be494144 --- /dev/null +++ b/docs/repro/store_pad8/scripts/run_cannsim.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# Run store_pad8 correctness under cannsim. +# Usage: STORE_PAD8_CASE=large|small ./scripts/run_cannsim.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +KERNEL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/env.sh" + +BACKEND="${1:-vmi}" +export TLVF_VMI_BACKEND="${BACKEND}" +STORE_PAD8_CASE="${STORE_PAD8_CASE:-large}" +export STORE_PAD8_CASE + +force_build() { + python3 -c " +import sys +sys.path.insert(0, '${KERNEL_ROOT}') +from common.cce_vf_build import build_cce_root +from pathlib import Path +build_cce_root(Path('${KERNEL_ROOT}/cce'), force=True) +" +} + +TEST="${KERNEL_ROOT}/test/test_store_pad8.py" +if [ "${BACKEND}" = "cce" ]; then + force_build +fi +OUT="${KERNEL_ROOT}/sim_outputs/store_pad8_${STORE_PAD8_CASE}_${BACKEND}" +mkdir -p "${OUT}" +exec "${SCRIPT_DIR}/run_sim.sh" "${TEST}" "${OUT}" diff --git a/docs/repro/store_pad8/scripts/run_sim.sh b/docs/repro/store_pad8/scripts/run_sim.sh new file mode 100755 index 0000000000..9755a264e5 --- /dev/null +++ b/docs/repro/store_pad8/scripts/run_sim.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# cannsim record + report. Usage: run_sim.sh [out_dir] +set -euo pipefail +SHARED="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SCRIPT="$(realpath "${1:?Usage: run_sim.sh [out_dir]}")" +KERNEL_ROOT="$(cd "$(dirname "${SCRIPT}")/.." && pwd)" +OUT_DIR="${2:-${KERNEL_ROOT}/sim_outputs/$(basename "$SCRIPT" .py)}" +ENTRY="${SHARED}/run_sim_entry.sh" +mkdir -p "${OUT_DIR}" + +# cannsim mkdirs log_ca next to the entry script before the entry runs. +for asset in log_ca instr.bin; do + link="${SHARED}/${asset}" + if [ -L "${link}" ] || [ -e "${link}" ]; then + rm -rf "${link}" + fi +done +rm -rf "${KERNEL_ROOT}/log_ca" "${KERNEL_ROOT}/instr.bin" +mkdir -p "${KERNEL_ROOT}/log_ca" + +echo "==> cannsim record ${SCRIPT} -> ${OUT_DIR}" +set +e +cannsim record "${ENTRY}" -s Ascend950 --gen-report -o "${OUT_DIR}" -u "${SCRIPT}" +rc=$? +set -e +echo "cannsim exit code: ${rc}" + +RUN="$(find "${OUT_DIR}" -maxdepth 1 -type d -name 'cannsim_*' 2>/dev/null | sort | tail -1 || true)" +if [ -n "${RUN}" ]; then + [ -d "${KERNEL_ROOT}/log_ca" ] && [ ! -d "${RUN}/log_ca" ] && cp -a "${KERNEL_ROOT}/log_ca" "${RUN}/log_ca" + if [ -f "${KERNEL_ROOT}/instr.bin" ]; then + dest="${RUN}/instr.bin" + if [ ! -f "${dest}" ] || [ -L "${dest}" ] || [ "${dest}" -ef "${KERNEL_ROOT}/instr.bin" ]; then + cp -f "${KERNEL_ROOT}/instr.bin" "${RUN}/.instr.bin.tmp" && mv -f "${RUN}/.instr.bin.tmp" "${dest}" + fi + fi + if [ -f "${RUN}/instr.bin" ]; then + mkdir -p "${RUN}/report" + cannsim report -e "${RUN}" -o "${RUN}/report" -n 0 || true + fi +fi + +python3 "${SHARED}/cannsim_metrics.py" "${OUT_DIR}" || true +exit 0 diff --git a/docs/repro/store_pad8/scripts/run_sim_entry.sh b/docs/repro/store_pad8/scripts/run_sim_entry.sh new file mode 100755 index 0000000000..ca8c92179e --- /dev/null +++ b/docs/repro/store_pad8/scripts/run_sim_entry.sh @@ -0,0 +1,16 @@ +#!/usr/bin/env bash +# Real file (not a symlink): cannsim collects log_ca/instr.bin from this directory. +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +KERNEL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" + +mkdir -p "${KERNEL_ROOT}/log_ca" +rm -f "${KERNEL_ROOT}/instr.bin" +find "${KERNEL_ROOT}/log_ca" -mindepth 1 -delete 2>/dev/null || true + +for asset in log_ca instr.bin; do + ln -sfn "${KERNEL_ROOT}/${asset}" "${SCRIPT_DIR}/${asset}" +done + +cd "${KERNEL_ROOT}" +exec python3 "$@" diff --git a/docs/repro/store_pad8/scripts/run_three_way.sh b/docs/repro/store_pad8/scripts/run_three_way.sh new file mode 100755 index 0000000000..7dafbc5eda --- /dev/null +++ b/docs/repro/store_pad8/scripts/run_three_way.sh @@ -0,0 +1,36 @@ +#!/usr/bin/env bash +# Four-way reduce store: CCE ONEPT vs VMI pad8 / mask1 / group1. +# Usage: STORE_PAD8_CASE=large|small ./scripts/run_three_way.sh +set -euo pipefail +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +KERNEL_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +# shellcheck disable=SC1091 +source "${SCRIPT_DIR}/env.sh" + +CASE="${STORE_PAD8_CASE:-large}" +export STORE_PAD8_CASE="${CASE}" +TEST="${KERNEL_ROOT}/test/test_store_pad8.py" + +echo "=== CCE ONEPT STORE_PAD8_CASE=${CASE} ===" +STORE_PAD8_CASE="${CASE}" "${SCRIPT_DIR}/run_cannsim.sh" cce || true + +echo "=== VMI pad8 STORE_PAD8_CASE=${CASE} ===" +export STORE_PAD8_VARIANT=pad8 +STORE_PAD8_CASE="${CASE}" "${SCRIPT_DIR}/run_cannsim.sh" vmi || true + +for variant in mask1 group1; do + echo "=== VMI ${variant} STORE_PAD8_CASE=${CASE} ===" + export STORE_PAD8_VARIANT="${variant}" + OUT="${KERNEL_ROOT}/sim_outputs/store_pad8_${CASE}_vmi_${variant}" + mkdir -p "${OUT}" + TLVF_VMI_BACKEND=vmi STORE_PAD8_VARIANT="${variant}" STORE_PAD8_CASE="${CASE}" \ + "${SCRIPT_DIR}/run_sim.sh" "${TEST}" "${OUT}" || true +done + +echo "=== metrics ===" +python3 "${SCRIPT_DIR}/cannsim_metrics.py" \ + "${KERNEL_ROOT}/sim_outputs/store_pad8_${CASE}_cce" \ + "${KERNEL_ROOT}/sim_outputs/store_pad8_${CASE}_vmi" \ + "${KERNEL_ROOT}/sim_outputs/store_pad8_${CASE}_vmi_mask1" \ + "${KERNEL_ROOT}/sim_outputs/store_pad8_${CASE}_vmi_group1" \ + --table || true diff --git a/docs/repro/store_pad8/test/test_store_pad8.py b/docs/repro/store_pad8/test/test_store_pad8.py new file mode 100644 index 0000000000..3448bfb84f --- /dev/null +++ b/docs/repro/store_pad8/test/test_store_pad8.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Correctness test for the store_pad8 microbench (CCE / VMI).""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + +from common.golden import LARGE, SMALL, extract_padded, generate_case # noqa: E402 +from common.launcher import launch # noqa: E402 +from common.torch_runtime import device_str, init_torch_npu # noqa: E402 + + +def _check(backend: str, case, variant: str) -> tuple[bool, float]: + ref = generate_case(case) + acc = __import__("torch").from_numpy(ref["acc"]).to(device_str()) + got = launch(acc, case.n_acc, backend=backend, variant=variant).cpu().numpy() + if backend == "vmi" and variant == "pad8": + got = extract_padded(got, case.n_acc) + diff = float(np.max(np.abs(got.astype(np.float32) - ref["reduced"]))) + atol = float(os.environ.get("STORE_PAD8_ATOL", 2e-3)) + ok = diff <= atol + tag = f"{backend}/{variant}" if backend == "vmi" else backend + print( + f"[{tag}] store_pad8 {case.name} (n_acc={case.n_acc}): maxDiff={diff:.6g} ok={ok}", + flush=True, + ) + return ok, diff + + +def main() -> int: + init_torch_npu() + backend = os.environ.get("TLVF_VMI_BACKEND", "vmi").lower() + variant = os.environ.get("STORE_PAD8_VARIANT", "pad8").lower() + case_name = os.environ.get("STORE_PAD8_CASE", "large").lower() + case = LARGE if case_name == "large" else SMALL + + ok, _ = _check(backend, case, variant) + tag = f"{backend}/{variant}" if backend == "vmi" else backend + if ok: + print(f"All store_pad8 [{tag}/{case.name}] tests PASSED", flush=True) + os._exit(0) + print(f"store_pad8 [{tag}/{case.name}] tests FAILED", flush=True) + os._exit(1) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/repro/store_pad8/vmi/store_group1_vmi.py b/docs/repro/store_pad8/vmi/store_group1_vmi.py new file mode 100644 index 0000000000..833c29481b --- /dev/null +++ b/docs/repro/store_pad8/vmi/store_group1_vmi.py @@ -0,0 +1,56 @@ +"""PTO VMI reduce epilogue: vcadd(group=1) + vstore(group=1) compact [N_ACC]. + +Hinted ONEPT / 1PT path: reduce to a 1-lane score, then unit-stride group +store (no mask). On a PTOAS build that lowers compact reduction stores to +point stores, this should become vsts ... 1PT_B32 — unlike mask1 (masked +1-lane store) or pad-8 (vbrc×8). +""" + +from ptodsl import pto + +_VL = 64 + +_UB = { + "acc": 0x0000, + "out": 0x2000, +} + + +def _build(n_acc: int, name: str): + @pto.jit( + name=name, + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, + ) + def kernel( + acc_gm: pto.ptr(pto.f32, "gm"), + out_gm: pto.ptr(pto.f32, "gm"), + ): + acc = pto.castptr(pto.const(_UB["acc"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + out = pto.castptr(pto.const(_UB["out"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + + nbytes_acc = n_acc * _VL * 4 + nbytes_out = n_acc * 4 + pto.mte_gm_ub(acc_gm, acc, 0, nbytes_acc, nburst=(1, nbytes_acc, nbytes_acc)) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + mask64 = pto.vmi.create_mask(_VL, size=_VL) + + for i in pto.static_range(n_acc): + acc_v = pto.vmi.vload(acc, i * _VL, size=_VL) + red = pto.vmi.vcadd(acc_v, mask64, group=1, reassoc=True) + pto.vmi.vstore(red, out, i, group=1, stride=1) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.mte_ub_gm(out, out_gm, nbytes_out, nburst=(1, nbytes_out, nbytes_out)) + + return kernel + + +store_group1_vmi_large = _build(20, "store_group1_vmi_large") +store_group1_vmi_small = _build(4, "store_group1_vmi_small") diff --git a/docs/repro/store_pad8/vmi/store_mask1_vmi.py b/docs/repro/store_pad8/vmi/store_mask1_vmi.py new file mode 100644 index 0000000000..b45f49e8b4 --- /dev/null +++ b/docs/repro/store_pad8/vmi/store_mask1_vmi.py @@ -0,0 +1,59 @@ +"""PTO VMI reduce epilogue: vcadd + 1-lane masked store (compact [N_ACC]). + +Same math as the pad-8 workaround and CCE ONEPT: sum each of N_ACC length-64 +f32 rows to one scalar. Unlike pad-8, this does not vbrc to 8 lanes — it stores +under a 1-lane mask into a compact [N_ACC] buffer (CCE GM layout). + +On vmi-v0.1.3 this lowers (vdup/vadd + masked vsts) but is not CCE ONEPT / +1PT_B32. Use it to separate kernel formulation from compiler lowering cost. +""" + +from ptodsl import pto + +_VL = 64 + +_UB = { + "acc": 0x0000, + "out": 0x2000, +} + + +def _build(n_acc: int, name: str): + @pto.jit( + name=name, + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, + ) + def kernel( + acc_gm: pto.ptr(pto.f32, "gm"), + out_gm: pto.ptr(pto.f32, "gm"), + ): + acc = pto.castptr(pto.const(_UB["acc"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + out = pto.castptr(pto.const(_UB["out"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + + nbytes_acc = n_acc * _VL * 4 + nbytes_out = n_acc * 4 + pto.mte_gm_ub(acc_gm, acc, 0, nbytes_acc, nburst=(1, nbytes_acc, nbytes_acc)) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + mask64 = pto.vmi.create_mask(_VL, size=_VL) + mask1 = pto.vmi.create_mask(1, size=1) + + for i in pto.static_range(n_acc): + acc_v = pto.vmi.vload(acc, i * _VL, size=_VL) + red = pto.vmi.vcadd(acc_v, mask64, reassoc=True) + pto.vmi.vstore(red, out, i, mask1) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.mte_ub_gm(out, out_gm, nbytes_out, nburst=(1, nbytes_out, nbytes_out)) + + return kernel + + +store_mask1_vmi_large = _build(20, "store_mask1_vmi_large") +store_mask1_vmi_small = _build(4, "store_mask1_vmi_small") diff --git a/docs/repro/store_pad8/vmi/store_pad8_vmi.py b/docs/repro/store_pad8/vmi/store_pad8_vmi.py new file mode 100644 index 0000000000..84e2f07cff --- /dev/null +++ b/docs/repro/store_pad8/vmi/store_pad8_vmi.py @@ -0,0 +1,61 @@ +"""PTO VMI port of the store_pad8 reduce epilogue (pad-8 workaround). + +ASC/CCE: vcadd + scalar vsts ONEPT_B32 per accumulator (see +../cce/csrc/store_pad8_vf_sim_kernel.cpp). Product VMI workaround: broadcast +the reduced value to 8 lanes (vbrc size=8) and store under an 8-wide mask. +See store_mask1_vmi.py for the compact 1-lane store form. + +Output is the full padded [N_ACC, 8] buffer (lane 0 of each row holds the +reduced value). Host-side compaction is out of scope here. +""" + +from ptodsl import pto + +_VL = 64 +_REDUCE_PAD = 8 + +_UB = { + "acc": 0x0000, + "pad": 0x2000, +} + + +def _build(n_acc: int, name: str): + @pto.jit( + name=name, + target="a5", + backend="vpto", + mode="explicit", + kernel_kind="vector", + insert_sync=False, + ) + def kernel( + acc_gm: pto.ptr(pto.f32, "gm"), + pad_gm: pto.ptr(pto.f32, "gm"), + ): + acc = pto.castptr(pto.const(_UB["acc"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + pad = pto.castptr(pto.const(_UB["pad"], dtype=pto.ui64), pto.ptr(pto.f32, "ub")) + + nbytes_acc = n_acc * _VL * 4 + nbytes_pad = n_acc * _REDUCE_PAD * 4 + pto.mte_gm_ub(acc_gm, acc, 0, nbytes_acc, nburst=(1, nbytes_acc, nbytes_acc)) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.V, event_id=0) + + mask64 = pto.vmi.create_mask(_VL, size=_VL) + mask8 = pto.vmi.create_mask(_REDUCE_PAD, size=_REDUCE_PAD) + + for i in pto.static_range(n_acc): + acc_v = pto.vmi.vload(acc, i * _VL, size=_VL) + red = pto.vmi.vcadd(acc_v, mask64, reassoc=True) + pto.vmi.vstore(pto.vmi.vbrc(red, size=_REDUCE_PAD), pad, i * _REDUCE_PAD, mask8) + + pto.set_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=1) + pto.mte_ub_gm(pad, pad_gm, nbytes_pad, nburst=(1, nbytes_pad, nbytes_pad)) + + return kernel + + +store_pad8_vmi_large = _build(20, "store_pad8_vmi_large") +store_pad8_vmi_small = _build(4, "store_pad8_vmi_small")