diff --git a/docs/transformers/llm.md b/docs/transformers/llm.md index 922b2901cc..fd0a211b19 100644 --- a/docs/transformers/llm.md +++ b/docs/transformers/llm.md @@ -179,6 +179,39 @@ python llmexport.py \ cmake .. -DMNN_BUILD_CONVERTER=ON && make -j16 ``` 编译完成后 `build/` 目录下会生成 `MNNConvert` 可执行文件，`llmexport.py` 默认会在 `../../../build/` 下查找该工具；也可以通过 `--mnnconvert` 选项显式指定 MNNConvert 路径。若未提供本地 MNNConvert，脚本会回退到 pymnn（需先安装 `pip install MNN`）。此方案目前支持导出4bit和8bit模型。 +- 导出 segment 形式的 MNN LLM，使用 `--export mnn --segment`。该模式从 safetensors 权重和 workflow JSON 直接生成多个 MNN 子图，跳过 ONNX 中间文件，适合在 Metal 等后端上复用 decoder、logit、embedding 等 segment 模型。默认会在 `resource/*.json` 中查找匹配的 workflow，也可以通过 `--workflow /path/to/workflow.json` 显式指定。 + + ```bash + cd transformers/llm/export + python3 llmexport.py \ + --path /path/to/Qwen3-0.6B \ + --export mnn \ + --segment \ + --dst_path ./model + ``` + + segment 导出目录包含： + + ```text + model/ + ├── config.json # llm_demo 入口配置，包含 "mnn_llm_version": "segment" + ├── llm_config.json # 模型结构和模板配置 + ├── tokenizer.mtok + ├── embed.mnn + ├── decoder.mnn + ├── decoder.mnn.weight + ├── logit.mnn + ├── logit.mnn.weight + └── logit_topkv_1.mnn + ``` + + 运行 segment 模型时需要使用生成的 `config.json`： + + ```bash + ./llm_demo transformers/llm/export/model/config.json /path/to/prompt.txt + ``` + + C++ 运行时需启用 `MNN_BUILD_LLM=ON`，并打开 `MNN_LLM_SUPPORT_SEGMENT`（默认开启）。segment 路径当前仅支持 `--export mnn`，不支持 `--export onnx`。 - 如果直接转为mnn模型遇到问题，或者需要其他bits数的量化（如5bit/6bit），可以先将模型先转为onnx模型，使用`--export onnx`，然后使用./MNNConvert工具将onnx模型转为mnn模型: ``` @@ -197,7 +230,7 @@ usage: llmexport.py [-h] --path PATH [--type TYPE] [--tokenizer_path TOKENIZER_P [--gptq_path GPTQ_PATH] [--dst_path DST_PATH] [--verbose] [--test TEST] [--export EXPORT] [--onnx_slim] [--quant_bit QUANT_BIT] [--quant_block QUANT_BLOCK] [--lm_quant_bit LM_QUANT_BIT] [--mnnconvert MNNCONVERT] [--ppl] [--awq] [--omni] [--sym] [--seperate_embed] - [--lora_split] + [--lora_split] [--segment] [--workflow WORKFLOW] llm_exporter @@ -219,6 +252,8 @@ optional arguments: --verbose Whether or not to print verbose. --test TEST test model inference with query `TEST`. --export EXPORT export model to an onnx/mnn model. + --segment export segment MNN LLM from safetensors workflow directly, without ONNX export. + --workflow WORKFLOW workflow json for --segment safetensors conversion. If absent, search resource/*.json. --onnx_slim Whether or not to use onnx-slim. --quant_bit QUANT_BIT mnn quant bit, 4 or 8, default is 4. @@ -1158,4 +1193,4 @@ adb push model /data/local/tmp/MNN/model ``` cd ${MNN_ROOT} project/android/testCommon.sh ./llm_demo model/config_mlda.json -``` \ No newline at end of file +``` diff --git a/resource/qwen3_hf_0.6b.json b/resource/qwen3_hf_0.6b.json new file mode 100644 index 0000000000..33c5f1c3eb --- /dev/null +++ b/resource/qwen3_hf_0.6b.json @@ -0,0 +1,34 @@ +{ + "models": [ + { + "name": "hf_decoder", + "blocks": [ + { + "type": "QwenTransformer", + "hiddenSize": 1024, + "headDim": 128, + "numHead": 16, + "kvNumHead": 8, + "number": 28, + "max_position_embeddings": 40960 + } + ] + }, + { + "name": "logit", + "blocks": [ + { + "type": "InnerProduct", + "prefix": "lm_head" + }, + { + "type": "TieEmbedding" + }, + { + "type": "TopKV", + "K": [1, 5] + } + ] + } + ] +} diff --git a/skills/support-new-llm/SKILL.md b/skills/support-new-llm/SKILL.md index 0ea0629217..5620553093 100644 --- a/skills/support-new-llm/SKILL.md +++ b/skills/support-new-llm/SKILL.md @@ -1,11 +1,11 @@ --- name: support-new-llm -description: 为 MNN 框架添加新的 LLM 模型支持。支持从 HuggingFace/ModelScope 下载模型，分析架构，添加映射，Hook 对齐测试，导出 MNN 模型。采用 TDD 模式，分 6 步执行，每步有独立测试标准。 +description: 为 MNN 框架添加新的 LLM 模型支持。支持从 HuggingFace/ModelScope 下载模型，分析架构，添加映射，Hook 对齐测试，导出 MNN 模型；当用户明确要求 safetensors/segment/workflow/MNNConvert -f ST 时，走 safetensors segment 补充分支。采用 TDD 模式，分 6 步执行，每步有独立测试标准。 --- # MNN LLM 新模型支持 SKILL -> **触发条件**：当用户请求支持/添加/适配一个新的 LLM 模型时触发。常见表述包括："支持xxx模型"、"添加xxx模型支持"、"适配xxx"、"导出xxx模型"等。 +> **触发条件**：当用户请求支持/添加/适配一个新的 LLM 模型时触发。常见表述包括："支持xxx模型"、"添加xxx模型支持"、"适配xxx"、"导出xxx模型"等。若用户明确提到 `safetensors`、`--segment`、`workflow.json`、`MNNConvert -f ST` 或“绕过 ONNX 直接转换”，先读 `safetensors-segment.md`。 ## 概述 @@ -20,6 +20,8 @@ MNN 的模型导出本质上是**对照 HuggingFace transformers 库中原始模 3. 用 Python `--test` 验证映射正确性 4. 导出 MNN 模型并用 C++ 引擎验证 +默认导出链路是 `llmexport.py --export mnn`。如果目标是 safetensors segment 格式，则使用 `llmexport.py --export mnn --segment`，按 `safetensors-segment.md` 先校验 workflow、safetensors key 和 builder 约定。 + ### 注意事项 > **🚨 严禁将输出错误归因于"量化精度不够"**：4bit 量化的 0.5B 小模型都能正确输出。如果 C++ 输出完全不对（如图片识别不出、输出乱码），**一定是实现细节没有与 HF 对齐**，必须逐步 dump 数据对比定位，不要靠猜。 @@ -44,6 +46,9 @@ MNN 的模型导出本质上是**对照 HuggingFace transformers 库中原始模 | `transformers/llm/export/utils/audio.py` | Audio Encoder 实现 | 音频模型 | | `transformers/llm/export/utils/custom_op.py` | 自定义算子导出 | 新算子时 | | `transformers/llm/export/llmexport.py` | 导出主流程入口 | 偶尔 | +| `transformers/llm/export/segment.py` | safetensors segment 导出入口 | segment 分支 | +| `resource/*.json` | safetensors workflow 模板 | segment 分支 | +| `tools/converter/source/safetensors/*.cpp` | safetensors converter / builder 实现 | segment 分支 | --- @@ -98,6 +103,7 @@ MNN 的模型导出本质上是**对照 HuggingFace transformers 库中原始模 | Tier 4 (音频模型) | 1 → 2 → 3 → 5 → 4 | 需要 audio.py | | Tier 5 (视觉模型) | 1 → 2 → 3 → 5 → 4 | 需要 vision.py | | Tier 6 (全新架构) | 1 → 2 → 6 → 3 → 4 | 需要新算子（如叠加 Tier 4/5 则加入 step5） | +| Safetensors segment | 1 → S1/S2/S3 → S4/S5 | 明确要求 `--segment` / workflow / `MNNConvert -f ST` 时，参见 `safetensors-segment.md` | --- @@ -185,8 +191,10 @@ modeling_*.py 中是否有全新的 Attention 类型（非标准 SDPA）? **在开始之前，建议先浏览 `common-pitfalls.md`**，了解已知的常见问题和解决方案（RoPE 变体、dtype 级联、Jinja 限制、stop token、残差模式、MoE 支持要点、FakeLinear axis 陷阱、**do_map 静默失败与 rope_theta 间接存储**、非标准模型加载等）。 +**Safetensors segment 分支的常见问题**：workflow 超参与权重 shape 不匹配、builder 预期 key 前缀不存在、自动 workflow 匹配选错、segment runtime 没启用。详见 `safetensors-segment.md`。 + --- ## 开始执行 -**现在请打开 `skills/support-new-llm/step1-analyze.md`，开始步骤 1。** +**现在请打开 `skills/support-new-llm/step1-analyze.md`，开始步骤 1。若用户明确要求 safetensors segment 导出，同时打开 `skills/support-new-llm/safetensors-segment.md`。** diff --git a/skills/support-new-llm/safetensors-segment.md b/skills/support-new-llm/safetensors-segment.md new file mode 100644 index 0000000000..cead183cf3 --- /dev/null +++ b/skills/support-new-llm/safetensors-segment.md @@ -0,0 +1,280 @@ +# Safetensors Segment 导出补充 + +> **适用场景**：用户明确提到 `safetensors`、`--segment`、`workflow.json`、`MNNConvert -f ST`，或要求绕过 ONNX、直接从 safetensors 生成 segment 格式 MNN LLM。 + +本补充从 `skills/segment-new-llm` 的 safetensors 流程提取而来，并按当前 MNN 仓库路径修正。默认 `support-new-llm` 流程仍是 `llmexport.py --export mnn` 的标准导出；只有命中上述场景时才切到本分支。 + +不要照搬其他仓库/旧 skill 中的 PantherLLM 路径（如 `converter/resource/*.json`、`converter/mnn_safetensors_plugin`、`--customOpLibs libpantherllm_safetensors_plugin`）。当前 MNN 仓库的 segment 分支以 `transformers/llm/export/segment.py`、`resource/*.json` 和 `tools/converter/source/safetensors` 为准。 + +--- + +## 入口与核心文件 + +Segment 分支的主路径是： + +```text +HF / ModelScope model dir + | + v +safetensors weights + workflow JSON + | + v +llmexport.py --export mnn --segment + | + v +MNNConvert -f ST + | + v +segment model dir + | + v +llm_demo /config.json prompt.txt +``` + +核心文件： + +| 文件路径 | 作用 | +|---------|------| +| `transformers/llm/export/segment.py` | segment 导出入口；解析 workflow、safetensors 和导出配置 | +| `transformers/llm/export/llmexport.py` | `--segment` / `--workflow` 参数入口 | +| `resource/*.json` | workflow 模板；当前典型样例是 `resource/qwen3_hf_0.6b.json` | +| `tools/converter/source/safetensors/*.cpp` | safetensors builder / converter 实现 | +| `tools/converter/source/safetensors/SafetensorModelRegistry.hpp` | `REGISTER_SAFETENSOR_MODEL_BUILDER` 注册机制 | +| `transformers/llm/engine/src/segment.cpp` | C++ runtime 的 segment 加载路径 | + +--- + +## 步骤 S1：确认输入形态 + +先判断用户给的是哪种输入： + +| 输入形态 | 处理方式 | +|---------|---------| +| 模型目录，包含 `*.safetensors` 和 `config.json` | 可直接作为 `--path` | +| 单个 `.safetensors` 文件 | 可直接作为 `--path`，但仍需要 tokenizer/config 来源 | +| sharded safetensors + `*.safetensors.index.json` | `segment.py` 会按 index 中的 `weight_map` 顺序传给 `MNNConvert` | +| 只有 PyTorch `state_dict` / `.bin` | 先转换为 safetensors，再进入本流程 | +| 已有 workflow JSON | 显式传 `--workflow /path/to/workflow.json` | +| 没有 workflow JSON | 先从 `resource/*.json` 找最接近模板；不要盲目依赖自动匹配 | + +必须记录： + +- `model_type` +- `hidden_size` +- `num_hidden_layers` +- `num_attention_heads` +- `num_key_value_heads` +- `head_dim` +- `max_position_embeddings` +- embedding / blocks / norm / lm_head 的实际 safetensors key +- tokenizer 文件是否完整 + +--- + +## 步骤 S2：选择 workflow 与 builder + +Segment 分支不是在 `model_mapper.py` 中添加 Python 映射，而是靠 **workflow + safetensors builder**。 + +Workflow 关键点： + +- 顶层 `models[].name` 决定调用哪个 builder。 +- `blocks[]` 描述结构和超参，例如 `hiddenSize`、`headDim`、`numHead`、`kvNumHead`、`number`。 +- 当前可优先参考 `resource/qwen3_hf_0.6b.json`。 + +Builder 关键点： + +- 注册点使用 `REGISTER_SAFETENSOR_MODEL_BUILDER("name", builderFunc)`。 +- 当前文本 decoder 典型实现是 `tools/converter/source/safetensors/HuggingFaceQwen3.cpp`。 +- `logit` 典型实现是 `tools/converter/source/safetensors/Logit.cpp`。 + +判断是否能复用 workflow： + +- 权重命名与现有 builder 预期一致。 +- block 类型一致。 +- 只需要修改层数、hidden、head、kv head、head dim、max position。 +- 输出仍是 segment runtime 需要的 `embed.mnn`、`decoder.mnn`、`logit.mnn`、`logit_topkv_*.mnn` 等文件。 + +需要新增或修改 builder 的信号： + +- 权重前缀不同，现有 builder 找不到关键 tensor。 +- Attention / MLP / norm / residual 结构不同。 +- 需要新增 workflow block 字段才能表达模型结构。 +- 输出文件结构不能被现有 segment runtime 加载。 + +--- + +## 步骤 S3：转换前静态校验 + +在执行 `MNNConvert -f ST` 前，先验证 key、shape 和 workflow 超参。 + +### safetensors key 检查 + +```python +from safetensors import safe_open + +st_path = "/path/to/model.safetensors" +required_keys = [ + "model.embed_tokens.weight", + "model.layers.0.self_attn.q_proj.weight", + "model.layers.0.self_attn.k_proj.weight", + "model.layers.0.self_attn.v_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + "model.layers.0.mlp.gate_proj.weight", + "model.layers.0.mlp.up_proj.weight", + "model.layers.0.mlp.down_proj.weight", + "model.norm.weight", + "lm_head.weight", +] + +with safe_open(st_path, framework="pt", device="cpu") as f: + key_set = set(f.keys()) + print("tensor_count:", len(key_set)) + for key in required_keys: + if key in key_set: + print("OK ", key, f.get_tensor(key).shape) + else: + print("MISS", key) + +missing = [key for key in required_keys if key not in key_set] +if missing: + raise SystemExit(f"missing keys: {missing}") +``` + +### workflow 超参与权重 shape 检查 + +```python +import json +from safetensors import safe_open + +workflow_path = "/path/to/workflow.json" +st_path = "/path/to/model.safetensors" + +with open(workflow_path, "r", encoding="utf-8") as f: + workflow = json.load(f) + +with safe_open(st_path, framework="pt", device="cpu") as st: + q = st.get_tensor("model.layers.0.self_attn.q_proj.weight") + +for model in workflow.get("models", []): + print("model:", model.get("name")) + for block in model.get("blocks", []): + if block.get("type") in {"QwenTransformer", "GPT2Transformer"}: + hidden = block.get("hiddenSize") + head_dim = block.get("headDim") + num_head = block.get("numHead") + if hidden is not None: + assert q.shape[1] == hidden, (q.shape, hidden) + if head_dim is not None and num_head is not None: + assert head_dim * num_head == q.shape[0], (head_dim, num_head, q.shape) + +print("workflow contract looks OK") +``` + +通过标准： + +- builder 依赖的关键 key 全部存在。 +- workflow 中的层数、hidden、head 维度与权重 shape 对齐。 +- tokenizer/config 资源来源明确。 +- 已保存参考模型的输入 prompt、token ids 和输出，用于导出后对比。 + +--- + +## 步骤 S4：执行 segment 导出 + +### 构建要求 + +```bash +mkdir -p build +cd build +cmake .. -DMNN_BUILD_LLM=ON -DMNN_BUILD_CONVERTER=ON +make -j$(nproc) +``` + +`MNN_LLM_SUPPORT_SEGMENT` 默认开启；如果被关闭，segment runtime 不能加载 `"mnn_llm_version": "segment"` 的模型。 + +### 推荐命令：通过 llmexport.py + +```bash +cd transformers/llm/export +python3 llmexport.py \ + --path /path/to/model_dir_or_safetensors \ + --export mnn \ + --segment \ + --workflow /path/to/workflow.json \ + --dst_path ./MODEL \ + --quant_bit 4 \ + --quant_block 64 +``` + +如果省略 `--workflow`，`segment.py` 会在 `resource/` 和 `transformers/llm/resource/` 下搜索可匹配的 JSON。命中多个或找不到时，应显式传入 workflow。 + +### 调试命令：直接调用 MNNConvert + +```bash +build/MNNConvert \ + -f ST \ + -i /path/to/workflow.json \ + -i /path/to/model.safetensors \ + -o /path/to/out_dir \ + --allowCustomOp \ + --saveExternalData \ + --weightQuantBits 4 \ + --weightQuantBlock 64 +``` + +多 shard safetensors 时，对每个 shard 追加一个 `-i /path/to/shard.safetensors`，顺序应与 index 中的 `weight_map` 一致。 + +--- + +## 步骤 S5：检查产物并验证 + +典型输出： + +```text +MODEL/ +├── config.json # 包含 "mnn_llm_version": "segment" +├── llm_config.json +├── tokenizer.mtok +├── embed.mnn +├── decoder.mnn +├── decoder.mnn.weight +├── logit.mnn +├── logit.mnn.weight +└── logit_topkv_1.mnn +``` + +检查： + +```bash +ls -la /path/to/MODEL +cat /path/to/MODEL/config.json +``` + +运行： + +```bash +echo "你好" > /tmp/prompt.txt +build/llm_demo /path/to/MODEL/config.json /tmp/prompt.txt +``` + +通过标准： + +- `config.json` 存在且包含 `"mnn_llm_version": "segment"`。 +- `embed.mnn`、`decoder.mnn`、`logit.mnn` 等关键文件存在且大小 > 0。 +- `llm_demo` 能加载并生成合理文本。 +- 输出与步骤 S3 保留的参考输出方向一致。 + +--- + +## 常见失败 + +| 现象 | 排查顺序 | +|------|---------| +| `no suitable workflow json` | 显式传 `--workflow`；检查 workflow 超参是否与 config 匹配 | +| `multiple suitable workflow json files` | 显式传 `--workflow`，不要让自动匹配猜 | +| `missing tensor` | 回到步骤 S3，核对 safetensors key 和 builder 硬编码前缀 | +| `unknown builder` | 检查 `models[].name` 是否已在 `tools/converter/source/safetensors` 注册 | +| 转换成功但加载失败 | 检查 `config.json`、`llm_config.json`、`tokenizer.mtok` 和输出文件名 | +| 输出完全不对 | 先查 workflow 超参、权重前缀、builder 读权重/reshape/transpose，再考虑量化 | + +不要在没有 key/shape 证据的情况下把问题归因于量化精度。 diff --git a/skills/support-new-llm/step1-analyze.md b/skills/support-new-llm/step1-analyze.md index cf041ebd38..6e6f0dc581 100644 --- a/skills/support-new-llm/step1-analyze.md +++ b/skills/support-new-llm/step1-analyze.md @@ -11,6 +11,8 @@ 根据用户提供的输入，选择对应的方式： +> **Safetensors segment 分支**：如果用户明确要求 `safetensors`、`--segment`、`workflow.json` 或 `MNNConvert -f ST`，本步骤仍需下载/确认模型目录和参考推理，但后续映射与导出要按 `safetensors-segment.md` 执行，而不是默认 ONNX 导出路径。 + ### 情况 A：用户提供本地路径 ``` @@ -72,6 +74,17 @@ snapshot_download( - `*.safetensors` 或 `pytorch_model*.bin`（模型权重） - `tokenizer.json` 或 `tokenizer.model`（tokenizer 文件） +### Safetensors segment 输入补充 + +命中 segment 分支时，还需要记录： + +- safetensors 是单文件、`model.safetensors`，还是 sharded safetensors + `*.safetensors.index.json` +- 是否已有 workflow JSON；没有则先从 `resource/*.json` 中找最接近模板 +- embedding / blocks / norm / lm_head 的实际 safetensors key +- 是否需要显式传 `--workflow`，避免自动匹配选错 + +具体 key/shape 校验脚本见 `safetensors-segment.md`。 + --- ## 1.2 阅读模型 README 和 config.json diff --git a/skills/support-new-llm/step2-mapping.md b/skills/support-new-llm/step2-mapping.md index f29c9b74ca..02ea04c8f3 100644 --- a/skills/support-new-llm/step2-mapping.md +++ b/skills/support-new-llm/step2-mapping.md @@ -10,6 +10,8 @@ MNN 使用 4 层映射将 HuggingFace 模型结构转换为统一接口： +> **Safetensors segment 分支**：如果本次目标是 `--segment` 或 `MNNConvert -f ST`，不要把主要工作放在 `model_mapper.py`。segment 分支的映射单位是 `resource/*.json` workflow 和 `tools/converter/source/safetensors` builder，流程见 `safetensors-segment.md` 的步骤 S2。 + | 映射键 | 作用 | 说明 | |--------|------|------| | `config` | HF config.json 字段 → LlmConfig 属性 | 把模型配置正确读入 | diff --git a/skills/support-new-llm/step3-test-python.md b/skills/support-new-llm/step3-test-python.md index f70d12b548..7bc8b13f86 100644 --- a/skills/support-new-llm/step3-test-python.md +++ b/skills/support-new-llm/step3-test-python.md @@ -10,6 +10,8 @@ 仅看最终输出文本是否"合理"是不够的。本步骤通过 **hook 机制**在两个模型的关键位置截取中间结果，逐层对比，精确定位映射或实现中的错误。 +> **Safetensors segment 分支**：如果目标是 `--segment` / `MNNConvert -f ST`，本步骤需要先做 `safetensors-segment.md` 中的 S3 静态校验（key、shape、workflow 超参、tokenizer/config 资源）。只有当 segment 分支仍修改了 Python LlmModel 或 transformers 逻辑时，才继续执行本文的 hook 对齐。 + **对比的两套模型**： 1. **原始 transformers 模型**：步骤 1 中加载的 `AutoModelForCausalLM`（标准答案） 2. **MNN LlmModel**：步骤 2 中映射转换后的 `LlmModel`（需要验证） diff --git a/skills/support-new-llm/step4-export.md b/skills/support-new-llm/step4-export.md index 3ef1a56ca9..6246944d68 100644 --- a/skills/support-new-llm/step4-export.md +++ b/skills/support-new-llm/step4-export.md @@ -8,6 +8,8 @@ ## 4.1 导出 MNN 模型 +> **Safetensors segment 分支**：如果目标是 `--segment` / `MNNConvert -f ST`，不要使用本节默认 ONNX 导出路径，改用 `safetensors-segment.md` 的 S4/S5：`llmexport.py --export mnn --segment --workflow ...`，并用 `llm_demo /config.json prompt.txt` 验证 segment runtime。 + ```bash cd transformers/llm/export python3 llmexport.py \ diff --git a/skills/test-ci/SKILL.md b/skills/test-ci/SKILL.md index f01873fea7..411cfa8004 100644 --- a/skills/test-ci/SKILL.md +++ b/skills/test-ci/SKILL.md @@ -65,6 +65,9 @@ Valid filters: `all` (default) · `cpu` · `opencl` · `opencl-image` · * Combined stdout/stderr for every stage is saved under `logs/test_ci-/.log` — read the named log of a failing stage for the trailing output. `rc=137` ≈ OOM-kill, `rc=139` ≈ SIGSEGV. +* For GPU/OpenCL smoke tests, verify that the intended backend actually loaded + (for example, OpenCL tuning/backend logs are present). A correct model output + alone is not sufficient when CPU fallback is possible. ## Environment variables @@ -118,6 +121,10 @@ file explains every field and every `skip` entry's rationale. [`TESTING.md`](../../TESTING.md) § "How to add a new operator test". 2. If its name prefix matches an existing stage (e.g. `op/*`), it is picked up automatically — no JSON change needed. Otherwise add a dedicated stage. +3. Do not add backend-specific skips inside an operator test. If a configured + backend fails, fix the backend implementation or, for a confirmed driver + issue, put the exact test name in the stage `skip` list with a documented + rationale in `test_stages.json`. For deeper work on operators themselves, see the [`add-new-op`](../add-new-op/SKILL.md) skill. diff --git a/tools/converter/CMakeLists.txt b/tools/converter/CMakeLists.txt index 1e757c9f82..1c0d9dd4a4 100644 --- a/tools/converter/CMakeLists.txt +++ b/tools/converter/CMakeLists.txt @@ -20,10 +20,12 @@ IF(MNN_BUILD_CONVERTER) SET(MNN_CONVERTER_BACKENDS_OBJECTS "") include_directories(${CMAKE_CURRENT_LIST_DIR}/include) include_directories(${CMAKE_CURRENT_LIST_DIR}/source/tflite/schema) + include_directories(${CMAKE_CURRENT_LIST_DIR}/source/safetensors) include_directories(${CMAKE_CURRENT_BINARY_DIR}) include(${CMAKE_CURRENT_LIST_DIR}/source/compression/CMakeLists.txt) include(${CMAKE_CURRENT_LIST_DIR}/source/tensorflow/CMakeLists.txt) include(${CMAKE_CURRENT_LIST_DIR}/source/onnx/CMakeLists.txt) + include(${CMAKE_CURRENT_LIST_DIR}/source/safetensors/CMakeLists.txt) include(${CMAKE_CURRENT_LIST_DIR}/source/caffe/CMakeLists.txt) include(${CMAKE_CURRENT_LIST_DIR}/source/MNN/CMakeLists.txt) include(${CMAKE_CURRENT_LIST_DIR}/source/optimizer/CMakeLists.txt) diff --git a/tools/converter/include/config.hpp b/tools/converter/include/config.hpp index 5f2c9931d6..f165a9b211 100644 --- a/tools/converter/include/config.hpp +++ b/tools/converter/include/config.hpp @@ -11,6 +11,7 @@ #include #include #include +#include struct PostTreatContext; class MNN_PUBLIC modelConfig { public: @@ -23,13 +24,14 @@ class MNN_PUBLIC modelConfig { saveHalfFloat(false){ } ~ modelConfig (); - enum MODEL_SOURCE { TENSORFLOW = 0, CAFFE, ONNX, MNN, TFLITE, TORCH, JSON, MAX_SOURCE }; + enum MODEL_SOURCE { TENSORFLOW = 0, CAFFE, ONNX, MNN, TFLITE, TORCH, JSON, SAFETENSORS, MAX_SOURCE }; // MNN model path std::string MNNModel; // if model is tensorflow, this value is NULL; std::string prototxtFile; // tensorflow pb, or caffe model + std::vector modelFiles; std::string modelFile; // bizCode std::string bizCode; diff --git a/tools/converter/source/common/RemoveParams.cpp b/tools/converter/source/common/RemoveParams.cpp index 0344cff8d1..badd328297 100644 --- a/tools/converter/source/common/RemoveParams.cpp +++ b/tools/converter/source/common/RemoveParams.cpp @@ -70,6 +70,8 @@ void RemoveAndStoreParam(std::unique_ptr& op, std::ofstream* fs, int64 } break; } +// It will some case cause error +#ifdef MNN_SUPPORT_CONST_EXTERNAL case MNN::OpParameter_Blob: { auto param = op->main.AsBlob(); size_t totalSize = 1; @@ -100,6 +102,7 @@ void RemoveAndStoreParam(std::unique_ptr& op, std::ofstream* fs, int64 } break; } +#endif default: break; } diff --git a/tools/converter/source/common/cli.cpp b/tools/converter/source/common/cli.cpp index d2dd9ab4d4..e1639ce2cd 100644 --- a/tools/converter/source/common/cli.cpp +++ b/tools/converter/source/common/cli.cpp @@ -39,6 +39,7 @@ #include #include #include "core/MemoryFormater.h" +#include "../safetensors/SafetensorConverter.hpp" modelConfig::~modelConfig() { if (nullptr != compressInfo) { delete compressInfo; @@ -148,12 +149,12 @@ bool Cli::initializeMNNConvertArgs(modelConfig &modelPath, int argc, char **argv "Convert Other Model Format To MNN Model\n")( std::make_pair("v", "version"), "show current version")(std::make_pair("f", "framework"), #ifdef MNN_BUILD_TORCH - "model type, ex: [TF,CAFFE,ONNX,TFLITE,MNN,TORCH,JSON]", + "model type, ex: [TF,CAFFE,ONNX,TFLITE,MNN,TORCH,JSON,ST]", #else - "model type, ex: [TF,CAFFE,ONNX,TFLITE,MNN,JSON]", + "model type, ex: [TF,CAFFE,ONNX,TFLITE,MNN,JSON,ST]", #endif cxxopts::value())( - "modelFile", "tensorflow Pb or caffeModel, ex: *.pb,*caffemodel", cxxopts::value())( + std::make_pair("i", "modelFile"), "tensorflow Pb or caffeModel, ex: *.pb,*caffemodel", cxxopts::value>())( "batch", "if model input's batch is not set, set as the batch size you set", cxxopts::value())( "keepInputFormat", "keep input dimension format or not, default: true", cxxopts::value())( "optimizeLevel", @@ -161,7 +162,7 @@ bool Cli::initializeMNNConvertArgs(modelConfig &modelPath, int argc, char **argv "every input case is right, 2: normally right but some case may be wrong, default 1", cxxopts::value())("optimizePrefer", "graph optimize option, 0 for normal, 1 for smalleset, 2 for fastest", cxxopts::value())("prototxt", "only used for caffe, ex: *.prototxt", - cxxopts::value())("MNNModel", "MNN model, ex: *.mnn", + cxxopts::value())(std::make_pair("o", "MNNModel"), "MNN model, ex: *.mnn", cxxopts::value())( "fp16", "save Conv's weight/bias in half_float data type")( "benchmarkModel", @@ -259,6 +260,8 @@ bool Cli::initializeMNNConvertArgs(modelConfig &modelPath, int argc, char **argv #endif } else if ("JSON" == frameWork) { modelPath.model = modelConfig::JSON; + } else if ("ST" == frameWork) { + modelPath.model = modelConfig::SAFETENSORS; } else { std::cout << "Framework Input ERROR or Not Support This Model Type Now!" << std::endl; return false; @@ -283,7 +286,13 @@ bool Cli::initializeMNNConvertArgs(modelConfig &modelPath, int argc, char **argv // model file path if (result.count("modelFile")) { - const std::string modelFile = result["modelFile"].as(); + auto files = result["modelFile"].as>(); + modelPath.modelFiles = files; + if (files.empty()) { + DLOG(INFO) << "modelFile Not set Invalid, use --modelFile to set!"; + return false; + } + const std::string modelFile = files[0]; if (CommonKit::FileIsExist(modelFile)) { modelPath.modelFile = modelFile; } else { @@ -562,6 +571,24 @@ bool Cli::convertModel(modelConfig& modelPath) { dumpModelInfo(modelPath.modelFile.c_str()); return true; } + if (modelPath.model == modelConfig::SAFETENSORS) { + std::cout << "Create Converter with config: " << modelPath.modelFiles[0] << std::endl; + MNN::SafeTensors::Converter converter(modelPath.modelFiles[0]); + auto models = converter.listModels(); + for (int i = 1; i < modelPath.modelFiles.size(); ++i) { + std::cout << "Load Safetensors " << modelPath.modelFiles[i] << std::endl; + converter.loadSafeTensors(modelPath.modelFiles[i]); + } + for (auto& name : models) { + auto newConfig = modelPath; + newConfig.MNNModel = modelPath.MNNModel + "/"; + if (!converter.convert(name, newConfig)) { + std::cout << "Convert " << name << " Failed" << std::endl; + return false; + } + } + return true; + } std::cout << "Start to Convert Other Model Format To MNN Model..., target version: " << modelPath.targetVersion << std::endl; std::unique_ptr netT = std::unique_ptr(new MNN::NetT()); int parseRes = 1; diff --git a/tools/converter/source/optimizer/postconvert/AddTensorFormatConverter.cpp b/tools/converter/source/optimizer/postconvert/AddTensorFormatConverter.cpp index 1e7b98afcd..cf2b8f950c 100644 --- a/tools/converter/source/optimizer/postconvert/AddTensorFormatConverter.cpp +++ b/tools/converter/source/optimizer/postconvert/AddTensorFormatConverter.cpp @@ -20,6 +20,7 @@ static void _setInputFormat(std::vector& tensorFormat, int inde enum FormatSetType { NC4HW4_SINGLE, // only first input / output is nc4hw4 NC4HW4_FULL, // all nc4hw4 + NC4HW4_OUTPUT, // output nc4hw4 COMPABILIT_SINGLE, // only first input / output is compability COMPABILIT_FULL, // all format should be same ORIGIN @@ -38,6 +39,18 @@ static FormatSetType _getFormatType(const OpT* op, MNN_DATA_FORMAT originFormat) case MNN::OpType_PReLU: case MNN::OpType_Dilation2D: return NC4HW4_SINGLE; + case MNN::OpType_Attention: + if (op->main.AsAttentionParam()->output_c4) { + return NC4HW4_OUTPUT; + } else { + return ORIGIN; + } + case MNN::OpType_LayerNorm: + if (op->defaultDimentionFormat == MNN_DATA_FORMAT_NC4HW4) { + return NC4HW4_FULL; + } else { + return ORIGIN; + } case MNN::OpType_ConvInt8: case MNN::OpType_Pooling: case MNN::OpType_Pooling3D: @@ -136,6 +149,8 @@ static MNN_DATA_FORMAT _getRequireFormat(FormatSetType type, int inputIndex, MNN return originFormat; case NC4HW4_FULL: return MNN_DATA_FORMAT_NC4HW4; + case NC4HW4_OUTPUT: + return originFormat; case NC4HW4_SINGLE: if (inputIndex == 0) { return MNN_DATA_FORMAT_NC4HW4; @@ -214,6 +229,16 @@ static bool _computeTensorFormat(std::vector& tensorFormat, std } return true; } + case NC4HW4_OUTPUT: + { + for (int i=0; iinputIndexes.size(); ++i) { + _setInputFormat(tensorFormat, op->inputIndexes[i], originFormat); + } + for (int i=0; ioutputIndexes.size(); ++i) { + tensorFormat[op->outputIndexes[i]] = MNN_DATA_FORMAT_NC4HW4; + } + return true; + } case COMPABILIT_SINGLE: { for (int i=1; iinputIndexes.size(); ++i) { diff --git a/tools/converter/source/safetensors/CMakeLists.txt b/tools/converter/source/safetensors/CMakeLists.txt new file mode 100644 index 0000000000..45678c3711 --- /dev/null +++ b/tools/converter/source/safetensors/CMakeLists.txt @@ -0,0 +1,12 @@ +file(GLOB SafeTensors_SRC CONFIGURE_DEPENDS + ${CMAKE_CURRENT_LIST_DIR}/*.cpp + ${CMAKE_CURRENT_LIST_DIR}/*.c + ${CMAKE_CURRENT_LIST_DIR}/*.cc + ${CMAKE_CURRENT_LIST_DIR}/*.h + ${CMAKE_CURRENT_LIST_DIR}/*.hpp +) + +add_library(MNNConverterSafeTensors OBJECT ${SafeTensors_SRC}) + +list(APPEND MNN_CONVERTER_BACKENDS_OBJECTS $) +list(APPEND MNN_CONVERTER_BACKENDS_TARGETS MNNConverterSafeTensors) diff --git a/tools/converter/source/safetensors/HuggingFaceQwen3.cpp b/tools/converter/source/safetensors/HuggingFaceQwen3.cpp new file mode 100644 index 0000000000..5686566c51 --- /dev/null +++ b/tools/converter/source/safetensors/HuggingFaceQwen3.cpp @@ -0,0 +1,401 @@ +#include +#include + +#include +#include "MNN_generated.h" + +#include "../optimizer/Global.hpp" +#include "SafetensorConverter.hpp" +#include "SafetensorModelRegistry.hpp" +#include "SafetensorUtils.hpp" +#include "WorkflowJson.hpp" +#include "HuggingFaceQwen3.hpp" + +using namespace MNN::Express; +using namespace MNN::Express::SafeTensorUtils; + +namespace MNN { +namespace SafeTensors { + +static VARP _linear2d(VARP x4d, VARP weightOI, VARP bias = nullptr) { + auto wInfo = weightOI->getInfo(); + if (nullptr == wInfo || wInfo->dim.size() < 2) { + return nullptr; + } + + const int outDim = wInfo->dim[0]; + const int inDim = wInfo->dim[1]; + if (inDim <= 0 || outDim <= 0) { + return nullptr; + } + + if (nullptr == weightOI->readMap()) { + weightOI = _Cast(weightOI); + weightOI.fix(VARP::CONSTANT); + } + + std::vector weightData(weightOI->getInfo()->size); + ::memcpy(weightData.data(), weightOI->readMap(), weightData.size() * sizeof(float)); + + std::vector biasData(outDim, 0.0f); + if (nullptr != bias) { + if (nullptr == bias->readMap()) { + bias = _Cast(bias); + bias.fix(VARP::CONSTANT); + } + ::memcpy(biasData.data(), bias->readMap(), outDim * sizeof(float)); + } + + return _Conv(std::move(weightData), std::move(biasData), x4d, {inDim, outDim}, {1, 1}); +} + +class HuggingFaceQwen3 { +public: + HuggingFaceQwen3(const Converter* converter) : mConverter(converter) {} + + struct BlockInfo { + int hiddenSize = 0; + int headDim = 0; + int numberHead = 0; + int ropeCutHeadDim = 0; + + VARP cosEven; + VARP cosOdd; + VARP sinEven; + VARP sinOdd; + + VARP shapeQKV; + }; + + std::pair makeBlock(VARP hiddenState, VARP add, const BlockInfo& info, VARP mask, int blockIndex) { + auto blockPrefix = std::string("model.layers.") + std::to_string(blockIndex) + "."; + auto attnPrefix = blockPrefix + "self_attn."; + auto mlpPrefix = blockPrefix + "mlp."; + + auto setName = [](const VARP& v, const std::string& name) { + if (nullptr != v.get()) { + v->setName(name); + } + }; + + auto load = [this](const std::string& name) { + if (mConverter->hasTensor(name)) { + return mConverter->loadTensor(name, false); + } + return MNN::Express::VARP(nullptr); + }; + + const int hiddenSize = info.hiddenSize; + const float ln_eps = 1.0e-6f; + const bool useC4Opt = true; + + // RMSNorm + QKV + auto ln1Weight = load(blockPrefix + "input_layernorm.weight"); + + VARP hiddenStateNorm; + if (nullptr != add.get()) { + auto res = _BinaryLayerNorm(hiddenState, add, {ln1Weight, nullptr, ln_eps, true, hiddenSize, true}); + hiddenStateNorm = res.second; + hiddenState = res.first; + } else { + hiddenStateNorm = _TransformerLayerNorm(hiddenState, {ln1Weight, nullptr, ln_eps, true, hiddenSize, true}); + } + setName(hiddenStateNorm, blockPrefix + "input_layernorm.out"); + auto hiddenStateNorm4d = hiddenStateNorm; + + auto qWeight = load(attnPrefix + "q_proj.weight"); + auto kWeight = load(attnPrefix + "k_proj.weight"); + auto vWeight = load(attnPrefix + "v_proj.weight"); + auto oWeight = load(attnPrefix + "o_proj.weight"); + + VARP qBias; + VARP kBias; + VARP vBias; + if (mConverter->hasTensor(attnPrefix + "q_proj.bias")) { + qBias = load(attnPrefix + "q_proj.bias"); + } + if (mConverter->hasTensor(attnPrefix + "k_proj.bias")) { + kBias = load(attnPrefix + "k_proj.bias"); + } + if (mConverter->hasTensor(attnPrefix + "v_proj.bias")) { + vBias = load(attnPrefix + "v_proj.bias"); + } + + auto qWeightInfo = qWeight->getInfo(); + auto kWeightInfo = kWeight->getInfo(); + auto vWeightInfo = vWeight->getInfo(); + if (nullptr == qWeightInfo || nullptr == kWeightInfo || nullptr == vWeightInfo) { + return {nullptr, nullptr}; + } + + const int queryHiddenSize = qWeightInfo->dim[0]; + int headDim = info.headDim; + int numHeads = info.numberHead > 0 ? info.numberHead : (queryHiddenSize / headDim); + if (headDim <= 0 || numHeads <= 0 || numHeads * headDim != queryHiddenSize) { + return {nullptr, nullptr}; + } + const int attnOutSize = headDim * numHeads; + + auto q = _linear2d(hiddenStateNorm4d, qWeight, qBias); + setName(q, attnPrefix + "q_proj.out"); + auto k = _linear2d(hiddenStateNorm4d, kWeight, kBias); + setName(k, attnPrefix + "k_proj.out"); + auto v = _linear2d(hiddenStateNorm4d, vWeight, vBias); + setName(v, attnPrefix + "v_proj.out"); + + auto shapeqkv = info.shapeQKV; + + q = _Reshape(q, shapeqkv); + setName(q, attnPrefix + "q_proj.out_reshape"); + RopeInfo ropeParam; + ropeParam.cutHeadDim = info.ropeCutHeadDim; + + if (mConverter->hasTensor(attnPrefix + "q_norm.weight")) { + auto qNorm = load(attnPrefix + "q_norm.weight"); + ropeParam.qNorm = {qNorm, nullptr, ln_eps, true}; + } + + k = _Reshape(k, shapeqkv); + setName(k, attnPrefix + "k_proj.out_reshape"); + if (mConverter->hasTensor(attnPrefix + "k_norm.weight")) { + auto kNorm = load(attnPrefix + "k_norm.weight"); + ropeParam.kNorm = {kNorm, nullptr, ln_eps, true}; + } + + v = _Reshape(v, shapeqkv); + setName(v, attnPrefix + "v_proj.out_reshape"); + + // RoPE + { + auto ropeOutputs = _TransformerRoPE(q, k, info.cosEven, info.cosOdd, info.sinEven, info.sinOdd, ropeParam); + q = ropeOutputs[0]; + k = ropeOutputs[1]; + setName(q, attnPrefix + "q_after_rope"); + setName(k, attnPrefix + "k_after_rope"); + } + + auto attn = _GPT2Attention(numHeads, headDim, q, k, v, nullptr, nullptr, nullptr, nullptr, mask, useC4Opt); + setName(attn, attnPrefix + "attention.out"); + + if (attnOutSize != numHeads * headDim) { + return {nullptr, nullptr}; + } + auto o = _linear2d(attn, oWeight); + setName(o, attnPrefix + "o_proj.out"); + + // RMSNorm + MLP + auto ln2Weight = load(blockPrefix + "post_attention_layernorm.weight"); + auto fuseLayerNorm = _BinaryLayerNorm(hiddenState, o, {ln2Weight, nullptr, ln_eps, true, hiddenSize, true}); + hiddenStateNorm = fuseLayerNorm.second; + hiddenState = fuseLayerNorm.first; + setName(hiddenState, blockPrefix + "resid1"); + setName(hiddenStateNorm, blockPrefix + "post_attention_layernorm.out"); + hiddenStateNorm4d = hiddenStateNorm; + + auto gateWeight = load(mlpPrefix + "gate_proj.weight"); + auto upWeight = load(mlpPrefix + "up_proj.weight"); + auto downWeight = load(mlpPrefix + "down_proj.weight"); + + auto gate = _linear2d(hiddenStateNorm4d, gateWeight); + setName(gate, mlpPrefix + "gate_proj.out"); + auto up = _linear2d(hiddenStateNorm4d, upWeight); + setName(up, mlpPrefix + "up_proj.out"); + + auto ffn = _MulSilu(up, gate); + setName(ffn, mlpPrefix + "mul_silu.out"); + + ffn = _linear2d(ffn, downWeight); + setName(ffn, mlpPrefix + "down_proj.out"); + + return {hiddenState, ffn}; + } + +private: + const Converter* mConverter = nullptr; +}; + +void HuggingFaceQwen3Convert(const Converter* converter, MNN::NetT* dst, const HuggingFaceQwen3Config& config) { + if (nullptr == converter || nullptr == dst) { + return; + } + + HuggingFaceQwen3 qwen3(converter); + + int blockSize = config.blockNumber; + if (blockSize <= 0) { + const int maxBlockSize = 256; + for (int blockIndex = 0; blockIndex < maxBlockSize; ++blockIndex) { + auto prefix = std::string("model.layers.") + std::to_string(blockIndex) + ".self_attn.q_proj.weight"; + if (!converter->hasTensor(prefix)) { + blockSize = blockIndex; + break; + } + } + } + + int hiddenSize = config.hiddenSize; + int headDim = config.headDim; + int numHead = config.numHead; + + if (hiddenSize <= 0 || headDim <= 0 || numHead <= 0) { + auto qWeight0 = converter->loadTensor("model.layers.0.self_attn.q_proj.weight"); + auto kWeight0 = converter->loadTensor("model.layers.0.self_attn.k_proj.weight"); + if (nullptr != qWeight0.get() && nullptr != qWeight0->getInfo() && qWeight0->getInfo()->dim.size() >= 2) { + const int queryHiddenSize = qWeight0->getInfo()->dim[0]; + const int inputHiddenSize = qWeight0->getInfo()->dim[1]; + if (hiddenSize <= 0) { + hiddenSize = inputHiddenSize; + } + + if (numHead > 0 && headDim <= 0 && queryHiddenSize % numHead == 0) { + headDim = queryHiddenSize / numHead; + } else if (headDim > 0 && numHead <= 0 && queryHiddenSize % headDim == 0) { + numHead = queryHiddenSize / headDim; + } else if (headDim <= 0 && numHead <= 0) { + int kvHiddenSize = 0; + if (nullptr != kWeight0.get() && nullptr != kWeight0->getInfo() && kWeight0->getInfo()->dim.size() >= 2) { + kvHiddenSize = kWeight0->getInfo()->dim[0]; + } + static const int candidates[] = {128, 96, 80, 72, 64, 48, 40, 32}; + for (int cand : candidates) { + if (cand <= 0) { + continue; + } + if (queryHiddenSize % cand != 0) { + continue; + } + if (kvHiddenSize > 0 && kvHiddenSize % cand != 0) { + continue; + } + int candHead = queryHiddenSize / cand; + if (candHead > 0 && candHead <= 64) { + headDim = cand; + numHead = candHead; + break; + } + } + } + } + } + + HuggingFaceQwen3::BlockInfo blockInfo; + blockInfo.hiddenSize = hiddenSize > 0 ? hiddenSize : 1024; + blockInfo.headDim = headDim > 0 ? headDim : 128; + blockInfo.numberHead = numHead > 0 ? numHead : 16; + blockInfo.ropeCutHeadDim = config.ropeCutHeadDim; + + auto embed = _Input({1, -1, blockInfo.hiddenSize}, NCHW, halide_type_of()); + embed->setName("input_embedding"); + + auto position = _Input({1, -1}, NCHW, halide_type_of()); + position->setName("position_ids"); + + auto mask = _Input({}, NCHW, halide_type_of()); + mask->setName("mask"); + + auto one = _Unsqueeze(_Scalar(1), {0}); + auto negone = _Unsqueeze(_Scalar(-1), {0}); + auto shapeHiddenState = _Shape(embed, true); + auto seqLenVar = _Slice(shapeHiddenState, _Unsqueeze(_Scalar(1), {0}), one); + auto batchLenVar = _Slice(shapeHiddenState, _Unsqueeze(_Scalar(0), {0}), one); + + auto headDimVar = _Unsqueeze(_Scalar(blockInfo.headDim), {0}); + + const int posEmbEnd = config.maxPositionEmbeddings > 0 ? config.maxPositionEmbeddings : 32768; + const float ropeTheta = config.ropeTheta > 0.0f ? config.ropeTheta : 100000.0f; + auto posEmb = _PrecomputePosEmbedding(blockInfo.headDim, posEmbEnd, ropeTheta); + posEmb.fix(VARP::CONSTANT); + posEmb->setName("precompute_posemb"); + + posEmb = _GatherV2(posEmb, position, _Scalar(1)); + auto cosAndsin = _Split(posEmb, {2}, 0); + + blockInfo.shapeQKV = _Concat({batchLenVar, seqLenVar, negone, headDimVar}, 0); + blockInfo.shapeQKV->setName("shape_qkv"); + + blockInfo.cosEven = _Squeeze(cosAndsin[0], {0}); + blockInfo.cosOdd = _Squeeze(cosAndsin[0], {0}); + blockInfo.sinEven = _Squeeze(cosAndsin[1], {0}); + blockInfo.sinOdd = _Squeeze(cosAndsin[1], {0}); + + auto hiddenState = _Reshape(embed, {-1, hiddenSize, 1, 1}); + hiddenState = _Convert(hiddenState, NC4HW4); + VARP add = nullptr; + for (int blockIndex = 0; blockIndex < blockSize; ++blockIndex) { + auto res = qwen3.makeBlock(hiddenState, add, blockInfo, mask, blockIndex); + hiddenState = res.first; + add = res.second; + hiddenState->setName("block" + std::to_string(blockIndex)); + } + + // Final RMSNorm + if (add.get() != nullptr) { + hiddenState = _Add(hiddenState, add); + } + auto normWeight = converter->loadTensor("model.norm.weight"); + hiddenState = _TransformerLayerNorm(hiddenState, {normWeight, nullptr, 1.0e-6f, true, blockInfo.hiddenSize, true}); + hiddenState = _Reshape(hiddenState, shapeHiddenState); + hiddenState->setName("hidden_state"); + + std::vector outputs = {hiddenState}; + std::vector outputNames = {"hidden_state"}; + if (config.outputLastHiddenState) { + auto lastHiddenState = _MakeLastHiddenStateOutput(hiddenState, blockInfo.hiddenSize); + outputs.emplace_back(lastHiddenState); + outputNames.emplace_back("last_hidden_state"); + } + + Variable::save(outputs, dst); + dst->sourceType = NetSource_ONNX; + dst->outputName = std::move(outputNames); +} + +namespace { +static bool _convertHuggingFaceDecoderModel(const Converter* converter, const rapidjson::Value* model, modelConfig& modelPath) { + if (nullptr == converter) { + return false; + } + + auto netT = std::unique_ptr(new MNN::NetT); + HuggingFaceQwen3Config config; + + if (nullptr != model && model->IsObject()) { + auto blocks = WorkflowJson::getArray(*model, "blocks"); + if (nullptr != blocks) { + for (auto& block : blocks->GetArray()) { + if (!block.IsObject()) { + continue; + } + auto type = WorkflowJson::getString(block, "type"); + if (type == "QwenTransformer") { + config.hiddenSize = WorkflowJson::getInt(block, "hiddenSize", config.hiddenSize); + config.headDim = WorkflowJson::getInt(block, "headDim", config.headDim); + config.numHead = WorkflowJson::getInt(block, "numHead", config.numHead); + config.kvNumHead = WorkflowJson::getInt(block, "kvNumHead", config.kvNumHead); + config.blockNumber = WorkflowJson::getInt(block, "number", config.blockNumber); + config.maxPositionEmbeddings = WorkflowJson::getInt(block, "maxPositionEmbeddings", config.maxPositionEmbeddings); + config.maxPositionEmbeddings = WorkflowJson::getInt(block, "max_position_embeddings", config.maxPositionEmbeddings); + // backward compatible field name (legacy) + config.maxPositionEmbeddings = WorkflowJson::getInt(block, "bit", config.maxPositionEmbeddings); + config.ropeTheta = WorkflowJson::getFloat(block, "ropeTheta", config.ropeTheta); + config.ropeTheta = WorkflowJson::getFloat(block, "rope_theta", config.ropeTheta); + config.ropeCutHeadDim = WorkflowJson::getInt(block, "ropeCutHeadDim", config.ropeCutHeadDim); + config.ropeCutHeadDim = WorkflowJson::getInt(block, "rope_cut_head_dim", config.ropeCutHeadDim); + break; + } + } + } + } + + auto path = modelPath.MNNModel; + modelPath.MNNModel = path + "decoder.mnn"; + HuggingFaceQwen3Convert(converter, netT.get(), config); + optimizeAndWrite(modelPath, netT); + return true; +} + +REGISTER_SAFETENSOR_MODEL_BUILDER("hf_decoder", _convertHuggingFaceDecoderModel); +} // namespace + +} // namespace SafeTensors +} // namespace MNN diff --git a/tools/converter/source/safetensors/HuggingFaceQwen3.hpp b/tools/converter/source/safetensors/HuggingFaceQwen3.hpp new file mode 100644 index 0000000000..29d0d1387c --- /dev/null +++ b/tools/converter/source/safetensors/HuggingFaceQwen3.hpp @@ -0,0 +1,26 @@ +#ifndef HuggingFaceQwen3_hpp +#define HuggingFaceQwen3_hpp + +#include "SafetensorConverter.hpp" + +namespace MNN { +namespace SafeTensors { + +struct HuggingFaceQwen3Config { + int hiddenSize = 0; + int headDim = 0; + int numHead = 0; + int kvNumHead = 0; + int blockNumber = 0; + int maxPositionEmbeddings = 0; + float ropeTheta = 0.0f; + int ropeCutHeadDim = 0; + bool outputLastHiddenState = true; +}; + +void HuggingFaceQwen3Convert(const Converter* converter, MNN::NetT* dst, const HuggingFaceQwen3Config& config); + +} // namespace SafeTensors +} // namespace MNN + +#endif diff --git a/tools/converter/source/safetensors/Logit.cpp b/tools/converter/source/safetensors/Logit.cpp new file mode 100644 index 0000000000..970586505a --- /dev/null +++ b/tools/converter/source/safetensors/Logit.cpp @@ -0,0 +1,513 @@ +#include +#include +#include +#include +#include +#include "MNN_generated.h" +#include + +#include "SafetensorConverter.hpp" +#include "Logit.hpp" +#include "SafetensorModelRegistry.hpp" +#include "SafetensorUtils.hpp" +#include "WorkflowJson.hpp" + +using namespace MNN::Express; +using namespace MNN::Express::SafeTensorUtils; + +namespace MNN { +namespace SafeTensors { + +static inline void _setNameIfEmpty(const VARP& v, const std::string& name) { + if (nullptr != v.get() && v->name().empty()) { + v->setName(name); + } +} + +static VARP _linear2d(VARP x4d, VARP weightOI, VARP bias = nullptr) { + auto wInfo = weightOI->getInfo(); + if (nullptr == wInfo || wInfo->dim.size() < 2) { + return nullptr; + } + + const int outDim = wInfo->dim[0]; + const int inDim = wInfo->dim[1]; + if (inDim <= 0 || outDim <= 0) { + return nullptr; + } + + std::vector weightData(weightOI->getInfo()->size); + ::memcpy(weightData.data(), weightOI->readMap(), weightData.size() * sizeof(float)); + std::vector biasData(outDim, 0.0f); + if (nullptr != bias) { + ::memcpy(biasData.data(), bias->readMap(), outDim * sizeof(float)); + } + + return _Conv(std::move(weightData), std::move(biasData), x4d, {inDim, outDim}, {1, 1}); +} + +// Deep-copy `src` into `dst` via flatbuffers Pack/UnPack and strip Convolution2D +// weight payloads — the runtime reuses the original logit's quantized weights. +static void _cloneLogitNet(const MNN::NetT* src, MNN::NetT* dst) { + flatbuffers::FlatBufferBuilder fbb; + fbb.Finish(MNN::CreateNet(fbb, src)); + std::unique_ptr cloned(flatbuffers::GetRoot(fbb.GetBufferPointer())->UnPack()); + *dst = std::move(*cloned); + for (auto& op : dst->oplists) { + if (op && op->main.type == OpParameter_Convolution2D) { + auto conv = op->main.AsConvolution2D(); + conv->weight.clear(); + conv->bias.clear(); + conv->quanParameter.reset(); + conv->external.clear(); + } + } +} + +// Locate logits tensor index — prefer outputName mapping, fall back to the last +// op with an output index. +static int _findLogitsIndex(const MNN::NetT* net) { + if (!net->outputName.empty()) { + const auto& outName = net->outputName[0]; + for (int i = 0; i < (int)net->tensorName.size(); ++i) { + if (net->tensorName[i] == outName) return i; + } + } + for (int i = (int)net->oplists.size() - 1; i >= 0; --i) { + if (net->oplists[i] && !net->oplists[i]->outputIndexes.empty()) { + return net->oplists[i]->outputIndexes[0]; + } + } + return -1; +} + +static int _addTensor(MNN::NetT* net, const std::string& name) { + int idx = (int)net->tensorName.size(); + net->tensorName.push_back(name); + return idx; +} + +static int _appendConstInt(MNN::NetT* net, const std::string& opName, const std::string& tensorName, int value) { + int idx = _addTensor(net, tensorName); + std::unique_ptr op(new OpT); + op->type = OpType_Const; + op->main.type = OpParameter_Blob; + op->main.value = new BlobT; + auto blob = op->main.AsBlob(); + blob->dataFormat = MNN_DATA_FORMAT_NCHW; + blob->dataType = DataType_DT_INT32; + blob->dims = {1}; + blob->int32s = {value}; + op->name = opName; + op->outputIndexes = {idx}; + net->oplists.emplace_back(std::move(op)); + return idx; +} + +static int _appendSoftmaxOp(MNN::NetT* net, int inputIdx, const std::string& opName, const std::string& tensorName, int axis = -1) { + int idx = _addTensor(net, tensorName); + std::unique_ptr op(new OpT); + op->type = OpType_Softmax; + op->main.type = OpParameter_Axis; + op->main.value = new AxisT; + op->main.AsAxis()->axis = axis; + op->name = opName; + op->inputIndexes = {inputIdx}; + op->outputIndexes = {idx}; + net->oplists.emplace_back(std::move(op)); + return idx; +} + +// Returns {valuesIdx, indicesIdx}. Default largest=true (no main parameter). +static std::pair _appendTopK2Op(MNN::NetT* net, int inputIdx, int kIdx, + const std::string& opName, + const std::string& valuesName, + const std::string& indicesName) { + int valuesIdx = _addTensor(net, valuesName); + int indicesIdx = _addTensor(net, indicesName); + std::unique_ptr op(new OpT); + op->type = OpType_TopKV2; + op->name = opName; + op->inputIndexes = {inputIdx, kIdx}; + op->outputIndexes = {valuesIdx, indicesIdx}; + net->oplists.emplace_back(std::move(op)); + return {valuesIdx, indicesIdx}; +} + +static int _appendUnaryOp(MNN::NetT* net, int inputIdx, UnaryOpOperation kind, + const std::string& opName, const std::string& tensorName) { + int idx = _addTensor(net, tensorName); + std::unique_ptr op(new OpT); + op->type = OpType_UnaryOp; + op->main.type = OpParameter_UnaryOp; + op->main.value = new UnaryOpT; + op->main.AsUnaryOp()->opType = kind; + op->name = opName; + op->inputIndexes = {inputIdx}; + op->outputIndexes = {idx}; + net->oplists.emplace_back(std::move(op)); + return idx; +} + +// Clone `logit` into `dst` and resolve the logits tensor index. Returns -1 on +// failure (also logs the error tagged with `fnTag`). +static int _cloneAndFindLogits(const MNN::NetT* logit, MNN::NetT* dst, const char* fnTag) { + if (nullptr == logit || nullptr == dst) { + return -1; + } + _cloneLogitNet(logit, dst); + int idx = _findLogitsIndex(dst); + if (idx < 0) { + MNN_ERROR("%s: 未找到 logits 输出\n", fnTag); + } + return idx; +} + +void LogitConvert(const Converter* converter, MNN::NetT* dst, const LogitConfig& config) { + if (nullptr == converter || nullptr == dst) { + return; + } + + auto weight = converter->loadTensor(config.wteWeightName); + if (nullptr == weight.get() && config.wteWeightName.size() > 7 && config.wteWeightName.substr(0, 7) == "module.") { + weight = converter->loadTensor(config.wteWeightName.substr(7)); + } + if (nullptr == weight.get() || nullptr == weight->getInfo() || weight->getInfo()->dim.size() < 2) { + MNN_ERROR("LogitConvert: missing/invalid %s\n", config.wteWeightName.c_str()); + return; + } + + const int d0 = weight->getInfo()->dim[0]; + const int d1 = weight->getInfo()->dim[1]; + if (d0 <= 0 || d1 <= 0) { + MNN_ERROR("LogitConvert: invalid wte weight shape\n"); + return; + } + + int hiddenSize = config.hiddenSize; + int vocabSize = 0; + bool needTranspose = false; + + if (hiddenSize > 0) { + if (d1 == hiddenSize) { + vocabSize = d0; + } else if (d0 == hiddenSize) { + vocabSize = d1; + needTranspose = true; + } else { + // Fallback: assume [vocab, hidden] + hiddenSize = d1; + vocabSize = d0; + } + } else { + // Fallback heuristics: vocab is usually larger than hidden + if (d0 >= d1) { + vocabSize = d0; + hiddenSize = d1; + } else { + vocabSize = d1; + hiddenSize = d0; + needTranspose = true; + } + } + + if (needTranspose) { + weight = _Transpose(weight, {1, 0}); // -> [vocab, hidden] + } + + // Input hidden state: [B, S, H] + auto hiddenState = _Input({1, -1, hiddenSize}, NCHW, halide_type_of()); + hiddenState->setName(config.inputName); + + auto shapeHidden = _Shape(hiddenState, true); + auto one = _Unsqueeze(_Scalar(1), {0}); + auto batchVar = _Slice(shapeHidden, _Unsqueeze(_Scalar(0), {0}), one); + auto seqVar = _Slice(shapeHidden, _Unsqueeze(_Scalar(1), {0}), one); + + auto hidden2d = _Reshape(hiddenState, {-1, hiddenSize, 1, 1}); + + // Optional bias + VARP bias = nullptr; + auto prefix = config.wteWeightName; + const std::string suffix = ".weight"; + if (prefix.size() > suffix.size() && prefix.compare(prefix.size() - suffix.size(), suffix.size(), suffix) == 0) { + prefix.resize(prefix.size() - suffix.size()); + } + auto biasName = prefix + ".bias"; + if (converter->hasTensor(biasName)) { + bias = converter->loadTensor(biasName); + } else if (biasName.size() > 7 && biasName.substr(0, 7) == "module." && converter->hasTensor(biasName.substr(7))) { + bias = converter->loadTensor(biasName.substr(7)); + } + + // Quantized path if weight_qscale exists + VARP logits2d = nullptr; + auto wScaleName = config.wteWeightName + "_qscale"; + std::string textScaleName = config.wteWeightName; + if (textScaleName.find(".weight") != std::string::npos) { + textScaleName.replace(textScaleName.find(".weight"), 7, ".text_embedding.weight_qscale"); + } + + auto loadScale = [&](const std::string& name) -> VARP { + if (converter->hasTensor(name)) { + return converter->loadTensor(name); + } + if (name.size() > 7 && name.substr(0, 7) == "module." && converter->hasTensor(name.substr(7))) { + return converter->loadTensor(name.substr(7)); + } + return nullptr; + }; + + VARP wScale = loadScale(wScaleName); + if (nullptr == wScale.get()) { + wScale = loadScale(textScaleName); + } + + if (wScale.get() != nullptr) { + logits2d = _QConvolution1x1(hiddenSize, hidden2d, nullptr, nullptr, weight, wScale, nullptr, bias, vocabSize); + } else { + // Float path + if (weight->getInfo()->type.code != halide_type_float) { + MNN_ERROR("LogitConvert: wte weight is not float and no qscale found\n"); + return; + } + logits2d = _linear2d(hidden2d, weight, bias); + } + + if (nullptr == logits2d.get()) { + MNN_ERROR("LogitConvert: build logits failed\n"); + return; + } + _setNameIfEmpty(logits2d, prefix + ".out2d"); + + auto vocabVar = _Unsqueeze(_Scalar(vocabSize), {0}); + auto logits3d = _Reshape(logits2d, _Concat({batchVar, seqVar, vocabVar}, 0)); + logits3d->setName(config.outputName); + + Variable::save({logits3d}, dst); + dst->sourceType = NetSource_ONNX; + dst->outputName = {config.outputName}; + +} + +void MakeTieEmbedding(const Converter* converter, const MNN::NetT* src, MNN::NetT* dst) { + if (nullptr == converter || nullptr == src || nullptr == dst) { + return; + } + + std::string sharedName; + int ic = 0; + int oc = 0; + + for (auto& op : src->oplists) { + if (nullptr == op) { + continue; + } + if (op->type != OpType_Convolution) { + continue; + } + auto conv = op->main.AsConvolution2D(); + if (nullptr == conv || nullptr == conv->common) { + continue; + } + if (conv->common->inputCount <= 0 || conv->common->outputCount <= 0) { + continue; + } + // Keep the last conv as shared weight provider (align with makeSharedGather.py). + sharedName = op->name; + ic = conv->common->inputCount; + oc = conv->common->outputCount; + } + + if (sharedName.empty() || ic <= 0 || oc <= 0) { + MNN_ERROR("MakeTieEmbedding: can't find valid convolution in src\n"); + return; + } + + // Indices input. + auto input = _Input({-1}, NCHW, halide_type_of()); + input->setName("x"); + + // GatherV2 with OpParameter_Input main is a special form that lets runtime reuse + // the quantized weights from base model's convolution execution. + std::unique_ptr gather(new OpT); + gather->type = OpType_GatherV2; + gather->main.type = OpParameter_Input; + gather->main.value = new InputT; + gather->main.AsInput()->dims = {oc, ic}; + gather->main.AsInput()->dtype = DataType_DT_FLOAT; + gather->main.AsInput()->dformat = MNN_DATA_FORMAT_NCHW; + + auto gatherExpr = Expr::create(gather.get(), {input}); + gatherExpr->setName(sharedName); + + auto output = Variable::create(gatherExpr); + output->setName(sharedName); + + Variable::save({output}, dst); + dst->sourceType = NetSource_ONNX; + dst->outputName = {sharedName}; +} + + +// Clone `logit` into `dst` and append a TopKV2 producing top-K indices. +void MakeTopKV(const Converter* /*converter*/, const MNN::NetT* logit, MNN::NetT* dst, int K) { + int logitsIdx = _cloneAndFindLogits(logit, dst, "MakeTopKV"); + if (logitsIdx < 0) return; + + int kIdx = _appendConstInt(dst, "const_topk_k", "topk_k", K); + auto vi = _appendTopK2Op(dst, logitsIdx, kIdx, "TopKV2", "topk_values", "topk_indices"); + dst->outputName = {dst->tensorName[vi.second]}; +} + +// Clone `logit` into `dst` and append a Softmax (axis=-1) as the new output. +void MakeSoftmax(const Converter* /*converter*/, const MNN::NetT* logit, MNN::NetT* dst) { + int logitsIdx = _cloneAndFindLogits(logit, dst, "MakeSoftmax"); + if (logitsIdx < 0) return; + + int smIdx = _appendSoftmaxOp(dst, logitsIdx, "LogitSoftmax", "logit_softmax"); + dst->outputName = {dst->tensorName[smIdx]}; +} + +// Clone `logit` into `dst`, append Softmax → TopKV2 → Log(values) for beam search. +// Outputs {log(values), indices}. +void MakeBeamTopKV(const Converter* /*converter*/, const MNN::NetT* logit, MNN::NetT* dst, int K) { + int logitsIdx = _cloneAndFindLogits(logit, dst, "MakeBeamTopKV"); + if (logitsIdx < 0) return; + + int smIdx = _appendSoftmaxOp(dst, logitsIdx, "BeamSoftmax", "beam_softmax"); + int kIdx = _appendConstInt(dst, "const_beam_topk_k", "beam_topk_k", K); + auto vi = _appendTopK2Op(dst, smIdx, kIdx, "BeamTopKV2", "beam_topk_values", "beam_topk_indices"); + int logIdx = _appendUnaryOp(dst, vi.first, UnaryOpOperation_LOG, "BeamTopKV2_Log", "beam_topk_log_values"); + dst->outputName = {dst->tensorName[logIdx], dst->tensorName[vi.second]}; +} + +namespace { + +// Parse the K parameter, accepting either an int array or a single int (legacy). +// Always returns a non-empty list (defaults to {1}). +static std::vector _parseKList(const rapidjson::Value& block) { + std::vector kList; + if (auto kArr = WorkflowJson::getArray(block, "K")) { + for (auto& kv : kArr->GetArray()) { + if (kv.IsInt() && kv.GetInt() > 0) kList.push_back(kv.GetInt()); + } + } else { + int K = WorkflowJson::getInt(block, "K", 1); + if (K > 0) kList.push_back(K); + } + if (kList.empty()) kList.push_back(1); + return kList; +} + +// Build and write a per-K logit variant for each entry in `kList`. External +// data is force-disabled for variants since they reuse the base logit's +// quantized weights. +template +static void _saveLogitVariants(const Converter* converter, MNN::NetT* logitNet, + modelConfig& modelPath, const std::string& path, + const std::vector& kList, + const std::string& filePrefix, Builder build) { + auto originExternal = modelPath.saveExternalData; + modelPath.saveExternalData = false; + for (int K : kList) { + auto net = std::unique_ptr(new MNN::NetT); + build(converter, logitNet, net.get(), K); + modelPath.MNNModel = path + filePrefix + std::to_string(K) + ".mnn"; + optimizeAndWrite(modelPath, net); + } + modelPath.saveExternalData = originExternal; +} + +static bool _convertLogitModel(const Converter* converter, const rapidjson::Value* model, modelConfig& modelPath) { + if (nullptr == converter) { + return false; + } + + LogitConfig config; + auto path = modelPath.MNNModel; + + auto logitNet = std::unique_ptr(new MNN::NetT); + std::unique_ptr embeddingNet; + + auto ensureLogits = [&]() { + if (logitNet->oplists.empty()) { + LogitConvert(converter, logitNet.get(), config); + } + }; + + if (nullptr != model && model->IsObject()) { + auto blocks = WorkflowJson::getArray(*model, "blocks"); + if (nullptr != blocks) { + for (auto& block : blocks->GetArray()) { + if (!block.IsObject()) { + continue; + } + + auto prefix = WorkflowJson::getString(block, "prefix"); + if (!prefix.empty()) { + auto wteWeightName = prefix; + if (wteWeightName.find(".weight") == std::string::npos) { + auto candidate = wteWeightName + ".weight"; + if (converter->hasTensor(candidate)) { + wteWeightName = candidate; + } + } + if (converter->hasTensor(wteWeightName)) { + config.wteWeightName = wteWeightName; + } + } + + auto type = WorkflowJson::getString(block, "type"); + if (type == "InnerProduct") { + LogitConvert(converter, logitNet.get(), config); + continue; + } + if (type == "TieEmbedding") { + embeddingNet = std::unique_ptr(new MNN::NetT); + MakeTieEmbedding(converter, logitNet.get(), embeddingNet.get()); + continue; + } + if (type == "TopKV") { + ensureLogits(); + _saveLogitVariants(converter, logitNet.get(), modelPath, path, + _parseKList(block), "logit_topkv_", MakeTopKV); + continue; + } + if (type == "Softmax") { + ensureLogits(); + auto originExternal = modelPath.saveExternalData; + modelPath.saveExternalData = false; + auto softmaxNet = std::unique_ptr(new MNN::NetT); + MakeSoftmax(converter, logitNet.get(), softmaxNet.get()); + modelPath.MNNModel = path + "logit_softmax.mnn"; + optimizeAndWrite(modelPath, softmaxNet); + modelPath.saveExternalData = originExternal; + continue; + } + if (type == "BeamTopKV") { + ensureLogits(); + _saveLogitVariants(converter, logitNet.get(), modelPath, path, + _parseKList(block), "logit_beam_", MakeBeamTopKV); + continue; + } + } + } + } + + modelPath.MNNModel = path + "logit.mnn"; + optimizeAndWrite(modelPath, logitNet); + + if (nullptr != embeddingNet.get()) { + modelPath.MNNModel = path + "embed.mnn"; + optimizeAndWrite(modelPath, embeddingNet); + } + return true; +} + +REGISTER_SAFETENSOR_MODEL_BUILDER("logit", _convertLogitModel); + +} // namespace + +} // namespace SafeTensors +} // namespace MNN diff --git a/tools/converter/source/safetensors/Logit.hpp b/tools/converter/source/safetensors/Logit.hpp new file mode 100644 index 0000000000..d946f97226 --- /dev/null +++ b/tools/converter/source/safetensors/Logit.hpp @@ -0,0 +1,30 @@ +#ifndef Logit_hpp +#define Logit_hpp + +#include +#include "SafetensorConverter.hpp" + +namespace MNN { +namespace SafeTensors { + +struct LogitConfig { + // Default to Qwen/GPT2 word embedding matrix + std::string wteWeightName = "module.gpt2.transformer.wte.weight"; + + // Optional: if provided, will be used to disambiguate weight layout + int hiddenSize = 0; + + std::string inputName = "hidden_state"; + std::string outputName = "output"; +}; + +void LogitConvert(const Converter* converter, MNN::NetT* dst, const LogitConfig& config); +void MakeTieEmbedding(const Converter* converter, const MNN::NetT* src, MNN::NetT* dst); +void MakeTopKV(const Converter* converter, const MNN::NetT* logit, MNN::NetT* dst, int K); +void MakeSoftmax(const Converter* converter, const MNN::NetT* logit, MNN::NetT* dst); +void MakeBeamTopKV(const Converter* converter, const MNN::NetT* logit, MNN::NetT* dst, int K); + +} // namespace SafeTensors +} // namespace MNN + +#endif diff --git a/tools/converter/source/safetensors/SafetensorConverter.cpp b/tools/converter/source/safetensors/SafetensorConverter.cpp new file mode 100644 index 0000000000..ab0a15768e --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorConverter.cpp @@ -0,0 +1,214 @@ +#include +#include +#include + +#include +#include + +#include + +#include "SafetensorConverter.hpp" +#include "SafetensorModelRegistry.hpp" +#include "WorkflowJson.hpp" + +#include "../common/CommonUtils.hpp" + +#define SAFETENSORS_CPP_IMPLEMENTATION +#include "safetensors.hh" +namespace MNN { +namespace SafeTensors { + +static halide_type_t _convertSafeTensorDType(safetensors::dtype dtype) { + switch (dtype) { + case safetensors::kBOOL: + // Safetensors stores BOOL as 1 byte. MNN's 1-bit bool type is not widely supported. + return halide_type_of(); + case safetensors::kUINT8: + return halide_type_of(); + case safetensors::kINT8: + return halide_type_of(); + case safetensors::kINT16: + return halide_type_of(); + case safetensors::kUINT16: + return halide_type_of(); + case safetensors::kINT32: + return halide_type_of(); + case safetensors::kUINT32: + return halide_type_of(); + case safetensors::kINT64: + return halide_type_of(); + case safetensors::kUINT64: + return halide_type_of(); + case safetensors::kFLOAT16: + return halide_type_t(halide_type_float, 16); + case safetensors::kBFLOAT16: + return halide_type_t(halide_type_bfloat, 16); + case safetensors::kFLOAT32: + return halide_type_of(); + case safetensors::kFLOAT64: + return halide_type_of(); + default: + break; + } + return halide_type_of(); +} + +struct Converter::Content { + rapidjson::Document mWorkFlow; + safetensors::safetensors_t mSt; +}; + +Converter::Converter(const std::string& jsonFile) { + mMain = new Content; + + std::ifstream fileNames(jsonFile); + std::ostringstream output; + output << fileNames.rdbuf(); + auto outputStr = output.str(); + + mMain->mWorkFlow.Parse(outputStr.c_str()); + if (mMain->mWorkFlow.HasParseError() || !mMain->mWorkFlow.IsObject()) { + MNN_ERROR("Invalid json\n"); + mMain->mWorkFlow.SetObject(); + return; + } +} + +Converter::~ Converter() { + delete mMain; +} + +std::vector Converter::listModels() const { + std::vector res; + if (nullptr == mMain) { + return res; + } + auto models = WorkflowJson::getArray(mMain->mWorkFlow, "models"); + if (nullptr == models) { + return res; + } + for (auto& model : models->GetArray()) { + if (!model.IsObject()) { + continue; + } + auto name = WorkflowJson::getString(model, "name"); + if (name.empty()) { + continue; + } + res.emplace_back(std::move(name)); + } + return res; +} +void Converter::loadSafeTensors(const std::string& safeTensorFile) { + std::string warn, err; + auto ret = safetensors::mmap_from_file(safeTensorFile, &mMain->mSt, &warn, &err); + if (warn.size()) { + FUNC_PRINT_ALL(warn.c_str(), s); + } + if (!ret) { + FUNC_PRINT_ALL(err.c_str(), s); + return; + } +} +bool Converter::convert(const std::string& name, modelConfig& modelPath) { + auto builder = SafetensorModelRegistry::get()->find(name); + if (nullptr == builder) { + MNN_ERROR("SafetensorConverter: unsupported model %s\n", name.c_str()); + return false; + } + + const rapidjson::Value* model = nullptr; + if (mMain != nullptr && mMain->mWorkFlow.IsObject()) { + auto models = WorkflowJson::getArray(mMain->mWorkFlow, "models"); + if (nullptr != models) { + for (auto& item : models->GetArray()) { + if (!item.IsObject()) { + continue; + } + auto modelName = WorkflowJson::getString(item, "name"); + MNN_PRINT("Checking model config for: %s (target: %s)\n", modelName.c_str(), name.c_str()); + if (!modelName.empty() && modelName == name) { + model = &item; + int weightQuantBits = WorkflowJson::getInt(item, "weightQuantBits", -1); + if (weightQuantBits >= 0) { + MNN_PRINT("Override weightQuantBits to %d for model %s\n", weightQuantBits, name.c_str()); + modelPath.weightQuantBits = weightQuantBits; + } + break; + } + } + } + } + + return builder(this, model, modelPath); +} +bool Converter::hasTensor(const std::string& name) const { + safetensors::tensor_t t; + if (mMain->mSt.tensors.at(name, &t)) { + return true; + } + return false; +} + +MNN::Express::VARP Converter::loadTensor(const std::string& name, bool print) const { + safetensors::tensor_t t; + bool find = mMain->mSt.tensors.at(name, &t); + if (!find) { + if (print) { + FUNC_PRINT_ALL(name.c_str(), s); + } + return nullptr; + } + + const uint8_t* dataBufferAddr = nullptr; + size_t dataBufferSize = 0; + if (mMain->mSt.mmaped) { + dataBufferAddr = mMain->mSt.databuffer_addr; + dataBufferSize = mMain->mSt.databuffer_size; + } else { + dataBufferAddr = mMain->mSt.storage.data(); + dataBufferSize = mMain->mSt.storage.size(); + } + if (nullptr == dataBufferAddr || dataBufferSize == 0) { + MNN_ERROR("Safetensors databuffer is empty, please call loadSafeTensors first\n"); + return nullptr; + } + + const size_t offsetBegin = t.data_offsets[0]; + const size_t offsetEnd = t.data_offsets[1]; + if (offsetBegin > offsetEnd || offsetEnd > dataBufferSize) { + MNN_ERROR("Invalid tensor offsets for %s: [%zu, %zu), databuffer=%zu\n", name.c_str(), offsetBegin, offsetEnd, dataBufferSize); + return nullptr; + } + + const size_t nitems = safetensors::get_shape_size(t); + const size_t itemBytes = safetensors::get_dtype_bytes(t.dtype); + const size_t expectedBytes = nitems * itemBytes; + const size_t actualBytes = offsetEnd - offsetBegin; + if (expectedBytes != actualBytes) { + MNN_ERROR("Invalid tensor %s: expected %zu bytes(%zu*%zu), got %zu\n", name.c_str(), expectedBytes, nitems, itemBytes, actualBytes); + return nullptr; + } + + MNN::Express::INTS shape; + shape.reserve(t.shape.size()); + for (auto dim : t.shape) { + if (dim > static_cast(std::numeric_limits::max())) { + MNN_ERROR("Tensor %s has too large shape dim: %zu\n", name.c_str(), dim); + return nullptr; + } + shape.emplace_back(static_cast(dim)); + } + + auto tensorStart = dataBufferAddr + offsetBegin; + auto dtype = _convertSafeTensorDType(t.dtype); + auto var = MNN::Express::_Const(tensorStart, std::move(shape), MNN::Express::NCHW, dtype); + if (t.dtype == safetensors::kBFLOAT16) { + var = MNN::Express::_Cast(var); + var.fix(MNN::Express::VARP::CONSTANT); + } + return var; +} + +}; +}; diff --git a/tools/converter/source/safetensors/SafetensorConverter.hpp b/tools/converter/source/safetensors/SafetensorConverter.hpp new file mode 100644 index 0000000000..8305372060 --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorConverter.hpp @@ -0,0 +1,27 @@ +#ifndef SafetensorConverter_hpp +#define SafetensorConverter_hpp +#include +#include +#include +#include +#include "config.hpp" +namespace MNN { +namespace SafeTensors { +class MNN_PUBLIC Converter { +public: + Converter(const std::string& jsonFile); + ~ Converter(); + std::vector listModels() const; + void loadSafeTensors(const std::string& safeTensorFile); + bool convert(const std::string& name, modelConfig& modelPath); + MNN::Express::VARP loadTensor(const std::string& name, bool printNotFound = true) const; + bool hasTensor(const std::string& name) const; + struct Content; +private: + Content* mMain = nullptr; +}; +}; +}; + +#endif + diff --git a/tools/converter/source/safetensors/SafetensorModelRegistry.cpp b/tools/converter/source/safetensors/SafetensorModelRegistry.cpp new file mode 100644 index 0000000000..4f9478da12 --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorModelRegistry.cpp @@ -0,0 +1,83 @@ +#include "SafetensorModelRegistry.hpp" + +#include + +#include + +#include "MNN_generated.h" +#include "PostConverter.hpp" +#include "writeFb.hpp" +#include "../common/CommonUtils.hpp" + +namespace MNN { +namespace SafeTensors { + +struct SafetensorModelRegistry::Impl { + std::unordered_map builders; +}; + +SafetensorModelRegistry* SafetensorModelRegistry::get() { + static SafetensorModelRegistry gRegistry; + if (!gRegistry.mImpl) { + gRegistry.mImpl.reset(new Impl); + } + return &gRegistry; +} + +void SafetensorModelRegistry::insert(const std::string& name, ModelBuilder builder) { + if (name.empty() || builder == nullptr) { + return; + } + if (!mImpl) { + mImpl.reset(new Impl); + } + auto iter = mImpl->builders.find(name); + if (iter != mImpl->builders.end()) { + MNN_PRINT("SafetensorModelRegistry: override builder for %s\n", name.c_str()); + } + mImpl->builders[name] = builder; +} + +ModelBuilder SafetensorModelRegistry::find(const std::string& name) const { + if (!mImpl) { + return nullptr; + } + auto iter = mImpl->builders.find(name); + if (iter == mImpl->builders.end()) { + return nullptr; + } + return iter->second; +} + +SafetensorModelRegister::SafetensorModelRegister(const char* name, ModelBuilder builder) { + if (nullptr == name || builder == nullptr) { + return; + } + SafetensorModelRegistry::get()->insert(name, builder); +} + +MNN_PUBLIC void optimizeAndWrite(modelConfig& modelPath, std::unique_ptr& netT) { + if (nullptr == netT.get()) { + return; + } + + std::unique_ptr metaOp(new MNN::OpT); + metaOp->type = MNN::OpType_Extra; + metaOp->main.value = new MNN::ExtraT; + metaOp->main.type = MNN::OpParameter_Extra; + metaOp->main.AsExtra()->type = "Meta"; + metaOp->main.AsExtra()->engine = "MNN"; + + std::vector expectedPass; + CommonKit::loadCompress(modelPath); + + std::unique_ptr newNet = optimizeNet(netT, modelPath.forTraining, modelPath, expectedPass); + if (nullptr != newNet) { + (void)writeFb(newNet, modelPath, std::move(metaOp)); + } else { + MNN_ERROR("SafetensorModelRegistry: optimizeNet failed, skip writing %s\n", modelPath.MNNModel.c_str()); + } +} + +} // namespace SafeTensors +} // namespace MNN diff --git a/tools/converter/source/safetensors/SafetensorModelRegistry.hpp b/tools/converter/source/safetensors/SafetensorModelRegistry.hpp new file mode 100644 index 0000000000..de1f46194b --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorModelRegistry.hpp @@ -0,0 +1,53 @@ +#ifndef SafetensorModelRegistry_hpp +#define SafetensorModelRegistry_hpp + +#include +#include + +#include "config.hpp" + +#include + +#include + +namespace MNN { +struct NetT; + +namespace SafeTensors { + +class Converter; + +using ModelBuilder = bool (*)(const Converter* converter, const rapidjson::Value* model, modelConfig& modelPath); + +class MNN_PUBLIC SafetensorModelRegistry { +public: + static SafetensorModelRegistry* get(); + + void insert(const std::string& name, ModelBuilder builder); + ModelBuilder find(const std::string& name) const; + +private: + SafetensorModelRegistry() = default; + + struct Impl; + std::unique_ptr mImpl; +}; + +class MNN_PUBLIC SafetensorModelRegister { +public: + SafetensorModelRegister(const char* name, ModelBuilder builder); +}; + +MNN_PUBLIC void optimizeAndWrite(modelConfig& modelPath, std::unique_ptr& netT); + +} // namespace SafeTensors +} // namespace MNN + +#define MNN_SAFETENSOR_JOIN_INNER(x, y) x##y +#define MNN_SAFETENSOR_JOIN(x, y) MNN_SAFETENSOR_JOIN_INNER(x, y) + +#define REGISTER_SAFETENSOR_MODEL_BUILDER(modelName, builderFunc) \ + static ::MNN::SafeTensors::SafetensorModelRegister \ + MNN_SAFETENSOR_JOIN(__mnn_safetensor_model_register_, __COUNTER__)(modelName, builderFunc) + +#endif diff --git a/tools/converter/source/safetensors/SafetensorUtils.cpp b/tools/converter/source/safetensors/SafetensorUtils.cpp new file mode 100644 index 0000000000..1bd642f39b --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorUtils.cpp @@ -0,0 +1,437 @@ +#include "SafetensorUtils.hpp" + +#include +#include +#include + +#include + +#include "MNN_generated.h" +#include "core/IDSTEncoder.hpp" + +namespace MNN { +namespace Express { +namespace SafeTensorUtils { + +VARP _MakeLastHiddenStateOutput(VARP hiddenState, int hiddenSize) { + if (nullptr == hiddenState.get()) { + return nullptr; + } + std::vector sizes = {1, 1, hiddenSize}; + auto sizeVar = _Const(sizes.data(), {3}, NCHW, halide_type_of()); + std::vector begins = {0, -1, 0}; + auto beginVar = _Const(begins.data(), {3}, NCHW, halide_type_of()); + auto output = _Slice(hiddenState, beginVar, sizeVar); + output->setName("last_hidden_state"); + return output; +} + +VARP _GPT2Attention(int numHead, int headDim, VARP q, VARP k, VARP v, VARP qk_scale_q, VARP qk_scale_k, + VARP sv_scale_s, VARP sv_scale_v, VARP mask, bool supportC4Opt, float attnScale) { + std::unique_ptr op(new OpT); + op->type = OpType_Attention; + op->main.value = new AttentionParamT; + op->main.type = OpParameter_AttentionParam; + op->main.AsAttentionParam()->kv_cache = true; + op->main.AsAttentionParam()->attnScale = attnScale; + bool supportC4 = (headDim % 16 == 0) && supportC4Opt; + op->main.AsAttentionParam()->output_c4 = supportC4; + if (nullptr != qk_scale_q || nullptr != qk_scale_k) { + op->main.AsAttentionParam()->mhq_quant.resize(4); + for (int i = 0; i < 4; ++i) { + op->main.AsAttentionParam()->mhq_quant[i].reset(new TensorQuantInfoT); + op->main.AsAttentionParam()->mhq_quant[i]->scale = 0.0f; + } + auto& mhqQuant = op->main.AsAttentionParam()->mhq_quant; + if (nullptr != qk_scale_q) { + mhqQuant[0]->scale = qk_scale_q->readMap()[0]; + } + if (nullptr != qk_scale_k) { + mhqQuant[1]->scale = qk_scale_k->readMap()[0]; + } + if (nullptr != sv_scale_s) { + mhqQuant[2]->scale = sv_scale_s->readMap()[0]; + } + if (nullptr != sv_scale_v) { + mhqQuant[3]->scale = sv_scale_v->readMap()[0]; + } + } + VARP output; + if (nullptr != mask.get()) { + output = Variable::create(Expr::create(op.get(), {q, k, v, mask})); + } else { + output = Variable::create(Expr::create(op.get(), {q, k, v})); + } + if (!supportC4) { + output = _Reshape(output, {-1, numHead * headDim, 1, 1}); + } + return output; +} + +static void _splitBufToArray(const uint8_t* buf, uint8_t* arr, size_t arrLen, size_t needBits) { + unsigned char mask = (1 << needBits) - 1; + unsigned char* tmp = (unsigned char*)buf; + int offset = 0; + for (size_t i = 0; i < arrLen; ++i) { + unsigned char idx = 0; + long shift = 8 - needBits - offset % 8; + if (shift < 0) { + idx = (tmp[offset / 8] << (0 - shift)) & mask; + idx |= (tmp[(offset / 8) + 1] >> (8 + shift)) & mask; + } else { + idx = (tmp[offset / 8] >> shift) & mask; + } + offset += needBits; + if (offset % 8 == 0) { + tmp += offset / 8; + offset = 0; + } + arr[i] = idx; + } +} + +VARP _QConvolution1x1(int inputCount, VARP input, VARP inputScale, VARP inputZero, VARP weight, VARP wscale, + VARP wzeropoint, VARP bias, int outputCount, bool scaleInputCount, int weightBit) { + std::unique_ptr conv(new OpT); + conv->type = OpType_Convolution; + conv->main.type = OpParameter_Convolution2D; + conv->main.value = new Convolution2DT; + auto parm = conv->main.AsConvolution2D(); + parm->common.reset(new Convolution2DCommonT); + if (outputCount > 0) { + parm->common->outputCount = outputCount; + } else { + parm->common->outputCount = (int)bias->getInfo()->size; + outputCount = parm->common->outputCount; + } + auto weightSize = weight->getInfo()->size; + auto weightInputCount = weightSize / parm->common->outputCount; + if (0 == weightBit) { + weightBit = 8 * (int)weightInputCount / inputCount; + } + MNN_ASSERT(weightBit <= 8); + if (nullptr == wscale.get()) { + parm->weight.resize(weightSize); + auto ptr = weight->readMap(); + if (nullptr == ptr) { + MNN_ERROR("_QConvolution1x1: weight->readMap() is nullptr!\n"); + return nullptr; + } + ::memcpy(parm->weight.data(), ptr, weightSize * sizeof(float)); + parm->common->inputCount = inputCount; + parm->bias.resize(parm->common->outputCount); + if (nullptr != bias) { + auto bptr = bias->readMap(); + if (nullptr == bptr) { + MNN_ERROR("_QConvolution1x1: bias->readMap() is nullptr!\n"); + return nullptr; + } + ::memcpy(parm->bias.data(), bptr, bias->getInfo()->size * sizeof(float)); + } else { + ::memset(parm->bias.data(), 0, parm->bias.size() * sizeof(float)); + } + return Variable::create(Expr::create(conv.get(), {input})); + } + + if (scaleInputCount) { + std::vector scales(inputCount); + auto scalePtr = wscale->readMap(); + if (nullptr == scalePtr) { + MNN_ERROR("_QConvolution1x1: wscale->readMap() is nullptr!\n"); + return nullptr; + } + ::memcpy(scales.data(), scalePtr, inputCount * sizeof(float)); + std::vector emptyBias; + input = _Scale(input, inputCount, std::move(scales), std::move(emptyBias)); + wscale = _Const(1.0f, {outputCount}, NCHW); + } + if (wscale->getInfo()->size == 1 && parm->common->outputCount > 1) { + auto scalePtr = wscale->readMap(); + if (nullptr == scalePtr) { + MNN_ERROR("_QConvolution1x1: scalar wscale->readMap() is nullptr!\n"); + return nullptr; + } + std::vector scales(parm->common->outputCount, scalePtr[0]); + wscale = _Const(scales.data(), {parm->common->outputCount}, NCHW); + } + auto scaleSize = wscale->getInfo()->size; + if (parm->common->outputCount > scaleSize) { + MNN_ERROR("scaleSize %zu <= outputCount %d\n", scaleSize, parm->common->outputCount); + return nullptr; + } + + parm->common->inputCount = inputCount; + parm->bias.resize(parm->common->outputCount); + if (nullptr != bias) { + auto bptr = bias->readMap(); + if (nullptr == bptr) { + MNN_ERROR("_QConvolution1x1 quant: bias->readMap() is nullptr!\n"); + return nullptr; + } + ::memcpy(parm->bias.data(), bptr, bias->getInfo()->size * sizeof(float)); + } else { + ::memset(parm->bias.data(), 0, parm->bias.size() * sizeof(float)); + } + if (nullptr != inputScale) { + auto scale = inputScale->readMap()[0]; + float zeroPoint = 0.0f; + if (nullptr != inputZero) { + zeroPoint = inputZero->readMap()[0]; + } + input->writeScaleMap(scale, zeroPoint); + } + + int n = parm->common->outputCount; + int k = parm->common->inputCount; + std::vector weightInt8(n * k); + if (4 == weightBit) { + int kDiv8 = k / 8; + auto weightSrcInt8 = weight->readMap(); + for (int i = 0; i < kDiv8; ++i) { + for (int u = 0; u < 4; ++u) { + for (int v = 0; v < n; ++v) { + auto packed = weightSrcInt8[(i * 4 + u) * n + v]; + int8_t item1 = packed >> 4; + int8_t item0 = packed - item1 * 16; + if (item0 >= 8) { + item0 -= 16; + } + MNN_ASSERT(item1 <= 7 && item1 >= -8); + weightInt8[v * k + i * 8 + u] = item0; + weightInt8[v * k + i * 8 + u + 4] = item1; + } + } + } + } else if (weightBit == 8) { + ::memcpy(weightInt8.data(), weight->readMap(), n * k); + } else { + auto weightSrcUInt8 = weight->readMap(); + auto weightUInt8 = (uint8_t*)weightInt8.data(); + _splitBufToArray(weightSrcUInt8, weightUInt8, n * k, weightBit); + int offset = 1 << (weightBit - 1); + for (int i = 0; i < n * k; ++i) { + weightInt8[i] = (int)weightUInt8[i] - offset; + } + } + + int dstWeightBit = weightBit; + if (8 == weightBit) { + int maxV = -256; + int minV = 256; + for (int v = 0; v < n * k; ++v) { + auto q = weightInt8[v]; + if (q > maxV) { + maxV = q; + } + if (q < minV) { + minV = q; + } + } + int targetBit = 0; + if (maxV >= 0) { + targetBit = (int)ceil(log(maxV + 1) / log(2)) + 1; + } + if (minV < 0) { + auto d1 = (int)ceil(log(-minV) / log(2)) + 1; + if (d1 > targetBit) { + targetBit = d1; + } + } + dstWeightBit = targetBit; + } + if (dstWeightBit > 4) { + dstWeightBit = 8; + } else if (dstWeightBit > 1) { + dstWeightBit = 4; + } else { + dstWeightBit = 1; + } + + std::vector scale; + bool async = false; + if (nullptr != wzeropoint) { + if (_ReduceMax(_Abs(_Cast(wzeropoint)))->readMap()[0] >= 1e-11f) { + async = true; + } + } + if (async) { + scale.resize(2 * scaleSize); + if (wzeropoint->getInfo()->type.code == halide_type_float) { + auto zeroPoint = wzeropoint->readMap(); + auto scalePtr = wscale->readMap(); + for (int i = 0; i < scaleSize; ++i) { + scale[2 * i + 1] = scalePtr[i]; + scale[2 * i + 0] = zeroPoint[i]; + } + } else { + auto zeroPoint = wzeropoint->readMap(); + auto scalePtr = wscale->readMap(); + for (int i = 0; i < scaleSize; ++i) { + scale[2 * i + 1] = scalePtr[i]; + scale[2 * i + 0] = -scalePtr[i] * zeroPoint[i]; + } + } + } else { + scale.resize(scaleSize); + ::memcpy(scale.data(), wscale->readMap(), scaleSize * sizeof(float)); + } + auto kernelSize = n * k / scaleSize; + parm->quanParameter = IDSTEncoder::encode(nullptr, scale, kernelSize, scaleSize, async, weightInt8.data(), 1, + {dstWeightBit, false}); + return Variable::create(Expr::create(conv.get(), {input})); +} + +static std::unique_ptr _makeLayerNorm(const LayerNormInfo& info) { + auto inputDim = info.hiddenSize; + if (0 == inputDim) { + if (nullptr != info.inputLayerNormWeight && nullptr != info.inputLayerNormWeight->getInfo()) { + inputDim = (int)info.inputLayerNormWeight->getInfo()->size; + } else { + MNN_ERROR("_TransformerLayerNorm: hiddenSize is 0 and inputLayerNormWeight is missing!\n"); + } + } + std::unique_ptr layerNorm(new OpT); + layerNorm->type = OpType_LayerNorm; + layerNorm->main.value = new LayerNormT; + layerNorm->main.type = OpParameter_LayerNorm; + layerNorm->main.AsLayerNorm()->axis = {-1}; + layerNorm->main.AsLayerNorm()->group = 1; + layerNorm->main.AsLayerNorm()->epsilon = info.ln_eps; + layerNorm->main.AsLayerNorm()->useRMSNorm = info.useRMSNorm; + if (info.useC4) { + layerNorm->defaultDimentionFormat = MNN_DATA_FORMAT_NC4HW4; + } + if (nullptr != info.inputLayerNormWeight) { + layerNorm->main.AsLayerNorm()->beta.resize(inputDim); + layerNorm->main.AsLayerNorm()->gamma.resize(inputDim); + if (nullptr != info.inputLayerNormBias) { + ::memcpy(layerNorm->main.AsLayerNorm()->beta.data(), info.inputLayerNormBias->readMap(), + inputDim * sizeof(float)); + } else { + ::memset(layerNorm->main.AsLayerNorm()->beta.data(), 0, inputDim * sizeof(float)); + } + ::memcpy(layerNorm->main.AsLayerNorm()->gamma.data(), info.inputLayerNormWeight->readMap(), + inputDim * sizeof(float)); + } + return layerNorm; +} + +std::pair _BinaryLayerNorm(VARP r0, VARP r1, const LayerNormInfo& info) { + std::unique_ptr layerNorm = _makeLayerNorm(info); + auto expr = Expr::create(layerNorm.get(), {r0, r1}, 2); + return {Variable::create(expr, 0), Variable::create(expr, 1)}; +} + +VARP _TransformerLayerNorm(VARP hiddenState, const LayerNormInfo& info) { + std::unique_ptr layerNorm = _makeLayerNorm(info); + return Variable::create(Expr::create(layerNorm.get(), {hiddenState})); +} + +static void _fillRopeTable(float* dst, const std::vector& cosTable, const std::vector& sinTable, + int end, int halfDim) { + const int tableSize = end * halfDim; + for (int t = 0; t < end; ++t) { + for (int i = 0; i < halfDim; ++i) { + const int evenIndex = (2 * i) % halfDim; + const int oddIndex = (2 * i + 1) % halfDim; + const int srcEven = t * halfDim + evenIndex; + const int srcOdd = t * halfDim + oddIndex; + const int dstIndex = t * halfDim + i; + dst[dstIndex] = cosTable[srcEven]; + dst[tableSize + dstIndex] = cosTable[srcOdd]; + dst[2 * tableSize + dstIndex] = sinTable[srcEven]; + dst[3 * tableSize + dstIndex] = sinTable[srcOdd]; + } + } +} + +VARP _PrecomputePosEmbedding(int dim, int end, float theta, bool interleaved) { + if (dim % 2 != 0 || end <= 0 || theta <= 0.0f) { + return nullptr; + } + + const int halfDim = dim / 2; + const int tableSize = end * halfDim; + std::vector cosTable(tableSize); + std::vector sinTable(tableSize); + for (int t = 0; t < end; ++t) { + for (int i = 0; i < halfDim; ++i) { + const float exponent = static_cast(2 * i) / static_cast(dim); + const float invFreq = 1.0f / std::pow(theta, exponent); + const float freq = static_cast(t) * invFreq; + const int offset = t * halfDim + i; + cosTable[offset] = std::cos(freq); + sinTable[offset] = std::sin(freq); + } + } + if (!interleaved) { + std::vector freqsCis(2 * tableSize); + ::memcpy(freqsCis.data(), cosTable.data(), tableSize * sizeof(float)); + ::memcpy(freqsCis.data() + tableSize, sinTable.data(), tableSize * sizeof(float)); + auto res = _Const(freqsCis.data(), {2, end, halfDim}, NCHW, halide_type_of()); + res.fix(VARP::CONSTANT); + return res; + } + + std::vector ropeTables(4 * tableSize); + _fillRopeTable(ropeTables.data(), cosTable, sinTable, end, halfDim); + auto res = _Const(ropeTables.data(), {4, end, halfDim}, NCHW, halide_type_of()); + res.fix(VARP::CONSTANT); + return res; +} + +VARPS _TransformerRoPE(VARP q, VARP k, VARP cosEven, VARP cosOdd, VARP sinEven, VARP sinOdd, const RopeInfo& info) { + std::unique_ptr qnorm; + std::unique_ptr knorm; + if (nullptr != info.qNorm.inputLayerNormWeight.get()) { + qnorm = _makeLayerNorm(info.qNorm); + } + if (nullptr != info.kNorm.inputLayerNormWeight.get()) { + knorm = _makeLayerNorm(info.kNorm); + } + + std::unique_ptr ropeOp(new OpT); + ropeOp->type = OpType_RoPE; + ExtraT* extra = nullptr; + if (info.cutHeadDim > 0 || nullptr != qnorm || nullptr != knorm) { + ropeOp->main.type = OpParameter_Extra; + extra = new ExtraT; + extra->type = "RoPE"; + extra->engine = "MNN"; + ropeOp->main.value = extra; + } + if (nullptr != qnorm) { + std::unique_ptr attr(new AttributeT); + flatbuffers::FlatBufferBuilder builder; + builder.Finish(Op::Pack(builder, qnorm.get())); + attr->key = "q_norm"; + attr->tensor.reset(new BlobT); + attr->tensor->dataType = DataType_DT_INT8; + attr->tensor->int8s.resize(builder.GetSize()); + ::memcpy(attr->tensor->int8s.data(), builder.GetBufferPointer(), builder.GetSize()); + extra->attr.emplace_back(std::move(attr)); + } + if (nullptr != knorm) { + std::unique_ptr attr(new AttributeT); + flatbuffers::FlatBufferBuilder builder; + builder.Finish(Op::Pack(builder, knorm.get())); + attr->key = "k_norm"; + attr->tensor.reset(new BlobT); + attr->tensor->dataType = DataType_DT_INT8; + attr->tensor->int8s.resize(builder.GetSize()); + ::memcpy(attr->tensor->int8s.data(), builder.GetBufferPointer(), builder.GetSize()); + extra->attr.emplace_back(std::move(attr)); + } + if (info.cutHeadDim > 0) { + std::unique_ptr attr(new AttributeT); + attr->key = "rope_cut_head_dim"; + attr->i = info.cutHeadDim; + extra->attr.emplace_back(std::move(attr)); + } + auto ropeExpr = Expr::create(ropeOp.get(), {q, k, cosEven, cosOdd, sinEven, sinOdd}, 2); + return {Variable::create(ropeExpr, 0), Variable::create(ropeExpr, 1)}; +} + +} // namespace SafeTensorUtils +} // namespace Express +} // namespace MNN diff --git a/tools/converter/source/safetensors/SafetensorUtils.hpp b/tools/converter/source/safetensors/SafetensorUtils.hpp new file mode 100644 index 0000000000..a400ecfc96 --- /dev/null +++ b/tools/converter/source/safetensors/SafetensorUtils.hpp @@ -0,0 +1,50 @@ +#ifndef SafetensorUtils_hpp +#define SafetensorUtils_hpp + +#include + +#include +#include + +namespace MNN { +namespace Express { +namespace SafeTensorUtils { + +struct LayerNormInfo { + VARP inputLayerNormWeight; + VARP inputLayerNormBias; + float ln_eps = 0.0f; + bool useRMSNorm = false; + int hiddenSize = 0; + bool useC4 = false; + + LayerNormInfo() = default; + LayerNormInfo(VARP weight, VARP bias, float eps, bool rms, int hidden = 0, bool c4 = false) + : inputLayerNormWeight(weight), inputLayerNormBias(bias), ln_eps(eps), useRMSNorm(rms), hiddenSize(hidden), useC4(c4) { + } +}; + +struct RopeInfo { + LayerNormInfo qNorm; + LayerNormInfo kNorm; + int cutHeadDim = 0; +}; + +MNN_PUBLIC VARP _QConvolution1x1(int inputCount, VARP input, VARP inputScale, VARP inputZero, VARP weight, + VARP wscale, VARP wzeropoint, VARP bias, int outputcount = 0, + bool scaleInputCount = false, int weightBits = 0); +MNN_PUBLIC VARP _TransformerLayerNorm(VARP hiddenState, const LayerNormInfo& info); +MNN_PUBLIC std::pair _BinaryLayerNorm(VARP r0, VARP r1, const LayerNormInfo& info); +MNN_PUBLIC VARP _GPT2Attention(int numHead, int headDim, VARP q, VARP k, VARP v, VARP qk_scale_q, VARP qk_scale_k, + VARP sv_scale_s, VARP sv_scale_v, VARP mask, bool supportC4Opt = false, + float attnScale = 0.0f); +MNN_PUBLIC VARP _PrecomputePosEmbedding(int dim, int end, float theta = 1000000.0f, bool interleaved = false); +MNN_PUBLIC VARPS _TransformerRoPE(VARP q, VARP k, VARP cosEven, VARP cosOdd, VARP sinEven, VARP sinOdd, + const RopeInfo& info); +MNN_PUBLIC VARP _MakeLastHiddenStateOutput(VARP hiddenState, int hiddenSize); + +} // namespace SafeTensorUtils +} // namespace Express +} // namespace MNN + +#endif diff --git a/tools/converter/source/safetensors/WorkflowJson.hpp b/tools/converter/source/safetensors/WorkflowJson.hpp new file mode 100644 index 0000000000..6d54d9e120 --- /dev/null +++ b/tools/converter/source/safetensors/WorkflowJson.hpp @@ -0,0 +1,124 @@ +#ifndef WorkflowJson_hpp +#define WorkflowJson_hpp + +#include + +#include + +namespace MNN { +namespace SafeTensors { +namespace WorkflowJson { + +inline const rapidjson::Value* _findMember(const rapidjson::Value& obj, const char* key) { + if (!obj.IsObject() || nullptr == key) { + return nullptr; + } + auto it = obj.FindMember(key); + if (it == obj.MemberEnd()) { + return nullptr; + } + return &it->value; +} + +inline std::string getString(const rapidjson::Value& obj, const char* key, const std::string& defaultValue = "") { + auto v = _findMember(obj, key); + if (nullptr == v || !v->IsString()) { + return defaultValue; + } + return v->GetString(); +} + +inline bool getBool(const rapidjson::Value& obj, const char* key, bool defaultValue = false) { + auto v = _findMember(obj, key); + if (nullptr == v) { + return defaultValue; + } + return v->GetBool(); +} +inline int getInt(const rapidjson::Value& obj, const char* key, int defaultValue = 0) { + auto v = _findMember(obj, key); + if (nullptr == v || !v->IsInt()) { + return defaultValue; + } + return v->GetInt(); +} + +inline float getFloat(const rapidjson::Value& obj, const char* key, float defaultValue = 0.0f) { + auto v = _findMember(obj, key); + if (nullptr == v) { + return defaultValue; + } + if (v->IsFloat()) { + return v->GetFloat(); + } + if (v->IsDouble()) { + return static_cast(v->GetDouble()); + } + return defaultValue; +} + +inline const rapidjson::Value* getArray(const rapidjson::Value& obj, const char* key) { + auto v = _findMember(obj, key); + if (nullptr == v || !v->IsArray()) { + return nullptr; + } + return v; +} + +inline bool firstArrayStringEquals(const rapidjson::Value& obj, const char* key, const char* expected) { + if (nullptr == expected) { + return false; + } + auto v = getArray(obj, key); + if (nullptr == v || v->Empty()) { + return false; + } + auto& first = (*v)[0]; + if (!first.IsString()) { + return false; + } + return first.GetString() == std::string(expected); +} + +inline bool arrayStringContains(const rapidjson::Value& obj, const char* key, const char* expected) { + if (nullptr == expected) { + return false; + } + auto v = getArray(obj, key); + if (nullptr == v) { + return false; + } + const std::string target(expected); + for (auto& item : v->GetArray()) { + if (item.IsString() && item.GetString() == target) { + return true; + } + } + return false; +} + +inline const rapidjson::Value* findFirstBlockByType(const rapidjson::Value& model, const char* type) { + if (nullptr == type) { + return nullptr; + } + auto blocks = getArray(model, "blocks"); + if (nullptr == blocks) { + return nullptr; + } + for (auto& item : blocks->GetArray()) { + if (!item.IsObject()) { + continue; + } + auto t = _findMember(item, "type"); + if (nullptr != t && t->IsString() && t->GetString() == std::string(type)) { + return &item; + } + } + return nullptr; +} + +} // namespace WorkflowJson +} // namespace SafeTensors +} // namespace MNN + +#endif diff --git a/tools/converter/source/safetensors/safetensors.hh b/tools/converter/source/safetensors/safetensors.hh new file mode 100644 index 0000000000..e0add85f6e --- /dev/null +++ b/tools/converter/source/safetensors/safetensors.hh @@ -0,0 +1,4865 @@ +// SPDX-License-Identifier: MIT +// Copyright 2023 - Present, Syoyo Fujita. +// Inspired from: +// https://gist.github.com/Narsil/5d6bf307995158ad2c4994f323967284 +#pragma once + +#include +#include +#include +#include +#include + +#ifdef __ANDROID__ +#ifdef SAFETENSORS_CPP_ANDROID_LOAD_FROM_ASSETS +#include +#endif + +#ifdef SAFETENSORS_CPP_IMPLEMENTATION +AAssetManager *asset_manager = nullptr; +#else +extern AAssetManager *asset_manager; +#endif +#endif + + +namespace safetensors { + +constexpr size_t kMaxDim = + 8; // must be equal to SAFETENSORS_C_MAX_DIM in `safetensors-c.h` + +enum dtype { + kBOOL, + kUINT8, + kINT8, + kINT16, + kUINT16, + kFLOAT16, + kBFLOAT16, + kINT32, + kUINT32, + kFLOAT32, + kFLOAT64, + kINT64, + kUINT64, +}; + +namespace minijson { + +// Simple C++ implementation of Python's OrderedDict like dictonary +// (preserves key insertion order) +// Modified for JSON: +// - No duplicated key allowed + +template +class ordered_dict { + public: + bool at(const size_t idx, T *dst) const { + if (idx >= _keys.size()) { + return false; + } + + if (!_m.count(_keys[idx])) { + // This should not happen though. + return false; + } + + (*dst) = _m.at(_keys[idx]); + + return true; + } + + bool count(const std::string &key) const { return _m.count(key); } + + void insert(const std::string &key, const T &value) { + if (_m.count(key)) { + // overwrite existing value + } else { + _keys.push_back(key); + } + + _m[key] = value; + } + + void insert(const std::string &key, T &&value) { + if (_m.count(key)) { + // overwrite existing value + } else { + _keys.push_back(key); + } + + _m[key] = std::move(value); + } + + bool at(const std::string &key, T *dst) const { + if (!_m.count(key)) { + // This should not happen though. + return false; + } + + (*dst) = _m.at(key); + + return true; + } + + const std::vector &keys() const { return _keys; } + + size_t size() const { return _m.size(); } + + bool erase(const std::string &key) { + // simple linear search + for (size_t i = 0; i < _keys.size(); i++) { + if (_keys[i] == key) { + _keys.erase(_keys.begin() + i); + _m.erase(key); + return true; + } + } + + return false; + } + + private: + std::vector _keys; + std::map _m; +}; + +} // namespace minijson + +template +using ordered_dict = minijson::ordered_dict; + +struct tensor_t { + safetensors::dtype dtype; + std::vector shape; + std::array data_offsets; +}; + +struct safetensors_t { + // we need ordered dict(preserves the order of key insertion) + // as done in Python's OrderedDict, since JSON data may not be sorted by its key string. + ordered_dict tensors; + ordered_dict metadata; + std::vector storage; // empty when mmap'ed + size_t header_size{0}; // JSON size + + bool mmaped{false}; + + // + // Following members are set when mmaped. + // + const uint8_t *mmap_addr{nullptr}; + size_t mmap_size{0}; + const uint8_t *databuffer_addr{nullptr}; // [mmap_addr + header_size + 8] + size_t databuffer_size{0}; // mmap_size - header_size - 8 + // opaque pointer to safetensors_file and safetensors_mmap + void *st_file{nullptr}; + void *st_mmap{nullptr}; + + ~safetensors_t(); +}; + +// +// Load safetensors from file. +// databuffer is copied to `safetensors_t::storage`. +// +// @param[in] filename Filepath. Assume UTF-8 filepath. +// @param[out] st safetensors data. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +bool load_from_file(const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err); + +// +// Load safetensors data from memory. +// databuffer is copied to `safetensors_t::storage`. +// +// @param[in] addr Memory address of safetensors data. +// @param[in] nbytes The size in bytes. +// @param[in] filename Filename of corresponding memory data. Can be empty. +// @param[out] st safetensors data. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +// +bool load_from_memory(const uint8_t *addr, const size_t nbytes, + const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err); + +// +// Load safetensors with memory mapping(i.e. zero-copy). +// databuffer is not copied to `safetensors_t` object, thus the app must hold +// file during `safetensor_t` object is live. +// +// @param[in] filename Filepath. Assume UTF-8 filepath. +// @param[out] st safetensors data. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +bool mmap_from_file(const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err); + +// +// Load safetensors from mmaped region. +// databuffer is not copied to `safetensors_t` object, thus the app must not +// free/unmap `addr` during `safetensor_t` object is live. +// +// @param[in] addr mmaped memory address of safetensors data. +// @param[in] nbytes mmap bytes. +// @param[in] filename Filename of corresponding memory data. Can be empty. +// @param[out] st safetensors data. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +bool mmap_from_memory(const uint8_t *arr, const size_t nbytes, + const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err); + +// +// Save safetensors to file. +// +// @param[in] st safetensors data. +// @param[in] filename Filepath. Assume UTF-8 filepath. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +bool save_to_file(const safetensors_t &st, const std::string &filename, + std::string *warn, std::string *err); + +// +// Save safetensors to memory. +// +// @param[in] st safetensors data. +// @param[out] data_out Serialized safetensor data. +// @param[out] warn Warning message buffer(can be nullptr if you don't need +// warning message) +// @param[out] err Error message buffer(can be nullptr if you don't need error +// message) +// +// @return true upon success. `err` will be filled when false. +bool save_to_memory(const std::string &filename, std::vector *data_out, + std::string *warn, std::string *err); + +// +// Utility functions +// + +// Returns shape[0] * shape[1] * ... +// Empty Tensor(any shape[i] is 0) returns 0. +// Zero-rank tensor([]) return 1. +size_t get_shape_size(const tensor_t &t); + +// Returns dtype size in bytes. +size_t get_dtype_bytes(const safetensors::dtype dtype); +std::string get_dtype_str(const safetensors::dtype dtype); + +// Validate data_offsets of all tensors in safetensors_t. +bool validate_data_offsets(const safetensors_t &st, std::string &err); + +uint16_t float_to_bfloat16(float x); +float bfloat16_to_float(uint16_t x); + +uint16_t float_to_fp16(float x); +float fp16_to_float(uint16_t x); + +} // namespace safetensors + +#if defined(SAFETENSORS_CPP_IMPLEMENTATION) + +#include +#include +#include + +#ifdef __has_include +#if __has_include() +#include +#if defined(_POSIX_MAPPED_FILES) +#include +#endif +#if defined(_POSIX_MEMLOCK_RANGE) +#include +#endif +#endif +#endif + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#include // for _fseeki64 +#include +#endif + +#if !defined(MINIJSON_IMPLEMENTATION) +#define MINIJSON_IMPLEMENTATION +#endif + +// minijson: https://github.com/syoyo/minijson + +/* + * JSON parser: C++ oriented JSON parser. + */ + +#include +#include +#include +#include +#include + +//#define __MINIJSON_LIBERAL + +// We recommended to use simdjson from_chars. +// Using strtod() is a fallback +#if defined(MINIJSON_USE_STRTOD) +// Use stdlib's strtod +#include +#else + +namespace minijson { +namespace simdjson { +namespace internal { + +double from_chars(const char *first) noexcept; +double from_chars(const char *first, const char *end) noexcept; + +char *to_chars(char *first, const char *last, double value); + +} // namespace internal +} // namespace simdjson +} // namesapce minijson + +#endif + +namespace minijson { + +namespace detail { + +double from_chars(const char *p); +const char *my_strchr(const char *p, int ch); + +} // namespace detail + +namespace detail { + +// +// Usage: +// - set_input() +// - scan_string() +// - success: use `token_buffer` string +// - error: use `error_message` +// +struct string_parser { + // input string must be UTF-8 + void set_input(const std::string &s) { _input = s; } + + bool scan_string(); + + void reset() { + if (_input.size()) { + current = _input[0]; + } else { + current = '\0'; + } + curr_idx = 0; + token_buffer.clear(); + } + + // fetch next token. + unsigned char get() { + if ((curr_idx + 1) < _input.size()) { + curr_idx++; + current = _input[curr_idx]; + return current; + } + current = '\0'; + return current; + } + + bool eof() { + if (_input.empty()) { + return true; + } + + if (curr_idx >= _input.size()) { + return true; + } + + return false; + } + + void add(const unsigned char c) { token_buffer += c; } + + void add(const int i) { + // use lower 8bit + token_buffer += static_cast(i & 0xff); + } + + int get_codepoint(); + + bool next_byte_in_range(const std::initializer_list ranges); + + std::string error_message; + std::string token_buffer; // output + + unsigned char current{'\0'}; + size_t curr_idx{0}; + std::string _input; +}; + +} // namespace detail + +typedef enum { + unknown_type, + null_type, + boolean_type, + number_type, + string_type, + array_type, + object_type, +} type; + +typedef enum { + no_error, + undefined_error, + invalid_token_error, + unknown_type_error, + memory_allocation_error, + corrupted_json_error, + duplicated_key_error, +} error; + +class value; + +typedef bool boolean; +typedef double number; +typedef std::string string; +typedef safetensors::ordered_dict object; +typedef std::vector array; +typedef struct { +} null_t; + +// null_t null; + +template +struct TypeTraits; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 0; } +}; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 1; } +}; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 2; } +}; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 3; } +}; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 4; } +}; + +template <> +struct TypeTraits { + static constexpr uint32_t type_id() { return 5; } +}; + +class value { + private: + type t; + union { + null_t n; + boolean b; + number d; + std::string *s; + array *a; + object *o; + } u; + + void _free_u() { + if (t == string_type) { + delete this->u.s; + this->u.s = nullptr; + } + if (t == array_type) { + delete this->u.a; + this->u.a = nullptr; + } + if (t == object_type) { + delete this->u.o; + this->u.o = nullptr; + } + } + + public: + value() : t(unknown_type), u() {} + value(null_t n) : t(null_type), u() { u.n = n; } + value(boolean b) : t(boolean_type), u() { u.b = b; } + value(number d) : t(boolean_type), u() { u.d = d; } + value(const char *s) : t(string_type), u() { u.s = new std::string(s); } + value(const std::string &s) : t(string_type), u() { + u.s = new std::string(s); + } + value(const array &a) : t(array_type), u() { u.a = new array(a); } + value(const object &o) : t(object_type), u() { u.o = new object(o); } + value(const value &v) : t(v.t), u() { + if (t == array_type) { + u.a = new array(); + *u.a = *v.u.a; + } else if (t == object_type) { + u.o = new object(); + *u.o = *v.u.o; + } else if (t == string_type) { + u.s = new std::string(); + *u.s = *v.u.s; + } else + u.d = v.u.d; + } + ~value() { _free_u(); } + + template + bool is() const { + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == null_type) + return true; + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == boolean_type) + return true; + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == number_type) + return true; + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == string_type) + return true; + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == array_type) + return true; + if (TypeTraits::type_id() == TypeTraits::type_id() && + t == object_type) + return true; + return false; + } + + template + const T *as() const { + if ((t == array_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.a); + } + + if ((t == object_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.o); + } + + if ((t == string_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.s); + } + + if ((t == null_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.n); + } + + if ((t == boolean_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.b); + } + + if ((t == number_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.d); + } + + return nullptr; + } + + template + T *as() { + if ((t == array_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.a); + } + + if ((t == object_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.o); + } + + if ((t == string_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(u.s); + } + + if ((t == null_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.n); + } + + if ((t == boolean_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.b); + } + + if ((t == number_type) && + (TypeTraits::type_id() == TypeTraits::type_id())) { + return reinterpret_cast(&u.d); + } + + return nullptr; + } + + null_t &operator=(null_t &n) { + t = null_type; + u.n = n; + return u.n; + } + boolean &operator=(boolean b) { + t = boolean_type; + u.b = b; + return u.b; + } + number &operator=(number d) { + t = number_type; + u.d = d; + return u.d; + } + const std::string &operator=(const char *s) { + _free_u(); + t = string_type; + u.s = new std::string(s); + return *u.s; + } + const std::string &operator=(const std::string &s) { + _free_u(); + t = string_type; + u.s = new std::string(s); + return *u.s; + } + const object &operator=(const object &o) { + _free_u(); + t = object_type; + u.o = new object(o); + return *u.o; + } + const array &operator=(const array &a) { + _free_u(); + t = array_type; + u.a = new array(a); + return *u.a; + } + const value &operator=(const value &v) { + _free_u(); + t = v.t; + if (t == array_type) { + u.a = new array(*v.u.a); + } else if (t == object_type) { + u.o = new object(*v.u.o); + } else if (t == string_type) { + u.s = new std::string(*v.u.s); + } else + u.d = v.u.d; + return *this; + } + + std::string type_name() const { + if (t == array_type) { + return "array"; + } + + if (t == object_type) { + return "object"; + } + + if (t == string_type) { + return "string"; + } + + if (t == null_type) { + return "null"; + } + + if (t == boolean_type) { + return "boolean"; + } + + if (t == number_type) { + return "number"; + } + + return "[[invalid]]"; + } + + std::string str(const char *p) const { + std::stringstream ss; + ss << '"'; + while (*p) { + if (*p == '\n') { + ss << "\\n"; + } else if (*p == '\r') { + ss << "\\r"; + } else if (*p == '\t') { + ss << "\\t"; + } else if (detail::my_strchr("\"", *p)) { + ss << "\\" << *p; + } else { + ss << *p; + } + p++; + } + ss << '"'; + return ss.str(); + } + + std::string str() const { + std::stringstream ss; + if (t == unknown_type) { + ss << "undefined"; + } else if (t == null_type) { + ss << "null"; + } else if (t == boolean_type) { + ss << (u.b ? "true" : "false"); + } else if (t == number_type) { + ss << double(u.d); + } else if (t == string_type) { + ss << str(u.s->c_str()); + } else if (const array *pa = as()) { + array::const_iterator i; + ss << "["; + // array a = get(); + for (i = pa->begin(); i != pa->end(); i++) { + if (i != pa->begin()) ss << ", "; + ss << i->str(); + } + ss << "]"; + } else if (auto po = as()) { + // object::const_iterator i; + ss << "{"; + // object o = get(); + for (size_t i = 0; i < po->size(); i++) { + if (i > 0) ss << ", "; + ss << "\"" << po->keys()[i] << "\""; + + value v; + if (po->at(i, &v)) { + ss << ": " << v.str(); + } else { + // TODO: report error + ss << ": null"; + } + } + ss << "}"; + } + return ss.str(); + } +}; + +#define MINIJSON_SKIP(i) \ + while (*i && detail::my_strchr("\r\n \t", *i)) { \ + i++; \ + } + +template +inline error parse_object(Iter &i, value &v) { + object o; + i++; + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } + if (*i != '\x7d') { + while (*i) { + value vk, vv; + error e = parse_string(i, vk); + if (e != no_error) return e; + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } + if (*i != ':') return invalid_token_error; + i++; + e = parse_any(i, vv); + if (e != no_error) return e; + + auto ps = vk.as(); + if (!ps) { + return unknown_type_error; + } + + if (o.count(*ps)) { + return duplicated_key_error; + } + o.insert(*ps, vv); + + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } + if (*i == '\x7d') break; + if (*i != ',') return invalid_token_error; + i++; + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } +#ifdef __MINIJSON_LIBERAL + if (*i == '\x7d') break; +#endif + } + } + v = value(o); + i++; + return no_error; +} + +template +inline error parse_array(Iter &i, value &v) { + array a; + i++; + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } + if (*i != ']') { + while (*i) { + value va; + error e = parse_any(i, va); + if (e != no_error) return e; + a.push_back(va); + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } + if (*i == ']') break; + if (*i != ',') return invalid_token_error; + i++; + MINIJSON_SKIP(i) + if (!(*i)) { + return corrupted_json_error; + } +#ifdef __MINIJSON_LIBERAL + if (*i == '\x7d') break; +#endif + } + } + v = value(a); + i++; + return no_error; +} + +template +inline error parse_null(Iter &i, value &v) { + Iter p = i; + if (*i == 'n' && *(i + 1) == 'u' && *(i + 2) == 'l' && *(i + 3) == 'l') { + i += 4; + v = null_t(); + } + if (*i && nullptr == detail::my_strchr(":,\x7d]\r\n ", *i)) { + i = p; + return undefined_error; + } + return no_error; +} + +template +inline error parse_boolean(Iter &i, value &v) { + Iter p = i; + if (*i == 't' && *(i + 1) == 'r' && *(i + 2) == 'u' && *(i + 3) == 'e') { + i += 4; + v = static_cast(true); + } else if (*i == 'f' && *(i + 1) == 'a' && *(i + 2) == 'l' && + *(i + 3) == 's' && *(i + 4) == 'e') { + i += 5; + v = static_cast(false); + } + if (*i && nullptr == detail::my_strchr(":,\x7d]\r\n ", *i)) { + i = p; + return undefined_error; + } + return no_error; +} + +template +inline error parse_number(Iter &i, value &v) { + Iter p = i; + + if (*i == '-') { + i++; + } + +#define MINIJSON_IS_NUM(x) ('0' <= x && x <= '9') +#define MINIJSON_IS_ALNUM(x) \ + (('0' <= x && x <= '9') || ('a' <= x && x <= 'f') || ('A' <= x && x <= 'F')) + if (*i == '0' && *(i + 1) == 'x' && MINIJSON_IS_ALNUM(*(i + 2))) { + i += 3; + while (MINIJSON_IS_ALNUM(*i)) i++; + v = static_cast(detail::from_chars(p)); + } else { + while (MINIJSON_IS_NUM(*i)) i++; + if (*i == '.') { + i++; + if (!MINIJSON_IS_NUM(*i)) { + i = p; + return invalid_token_error; + } + while (MINIJSON_IS_NUM(*i)) i++; + } + if (*i == 'e') { + i++; + if (!MINIJSON_IS_NUM(*i)) { + i = p; + return invalid_token_error; + } + while (MINIJSON_IS_NUM(*i)) i++; + } + v = static_cast(detail::from_chars(p)); + } + if (*i && nullptr == detail::my_strchr(":,\x7d]\r\n ", *i)) { + i = p; + return invalid_token_error; + } + return no_error; +} + +template +inline error parse_string(Iter &i, value &v) { + if (*i != '"') return invalid_token_error; + + Iter s = i; + char t = *i++; // = '"' + Iter p = i; + +#if 0 + std::stringstream ss; + while (*i && *i != t) { + if (*i == '\\' && *(i + 1)) { + i++; + if (*i == 'n') + ss << "\n"; + else if (*i == 'r') + ss << "\r"; + else if (*i == 't') + ss << "\t"; + else + ss << *i; + } else { + ss << *i; + } + i++; + } +#else + // read until '"' + while (*i && *i != t) { + if (*i == '\\' && *(i + 1)) { + i++; + } + i++; + } + +#endif + if (!*i) return invalid_token_error; + if (i < p) { + return corrupted_json_error; + } + +#if 0 + v = std::string(p, size_t(i - p)); + + i++; + if (*i && nullptr == detail::my_strchr(":,\x7d]\r\n ", *i)) { + i = p; + return invalid_token_error; + } + +#else + + i++; + if (*i && nullptr == detail::my_strchr(":,\x7d]\r\n ", *i)) { + i = p; + return invalid_token_error; + } + + // include first and last '"' char + std::string buf(s, size_t(i - s)); + + detail::string_parser str_parser; + str_parser.set_input(buf); + + if (!str_parser.scan_string()) { + // TODO: error message + // str_parser.error_message; + return invalid_token_error; + } else { + v = str_parser.token_buffer; + } + +#endif + + return no_error; +} + +template +inline error parse_any(Iter &i, value &v) { + MINIJSON_SKIP(i) + if (*i == '\x7b') return parse_object(i, v); + if (*i == '[') return parse_array(i, v); + if (*i == 't' || *i == 'f') return parse_boolean(i, v); + if (*i == 'n') return parse_null(i, v); + if ((*i == '-') || ('0' <= *i && *i <= '9')) return parse_number(i, v); + if (*i == '"') return parse_string(i, v); + return invalid_token_error; +} + +template +inline error parse(Iter &i, value &v) { + return parse_any(i, v); +} + +#undef MINIJSON_SKIP + +inline const char *errstr(error e) { + const char *s = "unknown error"; + switch (e) { + case no_error: { + s = "no error"; + break; + } + case undefined_error: { + s = "undefined"; + break; + } + case invalid_token_error: { + s = "invalid token"; + break; + } + case unknown_type_error: { + s = "unknown type"; + break; + } + case memory_allocation_error: { + s = "memory allocation error"; + break; + } + case corrupted_json_error: { + s = "input is corrupted"; + break; + } + case duplicated_key_error: { + s = "duplicated key found"; + break; + } + // default: return "unknown error"; + } + + return s; +} + +} // namespace minijson + +#if defined(MINIJSON_IMPLEMENTATION) + +namespace minijson { + +namespace detail { + +// clang-format off +// +// From json.hpp --------------------------------------------------------- +// __ _____ _____ _____ +// __| | __| | | | JSON for Modern C++ +// | | |__ | | | | | | version 3.11.3 +// |_____|_____|_____|_|___| https://github.com/nlohmann/json +// +// SPDX-FileCopyrightText: 2013-2023 Niels Lohmann +// SPDX-License-Identifier: MIT + +#if 1 + #define JSON_HEDLEY_UNLIKELY(cond) (cond) + #define JSON_HEDLEY_LIKELY(cond) (cond) + + /*! + @brief get codepoint from 4 hex characters following `\u` + + For input "\u c1 c2 c3 c4" the codepoint is: + (c1 * 0x1000) + (c2 * 0x0100) + (c3 * 0x0010) + c4 + = (c1 << 12) + (c2 << 8) + (c3 << 4) + (c4 << 0) + + Furthermore, the possible characters '0'..'9', 'A'..'F', and 'a'..'f' + must be converted to the integers 0x0..0x9, 0xA..0xF, 0xA..0xF, resp. The + conversion is done by subtracting the offset (0x30, 0x37, and 0x57) + between the ASCII value of the character and the desired integer value. + + @return codepoint (0x0000..0xFFFF) or -1 in case of an error (e.g. EOF or + non-hex character) + */ + int string_parser::get_codepoint() + { + // this function only makes sense after reading `\u` + //JSON_ASSERT(current == 'u'); + if (current != 'u') { + return -1; + } + int codepoint = 0; + + const auto factors = { 12u, 8u, 4u, 0u }; + for (const auto factor : factors) + { + get(); + + if (current >= '0' && current <= '9') + { + codepoint += static_cast((static_cast(current) - 0x30u) << factor); + } + else if (current >= 'A' && current <= 'F') + { + codepoint += static_cast((static_cast(current) - 0x37u) << factor); + } + else if (current >= 'a' && current <= 'f') + { + codepoint += static_cast((static_cast(current) - 0x57u) << factor); + } + else + { + return -1; + } + } + + if (0x0000 <= codepoint && codepoint <= 0xFFFF) { + } else { + return -1; + } + return codepoint; + } + + /*! + @brief check if the next byte(s) are inside a given range + + Adds the current byte and, for each passed range, reads a new byte and + checks if it is inside the range. If a violation was detected, set up an + error message and return false. Otherwise, return true. + + @param[in] ranges list of integers; interpreted as list of pairs of + inclusive lower and upper bound, respectively + + @pre The passed list @a ranges must have 2, 4, or 6 elements; that is, + 1, 2, or 3 pairs. This precondition is enforced by an assertion. + + @return true if and only if no range violation was detected + */ + bool string_parser::next_byte_in_range(const std::initializer_list ranges) + { + if (ranges.size() == 2 || ranges.size() == 4 || ranges.size() == 6) { + } else { + return false; + } + + add(current); + + for (auto range = ranges.begin(); range != ranges.end(); ++range) + { + get(); + if (JSON_HEDLEY_LIKELY(*range <= current && current <= *(++range))) // NOLINT(bugprone-inc-dec-in-conditions) + { + add(current); + } + else + { + error_message = "invalid string: ill-formed UTF-8 byte"; + return false; + } + } + + return true; + } + /*! + @brief scan a string literal + + This function scans a string according to Sect. 7 of RFC 8259. While + scanning, bytes are escaped and copied into buffer token_buffer. Then the + function returns successfully, token_buffer is *not* null-terminated (as it + may contain \0 bytes), and token_buffer.size() is the number of bytes in the + string. + + @return true if string could be successfully scanned, + false otherwise + + @note In case of errors, variable error_message contains a textual + description. + */ + bool string_parser::scan_string() + { + // reset token_buffer (ignore opening quote) + reset(); + + // we entered the function by reading an open quote + //JSON_ASSERT(current == '\"'); + if (current != '\"') { + error_message = "first character must be '\"'"; + return false; + } + + + while (!eof()) + { + + // get next character + switch (get()) + { + + // closing quote + case '\"': + { + return true; + } + + // escapes + case '\\': + { + switch (get()) + { + // quotation mark + case '\"': + add('\"'); + break; + // reverse solidus + case '\\': + add('\\'); + break; + // solidus + case '/': + add('/'); + break; + // backspace + case 'b': + add('\b'); + break; + // form feed + case 'f': + add('\f'); + break; + // line feed + case 'n': + add('\n'); + break; + // carriage return + case 'r': + add('\r'); + break; + // tab + case 't': + add('\t'); + break; + + // unicode escapes + case 'u': + { + const int codepoint1 = get_codepoint(); + int codepoint = codepoint1; // start with codepoint1 + + if (JSON_HEDLEY_UNLIKELY(codepoint1 == -1)) + { + error_message = "invalid string: '\\u' must be followed by 4 hex digits"; + return false; + } + + // check if code point is a high surrogate + if (0xD800 <= codepoint1 && codepoint1 <= 0xDBFF) + { + // expect next \uxxxx entry + if (JSON_HEDLEY_LIKELY(get() == '\\' && get() == 'u')) + { + const int codepoint2 = get_codepoint(); + + if (JSON_HEDLEY_UNLIKELY(codepoint2 == -1)) + { + error_message = "invalid string: '\\u' must be followed by 4 hex digits"; + return false; + } + + // check if codepoint2 is a low surrogate + if (JSON_HEDLEY_LIKELY(0xDC00 <= codepoint2 && codepoint2 <= 0xDFFF)) + { + // overwrite codepoint + codepoint = static_cast( + // high surrogate occupies the most significant 22 bits + (static_cast(codepoint1) << 10u) + // low surrogate occupies the least significant 15 bits + + static_cast(codepoint2) + // there is still the 0xD800, 0xDC00 and 0x10000 noise + // in the result, so we have to subtract with: + // (0xD800 << 10) + DC00 - 0x10000 = 0x35FDC00 + - 0x35FDC00u); + } + else + { + error_message = "invalid string: surrogate U+D800..U+DBFF must be followed by U+DC00..U+DFFF"; + return false; + } + } + else + { + error_message = "invalid string: surrogate U+D800..U+DBFF must be followed by U+DC00..U+DFFF"; + return false; + } + } + else + { + if (JSON_HEDLEY_UNLIKELY(0xDC00 <= codepoint1 && codepoint1 <= 0xDFFF)) + { + error_message = "invalid string: surrogate U+DC00..U+DFFF must follow U+D800..U+DBFF"; + return false; + } + } + + // result of the above calculation yields a proper codepoint + //JSON_ASSERT(0x00 <= codepoint && codepoint <= 0x10FFFF); + if (0x00 <= codepoint && codepoint <= 0x10FFFF) { + } else { + error_message = "invalid string: invalid codepoint"; + return false; + } + + // translate codepoint into bytes + if (codepoint < 0x80) + { + // 1-byte characters: 0xxxxxxx (ASCII) + add(static_cast(codepoint)); + } + else if (codepoint <= 0x7FF) + { + // 2-byte characters: 110xxxxx 10xxxxxx + add(static_cast(0xC0u | (static_cast(codepoint) >> 6u))); + add(static_cast(0x80u | (static_cast(codepoint) & 0x3Fu))); + } + else if (codepoint <= 0xFFFF) + { + // 3-byte characters: 1110xxxx 10xxxxxx 10xxxxxx + add(static_cast(0xE0u | (static_cast(codepoint) >> 12u))); + add(static_cast(0x80u | ((static_cast(codepoint) >> 6u) & 0x3Fu))); + add(static_cast(0x80u | (static_cast(codepoint) & 0x3Fu))); + } + else + { + // 4-byte characters: 11110xxx 10xxxxxx 10xxxxxx 10xxxxxx + add(static_cast(0xF0u | (static_cast(codepoint) >> 18u))); + add(static_cast(0x80u | ((static_cast(codepoint) >> 12u) & 0x3Fu))); + add(static_cast(0x80u | ((static_cast(codepoint) >> 6u) & 0x3Fu))); + add(static_cast(0x80u | (static_cast(codepoint) & 0x3Fu))); + } + + break; + } + + // other characters after escape + default: + error_message = "invalid string: forbidden character after backslash"; + return false; + } + + break; + } + + // invalid control characters + case 0x00: + { + error_message = "invalid string: control character U+0000 (NUL) must be escaped to \\u0000"; + return false; + } + + case 0x01: + { + error_message = "invalid string: control character U+0001 (SOH) must be escaped to \\u0001"; + return false; + } + + case 0x02: + { + error_message = "invalid string: control character U+0002 (STX) must be escaped to \\u0002"; + return false; + } + + case 0x03: + { + error_message = "invalid string: control character U+0003 (ETX) must be escaped to \\u0003"; + return false; + } + + case 0x04: + { + error_message = "invalid string: control character U+0004 (EOT) must be escaped to \\u0004"; + return false; + } + + case 0x05: + { + error_message = "invalid string: control character U+0005 (ENQ) must be escaped to \\u0005"; + return false; + } + + case 0x06: + { + error_message = "invalid string: control character U+0006 (ACK) must be escaped to \\u0006"; + return false; + } + + case 0x07: + { + error_message = "invalid string: control character U+0007 (BEL) must be escaped to \\u0007"; + return false; + } + + case 0x08: + { + error_message = "invalid string: control character U+0008 (BS) must be escaped to \\u0008 or \\b"; + return false; + } + + case 0x09: + { + error_message = "invalid string: control character U+0009 (HT) must be escaped to \\u0009 or \\t"; + return false; + } + + case 0x0A: + { + error_message = "invalid string: control character U+000A (LF) must be escaped to \\u000A or \\n"; + return false; + } + + case 0x0B: + { + error_message = "invalid string: control character U+000B (VT) must be escaped to \\u000B"; + return false; + } + + case 0x0C: + { + error_message = "invalid string: control character U+000C (FF) must be escaped to \\u000C or \\f"; + return false; + } + + case 0x0D: + { + error_message = "invalid string: control character U+000D (CR) must be escaped to \\u000D or \\r"; + return false; + } + + case 0x0E: + { + error_message = "invalid string: control character U+000E (SO) must be escaped to \\u000E"; + return false; + } + + case 0x0F: + { + error_message = "invalid string: control character U+000F (SI) must be escaped to \\u000F"; + return false; + } + + case 0x10: + { + error_message = "invalid string: control character U+0010 (DLE) must be escaped to \\u0010"; + return false; + } + + case 0x11: + { + error_message = "invalid string: control character U+0011 (DC1) must be escaped to \\u0011"; + return false; + } + + case 0x12: + { + error_message = "invalid string: control character U+0012 (DC2) must be escaped to \\u0012"; + return false; + } + + case 0x13: + { + error_message = "invalid string: control character U+0013 (DC3) must be escaped to \\u0013"; + return false; + } + + case 0x14: + { + error_message = "invalid string: control character U+0014 (DC4) must be escaped to \\u0014"; + return false; + } + + case 0x15: + { + error_message = "invalid string: control character U+0015 (NAK) must be escaped to \\u0015"; + return false; + } + + case 0x16: + { + error_message = "invalid string: control character U+0016 (SYN) must be escaped to \\u0016"; + return false; + } + + case 0x17: + { + error_message = "invalid string: control character U+0017 (ETB) must be escaped to \\u0017"; + return false; + } + + case 0x18: + { + error_message = "invalid string: control character U+0018 (CAN) must be escaped to \\u0018"; + return false; + } + + case 0x19: + { + error_message = "invalid string: control character U+0019 (EM) must be escaped to \\u0019"; + return false; + } + + case 0x1A: + { + error_message = "invalid string: control character U+001A (SUB) must be escaped to \\u001A"; + return false; + } + + case 0x1B: + { + error_message = "invalid string: control character U+001B (ESC) must be escaped to \\u001B"; + return false; + } + + case 0x1C: + { + error_message = "invalid string: control character U+001C (FS) must be escaped to \\u001C"; + return false; + } + + case 0x1D: + { + error_message = "invalid string: control character U+001D (GS) must be escaped to \\u001D"; + return false; + } + + case 0x1E: + { + error_message = "invalid string: control character U+001E (RS) must be escaped to \\u001E"; + return false; + } + + case 0x1F: + { + error_message = "invalid string: control character U+001F (US) must be escaped to \\u001F"; + return false; + } + + // U+0020..U+007F (except U+0022 (quote) and U+005C (backspace)) + case 0x20: + case 0x21: + case 0x23: + case 0x24: + case 0x25: + case 0x26: + case 0x27: + case 0x28: + case 0x29: + case 0x2A: + case 0x2B: + case 0x2C: + case 0x2D: + case 0x2E: + case 0x2F: + case 0x30: + case 0x31: + case 0x32: + case 0x33: + case 0x34: + case 0x35: + case 0x36: + case 0x37: + case 0x38: + case 0x39: + case 0x3A: + case 0x3B: + case 0x3C: + case 0x3D: + case 0x3E: + case 0x3F: + case 0x40: + case 0x41: + case 0x42: + case 0x43: + case 0x44: + case 0x45: + case 0x46: + case 0x47: + case 0x48: + case 0x49: + case 0x4A: + case 0x4B: + case 0x4C: + case 0x4D: + case 0x4E: + case 0x4F: + case 0x50: + case 0x51: + case 0x52: + case 0x53: + case 0x54: + case 0x55: + case 0x56: + case 0x57: + case 0x58: + case 0x59: + case 0x5A: + case 0x5B: + case 0x5D: + case 0x5E: + case 0x5F: + case 0x60: + case 0x61: + case 0x62: + case 0x63: + case 0x64: + case 0x65: + case 0x66: + case 0x67: + case 0x68: + case 0x69: + case 0x6A: + case 0x6B: + case 0x6C: + case 0x6D: + case 0x6E: + case 0x6F: + case 0x70: + case 0x71: + case 0x72: + case 0x73: + case 0x74: + case 0x75: + case 0x76: + case 0x77: + case 0x78: + case 0x79: + case 0x7A: + case 0x7B: + case 0x7C: + case 0x7D: + case 0x7E: + case 0x7F: + { + add(current); + break; + } + + // U+0080..U+07FF: bytes C2..DF 80..BF + case 0xC2: + case 0xC3: + case 0xC4: + case 0xC5: + case 0xC6: + case 0xC7: + case 0xC8: + case 0xC9: + case 0xCA: + case 0xCB: + case 0xCC: + case 0xCD: + case 0xCE: + case 0xCF: + case 0xD0: + case 0xD1: + case 0xD2: + case 0xD3: + case 0xD4: + case 0xD5: + case 0xD6: + case 0xD7: + case 0xD8: + case 0xD9: + case 0xDA: + case 0xDB: + case 0xDC: + case 0xDD: + case 0xDE: + case 0xDF: + { + if (JSON_HEDLEY_UNLIKELY(!next_byte_in_range({0x80, 0xBF}))) + { + return false; + } + break; + } + + // U+0800..U+0FFF: bytes E0 A0..BF 80..BF + case 0xE0: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0xA0, 0xBF, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // U+1000..U+CFFF: bytes E1..EC 80..BF 80..BF + // U+E000..U+FFFF: bytes EE..EF 80..BF 80..BF + case 0xE1: + case 0xE2: + case 0xE3: + case 0xE4: + case 0xE5: + case 0xE6: + case 0xE7: + case 0xE8: + case 0xE9: + case 0xEA: + case 0xEB: + case 0xEC: + case 0xEE: + case 0xEF: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0x80, 0xBF, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // U+D000..U+D7FF: bytes ED 80..9F 80..BF + case 0xED: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0x80, 0x9F, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // U+10000..U+3FFFF F0 90..BF 80..BF 80..BF + case 0xF0: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0x90, 0xBF, 0x80, 0xBF, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // U+40000..U+FFFFF F1..F3 80..BF 80..BF 80..BF + case 0xF1: + case 0xF2: + case 0xF3: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0x80, 0xBF, 0x80, 0xBF, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // U+100000..U+10FFFF F4 80..8F 80..BF 80..BF + case 0xF4: + { + if (JSON_HEDLEY_UNLIKELY(!(next_byte_in_range({0x80, 0x8F, 0x80, 0xBF, 0x80, 0xBF})))) + { + return false; + } + break; + } + + // remaining bytes (80..C1 and F5..FF) are ill-formed + default: + { + error_message = "invalid string: ill-formed UTF-8 byte"; + return false; + } + } + } + + error_message = "invalid string: missing closing quote"; + return false; + } +#endif +// end json.hpp +// clang-format on + +} // namespace detail + +namespace detail { + +double from_chars(const char *p) { +#if defined(MINIJSON_USE_STRTOD) + return strtod(p, nullptr); +#else + return simdjson::internal::from_chars(p); +#endif +} + +const char *my_strchr(const char *p, int ch) { + char c; + + constexpr uint64_t kMaxCount = 1024ull * 1024ull; // up to 1M chars + + uint64_t cnt{0}; + + c = ch; + for (;; ++p, cnt++) { + if (cnt > kMaxCount) { + return nullptr; + } + + if (*p == c) { + return (p); + } + if (*p == '\0') { + return (nullptr); + } + } +} + +} // namespace detail +} // namespace minijson + +#if !defined(MINIJSON_USE_STRTOD) + +#include +#include + +namespace minijson { +namespace simdjson { +namespace internal { + +/** + * The code in the internal::from_chars function is meant to handle the + *floating-point number parsing when we have more than 19 digits in the decimal + *mantissa. This should only be seen in adversarial scenarios: we do not expect + *production systems to even produce such floating-point numbers. + * + * The parser is based on work by Nigel Tao (at + *https://github.com/google/wuffs/) who credits Ken Thompson for the design (via + *a reference to the Go source code). See + * https://github.com/google/wuffs/blob/aa46859ea40c72516deffa1b146121952d6dfd3b/internal/cgen/base/floatconv-submodule-data.c + * https://github.com/google/wuffs/blob/46cd8105f47ca07ae2ba8e6a7818ef9c0df6c152/internal/cgen/base/floatconv-submodule-code.c + * It is probably not very fast but it is a fallback that should almost never be + * called in real life. Google Wuffs is published under APL 2.0. + **/ + +namespace { +constexpr uint32_t max_digits = 768; +constexpr int32_t decimal_point_range = 2047; +} // namespace + +struct adjusted_mantissa { + uint64_t mantissa; + int power2; + adjusted_mantissa() : mantissa(0), power2(0) {} +}; + +struct decimal { + uint32_t num_digits; + int32_t decimal_point; + bool negative; + bool truncated; + uint8_t digits[max_digits]; +}; + +template +struct binary_format { + static constexpr int mantissa_explicit_bits(); + static constexpr int minimum_exponent(); + static constexpr int infinite_power(); + static constexpr int sign_index(); +}; + +template <> +constexpr int binary_format::mantissa_explicit_bits() { + return 52; +} + +template <> +constexpr int binary_format::minimum_exponent() { + return -1023; +} +template <> +constexpr int binary_format::infinite_power() { + return 0x7FF; +} + +template <> +constexpr int binary_format::sign_index() { + return 63; +} + +inline bool is_integer(char c) noexcept { return (c >= '0' && c <= '9'); } + +// This should always succeed since it follows a call to parse_number. +static decimal parse_decimal(const char *&p) noexcept { + decimal answer; + answer.num_digits = 0; + answer.decimal_point = 0; + answer.truncated = false; + answer.negative = (*p == '-'); + if ((*p == '-') || (*p == '+')) { + ++p; + } + + while (*p == '0') { + ++p; + } + while (is_integer(*p)) { + if (answer.num_digits < max_digits) { + answer.digits[answer.num_digits] = uint8_t(*p - '0'); + } + answer.num_digits++; + ++p; + } + if (*p == '.') { + ++p; + const char *first_after_period = p; + // if we have not yet encountered a zero, we have to skip it as well + if (answer.num_digits == 0) { + // skip zeros + while (*p == '0') { + ++p; + } + } + while (is_integer(*p)) { + if (answer.num_digits < max_digits) { + answer.digits[answer.num_digits] = uint8_t(*p - '0'); + } + answer.num_digits++; + ++p; + } + answer.decimal_point = int32_t(first_after_period - p); + } + if (answer.num_digits > 0) { + const char *preverse = p - 1; + int32_t trailing_zeros = 0; + while ((*preverse == '0') || (*preverse == '.')) { + if (*preverse == '0') { + trailing_zeros++; + } + --preverse; + } + answer.decimal_point += int32_t(answer.num_digits); + answer.num_digits -= uint32_t(trailing_zeros); + } + if (answer.num_digits > max_digits) { + answer.num_digits = max_digits; + answer.truncated = true; + } + if (('e' == *p) || ('E' == *p)) { + ++p; + bool neg_exp = false; + if ('-' == *p) { + neg_exp = true; + ++p; + } else if ('+' == *p) { + ++p; + } + int32_t exp_number = 0; // exponential part + while (is_integer(*p)) { + uint8_t digit = uint8_t(*p - '0'); + if (exp_number < 0x10000) { + exp_number = 10 * exp_number + digit; + } + ++p; + } + answer.decimal_point += (neg_exp ? -exp_number : exp_number); + } + return answer; +} + +// This should always succeed since it follows a call to parse_number. +// Will not read at or beyond the "end" pointer. +static decimal parse_decimal(const char *&p, const char *end) noexcept { + decimal answer; + answer.num_digits = 0; + answer.decimal_point = 0; + answer.truncated = false; + if (p == end) { + return answer; + } // should never happen + answer.negative = (*p == '-'); + if ((*p == '-') || (*p == '+')) { + ++p; + } + + while ((p != end) && (*p == '0')) { + ++p; + } + while ((p != end) && is_integer(*p)) { + if (answer.num_digits < max_digits) { + answer.digits[answer.num_digits] = uint8_t(*p - '0'); + } + answer.num_digits++; + ++p; + } + if ((p != end) && (*p == '.')) { + ++p; + if (p == end) { + return answer; + } // should never happen + const char *first_after_period = p; + // if we have not yet encountered a zero, we have to skip it as well + if (answer.num_digits == 0) { + // skip zeros + while (*p == '0') { + ++p; + } + } + while ((p != end) && is_integer(*p)) { + if (answer.num_digits < max_digits) { + answer.digits[answer.num_digits] = uint8_t(*p - '0'); + } + answer.num_digits++; + ++p; + } + answer.decimal_point = int32_t(first_after_period - p); + } + if (answer.num_digits > 0) { + const char *preverse = p - 1; + int32_t trailing_zeros = 0; + while ((*preverse == '0') || (*preverse == '.')) { + if (*preverse == '0') { + trailing_zeros++; + } + --preverse; + } + answer.decimal_point += int32_t(answer.num_digits); + answer.num_digits -= uint32_t(trailing_zeros); + } + if (answer.num_digits > max_digits) { + answer.num_digits = max_digits; + answer.truncated = true; + } + if ((p != end) && (('e' == *p) || ('E' == *p))) { + ++p; + if (p == end) { + return answer; + } // should never happen + bool neg_exp = false; + if ('-' == *p) { + neg_exp = true; + ++p; + } else if ('+' == *p) { + ++p; + } + int32_t exp_number = 0; // exponential part + while ((p != end) && is_integer(*p)) { + uint8_t digit = uint8_t(*p - '0'); + if (exp_number < 0x10000) { + exp_number = 10 * exp_number + digit; + } + ++p; + } + answer.decimal_point += (neg_exp ? -exp_number : exp_number); + } + return answer; +} + +namespace { + +// remove all final zeroes +inline void trim(decimal &h) { + while ((h.num_digits > 0) && (h.digits[h.num_digits - 1] == 0)) { + h.num_digits--; + } +} + +uint32_t number_of_digits_decimal_left_shift(decimal &h, uint32_t shift) { + shift &= 63; + const static uint16_t number_of_digits_decimal_left_shift_table[65] = { + 0x0000, 0x0800, 0x0801, 0x0803, 0x1006, 0x1009, 0x100D, 0x1812, 0x1817, + 0x181D, 0x2024, 0x202B, 0x2033, 0x203C, 0x2846, 0x2850, 0x285B, 0x3067, + 0x3073, 0x3080, 0x388E, 0x389C, 0x38AB, 0x38BB, 0x40CC, 0x40DD, 0x40EF, + 0x4902, 0x4915, 0x4929, 0x513E, 0x5153, 0x5169, 0x5180, 0x5998, 0x59B0, + 0x59C9, 0x61E3, 0x61FD, 0x6218, 0x6A34, 0x6A50, 0x6A6D, 0x6A8B, 0x72AA, + 0x72C9, 0x72E9, 0x7B0A, 0x7B2B, 0x7B4D, 0x8370, 0x8393, 0x83B7, 0x83DC, + 0x8C02, 0x8C28, 0x8C4F, 0x9477, 0x949F, 0x94C8, 0x9CF2, 0x051C, 0x051C, + 0x051C, 0x051C, + }; + uint32_t x_a = number_of_digits_decimal_left_shift_table[shift]; + uint32_t x_b = number_of_digits_decimal_left_shift_table[shift + 1]; + uint32_t num_new_digits = x_a >> 11; + uint32_t pow5_a = 0x7FF & x_a; + uint32_t pow5_b = 0x7FF & x_b; + const static uint8_t + number_of_digits_decimal_left_shift_table_powers_of_5[0x051C] = { + 5, 2, 5, 1, 2, 5, 6, 2, 5, 3, 1, 2, 5, 1, 5, 6, 2, 5, 7, 8, 1, 2, 5, + 3, 9, 0, 6, 2, 5, 1, 9, 5, 3, 1, 2, 5, 9, 7, 6, 5, 6, 2, 5, 4, 8, 8, + 2, 8, 1, 2, 5, 2, 4, 4, 1, 4, 0, 6, 2, 5, 1, 2, 2, 0, 7, 0, 3, 1, 2, + 5, 6, 1, 0, 3, 5, 1, 5, 6, 2, 5, 3, 0, 5, 1, 7, 5, 7, 8, 1, 2, 5, 1, + 5, 2, 5, 8, 7, 8, 9, 0, 6, 2, 5, 7, 6, 2, 9, 3, 9, 4, 5, 3, 1, 2, 5, + 3, 8, 1, 4, 6, 9, 7, 2, 6, 5, 6, 2, 5, 1, 9, 0, 7, 3, 4, 8, 6, 3, 2, + 8, 1, 2, 5, 9, 5, 3, 6, 7, 4, 3, 1, 6, 4, 0, 6, 2, 5, 4, 7, 6, 8, 3, + 7, 1, 5, 8, 2, 0, 3, 1, 2, 5, 2, 3, 8, 4, 1, 8, 5, 7, 9, 1, 0, 1, 5, + 6, 2, 5, 1, 1, 9, 2, 0, 9, 2, 8, 9, 5, 5, 0, 7, 8, 1, 2, 5, 5, 9, 6, + 0, 4, 6, 4, 4, 7, 7, 5, 3, 9, 0, 6, 2, 5, 2, 9, 8, 0, 2, 3, 2, 2, 3, + 8, 7, 6, 9, 5, 3, 1, 2, 5, 1, 4, 9, 0, 1, 1, 6, 1, 1, 9, 3, 8, 4, 7, + 6, 5, 6, 2, 5, 7, 4, 5, 0, 5, 8, 0, 5, 9, 6, 9, 2, 3, 8, 2, 8, 1, 2, + 5, 3, 7, 2, 5, 2, 9, 0, 2, 9, 8, 4, 6, 1, 9, 1, 4, 0, 6, 2, 5, 1, 8, + 6, 2, 6, 4, 5, 1, 4, 9, 2, 3, 0, 9, 5, 7, 0, 3, 1, 2, 5, 9, 3, 1, 3, + 2, 2, 5, 7, 4, 6, 1, 5, 4, 7, 8, 5, 1, 5, 6, 2, 5, 4, 6, 5, 6, 6, 1, + 2, 8, 7, 3, 0, 7, 7, 3, 9, 2, 5, 7, 8, 1, 2, 5, 2, 3, 2, 8, 3, 0, 6, + 4, 3, 6, 5, 3, 8, 6, 9, 6, 2, 8, 9, 0, 6, 2, 5, 1, 1, 6, 4, 1, 5, 3, + 2, 1, 8, 2, 6, 9, 3, 4, 8, 1, 4, 4, 5, 3, 1, 2, 5, 5, 8, 2, 0, 7, 6, + 6, 0, 9, 1, 3, 4, 6, 7, 4, 0, 7, 2, 2, 6, 5, 6, 2, 5, 2, 9, 1, 0, 3, + 8, 3, 0, 4, 5, 6, 7, 3, 3, 7, 0, 3, 6, 1, 3, 2, 8, 1, 2, 5, 1, 4, 5, + 5, 1, 9, 1, 5, 2, 2, 8, 3, 6, 6, 8, 5, 1, 8, 0, 6, 6, 4, 0, 6, 2, 5, + 7, 2, 7, 5, 9, 5, 7, 6, 1, 4, 1, 8, 3, 4, 2, 5, 9, 0, 3, 3, 2, 0, 3, + 1, 2, 5, 3, 6, 3, 7, 9, 7, 8, 8, 0, 7, 0, 9, 1, 7, 1, 2, 9, 5, 1, 6, + 6, 0, 1, 5, 6, 2, 5, 1, 8, 1, 8, 9, 8, 9, 4, 0, 3, 5, 4, 5, 8, 5, 6, + 4, 7, 5, 8, 3, 0, 0, 7, 8, 1, 2, 5, 9, 0, 9, 4, 9, 4, 7, 0, 1, 7, 7, + 2, 9, 2, 8, 2, 3, 7, 9, 1, 5, 0, 3, 9, 0, 6, 2, 5, 4, 5, 4, 7, 4, 7, + 3, 5, 0, 8, 8, 6, 4, 6, 4, 1, 1, 8, 9, 5, 7, 5, 1, 9, 5, 3, 1, 2, 5, + 2, 2, 7, 3, 7, 3, 6, 7, 5, 4, 4, 3, 2, 3, 2, 0, 5, 9, 4, 7, 8, 7, 5, + 9, 7, 6, 5, 6, 2, 5, 1, 1, 3, 6, 8, 6, 8, 3, 7, 7, 2, 1, 6, 1, 6, 0, + 2, 9, 7, 3, 9, 3, 7, 9, 8, 8, 2, 8, 1, 2, 5, 5, 6, 8, 4, 3, 4, 1, 8, + 8, 6, 0, 8, 0, 8, 0, 1, 4, 8, 6, 9, 6, 8, 9, 9, 4, 1, 4, 0, 6, 2, 5, + 2, 8, 4, 2, 1, 7, 0, 9, 4, 3, 0, 4, 0, 4, 0, 0, 7, 4, 3, 4, 8, 4, 4, + 9, 7, 0, 7, 0, 3, 1, 2, 5, 1, 4, 2, 1, 0, 8, 5, 4, 7, 1, 5, 2, 0, 2, + 0, 0, 3, 7, 1, 7, 4, 2, 2, 4, 8, 5, 3, 5, 1, 5, 6, 2, 5, 7, 1, 0, 5, + 4, 2, 7, 3, 5, 7, 6, 0, 1, 0, 0, 1, 8, 5, 8, 7, 1, 1, 2, 4, 2, 6, 7, + 5, 7, 8, 1, 2, 5, 3, 5, 5, 2, 7, 1, 3, 6, 7, 8, 8, 0, 0, 5, 0, 0, 9, + 2, 9, 3, 5, 5, 6, 2, 1, 3, 3, 7, 8, 9, 0, 6, 2, 5, 1, 7, 7, 6, 3, 5, + 6, 8, 3, 9, 4, 0, 0, 2, 5, 0, 4, 6, 4, 6, 7, 7, 8, 1, 0, 6, 6, 8, 9, + 4, 5, 3, 1, 2, 5, 8, 8, 8, 1, 7, 8, 4, 1, 9, 7, 0, 0, 1, 2, 5, 2, 3, + 2, 3, 3, 8, 9, 0, 5, 3, 3, 4, 4, 7, 2, 6, 5, 6, 2, 5, 4, 4, 4, 0, 8, + 9, 2, 0, 9, 8, 5, 0, 0, 6, 2, 6, 1, 6, 1, 6, 9, 4, 5, 2, 6, 6, 7, 2, + 3, 6, 3, 2, 8, 1, 2, 5, 2, 2, 2, 0, 4, 4, 6, 0, 4, 9, 2, 5, 0, 3, 1, + 3, 0, 8, 0, 8, 4, 7, 2, 6, 3, 3, 3, 6, 1, 8, 1, 6, 4, 0, 6, 2, 5, 1, + 1, 1, 0, 2, 2, 3, 0, 2, 4, 6, 2, 5, 1, 5, 6, 5, 4, 0, 4, 2, 3, 6, 3, + 1, 6, 6, 8, 0, 9, 0, 8, 2, 0, 3, 1, 2, 5, 5, 5, 5, 1, 1, 1, 5, 1, 2, + 3, 1, 2, 5, 7, 8, 2, 7, 0, 2, 1, 1, 8, 1, 5, 8, 3, 4, 0, 4, 5, 4, 1, + 0, 1, 5, 6, 2, 5, 2, 7, 7, 5, 5, 5, 7, 5, 6, 1, 5, 6, 2, 8, 9, 1, 3, + 5, 1, 0, 5, 9, 0, 7, 9, 1, 7, 0, 2, 2, 7, 0, 5, 0, 7, 8, 1, 2, 5, 1, + 3, 8, 7, 7, 7, 8, 7, 8, 0, 7, 8, 1, 4, 4, 5, 6, 7, 5, 5, 2, 9, 5, 3, + 9, 5, 8, 5, 1, 1, 3, 5, 2, 5, 3, 9, 0, 6, 2, 5, 6, 9, 3, 8, 8, 9, 3, + 9, 0, 3, 9, 0, 7, 2, 2, 8, 3, 7, 7, 6, 4, 7, 6, 9, 7, 9, 2, 5, 5, 6, + 7, 6, 2, 6, 9, 5, 3, 1, 2, 5, 3, 4, 6, 9, 4, 4, 6, 9, 5, 1, 9, 5, 3, + 6, 1, 4, 1, 8, 8, 8, 2, 3, 8, 4, 8, 9, 6, 2, 7, 8, 3, 8, 1, 3, 4, 7, + 6, 5, 6, 2, 5, 1, 7, 3, 4, 7, 2, 3, 4, 7, 5, 9, 7, 6, 8, 0, 7, 0, 9, + 4, 4, 1, 1, 9, 2, 4, 4, 8, 1, 3, 9, 1, 9, 0, 6, 7, 3, 8, 2, 8, 1, 2, + 5, 8, 6, 7, 3, 6, 1, 7, 3, 7, 9, 8, 8, 4, 0, 3, 5, 4, 7, 2, 0, 5, 9, + 6, 2, 2, 4, 0, 6, 9, 5, 9, 5, 3, 3, 6, 9, 1, 4, 0, 6, 2, 5, + }; + const uint8_t *pow5 = + &number_of_digits_decimal_left_shift_table_powers_of_5[pow5_a]; + uint32_t i = 0; + uint32_t n = pow5_b - pow5_a; + for (; i < n; i++) { + if (i >= h.num_digits) { + return num_new_digits - 1; + } else if (h.digits[i] == pow5[i]) { + continue; + } else if (h.digits[i] < pow5[i]) { + return num_new_digits - 1; + } else { + return num_new_digits; + } + } + return num_new_digits; +} + +} // end of anonymous namespace + +static uint64_t round(decimal &h) { + if ((h.num_digits == 0) || (h.decimal_point < 0)) { + return 0; + } else if (h.decimal_point > 18) { + return UINT64_MAX; + } + // at this point, we know that h.decimal_point >= 0 + uint32_t dp = uint32_t(h.decimal_point); + uint64_t n = 0; + for (uint32_t i = 0; i < dp; i++) { + n = (10 * n) + ((i < h.num_digits) ? h.digits[i] : 0); + } + bool round_up = false; + if (dp < h.num_digits) { + round_up = h.digits[dp] >= 5; // normally, we round up + // but we may need to round to even! + if ((h.digits[dp] == 5) && (dp + 1 == h.num_digits)) { + round_up = h.truncated || ((dp > 0) && (1 & h.digits[dp - 1])); + } + } + if (round_up) { + n++; + } + return n; +} + +// computes h * 2^-shift +static void decimal_left_shift(decimal &h, uint32_t shift) { + if (h.num_digits == 0) { + return; + } + uint32_t num_new_digits = number_of_digits_decimal_left_shift(h, shift); + int32_t read_index = int32_t(h.num_digits - 1); + uint32_t write_index = h.num_digits - 1 + num_new_digits; + uint64_t n = 0; + + while (read_index >= 0) { + n += uint64_t(h.digits[read_index]) << shift; + uint64_t quotient = n / 10; + uint64_t remainder = n - (10 * quotient); + if (write_index < max_digits) { + h.digits[write_index] = uint8_t(remainder); + } else if (remainder > 0) { + h.truncated = true; + } + n = quotient; + write_index--; + read_index--; + } + while (n > 0) { + uint64_t quotient = n / 10; + uint64_t remainder = n - (10 * quotient); + if (write_index < max_digits) { + h.digits[write_index] = uint8_t(remainder); + } else if (remainder > 0) { + h.truncated = true; + } + n = quotient; + write_index--; + } + h.num_digits += num_new_digits; + if (h.num_digits > max_digits) { + h.num_digits = max_digits; + } + h.decimal_point += int32_t(num_new_digits); + trim(h); +} + +// computes h * 2^shift +static void decimal_right_shift(decimal &h, uint32_t shift) { + uint32_t read_index = 0; + uint32_t write_index = 0; + + uint64_t n = 0; + + while ((n >> shift) == 0) { + if (read_index < h.num_digits) { + n = (10 * n) + h.digits[read_index++]; + } else if (n == 0) { + return; + } else { + while ((n >> shift) == 0) { + n = 10 * n; + read_index++; + } + break; + } + } + h.decimal_point -= int32_t(read_index - 1); + if (h.decimal_point < -decimal_point_range) { // it is zero + h.num_digits = 0; + h.decimal_point = 0; + h.negative = false; + h.truncated = false; + return; + } + uint64_t mask = (uint64_t(1) << shift) - 1; + while (read_index < h.num_digits) { + uint8_t new_digit = uint8_t(n >> shift); + n = (10 * (n & mask)) + h.digits[read_index++]; + h.digits[write_index++] = new_digit; + } + while (n > 0) { + uint8_t new_digit = uint8_t(n >> shift); + n = 10 * (n & mask); + if (write_index < max_digits) { + h.digits[write_index++] = new_digit; + } else if (new_digit > 0) { + h.truncated = true; + } + } + h.num_digits = write_index; + trim(h); +} + +template +adjusted_mantissa compute_float(decimal &d) { + adjusted_mantissa answer; + if (d.num_digits == 0) { + // should be zero + answer.power2 = 0; + answer.mantissa = 0; + return answer; + } + // At this point, going further, we can assume that d.num_digits > 0. + // We want to guard against excessive decimal point values because + // they can result in long running times. Indeed, we do + // shifts by at most 60 bits. We have that log(10**400)/log(2**60) ~= 22 + // which is fine, but log(10**299995)/log(2**60) ~= 16609 which is not + // fine (runs for a long time). + // + if (d.decimal_point < -324) { + // We have something smaller than 1e-324 which is always zero + // in binary64 and binary32. + // It should be zero. + answer.power2 = 0; + answer.mantissa = 0; + return answer; + } else if (d.decimal_point >= 310) { + // We have something at least as large as 0.1e310 which is + // always infinite. + answer.power2 = binary::infinite_power(); + answer.mantissa = 0; + return answer; + } + + static const uint32_t max_shift = 60; + static const uint32_t num_powers = 19; + static const uint8_t powers[19] = { + 0, 3, 6, 9, 13, 16, 19, 23, 26, 29, // + 33, 36, 39, 43, 46, 49, 53, 56, 59, // + }; + int32_t exp2 = 0; + while (d.decimal_point > 0) { + uint32_t n = uint32_t(d.decimal_point); + uint32_t shift = (n < num_powers) ? powers[n] : max_shift; + decimal_right_shift(d, shift); + if (d.decimal_point < -decimal_point_range) { + // should be zero + answer.power2 = 0; + answer.mantissa = 0; + return answer; + } + exp2 += int32_t(shift); + } + // We shift left toward [1/2 ... 1]. + while (d.decimal_point <= 0) { + uint32_t shift; + if (d.decimal_point == 0) { + if (d.digits[0] >= 5) { + break; + } + shift = (d.digits[0] < 2) ? 2 : 1; + } else { + uint32_t n = uint32_t(-d.decimal_point); + shift = (n < num_powers) ? powers[n] : max_shift; + } + decimal_left_shift(d, shift); + if (d.decimal_point > decimal_point_range) { + // we want to get infinity: + answer.power2 = 0xFF; + answer.mantissa = 0; + return answer; + } + exp2 -= int32_t(shift); + } + // We are now in the range [1/2 ... 1] but the binary format uses [1 ... 2]. + exp2--; + constexpr int32_t minimum_exponent = binary::minimum_exponent(); + while ((minimum_exponent + 1) > exp2) { + uint32_t n = uint32_t((minimum_exponent + 1) - exp2); + if (n > max_shift) { + n = max_shift; + } + decimal_right_shift(d, n); + exp2 += int32_t(n); + } + if ((exp2 - minimum_exponent) >= binary::infinite_power()) { + answer.power2 = binary::infinite_power(); + answer.mantissa = 0; + return answer; + } + + const int mantissa_size_in_bits = binary::mantissa_explicit_bits() + 1; + decimal_left_shift(d, mantissa_size_in_bits); + + uint64_t mantissa = round(d); + // It is possible that we have an overflow, in which case we need + // to shift back. + if (mantissa >= (uint64_t(1) << mantissa_size_in_bits)) { + decimal_right_shift(d, 1); + exp2 += 1; + mantissa = round(d); + if ((exp2 - minimum_exponent) >= binary::infinite_power()) { + answer.power2 = binary::infinite_power(); + answer.mantissa = 0; + return answer; + } + } + answer.power2 = exp2 - binary::minimum_exponent(); + if (mantissa < (uint64_t(1) << binary::mantissa_explicit_bits())) { + answer.power2--; + } + answer.mantissa = + mantissa & ((uint64_t(1) << binary::mantissa_explicit_bits()) - 1); + return answer; +} + +template +adjusted_mantissa parse_long_mantissa(const char *first) { + decimal d = parse_decimal(first); + return compute_float(d); +} + +template +adjusted_mantissa parse_long_mantissa(const char *first, const char *end) { + decimal d = parse_decimal(first, end); + return compute_float(d); +} + +double from_chars(const char *first) noexcept { + bool negative = first[0] == '-'; + if (negative) { + first++; + } + adjusted_mantissa am = parse_long_mantissa>(first); + uint64_t word = am.mantissa; + word |= uint64_t(am.power2) + << binary_format::mantissa_explicit_bits(); + word = negative ? word | (uint64_t(1) << binary_format::sign_index()) + : word; + double value; + std::memcpy(&value, &word, sizeof(double)); + return value; +} + +double from_chars(const char *first, const char *end) noexcept { + bool negative = first[0] == '-'; + if (negative) { + first++; + } + adjusted_mantissa am = parse_long_mantissa>(first, end); + uint64_t word = am.mantissa; + word |= uint64_t(am.power2) + << binary_format::mantissa_explicit_bits(); + word = negative ? word | (uint64_t(1) << binary_format::sign_index()) + : word; + double value; + std::memcpy(&value, &word, sizeof(double)); + return value; +} + +} // namespace internal +} // namespace simdjson +} // namespace minijson + +namespace minijson { +namespace simdjson { +namespace internal { +/*! +implements the Grisu2 algorithm for binary to decimal floating-point +conversion. +Adapted from JSON for Modern C++ + +This implementation is a slightly modified version of the reference +implementation which may be obtained from +http://florian.loitsch.com/publications (bench.tar.gz). +The code is distributed under the MIT license, Copyright (c) 2009 Florian +Loitsch. For a detailed description of the algorithm see: [1] Loitsch, "Printing +Floating-Point Numbers Quickly and Accurately with Integers", Proceedings of the +ACM SIGPLAN 2010 Conference on Programming Language Design and Implementation, +PLDI 2010 [2] Burger, Dybvig, "Printing Floating-Point Numbers Quickly and +Accurately", Proceedings of the ACM SIGPLAN 1996 Conference on Programming +Language Design and Implementation, PLDI 1996 +*/ +namespace dtoa_impl { + +template +Target reinterpret_bits(const Source source) { + static_assert(sizeof(Target) == sizeof(Source), "size mismatch"); + + Target target; + std::memcpy(&target, &source, sizeof(Source)); + return target; +} + +struct diyfp // f * 2^e +{ + static constexpr int kPrecision = 64; // = q + + std::uint64_t f = 0; + int e = 0; + + constexpr diyfp(std::uint64_t f_, int e_) noexcept : f(f_), e(e_) {} + + /*! + @brief returns x - y + @pre x.e == y.e and x.f >= y.f + */ + static diyfp sub(const diyfp &x, const diyfp &y) noexcept { + return {x.f - y.f, x.e}; + } + + /*! + @brief returns x * y + @note The result is rounded. (Only the upper q bits are returned.) + */ + static diyfp mul(const diyfp &x, const diyfp &y) noexcept { + static_assert(kPrecision == 64, "internal error"); + + // Computes: + // f = round((x.f * y.f) / 2^q) + // e = x.e + y.e + q + + // Emulate the 64-bit * 64-bit multiplication: + // + // p = u * v + // = (u_lo + 2^32 u_hi) (v_lo + 2^32 v_hi) + // = (u_lo v_lo ) + 2^32 ((u_lo v_hi ) + (u_hi v_lo )) + + // 2^64 (u_hi v_hi ) = (p0 ) + 2^32 ((p1 ) + (p2 )) + // + 2^64 (p3 ) = (p0_lo + 2^32 p0_hi) + 2^32 ((p1_lo + + // 2^32 p1_hi) + (p2_lo + 2^32 p2_hi)) + 2^64 (p3 ) = + // (p0_lo ) + 2^32 (p0_hi + p1_lo + p2_lo ) + 2^64 (p1_hi + + // p2_hi + p3) = (p0_lo ) + 2^32 (Q ) + 2^64 (H ) = (p0_lo ) + + // 2^32 (Q_lo + 2^32 Q_hi ) + 2^64 (H ) + // + // (Since Q might be larger than 2^32 - 1) + // + // = (p0_lo + 2^32 Q_lo) + 2^64 (Q_hi + H) + // + // (Q_hi + H does not overflow a 64-bit int) + // + // = p_lo + 2^64 p_hi + + const std::uint64_t u_lo = x.f & 0xFFFFFFFFu; + const std::uint64_t u_hi = x.f >> 32u; + const std::uint64_t v_lo = y.f & 0xFFFFFFFFu; + const std::uint64_t v_hi = y.f >> 32u; + + const std::uint64_t p0 = u_lo * v_lo; + const std::uint64_t p1 = u_lo * v_hi; + const std::uint64_t p2 = u_hi * v_lo; + const std::uint64_t p3 = u_hi * v_hi; + + const std::uint64_t p0_hi = p0 >> 32u; + const std::uint64_t p1_lo = p1 & 0xFFFFFFFFu; + const std::uint64_t p1_hi = p1 >> 32u; + const std::uint64_t p2_lo = p2 & 0xFFFFFFFFu; + const std::uint64_t p2_hi = p2 >> 32u; + + std::uint64_t Q = p0_hi + p1_lo + p2_lo; + + // The full product might now be computed as + // + // p_hi = p3 + p2_hi + p1_hi + (Q >> 32) + // p_lo = p0_lo + (Q << 32) + // + // But in this particular case here, the full p_lo is not required. + // Effectively we only need to add the highest bit in p_lo to p_hi (and + // Q_hi + 1 does not overflow). + + Q += std::uint64_t{1} << (64u - 32u - 1u); // round, ties up + + const std::uint64_t h = p3 + p2_hi + p1_hi + (Q >> 32u); + + return {h, x.e + y.e + 64}; + } + + /*! + @brief normalize x such that the significand is >= 2^(q-1) + @pre x.f != 0 + */ + static diyfp normalize(diyfp x) noexcept { + while ((x.f >> 63u) == 0) { + x.f <<= 1u; + x.e--; + } + + return x; + } + + /*! + @brief normalize x such that the result has the exponent E + @pre e >= x.e and the upper e - x.e bits of x.f must be zero. + */ + static diyfp normalize_to(const diyfp &x, + const int target_exponent) noexcept { + const int delta = x.e - target_exponent; + + return {x.f << delta, target_exponent}; + } +}; + +struct boundaries { + diyfp w; + diyfp minus; + diyfp plus; +}; + +/*! +Compute the (normalized) diyfp representing the input number 'value' and its +boundaries. +@pre value must be finite and positive +*/ +template +boundaries compute_boundaries(FloatType value) { + // Convert the IEEE representation into a diyfp. + // + // If v is denormal: + // value = 0.F * 2^(1 - bias) = ( F) * 2^(1 - bias - (p-1)) + // If v is normalized: + // value = 1.F * 2^(E - bias) = (2^(p-1) + F) * 2^(E - bias - (p-1)) + + static_assert(std::numeric_limits::is_iec559, + "internal error: dtoa_short requires an IEEE-754 " + "floating-point implementation"); + + constexpr int kPrecision = + std::numeric_limits::digits; // = p (includes the hidden bit) + constexpr int kBias = + std::numeric_limits::max_exponent - 1 + (kPrecision - 1); + constexpr int kMinExp = 1 - kBias; + constexpr std::uint64_t kHiddenBit = std::uint64_t{1} + << (kPrecision - 1); // = 2^(p-1) + + using bits_type = typename std::conditional::type; + + const std::uint64_t bits = reinterpret_bits(value); + const std::uint64_t E = bits >> (kPrecision - 1); + const std::uint64_t F = bits & (kHiddenBit - 1); + + const bool is_denormal = E == 0; + const diyfp v = is_denormal + ? diyfp(F, kMinExp) + : diyfp(F + kHiddenBit, static_cast(E) - kBias); + + // Compute the boundaries m- and m+ of the floating-point value + // v = f * 2^e. + // + // Determine v- and v+, the floating-point predecessor and successor if v, + // respectively. + // + // v- = v - 2^e if f != 2^(p-1) or e == e_min (A) + // = v - 2^(e-1) if f == 2^(p-1) and e > e_min (B) + // + // v+ = v + 2^e + // + // Let m- = (v- + v) / 2 and m+ = (v + v+) / 2. All real numbers _strictly_ + // between m- and m+ round to v, regardless of how the input rounding + // algorithm breaks ties. + // + // ---+-------------+-------------+-------------+-------------+--- (A) + // v- m- v m+ v+ + // + // -----------------+------+------+-------------+-------------+--- (B) + // v- m- v m+ v+ + + const bool lower_boundary_is_closer = F == 0 && E > 1; + const diyfp m_plus = diyfp(2 * v.f + 1, v.e - 1); + const diyfp m_minus = lower_boundary_is_closer + ? diyfp(4 * v.f - 1, v.e - 2) // (B) + : diyfp(2 * v.f - 1, v.e - 1); // (A) + + // Determine the normalized w+ = m+. + const diyfp w_plus = diyfp::normalize(m_plus); + + // Determine w- = m- such that e_(w-) = e_(w+). + const diyfp w_minus = diyfp::normalize_to(m_minus, w_plus.e); + + return {diyfp::normalize(v), w_minus, w_plus}; +} + +// Given normalized diyfp w, Grisu needs to find a (normalized) cached +// power-of-ten c, such that the exponent of the product c * w = f * 2^e lies +// within a certain range [alpha, gamma] (Definition 3.2 from [1]) +// +// alpha <= e = e_c + e_w + q <= gamma +// +// or +// +// f_c * f_w * 2^alpha <= f_c 2^(e_c) * f_w 2^(e_w) * 2^q +// <= f_c * f_w * 2^gamma +// +// Since c and w are normalized, i.e. 2^(q-1) <= f < 2^q, this implies +// +// 2^(q-1) * 2^(q-1) * 2^alpha <= c * w * 2^q < 2^q * 2^q * 2^gamma +// +// or +// +// 2^(q - 2 + alpha) <= c * w < 2^(q + gamma) +// +// The choice of (alpha,gamma) determines the size of the table and the form of +// the digit generation procedure. Using (alpha,gamma)=(-60,-32) works out well +// in practice: +// +// The idea is to cut the number c * w = f * 2^e into two parts, which can be +// processed independently: An integral part p1, and a fractional part p2: +// +// f * 2^e = ( (f div 2^-e) * 2^-e + (f mod 2^-e) ) * 2^e +// = (f div 2^-e) + (f mod 2^-e) * 2^e +// = p1 + p2 * 2^e +// +// The conversion of p1 into decimal form requires a series of divisions and +// modulos by (a power of) 10. These operations are faster for 32-bit than for +// 64-bit integers, so p1 should ideally fit into a 32-bit integer. This can be +// achieved by choosing +// +// -e >= 32 or e <= -32 := gamma +// +// In order to convert the fractional part +// +// p2 * 2^e = p2 / 2^-e = d[-1] / 10^1 + d[-2] / 10^2 + ... +// +// into decimal form, the fraction is repeatedly multiplied by 10 and the digits +// d[-i] are extracted in order: +// +// (10 * p2) div 2^-e = d[-1] +// (10 * p2) mod 2^-e = d[-2] / 10^1 + ... +// +// The multiplication by 10 must not overflow. It is sufficient to choose +// +// 10 * p2 < 16 * p2 = 2^4 * p2 <= 2^64. +// +// Since p2 = f mod 2^-e < 2^-e, +// +// -e <= 60 or e >= -60 := alpha + +constexpr int kAlpha = -60; +constexpr int kGamma = -32; + +struct cached_power // c = f * 2^e ~= 10^k +{ + std::uint64_t f; + int e; + int k; +}; + +/*! +For a normalized diyfp w = f * 2^e, this function returns a (normalized) cached +power-of-ten c = f_c * 2^e_c, such that the exponent of the product w * c +satisfies (Definition 3.2 from [1]) + alpha <= e_c + e + q <= gamma. +*/ +inline cached_power get_cached_power_for_binary_exponent(int e) { + // Now + // + // alpha <= e_c + e + q <= gamma (1) + // ==> f_c * 2^alpha <= c * 2^e * 2^q + // + // and since the c's are normalized, 2^(q-1) <= f_c, + // + // ==> 2^(q - 1 + alpha) <= c * 2^(e + q) + // ==> 2^(alpha - e - 1) <= c + // + // If c were an exact power of ten, i.e. c = 10^k, one may determine k as + // + // k = ceil( log_10( 2^(alpha - e - 1) ) ) + // = ceil( (alpha - e - 1) * log_10(2) ) + // + // From the paper: + // "In theory the result of the procedure could be wrong since c is rounded, + // and the computation itself is approximated [...]. In practice, however, + // this simple function is sufficient." + // + // For IEEE double precision floating-point numbers converted into + // normalized diyfp's w = f * 2^e, with q = 64, + // + // e >= -1022 (min IEEE exponent) + // -52 (p - 1) + // -52 (p - 1, possibly normalize denormal IEEE numbers) + // -11 (normalize the diyfp) + // = -1137 + // + // and + // + // e <= +1023 (max IEEE exponent) + // -52 (p - 1) + // -11 (normalize the diyfp) + // = 960 + // + // This binary exponent range [-1137,960] results in a decimal exponent + // range [-307,324]. One does not need to store a cached power for each + // k in this range. For each such k it suffices to find a cached power + // such that the exponent of the product lies in [alpha,gamma]. + // This implies that the difference of the decimal exponents of adjacent + // table entries must be less than or equal to + // + // floor( (gamma - alpha) * log_10(2) ) = 8. + // + // (A smaller distance gamma-alpha would require a larger table.) + + // NB: + // Actually this function returns c, such that -60 <= e_c + e + 64 <= -34. + + constexpr int kCachedPowersMinDecExp = -300; + constexpr int kCachedPowersDecStep = 8; + + static constexpr std::array kCachedPowers = {{ + {0xAB70FE17C79AC6CA, -1060, -300}, {0xFF77B1FCBEBCDC4F, -1034, -292}, + {0xBE5691EF416BD60C, -1007, -284}, {0x8DD01FAD907FFC3C, -980, -276}, + {0xD3515C2831559A83, -954, -268}, {0x9D71AC8FADA6C9B5, -927, -260}, + {0xEA9C227723EE8BCB, -901, -252}, {0xAECC49914078536D, -874, -244}, + {0x823C12795DB6CE57, -847, -236}, {0xC21094364DFB5637, -821, -228}, + {0x9096EA6F3848984F, -794, -220}, {0xD77485CB25823AC7, -768, -212}, + {0xA086CFCD97BF97F4, -741, -204}, {0xEF340A98172AACE5, -715, -196}, + {0xB23867FB2A35B28E, -688, -188}, {0x84C8D4DFD2C63F3B, -661, -180}, + {0xC5DD44271AD3CDBA, -635, -172}, {0x936B9FCEBB25C996, -608, -164}, + {0xDBAC6C247D62A584, -582, -156}, {0xA3AB66580D5FDAF6, -555, -148}, + {0xF3E2F893DEC3F126, -529, -140}, {0xB5B5ADA8AAFF80B8, -502, -132}, + {0x87625F056C7C4A8B, -475, -124}, {0xC9BCFF6034C13053, -449, -116}, + {0x964E858C91BA2655, -422, -108}, {0xDFF9772470297EBD, -396, -100}, + {0xA6DFBD9FB8E5B88F, -369, -92}, {0xF8A95FCF88747D94, -343, -84}, + {0xB94470938FA89BCF, -316, -76}, {0x8A08F0F8BF0F156B, -289, -68}, + {0xCDB02555653131B6, -263, -60}, {0x993FE2C6D07B7FAC, -236, -52}, + {0xE45C10C42A2B3B06, -210, -44}, {0xAA242499697392D3, -183, -36}, + {0xFD87B5F28300CA0E, -157, -28}, {0xBCE5086492111AEB, -130, -20}, + {0x8CBCCC096F5088CC, -103, -12}, {0xD1B71758E219652C, -77, -4}, + {0x9C40000000000000, -50, 4}, {0xE8D4A51000000000, -24, 12}, + {0xAD78EBC5AC620000, 3, 20}, {0x813F3978F8940984, 30, 28}, + {0xC097CE7BC90715B3, 56, 36}, {0x8F7E32CE7BEA5C70, 83, 44}, + {0xD5D238A4ABE98068, 109, 52}, {0x9F4F2726179A2245, 136, 60}, + {0xED63A231D4C4FB27, 162, 68}, {0xB0DE65388CC8ADA8, 189, 76}, + {0x83C7088E1AAB65DB, 216, 84}, {0xC45D1DF942711D9A, 242, 92}, + {0x924D692CA61BE758, 269, 100}, {0xDA01EE641A708DEA, 295, 108}, + {0xA26DA3999AEF774A, 322, 116}, {0xF209787BB47D6B85, 348, 124}, + {0xB454E4A179DD1877, 375, 132}, {0x865B86925B9BC5C2, 402, 140}, + {0xC83553C5C8965D3D, 428, 148}, {0x952AB45CFA97A0B3, 455, 156}, + {0xDE469FBD99A05FE3, 481, 164}, {0xA59BC234DB398C25, 508, 172}, + {0xF6C69A72A3989F5C, 534, 180}, {0xB7DCBF5354E9BECE, 561, 188}, + {0x88FCF317F22241E2, 588, 196}, {0xCC20CE9BD35C78A5, 614, 204}, + {0x98165AF37B2153DF, 641, 212}, {0xE2A0B5DC971F303A, 667, 220}, + {0xA8D9D1535CE3B396, 694, 228}, {0xFB9B7CD9A4A7443C, 720, 236}, + {0xBB764C4CA7A44410, 747, 244}, {0x8BAB8EEFB6409C1A, 774, 252}, + {0xD01FEF10A657842C, 800, 260}, {0x9B10A4E5E9913129, 827, 268}, + {0xE7109BFBA19C0C9D, 853, 276}, {0xAC2820D9623BF429, 880, 284}, + {0x80444B5E7AA7CF85, 907, 292}, {0xBF21E44003ACDD2D, 933, 300}, + {0x8E679C2F5E44FF8F, 960, 308}, {0xD433179D9C8CB841, 986, 316}, + {0x9E19DB92B4E31BA9, 1013, 324}, + }}; + + // This computation gives exactly the same results for k as + // k = ceil((kAlpha - e - 1) * 0.30102999566398114) + // for |e| <= 1500, but doesn't require floating-point operations. + // NB: log_10(2) ~= 78913 / 2^18 + const int f = kAlpha - e - 1; + const int k = (f * 78913) / (1 << 18) + static_cast(f > 0); + + const int index = (-kCachedPowersMinDecExp + k + (kCachedPowersDecStep - 1)) / + kCachedPowersDecStep; + + const cached_power cached = kCachedPowers[static_cast(index)]; + + return cached; +} + +/*! +For n != 0, returns k, such that pow10 := 10^(k-1) <= n < 10^k. +For n == 0, returns 1 and sets pow10 := 1. +*/ +inline int find_largest_pow10(const std::uint32_t n, std::uint32_t &pow10) { + // LCOV_EXCL_START + if (n >= 1000000000) { + pow10 = 1000000000; + return 10; + } + // LCOV_EXCL_STOP + else if (n >= 100000000) { + pow10 = 100000000; + return 9; + } else if (n >= 10000000) { + pow10 = 10000000; + return 8; + } else if (n >= 1000000) { + pow10 = 1000000; + return 7; + } else if (n >= 100000) { + pow10 = 100000; + return 6; + } else if (n >= 10000) { + pow10 = 10000; + return 5; + } else if (n >= 1000) { + pow10 = 1000; + return 4; + } else if (n >= 100) { + pow10 = 100; + return 3; + } else if (n >= 10) { + pow10 = 10; + return 2; + } else { + pow10 = 1; + return 1; + } +} + +inline void grisu2_round(char *buf, int len, std::uint64_t dist, + std::uint64_t delta, std::uint64_t rest, + std::uint64_t ten_k) { + // <--------------------------- delta ----> + // <---- dist ---------> + // --------------[------------------+-------------------]-------------- + // M- w M+ + // + // ten_k + // <------> + // <---- rest ----> + // --------------[------------------+----+--------------]-------------- + // w V + // = buf * 10^k + // + // ten_k represents a unit-in-the-last-place in the decimal representation + // stored in buf. + // Decrement buf by ten_k while this takes buf closer to w. + + // The tests are written in this order to avoid overflow in unsigned + // integer arithmetic. + + while (rest < dist && delta - rest >= ten_k && + (rest + ten_k < dist || dist - rest > rest + ten_k - dist)) { + buf[len - 1]--; + rest += ten_k; + } +} + +/*! +Generates V = buffer * 10^decimal_exponent, such that M- <= V <= M+. +M- and M+ must be normalized and share the same exponent -60 <= e <= -32. +*/ +inline void grisu2_digit_gen(char *buffer, int &length, int &decimal_exponent, + diyfp M_minus, diyfp w, diyfp M_plus) { + static_assert(kAlpha >= -60, "internal error"); + static_assert(kGamma <= -32, "internal error"); + + // Generates the digits (and the exponent) of a decimal floating-point + // number V = buffer * 10^decimal_exponent in the range [M-, M+]. The diyfp's + // w, M- and M+ share the same exponent e, which satisfies alpha <= e <= + // gamma. + // + // <--------------------------- delta ----> + // <---- dist ---------> + // --------------[------------------+-------------------]-------------- + // M- w M+ + // + // Grisu2 generates the digits of M+ from left to right and stops as soon as + // V is in [M-,M+]. + + std::uint64_t delta = + diyfp::sub(M_plus, M_minus) + .f; // (significand of (M+ - M-), implicit exponent is e) + std::uint64_t dist = + diyfp::sub(M_plus, w) + .f; // (significand of (M+ - w ), implicit exponent is e) + + // Split M+ = f * 2^e into two parts p1 and p2 (note: e < 0): + // + // M+ = f * 2^e + // = ((f div 2^-e) * 2^-e + (f mod 2^-e)) * 2^e + // = ((p1 ) * 2^-e + (p2 )) * 2^e + // = p1 + p2 * 2^e + + const diyfp one(std::uint64_t{1} << -M_plus.e, M_plus.e); + + auto p1 = static_cast( + M_plus.f >> + -one.e); // p1 = f div 2^-e (Since -e >= 32, p1 fits into a 32-bit int.) + std::uint64_t p2 = M_plus.f & (one.f - 1); // p2 = f mod 2^-e + + // 1) + // + // Generate the digits of the integral part p1 = d[n-1]...d[1]d[0] + + std::uint32_t pow10; + const int k = find_largest_pow10(p1, pow10); + + // 10^(k-1) <= p1 < 10^k, pow10 = 10^(k-1) + // + // p1 = (p1 div 10^(k-1)) * 10^(k-1) + (p1 mod 10^(k-1)) + // = (d[k-1] ) * 10^(k-1) + (p1 mod 10^(k-1)) + // + // M+ = p1 + p2 * 2^e + // = d[k-1] * 10^(k-1) + (p1 mod 10^(k-1)) + p2 * 2^e + // = d[k-1] * 10^(k-1) + ((p1 mod 10^(k-1)) * 2^-e + p2) * 2^e + // = d[k-1] * 10^(k-1) + ( rest) * 2^e + // + // Now generate the digits d[n] of p1 from left to right (n = k-1,...,0) + // + // p1 = d[k-1]...d[n] * 10^n + d[n-1]...d[0] + // + // but stop as soon as + // + // rest * 2^e = (d[n-1]...d[0] * 2^-e + p2) * 2^e <= delta * 2^e + + int n = k; + while (n > 0) { + // Invariants: + // M+ = buffer * 10^n + (p1 + p2 * 2^e) (buffer = 0 for n = k) + // pow10 = 10^(n-1) <= p1 < 10^n + // + const std::uint32_t d = p1 / pow10; // d = p1 div 10^(n-1) + const std::uint32_t r = p1 % pow10; // r = p1 mod 10^(n-1) + // + // M+ = buffer * 10^n + (d * 10^(n-1) + r) + p2 * 2^e + // = (buffer * 10 + d) * 10^(n-1) + (r + p2 * 2^e) + // + buffer[length++] = static_cast('0' + d); // buffer := buffer * 10 + d + // + // M+ = buffer * 10^(n-1) + (r + p2 * 2^e) + // + p1 = r; + n--; + // + // M+ = buffer * 10^n + (p1 + p2 * 2^e) + // pow10 = 10^n + // + + // Now check if enough digits have been generated. + // Compute + // + // p1 + p2 * 2^e = (p1 * 2^-e + p2) * 2^e = rest * 2^e + // + // Note: + // Since rest and delta share the same exponent e, it suffices to + // compare the significands. + const std::uint64_t rest = (std::uint64_t{p1} << -one.e) + p2; + if (rest <= delta) { + // V = buffer * 10^n, with M- <= V <= M+. + + decimal_exponent += n; + + // We may now just stop. But instead look if the buffer could be + // decremented to bring V closer to w. + // + // pow10 = 10^n is now 1 ulp in the decimal representation V. + // The rounding procedure works with diyfp's with an implicit + // exponent of e. + // + // 10^n = (10^n * 2^-e) * 2^e = ulp * 2^e + // + const std::uint64_t ten_n = std::uint64_t{pow10} << -one.e; + grisu2_round(buffer, length, dist, delta, rest, ten_n); + + return; + } + + pow10 /= 10; + // + // pow10 = 10^(n-1) <= p1 < 10^n + // Invariants restored. + } + + // 2) + // + // The digits of the integral part have been generated: + // + // M+ = d[k-1]...d[1]d[0] + p2 * 2^e + // = buffer + p2 * 2^e + // + // Now generate the digits of the fractional part p2 * 2^e. + // + // Note: + // No decimal point is generated: the exponent is adjusted instead. + // + // p2 actually represents the fraction + // + // p2 * 2^e + // = p2 / 2^-e + // = d[-1] / 10^1 + d[-2] / 10^2 + ... + // + // Now generate the digits d[-m] of p1 from left to right (m = 1,2,...) + // + // p2 * 2^e = d[-1]d[-2]...d[-m] * 10^-m + // + 10^-m * (d[-m-1] / 10^1 + d[-m-2] / 10^2 + ...) + // + // using + // + // 10^m * p2 = ((10^m * p2) div 2^-e) * 2^-e + ((10^m * p2) mod 2^-e) + // = ( d) * 2^-e + ( r) + // + // or + // 10^m * p2 * 2^e = d + r * 2^e + // + // i.e. + // + // M+ = buffer + p2 * 2^e + // = buffer + 10^-m * (d + r * 2^e) + // = (buffer * 10^m + d) * 10^-m + 10^-m * r * 2^e + // + // and stop as soon as 10^-m * r * 2^e <= delta * 2^e + + int m = 0; + for (;;) { + // Invariant: + // M+ = buffer * 10^-m + 10^-m * (d[-m-1] / 10 + d[-m-2] / 10^2 + ...) + // * 2^e + // = buffer * 10^-m + 10^-m * (p2 ) + // * 2^e = buffer * 10^-m + 10^-m * (1/10 * (10 * p2) ) * 2^e = + // buffer * 10^-m + 10^-m * (1/10 * ((10*p2 div 2^-e) * 2^-e + + // (10*p2 mod 2^-e)) * 2^e + // + p2 *= 10; + const std::uint64_t d = p2 >> -one.e; // d = (10 * p2) div 2^-e + const std::uint64_t r = p2 & (one.f - 1); // r = (10 * p2) mod 2^-e + // + // M+ = buffer * 10^-m + 10^-m * (1/10 * (d * 2^-e + r) * 2^e + // = buffer * 10^-m + 10^-m * (1/10 * (d + r * 2^e)) + // = (buffer * 10 + d) * 10^(-m-1) + 10^(-m-1) * r * 2^e + // + buffer[length++] = static_cast('0' + d); // buffer := buffer * 10 + d + // + // M+ = buffer * 10^(-m-1) + 10^(-m-1) * r * 2^e + // + p2 = r; + m++; + // + // M+ = buffer * 10^-m + 10^-m * p2 * 2^e + // Invariant restored. + + // Check if enough digits have been generated. + // + // 10^-m * p2 * 2^e <= delta * 2^e + // p2 * 2^e <= 10^m * delta * 2^e + // p2 <= 10^m * delta + delta *= 10; + dist *= 10; + if (p2 <= delta) { + break; + } + } + + // V = buffer * 10^-m, with M- <= V <= M+. + + decimal_exponent -= m; + + // 1 ulp in the decimal representation is now 10^-m. + // Since delta and dist are now scaled by 10^m, we need to do the + // same with ulp in order to keep the units in sync. + // + // 10^m * 10^-m = 1 = 2^-e * 2^e = ten_m * 2^e + // + const std::uint64_t ten_m = one.f; + grisu2_round(buffer, length, dist, delta, p2, ten_m); + + // By construction this algorithm generates the shortest possible decimal + // number (Loitsch, Theorem 6.2) which rounds back to w. + // For an input number of precision p, at least + // + // N = 1 + ceil(p * log_10(2)) + // + // decimal digits are sufficient to identify all binary floating-point + // numbers (Matula, "In-and-Out conversions"). + // This implies that the algorithm does not produce more than N decimal + // digits. + // + // N = 17 for p = 53 (IEEE double precision) + // N = 9 for p = 24 (IEEE single precision) +} + +/*! +v = buf * 10^decimal_exponent +len is the length of the buffer (number of decimal digits) +The buffer must be large enough, i.e. >= max_digits10. +*/ +inline void grisu2(char *buf, int &len, int &decimal_exponent, diyfp m_minus, + diyfp v, diyfp m_plus) { + // --------(-----------------------+-----------------------)-------- (A) + // m- v m+ + // + // --------------------(-----------+-----------------------)-------- (B) + // m- v m+ + // + // First scale v (and m- and m+) such that the exponent is in the range + // [alpha, gamma]. + + const cached_power cached = get_cached_power_for_binary_exponent(m_plus.e); + + const diyfp c_minus_k(cached.f, cached.e); // = c ~= 10^-k + + // The exponent of the products is = v.e + c_minus_k.e + q and is in the range + // [alpha,gamma] + const diyfp w = diyfp::mul(v, c_minus_k); + const diyfp w_minus = diyfp::mul(m_minus, c_minus_k); + const diyfp w_plus = diyfp::mul(m_plus, c_minus_k); + + // ----(---+---)---------------(---+---)---------------(---+---)---- + // w- w w+ + // = c*m- = c*v = c*m+ + // + // diyfp::mul rounds its result and c_minus_k is approximated too. w, w- and + // w+ are now off by a small amount. + // In fact: + // + // w - v * 10^k < 1 ulp + // + // To account for this inaccuracy, add resp. subtract 1 ulp. + // + // --------+---[---------------(---+---)---------------]---+-------- + // w- M- w M+ w+ + // + // Now any number in [M-, M+] (bounds included) will round to w when input, + // regardless of how the input rounding algorithm breaks ties. + // + // And digit_gen generates the shortest possible such number in [M-, M+]. + // Note that this does not mean that Grisu2 always generates the shortest + // possible number in the interval (m-, m+). + const diyfp M_minus(w_minus.f + 1, w_minus.e); + const diyfp M_plus(w_plus.f - 1, w_plus.e); + + decimal_exponent = -cached.k; // = -(-k) = k + + grisu2_digit_gen(buf, len, decimal_exponent, M_minus, w, M_plus); +} + +/*! +v = buf * 10^decimal_exponent +len is the length of the buffer (number of decimal digits) +The buffer must be large enough, i.e. >= max_digits10. +*/ +template +void grisu2(char *buf, int &len, int &decimal_exponent, FloatType value) { + static_assert(diyfp::kPrecision >= std::numeric_limits::digits + 3, + "internal error: not enough precision"); + + // If the neighbors (and boundaries) of 'value' are always computed for + // double-precision numbers, all float's can be recovered using strtod (and + // strtof). However, the resulting decimal representations are not exactly + // "short". + // + // The documentation for 'std::to_chars' + // (https://en.cppreference.com/w/cpp/utility/to_chars) says "value is + // converted to a string as if by std::sprintf in the default ("C") locale" + // and since sprintf promotes float's to double's, I think this is exactly + // what 'std::to_chars' does. On the other hand, the documentation for + // 'std::to_chars' requires that "parsing the representation using the + // corresponding std::from_chars function recovers value exactly". That + // indicates that single precision floating-point numbers should be recovered + // using 'std::strtof'. + // + // NB: If the neighbors are computed for single-precision numbers, there is a + // single float + // (7.0385307e-26f) which can't be recovered using strtod. The resulting + // double precision value is off by 1 ulp. +#if 0 + const boundaries w = compute_boundaries(static_cast(value)); +#else + const boundaries w = compute_boundaries(value); +#endif + + grisu2(buf, len, decimal_exponent, w.minus, w.w, w.plus); +} + +/*! +@brief appends a decimal representation of e to buf +@return a pointer to the element following the exponent. +@pre -1000 < e < 1000 +*/ +inline char *append_exponent(char *buf, int e) { + if (e < 0) { + e = -e; + *buf++ = '-'; + } else { + *buf++ = '+'; + } + + auto k = static_cast(e); + if (k < 10) { + // Always print at least two digits in the exponent. + // This is for compatibility with printf("%g"). + *buf++ = '0'; + *buf++ = static_cast('0' + k); + } else if (k < 100) { + *buf++ = static_cast('0' + k / 10); + k %= 10; + *buf++ = static_cast('0' + k); + } else { + *buf++ = static_cast('0' + k / 100); + k %= 100; + *buf++ = static_cast('0' + k / 10); + k %= 10; + *buf++ = static_cast('0' + k); + } + + return buf; +} + +/*! +@brief prettify v = buf * 10^decimal_exponent +If v is in the range [10^min_exp, 10^max_exp) it will be printed in fixed-point +notation. Otherwise it will be printed in exponential notation. +@pre min_exp < 0 +@pre max_exp > 0 +*/ +inline char *format_buffer(char *buf, int len, int decimal_exponent, + int min_exp, int max_exp) { + const int k = len; + const int n = len + decimal_exponent; + + // v = buf * 10^(n-k) + // k is the length of the buffer (number of decimal digits) + // n is the position of the decimal point relative to the start of the buffer. + + if (k <= n && n <= max_exp) { + // digits[000] + // len <= max_exp + 2 + + std::memset(buf + k, '0', static_cast(n) - static_cast(k)); + // Make it look like a floating-point number (#362, #378) + buf[n + 0] = '.'; + buf[n + 1] = '0'; + return buf + (static_cast(n)) + 2; + } + + if (0 < n && n <= max_exp) { + // dig.its + // len <= max_digits10 + 1 + std::memmove(buf + (static_cast(n) + 1), buf + n, + static_cast(k) - static_cast(n)); + buf[n] = '.'; + return buf + (static_cast(k) + 1U); + } + + if (min_exp < n && n <= 0) { + // 0.[000]digits + // len <= 2 + (-min_exp - 1) + max_digits10 + + std::memmove(buf + (2 + static_cast(-n)), buf, + static_cast(k)); + buf[0] = '0'; + buf[1] = '.'; + std::memset(buf + 2, '0', static_cast(-n)); + return buf + (2U + static_cast(-n) + static_cast(k)); + } + + if (k == 1) { + // dE+123 + // len <= 1 + 5 + + buf += 1; + } else { + // d.igitsE+123 + // len <= max_digits10 + 1 + 5 + + std::memmove(buf + 2, buf + 1, static_cast(k) - 1); + buf[1] = '.'; + buf += 1 + static_cast(k); + } + + *buf++ = 'e'; + return append_exponent(buf, n - 1); +} + +} // namespace dtoa_impl + +/*! +The format of the resulting decimal representation is similar to printf's %g +format. Returns an iterator pointing past-the-end of the decimal representation. +@note The input number must be finite, i.e. NaN's and Inf's are not supported. +@note The buffer must be large enough. +@note The result is NOT null-terminated. +*/ +char *to_chars(char *first, const char *last, double value) { + static_cast(last); // maybe unused - fix warning + + // bool negative = std::signbit(value); + bool negative = (*reinterpret_cast(&value)) & (1 << 31ull); + if (negative) { + value = -value; + *first++ = '-'; + } + +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wfloat-equal" +#endif + + if (value == 0) // +-0 + { + *first++ = '0'; + // Make it look like a floating-point number (#362, #378) + *first++ = '.'; + *first++ = '0'; + return first; + } + +#ifdef __clang__ +#pragma clang diagnostic pop +#endif + + // Compute v = buffer * 10^decimal_exponent. + // The decimal digits are stored in the buffer, which needs to be interpreted + // as an unsigned decimal integer. + // len is the length of the buffer, i.e. the number of decimal digits. + int len = 0; + int decimal_exponent = 0; + dtoa_impl::grisu2(first, len, decimal_exponent, value); + // Format the buffer like printf("%.*g", prec, value) + constexpr int kMinExp = -4; + constexpr int kMaxExp = std::numeric_limits::digits10; + + return dtoa_impl::format_buffer(first, len, decimal_exponent, kMinExp, + kMaxExp); +} +} // namespace internal +} // namespace simdjson +} // namespace minijson + +#endif // !MINIJSON_USE_STRTOD + +#endif // MINIJSON_IMPLEMENTATION + + +namespace safetensors { + +// Max header(JSON) size. 100 MB as done in original safetensors implementation. +constexpr size_t kMaxJSONSize = 1024ull * 1024ull * 100ull; + +namespace detail { + +#ifdef _WIN32 +std::wstring UTF8ToWchar(const std::string &str) { + int wstr_size = + MultiByteToWideChar(CP_UTF8, 0, str.data(), int(str.size()), nullptr, 0); + std::wstring wstr(size_t(wstr_size), 0); + MultiByteToWideChar(CP_UTF8, 0, str.data(), int(str.size()), &wstr[0], + int(wstr.size())); + return wstr; +} + +std::string WcharToUTF8(const std::wstring &wstr) { + int str_size = WideCharToMultiByte(CP_UTF8, 0, wstr.data(), int(wstr.size()), + nullptr, 0, nullptr, nullptr); + std::string str(size_t(str_size), 0); + WideCharToMultiByte(CP_UTF8, 0, wstr.data(), int(wstr.size()), &str[0], + int(str.size()), nullptr, nullptr); + return str; +} +#endif + +bool ReadWholeFile(std::vector *out, std::string *err, + const std::string &filepath, void *) { +#ifdef SAFETENSORS_CPP_ANDROID_LOAD_FROM_ASSETS + if (asset_manager) { + AAsset *asset = AAssetManager_open(asset_manager, filepath.c_str(), + AASSET_MODE_STREAMING); + if (!asset) { + if (err) { + (*err) += "File open error : " + filepath + "\n"; + } + return false; + } + size_t size = AAsset_getLength(asset); + if (size == 0) { + if (err) { + (*err) += "Invalid file size : " + filepath + + " (does the path point to a directory?)"; + } + return false; + } + out->resize(size); + AAsset_read(asset, reinterpret_cast(&out->at(0)), size); + AAsset_close(asset); + return true; + } else { + if (err) { + (*err) += "No asset manager specified : " + filepath + "\n"; + } + return false; + } +#else +#ifdef _WIN32 +#if defined(__GLIBCXX__) // mingw + int file_descriptor = + _wopen(UTF8ToWchar(filepath).c_str(), _O_RDONLY | _O_BINARY); + __gnu_cxx::stdio_filebuf wfile_buf(file_descriptor, std::ios_base::in); + std::istream f(&wfile_buf); +#elif defined(_MSC_VER) || defined(_LIBCPP_VERSION) + // For libcxx, assume _LIBCPP_HAS_OPEN_WITH_WCHAR is defined to accept + // `wchar_t *` + std::ifstream f(UTF8ToWchar(filepath).c_str(), std::ifstream::binary); +#else + // Unknown compiler/runtime + std::ifstream f(filepath.c_str(), std::ifstream::binary); +#endif +#else + std::ifstream f(filepath.c_str(), std::ifstream::binary); +#endif + if (!f) { + if (err) { + (*err) += "File open error : " + filepath + "\n"; + } + return false; + } + + // For directory(and pipe?), peek() will fail(Posix gnustl/libc++ only) + f.peek(); + if (!f) { + if (err) { + (*err) += + "File read error. Maybe empty file or invalid file : " + filepath + + "\n"; + } + return false; + } + + f.seekg(0, f.end); + size_t sz = static_cast(f.tellg()); + + // std::cout << "sz = " << sz << "\n"; + f.seekg(0, f.beg); + + if (int64_t(sz) < 0) { + if (err) { + (*err) += "Invalid file size : " + filepath + + " (does the path point to a directory?)"; + } + return false; + } else if (sz == 0) { + if (err) { + (*err) += "File is empty : " + filepath + "\n"; + } + return false; + } else if (sz >= (std::numeric_limits::max)()) { + if (err) { + (*err) += "Invalid file size : " + filepath + "\n"; + } + return false; + } + + out->resize(sz); + f.read(reinterpret_cast(&out->at(0)), + static_cast(sz)); + + return true; +#endif +} + +bool parse_metadata(const ::minijson::value &v, + ordered_dict &dst, std::string *err) { + if (auto po = v.as<::minijson::object>()) { + for (size_t i = 0; i < po->size(); i++) { + ::minijson::value ov; + if (!po->at(i, &ov)) { + if (err) { + (*err) += + "[Internal error] Invalid object found in __metadata__, at index " + std::to_string(i) + ".\n"; + } + return false; + } + + if (auto so = ov.as()) { + if (dst.count(po->keys()[i])) { + // This should not be happen though + if (err) { + (*err) += + "Duplicate key `" + po->keys()[i] + "` found in __metadata__.\n"; + } + return false; + } + + dst.insert(po->keys()[i], *so); + } else { + if (err) { + (*err) += "`" + po->keys()[i] + "` must be string value.\n"; + } + return false; + } + } + } else { + if (err) { + (*err) += "`__metadata__` value must be JSON object.\n"; + } + return false; + } + + return true; +} + +bool parse_dtype(const ::minijson::value &v, safetensors::dtype &dtype, + std::string *err) { + if (auto so = v.as()) { + if ((*so) == "BOOL") { + dtype = safetensors::dtype::kBOOL; + } else if ((*so) == "U8") { + dtype = safetensors::dtype::kUINT8; + } else if ((*so) == "I8") { + dtype = safetensors::dtype::kINT8; + } else if ((*so) == "U16") { + dtype = safetensors::dtype::kUINT16; + } else if ((*so) == "I16") { + dtype = safetensors::dtype::kINT16; + } else if ((*so) == "U32") { + dtype = safetensors::dtype::kUINT32; + } else if ((*so) == "I32") { + dtype = safetensors::dtype::kINT32; + } else if ((*so) == "U64") { + dtype = safetensors::dtype::kUINT64; + } else if ((*so) == "I64") { + dtype = safetensors::dtype::kINT64; + } else if ((*so) == "F16") { + dtype = safetensors::dtype::kFLOAT16; + } else if ((*so) == "BF16") { + dtype = safetensors::dtype::kBFLOAT16; + } else if ((*so) == "F32") { + dtype = safetensors::dtype::kFLOAT32; + } else if ((*so) == "F64") { + dtype = safetensors::dtype::kFLOAT64; + } else { + if (err) { + (*err) += "Unknown `dtype` string: " + *so + ".\n"; + } + return false; + } + } else { + if (err) { + (*err) += + "`dtype` item should be string type but got " + v.type_name() + ".\n"; + } + return false; + } + + return true; +} + +bool parse_shape(const ::minijson::value &v, std::vector &dst, + std::string *err) { + // NOTE: + // - Empty tensors (tensors with 1 dimension being 0) are allowed + // - [] is allowed(0-Rank tensor = merely a scalar) + if (auto pa = v.as<::minijson::array>()) { + ::minijson::array::const_iterator i; + + for (i = pa->begin(); i != pa->end(); i++) { + if (auto pn = i->as<::minijson::number>()) { + if (dst.size() >= kMaxDim) { + if (err) { + (*err) += "`shape` length must be less than " + + std::to_string(kMaxDim) + " but got " + + std::to_string(dst.size()) + ".\n"; + } + return false; + } + + dst.push_back(size_t(*pn)); + + } else { + if (err) { + (*err) += "Array item in `shape` must be number type, but got " + + i->type_name() + ".\n"; + } + return false; + } + } + } else { + if (err) { + (*err) += + "`shape` value must be JSON array, but got " + v.type_name() + ".\n"; + } + return false; + } + + return true; +} + +bool parse_data_offsets(const ::minijson::value &v, std::array &dst, + std::string *err) { + if (auto pa = v.as<::minijson::array>()) { + ::minijson::array::const_iterator i; + size_t cnt = 0; + + for (i = pa->begin(); i != pa->end(); i++) { + if (auto pn = i->as<::minijson::number>()) { + if (cnt >= 2) { + if (err) { + (*err) += "`data_offsets` length must be 2.\n"; + } + return false; + } + + dst[cnt] = size_t(*pn); + + cnt++; + + } else { + if (err) { + (*err) += + "Array item in `data_offsets` must be number type, but got " + + i->type_name() + ".\n"; + } + return false; + } + } + + if (cnt != 2) { + if (err) { + (*err) += "`data_offsets` length must be 2.\n"; + } + return false; + } + } else { + if (err) { + (*err) += "`data_offsets` value must be JSON array, but got " + + v.type_name() + ".\n"; + } + return false; + } + + return true; +} + +bool parse_tensor(const std::string &name, const ::minijson::value &v, + tensor_t &tensor, std::string *err) { + if (auto po = v.as<::minijson::object>()) { + + bool dtype_found{false}; + bool shape_found{false}; + bool data_offsets_found{false}; + + dtype dtype; + std::vector shape; + std::array data_offsets{}; + + for (size_t i = 0; i < po->size(); i++) { + std::string key = po->keys()[i]; + + if (key == "dtype") { + ::minijson::value value; + if (!po->at(i, &value)) { + if (err) { + (*err) += "Internal error. `dtype` has invalid object.\n"; + } + return false; + } + + if (!parse_dtype(value, dtype, err)) { + return false; + } + + dtype_found = true; + } else if (key == "shape") { + ::minijson::value value; + if (!po->at(i, &value)) { + if (err) { + (*err) += "Internal error. `shape` has invalid object.\n"; + } + return false; + } + + if (!parse_shape(value, shape, err)) { + return false; + } + + shape_found = true; + } else if (key == "data_offsets") { + ::minijson::value value; + if (!po->at(i, &value)) { + if (err) { + (*err) += "Internal error. `data_offsets` has invalid object.\n"; + } + return false; + } + if (!parse_data_offsets(value, data_offsets, err)) { + return false; + } + + data_offsets_found = true; + } else { + // Unknown key. Report error? + } + } + + if (!dtype_found) { + if (err) { + (*err) += "`" + name + "` does not have `dtype` item.\n"; + } + return false; + } + + if (!shape_found) { + if (err) { + (*err) += "`" + name + "` does not have `shape` item.\n"; + } + return false; + } + + bool is_empty_tensor{false}; + if ((shape.size() > 0)) { + for (size_t i = 0; i < shape.size(); i++) { + if (shape[i] == 0) { + is_empty_tensor = true; + break; + } + } + } + + if (is_empty_tensor) { + // They are not storing any data in the databuffer, yet retaining size in + // the header. So ignore data_offsets + if (data_offsets_found) { + // TODO: make this warn instead of err? + if (err) { + (*err) += + "`" + name + + "` is empty tensors(tensors with 1 dimension being 0), and no " + "data in databuffer, but `data_offsets` item is provided.\n"; + } + // DO NOT RETURN FALSE, JUST CONTINUE + } + } else { + if (!data_offsets_found) { + if (err) { + (*err) += "`" + name + "` does not have `data_offsets` item.\n"; + } + return false; + } + } + + tensor.dtype = dtype; + tensor.shape = shape; + tensor.data_offsets = data_offsets; + + } else { + if (err) { + (*err) += "`" + name + "` value must be JSON object.\n"; + } + return false; + } + + return true; +} + +// From llama.cpp +#if defined(_WIN32) +static std::string safetensors_format_win_err(DWORD err) { + LPSTR buf; + size_t size = FormatMessageA( + FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | + FORMAT_MESSAGE_IGNORE_INSERTS, + NULL, err, MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPSTR)&buf, 0, + NULL); + if (!size) { + return "FormatMessageA failed"; + } + std::string ret(buf, size); + LocalFree(buf); + return ret; +} +#endif + +struct safetensors_file { + // use FILE * so we don't have to re-open the file to mmap + FILE *fp{nullptr}; + size_t size{0}; + mutable bool _valid{false}; + std::string _err; + + safetensors_file(const char *fname, const char *mode) { + fp = std::fopen(fname, mode); + if (fp == nullptr) { + _err = "failed to open " + std::string(fname) + ":" + + std::string(strerror(errno)) + "\n"; + _valid = false; + } else { + seek(0, SEEK_END); + size = tell(); + seek(0, SEEK_SET); + _valid = true; + } + } + + ~safetensors_file() { + if (fp) { + std::fclose(fp); + fp = nullptr; + } + } + + size_t tell() const { +#ifdef _WIN32 + __int64 ret = _ftelli64(fp); +#else + long ret = std::ftell(fp); +#endif + if (ret == -1) { + // this really shouldn't fail + _valid = false; + return 0; + } + + return (size_t)ret; + } + + void seek(size_t offset, int whence) const { +#ifdef _WIN32 + int ret = _fseeki64(fp, (__int64)offset, whence); +#else + int ret = std::fseek(fp, (long)offset, whence); +#endif + if (ret == 0) { + _valid = false; + } + } + + bool &is_valid() const { return _valid; } + + const std::string &get_error() const { return _err; } +}; + +struct safetensors_mmap { + uint8_t *addr{nullptr}; + size_t size{0}; + + bool _valid{false}; + std::string _warn; + std::string _err; + + const bool is_valid() const { return _valid; } + + const std::string &get_error() const { return _err; } + + const std::string &get_warning() const { return _warn; } + + safetensors_mmap(const safetensors_mmap &) = delete; + +#ifdef _POSIX_MAPPED_FILES + static constexpr bool SUPPORTED = true; + + safetensors_mmap(struct safetensors_file *file, + size_t prefetch = (size_t)-1 /* -1 = max value */, + bool numa = false) { + size = file->size; + int fd = fileno(file->fp); + int flags = MAP_SHARED; + // prefetch/readahead impairs performance on NUMA systems + if (numa) { + prefetch = 0; + } +#ifdef __linux__ + if (prefetch) { + flags |= MAP_POPULATE; + } +#endif + addr = reinterpret_cast( + mmap(NULL, file->size, PROT_READ, flags, fd, 0)); + if (addr == MAP_FAILED) { + _valid = false; + _err = "mmap failed: " + std::string(strerror(errno)) + "\n"; + + size = 0; + addr = nullptr; + + return; + } + + if (prefetch > 0) { + // Advise the kernel to preload the mapped memory + if (posix_madvise(addr, std::min(file->size, prefetch), + POSIX_MADV_WILLNEED)) { + _warn += "posix_madvise(.., POSIX_MADV_WILLNEED) failed: " + + std::string(strerror(errno)) + "\n"; + } + } + if (numa) { + // advise the kernel not to use readahead + // (because the next page might not belong on the same node) + if (posix_madvise(addr, file->size, POSIX_MADV_RANDOM)) { + _warn += "posix_madvise(.., POSIX_MADV_RANDOM) failed: " + + std::string(strerror(errno)) + "\n"; + } + } + + _valid = true; + } + + ~safetensors_mmap() { + if (_valid) { + munmap(addr, size); + } + size = 0; + addr = nullptr; + _valid = false; + } + +#elif defined(_WIN32) + static constexpr bool SUPPORTED = true; + + safetensors_mmap(struct safetensors_file *file, bool prefetch = true, + bool numa = false) { + (void)numa; + + size = file->size; + + HANDLE hFile = (HANDLE)_get_osfhandle(_fileno(file->fp)); + + HANDLE hMapping = + CreateFileMappingA(hFile, NULL, PAGE_READONLY, 0, 0, NULL); + DWORD error = GetLastError(); + + if (hMapping == NULL) { + // TODO: get error message + _err = "CreateFileMappingA failed: " + safetensors_format_win_err(error) + + "\n"; + _valid = false; + size = 0; + addr = nullptr; + return; + } + + addr = reinterpret_cast( + MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0)); + error = GetLastError(); + CloseHandle(hMapping); + + if (addr == NULL) { + _err = + "MapViewOfFile failed: " + safetensors_format_win_err(error) + "\n"; + } + +#if _WIN32_WINNT >= _WIN32_WINNT_WIN8 + if (prefetch) { + // PrefetchVirtualMemory is only present on Windows 8 and above, so we + // dynamically load it + BOOL(WINAPI * pPrefetchVirtualMemory) + (HANDLE, ULONG_PTR, PWIN32_MEMORY_RANGE_ENTRY, ULONG); + HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll"); + + // may fail on pre-Windows 8 systems + pPrefetchVirtualMemory = + reinterpret_cast( + GetProcAddress(hKernel32, "PrefetchVirtualMemory")); + + if (pPrefetchVirtualMemory) { + // advise the kernel to preload the mapped memory + WIN32_MEMORY_RANGE_ENTRY range; + range.VirtualAddress = addr; + range.NumberOfBytes = (SIZE_T)size; + if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) { + _warn += "PrefetchVirtualMemory failed: " + + safetensors_format_win_err(GetLastError()) + "\n"; + } + } + } +#endif + } + ~safetensors_mmap() { + if (!UnmapViewOfFile(addr)) { + _warn += "UnmapViewOfFile failed: " + + safetensors_format_win_err(GetLastError()) + "\n"; + } + } +#else + static constexpr bool SUPPORTED = false; + + safetensors_mmap(struct safetensors_file *file, bool prefetch = true, + bool numa = false) { + (void)file; + (void)prefetch; + (void)numa; + + _valid = false; + _err = "mmap not supported\n"; + addr = nullptr; + size = 0; + } +#endif +}; + +// Based on MIOPen bfloat16 +// https://github.com/ROCmSoftwarePlatform/MIOpen/blob/master/src/kernels/bfloat16_dev.hpp + +/******************************************************************************* + * + * MIT License + * + * Copyright (c) 2019 Advanced Micro Devices, Inc. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + *all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + * SOFTWARE. + * + *******************************************************************************/ + +typedef union cvt_bf16_fp32 { + uint32_t u32; + uint16_t ushortvec[2]; + + float f32; +} cvt_bf16_fp32_t; + +float bfloat16_to_float(uint16_t src_val) { + cvt_bf16_fp32_t target_val; + + target_val.ushortvec[0] = 0; + target_val.ushortvec[1] = src_val; + + return target_val.f32; +} + +uint16_t float_to_bfloat16(float src_val) { + cvt_bf16_fp32_t target_val; + target_val.f32 = src_val; + // BF16 round and NaN preservation code matches + // https://github.com/ROCmSoftwarePlatform/rocBLAS/blob/develop/library/include/rocblas_bfloat16.h + if ((~target_val.u32 & 0x7f800000) == 0) // Inf or NaN + { + // When all of the exponent bits are 1, the value is Inf or NaN. + // Inf is indicated by a zero mantissa. NaN is indicated by any nonzero + // mantissa bit. Quiet NaN is indicated by the most significant mantissa + // bit being 1. Signaling NaN is indicated by the most significant + // mantissa bit being 0 but some other bit(s) being 1. If any of the + // lower 16 bits of the mantissa are 1, we set the least significant bit + // of the bfloat16 mantissa, in order to preserve signaling NaN in case + // the bloat16's mantissa bits are all 0. + if ((target_val.u32 & 0xffff) != 0) { + target_val.u32 |= 0x10000; // Preserve signaling NaN + } + } else { +#if 1 // MIOPEN_USE_RNE_BFLOAT16 + // When the exponent bits are not all 1s, then the value is zero, normal, + // or subnormal. We round the bfloat16 mantissa up by adding 0x7FFF, plus + // 1 if the least significant bit of the bfloat16 mantissa is 1 (odd). + // This causes the bfloat16's mantissa to be incremented by 1 if the 16 + // least significant bits of the float mantissa are greater than 0x8000, + // or if they are equal to 0x8000 and the least significant bit of the + // bfloat16 mantissa is 1 (odd). This causes it to be rounded to even when + // the lower 16 bits are exactly 0x8000. If the bfloat16 mantissa already + // has the value 0x7f, then incrementing it causes it to become 0x00 and + // the exponent is incremented by one, which is the next higher FP value + // to the unrounded bfloat16 value. When the bfloat16 value is subnormal + // with an exponent of 0x00 and a mantissa of 0x7F, it may be rounded up + // to a normal value with an exponent of 0x01 and a mantissa of 0x00. + // When the bfloat16 value has an exponent of 0xFE and a mantissa of 0x7F, + // incrementing it causes it to become an exponent of 0xFF and a mantissa + // of 0x00, which is Inf, the next higher value to the unrounded value. + target_val.u32 += (0x7fff + (target_val.ushortvec[1] & 1)); +#endif // MIOPEN_USE_RNE_BFLOAT16 + } + + return target_val.ushortvec[1]; +} + +// half <-> float conversion based on: https://gist.github.com/rygorous/2156668 +// (CC0 license) +// + +// Little endian +union FP32le { + unsigned int u; + float f; + struct { + unsigned int Mantissa : 23; + unsigned int Exponent : 8; + unsigned int Sign : 1; + } s; +}; + +// Little endian +union float16le { + unsigned short u; + struct { + unsigned int Mantissa : 10; + unsigned int Exponent : 5; + unsigned int Sign : 1; + } s; +}; + +float half_to_float_le(float16le h) { + static const FP32le magic = {113 << 23}; + static const unsigned int shifted_exp = 0x7c00 + << 13; // exponent mask after shift + FP32le o; + + o.u = (h.u & 0x7fffU) << 13U; // exponent/mantissa bits + unsigned int exp_ = shifted_exp & o.u; // just the exponent + o.u += (127 - 15) << 23; // exponent adjust + + // handle exponent special cases + if (exp_ == shifted_exp) // Inf/NaN? + o.u += (128 - 16) << 23; // extra exp adjust + else if (exp_ == 0) // Zero/Denormal? + { + o.u += 1 << 23; // extra exp adjust + o.f -= magic.f; // renormalize + } + + o.u |= (h.u & 0x8000U) << 16U; // sign bit + return o.f; +} + +uint16_t float_to_half_full_le(float _f) { + FP32le f; + f.f = _f; + float16le o = {0}; + + // Based on ISPC reference code (with minor modifications) + if (f.s.Exponent == 0) // Signed zero/denormal (which will underflow) + o.s.Exponent = 0; + else if (f.s.Exponent == 255) // Inf or NaN (all exponent bits set) + { + o.s.Exponent = 31; + o.s.Mantissa = f.s.Mantissa ? 0x200 : 0; // NaN->qNaN and Inf->Inf + } else // Normalized number + { + // Exponent unbias the single, then bias the halfp + int newexp = f.s.Exponent - 127 + 15; + if (newexp >= 31) // Overflow, return signed infinity + o.s.Exponent = 31; + else if (newexp <= 0) // Underflow + { + if ((14 - newexp) <= 24) // Mantissa might be non-zero + { + unsigned int mant = f.s.Mantissa | 0x800000; // Hidden 1 bit + o.s.Mantissa = mant >> (14 - newexp); + if ((mant >> (13 - newexp)) & 1) // Check for rounding + o.u++; // Round, might overflow into exp bit, but this is OK + } + } else { + o.s.Exponent = static_cast(newexp); + o.s.Mantissa = f.s.Mantissa >> 13; + if (f.s.Mantissa & 0x1000) // Check for rounding + o.u++; // Round, might overflow to inf, this is OK + } + } + + o.s.Sign = f.s.Sign; + + return o.u; +} + +bool parse_safetensors_header(const uint8_t *addr, const size_t nbytes, + const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err) { + if (nbytes < 16) { + if (err) { + (*err) += "Size is too short.\n"; + } + return false; + } + + uint64_t header_size{0}; + memcpy(reinterpret_cast(&header_size), addr, + sizeof(uint64_t)); + + if (header_size < 4) { + if (err) { + (*err) += "Header size is too short.\n"; + } + return false; + } + + if ((8 + header_size) > nbytes) { + if (err) { + (*err) += "Header size " + std::to_string(header_size) + + " + 8 exceeds input size " + std::to_string(nbytes) + " .\n"; + } + return false; + } + + if (header_size > kMaxJSONSize) { + if (err) { + (*err) += "Header JSON size exceeds the limit(" + + std::to_string(kMaxJSONSize) + ").\n"; + } + return false; + } + + // assume JSON data is small enough. + std::string json_str(reinterpret_cast(&addr[8]), header_size); + const char *p = json_str.c_str(); + + ::minijson::value v; + ::minijson::error e = ::minijson::parse(p, v); + + if (e != ::minijson::no_error) { + if (err) { + std::string json_err(::minijson::errstr(e)); + (*err) += "JSON parse error: " + json_err + "\n"; + } + + return false; + } + + ordered_dict tensors; + ordered_dict metadata; + + // root element must be dict. + if (auto po = v.as<::minijson::object>()) { + for (size_t i = 0; i < po->size(); i++) { + std::string key = po->keys()[i]; + + if (key == "__metadata__") { + ::minijson::value value; + if (!po->at(i, &value)) { + if (err) { + (*err) += "Internal error. Invalid object in __metadata__.\n"; + } + return false; + } + + if (!detail::parse_metadata(value, metadata, err)) { + return false; + } + } else { + // tensor + + if (tensors.count(key)) { + if (err) { + (*err) += "Duplicate key `" + key + "` found.\n"; + } + return false; + } + + ::minijson::value value; + if (!po->at(i, &value)) { + if (err) { + (*err) += "Internal error. Invalid object in `" + key + "`.\n"; + } + return false; + } + + tensor_t tensor; + if (!detail::parse_tensor(key, value, tensor, err)) { + return false; + } + + tensors.insert(key, std::move(tensor)); + } + } + } else { + if (err) { + (*err) += "JSON root elements must be object(dict)\n"; + } + } + + st->tensors = std::move(tensors); + st->metadata = std::move(metadata); + st->header_size = header_size; + +#if 0 + size_t databuffer_size = nbytes - header_size - 8; + + st->storage.resize(nbytes); + memcpy(st->storage.data(), addr + 8 + header_size, nbytes); + + st->mmaped = false; + st->mmap_addr = addr + 8 + header_size; + st->mmap_size = 0; +#endif + + return true; +} + +} // namespace detail + +safetensors_t::~safetensors_t() { + if (st_mmap) { + detail::safetensors_mmap *p = + reinterpret_cast(st_mmap); + delete p; + st_mmap = nullptr; + } + + if (st_file) { + detail::safetensors_file *p = + reinterpret_cast(st_file); + delete p; + st_file = nullptr; + } +} + +// +// - 8byte: header_size +// - json data(header_size bytes) +// - tensor data(filesize - header_size) +// + +bool load_from_file(const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err) { + std::vector data; + if (!detail::ReadWholeFile(&data, err, filename, nullptr)) { + return false; + } + + return load_from_memory(reinterpret_cast(data.data()), + data.size(), filename, st, warn, err); +} + +bool load_from_memory(const uint8_t *addr, const size_t nbytes, + const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err) { + if (nbytes < 16) { + if (err) { + (*err) += "Size is too short.\n"; + } + return false; + } + + if (!detail::parse_safetensors_header(addr, nbytes, filename, st, warn, + err)) { + return false; + } + + size_t databuffer_size = nbytes - st->header_size - 8; + + st->storage.resize(databuffer_size); + memcpy(st->storage.data(), addr + 8 + st->header_size, databuffer_size); + + st->mmaped = false; + st->mmap_addr = nullptr; + st->mmap_size = 0; + st->databuffer_addr = nullptr; + st->databuffer_size = 0; + + return true; +} + +bool mmap_from_file(const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err) { + if (!st) { + return false; + } + + detail::safetensors_file *pf = + new detail::safetensors_file(filename.c_str(), "rb"); + if (!pf->is_valid()) { + if (err) { + (*err) += pf->get_error(); + } + delete pf; + return false; + } + + // TODO: prefetch, numa + detail::safetensors_mmap *pm = new detail::safetensors_mmap(pf); + + bool ret = mmap_from_memory(pm->addr, pm->size, filename, st, warn, err); + + if (!ret) { + delete pm; + delete pf; + + return false; + } + + st->mmap_addr = pm->addr; + st->mmap_size = pm->size; + + st->databuffer_addr = st->mmap_addr + 8 + st->header_size; + st->databuffer_size = st->mmap_size - (8 + st->header_size); + + // retain pointer + st->st_file = pf; + st->st_mmap = pm; + + st->mmaped = true; + + return true; +} + +bool mmap_from_memory(const uint8_t *addr, const size_t nbytes, + const std::string &filename, safetensors_t *st, + std::string *warn, std::string *err) { + if (!addr) { + return false; + } + + if (nbytes < 16) { + return false; + } + + if (!st) { + return false; + } + + if (!detail::parse_safetensors_header(addr, nbytes, filename, st, warn, + err)) { + return false; + } + + size_t databuffer_size = nbytes - st->header_size - 8; + + st->mmaped = true; + + st->mmap_addr = addr; + st->mmap_size = nbytes; + + st->databuffer_addr = st->mmap_addr + 8 + st->header_size; + st->databuffer_size = st->mmap_size - (8 + st->header_size); + + return true; +} + +float bfloat16_to_float(uint16_t x) { return detail::bfloat16_to_float(x); } + +uint16_t float_to_bfloat16(float x) { return detail::float_to_bfloat16(x); } + +float fp16_to_float(uint16_t x) { + detail::float16le src; + src.u = x; + return detail::half_to_float_le(src); +} + +uint16_t float_to_fp16(float x) { return detail::float_to_half_full_le(x); } + +size_t get_dtype_bytes(const safetensors::dtype dtype) { + size_t sz = 0; + + switch (dtype) { + case safetensors::dtype::kBOOL: + // Original Rust implementaion uses 1. + sz = 1; + break; + case safetensors::dtype::kUINT8: + sz = 1; + break; + case safetensors::dtype::kINT8: + sz = 1; + break; + case safetensors::dtype::kUINT16: + sz = 2; + break; + case safetensors::dtype::kINT16: + sz = 2; + break; + case safetensors::dtype::kINT32: + sz = 4; + break; + case safetensors::dtype::kUINT32: + sz = 4; + break; + case safetensors::dtype::kFLOAT16: + sz = 2; + break; + case safetensors::dtype::kBFLOAT16: + sz = 2; + break; + case safetensors::dtype::kFLOAT32: + sz = 4; + break; + case safetensors::dtype::kFLOAT64: + sz = 8; + break; + case safetensors::dtype::kINT64: + sz = 8; + break; + case safetensors::dtype::kUINT64: + sz = 8; + break; + } + + return sz; +} + +std::string get_dtype_str(const safetensors::dtype dtype) { + switch (dtype) { + case safetensors::dtype::kBOOL: + return "BOOL"; + case safetensors::dtype::kUINT8: + return "U8"; + case safetensors::dtype::kINT8: + return "I8"; + case safetensors::dtype::kUINT16: + return "U16"; + case safetensors::dtype::kINT16: + return "I16"; + case safetensors::dtype::kINT32: + return "I32"; + case safetensors::dtype::kUINT32: + return "U32"; + case safetensors::dtype::kFLOAT16: + return "F16"; + case safetensors::dtype::kBFLOAT16: + return "BF16"; + case safetensors::dtype::kFLOAT32: + return "F32"; + case safetensors::dtype::kFLOAT64: + return "F64"; + case safetensors::dtype::kINT64: + return "I64"; + case safetensors::dtype::kUINT64: + return "U64"; + } + return "???"; +} + +// Empty Tensor returns 0. +// Zero-rank Tensor reuturns 1(scalar) +size_t get_shape_size(const tensor_t &t) { + if (t.shape.empty()) { + return 1; + } + + if (t.shape.size() >= kMaxDim) { // invalid ndim + return 0; + } + + size_t sz = 1; + + for (size_t i = 0; i < t.shape.size(); i++) { + sz *= t.shape[i]; + } + + return sz; +} + +bool validate_data_offsets(const safetensors_t &st, std::string &err) { + bool valid{true}; + + std::stringstream ss; + + size_t databuffersize; + if (st.mmaped) { + databuffersize = st.databuffer_size; + } else { + databuffersize = st.storage.size(); + } + + size_t ntensors{0}; + // Iterate with key insertion order. + for (size_t i =0 ;i < st.tensors.size(); i++) { + + std::string key = st.tensors.keys()[i]; + + tensor_t tensor; + if (!st.tensors.at(i, &tensor)) { + ss << "Internal error: Failed to get tensor at [" << i << "]\n"; + valid = false; + continue; + } + + if (tensor.data_offsets[0] > tensor.data_offsets[1]) { + ss << key << ".data_offsets.BEGIN " << tensor.data_offsets[0] + << " must be less than or equal to data_offsets.END " + << tensor.data_offsets[1] << "\n"; + valid = false; + } + + size_t tensor_size = get_dtype_bytes(tensor.dtype) * get_shape_size(tensor); + + if (tensor_size == 0) { + // OK + continue; + } + + // data_offsets are absolute offset from the databuffer(file) + if (tensor.data_offsets[0] > databuffersize) { + ss << "Tensor `" << key << "`.data_offset.BEGIN " + << tensor.data_offsets[0] << " exceeds databuffer size " + << databuffersize << ".\n"; + valid = false; + } + + if (tensor.data_offsets[1] > databuffersize) { + ss << "Tensor `" << key << "`.data_offset.END " + << tensor.data_offsets[1] << " exceeds databuffer size " + << databuffersize << ".\n"; + valid = false; + } + + size_t data_size = tensor.data_offsets[1] - tensor.data_offsets[0]; + + if (tensor_size != data_size) { + ss << "Data size mismatch. The size in Tensor `" << key << "` is " + << tensor_size << ", but the size from data_offsets is " << data_size + << "\n"; + valid = false; + } + + ntensors++; + if (ntensors == st.tensors.size()) { + // Last element's data_offsets[1] must be equal to databuffer size. + if (tensor.data_offsets[1] != databuffersize) { + ss << "The last tensor's data_offset.END(" << tensor.data_offsets[1] + << ") must be equal to databufer size " << databuffersize << ".\n"; + valid = false; + } + } + } + + if (!valid) { + err = ss.str(); + } + + return valid; +} + +bool save_to_memory(const safetensors_t &st, std::vector *dst, + std::string *warn, std::string *err) { + // directly serialize JSON string. + std::stringstream ss; + + // NOTE: The last offset **must** be the end of the file, + // so write __metadata__ first(if metadata part exists) + + std::string _err; + if (!validate_data_offsets(st, _err)) { + if (err) { + (*err) += "Invalid safensors is provided.\n"; + (*err) += _err; + } + return false; + } + + ss << "{"; + if (st.metadata.size()) { + ss << "\"__metadata__\": {"; + size_t nmeta = 0; + for (size_t i = 0; i < st.metadata.size(); i++) { + std::string key = st.metadata.keys()[i]; + std::string value; + st.metadata.at(i, &value); + + if (nmeta > 0) { + ss << ", "; + } + ss << "\"" + key + "\": \"" << value << "\""; + nmeta++; + } + ss << "}"; + + if (st.tensors.size()) { + ss << ", "; + } + } + + size_t ntensors = 0; + { + for (size_t i = 0; i < st.tensors.size(); i++) { + + std::string key = st.tensors.keys()[i]; + safetensors::tensor_t tensor; + st.tensors.at(i, &tensor); + + if (tensor.shape.size() > safetensors::kMaxDim) { + if (err) { + (*err) += key + ".shape is too large.\n"; + (*err) += _err; + } + return false; + } + + if (ntensors > 0) { + ss << ", "; + } + ss << "\"" << key << "\": {"; + ss << "\"dtype\": \"" << safetensors::get_dtype_str(tensor.dtype) + << "\", "; + ss << "\"shape\": ["; + for (size_t i = 0; i < tensor.shape.size(); i++) { + if (i > 0) { + ss << ", "; + } + ss << tensor.shape[i]; + } + ss << "]"; + ss << ", \"data_offsets\": [" << tensor.data_offsets[0] << ", " + << tensor.data_offsets[1] << "]"; + ss << "}"; + ntensors++; + } + } + ss << "}"; + + std::string header_str = ss.str(); + + uint64_t header_size = header_str.size(); // do not include '\n' + + const void *databuffer_addr{nullptr}; + size_t databuffer_size{0}; + if (st.mmaped) { + databuffer_size = st.databuffer_size; + databuffer_addr = st.databuffer_addr; + } else { + databuffer_size = st.storage.size(); + databuffer_addr = reinterpret_cast(st.storage.data()); + } + + // make databuffer addr start from the multiple of 8. + size_t pad_bytes = 0; + if ((header_size % 8) != 0) { + pad_bytes = 8 - (header_size % 8); + } + //printf("header_size = %d\n", int(header_size)); + //printf("pad_bytes = %d\n", int(pad_bytes)); + size_t padded_header_size = header_size + pad_bytes; + dst->resize(8 + padded_header_size + databuffer_size); + + // write padded header_size + memcpy(dst->data(), &padded_header_size, 8); + + // write header + memcpy(dst->data() + 8, header_str.data(), header_size); + + // Use whitespace for trailing padding. + memset(dst->data() + 8 + header_size, 0x20, pad_bytes); + + memcpy(dst->data() + 8 + padded_header_size, databuffer_addr, + databuffer_size); + + return true; +} + +bool save_to_file(const safetensors_t &st, const std::string &filename, + std::string *warn, std::string *err) { + // TODO: Use more reliable io. + std::ofstream ofs(filename, std::ios::binary); + + if (!ofs) { + if (err) { + (*err) += "Failed to open `" + filename + + "` to write. File is either existing directory or " + "write-protected, or disk is full?\n"; + } + return false; + } + + std::vector buf; + if (!save_to_memory(st, &buf, warn, err)) { + return false; + } + + ofs.write(reinterpret_cast(buf.data()), buf.size()); + if (!ofs) { + if (err) { + (*err) += "Failed to write safetensor data to `" + filename + + "`. Maybe no disk space available?(Required bytes : " + + std::to_string(buf.size()) + "\n"; + } + return false; + } + + return true; +} + +} // namespace safetensors + +#endif diff --git a/transformers/README.md b/transformers/README.md index 24ff424ceb..01fd10b927 100644 --- a/transformers/README.md +++ b/transformers/README.md @@ -57,6 +57,26 @@ The directory structure is as follows: + Direct Conversion to MNN Model Use `--export mnn` to directly convert to an MNN model. Note that you need to either install pymnn or specify the path to the MNNConvert tool using the `--mnnconvert` option. At least one of these conditions must be met. If pymnn is not installed and the MNNConvert tool's path is not specified via --mnnconvert, the llmexport.py script will search for the MNNConvert tool in the directory "../../../build/". Ensure that the MNNConvert file exists in this directory. This method currently supports exporting 4-bit and 8-bit models. ++ Segment MNN Export +Use `--export mnn --segment` to export a segment-format MNN LLM directly from safetensors weights and a workflow JSON, without generating ONNX first. If `--workflow` is not specified, `llmexport.py` searches `resource/*.json` for a matching workflow. + +``` +cd transformers/llm/export +python3 llmexport.py \ + --path /path/to/Qwen3-0.6B \ + --export mnn \ + --segment \ + --dst_path ./model +``` + +The output directory contains `config.json` with `"mnn_llm_version": "segment"`, `llm_config.json`, `tokenizer.mtok`, `embed.mnn`, `decoder.mnn`, `decoder.mnn.weight`, `logit.mnn`, `logit.mnn.weight`, and `logit_topkv_1.mnn`. Run the segment model with the generated `config.json`: + +``` +./llm_demo transformers/llm/export/model/config.json /path/to/prompt.txt +``` + +The C++ runtime must be built with `MNN_BUILD_LLM=ON` and `MNN_LLM_SUPPORT_SEGMENT=ON` (enabled by default). Segment export currently supports `--export mnn` only. + + If you encounter issues with directly converting to an MNN model or require quantization with other bit depths (e.g., 5-bit/6-bit), you can first convert the model to an ONNX model using `--export onnx`. Then, use the MNNConvert tool to convert the ONNX model to an MNN model with the following command: ``` @@ -72,7 +92,7 @@ Use `--export mnn` to directly convert to an MNN model. Note that you need to ei ``` usage: llmexport.py [-h] --path PATH [--type TYPE] [--lora_path LORA_PATH] [--dst_path DST_PATH] [--test TEST] [--export EXPORT] [--quant_bit QUANT_BIT] [--quant_block QUANT_BLOCK] [--lm_quant_bit LM_QUANT_BIT] - [--mnnconvert MNNCONVERT] + [--mnnconvert MNNCONVERT] [--segment] [--workflow WORKFLOW] llm_exporter @@ -89,6 +109,8 @@ options: --dst_path DST_PATH export onnx/mnn model to path, default is `./model`. --test TEST test model inference with query `TEST`. --export EXPORT export model to an onnx/mnn model. + --segment export segment MNN LLM from safetensors workflow directly, without ONNX export. + --workflow WORKFLOW workflow json for --segment safetensors conversion. If absent, search resource/*.json. --quant_bit QUANT_BIT mnn quant bit, 4 or 8, default is 4. --quant_block QUANT_BLOCK diff --git a/transformers/llm/engine/CMakeLists.txt b/transformers/llm/engine/CMakeLists.txt index 1b005a7858..8fbe34d874 100644 --- a/transformers/llm/engine/CMakeLists.txt +++ b/transformers/llm/engine/CMakeLists.txt @@ -1,6 +1,7 @@ option(BUILD_MLS "Build PC Commandline." OFF) option(MNN_LLM_BUILD_DEMO "Build LLM demo" ON) option(LLM_SUPPORT_HTTP_RESOURCE "Support HTTP resource download" ON) +option(MNN_LLM_SUPPORT_SEGMENT "Enable mnn_llm_version=segment runtime path." ON) set(LLM_DEPS ${MNN_DEPS}) if (MNN_BUILD_OPENCV) @@ -44,6 +45,9 @@ else() endif() # jinja.cpp template engine (always enabled, header-only) target_compile_definitions(llm PRIVATE LLM_USE_JINJA) +if (MNN_LLM_SUPPORT_SEGMENT) + target_compile_definitions(llm PRIVATE MNN_LLM_SUPPORT_SEGMENT) +endif() # Option to store MNN_PRINT/MNN_ERROR output into a string buffer. # Only enabled on Android. Pass -DLLM_LOG_TO_STRING=ON at cmake configure time. @@ -166,4 +170,4 @@ set_property(TARGET mls PROPERTY CXX_STANDARD_REQUIRED ON) # target_compile_options(mls PRIVATE -std=c++17) target_link_libraries(mls PRIVATE ${LLM_DEPS}) target_compile_definitions(mls PRIVATE CPPHTTPLIB_OPENSSL_SUPPORT) -endif() \ No newline at end of file +endif() diff --git a/transformers/llm/engine/include/llm/llm.hpp b/transformers/llm/engine/include/llm/llm.hpp index b7a06db562..ff0e52da0a 100644 --- a/transformers/llm/engine/include/llm/llm.hpp +++ b/transformers/llm/engine/include/llm/llm.hpp @@ -164,7 +164,7 @@ class MNN_PUBLIC Llm { void response(const ChatMessages& chat_prompts, std::ostream* os = &std::cout, const char* end_with = nullptr, int max_new_tokens = -1); void response(MNN::Express::VARP input_embeds, std::ostream* os = &std::cout, const char* end_with = nullptr, int max_new_tokens = -1); virtual void generate_init(std::ostream* os = nullptr, const char* end_with = nullptr); - void generate(int max_token); + virtual void generate(int max_token); std::vector generate(const std::vector& input_ids, int max_new_tokens = -1); std::vector generate(MNN::Express::VARP input_embeds, int max_tokens = -1); bool stoped(); @@ -278,4 +278,4 @@ class MNN_PUBLIC Embedding : public Llm { } } -#endif // LLM_hpp \ No newline at end of file +#endif // LLM_hpp diff --git a/transformers/llm/engine/src/llm.cpp b/transformers/llm/engine/src/llm.cpp index 2f163a4906..4e2c21c51a 100644 --- a/transformers/llm/engine/src/llm.cpp +++ b/transformers/llm/engine/src/llm.cpp @@ -23,6 +23,9 @@ #include "diskembedding.hpp" #include "sampler.hpp" #include "omni.hpp" +#ifdef MNN_LLM_SUPPORT_SEGMENT +#include "segment.hpp" +#endif #include "speculative_decoding/generate.hpp" #include "core/MNNFileUtils.h" @@ -87,6 +90,16 @@ static inline void _llmOrigError(const char* msg) { Llm* Llm::createLLM(const std::string& config_path) { std::shared_ptr config(new LlmConfig(config_path)); Llm* llm = nullptr; + const auto llmVersion = config->mnn_llm_version(); +#ifdef MNN_LLM_SUPPORT_SEGMENT + if (llmVersion == "segment") { + return createSegmentLlm(config); + } +#else + if (llmVersion == "segment") { + MNN_ERROR("[Error]: mnn_llm_version=segment requires MNN_LLM_SUPPORT_SEGMENT.\n"); + } +#endif if (config->is_visual() || config->is_audio() || config->has_talker()) { llm = new Omni(config); } else { @@ -1617,4 +1630,4 @@ bool Llm::is_stop(int token_id) { return stop; } } // namespace Transformer -} // namespace MNN \ No newline at end of file +} // namespace MNN diff --git a/transformers/llm/engine/src/llmconfig.hpp b/transformers/llm/engine/src/llmconfig.hpp index e95c672352..3cad3af05e 100644 --- a/transformers/llm/engine/src/llmconfig.hpp +++ b/transformers/llm/engine/src/llmconfig.hpp @@ -81,6 +81,18 @@ class LlmConfig { } else { config_ = ujson::json::parse("{}"); base_dir_ = path; + if (!base_dir_.empty() && base_dir_.back() != '/' && base_dir_.back() != '\\') { + base_dir_ += "/"; + } + std::ifstream config_file(base_dir_ + "config.json"); + if (config_file.is_open()) { + std::ostringstream ostr; + ostr << config_file.rdbuf(); + auto model_config = ujson::json::parse(ostr.str()); + if (model_config.contains("mnn_llm_version")) { + config_.merge(model_config); + } + } } } // using config's base_dir @@ -146,6 +158,10 @@ class LlmConfig { std::string context_file() const { return base_dir_ + config_.value("context_file", "context.json"); } + + std::string mnn_llm_version() const { + return config_.value("mnn_llm_version", ""); + } // model file config end > // < generate config start @@ -634,4 +650,4 @@ class LlmConfig { } // Transformer } // MNN -#endif \ No newline at end of file +#endif diff --git a/transformers/llm/engine/src/segment.cpp b/transformers/llm/engine/src/segment.cpp new file mode 100644 index 0000000000..e64f961364 --- /dev/null +++ b/transformers/llm/engine/src/segment.cpp @@ -0,0 +1,505 @@ +#ifdef MNN_LLM_SUPPORT_SEGMENT + +#include "segment.hpp" + +#include +#include +#include + +#include +#include "core/MNNFileUtils.h" +#include "kvmeta.hpp" +#include "llmconfig.hpp" +#include "tokenizer/tokenizer.hpp" + +namespace MNN { +namespace Transformer { +namespace { + +using namespace Express; +using RuntimeManager = Express::Executor::RuntimeManager; + +static bool segmentCheckFile(const std::string& path, const char* name) { + if (!MNNFileExist(path.c_str())) { + MNN_ERROR("[Error]: segment %s not found: %s\n", name, path.c_str()); + return false; + } + std::ifstream f(path); + if (!f.is_open()) { + MNN_ERROR("[Error]: failed to open segment %s: %s\n", name, path.c_str()); + return false; + } + return true; +} + +static std::string segmentPath(const LlmConfig& config, const std::string& name) { + return config.base_dir_ + name; +} + +static MNNForwardType segmentForwardType(std::shared_ptr config) { + if (config->config_.contains("forwardtype")) { + return static_cast(config->config_.value("forwardtype", 0)); + } + const auto type = config->backend_type(); + if (type == "metal") + return MNN_FORWARD_METAL; + if (type == "cuda") + return MNN_FORWARD_CUDA; + if (type == "opencl") + return MNN_FORWARD_OPENCL; + if (type == "opengl") + return MNN_FORWARD_OPENGL; + if (type == "vulkan") + return MNN_FORWARD_VULKAN; + if (type == "npu") + return MNN_FORWARD_NN; + return MNN_FORWARD_CPU; +} + +static void segmentApplyBackendConfig(std::shared_ptr config, BackendConfig* backend) { + if (backend == nullptr) { + return; + } + if (config->config_.contains("precision")) { + backend->precision = static_cast(config->config_.value("precision", 2)); + } else if (config->precision() == "high") { + backend->precision = BackendConfig::Precision_High; + } else if (config->precision() == "low") { + backend->precision = BackendConfig::Precision_Low; + } + if (config->config_.contains("memory")) { + backend->memory = static_cast(config->config_.value("memory", 2)); + } else if (config->memory() == "high") { + backend->memory = BackendConfig::Memory_High; + } else if (config->memory() == "low") { + backend->memory = BackendConfig::Memory_Low; + } +} + +static VARP segmentTakeLastHidden(VARP hidden) { + if (hidden == nullptr) { + return nullptr; + } + auto info = hidden->getInfo(); + if (info == nullptr || info->dim.size() < 3) { + return hidden; + } + const int seqLen = info->dim[1]; + const int hiddenSize = info->dim[2]; + if (seqLen <= 0 || hiddenSize <= 0 || seqLen == 1) { + return hidden; + } + const size_t bytes = static_cast(hiddenSize) * info->type.bytes(); + const uint8_t* src = hidden->readMap(); + if (src == nullptr || bytes == 0) { + return hidden; + } + auto out = _Input({1, 1, hiddenSize}, info->order, info->type); + ::memcpy(out->writeMap(), src + static_cast(seqLen - 1) * bytes, bytes); + out.fix(VARP::CONSTANT); + return out; +} + +static void segmentWait(VARP var) { + if (var == nullptr || var->getTensor() == nullptr) { + return; + } + ((MNN::Tensor*)var->getTensor())->wait(MNN::Tensor::MAP_TENSOR_READ, true); +} + +} // namespace + +class SegmentLlm final : public Llm { +public: + explicit SegmentLlm(std::shared_ptr config) : Llm(config) { + mSeqLenIndex = 1; + mMeta->layer_nums = mConfig->config_.value("layer_nums", 0); + } + + bool load() override; + VARP embedding(const std::vector& input_ids) override; + VARP gen_attention_mask(int seq_len) override; + VARP gen_position_ids(int seq_len) override; + std::vector forwardRaw(VARP hiddenState, VARP mask, VARP inputPos, VARPS extraArgs = {}) override; + int sample(VARP logits, int offset = 0, int size = 0) override; + void response(const std::vector& input_ids, std::ostream* os = &std::cout, const char* end_with = nullptr, + int max_new_tokens = -1) override; + void generate(int max_token) override; + +private: + bool loadTokenizer(); + bool loadModules(); + bool prefill(const std::vector& input_ids); + int sampleFromHidden(VARP hidden); + VARP embeddingToken(int token); + VARP decodeAttentionMask(); + VARP decodePositionId(); + VARP decoderForward(VARP input, VARP mask = nullptr, VARP positionIds = nullptr); + void updateSegmentContext(int seqLen, int genLen); + +private: + std::shared_ptr mEmbedModule; + std::shared_ptr mDecoderModule; + std::shared_ptr mDecoderPrefillModule; + std::shared_ptr mLogitBaseModule; + std::shared_ptr mLogitModule; + VARP mLastHidden; + VARP mTokenInput; + VARP mDecodeMaskInput; + VARP mDecodePositionInput; + int mMaxDecodeTokens = 1024; +}; + +bool SegmentLlm::loadTokenizer() { + std::string tokenizerPath = mConfig->tokenizer_file(); + if (!segmentCheckFile(tokenizerPath, "tokenizer file")) { + return false; + } + mTokenizer.reset(Tokenizer::createTokenizer(tokenizerPath)); + if (mTokenizer == nullptr) { + MNN_ERROR("[Error]: failed to load segment tokenizer: %s\n", tokenizerPath.c_str()); + return false; + } + + if (mConfig->config_.contains("jinja")) { + setChatTemplate(); + return true; + } + + std::ifstream tokenConfig(segmentPath(*mConfig, "token_config.json")); + if (tokenConfig.is_open()) { + std::ostringstream ostr; + ostr << tokenConfig.rdbuf(); + auto json = ujson::json::parse(ostr.str()); + if (json.contains("chat_template")) { + mTokenizer->set_chat_template(json["chat_template"].get(), json.value("eos_token", "")); + } + } + return true; +} + +bool SegmentLlm::loadModules() { + const std::string embedPath = segmentPath(*mConfig, "embed.mnn"); + const std::string decoderPath = segmentPath(*mConfig, "decoder.mnn"); + const std::string decoderWeightPath = decoderPath + ".weight"; + const std::string logitPath = segmentPath(*mConfig, "logit.mnn"); + const std::string logitWeightPath = logitPath + ".weight"; + const std::string logitTopkPath = segmentPath(*mConfig, "logit_topkv_1.mnn"); + + if (!segmentCheckFile(embedPath, "embed model") || !segmentCheckFile(decoderPath, "decoder model") || + !segmentCheckFile(decoderWeightPath, "decoder weight") || !segmentCheckFile(logitPath, "logit model") || + !segmentCheckFile(logitWeightPath, "logit weight") || !segmentCheckFile(logitTopkPath, "logit topk model")) { + return false; + } + + BackendConfig backendConfig; + segmentApplyBackendConfig(mConfig, &backendConfig); + + ScheduleConfig decoderSchedule; + decoderSchedule.backendConfig = &backendConfig; + decoderSchedule.type = segmentForwardType(mConfig); + decoderSchedule.numThread = mConfig->config_.value("thread_num", 1); + if (decoderSchedule.type == MNN_FORWARD_OPENCL) { + decoderSchedule.numThread |= 64; + } + + mRuntimeManager.reset(RuntimeManager::createRuntimeManager(decoderSchedule), RuntimeManager::destroy); + mRuntimeManager->setHintPtr(Interpreter::KVCACHE_INFO, mMeta.get()); + Module::Config decoderConfig; + decoderConfig.rearrange = true; + if (decoderSchedule.type == MNN_FORWARD_OPENCL || decoderSchedule.type == MNN_FORWARD_VULKAN) { + decoderConfig.shapeMutable = false; + } + + const std::vector decoderInputs = {"input_embedding", "mask", "position_ids"}; + mDecoderModule.reset( + Module::load(decoderInputs, {"last_hidden_state"}, decoderPath.c_str(), mRuntimeManager, &decoderConfig)); + if (!mDecoderModule) { + mDecoderModule.reset( + Module::load(decoderInputs, {"hidden_state"}, decoderPath.c_str(), mRuntimeManager, &decoderConfig)); + } + if (!mDecoderModule) { + MNN_ERROR("[Error]: load segment decoder.mnn failed\n"); + return false; + } + mDecoderPrefillModule.reset(Module::clone(mDecoderModule.get())); + if (!mDecoderPrefillModule) { + MNN_ERROR("[Error]: clone segment decoder prefill module failed\n"); + return false; + } + + BackendConfig otherBackendConfig = backendConfig; + if (mConfig->config_.contains("otherPrecision")) { + otherBackendConfig.precision = static_cast( + mConfig->config_.value("otherPrecision", (int)otherBackendConfig.precision)); + } + ScheduleConfig otherSchedule = decoderSchedule; + otherSchedule.backendConfig = &otherBackendConfig; + mProcessorRuntimeManager.reset(RuntimeManager::createRuntimeManager(otherSchedule), RuntimeManager::destroy); + mProcessorRuntimeManager->setHintPtr(Interpreter::KVCACHE_INFO, nullptr); + + Module::Config moduleConfig; + moduleConfig.rearrange = true; + mLogitBaseModule.reset(Module::load({}, {}, logitPath.c_str(), mProcessorRuntimeManager, &moduleConfig)); + if (!mLogitBaseModule) { + MNN_ERROR("[Error]: load segment logit.mnn failed\n"); + return false; + } + + Module::Config depConfig = moduleConfig; + depConfig.base = mLogitBaseModule.get(); + mLogitModule.reset(Module::load({}, {}, logitTopkPath.c_str(), mProcessorRuntimeManager, &depConfig)); + mEmbedModule.reset(Module::load({}, {}, embedPath.c_str(), mProcessorRuntimeManager, &depConfig)); + if (!mLogitModule || !mEmbedModule) { + MNN_ERROR("[Error]: load segment logit_topkv_1.mnn/embed.mnn failed\n"); + return false; + } + return true; +} + +bool SegmentLlm::load() { + MNN::Express::ExecutorScope s(mExecutor); + Timer _t; + mMaxDecodeTokens = mConfig->config_.value("max_decode_tokens", mConfig->max_new_tokens()); + if (!loadTokenizer() || !loadModules()) { + return false; + } + mContext->load_us += _t.durationInUs(); + mContext->status = LlmStatus::RUNNING; + return true; +} + +VARP SegmentLlm::embedding(const std::vector& input_ids) { + MNN::Express::ExecutorScope s(mExecutor); + if (input_ids.empty() || !mEmbedModule) { + return nullptr; + } + auto var = _Input({1, static_cast(input_ids.size())}, NCHW, halide_type_of()); + ::memcpy(var->writeMap(), input_ids.data(), input_ids.size() * sizeof(int)); + auto outputs = mEmbedModule->onForward({var}); + return outputs.empty() ? nullptr : outputs[0]; +} + +VARP SegmentLlm::embeddingToken(int token) { + MNN::Express::ExecutorScope s(mExecutor); + if (!mEmbedModule) { + return nullptr; + } + if (mTokenInput == nullptr) { + mTokenInput = _Input({1, 1}, NCHW, halide_type_of()); + } + *mTokenInput->writeMap() = token; + auto outputs = mEmbedModule->onForward({mTokenInput}); + return outputs.empty() ? nullptr : outputs[0]; +} + +VARP SegmentLlm::gen_attention_mask(int seq_len) { + auto mask = _Input({}, NCHW, halide_type_of()); + *mask->writeMap() = 0.0f; + mask.fix(VARP::CONSTANT); + return mask; +} + +VARP SegmentLlm::gen_position_ids(int seq_len) { + auto positionIds = _Input({1, seq_len}, NCHW, halide_type_of()); + auto ptr = positionIds->writeMap(); + const int start = static_cast(mMeta->previous) - static_cast(mMeta->remove); + for (int i = 0; i < seq_len; ++i) { + ptr[i] = start + i; + } + positionIds.fix(VARP::CONSTANT); + return positionIds; +} + +VARP SegmentLlm::decodeAttentionMask() { + if (mDecodeMaskInput == nullptr) { + mDecodeMaskInput = _Input({}, NCHW, halide_type_of()); + *mDecodeMaskInput->writeMap() = 0.0f; + mDecodeMaskInput.fix(VARP::CONSTANT); + } + return mDecodeMaskInput; +} + +VARP SegmentLlm::decodePositionId() { + if (mDecodePositionInput == nullptr) { + mDecodePositionInput = _Input({1, 1}, NCHW, halide_type_of()); + } + const int start = static_cast(mMeta->previous) - static_cast(mMeta->remove); + *mDecodePositionInput->writeMap() = start; + return mDecodePositionInput; +} + +VARP SegmentLlm::decoderForward(VARP input, VARP mask, VARP positionIds) { + if (input == nullptr) { + return nullptr; + } + auto info = input->getInfo(); + if (info == nullptr || info->dim.size() < 3) { + return nullptr; + } + const int seqLen = info->dim[1]; + if (mask == nullptr) { + mask = (seqLen == 1) ? decodeAttentionMask() : gen_attention_mask(seqLen); + } + if (positionIds == nullptr) { + positionIds = (seqLen == 1) ? decodePositionId() : gen_position_ids(seqLen); + } + auto module = (seqLen == 1) ? mDecoderModule : mDecoderPrefillModule; + mMeta->add = seqLen; + auto outputs = module->onForward({input, mask, positionIds}); + mMeta->sync(); + if (outputs.empty()) { + mContext->status = LlmStatus::INTERNAL_ERROR; + return nullptr; + } + segmentWait(outputs[0]); + return outputs[0]; +} + +std::vector SegmentLlm::forwardRaw(VARP hiddenState, VARP mask, VARP inputPos, VARPS extraArgs) { + auto hidden = decoderForward(hiddenState, mask, inputPos); + if (hidden == nullptr || !mLogitModule) { + mContext->status = LlmStatus::INTERNAL_ERROR; + return {}; + } + mLastHidden = segmentTakeLastHidden(hidden); + auto outputs = mLogitModule->onForward({hidden}); + if (outputs.empty()) { + mContext->status = LlmStatus::INTERNAL_ERROR; + return {}; + } + return outputs; +} + +int SegmentLlm::sample(VARP logits, int offset, int size) { + if (logits == nullptr) { + return -1; + } + auto info = logits->getInfo(); + if (info == nullptr || info->size <= 0) { + return -1; + } + const int* topk = logits->readMap(); + return topk == nullptr ? -1 : topk[info->size - 1]; +} + +int SegmentLlm::sampleFromHidden(VARP hidden) { + if (hidden == nullptr || !mLogitModule) { + return -1; + } + auto outputs = mLogitModule->onForward({hidden}); + if (outputs.empty()) { + return -1; + } + return sample(outputs[0]); +} + +void SegmentLlm::updateSegmentContext(int seqLen, int genLen) { + mContext->all_seq_len += seqLen; + mContext->gen_seq_len += genLen; +} + +bool SegmentLlm::prefill(const std::vector& input_ids) { + if (input_ids.empty()) { + return false; + } + mContext->history_tokens.insert(mContext->history_tokens.end(), input_ids.begin(), input_ids.end()); + Timer _t; + auto emb = embedding(input_ids); + auto hidden = decoderForward(emb); + if (hidden == nullptr) { + mContext->status = LlmStatus::INTERNAL_ERROR; + return false; + } + mLastHidden = segmentTakeLastHidden(hidden); + if (mLastHidden.get() != nullptr) { + mLastHidden.fix(VARP::CONSTANT); + } + updateSegmentContext(static_cast(input_ids.size()), 0); + mContext->prompt_len = static_cast(input_ids.size()); + mContext->prefill_us += _t.durationInUs(); + return true; +} + +void SegmentLlm::generate(int max_token) { + CHECK_LLM_RUNNING(mContext); + MNN::Express::ExecutorScope s(mExecutor); + if (max_token < 0) { + max_token = mMaxDecodeTokens; + } + max_token = std::min(max_token, mMaxDecodeTokens); + int len = 0; + while (len < max_token) { + if (mContext->status == LlmStatus::USER_CANCEL || mContext->status == LlmStatus::INTERNAL_ERROR) { + break; + } + Timer _t; + int token = sampleFromHidden(mLastHidden); + if (token < 0) { + mContext->decode_us += _t.durationInUs(); + mContext->status = LlmStatus::INTERNAL_ERROR; + break; + } + mContext->current_token = token; + if (is_stop(token)) { + mContext->decode_us += _t.durationInUs(); + if (mContext->os != nullptr) { + *mContext->os << mContext->end_with << std::flush; + } + break; + } + + mContext->history_tokens.push_back(token); + mContext->output_tokens.push_back(token); + auto decodeStr = tokenizer_decode(token); + mContext->generate_str += decodeStr; + if (mContext->os != nullptr) { + *mContext->os << decodeStr << std::flush; + } + + auto emb = embeddingToken(token); + auto hidden = decoderForward(emb); + if (hidden == nullptr) { + mContext->decode_us += _t.durationInUs(); + mContext->status = LlmStatus::INTERNAL_ERROR; + break; + } + mLastHidden = segmentTakeLastHidden(hidden); + if (mLastHidden.get() != nullptr) { + mLastHidden.fix(VARP::CONSTANT); + } + updateSegmentContext(1, 1); + mContext->decode_us += _t.durationInUs(); + ++len; + } + if (len >= max_token) { + mContext->status = LlmStatus::MAX_TOKENS_FINISHED; + } +} + +void SegmentLlm::response(const std::vector& input_ids, std::ostream* os, const char* end_with, + int max_new_tokens) { + MNN::Express::ExecutorScope s(mExecutor); + if (!end_with) { + end_with = "\n"; + } + generate_init(os, end_with); + if (!prefill(input_ids)) { + return; + } + if (max_new_tokens < 0) { + max_new_tokens = mMaxDecodeTokens; + } + if (max_new_tokens > 0) { + generate(max_new_tokens); + } +} + +Llm* createSegmentLlm(std::shared_ptr config) { + return new SegmentLlm(std::move(config)); +} + +} // namespace Transformer +} // namespace MNN + +#endif // MNN_LLM_SUPPORT_SEGMENT diff --git a/transformers/llm/engine/src/segment.hpp b/transformers/llm/engine/src/segment.hpp new file mode 100644 index 0000000000..8c4663399c --- /dev/null +++ b/transformers/llm/engine/src/segment.hpp @@ -0,0 +1,20 @@ +#ifdef MNN_LLM_SUPPORT_SEGMENT + +#ifndef LLM_SEGMENT_HPP +#define LLM_SEGMENT_HPP + +#include + +#include "llm/llm.hpp" + +namespace MNN { +namespace Transformer { + +Llm* createSegmentLlm(std::shared_ptr config); + +} // namespace Transformer +} // namespace MNN + +#endif // LLM_SEGMENT_HPP + +#endif // MNN_LLM_SUPPORT_SEGMENT diff --git a/transformers/llm/engine/tools/llm_bench.cpp b/transformers/llm/engine/tools/llm_bench.cpp index 101a8a73d2..a6369903c8 100644 --- a/transformers/llm/engine/tools/llm_bench.cpp +++ b/transformers/llm/engine/tools/llm_bench.cpp @@ -2,7 +2,6 @@ #include "core/MNNFileUtils.h" #include #include -#include "Profiler.hpp" #include #include #include @@ -208,6 +207,7 @@ struct TestInstance { std::transform(cost_us.begin(), cost_us.end(), std::back_inserter(ts), [n_tokens](int64_t t) { return 1e6 * n_tokens / t; }); return ts; } + std::vector getTokensPerSecond(std::vector n_tokens, std::vector cost_us) const { std::vector ts(n_tokens.size()); for (int i = 0; i < n_tokens.size(); ++i) { @@ -380,8 +380,8 @@ struct markdownPrinter : public Printer { value = buf; } else if (field == "backend") { if (t.backend == 1) value = "METAL"; - else if (t.backend == 2) value = "CUDA"; else if (t.backend == 3) value = "OPENCL"; + else if (t.backend == 7) value = "VULKAN"; else value = "CPU"; } else if (field == "test") { if (t.nPrompt > 0 && t.nGenerate == 0) { @@ -487,6 +487,7 @@ struct jsonAggregator : public Printer { writer.Key("backend"); if (t.backend == 1) writer.String("METAL"); else if (t.backend == 3) writer.String("OPENCL"); + else if (t.backend == 7) writer.String("VULKAN"); else writer.String("CPU"); writer.Key("threads"); @@ -539,7 +540,7 @@ struct jsonAggregator : public Printer { std::vector speed; if (!inst.decodeUs.empty()) { - speed = inst.getTokensPerSecond(inst.nGenerates, inst.decodeUs); + speed = inst.getTokensPerSecond(inst.nGenerate, inst.decodeUs); } else if (!inst.samplesUs.empty()) { speed = inst.getTokensPerSecond(inst.nGenerate, inst.samplesUs); } @@ -568,7 +569,7 @@ struct jsonAggregator : public Printer { writer.Double(inst.getStdevUs(prefill_speed)); } if (!inst.decodeUs.empty()) { - auto decode_speed = inst.getTokensPerSecond(inst.nGenerates, inst.decodeUs); + auto decode_speed = inst.getTokensPerSecond(inst.nGenerate, inst.decodeUs); writer.Key("decode_tps"); writer.Double(inst.getAvgUs(decode_speed)); writer.Key("decode_std"); @@ -772,28 +773,58 @@ static std::vector get_cmd_params_instances(const Run return instances; } -std::string getDirectoryOf(const std::string& file_path, std::string& modelname) { - // weight filename - std::string weight_name = "llm.mnn.weight"; +static uint64_t getFileSizeIfExists(const std::string& path) { + if (!MNNFileExist(path.c_str())) { + return 0; + } + file_t file = MNNOpenFile(path.c_str(), MNN_FILE_READ); + if (file == INVALID_FILE) { + return 0; + } + auto size = MNNGetFileSize(file); + MNNCloseFile(file); + return size == INVALID_SIZE ? 0 : size; +} + +static bool isSegmentConfig(const std::string& file_path) { std::ifstream file(file_path.c_str()); + if (!file.is_open()) { + return false; + } std::string json_str((std::istreambuf_iterator(file)), std::istreambuf_iterator()); - rapidjson::Document doc; doc.Parse(json_str.c_str()); + return doc.HasMember("mnn_llm_version") && doc["mnn_llm_version"].IsString() && + std::string(doc["mnn_llm_version"].GetString()) == "segment"; +} - if (doc.HasMember("llm_weight") && doc["llm_weight"].IsString()) { - weight_name = doc["llm_weight"].GetString(); - } +uint64_t getModelSize(const std::string& file_path, std::string& modelname) { + std::ifstream file(file_path.c_str()); + std::string json_str((std::istreambuf_iterator(file)), std::istreambuf_iterator()); + + rapidjson::Document doc; + doc.Parse(json_str.c_str()); size_t pos = file_path.find_last_of("/\\"); if (pos == std::string::npos) { MNN_ERROR("Invalid model config path\n"); - return ""; + return 0; } auto dir = file_path.substr(0, pos); pos = dir.find_last_of("/\\"); modelname = dir.substr(pos + 1, -1); - return MNNFilePathConcat(dir, weight_name); + + if (isSegmentConfig(file_path)) { + return getFileSizeIfExists(MNNFilePathConcat(dir, "decoder.mnn.weight")) + + getFileSizeIfExists(MNNFilePathConcat(dir, "logit.mnn.weight")) + + getFileSizeIfExists(MNNFilePathConcat(dir, "embed.mnn.weight")); + } + + std::string weight_name = "llm.mnn.weight"; + if (doc.HasMember("llm_weight") && doc["llm_weight"].IsString()) { + weight_name = doc["llm_weight"].GetString(); + } + return getFileSizeIfExists(MNNFilePathConcat(dir, weight_name)); } static void printUsage(int /* argc */, char ** argv) { @@ -802,7 +833,7 @@ static void printUsage(int /* argc */, char ** argv) { printf("options:\n"); printf(" -h, --help\n"); printf(" -m, --model (default: ./Qwen2.5-1.5B-Instruct/config.json)\n"); - printf(" -a, --backends (default: %s)\n", "cpu"); + printf(" -a, --backends (default: %s)\n", "cpu"); printf(" -c, --precision (default: %s) | Note: (0:Normal(for cpu bakend, 'Normal' is 'High'),1:High,2:Low)\n", join(runtimeParamsDefaults.precision, ",").c_str()); printf(" -t, --threads (default: %s)\n", join(runtimeParamsDefaults.threads, ",").c_str()); printf(" -p, --n-prompt (default: %s)\n", join(testParamsDefaults.nPrompt, ",").c_str()); @@ -818,11 +849,10 @@ static void printUsage(int /* argc */, char ** argv) { printf(" -mr, --mixedSme2NeonRatio (default: 41) | Note: This parameter is intended to optimize multi-threaded inference performance on backends that support Arm SME instructions. The optimal ratio may vary across different models; we recommend trying values such as 41, 49, 33.\n"); printf(" -qatten, --quant-attention <0|1> (default: 0) | Note: if 1, quantize attention's key value to int8; default 0\n"); printf(" -j, --json (default: llm_bench.json) | Note: if set, output result to a JSON file\n"); - printf(" --profile Enable operator-level profiling to print detailed timing statistics\n"); } -static bool parseCmdParams(int argc, char ** argv, RuntimeParameters & runtimeParams, TestParameters & testParams, FILE** outfile, bool& helpInfo, bool& jsonMode, std::string& jsonFile, bool& enableProfile) { +static bool parseCmdParams(int argc, char ** argv, RuntimeParameters & runtimeParams, TestParameters & testParams, FILE** outfile, bool& helpInfo, bool& jsonMode, std::string& jsonFile) { std::string arg; bool invalidParam = false; const std::string argPrefix = "--"; @@ -884,10 +914,10 @@ static bool parseCmdParams(int argc, char ** argv, RuntimeParameters & runtimePa for (auto& type: ba) { if (type == "metal") { p.emplace_back(1); - } else if (type == "cuda") { - p.emplace_back(2); } else if (type == "opencl") { p.emplace_back(3); + } else if (type == "vulkan") { + p.emplace_back(7); } else { p.emplace_back(0); } @@ -994,8 +1024,6 @@ static bool parseCmdParams(int argc, char ** argv, RuntimeParameters & runtimePa if (i + 1 < argc && argv[i+1][0] != '-') { jsonFile = argv[++i]; } - } else if (arg == "--profile") { - enableProfile = true; } else { invalidParam = true; @@ -1058,7 +1086,7 @@ static bool parseCmdParams(int argc, char ** argv, RuntimeParameters & runtimePa } -static Llm* buildLLM(const std::string& config_path, int backend, int memory, int precision, int threads, int power, int dynamic_option, bool use_mmap, int divisionRatioSme2Neon, int smeCoreNum, int promptLen, int attention_mode) { +static Llm* buildLLM(const std::string& config_path, int backend, int memory, int precision, int threads, int power, int dynamic_option, bool use_mmap, int divisionRatioSme2Neon, int smeCoreNum, int promptLen, int attention_mode, bool isSegment) { auto llmPtr = Llm::createLLM(config_path); llmPtr->set_config(R"({ "async":false @@ -1067,16 +1095,24 @@ static Llm* buildLLM(const std::string& config_path, int backend, int memory, in // Otherwise, mContext->history_tokens retains data after the first run, skewing true prefill performance metrics." llmPtr->set_config(R"({"reuse_kv":false})"); std::map lever = {{0,"normal"}, {1, "high"}, {2, "low"}}; - std::map backend_type = {{0, "cpu"}, {1, "metal"}, {2, "cuda"}, {3, "opencl"}}; + std::map backend_type = {{0, "cpu"}, {1, "metal"}, {3, "opencl"}, {7, "vulkan"}}; std::map mmap = {{true,"true"}, {false, "false"}}; bool setSuccess = true; - setSuccess &= llmPtr->set_config("{\"precision\":\"" + lever[precision] + "\"}"); + if (isSegment) { + setSuccess &= llmPtr->set_config("{\"precision\":" + std::to_string(precision) + "}"); + } else { + setSuccess &= llmPtr->set_config("{\"precision\":\"" + lever[precision] + "\"}"); + } if (!setSuccess) { MNN_ERROR("precison for LLM config set error\n"); return nullptr; } - setSuccess &= llmPtr->set_config("{\"memory\":\"" + lever[memory] + "\"}"); + if (isSegment) { + setSuccess &= llmPtr->set_config("{\"memory\":" + std::to_string(memory) + "}"); + } else { + setSuccess &= llmPtr->set_config("{\"memory\":\"" + lever[memory] + "\"}"); + } if (!setSuccess) { MNN_ERROR("memory for LLM config set error\n"); return nullptr; @@ -1091,6 +1127,13 @@ static Llm* buildLLM(const std::string& config_path, int backend, int memory, in MNN_ERROR("backend_type for LLM config set error\n"); return nullptr; } + if (isSegment) { + setSuccess &= llmPtr->set_config("{\"forwardtype\":" + std::to_string(backend) + "}"); + if (!setSuccess) { + MNN_ERROR("forwardtype for LLM config set error\n"); + return nullptr; + } + } setSuccess &= llmPtr->set_config("{\"thread_num\":" + std::to_string(threads) + "}"); if (!setSuccess) { MNN_ERROR("thread_num for LLM config set error\n"); @@ -1141,8 +1184,7 @@ int main(int argc, char ** argv) { bool helpInfo = false; bool jsonMode = false; std::string jsonFile = "llm_bench.json"; - bool enableProfile = false; - bool parseSuccess = parseCmdParams(argc, argv, runtimeParams, testParams, &outfile, helpInfo, jsonMode, jsonFile, enableProfile); + bool parseSuccess = parseCmdParams(argc, argv, runtimeParams, testParams, &outfile, helpInfo, jsonMode, jsonFile); if (!parseSuccess) { MNN_ERROR("Parse arguments error\n"); return -1; @@ -1178,48 +1220,39 @@ int main(int argc, char ** argv) { for (const auto & instance: paramsInstances) { TestInstance t(instance); - auto llmWeightPath = getDirectoryOf(t.modelConfigFile, t.modelType); // To check path - - file_t file = MNNOpenFile(llmWeightPath.c_str(), MNN_FILE_READ); - t.modelSize = MNNGetFileSize(file); + t.modelSize = getModelSize(t.modelConfigFile, t.modelType); + const bool isSegment = isSegmentConfig(t.modelConfigFile); MNN::BackendConfig backendConfig; - // Map backend parameter to MNN forward type (0=CPU, 1=METAL, 2=CUDA, 3=OPENCL) - MNNForwardType forwardType = static_cast(instance.mCmdParam.backend); + MNNForwardType forwardType = isSegment ? MNN_FORWARD_CPU : static_cast(instance.mCmdParam.backend); auto executor = MNN::Express::Executor::newExecutor(forwardType, backendConfig, 1); MNN::Express::ExecutorScope scope(executor); - auto llmPtr = buildLLM(instance.mCmdParam.model, instance.mCmdParam.backend, instance.mCmdParam.memory, instance.mCmdParam.precision, instance.mCmdParam.threads, instance.mCmdParam.power, instance.mCmdParam.dynamicOption, instance.mCmdParam.useMmap, instance.mCmdParam.divisionRatioSme2Neon, instance.mCmdParam.smeCoreNum, instance.mCmdParam.nPrompt, instance.mCmdParam.attentionOption); - std::unique_ptr llm(llmPtr); - if (enableProfile) { - llm->set_config(R"({"enable_debug":true})"); - auto profiler = MNN::Profiler::getInstance(); - llm->setDebugCallback( - [profiler](const std::vector& inputs, const MNN::OperatorInfo* info) { - profiler->start(info); - return true; - }, - [profiler](const std::vector& outputs, const MNN::OperatorInfo* info) { - for (auto o : outputs) { - o->wait(MNN::Tensor::MAP_TENSOR_READ, true); - } - profiler->end(info); - return true; - } - ); - } - if (instance.mCmdParam.loadingTime == "true") { + auto createLLM = [&]() { + return std::unique_ptr(buildLLM(instance.mCmdParam.model, instance.mCmdParam.backend, + instance.mCmdParam.memory, instance.mCmdParam.precision, + instance.mCmdParam.threads, instance.mCmdParam.power, + instance.mCmdParam.dynamicOption, instance.mCmdParam.useMmap, + instance.mCmdParam.divisionRatioSme2Neon, + instance.mCmdParam.smeCoreNum, instance.mCmdParam.nPrompt, + instance.mCmdParam.attentionOption, isSegment)); + }; + auto measureLoadingTime = [&]() { + if (instance.mCmdParam.loadingTime != "true" || !t.loadingS.empty()) { + return; + } for (int k = 0; k < 3; ++k) { + auto loadLLM = createLLM(); Timer loadingCost; - llm->load(); + loadLLM->load(); t.loadingS.push_back((double)loadingCost.durationInUs() / 1e6); } - } else { - llm->load(); - } + }; + auto llm = createLLM(); + llm->load(); tuning_prepare(llm.get()); auto context = llm->getContext(); - // Ensure GPU sync for accurate timing + // Ensure GPU sync for accurate timing. llm->set_config("{\"async\":false}"); if (instance.mCmdParam.nGenerate > 0) { llm->set_config("{\"max_new_tokens\":1}"); @@ -1250,6 +1283,7 @@ int main(int argc, char ** argv) { printer_->printHeader(runtimeParams, testParams); printHeader = false; } + measureLoadingTime(); printer_->printPerformance(t); // Cool std::this_thread::sleep_for(std::chrono::milliseconds(5)); @@ -1280,6 +1314,7 @@ int main(int argc, char ** argv) { printer_->printHeader(runtimeParams, testParams); printHeader = false; } + measureLoadingTime(); printer_->printPerformance(t); // Cool std::this_thread::sleep_for(std::chrono::milliseconds(5)); @@ -1287,13 +1322,6 @@ int main(int argc, char ** argv) { } } - if (enableProfile) { - auto profiler = MNN::Profiler::getInstance(); - fprintf(stdout, "\n========== Operator Profile Results ==========\n"); - // profiler->printTimeByName(1); - profiler->printTimeByType(1); - } - fprintf(stdout, "\n"); if (outfile != stdout) { fclose(outfile); diff --git a/transformers/llm/export/llmexport.py b/transformers/llm/export/llmexport.py index 8835b9c6e9..ccf1f09296 100644 --- a/transformers/llm/export/llmexport.py +++ b/transformers/llm/export/llmexport.py @@ -21,6 +21,7 @@ from utils.smooth_quantizer import SmoothQuantizer from utils.omni_quantizer import OmniQuantizer from utils.torch_utils import onnx_export +import segment as segment_export class LlmExporter(torch.nn.Module): ''' @@ -29,7 +30,10 @@ class LlmExporter(torch.nn.Module): def __init__(self, args): super().__init__() self.init_from_args(args) - self.load_model(args.path) + if segment_export.enabled(args) and getattr(args, 'test', None) is None: + segment_export.load_metadata(self, args.path) + else: + self.load_model(args.path) def init_from_args(self, args): self.args = args @@ -47,7 +51,7 @@ def init_from_args(self, args): # init export dst dir if not os.path.exists(self.args.dst_path): os.makedirs(self.args.dst_path) - if not os.path.exists(self.onnx_path): + if not segment_export.enabled(self.args) and not os.path.exists(self.onnx_path): os.makedirs(self.onnx_path) @spinner_run(f'load pretrained model ', True) @@ -679,6 +683,9 @@ def export_language(self): self.onnx_load_param(onnx_model) def export(self, export_type): + if segment_export.enabled(self.args): + segment_export.export(self, export_type) + return if not self.args.skip_weight: if self.args.omni: self.omni_quant() @@ -877,6 +884,8 @@ def build_args(parser): parser.add_argument('--quant_config', type=str, default=None, help='path to the JSON file for op-wise quantization configuration.') parser.add_argument('--generate_for_npu', action='store_true', help='Whether or not to generate model for NPU deployment, default is False.') parser.add_argument('--skip_weight', action='store_true', help='Whether or not to skip loading model weights, useful for testing export flow.') + parser.add_argument('--segment', action='store_true', help='Export segment MNN LLM from safetensors workflow directly, without ONNX export.') + parser.add_argument('--workflow', type=str, default=None, help='workflow json for --segment safetensors conversion. If absent, search resource/*.json.') # omni quant parser.add_argument('--omni_epochs', type=int, default=20, help='OmniQuant 优化的轮数') parser.add_argument('--omni_lr', type=float, default=5e-3, help='OmniQuant 的学习率') @@ -916,4 +925,4 @@ def main(): llm_exporter.export(args.export) if __name__ == '__main__': - main() \ No newline at end of file + main() diff --git a/transformers/llm/export/segment.py b/transformers/llm/export/segment.py new file mode 100644 index 0000000000..582a0b0078 --- /dev/null +++ b/transformers/llm/export/segment.py @@ -0,0 +1,288 @@ +import glob +import json +import os + +from utils.config import LlmConfig +from utils.mnn_converter import MNNConverter +from utils.spinner import spinner_run +from utils.tokenizer import LlmTokenizer + + +def enabled(args): + return getattr(args, 'segment', False) + + +@spinner_run(f'load segment export metadata ', True) +def load_metadata(exporter, model_path): + model_path = os.path.abspath(os.path.expanduser(model_path)) + if exporter.args.tokenizer_path == exporter.args.path: + exporter.args.tokenizer_path = model_path + else: + tokenizer_path = os.path.expanduser(exporter.args.tokenizer_path) + if os.path.exists(tokenizer_path): + tokenizer_path = os.path.abspath(tokenizer_path) + exporter.args.tokenizer_path = tokenizer_path + + exporter.config = LlmConfig.from_pretrained(model_path) + exporter.model_type = exporter.config.model_type + exporter.tokenizer = LlmTokenizer.from_pretrained( + exporter.args.tokenizer_path, + model_type=exporter.model_type + ) + exporter.model = None + exporter.visual = None + exporter.audio = None + exporter.talker = None + exporter.mtp = None + exporter.scale_emb = None + exporter.llm_config = { + 'model_type': exporter.config.model_type, + 'hidden_size': exporter.config.hidden_size, + 'layer_nums': exporter.config.num_hidden_layers, + 'attention_mask': 'float', + 'attention_type': exporter.config.attention_type, + 'is_mrope': False + } + if exporter.config.sliding_window > 0: + exporter.llm_config['sliding_window'] = exporter.config.sliding_window + if hasattr(exporter.tokenizer, 'get_chat_template'): + chat_template = exporter.tokenizer.get_chat_template() + if chat_template is not None: + exporter.llm_config['jinja'] = { + 'chat_template': chat_template + } + if exporter.tokenizer.bos_token: + exporter.llm_config['jinja']['bos'] = exporter.tokenizer.bos_token + if exporter.tokenizer.eos_token: + exporter.llm_config['jinja']['eos'] = exporter.tokenizer.eos_token + if exporter.model_type == 'glm_ocr': + exporter.llm_config['jinja'] = { + 'chat_template': "[gMASK]{% for message in messages %}{% if message.role == \"user\" %}<|user|>\n{{ message.content }}{% elif message.role == \"assistant\" %}<|assistant|>\n{{ message.content }}{% elif message.role == \"system\" %}<|system|>\n{{ message.content }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>\n{% endif %}", + 'eos': '<|endoftext|>' + } + source_llm_config = os.path.join(model_path, 'llm_config.json') + if os.path.exists(source_llm_config): + with open(source_llm_config, 'r', encoding='utf-8') as f: + exporter.llm_config.update(json.load(f)) + return model_path + + +def _resource_dirs(): + export_dir = os.path.dirname(os.path.abspath(__file__)) + repo_root = os.path.abspath(os.path.join(export_dir, '../../..')) + candidates = [ + os.path.join(repo_root, 'resource'), + os.path.join(repo_root, 'transformers', 'llm', 'resource') + ] + return [path for path in candidates if os.path.isdir(path)] + + +def _workflow_score(exporter, workflow_path): + try: + with open(workflow_path, 'r', encoding='utf-8') as f: + workflow = json.load(f) + except Exception: + return None + models = workflow.get('models', []) + if not isinstance(models, list) or len(models) == 0: + return None + + model_names = [model.get('name', '') for model in models if isinstance(model, dict)] + lowered_names = [name.lower() for name in model_names] + score = 0 + if any(name in ('hf_decoder', 'decoder', 'gpt2_decoder') for name in lowered_names): + score += 20 + if any(name in ('logit', 'logit_mobile') for name in lowered_names): + score += 20 + if any(name in ('encoder', 'encoder_mobile', 'audio_proj', 'ntp1', 'wpe') for name in lowered_names): + score -= 30 + + filename = os.path.basename(workflow_path).lower() + model_type = str(getattr(exporter, 'model_type', '') or '').lower() + if model_type and model_type in filename: + score += 15 + if 'qwen' in model_type and 'qwen' in filename: + score += 10 + if 'hf' in filename and any(name == 'hf_decoder' for name in lowered_names): + score += 10 + + blocks = [] + for model in models: + if not isinstance(model, dict): + continue + for block in model.get('blocks', []): + if isinstance(block, dict): + blocks.append(block) + + cfg_pairs = { + 'hiddenSize': getattr(exporter.config, 'hidden_size', None), + 'number': getattr(exporter.config, 'num_hidden_layers', None), + 'headDim': getattr(exporter.config, 'head_dim', None), + 'numHead': getattr(exporter.config, 'num_attention_heads', None), + 'kvNumHead': getattr(exporter.config, 'num_key_value_heads', None), + 'max_position_embeddings': getattr(getattr(exporter.config, 'origin_config', None), 'max_position_embeddings', None) + } + weights = { + 'hiddenSize': 40, + 'number': 35, + 'headDim': 20, + 'numHead': 20, + 'kvNumHead': 20, + 'max_position_embeddings': 5 + } + for key, cfg_value in cfg_pairs.items(): + if cfg_value is None or isinstance(cfg_value, list): + continue + for block in blocks: + workflow_value = block.get(key) + if workflow_value is None and key == 'max_position_embeddings': + workflow_value = block.get('maxPositionEmbeddings') + if workflow_value == cfg_value: + score += weights[key] + break + return score + + +def _resolve_workflow(exporter): + workflow = getattr(exporter.args, 'workflow', None) + if workflow: + workflow = os.path.abspath(os.path.expanduser(workflow)) + if not os.path.exists(workflow): + raise FileNotFoundError(f'workflow json not found: {workflow}') + return workflow + + candidates = [] + for resource_dir in _resource_dirs(): + for path in glob.glob(os.path.join(resource_dir, '**', '*.json'), recursive=True): + score = _workflow_score(exporter, path) + if score is not None and score > 0: + candidates.append((score, os.path.abspath(path))) + candidates.sort(key=lambda item: (-item[0], item[1])) + if not candidates: + searched = ', '.join(_resource_dirs()) + raise RuntimeError(f'--workflow is not set and no suitable workflow json was found under: {searched}') + + best_score = candidates[0][0] + best = [path for score, path in candidates if score == best_score] + if len(best) > 1: + lines = '\n'.join([f' {path}' for path in best]) + raise RuntimeError(f'--workflow is not set and multiple suitable workflow json files were found:\n{lines}\nPlease pass --workflow explicitly.') + + workflow = candidates[0][1] + print(f'--workflow is not set, use workflow json: {workflow}') + return workflow + + +def _resolve_safetensors(model_path): + model_path = os.path.abspath(os.path.expanduser(model_path)) + if os.path.isfile(model_path): + if model_path.endswith('.safetensors'): + return [model_path] + raise RuntimeError(f'--segment expects --path to be a model directory or a .safetensors file, got: {model_path}') + + model_file = os.path.join(model_path, 'model.safetensors') + if os.path.exists(model_file): + return [model_file] + + index_files = sorted(glob.glob(os.path.join(model_path, '*.safetensors.index.json'))) + if index_files: + with open(index_files[0], 'r', encoding='utf-8') as f: + index = json.load(f) + ordered = [] + for filename in index.get('weight_map', {}).values(): + if filename not in ordered: + ordered.append(filename) + paths = [os.path.join(model_path, filename) for filename in ordered] + else: + paths = sorted(glob.glob(os.path.join(model_path, '*.safetensors'))) + + paths = [path for path in paths if os.path.exists(path)] + if not paths: + raise RuntimeError(f'no safetensors file found under: {model_path}') + if len(paths) > 1: + print(f'found {len(paths)} safetensors files, pass all of them to MNNConvert') + return paths + + +def _quant_args(exporter): + quant_bit = exporter.args.quant_bit + if quant_bit == 32: + return [] + if quant_bit == 16: + return ['--fp16'] + return [ + '--weightQuantBits', + str(quant_bit), + '--weightQuantBlock', + str(exporter.args.quant_block) + ] + + +@spinner_run(f'convert safetensors model to ') +def _convert_safetensors(exporter, workflow_path, safetensors_paths): + convert_args = [ + '', + '-f', + 'ST', + '-i', + str(workflow_path) + ] + for safetensors_path in safetensors_paths: + convert_args += ['-i', str(safetensors_path)] + convert_args += [ + '-o', + str(exporter.args.dst_path), + '--allowCustomOp' + ] + if exporter.args.transformer_fuse: + convert_args += ['--transformerFuse'] + if exporter.args.group_conv_native: + convert_args += ['--groupConvNative'] + if exporter.args.sym: + convert_args += ['--weightQuantAsymmetric=0'] + convert_args += ['--saveExternalData'] + if exporter.args.hqq: + convert_args += ['--hqq'] + convert_args += _quant_args(exporter) + MNNConverter(exporter).convert(convert_args) + return exporter.args.dst_path + + +def _export_config(exporter, tokenizer_file): + with open(f'{exporter.args.dst_path}/export_args.json', 'w', encoding='utf-8') as f: + json.dump(exporter.args.__dict__, f, ensure_ascii=False, indent=4) + config_json = f'{exporter.args.dst_path}/llm_config.json' + with open(config_json, 'w', encoding='utf-8') as f: + json.dump(exporter.llm_config, f, ensure_ascii=False, indent=4) + + stop_ids = getattr(exporter.tokenizer, 'stop_ids', []) + eos_token = getattr(exporter.tokenizer, 'eos_token_id', None) + if eos_token is None and len(stop_ids) > 0: + eos_token = stop_ids[0] + if isinstance(eos_token, list): + eos_token = eos_token[0] if len(eos_token) > 0 else None + config = { + 'forwardtype': 1, + 'precision': 2, + 'memory': 2, + 'speculative': 0, + 'draft_len': 1, + 'max_decode_tokens': exporter.max_new_tokens, + 'mnn_llm_version': 'segment', + 'tokenizer_file': os.path.basename(tokenizer_file) + } + if eos_token is not None: + config['eos_token'] = int(eos_token) + with open(f'{exporter.args.dst_path}/config.json', 'w', encoding='utf-8') as f: + json.dump(config, f, ensure_ascii=False, indent=4) + return config_json + + +def export(exporter, export_type): + if export_type != 'mnn': + raise RuntimeError('--segment only supports --export mnn') + workflow = _resolve_workflow(exporter) + safetensors_paths = _resolve_safetensors(exporter.args.path) + _convert_safetensors(exporter, workflow, safetensors_paths) + tokenizer_file = exporter.export_tokenizer() + _export_config(exporter, tokenizer_file)