|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# Copyright (c) PyPTO Contributors. |
| 3 | +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of |
| 4 | +# CANN Open Software License Agreement Version 2.0 (the "License"). |
| 5 | +# Please refer to the License for details. You may not use this file except in compliance with the License. |
| 6 | +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, |
| 7 | +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. |
| 8 | +# See LICENSE in the root of the software repository for the full text of the License. |
| 9 | +# ----------------------------------------------------------------------------------------------------------- |
| 10 | +"""args_dump profiling smoke — capture pipeline produces a usable |
| 11 | +``args_dump/`` directory. |
| 12 | +
|
| 13 | +Re-uses ``vector_example`` (5 submit_task calls). With ``--dump-args`` the |
| 14 | +AICPU writer captures task dump records into a unified manifest + raw-byte |
| 15 | +payload pair under ``<output_prefix>/args_dump/``. Smoke asserts: |
| 16 | +manifest exists + parses, the ``bin_file`` field it names exists, entries |
| 17 | +use the unified schema, and no legacy args-only manifest is emitted. |
| 18 | +""" |
| 19 | + |
| 20 | +import json |
| 21 | +import subprocess |
| 22 | +import sys |
| 23 | + |
| 24 | +import torch |
| 25 | +from simpler.task_interface import ArgDirection as D |
| 26 | + |
| 27 | +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test |
| 28 | +from simpler_setup.scene_test import _outputs_dir, _sanitize_for_filename |
| 29 | + |
| 30 | +KERNELS_BASE = "../../../../../../examples/a5/tensormap_and_ringbuffer/vector_example/kernels" |
| 31 | + |
| 32 | + |
| 33 | +@scene_test(level=2, runtime="tensormap_and_ringbuffer") |
| 34 | +class TestArgsDump(SceneTestCase): |
| 35 | + """args_dump capture smoke, level-aware on the ``--dump-args`` level. |
| 36 | +
|
| 37 | + Uses ``partial_dump_orch`` (5 tasks; four carry ``dump(...)`` markers) so a |
| 38 | + single orchestration exercises both modes: |
| 39 | +
|
| 40 | + - ``--dump-args 1`` (partial): only marked args are captured — task |
| 41 | + ``0x..00`` via no-arg ``dump()`` (all tensor + scalar args), task |
| 42 | + ``0x..01`` via ``dump(t2_addend)`` (scalar-only), task ``0x..02`` via |
| 43 | + ``dump(d, inter_ci, t3_count)`` (mixed tensor + scalar, input ``e`` |
| 44 | + excluded), and task ``0x..03`` via no-arg ``dump()`` (all tensor args). |
| 45 | + Mode is latched host-side before dispatch, so it is race-free regardless |
| 46 | + of submission order. |
| 47 | + - ``--dump-args 2`` (full): markers are ignored, every task is dumped. |
| 48 | +
|
| 49 | + The dump level comes straight from the CLI ``--dump-args`` value |
| 50 | + (no per-case override). |
| 51 | + """ |
| 52 | + |
| 53 | + CALLABLE = { |
| 54 | + "orchestration": { |
| 55 | + "source": "kernels/orchestration/partial_dump_orch.cpp", |
| 56 | + "function_name": "aicpu_orchestration_entry", |
| 57 | + "signature": [D.IN, D.IN, D.OUT], |
| 58 | + }, |
| 59 | + "incores": [ |
| 60 | + { |
| 61 | + "func_id": 0, |
| 62 | + "source": f"{KERNELS_BASE}/aiv/kernel_add.cpp", |
| 63 | + "core_type": "aiv", |
| 64 | + "signature": [D.IN, D.IN, D.OUT], |
| 65 | + }, |
| 66 | + { |
| 67 | + "func_id": 1, |
| 68 | + "source": f"{KERNELS_BASE}/aiv/kernel_add_scalar.cpp", |
| 69 | + "core_type": "aiv", |
| 70 | + "signature": [D.IN, D.OUT], |
| 71 | + }, |
| 72 | + { |
| 73 | + "func_id": 2, |
| 74 | + "source": f"{KERNELS_BASE}/aiv/kernel_mul.cpp", |
| 75 | + "core_type": "aiv", |
| 76 | + "signature": [D.IN, D.IN, D.OUT], |
| 77 | + }, |
| 78 | + ], |
| 79 | + } |
| 80 | + |
| 81 | + CASES = [ |
| 82 | + { |
| 83 | + "name": "default", |
| 84 | + "platforms": ["a5sim", "a5"], |
| 85 | + "config": {"aicpu_thread_num": 4, "block_dim": 3}, |
| 86 | + "params": {}, |
| 87 | + }, |
| 88 | + ] |
| 89 | + |
| 90 | + def generate_args(self, params): |
| 91 | + SIZE = 128 * 128 |
| 92 | + return TaskArgsBuilder( |
| 93 | + Tensor("a", torch.full((SIZE,), 2.0, dtype=torch.float32)), |
| 94 | + Tensor("b", torch.full((SIZE,), 3.0, dtype=torch.float32)), |
| 95 | + Tensor("f", torch.zeros(SIZE, dtype=torch.float32)), |
| 96 | + ) |
| 97 | + |
| 98 | + def compute_golden(self, args, params): |
| 99 | + args.f[:] = (args.a + args.b + 1) * (args.a + args.b + 2) + (args.a + args.b) |
| 100 | + |
| 101 | + def test_run(self, st_platform, st_worker, request): |
| 102 | + super().test_run(st_platform, st_worker, request) |
| 103 | + level = int(request.config.getoption("--dump-args", default=0)) |
| 104 | + if not level: |
| 105 | + return |
| 106 | + safe_label = _sanitize_for_filename("TestArgsDump_default") |
| 107 | + matches = sorted(_outputs_dir().glob(f"{safe_label}_*"), key=lambda p: p.stat().st_mtime) |
| 108 | + assert matches, "args dump output directory missing" |
| 109 | + dump_dir = matches[-1] / "args_dump" |
| 110 | + assert dump_dir.is_dir(), f"args_dump/ missing under {matches[-1]} — dump capture failed?" |
| 111 | + manifest = dump_dir / "args_dump.json" |
| 112 | + assert manifest.exists(), f"args_dump.json missing under {dump_dir} — collector finalize failed?" |
| 113 | + with manifest.open() as f: |
| 114 | + data = json.load(f) |
| 115 | + bin_name = data.get("bin_file") |
| 116 | + tensors = data.get("args", []) |
| 117 | + assert tensors, f"args_dump.json has no entries: {data}" |
| 118 | + if level == 3: |
| 119 | + # full_json_only: metadata only, no payload and no .bin file. |
| 120 | + assert bin_name is None, f"level 3 manifest should have bin_file=null: {data}" |
| 121 | + assert not (dump_dir / "args.bin").exists(), "level 3 must not write args.bin" |
| 122 | + assert all(t.get("bin_size") == 0 for t in tensors), tensors |
| 123 | + else: |
| 124 | + assert bin_name, f"manifest missing bin_file field: {data}" |
| 125 | + bin_path = dump_dir / bin_name |
| 126 | + assert bin_path.exists(), f"manifest names bin_file={bin_name!r} but {bin_path} not found" |
| 127 | + assert bin_path.stat().st_size > 0, "args.bin is empty" |
| 128 | + |
| 129 | + # Unified manifest (#792): tensors and scalar args share one |
| 130 | + # args_dump.json keyed by a "kind" field; no separate legacy sidecar files. |
| 131 | + assert not (dump_dir / "tensor_dump.json").exists(), "tensor_dump.json should not be emitted" |
| 132 | + assert not (dump_dir / "kernel_args_dump.json").exists(), "kernel_args_dump.json should not be emitted" |
| 133 | + assert all("kind" in t for t in tensors), tensors |
| 134 | + scalar_entries = [t for t in tensors if t.get("kind") == "scalar"] |
| 135 | + assert all(t.get("stage") == "before_dispatch" for t in scalar_entries), scalar_entries |
| 136 | + assert all(t.get("bin_size") == 0 for t in scalar_entries), scalar_entries |
| 137 | + assert all("value" in t for t in scalar_entries), scalar_entries |
| 138 | + |
| 139 | + # func_id array (#1181): every entry carries its task's active-subtask |
| 140 | + # membership. partial_dump_orch dispatches all tasks via single-kernel |
| 141 | + # rt_submit_aiv_task, so each func_id must be a one-element array holding |
| 142 | + # a declared func_id, consistent across that task's entries; and all three |
| 143 | + # dispatched kernels (0/1/2) appear. Without this a wholly broken func_id |
| 144 | + # emit still PASSes. (Exact task_id->func mapping is not pinned: it shifts |
| 145 | + # between partial/full modes and add/mul are structurally indistinguishable.) |
| 146 | + per_task_func = {} |
| 147 | + seen_func = set() |
| 148 | + for t in tensors: |
| 149 | + fid = t.get("func_id") |
| 150 | + assert isinstance(fid, list) and len(fid) == 1, ( |
| 151 | + f"{t['task_id']} arg {t.get('arg_index')} ({t.get('kind')}): " |
| 152 | + f"func_id={fid} — single-kernel task must be a one-element array" |
| 153 | + ) |
| 154 | + assert fid[0] in (0, 1, 2), f"{t['task_id']}: func_id {fid} outside declared range 0/1/2" |
| 155 | + prev = per_task_func.setdefault(t["task_id"], fid[0]) |
| 156 | + assert prev == fid[0], f"{t['task_id']}: func_id differs across entries ({prev} vs {fid[0]})" |
| 157 | + seen_func.add(fid[0]) |
| 158 | + assert seen_func == {0, 1, 2}, f"dump should cover all dispatched kernels 0/1/2, got {sorted(seen_func)}" |
| 159 | + |
| 160 | + # Level-aware checks operate on the tensor entries. |
| 161 | + tensor_entries = [t for t in tensors if t.get("kind") == "tensor"] |
| 162 | + task_ids = {t["task_id"] for t in tensor_entries} |
| 163 | + if level == 1: |
| 164 | + # Partial: only the selected tensor/scalar args, race-free (host-latched). |
| 165 | + assert len(tensor_entries) == 7, f"partial expected 7 tensor entries, got {len(tensor_entries)}" |
| 166 | + assert task_ids == { |
| 167 | + "0x0000000100000000", |
| 168 | + "0x0000000100000002", |
| 169 | + "0x0000000100000003", |
| 170 | + } |
| 171 | + # Task granularity: dump() captured all tensor args on task 0. |
| 172 | + t00 = sorted(t["arg_index"] for t in tensor_entries if t["task_id"] == "0x0000000100000000") |
| 173 | + assert t00 == [0, 1] |
| 174 | + # Mixed granularity: dump(d, inter_ci, t3_count) captured tensor args |
| 175 | + # 0 + 2, not arg 1 (e). |
| 176 | + t02 = sorted(t["arg_index"] for t in tensor_entries if t["task_id"] == "0x0000000100000002") |
| 177 | + assert t02 == [0, 2] |
| 178 | + # Tensor-only task granularity: dump() captured all three tensor args. |
| 179 | + t03 = sorted(t["arg_index"] for t in tensor_entries if t["task_id"] == "0x0000000100000003") |
| 180 | + assert t03 == [0, 1, 2] |
| 181 | + scalar_by_task = { |
| 182 | + task_id: sorted(t["arg_index"] for t in scalar_entries if t["task_id"] == task_id) |
| 183 | + for task_id in {t["task_id"] for t in scalar_entries} |
| 184 | + } |
| 185 | + assert scalar_by_task == { |
| 186 | + "0x0000000100000000": [2, 3], |
| 187 | + "0x0000000100000001": [2], |
| 188 | + "0x0000000100000002": [3], |
| 189 | + } |
| 190 | + ambiguous_scalars = [ |
| 191 | + (t["task_id"], t["arg_index"]) for t in scalar_entries if t.get("arg_index_ambiguous", False) |
| 192 | + ] |
| 193 | + assert ambiguous_scalars == [("0x0000000100000002", 3)] |
| 194 | + else: |
| 195 | + # Full (level 2 or 3): markers ignored — every one of the 5 tasks is dumped. |
| 196 | + assert len(task_ids) >= 5, f"full dump should cover all 5 tasks, got {sorted(task_ids)}" |
| 197 | + |
| 198 | + # ---- Tool smoke: dump_viewer ---- |
| 199 | + # Exit-code-only check; the no-filter default lists every captured |
| 200 | + # arg without exporting. A schema change that breaks the viewer |
| 201 | + # fires here in the same CI step that produced the dump. |
| 202 | + subprocess.run( |
| 203 | + [sys.executable, "-m", "simpler_setup.tools.dump_viewer", str(dump_dir)], |
| 204 | + check=True, |
| 205 | + timeout=60, |
| 206 | + ) |
| 207 | + |
| 208 | + |
| 209 | +if __name__ == "__main__": |
| 210 | + SceneTestCase.run_module(__name__) |
0 commit comments