Add: Insight Trace workspace generation for MindStudio profiling

simpler_setup/insight_trace/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) PyPTO Contributors.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+"""MindStudio Insight trace workspace generation."""

simpler_setup/insight_trace/arg_resolver.py

@@ -0,0 +1,40 @@
+# Copyright (c) PyPTO Contributors.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+from __future__ import annotations
+
+from pathlib import Path
+
+from .args_pkg import _load_arg_spec, load_kernel_dump_args, resolve_builtin_args, resolve_scene_test_args
+from .models import KernelSpec, SceneCaseContext, TraceArg
+
+
+def resolve_args(
+    context: SceneCaseContext,
+    kernel: KernelSpec,
+    arg_spec: Path | None = None,
+    dump_dir: Path | None = None,
+    dispatch_id: int | None = None,
+) -> tuple[TraceArg, ...]:
+    if arg_spec is not None:
+        return _load_arg_spec(arg_spec)
+    if dump_dir is not None:
+        if dispatch_id is None:
+            raise ValueError("--dispatch-id is required with --dump-dir")
+        return load_kernel_dump_args(dump_dir, kernel.func_id, dispatch_id)
+    if _should_use_builtin_fallback(context):
+        return resolve_builtin_args(context, kernel)
+    try:
+        return resolve_scene_test_args(context)
+    except Exception:
+        return resolve_builtin_args(context, kernel)
+
+
+def _should_use_builtin_fallback(context: SceneCaseContext) -> bool:
+    module_path = context.module_dir.as_posix()
+    return "paged_attention" in module_path or "spmd_multiblock_mix" in module_path

simpler_setup/insight_trace/args_pkg/__init__.py

@@ -0,0 +1,11 @@
+from .from_dump import load_kernel_dump_args
+from .from_scene_test import resolve_scene_test_args
+from .from_spec import _load_arg_spec
+from .recipes import resolve_builtin_args
+
+__all__ = [
+    "_load_arg_spec",
+    "load_kernel_dump_args",
+    "resolve_builtin_args",
+    "resolve_scene_test_args",
+]

simpler_setup/insight_trace/args_pkg/from_dump.py

@@ -0,0 +1,65 @@
+# Copyright (c) PyPTO Contributors.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ..models import TraceArg, TraceScalarArg, TraceTensorArg
+
+
+def load_kernel_dump_args(dump_dir: Path, func_id: int, dispatch_id: int) -> tuple[TraceArg, ...]:
+    dump_path = _kernel_dump_path(dump_dir)
+    raw = json.loads(dump_path.read_text())
+    dispatch = None
+    for item in raw.get("dispatches", []):
+        if int(item.get("func_id", -1)) == func_id and int(item.get("dispatch_id", -1)) == dispatch_id:
+            dispatch = item
+            break
+    if dispatch is None:
+        raise ValueError(f"No kernel args dump dispatch matches func_id={func_id}, dispatch_id={dispatch_id}")
+
+    result: list[TraceArg] = []
+    for item in dispatch.get("args", []):
+        index = int(item["arg_index"])
+        kind = item["kind"]
+        if kind == "tensor":
+            result.append(
+                TraceTensorArg(
+                    index=index,
+                    name=f"arg{index}",
+                    dtype=item["dtype"],
+                    shape=tuple(int(dim) for dim in item["shape"]),
+                )
+            )
+        elif kind == "scalar":
+            value = item["value"]
+            pack_mode = item.get("pack_mode", "value")
+            if (pack_mode == "bits" or item["dtype"] == "FLOAT32_BITS") and isinstance(value, float):
+                value = _f32_bits(value)
+            result.append(TraceScalarArg(index, f"arg{index}", item["dtype"], value, pack_mode))
+        elif kind in {"local_context", "global_context"}:
+            continue
+        else:
+            raise ValueError(f"Unknown kernel dump arg kind: {kind}")
+    return tuple(sorted(result, key=lambda arg: arg.index))
+
+
+def _kernel_dump_path(dump_dir: Path) -> Path:
+    candidates = (dump_dir / "kernel_args_dump.json", dump_dir / "tensor_dump" / "kernel_args_dump.json")
+    for path in candidates:
+        if path.is_file():
+            return path
+    raise ValueError(f"kernel_args_dump.json not found under {dump_dir}")
+
+
+def _f32_bits(value: float) -> int:
+    import struct
+
+    return struct.unpack("I", struct.pack("f", value))[0]

simpler_setup/insight_trace/args_pkg/from_scene_test.py

