From 17fa04aa4b2e71af2d66242e0571aa6ec9785bcd Mon Sep 17 00:00:00 2001 From: vegetabledoww <17863935975@163.com> Date: Thu, 11 Jun 2026 11:17:12 +0800 Subject: [PATCH] Fix typed scalar dump propagation + rollback unrelated changes Core fix: - Submit-side per-task scalar dtype table registered at submit time - AICPU writes dtype into each SCALAR dump record directly - Host formats JSON from record.dtype (not host-side join) Rollback (unrelated changes from PR 966): - Remove tensor_count from scalar dtype registration chain - Restore callable signature plumbing (signature/sig_count propagation) - Revert add_scalars(values, dtypes, count) overload to (values, count) - Restore NEON fast path in add_scalars_i32() and its arm_neon.h include - Revert dump_viewer.py list_tensors() table format (keep scalar value display) Co-Authored-By: Claude Sonnet 4.6 --- docs/dfx/tensor-dump.md | 26 +++-- simpler_setup/tools/dump_viewer.py | 11 ++ .../runtime/pto_orchestrator.cpp | 22 ++-- .../runtime/pto_types.h | 10 +- .../runtime/pto_orchestrator.cpp | 23 ++-- .../runtime/pto_types.h | 14 ++- .../include/aicpu/tensor_dump_aicpu.h | 39 ++++--- .../shared/aicpu/tensor_dump_aicpu.cpp | 104 +++++++++++++++++- .../shared/host/tensor_dump_collector.cpp | 53 ++++++++- src/common/task_interface/data_type.h | 41 ++++++- 10 files changed, 287 insertions(+), 56 deletions(-) diff --git a/docs/dfx/tensor-dump.md b/docs/dfx/tensor-dump.md index 71a7c39112..a57619b409 100644 --- a/docs/dfx/tensor-dump.md +++ b/docs/dfx/tensor-dump.md @@ -26,9 +26,11 @@ saw, without the timing distortion of inline printing. - **Manifest + binary payload.** `--dump-tensor` writes a JSON manifest plus one `.bin` payload per run; tensor entries carry `bin_offset` / `bin_size`, while scalar entries stay manifest-only. -- **Unified scalar args.** Scalar callable slots are emitted as +- **Unified scalar args.** Scalar values are emitted as `kind: scalar`, `stage: before_dispatch`, zero-dim records in - `tensor_dump.json`; there is no separate args-only manifest. + `tensor_dump.json`; there is no separate args-only manifest. Their + final `dtype` is registered at submit time in a dump-only per-task side table + and resolved by AICPU when writing each scalar dump record. - **Cross-architecture.** Same flags and on-disk layout family on `a2a3` and `a5`. Both runtimes are wired through. @@ -274,8 +276,9 @@ What you can read out of `tensor_dump.json` + `tensor_dump.bin`: - **Per-task input snapshots** (`role: input`, `stage: before_dispatch`) — what each kernel was given. - **Per-dispatch scalar args** (`kind: scalar`, `stage: - before_dispatch`) — the raw callable-slot scalar values AICPU - handed to the kernel. + before_dispatch`) — the raw per-task scalar-slot values AICPU + handed to the kernel, tagged with the dump-only dtype captured at + submit time. - **Per-task output snapshots** (`role: output`, `stage: after_completion`) — what each kernel produced. The barrier ensures these reflect the kernel's final writes. @@ -284,8 +287,8 @@ What you can read out of `tensor_dump.json` + `tensor_dump.bin`: - **Logical view reconstruction** — `shape` / `strides` / `start_offset` / `is_contiguous` plus the gathered logical-contiguous payload. -- **Per-task identity** — `task_id` / `subtask_id` / `func_id` - correlates dump entries with swimlane and PMU rows. +- **Per-task identity** — `task_id`, `stage`, `role`, and `arg_index` + identify each dumped argument within a task. - **Loss accounting** — `truncated` / `overwritten` per-record flags, plus aggregate `dropped_records` / `dropped_overwrite` in the summary. @@ -361,10 +364,13 @@ Each runtime's scheduler dispatch code calls └──────────────────────────────────────┘ ``` -`dump_tensors_for_task` walks the formal callable signature, -matches each non-scalar slot to a `TensorDumpInfo` (dtype + shape + -strides + start offset + device address), and calls -`dump_tensor_record` for slots that match the current stage. +`dump_tensors_for_task` walks the formal callable signature for +non-scalar tensor slots, matches each tensor slot to a +`TensorDumpInfo` (dtype + shape + strides + start offset + device +address), and calls `dump_tensor_record` for slots that match the +current stage. Scalar values are dumped separately from the flat +payload `scalars[]`; their dtype table is registered at submit time +and emitted as a dump-only metadata record at `BEFORE_DISPATCH`. When dump is enabled, AICore executors also issue `pipe_barrier(PIPE_ALL)` after kernel execution and before writing diff --git a/simpler_setup/tools/dump_viewer.py b/simpler_setup/tools/dump_viewer.py index c38d5e9c4b..a141133005 100644 --- a/simpler_setup/tools/dump_viewer.py +++ b/simpler_setup/tools/dump_viewer.py @@ -51,8 +51,11 @@ "int64": ("q", 8), "uint64": ("Q", 8), "int16": ("h", 2), + "uint16": ("H", 2), "int8": ("b", 1), "uint8": ("B", 1), + "uint32": ("I", 4), + "bool": ("?", 1), } @@ -103,6 +106,7 @@ def write_tensor(tensor: dict, bin_path: Path, out): out.write(f"# stage: {t['stage']}\n") out.write(f"# arg_index: {t['arg_index']}\n") out.write(f"# dtype: {t['dtype']}\n") + out.write(f"# kind: {t.get('kind', 'tensor')}\n") out.write(f"# is_contiguous: {t['is_contiguous']}\n") out.write(f"# shape: {t['shape']}\n") out.write(f"# strides: {t['strides']}\n") @@ -114,6 +118,13 @@ def write_tensor(tensor: dict, bin_path: Path, out): if t.get("truncated"): out.write("# DATA TRUNCATED (tensor too large for arena)\n") + if t.get("kind") == "scalar": + val = t.get("value") + if t.get("dtype", "").upper() == "BOOL": + val = "true" if val else "false" + out.write(f"# value: {val}\n") + return + bin_size = t.get("bin_size", 0) if bin_size == 0: out.write("# (no data)\n") diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 08602e4b93..a7ca47199f 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -36,10 +36,9 @@ #include "pto_types.h" #include "tensor.h" -extern "C" void set_dump_tensor_task_mask(uint64_t task_id, uint64_t mask); - #if PTO2_PROFILING #include "aicpu/scope_stats_collector_aicpu.h" +#include "aicpu/tensor_dump_aicpu.h" #endif // Verify the captured Tensor blob size in DepGenRecord matches the runtime @@ -648,12 +647,19 @@ static TaskOutputTensors submit_task_common( payload.init(args, result, prepared.alloc_result, layout); #if PTO2_PROFILING - // Selective vs full dump is latched at dump_tensor_init from DumpDataHeader - // (host-decided before any dispatch), so it is race-free regardless of - // submission order. Here we only record each marked task's arg mask, which - // selective collection consults. - if (args.tensor_dump_arg_mask() != 0) { - set_dump_tensor_task_mask(task_id.raw, args.tensor_dump_arg_mask()); + if (is_dump_tensor_enabled()) { + if (args.scalar_count() > 0) { + set_dump_tensor_task_scalar_dtypes( + task_id.raw, static_cast(args.scalar_count()), args.scalar_dtypes() + ); + } + // Selective vs full dump is latched at dump_tensor_init from DumpDataHeader + // (host-decided before any dispatch), so it is race-free regardless of + // submission order. Here we only record each marked task's arg mask, which + // selective collection consults. + if (args.tensor_dump_arg_mask() != 0) { + set_dump_tensor_task_mask(task_id.raw, args.tensor_dump_arg_mask()); + } } #endif diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index 09b7ba4bf8..1e563ab0b8 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -32,6 +32,7 @@ #include #endif +#include "data_type.h" #include "pto_submit_types.h" #include "task_args.h" #include "tensor.h" @@ -183,6 +184,7 @@ struct Arg : TaskArgsTpl>>(), ++scalar_count_), + ...); } void add_scalars(const uint64_t *values, int count) { @@ -370,15 +374,19 @@ struct Arg : TaskArgsTpl 0) { + set_dump_tensor_task_scalar_dtypes( + task_id.raw, static_cast(args.scalar_count()), args.scalar_dtypes() + ); + } + // Selective vs full dump is latched at dump_tensor_init from DumpDataHeader + // (host-decided before any dispatch), so it is race-free regardless of + // submission order. Here we only record each marked task's arg mask, which + // selective collection consults. + if (args.tensor_dump_arg_mask() != 0) { + set_dump_tensor_task_mask(task_id.raw, args.tensor_dump_arg_mask()); + } } #endif diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index 825d2998e1..815d20b22c 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -28,12 +28,13 @@ #include #include -#include - #if defined(__aarch64__) #include #endif +#include + +#include "data_type.h" #include "pto_submit_types.h" #include "task_args.h" #include "tensor.h" @@ -181,6 +182,7 @@ struct Arg : TaskArgsTpl>>(), ++scalar_count_), + ...); } void add_scalars(const uint64_t *values, int count) { @@ -374,15 +378,19 @@ struct Arg : TaskArgsTpl(slot_state.task->task_id.raw), total_tensor_args, pl.tensor_count ); } @@ -172,19 +174,26 @@ inline void dump_tensors_for_task( arg_index++; tensor_index++; } - if (stage == TensorDumpStage::BEFORE_DISPATCH && pl.scalar_count > 0) { - for (int32_t i = 0; i < pl.scalar_count; i++) { - TensorDumpInfo info = {}; - info.task_id = slot_state.task->task_id.raw; - info.role = TensorDumpRole::INPUT; - info.stage = stage; - info.dtype = static_cast(DataType::UINT64); - info.ndims = 0; - info.arg_index = static_cast(pl.tensor_count + i); - info.kind = static_cast(TensorDumpKind::SCALAR); - info.scalar_value = pl.scalars[i]; - dump_tensor_record(thread_idx, info); - } + } + + if (stage == TensorDumpStage::BEFORE_DISPATCH && pl.scalar_count > 0) { + uint8_t scalar_dtypes[CORE_MAX_SCALAR_ARGS] = {}; + uint32_t dtype_scalar_count = 0; + bool has_scalar_dtypes = + get_dump_tensor_task_scalar_dtypes(slot_state.task->task_id.raw, &dtype_scalar_count, scalar_dtypes); + for (int32_t scalar_index = 0; scalar_index < pl.scalar_count; scalar_index++) { + TensorDumpInfo info = {}; + info.task_id = slot_state.task->task_id.raw; + info.role = TensorDumpRole::INPUT; + info.stage = stage; + info.dtype = (has_scalar_dtypes && scalar_index < static_cast(dtype_scalar_count)) ? + scalar_dtypes[scalar_index] : + static_cast(DataType::UINT64); + info.ndims = 0; + info.arg_index = static_cast(pl.tensor_count + scalar_index); + info.kind = static_cast(TensorDumpKind::SCALAR); + info.scalar_value = pl.scalars[scalar_index]; + dump_tensor_record(thread_idx, info); } } } diff --git a/src/common/platform/shared/aicpu/tensor_dump_aicpu.cpp b/src/common/platform/shared/aicpu/tensor_dump_aicpu.cpp index 697bac22e9..fdef9468ac 100644 --- a/src/common/platform/shared/aicpu/tensor_dump_aicpu.cpp +++ b/src/common/platform/shared/aicpu/tensor_dump_aicpu.cpp @@ -59,9 +59,15 @@ struct DumpTaskMaskEntry { uint64_t task_id; TensorDumpArgMask mask; }; +struct DumpTaskScalarDtypeEntry { + uint64_t task_id; + uint32_t scalar_count; + uint8_t scalar_dtypes[32]; +}; static constexpr uint64_t DUMP_TASK_MASK_EMPTY_TASK_ID = UINT64_MAX; static constexpr uint32_t DUMP_TASK_MASK_TABLE_CAPACITY = 32768; static DumpTaskMaskEntry *g_dump_mask_table = nullptr; +static DumpTaskScalarDtypeEntry *g_dump_scalar_dtype_table = nullptr; static bool ensure_dump_mask_table() { if (g_dump_mask_table != nullptr) { return true; @@ -79,14 +85,101 @@ static bool ensure_dump_mask_table() { return true; } +static bool ensure_dump_scalar_dtype_table() { + if (g_dump_scalar_dtype_table != nullptr) { + return true; + } + g_dump_scalar_dtype_table = static_cast( + malloc(sizeof(DumpTaskScalarDtypeEntry) * DUMP_TASK_MASK_TABLE_CAPACITY) + ); + if (g_dump_scalar_dtype_table == nullptr) { + LOG_ERROR("Failed to allocate tensor dump scalar dtype table"); + return false; + } + for (uint32_t i = 0; i < DUMP_TASK_MASK_TABLE_CAPACITY; i++) { + g_dump_scalar_dtype_table[i].task_id = DUMP_TASK_MASK_EMPTY_TASK_ID; + g_dump_scalar_dtype_table[i].scalar_count = 0; + memset(g_dump_scalar_dtype_table[i].scalar_dtypes, 0, sizeof(g_dump_scalar_dtype_table[i].scalar_dtypes)); + } + return true; +} + static void clear_dump_mask_table() { - if (g_dump_mask_table == nullptr) { + if (g_dump_mask_table != nullptr) { + for (uint32_t i = 0; i < DUMP_TASK_MASK_TABLE_CAPACITY; i++) { + g_dump_mask_table[i].task_id = DUMP_TASK_MASK_EMPTY_TASK_ID; + g_dump_mask_table[i].mask = TENSOR_DUMP_ARG_MASK_NONE; + } + } + if (g_dump_scalar_dtype_table != nullptr) { + for (uint32_t i = 0; i < DUMP_TASK_MASK_TABLE_CAPACITY; i++) { + g_dump_scalar_dtype_table[i].task_id = DUMP_TASK_MASK_EMPTY_TASK_ID; + g_dump_scalar_dtype_table[i].scalar_count = 0; + memset(g_dump_scalar_dtype_table[i].scalar_dtypes, 0, sizeof(g_dump_scalar_dtype_table[i].scalar_dtypes)); + } + } +} + +static bool resolve_dump_task_slot(uint64_t task_id, uint32_t *idx_out) { + uint32_t ring_id = static_cast(task_id >> 32); + if (ring_id >= TENSOR_DUMP_MASK_POOL_MAX_RINGS) { + return false; + } + uint32_t slot = static_cast(task_id) & TENSOR_DUMP_MASK_POOL_DEFAULT_SLOT_MASK; + *idx_out = (ring_id * TENSOR_DUMP_MASK_POOL_MAX_SLOTS + slot) & (DUMP_TASK_MASK_TABLE_CAPACITY - 1); + return true; +} + +extern "C" void +set_dump_tensor_task_scalar_dtypes(uint64_t task_id, uint32_t scalar_count, const uint8_t *scalar_dtypes) { + if (scalar_count == 0 || scalar_dtypes == nullptr) { return; } - for (uint32_t i = 0; i < DUMP_TASK_MASK_TABLE_CAPACITY; i++) { - g_dump_mask_table[i].task_id = DUMP_TASK_MASK_EMPTY_TASK_ID; - g_dump_mask_table[i].mask = TENSOR_DUMP_ARG_MASK_NONE; + if (scalar_count > 32) { + LOG_ERROR("tensor dump scalar dtype count exceeds max (%u > 32)", scalar_count); + return; + } + if (!ensure_dump_scalar_dtype_table()) { + return; + } + uint32_t idx = 0; + if (!resolve_dump_task_slot(task_id, &idx)) { + return; + } + for (uint32_t probe = 0; probe < DUMP_TASK_MASK_TABLE_CAPACITY; probe++) { + DumpTaskScalarDtypeEntry &entry = + g_dump_scalar_dtype_table[(idx + probe) & (DUMP_TASK_MASK_TABLE_CAPACITY - 1)]; + if (entry.task_id == DUMP_TASK_MASK_EMPTY_TASK_ID || entry.task_id == task_id) { + entry.task_id = task_id; + entry.scalar_count = scalar_count; + memcpy(entry.scalar_dtypes, scalar_dtypes, scalar_count * sizeof(uint8_t)); + return; + } } + LOG_ERROR("tensor dump scalar dtype table is full"); +} + +extern "C" bool get_dump_tensor_task_scalar_dtypes(uint64_t task_id, uint32_t *scalar_count, uint8_t *scalar_dtypes) { + if (g_dump_scalar_dtype_table == nullptr || scalar_count == nullptr || scalar_dtypes == nullptr) { + return false; + } + uint32_t idx = 0; + if (!resolve_dump_task_slot(task_id, &idx)) { + return false; + } + for (uint32_t probe = 0; probe < DUMP_TASK_MASK_TABLE_CAPACITY; probe++) { + const DumpTaskScalarDtypeEntry &entry = + g_dump_scalar_dtype_table[(idx + probe) & (DUMP_TASK_MASK_TABLE_CAPACITY - 1)]; + if (entry.task_id == task_id) { + *scalar_count = entry.scalar_count; + memcpy(scalar_dtypes, entry.scalar_dtypes, entry.scalar_count * sizeof(uint8_t)); + return true; + } + if (entry.task_id == DUMP_TASK_MASK_EMPTY_TASK_ID) { + return false; + } + } + return false; } extern "C" void set_dump_tensor_enabled(bool enable) { @@ -509,13 +602,14 @@ int dump_tensor_record(int thread_idx, const TensorDumpInfo &info) { // Reserve space in arena // Compute actual tensor data size from shape (not buffer.size which may include padding) + TensorDumpKind kind = static_cast(info.kind); uint64_t actual_elements = get_tensor_dump_num_elements(info); uint64_t elem_sz = get_element_size(static_cast(info.dtype)); uint64_t bytes = actual_elements * elem_sz; uint64_t copy_bytes = bytes; bool truncated = false; bool is_contiguous = tensor_dump_is_contiguous(info); - bool is_scalar = static_cast(info.kind) == TensorDumpKind::SCALAR; + bool is_scalar = kind == TensorDumpKind::SCALAR; if (is_scalar) { copy_bytes = 0; diff --git a/src/common/platform/shared/host/tensor_dump_collector.cpp b/src/common/platform/shared/host/tensor_dump_collector.cpp index 2fd9836977..ab5d943783 100644 --- a/src/common/platform/shared/host/tensor_dump_collector.cpp +++ b/src/common/platform/shared/host/tensor_dump_collector.cpp @@ -27,9 +27,11 @@ #include "host/tensor_dump_collector.h" +#include "data_type.h" + #include #include -#include +#include #include #include #include @@ -226,13 +228,14 @@ void TensorDumpCollector::process_dump_buffer(const DumpReadyBufferInfo &info) { dt.is_contiguous = (rec.is_contiguous != 0); dt.truncated = (rec.truncated != 0); dt.overwritten = false; - if (dt.truncated && ++total_truncated_count_ == 1) { - LOG_WARN("Tensor dump truncation detected. Increase PLATFORM_DUMP_AVG_TENSOR_BYTES."); - } dt.start_offset = rec.start_offset; std::memcpy(dt.shapes, rec.shapes, sizeof(dt.shapes)); std::memcpy(dt.strides, rec.strides, sizeof(dt.strides)); + if (dt.truncated && ++total_truncated_count_ == 1) { + LOG_WARN("Tensor dump truncation detected. Increase PLATFORM_DUMP_AVG_TENSOR_BYTES."); + } + if (dt.kind == TensorDumpKind::TENSOR && thread_idx >= 0 && thread_idx < static_cast(arenas_.size())) { ArenaInfo &ai = arenas_[thread_idx]; char *arena_host = reinterpret_cast(ai.host_ptr); @@ -402,6 +405,43 @@ static const char *tensor_dump_kind_name(TensorDumpKind kind) { return "unknown"; } +static void write_scalar_json_value(std::ofstream &json, const DumpedTensor &dt) { + uint64_t raw = dt.scalar_value; + if (dt.dtype == static_cast(DataType::FLOAT32)) { + float f; + memcpy(&f, &raw, sizeof(float)); + if (std::isnan(f)) { + json << ", \"value\": null"; + } else if (std::isinf(f)) { + json << ", \"value\": " << (f < 0 ? "\"-$Inf\"" : "\"$Inf\""); + } else { + std::ostringstream val_ss; + val_ss << f; + std::string val_str = val_ss.str(); + if (val_str.find('.') == std::string::npos && val_str.find('e') == std::string::npos) { + val_str += ".0"; + } + json << ", \"value\": " << val_str; + } + } else if (dt.dtype == static_cast(DataType::INT32)) { + int32_t val; + memcpy(&val, &raw, sizeof(int32_t)); + json << ", \"value\": " << val; + } else if (dt.dtype == static_cast(DataType::UINT32)) { + uint32_t val; + memcpy(&val, &raw, sizeof(uint32_t)); + json << ", \"value\": " << val; + } else if (dt.dtype == static_cast(DataType::BOOL)) { + json << ", \"value\": " << (raw != 0 ? "true" : "false"); + } else if (dt.dtype == static_cast(DataType::INT64)) { + int64_t val; + memcpy(&val, &raw, sizeof(int64_t)); + json << ", \"value\": " << val; + } else { + json << ", \"value\": " << raw; + } +} + static std::string dims_to_string(const uint32_t dims[], int ndims) { std::ostringstream ss; ss << "["; @@ -555,14 +595,15 @@ int TensorDumpCollector::export_dump_files() { first_entry = false; json << " {\"task_id\": \"0x" << std::hex << std::setfill('0') << std::setw(16) << dt.task_id << std::dec - << "\", \"arg_index\": " << dt.arg_index << ", \"role\": \"" << tensor_dump_role_name(dt.role) + << "\""; + json << ", \"arg_index\": " << dt.arg_index << ", \"role\": \"" << tensor_dump_role_name(dt.role) << "\", \"stage\": \"" << tensor_dump_stage_name(dt.stage) << "\", \"kind\": \"" << tensor_dump_kind_name(dt.kind) << "\", \"dtype\": \"" << dtype_name << "\", \"is_contiguous\": " << (dt.is_contiguous ? "true" : "false") << ", \"shape\": " << shape_str << ", \"strides\": " << strides_str << ", \"start_offset\": " << dt.start_offset << ", \"numel\": " << numel; if (dt.kind == TensorDumpKind::SCALAR) { - json << ", \"value\": " << dt.scalar_value; + write_scalar_json_value(json, dt); } json << ", \"bin_offset\": " << dt.bin_offset << ", \"bin_size\": " << dt.payload_size << ", \"truncated\": " << (dt.truncated ? "true" : "false") diff --git a/src/common/task_interface/data_type.h b/src/common/task_interface/data_type.h index 902fa503ef..ae990ad350 100644 --- a/src/common/task_interface/data_type.h +++ b/src/common/task_interface/data_type.h @@ -43,6 +43,7 @@ enum class DataType : uint8_t { UINT64, // 8 bytes UINT16, // 2 bytes UINT32, // 4 bytes + BOOL, // 1 byte (stored as 1/0 in uint64_t slot) DATA_TYPE_NUM, }; @@ -68,6 +69,7 @@ inline uint64_t get_element_size(DataType dtype) { 8, // DataType::UINT64 2, // DataType::UINT16 4, // DataType::UINT32 + 1, // DataType::BOOL }; return data_type_size[static_cast(dtype)]; } @@ -102,6 +104,8 @@ inline const char *get_dtype_name(DataType dtype) { return "UINT16"; case DataType::UINT32: return "UINT32"; + case DataType::BOOL: + return "BOOL"; default: return "UNKNOWN"; } @@ -136,7 +140,42 @@ inline const char *get_dtype_name(DataType dtype) { // // Named convenience functions (from_u64_f32 etc.) are removed — use the // template form from_u64() / to_u64() directly. -// ----------------------------------------------------------------------------- + +/** + * Map a C++ type to its DataType enum value. + * + * Used to preserve original scalar type information through the + * Arg -> PTO2TaskPayload -> tensor dump pipeline. + */ +template +inline constexpr uint8_t dtype_of() { + static_assert(is_supported_scalar_arg_v, "dtype_of: type must be arithmetic or enum"); + if constexpr (std::is_same_v) { + return static_cast(DataType::FLOAT32); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::FLOAT32); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::INT32); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::UINT32); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::INT64); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::UINT64); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::INT16); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::UINT16); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::INT8); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::UINT8); + } else if constexpr (std::is_same_v) { + return static_cast(DataType::BOOL); + } else { + return static_cast(DataType::UINT64); + } +} /** * Pack a value into uint64_t storage (zero-extends smaller types).