@@ -0,0 +1,142 @@
+# Copyright (c) PyPTO Contributors.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+from __future__ import annotations
+
+import ctypes
+
+from simpler.task_interface import ArgDirection
+from simpler_setup.scene_test import Scalar, Tensor
+
+from ..models import SceneCaseContext, TraceArg, TraceScalarArg, TraceTensorArg
+
+_TORCH_DTYPE_MAP = {
+    "torch.float32": "FLOAT32",
+    "torch.float16": "FLOAT16",
+    "torch.bfloat16": "BFLOAT16",
+    "torch.int32": "INT32",
+    "torch.int64": "INT64",
+    "torch.uint8": "UINT8",
+    "torch.int8": "INT8",
+    "torch.bool": "BOOL",
+}
+
+_CTYPES_DTYPE_MAP = {
+    ctypes.c_bool: "BOOL",
+    ctypes.c_int8: "INT8",
+    ctypes.c_uint8: "UINT8",
+    ctypes.c_int16: "INT16",
+    ctypes.c_uint16: "UINT16",
+    ctypes.c_int32: "INT32",
+    ctypes.c_uint32: "UINT32",
+    ctypes.c_int64: "INT64",
+    ctypes.c_uint64: "UINT64",
+    ctypes.c_float: "FLOAT32_BITS",
+    ctypes.c_double: "FLOAT64_BITS",
+}
+
+_ROLE_MAP = {
+    ArgDirection.IN: "input",
+    ArgDirection.OUT: "output",
+    ArgDirection.INOUT: "inout",
+}
+
+
+def resolve_scene_test_args(context: SceneCaseContext) -> tuple[TraceArg, ...]:
+    builder = context.test_class().generate_args(context.case.get("params", {}))
+    orch_signature = context.callable_spec.get("orchestration", {}).get("signature")
+    if orch_signature is None:
+        raise ValueError("No orchestration signature available for generic insight trace arg inference")
+
+    result: list[TraceArg] = []
+    tensor_index = 0
+    arg_index = 0
+    for spec in builder.specs:
+        if isinstance(spec, Tensor):
+            if tensor_index >= len(orch_signature):
+                raise ValueError(
+                    f"Tensor '{spec.name}' at index {tensor_index} has no matching orchestration signature entry"
+                )
+            result.append(
+                TraceTensorArg(
+                    index=arg_index,
+                    name=spec.name,
+                    dtype=_tensor_dtype_name(spec.value),
+                    shape=tuple(int(dim) for dim in spec.value.shape),
+                    role=_ROLE_MAP.get(orch_signature[tensor_index], "input"),
+                )
+            )
+            tensor_index += 1
+            arg_index += 1
+            continue
+
+        if isinstance(spec, Scalar):
+            result.append(
+                TraceScalarArg(
+                    index=arg_index,
+                    name=spec.name,
+                    dtype=_scalar_dtype_name(spec.value),
+                    value=_scalar_value(spec.value),
+                    pack_mode=_scalar_pack_mode(spec.value),
+                )
+            )
+            arg_index += 1
+            continue
+
+        raise ValueError(f"Unsupported TaskArgsBuilder spec type: {type(spec).__name__}")
+
+    if tensor_index != len(orch_signature):
+        raise ValueError(
+            f"Orchestration signature length {len(orch_signature)} does not match tensor count {tensor_index}"
+        )
+    return tuple(result)
+
+
+def _tensor_dtype_name(value) -> str:
+    key = str(value.dtype)
+    try:
+        return _TORCH_DTYPE_MAP[key]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported tensor dtype for insight trace: {key}") from exc
+
+
+def _scalar_dtype_name(value) -> str:
+    if type(value) in _CTYPES_DTYPE_MAP:
+        return _CTYPES_DTYPE_MAP[type(value)]
+    if isinstance(value, bool):
+        return "BOOL"
+    if isinstance(value, int):
+        return "INT64"
+    if isinstance(value, float):
+        return "FLOAT32_BITS"
+    raise ValueError(f"Unsupported scalar type for insight trace: {type(value).__name__}")
+
+
+def _scalar_value(value):
+    if isinstance(value, ctypes._SimpleCData):
+        raw = value.value
+    else:
+        raw = value
+    if isinstance(raw, bool):
+        return int(raw)
+    if isinstance(raw, float):
+        return _f32_bits(raw)
+    return raw
+
+
+def _scalar_pack_mode(value) -> str:
+    dtype = _scalar_dtype_name(value)
+    if dtype.endswith("_BITS"):
+        return "bits"
+    return "value"
+
+
+def _f32_bits(value: float) -> int:
+    import struct
+
+    return struct.unpack("I", struct.pack("f", value))[0]

simpler_setup/insight_trace/args_pkg/from_spec.py

@@ -0,0 +1,54 @@
+# Copyright (c) PyPTO Contributors.
+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
+# CANN Open Software License Agreement Version 2.0 (the "License").
+# Please refer to the License for details. You may not use this file except in compliance with the License.
+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
+# See LICENSE in the root of the software repository for the full text of the License.
+# -----------------------------------------------------------------------------------------------------------
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from ..models import TraceArg, TraceScalarArg, TraceTensorArg
+
+
+def _load_arg_spec(path: Path) -> tuple[TraceArg, ...]:
+    raw = json.loads(path.read_text())
+    result: list[TraceArg] = []
+    for item in raw.get("args", raw):
+        if item["kind"] == "tensor":
+            result.append(
+                TraceTensorArg(
+                    index=int(item["index"]),
+                    name=item["name"],
+                    dtype=item["dtype"],
+                    shape=tuple(int(dim) for dim in item["shape"]),
+                    role=item.get("role", "input"),
+                    fill=item.get("fill", "zero"),
+                )
+            )
+        elif item["kind"] == "scalar":
+            value = item["value"]
+            pack_mode = item.get("pack_mode", "value")
+            if (pack_mode == "bits" or item["dtype"] == "FLOAT32_BITS") and isinstance(value, float):
+                value = _f32_bits(value)
+            result.append(
+                TraceScalarArg(
+                    index=int(item["index"]),
+                    name=item["name"],
+                    dtype=item["dtype"],
+                    value=value,
+                    pack_mode=pack_mode,
+                )
+            )
+        else:
+            raise ValueError(f"Unknown arg kind: {item['kind']}")
+    return tuple(sorted(result, key=lambda arg: arg.index))
+
+
+def _f32_bits(value: float) -> int:
+    import struct
+
+    return struct.unpack("I", struct.pack("f", value))[0]