diff --git a/CMakeLists.txt b/CMakeLists.txt index f634d39f..e430d5bb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -90,25 +90,20 @@ option(NN_ENABLE_TESTS "Build test/verification executables" OFF) # 收集实际构建的目标，供 GPU/CUDA 统一启用（避免引用未定义目标） set(NN_APP_TARGETS "") set(NN_TEST_TARGETS "") +set(NN_BENCH_TARGETS "") # 性能基准（仅测试构建，手工运行，不注册 ctest） -# ── MNIST 目标 ────────────────────────────────────────────────────────────── +# ── MNIST 目标（标准构建核心应用） ─────────────────────────────────────── nn_add_executable(mnist_train src/mnist_train.cpp) nn_add_executable(mnist_infer src/mnist_infer.cpp) -nn_add_executable(mnist_bench src/mnist_bench.cpp) -list(APPEND NN_APP_TARGETS mnist_train mnist_infer mnist_bench) +list(APPEND NN_APP_TARGETS mnist_train mnist_infer) -# ── 文本目标（GPT 训练/推理 + BPE 分词器训练/推理） ─────────────────── +# ── 文本目标（GPT 训练/推理 + BPE 分词器训练/推理，标准构建核心应用） ── nn_add_executable(text_train src/text_train.cpp) nn_add_executable(text_infer src/text_infer.cpp) nn_add_executable(tokenizer_train src/tokenizer_train.cpp) nn_add_executable(tokenizer_infer src/tokenizer_infer.cpp) list(APPEND NN_APP_TARGETS text_train text_infer tokenizer_train tokenizer_infer) -# ── 计算性能基准工具 ───────────────────────────────────────────────────── -nn_add_executable(compute_bench src/compute_bench.cpp) -nn_add_executable(bench_thresholds src/bench_thresholds.cpp) -list(APPEND NN_APP_TARGETS compute_bench bench_thresholds) - # ── 测试 / 验证目标（默认关闭，NN_ENABLE_TESTS=ON 时构建）───────────── if(NN_ENABLE_TESTS) @@ -176,6 +171,14 @@ if(NN_ENABLE_TESTS) nn_add_executable(expr_opt_test src/expr_opt_test.cpp) list(APPEND NN_TEST_TARGETS expr_opt_test) + # ── 图 IR + 融合分析测试（IR-C：expr_graph.hpp 纯函数，无 GPU） ───── + nn_add_executable(expr_graph_test src/expr_graph_test.cpp) + list(APPEND NN_TEST_TARGETS expr_graph_test) + + # ── 图 IR 融合 GPU 端到端（IR-C：begin_expr/end_expr 融合执行） ───── + nn_add_executable(expr_fuse_test src/expr_fuse_test.cpp) + list(APPEND NN_TEST_TARGETS expr_fuse_test) + # ── ModelSpec 架构校验一致性测试（load_model 文件头 vs 模型规格比对） ──── nn_add_executable(model_spec_validation_test src/model_spec_validation_test.cpp) list(APPEND NN_TEST_TARGETS model_spec_validation_test) @@ -184,14 +187,6 @@ if(NN_ENABLE_TESTS) nn_add_executable(softmax_gradcheck src/softmax_gradcheck.cpp) list(APPEND NN_TEST_TARGETS softmax_gradcheck) - # ── 注意力反向逐级调试探针（对比暴力 Matrix 参考与引擎原语） ──────────── - nn_add_executable(attn_debug src/attn_debug.cpp) - list(APPEND NN_TEST_TARGETS attn_debug) - - # ── batched_matmul 原语探针（四种转置组合 vs 手工参考） ─────────────────── - nn_add_executable(matmul_probe src/matmul_probe.cpp) - list(APPEND NN_TEST_TARGETS matmul_probe) - # ── M4 matmul 融合原语测试（bmm_reduce/bmm_denom/bmm_apply，CPU+GPU） ── # ── M6 两趟式注意力反向原语测试（bmm_q_backward/bmm_kv_backward，CPU+GPU） # ── M5 列式 softmax 融合原语测试（col_softmax_denom/col_softmax_sparse_forward） @@ -202,6 +197,13 @@ if(NN_ENABLE_TESTS) nn_add_executable(ce_fusion_test src/ce_fusion_test.cpp) list(APPEND NN_TEST_TARGETS ce_fusion_test) + # ── 性能基准工具（仅测试构建；手工运行，不注册 ctest） ── + nn_add_executable(mnist_bench src/mnist_bench.cpp) + nn_add_executable(compute_bench src/compute_bench.cpp) + nn_add_executable(bench_thresholds src/bench_thresholds.cpp) + nn_add_executable(reduce_bench src/reduce_bench.cpp) + list(APPEND NN_BENCH_TARGETS mnist_bench compute_bench bench_thresholds reduce_bench) + endif() # ── 测试目标统一处理：输出到 <默认可执行目录>/test，并注册到 ctest ────────── @@ -220,6 +222,11 @@ if(NN_ENABLE_TESTS) RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/test) add_test(NAME ${TARGET} COMMAND ${TARGET}) endforeach() + # 基准工具输出到 test/，但不注册 ctest（手工运行） + foreach(TARGET ${NN_BENCH_TARGETS}) + set_target_properties(${TARGET} PROPERTIES + RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/test) + endforeach() endif() # ── Vulkan GPU 加速（可选） ──────────────────────────────────────────────── @@ -307,6 +314,10 @@ if(Vulkan_FOUND) target_link_libraries(scan_exprs PRIVATE nn_core Vulkan::Vulkan) add_custom_command( OUTPUT ${NN_EXPR_SCAN_BIN} + # scan_exprs 是原生可执行文件，直接写文件、不会自建目录。 + # Ninja 会自动创建自定义命令的输出目录，但 Makefiles 生成器不会， + # 全新构建时 generated/ 尚未存在会写入失败（见 08-pitfalls）。 + COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_BINARY_DIR}/generated COMMAND scan_exprs ${NN_EXPR_SCAN_BIN} DEPENDS scan_exprs ${SCAN_SRC} COMMENT "[GPU] Scanning Layer expressions (scan_exprs) → expr_specs.bin" @@ -319,6 +330,8 @@ if(Vulkan_FOUND) target_link_libraries(gen_fused PRIVATE nn_core) add_custom_command( OUTPUT ${NN_FUSED_REGISTRY_HPP} + # 同 scan_exprs：gen_fused 也是原生可执行文件，需先确保输出目录存在 + COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_BINARY_DIR}/generated COMMAND gen_fused ${CMAKE_BINARY_DIR}/generated ${GLSLC} ${NN_EXPR_SCAN_BIN} DEPENDS gen_fused ${FUSED_GEN_SRC} ${NN_EXPR_SCAN_BIN} COMMENT "[GPU] Synthesizing fused shaders (gen_fused) → fused_registry.hpp" @@ -339,6 +352,7 @@ if(Vulkan_FOUND) set(NN_GPU_TARGETS ${NN_APP_TARGETS}) if(NN_ENABLE_TESTS) list(APPEND NN_GPU_TARGETS ${NN_TEST_TARGETS}) + list(APPEND NN_GPU_TARGETS ${NN_BENCH_TARGETS}) endif() foreach(TARGET ${NN_GPU_TARGETS}) diff --git a/cli_controllers.py b/cli_controllers.py index d196e712..6cb74789 100644 --- a/cli_controllers.py +++ b/cli_controllers.py @@ -383,7 +383,7 @@ class MnistTrainController(CLIController): MNIST训练CLI控制器 封装mnist_train.exe，支持以下功能： - - 架构选择：MLP或Transformer + - 架构选择：MLP、Transformer 或 CNN - 训练参数：轮数、学习率、批大小、优化器等 - 设备选择：CPU、GPU(Vulkan)、CUDA - 模型管理：保存、恢复 @@ -470,6 +470,19 @@ def _format_args(self, **kwargs) -> List[str]: if "eval_samples" in kwargs: args.extend(["--eval-samples", self._format_arg_value(kwargs["eval_samples"])]) + # CNN专用参数 (LeNet-5 风格) + if "cnn_channels" in kwargs: + args.extend(["--cnn-channels", self._format_arg_value(kwargs["cnn_channels"])]) + + if "cnn_kernels" in kwargs: + args.extend(["--cnn-kernels", self._format_arg_value(kwargs["cnn_kernels"])]) + + if "cnn_pool" in kwargs: + args.extend(["--cnn-pool", self._format_arg_value(kwargs["cnn_pool"])]) + + if "cnn_fc" in kwargs: + args.extend(["--cnn-fc", self._format_arg_value(kwargs["cnn_fc"])]) + # 振荡检测 if "osc_guard" in kwargs: args.extend(["--osc-guard", self._format_arg_value(kwargs["osc_guard"])]) diff --git a/docs/11-ir-optimization.md b/docs/11-ir-optimization.md index 4a0efbc9..9f611f26 100644 --- a/docs/11-ir-optimization.md +++ b/docs/11-ir-optimization.md @@ -1,6 +1,6 @@ # IR 优化（IR Optimization）—— 为融合算子引入中间表示 -> 状态：**IR-A（canonicalize：DCE/常量折叠/代数化简/稳定重编号）+ IR-B（CSE + 寄存器分配 liveness）✅ 已实施（2026-08-23）**；IR-C（图 IR）/IR-D（emitter）待立项 +> 状态：**IR-A（canonicalize：DCE/常量折叠/代数化简/稳定重编号）+ IR-B（CSE + 寄存器分配 liveness）✅ 已实施（2026-08-23）**；**IR-C（图 IR + begin_expr/end_expr 融合分析）✅ 已实施（基础版：逐元素链拼接，2026-08-24）**；**IR-D（后端 emitter 抽象）✅ 已实施（2026-08-24）**；IR-C 进阶（多输入 DAG 融合 / 跨 kernel 自动融合）待立项 > 关联文档：`09-operator-fusion.md`（算子融合）、`DEVELOPMENT_STANDARDS.md`（分层铁律）、`08-pitfalls-and-lessons.md` > 目标：在**不推翻现有 AOT 闭合世界**的前提下，把 `ExprSpec` 从"直通轻量 IR"演进为带优化 pass 的规范 IR，分阶段提升开发便利与代码质量。 @@ -168,6 +168,21 @@ IR → GlslEmitter / CudaEmitter / CpuEmitter > - 验证：新增 `src/expr_opt_test`（各 pass/确定性/幂等/语义等价/上限压力/归约区段互斥）；`expr_dsl_test`/`expr_reduce_test`/`tensor_expr_test`/`fused_gpu_test`/`matmul_fusion_test`/`ce_fusion_test`/gradcheck 系列/gpt_checkpoint_test 全绿；MNIST + text_train（CPU/GPU、含 checkpoint-every）端到端训练正常。 > - 顺带修复预存在 bug：`text_train --save-interval 0` 触发 `(step+1) % 0` 整数除零崩溃（HEAD 亦复现，与 IR 无关）。 +> **实施记录（2026-08-24，IR-C 基础版 + IR-D 已完成）**： +> - **IR-C（图 IR + 融合分析，`expr_graph.hpp`）**： +> - 图结构：`ExprGraphNode{ spec(canonical), rows/cols, vector_out, dep_of_input, input_tensors }` + `ExprGraph`；`add_node` 录制（依赖识别基于 Tensor 的 `virtual_tag`，占位 Tensor 标记前序节点输出）。 +> - 融合分析 `fuse_expr_graph`：贪心按节点序，维护"当前 kernel"；**逐元素链拼接**——B 以 Linear 视图消费 A 的输出且 A 无其他消费者、形状相同、均无归约 → 把 A 的指令内联进 B（A 输出寄存器作为 B 的操作数），中间结果不落显存、单 kernel dispatch。拼接后 canonicalize + validate；超限/非法 → 保守放弃融合（各自成 kernel）。归约节点/归约输出作为融合边界。 +> - `Tensor` 增加 `virtual_tag`（0=非节点输出；reshape 保留）。 +> - **引擎接入**：`GpuEngine::begin_expr/end_expr` 录制 → `end_expr` 融合执行（每 kernel 一次 AOT dispatch，输出写入末尾节点占位 buffer——`run_fused_gpu` 新增 `output_override` 参数复用占位）；`CpuEngine` 在 `NN_EXPR_SCAN` 下 begin/end 录制并登记融合后的复合 spec（闭合世界两端一致）；`dsl::compute/compute_reduce` 的 scan 分支在录制段内加入录制图。 +> - 演示 Layer：`FusedChainLayer`（forward 用 begin_expr/end_expr 写三段链式表达式，GPU 融合成单 kernel，CPU 逐节点参考）。 +> - **关键坑**：① 拼接时 B 的常量池必须追加在 A 之后（否则 Const 引用越界，canonicalize 产出垃圾索引）；② 录制段内单个表达式**不单独登记**（被融合进复合 spec），闭合世界要求 scan dry-run 覆盖录制段内所有可能的融合组合（含 backward）；③ 中间量不逃逸是融合语义约束（如 LayerNorm/RMSNorm 的 cache 中间结果不可包进 begin/end 段）。 +> - 验证：新增 `src/expr_graph_test`（链融合/三节点链/独立分支/多消费者/归约边界/上限压力/确定性/依赖识别/输入顺序）+ `src/expr_fuse_test`（GPU 端到端：FusedChain forward 融合 err=0、手写录制融合 err=0、backward/未录制回归）。 +> - **IR-D（emitter 抽象，`expr_emitter.hpp`）**： +> - `ExprEmitter` 纯接口（name/generate/generate_reduce）+ `emitter_registry`（按后端名选择工厂）。 +> - `GlslEmitter`（glsl_gen.hpp 现有 generate_glsl/generate_glsl_reduce 封装）+ `CpuEmitter`（cpu_emitter.hpp，生成可编译 C++ 直线代码，验证"一份 IR 多后端"；数值与 CPU 解释器一致）。 +> - `gen_fused` 经 emitter 注册表选择后端（默认 glsl）；`--list-backends` 展示可用后端（glsl/cpu）。 +> - 验证：ctest 24/24 全绿（含新增 expr_graph_test / expr_fuse_test）；`scan_exprs` 收集到 20 条融合表达式（含 FusedChain 融合复合 spec + backward）；`fused_gpu_test` 等既有回归无回归。 + **推荐**：先落地 IR-A + IR-B（约 1 周），解决真实痛点并铺好确定性地基；IR-C 视需求再投入。 --- diff --git a/docs/12-innovative-designs.md b/docs/12-innovative-designs.md new file mode 100644 index 00000000..350522cc --- /dev/null +++ b/docs/12-innovative-designs.md @@ -0,0 +1,270 @@ +# 🚀 neuralnet.cpp 创新设计总览 + +> 本文档汇总本项目中所有具有**原创性或系统性思考**的设计。它不是架构/性能/融合等专项文档的重复，而是从"为什么这么做、创新点在哪、带来什么收益"的角度，把这些设计串成一张全景图。 +> +> 关联文档：[01-architecture](./01-architecture.md) · [02-performance](./02-performance.md) · [09-operator-fusion](./09-operator-fusion.md) · [10-memory-optimization](./10-memory-optimization.md) · [11-ir-optimization](./11-ir-optimization.md) · [05-algorithm-reference](./05-algorithm-reference.md) + +--- + +## 0. 总览 + +一个好的深度学习框架，设计难点往往不在"算法本身"，而在于**如何把算法用干净、可维护、可并行、可扩展的方式组织起来**。本项目的创新集中在以下几个维度： + +| 维度 | 核心创新 | 一句话 | +|------|---------|--------| +| **架构** | 分层引擎化 | 一次编写 Layer，CPU/GPU 零修改双跑 | +| **架构** | 铁律式职责划分 | 算法留 Layer、原语进引擎、Shader 是内部实现 | +| **编译** | AOT 闭合世界表达系统 | 表达式文本唯一、构建期扫描生成、运行时按 key 精确分发 | +| **性能** | SmartPolicy + 零分配线程池 | 自适应并行、消除每次任务的堆分配 | +| **性能** | 表达式模板零临时矩阵 | `a + b * c` 不物化中间结果 | +| **显存** | 结构融合（而非手写 kernel） | 三粒度融合：逐元素链 / matmul+归约 / 两趟注意力 | +| **显存** | 两趟式注意力 + 反向重算 | 不物化 `O(seq²)` 分数矩阵 | +| **显存** | 稀疏 CE + 梯度检查点 + 池归还 | 把显存峰值从 29GB 压到 27GB（并持续下探） | +| **IR** | 带确定性 pass 的规范 IR | DCE/常量折叠/CSE/寄存器分配，key 定义在 canonical IR 上 | +| **算法** | 纯 BPE + 兜底词表自举 | 小模型也能稳定启动、支持中文 | + +下面逐条展开。 + +--- + +## 1. 核心架构创新 + +### 1.1 引擎化架构（Engine-Based） + +**创新点**：让 `Layer::forward/backward` 只写一次，通过传入的 `ComputeEngine` 自动适配 CPU / GPU / CUDA，而不是像多数教学框架那样为每后端各写一份。 + +```cpp +class Layer { + virtual Result forward(ComputeEngine& engine, const Tensor& input) = 0; + virtual Result backward(ComputeEngine& engine, const Tensor& grad) = 0; +}; +``` + +**收益**： +- 新增一个后端 = 只实现 `ComputeEngine` 抽象，零 Layer 改动。 +- 算法正确性在 CPU 上验证后可放心跑 GPU。 + +### 1.2 分层"铁律"（Red Lines） + +这是本项目最鲜明的**原创治理原则**，贯穿所有文档： + +``` +引擎只提供 op-level 原语（matmul/reduce/broadcast/exp…），绝不出现 softmax/attention/layernorm 命名接口 +算法文本只写在 Layer（ReLU/GeLU/Softmax/Attention 的公式只在 compute_layer.hpp / compute_loss.hpp） +Shader / 融合逻辑是引擎内部实现，用户不可见 +``` + +**创新点**：不是把"引擎是否认识 softmax"当实现细节，而是把它当一条**不可逾越的红线**。引擎只认**结构**（如 `reduce(matmul(A,B))`），绝不认**算法名**。这使得算子融合、IR 优化、新后端都能在一个稳定的契约上展开。 + +配套约束： +- `ComputEngine` 原语可"专"（matmul+归约、matmul+softmax分母都是合法原语）但必须通用可复用、不叫算法名。 +- 全程 `noexcept` / `Result`（`std::expected`），零手动内存管理（`-fno-exceptions`）。 + +### 1.3 统一张量 `Tensor` + 零拷贝 + +`Tensor` 用 `shared_ptr` 持有 CPU/GPU 两份互斥存储，拷贝=引用计数，跨层传递零拷贝。修改需显式 `clone`。这为 GPU 训练"全程驻留显存、仅 evaluate 时下载"提供了容器基础。 + +--- + +## 2. AOT 闭合世界表达系统（Compiler-Style） + +这是本项目**最具工程创新性**的设计，把"表达式"当成一门可编译的语言： + +### 2.1 单一来源（Single Source of Truth） + +算法表达式**直接写在 C++ 源码里**（如 `exp(x - row_max) / sum(...)`），而不是单独维护一份"计算图定义"。 + +### 2.2 构建期扫描 + 生成 + +``` +Layer 内联表达式 → to_expr_spec 折叠 → ExprSpec（扁平 SSA IR） + → expr_spec_key（规范结构 key）→ scan_exprs 收集 + → gen_fused 用 glsl_gen 合成 GLSL → glslc → 内联 SPIR-V → fused_registry.hpp +运行时：fold → key → find_fused(key) → 精确 dispatch（闭合世界） +``` + +**创新点**： +- **闭合世界（closed world）**：所有可能出现的融合 kernel 都在构建期预编译、内联进最终程序。运行时**没有编译器、没有 eager 分发**。 +- **派生而非手写**：`scan_exprs` 通过 dry-run 执行 Layer 的 forward/backward 触达表达式，把折叠后的**结构（派生物）**采集进注册表，而不是手工维护一份 kernel 清单。 +- **未命中硬报错**：GPU 上遇到未预生成的表达式结构，**硬报错**并且绝不静默回退 CPU，保持项目"不降级"哲学——据此暴露扫描覆盖缺口。 + +### 2.3 确定性 key 铁律 + +key 定义在 **canonical（优化后）IR** 上，scan 与 runtime 两端必须先 canonicalize 再算 key。且 pass 遍历顺序固定、CSE 胜利者选择固定、寄存器分配固定——避免跨编译器（Clang/MSVC）漂移。这条铁律来自一次真实事故：C++ 实参求值顺序未指定曾导致 key 跨编译器不稳定。 + +--- + +## 3. 算子融合（结构融合，而非手写 kernel） + +### 3.1 设计哲学 + +**不手写任何命名算法 kernel**（如不手写 flash-attention、不手写 softmax kernel），而是： +1. Layer 用 `eval_expr` / `begin_expr/end_expr` 录制表达算法； +2. 引擎/工具按**结构**合成融合 kernel； +3. 所有中间 Tensor 由融合 kernel 内部消解，不落显存。 + +### 3.2 三粒度融合 + +| 粒度 | 覆盖 | 消除的中间量 | +|------|------|-------------| +| **逐元素链** | ReLU/GeLU/SwiGLU 连续逐元素 | 链上的全尺寸中间 Tensor | +| **matmul + 归约** | softmax 分子分母、norm 统计 | `A·B` 全尺寸物化 | +| **两趟注意力 / 稀疏 CE** | attention、大词表交叉熵 | `O(seq²)` / `O(vocab×seq)` 物化 | + +### 3.3 三个 matmul 融合原语 + +为承载两趟注意力，新增了一组**通用、可复用**的原语（不是"attention"，而是"matmul 后接结构"）： + +- `batched_matmul_reduce`：matmul 后沿输出维度归约，不物化中间 `A·B`。 +- `batched_matmul_softmax_denom`：减行 max → exp → 按列求和（softmax 分母，数值稳定）。 +- `batched_matmul_softmax_apply`：行 softmax 归一化后与 V 相乘累加，逐 tile 流式、不物化权重矩阵。 + +反向同样有 `..._softmax_backward_q` / `..._softmax_backward_kv`，kernel 内部重算权重矩阵。 + +### 3.4 形状无关融合（Key 创新） + +RoPE 的 `RowMod/RotateHalf` 参数如果以**结构常量**折进 key，每个 `d_k` 就要一个新融合 shader → 闭合世界无法穷举。 + +**解决**：把形状相关参数（mod/block）从 key 里**拿出来变成运行时 push constant**。结构相同、形状不同共享同一融合 shader——任意 `d_k`（含非 2 的幂）都全融合，零额外显存。这是"闭合世界"与"形状多变"矛盾的一个优雅解法。 + +--- + +## 4. 两趟式注意力（内存高效 Attention） + +### 4.1 问题 + +传统 `scores/masked/attn_cache_` 各占一份 `H·batch·seq²`，且 `attn_cache_` 永久缓存用于反向。这是 GPT+Vulkan 训练显存高的元凶之一。 + +### 4.2 方案 + +把 `scores → mask → softmax → ×V` 拆成**不物化 `seq²` 矩阵**的算法： + +- **Forward**：`m = max of QᵀK`（bmm_reduce）→ `l = Σ exp(QᵀK - m)`（bmm_denom）→ `O = W·V` 逐 tile（bmm_apply）。只留下 `m/l`（`H·batch·seq`）与 `O`（`H·batch·d_k·seq`）。 +- **Backward**：默认**反向重算 W**（用原语再算一次），放弃 `attn_cache_`。代价是 2× FLOPs，换整份 `BH·seq²` 缓存。 + +### 4.3 收益 + +每层从 ~3×`BH·seq²` 降到 `O(BH·seq·d_k)`。这是典型的**用计算换显存**，在训练显存吃紧、而算力相对有余的场景下是最优取舍。 + +--- + +## 5. 稀疏交叉熵（Sparse CE） + +### 5.1 问题 + +大词表（vocab≈25k）交叉熵若物化整个 `(classes, total)` 的 softmax，显存与下载开销巨大。 + +### 5.2 方案 + +两个 op-level 原语，单 kernel 完成全部计算： +- `col_softmax_denom`：exp 后列内树形归约求分母，不物化 exp 张量。 +- `col_softmax_sparse_forward`：单 kernel 同时产出**稠密梯度**（每列只对合法列施加）与**标签位置的 log_softmax loss**，不从整张 softmax gather。 + +**创新点**：labels 以 `(1,N)` 浮点打包上传（`vocab ≤ 2²⁴` 时 float 可精确表示类别索引），kernel 内转 `uint` 读取；非法/被屏蔽列整列置 0，与 CPU 参考完全一致。显存从 ~3-4×`(classes,total)` 降到 ~2×。 + +--- + +## 6. IR 优化（带确定性 pass 的规范 IR） + +### 6.1 动机 + +`ExprSpec` 已是事实上的轻量 IR（SSA、扁平、可序列化、确定性 key），但没有优化 pass，直接导致真实问题： +- 子表达式 `grad*gamma` 重复 3 次 → 超 `EXPR_MAX_INPUTS=8`，被迫手工拆表达式。 +- 无 DCE/常量折叠，shader 携带冗余计算。 +- 无寄存器分配，`num_regs` 受 `EXPR_MAX_REGS=16` 约束。 + +### 6.2 方案 + +- **IR-A canonicalize**：DCE + 常量折叠 + 代数化简 + 稳定重编号（保守、不改变浮点语义）。 +- **IR-B CSE + 寄存器分配**：哈希指令去重，liveness 线性扫描确定性贪心分配寄存器。 +- **IR-C 图 IR + 融合分析**：`begin_expr/end_expr` 录制虚拟寄存器 DAG；**逐元素链拼接**——B 以 Linear 视图消费 A 输出且 A 无其他消费者、形状相同、均无归约 → 指令内联进 B，单 kernel。归约节点/归约输出作为融合边界。 +- **IR-D emitter 抽象**：`ExprEmitter` 接口 + 注册表，一份 canonical IR 产出 GLSL/CPU 多后端（验证"一份 IR 多后端"）。 + +### 6.3 关键不变量 + +canonicalize 不改变 views/inputs 的顺序与内容，只优化 instrs/consts/num_regs → 运行时输入绑定布局不变。寄存器"先声明后赋值"以兼容 liveness 复用同号寄存器。 + +--- + +## 7. 显存优化体系（三级） + +围绕 GPT 训练峰值显存（29GB → 27GB，持续下探）的三条独立路径： + +### L1 激活重计算（梯度检查点） +- `Layer` 契约扩展 `recompute_supported / forward_recompute`；`GPTModel` 实现**块级检查点**（每 N 个 GPTBlock 存一次输入，backward 重算中间）。 +- 与全存基线**逐位一致（max_abs=0）**。 + +### L2 内存池整块归还 + 统计 +- `MemoryPool` 新增 `PoolStats` / `release_idle_blocks()`（整块 `vkFreeMemory`，带保留阈值防抖动）。 +- **结构融合宣示**：中间 Tensor 的归还路径安全——每个原语录制都包裹全内存屏障（写入→`SHADER_READ` / 输出→`MEMORY_READ`），保证 batch 命令缓冲内前序写入对后续可见。 + +### L3（远期） +- 自动融合 + 无 m/v 优化器（Adafactor 类）降参数状态内存。 + +### 设计红线 +全程**不引入 f16/bf16**（数值统一 fp32）——"省内存但不牺牲精度"，这本身也是一个收敛的取舍决策。 + +--- + +## 8. CPU 性能优化（L1/L2 层） + +### 8.1 SmartPolicy 自适应并行 +根据数据规模自动决定串行/并行，避免小数据量时线程调度开销（约 50~200μs）大于计算收益。实测在 ~512K 元素处首次稳定 >1.5x，此前串行更优。 + +### 8.2 线程池 latch 零分配 +旧版每任务 = N 次堆分配 + N 次加锁 + N 个 future 同步；新版 = 1 次加锁 + 0 次堆分配 + 1 个原子计数器，配合三阶段等待（短自旋 → work-stealing → 条件变量短超时）与"调用者参与计算"。 + +### 8.3 缓存分块 + 栈分配 +`BLOCK_SIZE=64`（64×64×4B=16KB 装入 L1），B 块预取转置，b_block 栈分配零堆分配。 + +### 8.4 算子融合（原语级） +- `axpy_inplace`：`clone+scale+add` 三步并一步。 +- `elementwise_select_scalar_cond`：条件选择融合（ReLU backward）。 +- 多头注意力**批量化**：`rearrange_3d → 单次 batched_matmul → 转回`，把 H 次 matmul 融为 1 次。 +- 因果掩码 / 位置编码缓存：相同 `(batch,seq)` 只构造一次。 + +### 8.5 表达式模板零临时矩阵 +`a + b * c` 在编译期构造 AST，单次遍历执行，不创建 `b*c` 中间矩阵，编译器可整体内联向量化。 + +--- + +## 9. 分词器与算法创新 + +### 9.1 纯 BPE（CharBPE）支持中文 +默认 `bpe` 为纯 BPE（优先于 BBPE），**可支持中文**：每个汉字作为基础字符单元，学习常用字词合并——比跨语言字节切分更能捕捉词义；BBPE 作为备选。词表含 256 个 ASCII 兜底 token（ID 4~259）+ 4 个特殊符号，保证最小 `vocab_size=260` 也能稳定启动不崩溃。 + +### 9.2 GPT 训练约定 +每行作为独立样本，`[BOS]` 开头、`[EOS]` 仅行尾；`generate()` 带 `min_new_tokens` 防止小模型过早输出 EOS；`seq_len` 降到覆盖 >99% 样本的规模以避 `n²` 显存。 + +--- + +## 10. 创新背后的工程方法 + +除了具体技术，项目在**方法论**上也有鲜明的、可复用的设计： + +| 方法 | 体现 | +|------|------| +| **合规红线先行** | 每次演进（融合/IR/显存）都先写"不可逾越红线"表，再谈实现 | +| **结构 > 命名的哲学** | 引擎认结构不认算法名，是融合与 IR 能落地的共同前提 | +| **闭合世界 + 硬报错** | 不静默降级，宁可报错暴露覆盖缺口 | +| **恶意保守的正确性安全网** | GPU 融合未命中/失败回退原语组合；融合路径与回退路径数值一致并互相验证 | +| **数值确定性** | 全程 fp32、定点可复现，优化 pass 不改变浮点语义 | +| **可验证增量里程碑** | M1-M6 每步有独立测试验证（gradcheck / GRUD/CPU 对照），先小层数回归再全面放开 | + +--- + +## 11. 收益一栏 + +| 设计 | 量化收益（示例配置） | +|------|---------------------| +| 两趟注意力 | 每层 `~3×BH·seq²` → `O(BH·seq·d_k)` | +| 稀疏 CE | 全 softmax `(vocab×seq)` 物化 → 仅标签 gather | +| 算子融合 + 显存体系 | GPT 训练峰值 ~29GB → ~27GB（并持续下探） | +| IR CSE + 寄存器分配 | 消除"手工拆表达式"，`num_regs` 受控在 16 内 | +| SmartPolicy | 小矩阵 0.3x 退化 → 串行；大矩阵最高 ~7.8x 加速 | +| 融合 axpy | 每 step 减少 ~600 次 GPU buffer 分配 | +| 注意力批量化 | H 次 matmul → 1 次 batched_matmul | + +--- + +> 🌙 本文档为创新设计的**全景速览**，不替代各专项文档的细节。想要深入的读者请跳转文首的关联文档。 \ No newline at end of file diff --git a/docs/optimize.md b/docs/optimize.md deleted file mode 100644 index 5ff81338..00000000 --- a/docs/optimize.md +++ /dev/null @@ -1,56 +0,0 @@ -# 显存/速度优化实测记录 - -## 梯度检查点（--checkpoint-every N）实测 - -同一 batch、同一配置（6step/batch × 20batch 积累）： - -| N | time | vram | -|---|---|---| -| 0 | 22s | 26GB | -| 1 | 31s | 9GB | -| 2 | 27s | 16.6GB | - -结论： -- **要最快 → N=0**（前提显存够）。检查点引入的额外前向 FLOPs 会随 batch 放大，拉高 batch 无法抵消。 -- **要省显存换更大 batch → N=2** 性价比最优（几乎减半显存、只慢 ~23%）。 -- **N=1 只用于极限压显存**（9GB），但重计算代价最大（+1 整遍前向）。 - ---- - -## Activation Offload（--activation-offload，L1-offload） - -**思路**：检查点（重算）用 FLOPs 换显存；offload 用 **PCIe 带宽**换显存——把每块内部激活搬 host-visible 内存，backward 拷回再反向，**不重算**（FLOPs 保持 1.0×）。 - -### 预期收益（PCIe 3.0 x16 ~7GB/s，激活 ~5GB） -- 显存：~9GB 级（同 N=1） -- 每 step：~23s（vs N=1 的 31s），**快 ~8s，无重算** - -### 实现（已落地，阶段1+2+3） -- 引擎原语：`ComputeEngine::offload_store(activation)` / `offload_load(handle, rows, cols)`（CPU no-op；GPU 用**录制式** `vkCmdCopyBuffer(dev↔host_visible)`，batch 内不提交，避免打断 batch）。 -- 后端：`GpuBuffer::create_host_visible` / `GpuTensor::create_host_visible_empty`（HOST_VISIBLE 内存池）；`copy_buffer_region_gpu`（带偏移拷贝）。 -- 层：`Layer::activation_cache()` 枚举 backward 所需中间激活；GPTBlock `export_activations`（forward 后搬 host + 释放 GPU）/ `import_activations`（backward 前拷回，不重算）。掩码小而常驻不参与 offload。 -- **slab 复用**：每块用一块持久 host-visible slab（大小=激活总字节，跨 step 复用），激活按 offset 写入/读出 → RAM = 激活实际体积，避免 128MB 块碎片膨胀。 -- 接入：`text_train --activation-offload`（与 `--checkpoint-every` 互斥）、GUI/CLI。 - -### RAM 实测（batch 6） -- 早期（每 tensor 独立缓冲）：RAM **26GB**（128MB 块碎片膨胀）。 -- slab 修复后：RAM ≈ 激活实际体积 **~7GB**。 - -### 验证 -- `offload_test`：GPU 激活 → host → GPU 往返 **bit-exact**（非 batch + batch 路径）。 -- `gpt_offload_test`：GPU 上 offload 与全存基线 forward+全部梯度 **bit-exact**。 -- 全部 gradcheck/fusion 回归绿。 - -### 注意 -- 需足够主机 RAM（激活随 batch 线性增长）与 PCIe 带宽；PCIe 3.0/x8 收益减半。 -- CUDA 后端 offload 暂为 no-op（未接 pinned memory）。 - -## 显存实测诊断（doc_ids 场景） - -**问题**：doc_ids 强制注意力走旧路径，物化 (BH×seq, seq)=384MB/块 注意力矩阵，且曾把 `attn_cache_`+`softmax.output_cache_` 存两份 → slab 实测 **22GB**（理论 ~11GB）。 - -**修复**： -1. **去重**：Softmax 单缓冲，旧路径只用 `softmax_.output_cache()` 一份 → slab 22→**16GB**。 -2. **（M7，2026-08）`AttnBias` 组合偏置统一位置编码**：两趟式原语掩码契约升级为组合式描述子（causal + doc_ids 块对角 + ALiBi slopes，见 [09 §4](09-operator-fusion.md)），doc_ids / ALiBi 及其组合**不再回退旧路径**，全部在融合 kernel 内按 (b,h,i,j) 组合偏置。doc_ids 仅需 `O(batch·seq)` 小张量每步上传，不再物化 (BH·seq,seq) 注意力矩阵。因此 `--doc-aligned-windows`（丢弃跨文档窗口以退化为纯因果）已**删除**——所有窗口（含跨文档）均保留，块对角文档掩码在两趟式路径内生效。 - - diff --git a/gui.py b/gui.py index e7029db0..4f206603 100644 --- a/gui.py +++ b/gui.py @@ -48,7 +48,7 @@ BUILD_DIR = PROJECT_DIR / "build" # ---------- 枚举选项 ---------- -ARCH_OPTIONS = ["mlp", "transformer"] +ARCH_OPTIONS = ["mlp", "transformer", "cnn"] OPTIMIZER_OPTIONS = ["sgd", "sgd_momentum", "adam", "adamw", "muon"] ENGINE_OPTIONS = ["CPU", "GPU (Vulkan)", "CUDA"] LR_SCHEDULE_OPTIONS = ["fixed", "cosine"] @@ -788,6 +788,15 @@ def __init__(self, master): self._tf_widgets = [self.tf_sep_label, self.d_model, self.num_heads, self.num_layers, self.d_ff, self.patch_size, self.eval_samples] + # --- CNN 专用参数 (LeNet-5 风格) --- + self.cnn_sep_label = _make_label(p, "── CNN 参数 ──", r); r += 1 + self.cnn_channels = _make_entry_row(p, "cnn_channels", r, "6,16"); r += 1 + self.cnn_kernels = _make_entry_row(p, "cnn_kernels", r, "5,5"); r += 1 + self.cnn_pool = _make_entry_row(p, "cnn_pool", r, "2"); r += 1 + self.cnn_fc = _make_entry_row(p, "cnn_fc", r, "120,10"); r += 1 + self._cnn_widgets = [self.cnn_sep_label, self.cnn_channels, self.cnn_kernels, + self.cnn_pool, self.cnn_fc] + # --- LR 调度相关控件（fixed 时隐藏） --- self._lr_schedule_widgets = [self.min_lr, self.warmup_epochs] @@ -815,14 +824,17 @@ def build_pack_config(self): "device": dev, "data": data, "hyperparameters": args} def _on_arch_change(self, *args): - """根据架构选择显示 MLP 或 Transformer 参数""" - is_mlp = self.arch.get() == "mlp" + """根据架构选择显示 MLP / Transformer / CNN 参数""" + arch = self.arch.get() for w in self._mlp_widgets: w = getattr(w, "widget", w) - w.grid() if is_mlp else w.grid_remove() + w.grid() if arch == "mlp" else w.grid_remove() for w in self._tf_widgets: w = getattr(w, "widget", w) - w.grid() if not is_mlp else w.grid_remove() + w.grid() if arch == "transformer" else w.grid_remove() + for w in self._cnn_widgets: + w = getattr(w, "widget", w) + w.grid() if arch == "cnn" else w.grid_remove() def _on_lr_schedule_change(self, *args): """fixed 调度下隐藏 min_lr / warmup""" @@ -854,16 +866,22 @@ def collect_args(self): if self.min_lr.get(): args["min_lr"] = float(self.min_lr.get()) if self.warmup_epochs.get(): args["warmup_epochs"] = int(self.warmup_epochs.get()) # 架构特定参数 - if self.arch.get() == "mlp": + arch = self.arch.get() + if arch == "mlp": if self.layer_dims.get(): args["layer_dims"] = self.layer_dims.get() args["norm"] = self.norm.get() - else: + elif arch == "transformer": if self.d_model.get(): args["d_model"] = int(self.d_model.get()) if self.num_heads.get(): args["num_heads"] = int(self.num_heads.get()) if self.num_layers.get(): args["num_layers"] = int(self.num_layers.get()) if self.d_ff.get(): args["d_ff"] = int(self.d_ff.get()) if self.patch_size.get(): args["patch_size"] = int(self.patch_size.get()) if self.eval_samples.get(): args["eval_samples"] = int(self.eval_samples.get()) + elif arch == "cnn": + if self.cnn_channels.get(): args["cnn_channels"] = self.cnn_channels.get() + if self.cnn_kernels.get(): args["cnn_kernels"] = self.cnn_kernels.get() + if self.cnn_pool.get(): args["cnn_pool"] = int(self.cnn_pool.get()) + if self.cnn_fc.get(): args["cnn_fc"] = self.cnn_fc.get() return args diff --git a/include/neuralnet.cpp/backend/vk_backend.hpp b/include/neuralnet.cpp/backend/vk_backend.hpp index 86f6cbe9..edb19369 100644 --- a/include/neuralnet.cpp/backend/vk_backend.hpp +++ b/include/neuralnet.cpp/backend/vk_backend.hpp @@ -766,6 +766,8 @@ class GpuBackend std::unordered_map fused_reduce_axis_; // 运行时视图参数个数（RowMod/RotateHalf 的 vp push constant 槽数） std::unordered_map fused_view_param_counts_; + // 每线程处理元素数（1=标量, 4=vec4；决定 dispatch 宽度缩放） + std::unordered_map fused_vec_width_; std::unique_ptr memory_pool_; std::unique_ptr staging_ring_; @@ -1312,6 +1314,7 @@ class GpuBackend fused_pipelines_.emplace(fs.key, std::move(*fp_r)); fused_reduce_axis_.emplace(fs.key, fs.reduce_axis); fused_view_param_counts_.emplace(fs.key, fs.view_param_count); + fused_vec_width_.emplace(fs.key, fs.vec_width); } } #endif @@ -3042,7 +3045,8 @@ class GpuBackend std::span consts, std::size_t rows, std::size_t cols, bool vector_out = false, - std::span view_params = {}) + std::span view_params = {}, + GpuTensor* output_override = nullptr) { if (!initialized_) return std::unexpected(Error{"GPU backend not initialized"}); @@ -3073,13 +3077,30 @@ class GpuBackend const std::uint32_t count = static_cast(rows * cols); // 1. 分配输出 Tensor（vector_out：归约向量原生形状 (rows,1)/(1,cols)） + // 若调用方指定 output_override（IR-C 录制占位 buffer），则复用其 + // buffer（形状必须匹配），结果直接写入占位 → Layer 持有的 Tensor + // 在 end_expr 后即物化，避免额外拷贝。 const std::size_t out_rows = (vector_out && raxis == 0) ? rows : (vector_out && raxis == 1) ? 1 : rows; const std::size_t out_cols = (vector_out && raxis == 1) ? cols : (vector_out && raxis == 0) ? 1 : cols; - auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); - if (!output_res) return std::unexpected(output_res.error()); - GpuTensor output = std::move(*output_res); + std::optional owned_output; + GpuTensor* output_ptr = output_override; + if (output_override) + { + if (output_override->rows() != out_rows || output_override->cols() != out_cols) + return std::unexpected(Error{ + "run_fused_gpu: output_override shape mismatch (" + + std::to_string(out_rows) + "x" + std::to_string(out_cols) + ")"}); + } + else + { + auto output_res = GpuTensor::create_empty(out_rows, out_cols, *this); + if (!output_res) return std::unexpected(output_res.error()); + owned_output.emplace(std::move(*output_res)); + output_ptr = &*owned_output; + } + GpuTensor output = *output_ptr; // 2. 分配描述符集（N 输入 + 1 输出） auto ds_r = alloc_desc_set(pipeline.descriptor_layout()); @@ -3152,11 +3173,14 @@ class GpuBackend vkCmdPushConstants(cmd, pipeline.pipeline_layout(), VK_SHADER_STAGE_COMPUTE_BIT, 0, static_cast(pc.size()), pc.data()); - // dispatch：逐元素 = ceil(count/256)；行归约 = rows 个工作组；列归约 = cols 个工作组 + // dispatch：逐元素 = ceil(count/(256*vec_width))；行归约 = rows 个工作组； + // 列归约(tile) = ceil(cols/256) 个工作组（每工作组 256 列） + const std::uint32_t vec_width = fused_vec_width_.count(shader_name) + ? fused_vec_width_.at(shader_name) : 1u; const std::uint32_t wg_count = (raxis == 0) ? static_cast(rows) - : (raxis == 1) ? static_cast(cols) - : (count + 255) / 256; + : (raxis == 1) ? (static_cast(cols) + 255u) / 256u + : (count + 256u * vec_width - 1u) / (256u * vec_width); vkCmdDispatch(cmd, wg_count, 1, 1); record_output_barrier(cmd, output.buffer().impl()); diff --git a/include/neuralnet.cpp/compute_layer.hpp b/include/neuralnet.cpp/compute_layer.hpp index 61c612c3..cc1d069a 100644 --- a/include/neuralnet.cpp/compute_layer.hpp +++ b/include/neuralnet.cpp/compute_layer.hpp @@ -85,6 +85,11 @@ class Layer // 非可学习状态收集（默认空，BatchNorm 的 running 统计量等需要 override） [[nodiscard]] virtual std::vector extra_state() { return {}; } + // 引擎相关初始化（创建/上传权重张量），替换构造函数中的 NN_ASSERT 模式。 + // 默认实现空操作；各层在构造后由 Model::add() 调用。 + // 返回 Result 以正确传播引擎错误，而非在 Release 下吞掉。 + [[nodiscard]] virtual Result init(ComputeEngine& /*engine*/) { return {}; } + // 梯度检查点（激活重计算）契约 ────────────────────────────────── // checkpoint_mode_ = true 时，forward 不保留中间激活（供 L1 激活重计算）； // forward_recompute 重算 forward 并重建缓存（供 backward 使用）。 @@ -147,6 +152,8 @@ class Layer class Linear final : public Layer { private: + std::size_t in_features_; + std::size_t out_features_; Tensor w_; // 权重 (out_features, in_features) Tensor b_; // 偏置 (out_features, 1) Tensor grad_w_; // 权重梯度 @@ -156,33 +163,36 @@ class Linear final : public Layer inline static thread_local std::mt19937_64 rng_{std::random_device{}()}; public: - Linear(ComputeEngine& engine, - std::size_t in_features, std::size_t out_features) + Linear(std::size_t in_features, std::size_t out_features) + : in_features_(in_features), out_features_(out_features) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override { // ── 在 CPU 上初始化权重（Xavier 均匀分布） ── - Matrix w_cpu(out_features, in_features); - const Scalar limit = std::sqrt(6.0 / static_cast(in_features + out_features)); + Matrix w_cpu(out_features_, in_features_); + const Scalar limit = std::sqrt(6.0 / static_cast(in_features_ + out_features_)); std::uniform_real_distribution dist(-limit, limit); auto w_span = w_cpu.span(); for (std::size_t i = 0; i < w_cpu.size(); ++i) w_span[i] = dist(rng_); - Matrix b_cpu(out_features, 1); // 零初始化 + Matrix b_cpu(out_features_, 1); // 零初始化 // ── 通过 engine 上传到目标设备 ── auto w_res = engine.from_matrix(w_cpu); - NN_ASSERT(w_res, w_res ? "" : w_res.error().message.c_str()); + if (!w_res) return std::unexpected(w_res.error()); w_ = std::move(*w_res); auto b_res = engine.from_matrix(b_cpu); - NN_ASSERT(b_res, b_res ? "" : b_res.error().message.c_str()); + if (!b_res) return std::unexpected(b_res.error()); b_ = std::move(*b_res); // ── 创建梯度张量（零初始化） ── - grad_w_ = engine.create_tensor(out_features, in_features); - grad_b_ = engine.create_tensor(out_features, 1); - { auto r1 = engine.zero(grad_w_); NN_ASSERT(r1, r1 ? "" : r1.error().message.c_str()); } - { auto r2 = engine.zero(grad_b_); NN_ASSERT(r2, r2 ? "" : r2.error().message.c_str()); } + grad_w_ = engine.create_tensor(out_features_, in_features_); + grad_b_ = engine.create_tensor(out_features_, 1); + { auto r1 = engine.zero(grad_w_); if (!r1) return std::unexpected(r1.error()); } + { auto r2 = engine.zero(grad_b_); if (!r2) return std::unexpected(r2.error()); } + return {}; } [[nodiscard]] std::vector parameters() override @@ -591,26 +601,28 @@ class LayerNorm final : public Layer static constexpr Scalar EPSILON = 1e-5; public: - explicit LayerNorm(ComputeEngine& engine, - std::size_t normalized_shape, Scalar epsilon = EPSILON) - : normalized_shape_(normalized_shape), epsilon_(epsilon) + explicit LayerNorm(std::size_t normalized_shape, Scalar epsilon = EPSILON) + : normalized_shape_(normalized_shape), epsilon_(epsilon) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override { // gamma 初始化为 1, beta 初始化为 0 - Matrix gamma_cpu(normalized_shape, 1, Scalar{1}); - Matrix beta_cpu(normalized_shape, 1, Scalar{0}); + Matrix gamma_cpu(normalized_shape_, 1, Scalar{1}); + Matrix beta_cpu(normalized_shape_, 1, Scalar{0}); auto g = engine.from_matrix(gamma_cpu); - NN_ASSERT(g, g ? "" : g.error().message.c_str()); + if (!g) return std::unexpected(g.error()); gamma_ = std::move(*g); auto bv = engine.from_matrix(beta_cpu); - NN_ASSERT(bv, bv ? "" : bv.error().message.c_str()); + if (!bv) return std::unexpected(bv.error()); beta_ = std::move(*bv); - grad_gamma_ = engine.create_tensor(normalized_shape, 1); - grad_beta_ = engine.create_tensor(normalized_shape, 1); - { auto r1 = engine.zero(grad_gamma_); NN_ASSERT(r1, r1 ? "" : r1.error().message.c_str()); } - { auto r2 = engine.zero(grad_beta_); NN_ASSERT(r2, r2 ? "" : r2.error().message.c_str()); } + grad_gamma_ = engine.create_tensor(normalized_shape_, 1); + grad_beta_ = engine.create_tensor(normalized_shape_, 1); + { auto r1 = engine.zero(grad_gamma_); if (!r1) return std::unexpected(r1.error()); } + { auto r2 = engine.zero(grad_beta_); if (!r2) return std::unexpected(r2.error()); } + return {}; } [[nodiscard]] std::vector parameters() override @@ -801,18 +813,20 @@ class RMSNorm final : public Layer static constexpr Scalar EPSILON = 1e-5; public: - explicit RMSNorm(ComputeEngine& engine, - std::size_t normalized_shape, Scalar epsilon = EPSILON) - : normalized_shape_(normalized_shape), epsilon_(epsilon) + explicit RMSNorm(std::size_t normalized_shape, Scalar epsilon = EPSILON) + : normalized_shape_(normalized_shape), epsilon_(epsilon) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override { // gamma 初始化为 1（无 beta） - Matrix gamma_cpu(normalized_shape, 1, Scalar{1}); + Matrix gamma_cpu(normalized_shape_, 1, Scalar{1}); auto g = engine.from_matrix(gamma_cpu); - NN_ASSERT(g, g ? "" : g.error().message.c_str()); + if (!g) return std::unexpected(g.error()); gamma_ = std::move(*g); - grad_gamma_ = engine.create_tensor(normalized_shape, 1); - { auto r1 = engine.zero(grad_gamma_); NN_ASSERT(r1, r1 ? "" : r1.error().message.c_str()); } + grad_gamma_ = engine.create_tensor(normalized_shape_, 1); + { auto r1 = engine.zero(grad_gamma_); if (!r1) return std::unexpected(r1.error()); } + return {}; } [[nodiscard]] std::vector parameters() override @@ -932,13 +946,82 @@ class RMSNorm final : public Layer } }; +// ══════════════════════════════════════════════════════════════════════════ +// FusedChainLayer — IR-C 演示：begin_expr/end_expr 多表达式链式融合 +// +// forward 用 begin_expr/end_expr 把"无法写进一行的多步逐元素变换"写成多个 +// 表达式（t、u 为中间量，不逃逸——融合的语义约束）： +// t = x * 2 (Linear 输入 + 常量) +// u = t + 3 (消费 t：Linear 引用 → 内联为寄存器) +// out = u * gamma (row 广播参数 (F,1)) +// 引擎在 end_expr 时做融合分析（IR-C，expr_graph.hpp）：三段拼接成单个 +// kernel，中间结果 t/u 留在寄存器（不落显存、不额外 dispatch）。 +// - GPU：begin_expr 录制 → end_expr 融合 → AOT 融合 shader 匹配 dispatch +// （闭合世界：scan_exprs 必须 dry-run 本 Layer 登记融合后的复合 spec）。 +// - CPU：begin/end 为 no-op，各表达式直接求值（参考实现）。 +// - scan：NN_EXPR_SCAN 下 end_expr 登记融合后的复合 spec（两端一致）。 +// ══════════════════════════════════════════════════════════════════════════ +class FusedChainLayer final : public Layer +{ +private: + std::size_t features_; + Tensor gamma_; // (F, 1) 按行广播系数 + +public: + FusedChainLayer(std::size_t features) : features_(features) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override + { + Matrix g(features_, 1, Scalar{0.5f}); // 演示用固定系数 + auto g_t = engine.from_matrix(g); + if (!g_t) return std::unexpected(g_t.error()); + gamma_ = std::move(*g_t); + return {}; + } + + [[nodiscard]] Result forward( + ComputeEngine& engine, const Tensor& input) override + { + const std::size_t F = input.rows(); + const std::size_t B = input.cols(); + + auto b = engine.begin_expr(); + if (!b) return std::unexpected(b.error()); + auto t = dsl::compute(engine, dsl::leaf(input) * Scalar{2}, F, B); + if (!t) return std::unexpected(t.error()); + auto u = dsl::compute(engine, dsl::leaf(*t) + Scalar{3}, F, B); + if (!u) return std::unexpected(u.error()); + auto out = dsl::compute(engine, + dsl::leaf(*u) * dsl::row_broadcast(gamma_), F, B); + if (!out) return std::unexpected(out.error()); + auto e = engine.end_expr(); + if (!e) return std::unexpected(e.error()); + return out; + } + + // out = ((x*2)+3)*gamma → dout/dx = 2*gamma。 + // 非 IR-C 演示路径：单表达式直接求值（scan 单独登记，GPU 独立 AOT）。 + [[nodiscard]] Result backward( + ComputeEngine& engine, const Tensor& grad_output) override + { + const std::size_t F = grad_output.rows(); + const std::size_t B = grad_output.cols(); + return dsl::compute(engine, + dsl::leaf(grad_output) * dsl::row_broadcast(gamma_) * Scalar{2}, + F, B); + } + + [[nodiscard]] std::vector parameters() override { return {}; } + [[nodiscard]] std::vector param_gradients() override { return {}; } +}; + // ── 归一化层工厂：按 NormType 创建 LayerNorm 或 RMSNorm ────────────────── [[nodiscard]] inline std::unique_ptr make_norm_layer( - ComputeEngine& engine, std::size_t d_model, NormType norm_type) + std::size_t d_model, NormType norm_type) { if (norm_type == NormType::RMSNorm) - return std::make_unique(engine, d_model); - return std::make_unique(engine, d_model); + return std::make_unique(d_model); + return std::make_unique(d_model); } // ══════════════════════════════════════════════════════════════════════════ @@ -1085,8 +1168,7 @@ class Conv2D final : public Layer } public: - Conv2D(ComputeEngine& engine, - std::size_t in_channels, std::size_t out_channels, + Conv2D(std::size_t in_channels, std::size_t out_channels, std::size_t kernel, std::size_t stride = 1, std::size_t padding = 0, std::size_t in_h = 0, std::size_t in_w = 0) : in_channels_(in_channels), out_channels_(out_channels), @@ -1096,28 +1178,34 @@ class Conv2D final : public Layer const std::size_t fan_in = in_channels * kernel * kernel; out_h_ = (in_h_ + 2 * padding_ - kernel_) / stride_ + 1; out_w_ = (in_w_ + 2 * padding_ - kernel_) / stride_ + 1; + } + + [[nodiscard]] Result init(ComputeEngine& engine) override + { + const std::size_t fan_in = in_channels_ * kernel_ * kernel_; // 权重 (C_out, C_in*k*k) — He 风格均匀初始化 - Matrix w_cpu(out_channels, fan_in); - const Scalar limit = std::sqrt(6.0 / static_cast(fan_in + out_channels)); + Matrix w_cpu(out_channels_, fan_in); + const Scalar limit = std::sqrt(6.0 / static_cast(fan_in + out_channels_)); std::uniform_real_distribution dist(-limit, limit); auto w_span = w_cpu.span(); for (std::size_t i = 0; i < w_cpu.size(); ++i) w_span[i] = dist(rng_); - Matrix b_cpu(out_channels, 1); // 零初始化 + Matrix b_cpu(out_channels_, 1); // 零初始化 auto w_res = engine.from_matrix(w_cpu); - NN_ASSERT(w_res, w_res ? "" : w_res.error().message.c_str()); + if (!w_res) return std::unexpected(w_res.error()); w_ = std::move(*w_res); auto b_res = engine.from_matrix(b_cpu); - NN_ASSERT(b_res, b_res ? "" : b_res.error().message.c_str()); + if (!b_res) return std::unexpected(b_res.error()); b_ = std::move(*b_res); - grad_w_ = engine.create_tensor(out_channels, fan_in); - grad_b_ = engine.create_tensor(out_channels, 1); - { auto r1 = engine.zero(grad_w_); NN_ASSERT(r1, r1 ? "" : r1.error().message.c_str()); } - { auto r2 = engine.zero(grad_b_); NN_ASSERT(r2, r2 ? "" : r2.error().message.c_str()); } + grad_w_ = engine.create_tensor(out_channels_, fan_in); + grad_b_ = engine.create_tensor(out_channels_, 1); + { auto r1 = engine.zero(grad_w_); if (!r1) return std::unexpected(r1.error()); } + { auto r2 = engine.zero(grad_b_); if (!r2) return std::unexpected(r2.error()); } + return {}; } std::vector parameters() override { return {w_, b_}; } @@ -1426,18 +1514,19 @@ class RotaryEmbedding } // 重建全表 (d_k, seq)：cos/sin 按维度对交错重复 - void rebuild(ComputeEngine& engine, std::size_t seq) + [[nodiscard]] Result rebuild(ComputeEngine& engine, std::size_t seq) { Matrix c(d_k_, seq), s(d_k_, seq); for (std::size_t pos = 0; pos < seq; ++pos) fill_pos_column_(c, s, pos, pos); auto cr = engine.from_matrix(c); - NN_ASSERT(cr, cr ? "" : cr.error().message.c_str()); + if (!cr) return std::unexpected(cr.error()); cos_cache_ = std::move(*cr); auto sr = engine.from_matrix(s); - NN_ASSERT(sr, sr ? "" : sr.error().message.c_str()); + if (!sr) return std::unexpected(sr.error()); sin_cache_ = std::move(*sr); seq_cached_ = seq; + return {}; } // 构造 RoPE 内联表达式（forward 用 Add 结尾，backward 用 Sub 结尾） @@ -1446,11 +1535,13 @@ class RotaryEmbedding public: RotaryEmbedding() = default; - RotaryEmbedding(ComputeEngine& /*engine*/, std::size_t d_k) + explicit RotaryEmbedding(std::size_t d_k) : d_k_(d_k) { } + [[nodiscard]] Result init(ComputeEngine& /*engine*/) { return {}; } + [[nodiscard]] std::size_t d_k() const noexcept { return d_k_; } // 全表应用：q 为 (batch*H*d_k, seq)，cos/sin 为 (d_k, seq) 短表 @@ -1461,7 +1552,10 @@ class RotaryEmbedding if (d_k_ == 0 || d_k_ % 2 != 0) return std::unexpected(Error{"RotaryEmbedding::apply: d_k must be positive and even"}); if (seq != seq_cached_) - rebuild(engine, seq); + { + auto r = rebuild(engine, seq); + if (!r) return std::unexpected(r.error()); + } const std::uint32_t dk = static_cast(d_k_); // 内联 RoPE 表达式（LLaMA half-swap；backward 为旋转正交，逆 = 反角）： // forward: out = q·cos + rotate_half(q)·sin @@ -1604,25 +1698,37 @@ class AttentionBase : public Layer std::size_t /*batch*/, std::size_t /*seq*/) {} public: - AttentionBase(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + AttentionBase(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0, PosEncodingType pos_enc = PosEncodingType::Learned) : d_model_(d_model), num_heads_(num_heads), d_k_(d_model / num_heads), seq_len_(seq_len), scale_(Scalar{1} / std::sqrt(static_cast(d_model / num_heads))), - w_q_(engine, d_model, d_model), - w_k_(engine, d_model, d_model), - w_v_(engine, d_model, d_model), - w_o_(engine, d_model, d_model), + w_q_(d_model, d_model), + w_k_(d_model, d_model), + w_v_(d_model, d_model), + w_o_(d_model, d_model), use_rope_(pos_enc == PosEncodingType::RoPE), - rope_(engine, d_model / num_heads) + rope_(d_model / num_heads) { NN_ASSERT(d_model % num_heads == 0, "AttentionBase: d_model must be divisible by num_heads"); } + [[nodiscard]] Result init(ComputeEngine& engine) override + { + auto r1 = w_q_.init(engine); if (!r1) return std::unexpected(r1.error()); + auto r2 = w_k_.init(engine); if (!r2) return std::unexpected(r2.error()); + auto r3 = w_v_.init(engine); if (!r3) return std::unexpected(r3.error()); + auto r4 = w_o_.init(engine); if (!r4) return std::unexpected(r4.error()); + if (use_rope_) + { + auto r5 = rope_.init(engine); if (!r5) return std::unexpected(r5.error()); + } + return {}; + } + std::vector parameters() override { auto p = w_q_.parameters(); @@ -2091,10 +2197,9 @@ class AttentionBase : public Layer class MultiHeadAttention final : public AttentionBase { public: - MultiHeadAttention(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + MultiHeadAttention(std::size_t d_model, std::size_t num_heads, std::size_t seq_len = 0) - : AttentionBase(engine, d_model, num_heads, seq_len) {} + : AttentionBase(d_model, num_heads, seq_len) {} }; // ══════════════════════════════════════════════════════════════════════════ @@ -2120,7 +2225,7 @@ class PositionalEncoding final : public Layer Tensor encoding_cache_; std::size_t cached_total_ = 0; - void rebuild_encoding(ComputeEngine& engine, std::size_t total_len) + [[nodiscard]] Result rebuild_encoding(ComputeEngine& engine, std::size_t total_len) { const std::size_t base_len = (tile_size_ > 0) ? tile_size_ : total_len; const std::size_t batch = (tile_size_ > 0) ? total_len / tile_size_ : 1; @@ -2154,14 +2259,14 @@ class PositionalEncoding final : public Layer } } auto r = engine.from_matrix(enc); - NN_ASSERT(r, r ? "" : r.error().message.c_str()); + if (!r) return std::unexpected(r.error()); encoding_cache_ = std::move(*r); cached_total_ = total_len; + return {}; } public: - PositionalEncoding(ComputeEngine& /*engine*/, - std::size_t d_model, std::size_t max_len = 5000, + PositionalEncoding(std::size_t d_model, std::size_t max_len = 5000, std::size_t tile_size = 0) : d_model_(d_model), max_len_(max_len), tile_size_(tile_size) { @@ -2179,7 +2284,10 @@ class PositionalEncoding final : public Layer return std::unexpected(Error{"PE forward: seq_len exceeds max_len"}); if (cached_total_ != total_len) - rebuild_encoding(engine, total_len); + { + auto r = rebuild_encoding(engine, total_len); + if (!r) return std::unexpected(r.error()); + } return engine.elementwise_binary(BinaryOp::Add, input, encoding_cache_); } @@ -2204,15 +2312,21 @@ class FeedForward final : public Layer bool use_swiglu_ = false; public: - FeedForward(ComputeEngine& engine, - std::size_t d_model, std::size_t d_ff, + FeedForward(std::size_t d_model, std::size_t d_ff, ActivationType activation = ActivationType::GeLU) - : fc1_(engine, d_model, + : fc1_(d_model, activation == ActivationType::SwiGLU ? 2 * d_ff : d_ff), - fc2_(engine, d_ff, d_model), + fc2_(d_ff, d_model), swiglu_(d_ff), use_swiglu_(activation == ActivationType::SwiGLU) {} + [[nodiscard]] Result init(ComputeEngine& engine) override + { + auto r1 = fc1_.init(engine); if (!r1) return std::unexpected(r1.error()); + auto r2 = fc2_.init(engine); if (!r2) return std::unexpected(r2.error()); + return {}; + } + std::vector parameters() override { auto p = fc1_.parameters(); @@ -2311,13 +2425,21 @@ class TransformerEncoderLayer final : public Layer Tensor residual2_cache_; public: - TransformerEncoderLayer(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + TransformerEncoderLayer(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, std::size_t seq_len = 0) - : self_attn_(engine, d_model, num_heads, seq_len), - norm1_(engine, d_model), - ff_(engine, d_model, d_ff), - norm2_(engine, d_model) {} + : self_attn_(d_model, num_heads, seq_len), + norm1_(d_model), + ff_(d_model, d_ff), + norm2_(d_model) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override + { + auto r1 = self_attn_.init(engine); if (!r1) return std::unexpected(r1.error()); + auto r2 = norm1_.init(engine); if (!r2) return std::unexpected(r2.error()); + auto r3 = ff_.init(engine); if (!r3) return std::unexpected(r3.error()); + auto r4 = norm2_.init(engine); if (!r4) return std::unexpected(r4.error()); + return {}; + } std::vector parameters() override { @@ -2468,24 +2590,32 @@ class TransformerEncoder final : public Layer Tensor ones_row_; // (1, num_patches) 全1，用于 backward 池化梯度广播 public: - TransformerEncoder(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + TransformerEncoder(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, std::size_t num_layers, std::size_t num_patches) : d_model_(d_model), num_patches_(num_patches), inv_num_patches_(Scalar{1} / static_cast(num_patches)), // PE 启用 tiling：每个样本独立使用 (d_model, num_patches) 的编码 - pos_encoding_(engine, d_model, num_patches, num_patches) + pos_encoding_(d_model, num_patches, num_patches) { // 所有 EncoderLayer 传入 seq_len=num_patches，启用 MHA 批量化路径 for (std::size_t i = 0; i < num_layers; ++i) - layers_.emplace_back(engine, d_model, num_heads, d_ff, num_patches); + layers_.emplace_back(d_model, num_heads, d_ff, num_patches); + } + [[nodiscard]] Result init(ComputeEngine& engine) override + { + for (auto& layer : layers_) + { + auto r = layer.init(engine); + if (!r) return std::unexpected(r.error()); + } // 预创建 ones_row_ (1, num_patches) 全1，用于 backward 广播 Matrix ones_cpu(1, num_patches_, Scalar{1}); auto or_t = engine.from_matrix(ones_cpu); - NN_ASSERT(or_t, or_t ? "" : or_t.error().message.c_str()); + if (!or_t) return std::unexpected(or_t.error()); ones_row_ = std::move(*or_t); + return {}; } std::vector parameters() override @@ -2612,20 +2742,24 @@ class PatchEmbedding final : public Layer Tensor input_cache_; public: - PatchEmbedding(ComputeEngine& engine, - std::size_t img_size, std::size_t patch_size, + PatchEmbedding(std::size_t img_size, std::size_t patch_size, std::size_t d_model) : img_size_(img_size), patch_size_(patch_size), grid_size_(img_size / patch_size), num_patches_(grid_size_ * grid_size_), patch_dim_(patch_size * patch_size), d_model_(d_model), - projection_(engine, patch_dim_, d_model) + projection_(patch_dim_, d_model) { NN_ASSERT(img_size % patch_size == 0, "PatchEmbedding: img_size must be divisible by patch_size"); } + [[nodiscard]] Result init(ComputeEngine& engine) override + { + return projection_.init(engine); + } + [[nodiscard]] std::size_t num_patches() const noexcept { return num_patches_; } [[nodiscard]] std::size_t d_model() const noexcept { return d_model_; } @@ -2832,20 +2966,21 @@ class CausalSelfAttention final : public AttentionBase } } - void ensure_mask_(ComputeEngine& engine, std::size_t batch, std::size_t seq_len) + [[nodiscard]] Result ensure_mask_(ComputeEngine& engine, std::size_t batch, std::size_t seq_len) { // 用 (batch, seq_len) 作为缓存键 if (mask_cached_batch_ == batch && mask_cached_seq_ == seq_len) - return; + return {}; // 纯因果掩码：委托 build_attention_mask（doc_ids 为空） auto mask = build_attention_mask(batch, seq_len, num_heads_, use_alibi_, slopes_); auto r = engine.from_matrix(mask); - NN_ASSERT(r, r ? "" : r.error().message.c_str()); + if (!r) return std::unexpected(r.error()); mask_cache_ = std::move(*r); mask_cached_batch_ = batch; mask_cached_seq_ = seq_len; + return {}; } public: @@ -2879,7 +3014,10 @@ class CausalSelfAttention final : public AttentionBase if (!mt) return std::unexpected(mt.error()); return engine.elementwise_binary(BinaryOp::Add, scores, *mt); } - ensure_mask_(engine, batch, seq); + { + auto r = ensure_mask_(engine, batch, seq); + if (!r) return std::unexpected(r.error()); + } return engine.elementwise_binary(BinaryOp::Add, scores, mask_cache_); } @@ -2949,12 +3087,11 @@ class CausalSelfAttention final : public AttentionBase } public: - CausalSelfAttention(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + CausalSelfAttention(std::size_t d_model, std::size_t num_heads, std::size_t /*max_len*/ = 1024, std::size_t seq_len = 0, PosEncodingType pos_enc = PosEncodingType::Learned) - : AttentionBase(engine, d_model, num_heads, seq_len, pos_enc), + : AttentionBase(d_model, num_heads, seq_len, pos_enc), use_alibi_(pos_enc == PosEncodingType::ALiBi) { if (use_alibi_) @@ -2988,17 +3125,25 @@ class GPTBlock final : public Layer std::vector offload_offsets_; // 各激活在 slab 中的 float 偏移 public: - GPTBlock(ComputeEngine& engine, - std::size_t d_model, std::size_t num_heads, + GPTBlock(std::size_t d_model, std::size_t num_heads, std::size_t d_ff, std::size_t max_len = 1024, std::size_t seq_len = 0, PosEncodingType pos_enc = PosEncodingType::Learned, ActivationType activation = ActivationType::GeLU, NormType norm_type = NormType::LayerNorm) - : self_attn_(engine, d_model, num_heads, max_len, seq_len, pos_enc), - norm1_(make_norm_layer(engine, d_model, norm_type)), - ff_(engine, d_model, d_ff, activation), - norm2_(make_norm_layer(engine, d_model, norm_type)) {} + : self_attn_(d_model, num_heads, max_len, seq_len, pos_enc), + norm1_(make_norm_layer(d_model, norm_type)), + ff_(d_model, d_ff, activation), + norm2_(make_norm_layer(d_model, norm_type)) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override + { + auto r1 = self_attn_.init(engine); if (!r1) return std::unexpected(r1.error()); + if (norm1_) { auto r = norm1_->init(engine); if (!r) return std::unexpected(r.error()); } + auto r2 = ff_.init(engine); if (!r2) return std::unexpected(r2.error()); + if (norm2_) { auto r = norm2_->init(engine); if (!r) return std::unexpected(r.error()); } + return {}; + } std::vector parameters() override { @@ -3277,6 +3422,9 @@ class PositionEncoder ComputeEngine& engine, const Tensor& grad_T, std::size_t batch, std::size_t seq) = 0; + // 初始化（引擎相关操作，如创建张量等） + [[nodiscard]] virtual Result init(ComputeEngine& /*engine*/) { return {}; } + [[nodiscard]] virtual std::vector parameters() { return {}; } [[nodiscard]] virtual std::vector param_gradients() { return {}; } }; @@ -3297,20 +3445,21 @@ class AdditivePositionEncoder : public PositionEncoder std::size_t pos_indices_seq_ = 0; // 缓存键：seq_len // 用给定的位置编码矩阵初始化 pos_emb_（learnable 时额外分配梯度） - void init_(ComputeEngine& engine, Matrix&& pe, bool learnable) + [[nodiscard]] Result init_(ComputeEngine& engine, Matrix&& pe, bool learnable) { const std::size_t rows = pe.rows(); const std::size_t cols = pe.cols(); learnable_ = learnable; auto pe_r = engine.from_matrix(pe); - NN_ASSERT(pe_r, pe_r ? "" : pe_r.error().message.c_str()); + if (!pe_r) return std::unexpected(pe_r.error()); pos_emb_ = std::move(*pe_r); if (learnable_) { grad_pos_emb_ = engine.create_tensor(rows, cols); auto r = engine.zero(grad_pos_emb_); - NN_ASSERT(r, r ? "" : r.error().message.c_str()); + if (!r) return std::unexpected(r.error()); } + return {}; } // 确保 pos_indices 缓存有效（batch-major：i = b*seq + t → position=t） @@ -3395,16 +3544,23 @@ class AdditivePositionEncoder : public PositionEncoder // 与 token_emb_ 共享同一 rng 序列（保持与旧实现完全一致的可复现性）。 class LearnedPositionEncoder final : public AdditivePositionEncoder { + std::size_t d_model_; + std::size_t seq_len_; + std::mt19937_64* rng_; // 非拥有指针 + std::normal_distribution* dist_; // 非拥有指针 + public: - LearnedPositionEncoder(ComputeEngine& engine, std::size_t d_model, - std::size_t seq_len, + LearnedPositionEncoder(std::size_t d_model, std::size_t seq_len, std::mt19937_64& rng, std::normal_distribution& dist) + : d_model_(d_model), seq_len_(seq_len), rng_(&rng), dist_(&dist) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override { - Matrix pe(seq_len, d_model); + Matrix pe(seq_len_, d_model_); auto pe_s = pe.span(); - for (std::size_t i = 0; i < pe.size(); ++i) pe_s[i] = dist(rng); - init_(engine, std::move(pe), /*learnable=*/true); + for (std::size_t i = 0; i < pe.size(); ++i) pe_s[i] = (*dist_)(*rng_); + return init_(engine, std::move(pe), /*learnable=*/true); } }; @@ -3412,20 +3568,25 @@ class LearnedPositionEncoder final : public AdditivePositionEncoder // PE(pos, 2i) = sin(pos/10000^(2i/d)), PE(pos, 2i+1) = cos(...) class SinusoidalPositionEncoder final : public AdditivePositionEncoder { + std::size_t d_model_; + std::size_t seq_len_; + public: - SinusoidalPositionEncoder(ComputeEngine& engine, std::size_t d_model, - std::size_t seq_len) + SinusoidalPositionEncoder(std::size_t d_model, std::size_t seq_len) + : d_model_(d_model), seq_len_(seq_len) {} + + [[nodiscard]] Result init(ComputeEngine& engine) override { - Matrix pe(seq_len, d_model); + Matrix pe(seq_len_, d_model_); auto pe_s = pe.span(); - for (std::size_t pos = 0; pos < seq_len; ++pos) - for (std::size_t i = 0; i < d_model; ++i) + for (std::size_t pos = 0; pos < seq_len_; ++pos) + for (std::size_t i = 0; i < d_model_; ++i) { Scalar angle = static_cast(pos) / - std::pow(Scalar{10000}, static_cast(2 * (i / 2)) / static_cast(d_model)); - pe_s[pos * d_model + i] = (i % 2 == 0) ? std::sin(angle) : std::cos(angle); + std::pow(Scalar{10000}, static_cast(2 * (i / 2)) / static_cast(d_model_)); + pe_s[pos * d_model_ + i] = (i % 2 == 0) ? std::sin(angle) : std::cos(angle); } - init_(engine, std::move(pe), /*learnable=*/false); + return init_(engine, std::move(pe), /*learnable=*/false); } }; @@ -3501,52 +3662,76 @@ class GPTModel final : public Layer bool activation_offload_ = false; public: - GPTModel(ComputeEngine& engine, - std::size_t vocab_size, std::size_t d_model, std::size_t seq_len, + GPTModel(std::size_t vocab_size, std::size_t d_model, std::size_t seq_len, std::size_t num_heads, std::size_t d_ff, std::size_t num_layers, PosEncodingType pos_enc_type = PosEncodingType::Learned, ActivationType activation = ActivationType::GeLU, NormType norm_type = NormType::LayerNorm) : vocab_size_(vocab_size), d_model_(d_model), seq_len_(seq_len), - ln_f_(make_norm_layer(engine, d_model, norm_type)), - lm_head_(engine, d_model, vocab_size) + ln_f_(make_norm_layer(d_model, norm_type)), + lm_head_(d_model, vocab_size) { - // 初始化 token_emb_ - Matrix te(vocab_size, d_model); - constexpr Scalar emb_init_std = 0.02; - std::mt19937_64 rng{42}; - std::normal_distribution dist(0.0, emb_init_std); - auto te_s = te.span(); - for (std::size_t i = 0; i < te.size(); ++i) te_s[i] = dist(rng); - - auto te_r = engine.from_matrix(te); - NN_ASSERT(te_r, te_r ? "" : te_r.error().message.c_str()); - token_emb_ = std::move(*te_r); - - grad_token_emb_ = engine.create_tensor(vocab_size, d_model); - { auto r1 = engine.zero(grad_token_emb_); NN_ASSERT(r1, r1 ? "" : r1.error().message.c_str()); } + blocks_.reserve(num_layers); + for (std::size_t i = 0; i < num_layers; ++i) + blocks_.emplace_back(d_model, num_heads, d_ff, seq_len, seq_len, + pos_enc_type, activation, norm_type); // 初始化位置编码器（Learned / Sinusoidal / ALiBi / RoPE） - // 注意：Learned 与 token_emb_ 共享同一 rng，保证与旧实现完全一致的可复现性。 switch (pos_enc_type) { case PosEncodingType::Learned: - pos_encoder_ = std::make_unique( - engine, d_model, seq_len, rng, dist); + { + std::mt19937_64 rng{42}; + std::normal_distribution dist(0.0, 0.02); + pos_encoder_ = std::make_unique( + d_model, seq_len, rng, dist); + } break; case PosEncodingType::Sinusoidal: pos_encoder_ = std::make_unique( - engine, d_model, seq_len); + d_model, seq_len); break; default: // ALiBi / RoPE：位置信息由注意力层注入 pos_encoder_ = std::make_unique(); break; } + } - // GPTBlock 传入 seq_len、pos_enc_type、activation 和 norm_type - for (std::size_t i = 0; i < num_layers; ++i) - blocks_.emplace_back(engine, d_model, num_heads, d_ff, seq_len, seq_len, - pos_enc_type, activation, norm_type); + [[nodiscard]] Result init(ComputeEngine& engine) override + { + // 初始化 token_emb_ + Matrix te(vocab_size_, d_model_); + constexpr Scalar emb_init_std = 0.02; + std::mt19937_64 rng{42}; + std::normal_distribution dist(0.0, emb_init_std); + auto te_s = te.span(); + for (std::size_t i = 0; i < te.size(); ++i) te_s[i] = dist(rng); + + auto te_r = engine.from_matrix(te); + if (!te_r) return std::unexpected(te_r.error()); + token_emb_ = std::move(*te_r); + + grad_token_emb_ = engine.create_tensor(vocab_size_, d_model_); + { auto r1 = engine.zero(grad_token_emb_); if (!r1) return std::unexpected(r1.error()); } + + // 初始化子层 + if (pos_encoder_) + { + auto r = pos_encoder_->init(engine); + if (!r) return std::unexpected(r.error()); + } + for (auto& block : blocks_) + { + auto r = block.init(engine); + if (!r) return std::unexpected(r.error()); + } + if (ln_f_) + { + auto r = ln_f_->init(engine); + if (!r) return std::unexpected(r.error()); + } + { auto r = lm_head_.init(engine); if (!r) return std::unexpected(r.error()); } + return {}; } std::vector parameters() override diff --git a/include/neuralnet.cpp/compute_loss.hpp b/include/neuralnet.cpp/compute_loss.hpp index 5d567321..fed2b451 100644 --- a/include/neuralnet.cpp/compute_loss.hpp +++ b/include/neuralnet.cpp/compute_loss.hpp @@ -166,8 +166,14 @@ class CrossEntropyLoss final : public Loss // col_sum ≥ 1（因 max 元素 shifted=0 → exp=1），故 log(col_sum) 有限 auto log_col_sum = engine.elementwise_unary(UnaryOp::Log, *col_sum); if (!log_col_sum) return std::unexpected(log_col_sum.error()); - auto log_softmax = engine.elementwise_binary(BinaryOp::Sub, *shifted, *log_col_sum); - if (!log_softmax) return std::unexpected(log_softmax.error()); + + // log_softmax = shifted - log(col_sum)：这是**列广播**（每个元素减去 + // 对应列的 log(col_sum)），不能用 elementwise_binary —— log_col_sum + // 形状为 (1, batch) 而 shifted 为 (classes, batch)，两引擎均要求形状 + // 完全一致会报 shape mismatch。此处 shifted 已不再使用，可就地广播。 + r = engine.broadcast_col_inplace(*shifted, *log_col_sum, BinaryOp::Sub); + if (!r) return std::unexpected(r.error()); + auto log_softmax = std::move(shifted); return DenseSoftmax{/*softmax=*/std::move(*exp_shift), /*log_softmax=*/std::move(*log_softmax)}; diff --git a/include/neuralnet.cpp/compute_optimizer.hpp b/include/neuralnet.cpp/compute_optimizer.hpp index ea314472..b9ad2097 100644 --- a/include/neuralnet.cpp/compute_optimizer.hpp +++ b/include/neuralnet.cpp/compute_optimizer.hpp @@ -56,19 +56,15 @@ class Optimizer } // 为每个参数创建同形状的零初始化 Tensor（供 Momentum/Adam/Muon 复用） - // TODO(1.1, S2): 本方法及各 Layer 构造函数/LayerNorm/RMSNorm 构造器/RoPE::rebuild - // 中的 `NN_ASSERT(, ...)` 在 Release(NDEBUG) 下会吞掉引擎错误。 - // 正确做法（符合项目 Result/expected 规范）：把这些 void/构造函数改为返回 - // Result，并将 `NN_ASSERT(r, ...)` 改为 `if (!r) return std::unexpected(r.error());` - // （或前置条件用 NN_REQUIRE）。因属签名级重构，1.0.0 暂缓。 - std::vector create_zero_buffers_() const + [[nodiscard]] Result> create_zero_buffers_() const { std::vector buffers; buffers.reserve(params_.size()); for (auto& p : params_) { auto buf = engine_.create_tensor(p.get().rows(), p.get().cols()); - { auto r = engine_.zero(buf); NN_ASSERT(r, r ? "" : r.error().message.c_str()); } + auto r = engine_.zero(buf); + if (!r) return std::unexpected(r.error()); buffers.push_back(std::move(buf)); } return buffers; @@ -206,8 +202,17 @@ class SGDWithMomentum : public Optimizer std::vector grads, Scalar lr, Scalar beta = 0.9) : Optimizer(engine, std::move(params), std::move(grads)), - lr_(lr), beta_(beta), - velocities_(create_zero_buffers_()) {} + lr_(lr), beta_(beta) + { + auto v_r = create_zero_buffers_(); + if (!v_r) + { + std::fprintf(stderr, "SGDWithMomentum init failed: %s\n", + v_r.error().message.c_str()); + std::abort(); + } + velocities_ = std::move(*v_r); + } void set_lr(Scalar lr) override { lr_ = lr; } @@ -310,8 +315,22 @@ class Adam : public Optimizer void init_moments_() { - m_ = create_zero_buffers_(); - v_ = create_zero_buffers_(); + auto m_r = create_zero_buffers_(); + if (!m_r) + { + std::fprintf(stderr, "Adam init (m_) failed: %s\n", + m_r.error().message.c_str()); + std::abort(); + } + m_ = std::move(*m_r); + auto v_r = create_zero_buffers_(); + if (!v_r) + { + std::fprintf(stderr, "Adam init (v_) failed: %s\n", + v_r.error().message.c_str()); + std::abort(); + } + v_ = std::move(*v_r); } // 偏差修正系数（Adam/AdamW 共用）：给定下一步步数 t_next， @@ -539,8 +558,17 @@ class Muon : public Optimizer Scalar ns_eps = 1e-7f) : Optimizer(engine, std::move(params), std::move(grads)), lr_(lr), momentum_(momentum), nesterov_(nesterov), - ns_steps_(ns_steps), ns_eps_(ns_eps), - velocities_(create_zero_buffers_()) {} + ns_steps_(ns_steps), ns_eps_(ns_eps) + { + auto v_r = create_zero_buffers_(); + if (!v_r) + { + std::fprintf(stderr, "Muon init failed: %s\n", + v_r.error().message.c_str()); + std::abort(); + } + velocities_ = std::move(*v_r); + } void set_lr(Scalar lr) override { lr_ = lr; } diff --git a/include/neuralnet.cpp/cpu_emitter.hpp b/include/neuralnet.cpp/cpu_emitter.hpp new file mode 100644 index 00000000..f19372f2 --- /dev/null +++ b/include/neuralnet.cpp/cpu_emitter.hpp @@ -0,0 +1,252 @@ +#ifndef NN_CPU_EMITTER_HPP +#define NN_CPU_EMITTER_HPP + +// ═══════════════════════════════════════════════════════════════════════════ +// cpu_emitter.hpp — CPU 后端 emitter（IR-D 第二实现） +// +// 对应文档 `docs/11-ir-optimization.md` IR-D：一份 canonical IR → 多后端。 +// 本 emitter 把 ExprSpec 展开为一段可直接嵌入 C++ 的**逐元素直线代码** +// （独立函数，无解释器开销），作为 GlslEmitter 之外的第二实现，验证 +// emitter 抽象确实支持"同一 IR 多后端"。 +// +// 生成的函数签名： +// void (const float* b0, ..., const float* bN-1, float* bout, +// std::size_t count, std::size_t cols, +// std::uint32_t vp0, ..., float c0, ...); +// - 输入/输出为扁平 float 缓冲（row-major，与 Tensor 布局一致） +// - vpN：运行时视图参数（RowMod 周期 / RotateHalf 块大小，形状无关） +// - cN：常量池 +// 归约表达式（generate_reduce）：本 emitter 暂不生成归约 kernel（返回空串）， +// 归约语义由 CPU eval_expr 解释器承担（先正确后优化）。 +// ═══════════════════════════════════════════════════════════════════════════ + +#include +#include +#include + +#include "expr_spec.hpp" +#include "expr_emitter.hpp" + +namespace nn +{ + +// ── 二元算子 → C++ 运算符片段 ─────────────────────────────────────────── +inline const char* cpu_binary_op(ExprOp op) +{ + switch (op) + { + case ExprOp::Add: return "+"; + case ExprOp::Sub: return "-"; + case ExprOp::Mul: return "*"; + case ExprOp::Div: return "/"; + case ExprOp::Lt: return "<"; + case ExprOp::Le: return "<="; + case ExprOp::Gt: return ">"; + case ExprOp::Ge: return ">="; + case ExprOp::Eq: return "=="; + case ExprOp::Ne: return "!="; + default: return ""; + } +} + +// ── 视图 → C++ 索引表达式 ──────────────────────────────────────────────── +// 布局约定与 GLSL/CPU eval_expr 一致：row = i/cols, col = i%cols。 +// vp_slot：该视图的运行时参数槽位（同 glsl_gen 的 vp 约定）。 +inline void cpu_view_read(std::ostringstream& os, + const ExprView& v, std::uint32_t buf_id, + const std::string& idx_var, + const std::string& row_var, const std::string& col_var, + std::uint32_t vp_slot = 0) +{ + const std::string buf = "b" + std::to_string(buf_id); + switch (v.kind) + { + default: + case static_cast(ExprViewKind::Linear): + os << buf << "[" << idx_var << "]"; + return; + case static_cast(ExprViewKind::RotateHalf): + { + // 与 GLSL 生成器一致的索引映射（block 为运行时视图参数 vpN） + const std::string blk = "(" + row_var + " / vp" + std::to_string(vp_slot) + ")"; + const std::string rl = "(" + row_var + " % vp" + std::to_string(vp_slot) + ")"; + const std::string half = "(vp" + std::to_string(vp_slot) + " / 2u)"; + const std::string hi = "((" + blk + " * vp" + std::to_string(vp_slot) + " + " + + rl + " + " + half + ") * cols + " + col_var + ")"; + const std::string lo = "((" + blk + " * vp" + std::to_string(vp_slot) + " + " + + rl + " - " + half + ") * cols + " + col_var + ")"; + const std::string neg = v.negate_first_half ? "-" : ""; + os << "((" << rl << " < " << half << ") ? " + << neg << buf << "[" << hi << "] : " << buf << "[" << lo << "])"; + return; + } + case static_cast(ExprViewKind::RowMod): + os << buf << "[(" << row_var << " % vp" << std::to_string(vp_slot) + << ") * cols + " << col_var << "]"; + return; + case static_cast(ExprViewKind::RowBroadcast): + os << buf << "[" << row_var << "]"; + return; + case static_cast(ExprViewKind::ColBroadcast): + os << buf << "[" << col_var << "]"; + return; + } +} + +// ── 指令展开（直线代码，逐条输出到 os） ────────────────────────────────── +inline void cpu_emit_instrs(std::ostringstream& os, const ExprSpec& spec) +{ + const auto operand = [&](const ExprOperand& op) -> std::string + { + switch (op.kind) + { + default: + case static_cast(ExprOperandKind::Input): + return "v" + std::to_string(op.idx); + case static_cast(ExprOperandKind::Reg): + case static_cast(ExprOperandKind::Fanout): + return "r" + std::to_string(op.idx); + case static_cast(ExprOperandKind::Const): + return "c" + std::to_string(op.idx); + } + }; + + for (const auto& ins : spec.instrs) + { + const ExprOp op = static_cast(ins.op); + const std::string dst = "r" + std::to_string(ins.dst); + const std::string a = operand(ins.a); + switch (op) + { + case ExprOp::Add: case ExprOp::Sub: case ExprOp::Mul: case ExprOp::Div: + { + const char* s = cpu_binary_op(op); + os << " " << dst << " = " << a << " " << s << " " << operand(ins.b) << ";\n"; + break; + } + case ExprOp::Max: + os << " " << dst << " = std::max(" << a << ", " << operand(ins.b) << ");\n"; + break; + case ExprOp::Min: + os << " " << dst << " = std::min(" << a << ", " << operand(ins.b) << ");\n"; + break; + case ExprOp::Lt: case ExprOp::Le: case ExprOp::Gt: + case ExprOp::Ge: case ExprOp::Eq: case ExprOp::Ne: + { + const char* s = cpu_binary_op(op); + os << " " << dst << " = (" << a << " " << s << " " << operand(ins.b) + << ") ? 1.0f : 0.0f;\n"; + break; + } + case ExprOp::Neg: + os << " " << dst << " = -(" << a << ");\n"; + break; + case ExprOp::Exp: + os << " " << dst << " = std::exp(" << a << ");\n"; + break; + case ExprOp::Log: + os << " " << dst << " = std::log(" << a << ");\n"; + break; + case ExprOp::Sqrt: + os << " " << dst << " = std::sqrt(" << a << ");\n"; + break; + case ExprOp::Rsqrt: + os << " " << dst << " = 1.0f / std::sqrt(" << a << ");\n"; + break; + case ExprOp::Abs: + os << " " << dst << " = std::fabs(" << a << ");\n"; + break; + case ExprOp::Tanh: + os << " " << dst << " = std::tanh(" << a << ");\n"; + break; + case ExprOp::Select: + os << " " << dst << " = (" << a << " != 0.0f) ? " + << operand(ins.b) << " : " << operand(ins.c) << ";\n"; + break; + default: + os << " // 不支持的算子（归约指令不在此 emitter 展开）\n"; + break; + } + } +} + +// ── CpuEmitter — 生成 C++ 逐元素直线代码 ──────────────────────────────── +class CpuEmitter final : public ExprEmitter +{ +public: + [[nodiscard]] std::string_view name() const noexcept override + { return "cpu"; } + + [[nodiscard]] std::string generate( + const std::string& name_, const ExprSpec& spec) override + { + // 归约语义不在本 emitter 展开（见 generate_reduce） + if (expr_spec_reduce_axis(spec) >= 0) + return {}; + const std::size_t n_inputs = spec.views.size(); + std::ostringstream L; + + L << "// ── 自动生成（CpuEmitter），请勿手动编辑 ──\n"; + L << "// 表达式: " << name_ << "\n"; + L << "inline void " << name_ << "(\n"; + for (std::size_t i = 0; i < n_inputs; ++i) + L << " const float* b" << i << ", // 输入 " << i << "\n"; + L << " float* bout,\n"; + L << " std::size_t count, std::size_t cols"; + const std::uint32_t n_vp = expr_spec_runtime_view_param_count(spec); + for (std::uint32_t i = 0; i < n_vp; ++i) + L << ",\n std::uint32_t vp" << i; + for (std::size_t i = 0; i < spec.consts.size(); ++i) + L << ",\n float c" << i; + L << ")\n{\n"; + L << " for (std::size_t i = 0; i < count; ++i) {\n"; + L << " const std::size_t row = i / cols;\n"; + L << " const std::size_t col = i % cols;\n"; + + // 输入读取变量（每输入缓存一次；vp 槽 = 此前运行时参数视图个数） + for (std::size_t i = 0; i < n_inputs; ++i) + { + std::uint32_t vp = 0; + for (std::size_t j = 0; j < i; ++j) + if (expr_view_has_runtime_param( + static_cast(spec.views[j].kind))) + ++vp; + L << " const float v" << i << " = "; + cpu_view_read(L, spec.views[i], static_cast(i), + "i", "row", "col", vp); + L << ";\n"; + } + + // 寄存器声明（先声明后赋值，兼容 IR-B liveness 复用同号寄存器） + if (spec.num_regs > 0) + { + L << " float r0"; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) L << ", r" << r; + L << ";\n"; + } + + cpu_emit_instrs(L, spec); + + const std::uint32_t last_dst = spec.instrs.back().dst; + L << " bout[i] = r" << last_dst << ";\n"; + L << " }\n"; + L << "}\n"; + return L.str(); + } + + [[nodiscard]] std::string generate_reduce( + const std::string&, const ExprSpec&) override + { + // 归约 kernel 由 CPU eval_expr 解释器承担；本 emitter 暂不生成 + return {}; + } +}; + +// 登记到 emitter 注册表 +inline const bool kCpuEmitterRegistered = + emitter_registry::register_backend("cpu", + []() -> ExprEmitter* { return new CpuEmitter(); }); + +} // namespace nn + +#endif // NN_CPU_EMITTER_HPP diff --git a/include/neuralnet.cpp/cpu_engine.hpp b/include/neuralnet.cpp/cpu_engine.hpp index ccfcffb8..e3dfd445 100644 --- a/include/neuralnet.cpp/cpu_engine.hpp +++ b/include/neuralnet.cpp/cpu_engine.hpp @@ -11,11 +11,14 @@ #include #include #include +#include #include "compute_engine.hpp" #include "algebra_compute.hpp" #include "algebra_expr.hpp" #include "expr_opt.hpp" +#include "expr_graph.hpp" // IR-C：图 IR + 融合分析（scan 模式登记用） +#include "expr_registry.hpp" // scan 模式融合 kernel 登记（NN_EXPR_SCAN） namespace nn { @@ -58,6 +61,14 @@ inline std::pair attn_bias_at( // ══════════════════════════════════════════════════════════════════════════ class CpuEngine final : public ComputeEngine { +private: + // IR-C 图 IR 录制（仅构建期 scan 模式启用；普通 CPU 运行保持 no-op）。 + // scan 模式下 begin_expr 开启录制，期间 dsl::compute/compute_reduce 把 + // 表达式加入录制图（见 expr_dsl.hpp 的 scan 分支），end_expr 时融合分析 + // 并把每个融合 kernel 的复合 spec 登记进 global_registry —— 保证 GPU + // 运行时 begin_expr/end_expr 融合出的复合 spec 命中 AOT shader（闭合世界）。 + std::optional recording_; + public: [[nodiscard]] Device device() const noexcept override { return Device::CPU; } @@ -66,9 +77,33 @@ class CpuEngine final : public ComputeEngine [[nodiscard]] Result end_batch() override { return {}; } [[nodiscard]] Result flush_batch() override { return {}; } - // ── 表达式录制：CPU 为 no-op（各表达式直接求值，行为不变） ────────── - [[nodiscard]] Result begin_expr() override { return {}; } - [[nodiscard]] Result end_expr() override { return {}; } + // ── 表达式录制（IR-C）────────────────────────────────────────────── + // 普通 CPU 运行：no-op（各表达式直接求值，行为不变——融合是 GPU 优化）。 + // 构建期 scan（NN_EXPR_SCAN）：begin_expr 开启录制图；end_expr 融合分析 + // 并登记全部融合 kernel 的复合 spec（闭合世界两端一致）。 + [[nodiscard]] Result begin_expr() override + { +#ifdef NN_EXPR_SCAN + recording_.emplace(); + fused::recording_graph_ptr() = &*recording_; +#endif + return {}; + } + + [[nodiscard]] Result end_expr() override + { +#ifdef NN_EXPR_SCAN + if (!recording_) + return {}; + ExprGraph g = std::move(*recording_); + recording_.reset(); + fused::recording_graph_ptr() = nullptr; + auto kernels = fuse_expr_graph(g); + for (auto& k : kernels) + fused::global_registry().add(k.spec); +#endif + return {}; + } // ── 张量工厂 ────────────────────────────────────────────────────────── [[nodiscard]] Tensor create_tensor(std::size_t rows, std::size_t cols) override diff --git a/include/neuralnet.cpp/domain_cnn.hpp b/include/neuralnet.cpp/domain_cnn.hpp index b08b9f55..95571e46 100644 --- a/include/neuralnet.cpp/domain_cnn.hpp +++ b/include/neuralnet.cpp/domain_cnn.hpp @@ -87,15 +87,19 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; if (cv.kernel > h + 2 * cv.padding || cv.kernel > w + 2 * cv.padding) return std::unexpected(Error{"CNN: kernel larger than input spatial size"}); - model.add(engine, c, cv.out_channels, cv.kernel, - cv.stride, cv.padding, h, w); + { + auto r = model.add(c, cv.out_channels, cv.kernel, + cv.stride, cv.padding, h, w); + if (!r) return std::unexpected(r.error()); + } h = conv_out_size(h, cv.kernel, cv.stride, cv.padding); w = conv_out_size(w, cv.kernel, cv.stride, cv.padding); c = cv.out_channels; if (cfg.pool > 0) { - model.add(c, h, w, cfg.pool); + auto r = model.add(c, h, w, cfg.pool); + if (!r) return std::unexpected(r.error()); h = (h - cfg.pool) / cfg.pool + 1; w = (w - cfg.pool) / cfg.pool + 1; } @@ -106,11 +110,20 @@ inline const std::vector MNIST_CNN_FC = {120, 10}; return std::unexpected(Error{"CNN: flattened size is zero"}); // 全连接头：Linear(flatten → fc_dims[0]) → ReLU → ... → Linear(→ num_classes) - model.add(engine, flattened, cfg.fc_dims[0]); + { + auto r = model.add(flattened, cfg.fc_dims[0]); + if (!r) return std::unexpected(r.error()); + } for (std::size_t i = 1; i < cfg.fc_dims.size(); ++i) { - model.add(); - model.add(engine, cfg.fc_dims[i - 1], cfg.fc_dims[i]); + { + auto r = model.add(); + if (!r) return std::unexpected(r.error()); + } + { + auto r = model.add(cfg.fc_dims[i - 1], cfg.fc_dims[i]); + if (!r) return std::unexpected(r.error()); + } } return model; } diff --git a/include/neuralnet.cpp/domain_gpt.hpp b/include/neuralnet.cpp/domain_gpt.hpp index 29d26276..423bb5c7 100644 --- a/include/neuralnet.cpp/domain_gpt.hpp +++ b/include/neuralnet.cpp/domain_gpt.hpp @@ -56,9 +56,12 @@ struct GptConfig { return std::unexpected(Error{"GPT d_model must be divisible by num_heads"}); Model model(engine); - model.add(engine, cfg.vocab_size, cfg.d_model, cfg.seq_len, - cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.pos_enc, - cfg.activation, cfg.norm_type); + { + auto r = model.add(cfg.vocab_size, cfg.d_model, cfg.seq_len, + cfg.num_heads, cfg.d_ff, cfg.num_layers, cfg.pos_enc, + cfg.activation, cfg.norm_type); + if (!r) return std::unexpected(r.error()); + } return model; } diff --git a/include/neuralnet.cpp/domain_mnist.hpp b/include/neuralnet.cpp/domain_mnist.hpp index 51eb8b33..b5d756a4 100644 --- a/include/neuralnet.cpp/domain_mnist.hpp +++ b/include/neuralnet.cpp/domain_mnist.hpp @@ -110,12 +110,21 @@ inline const std::vector MNIST_LAYER_DIMS = { std::size_t in_dim = layer_dims[i]; std::size_t out_dim = layer_dims[i + 1]; - model.add(engine, in_dim, out_dim); + { + auto r = model.add(in_dim, out_dim); + if (!r) return std::unexpected(r.error()); + } if (i < layer_dims.size() - 2) { - model.add_layer(make_norm_layer(engine, out_dim, norm_type)) - .add(); + { + auto r = model.add_layer(make_norm_layer(out_dim, norm_type)); + if (!r) return std::unexpected(r.error()); + } + { + auto r = model.add(); + if (!r) return std::unexpected(r.error()); + } } } return model; @@ -146,9 +155,18 @@ inline const std::vector MNIST_LAYER_DIMS = { const std::size_t num_patches = grid_size * grid_size; Model model(engine); - model.add(engine, img_size, patch_size, d_model); - model.add(engine, d_model, num_heads, d_ff, num_layers, num_patches); - model.add(engine, d_model, MNIST_NUM_CLASSES); + { + auto r = model.add(img_size, patch_size, d_model); + if (!r) return std::unexpected(r.error()); + } + { + auto r = model.add(d_model, num_heads, d_ff, num_layers, num_patches); + if (!r) return std::unexpected(r.error()); + } + { + auto r = model.add(d_model, MNIST_NUM_CLASSES); + if (!r) return std::unexpected(r.error()); + } return model; } diff --git a/include/neuralnet.cpp/expr_dsl.hpp b/include/neuralnet.cpp/expr_dsl.hpp index b0995d49..dbc9dd61 100644 --- a/include/neuralnet.cpp/expr_dsl.hpp +++ b/include/neuralnet.cpp/expr_dsl.hpp @@ -33,6 +33,7 @@ #include "compute_engine.hpp" #include "expr_spec.hpp" #include "expr_registry.hpp" +#include "expr_graph.hpp" // IR-C：录制图（begin_expr/end_expr 扫描登记） #include "algebra_expr.hpp" // nn::Expression / nn::BoolExpression 概念 #include "algebra_ops.hpp" // nn::ops（唯一算子来源，含 op_id()） #include "algebra_matrix.hpp" // Matrix @@ -458,6 +459,16 @@ template auto [spec, inputs] = to_expr_spec(e); if (auto v = validate_expr_spec(spec, inputs.size()); !v) return std::unexpected(v.error()); + // IR-C：begin_expr/end_expr 录制段内，表达式加入录制图（而非直接登记）； + // end_expr 时 CpuEngine 融合分析并登记融合后的复合 spec（闭合世界）。 + if (auto* g = fused::recording_graph()) + { + const int node = g->add_node(spec, inputs, rows, cols, + /*vector_out=*/false); + Tensor t = Tensor::cpu(rows, cols); + t.set_virtual_tag(static_cast(node) + 1); + return t; + } fused::global_registry().add(spec); return Tensor::cpu(rows, cols); #else @@ -530,11 +541,24 @@ template auto [spec, inputs] = to_expr_spec(e); if (auto v = validate_expr_spec(spec, inputs.size()); !v) return std::unexpected(v.error()); - fused::global_registry().add(spec); const int raxis = expr_spec_reduce_axis(spec); - if (raxis == 0) return Tensor::cpu(rows, 1); - if (raxis == 1) return Tensor::cpu(1, cols); - return Tensor::cpu(rows, cols); + const auto placeholder = [&](int node) { + Tensor t = (raxis == 0) ? Tensor::cpu(rows, 1) + : (raxis == 1) ? Tensor::cpu(1, cols) + : Tensor::cpu(rows, cols); + if (node >= 0) + t.set_virtual_tag(static_cast(node) + 1); + return t; + }; + // IR-C：录制段内加入图（end_expr 时融合登记）；否则直接登记 + if (auto* g = fused::recording_graph()) + { + const int node = g->add_node(spec, inputs, rows, cols, + /*vector_out=*/true); + return placeholder(node); + } + fused::global_registry().add(spec); + return placeholder(-1); #else auto [spec, inputs] = to_expr_spec(e); if (auto v = validate_expr_spec(spec, inputs.size()); !v) diff --git a/include/neuralnet.cpp/expr_emitter.hpp b/include/neuralnet.cpp/expr_emitter.hpp new file mode 100644 index 00000000..f4af1a44 --- /dev/null +++ b/include/neuralnet.cpp/expr_emitter.hpp @@ -0,0 +1,103 @@ +#ifndef NN_EXPR_EMITTER_HPP +#define NN_EXPR_EMITTER_HPP + +// ═══════════════════════════════════════════════════════════════════════════ +// expr_emitter.hpp — 后端 emitter 抽象（IR-D） +// +// 对应文档 `docs/11-ir-optimization.md` IR-D：把后端代码生成从 GLSL 专用 +// 抽象为 emitter 接口，实现"一份 canonical IR → 多后端代码"： +// +// IR (canonical ExprSpec) → GlslEmitter / CudaEmitter / CpuEmitter +// +// 设计： +// - ExprEmitter 是纯接口：给定 name + canonical ExprSpec，产出目标后端 +// 源码字符串（失败返回空串，语义与 glsl_gen 的约定一致）。 +// - 现有 generate_glsl / generate_glsl_reduce 成为 GlslEmitter 的实现 +// （见 glsl_gen.hpp），gen_fused 等消费方经接口调用，不再与 GLSL 绑定。 +// - 具体 emitter 类在各自头文件（glsl_gen.hpp / cpu_emitter.hpp）提供， +// 本头只定义接口 + 注册表（按后端名选择 emitter 工厂）。 +// 注意：本头供构建期生成器/工具使用，运行时无需包含（与 glsl_gen 一致）。 +// ═══════════════════════════════════════════════════════════════════════════ + +#include +#include +#include +#include +#include + +#include "expr_spec.hpp" + +namespace nn +{ + +// ── emitter 抽象接口 ────────────────────────────────────────────────────── +// 实现约定（与现有 glsl_gen 语义对齐）： +// - generate：逐元素 kernel（无归约视图/归约指令）。 +// - generate_reduce：归约 kernel（expr_spec_reduce_axis >= 0）； +// 混合轴（== -2）或不支持的归约结构返回空串。 +// - 产物均以"自动生成，请勿手动编辑"注释开头，便于区分手工代码。 +class ExprEmitter +{ +public: + virtual ~ExprEmitter() = default; + + // 后端名（"glsl" / "cpu" / ...），供注册表/日志识别 + [[nodiscard]] virtual std::string_view name() const noexcept = 0; + + // 逐元素 kernel 源码；失败返回空串（调用方报错） + [[nodiscard]] virtual std::string generate( + const std::string& name_, const ExprSpec& spec) = 0; + + // 归约 kernel 源码；不支持（混合轴/超槽）返回空串 + [[nodiscard]] virtual std::string generate_reduce( + const std::string& name_, const ExprSpec& spec) = 0; +}; + +// ── 简单注册表：按后端名选择 emitter 工厂 ──────────────────────────────── +// 供生成器工具（gen_fused 等）按 --backend 选择目标 emitter。 +// 具体 emitter 类型在各自头文件中向本注册表登记（见 GlslEmitter / CpuEmitter）。 +// 线程安全：初始化期单线程登记（静态初始化），运行期只读查询。 +namespace emitter_registry +{ + +using Factory = ExprEmitter* (*)(); + +// 后端表（静态初始化，见各 emitter 头文件的登记语句） +inline std::unordered_map& backends() +{ + static std::unordered_map m; + return m; +} + +// 登记后端（返回 false 表示已存在同名后端，拒绝覆盖） +inline bool register_backend(std::string_view name, Factory factory) +{ + auto& m = backends(); + return m.emplace(std::string(name), factory).second; +} + +// 按名创建 emitter；未登记返回 nullptr +[[nodiscard]] inline std::unique_ptr make(std::string_view name) +{ + const auto& m = backends(); + const auto it = m.find(std::string(name)); + if (it == m.end()) + return nullptr; + return std::unique_ptr(it->second()); +} + +// 所有已登记后端名（调试/日志） +[[nodiscard]] inline std::vector names() +{ + std::vector out; + out.reserve(backends().size()); + for (const auto& [k, v] : backends()) + out.push_back(k); + return out; +} + +} // namespace emitter_registry + +} // namespace nn + +#endif // NN_EXPR_EMITTER_HPP diff --git a/include/neuralnet.cpp/expr_graph.hpp b/include/neuralnet.cpp/expr_graph.hpp new file mode 100644 index 00000000..317f952e --- /dev/null +++ b/include/neuralnet.cpp/expr_graph.hpp @@ -0,0 +1,304 @@ +#ifndef NN_EXPR_GRAPH_HPP +#define NN_EXPR_GRAPH_HPP + +// ═══════════════════════════════════════════════════════════════════════════ +// expr_graph.hpp — 图 IR + 融合分析（IR-C） +// +// 对应文档 `docs/11-ir-optimization.md` IR-C：把扁平 ExprSpec 演进为图 IR +// （DAG），为 begin_expr/end_expr 提供多表达式融合分析。 +// +// 图结构： +// - 节点 = 一次 eval_expr / eval_expr_reduce（一个 canonical ExprSpec）。 +// - 边 = 数据依赖：节点的某个输入槽引用前序节点的输出（虚拟寄存器）。 +// - 录制：begin_expr 开启录制图；期间每次表达式求值把节点加入图并返回 +// 携带 virtual_tag 的占位 Tensor；end_expr 做融合分析并产出 kernel 序列。 +// +// 融合策略（本阶段落地：逐元素链拼接）： +// - 相邻节点 A → B（B 以 Linear 视图消费 A 的输出，且 A 无其他消费者） +// 且 A、B 均为逐元素（无归约）、形状相同 → 把 A 的指令内联进 B， +// A 的输出寄存器直接作为 B 的操作数 → 中间结果不落显存，一次 dispatch。 +// - 归约表达式 / 归约输出 / matmul 等作为融合边界 → 独立 kernel。 +// - 拼接后 canonicalize + validate；超限（instrs/regs/inputs/consts） +// 或非法 → 放弃融合（保守，保证正确性优先）。 +// +// 确定性铁律（与 expr_opt.hpp 一致）：贪心按节点序从前到后；拼接重映射 +// 顺序固定；结果经 canonicalize_expr_spec 统一规范化（两端一致）。 +// ═══════════════════════════════════════════════════════════════════════════ + +#include +#include +#include +#include +#include +#include + +#include "tensor.hpp" +#include "expr_spec.hpp" +#include "expr_opt.hpp" + +namespace nn +{ + +// ── 图 IR 节点 ──────────────────────────────────────────────────────────── +struct ExprGraphNode +{ + ExprSpec spec; // canonical IR（录制时已 canonicalize） + std::size_t rows = 0; // 输出形状（统一网格 (rows, cols)） + std::size_t cols = 0; + bool vector_out = false; // 输出为归约向量（(rows,1)/(1,cols)） + // 依赖：dep_of_input[k] = 该输入槽依赖的图节点下标；-1 = 外部输入。 + // 长度 == spec.views.size()（与输入一一对应）。 + std::vector dep_of_input; + // 外部输入张量：dep_of_input[k] < 0 时有效（Layer 传入的真实输入）。 + std::vector input_tensors; +}; + +// ── 图（录制期累积） ───────────────────────────────────────────────────── +struct ExprGraph +{ + std::vector nodes; + // 占位 Tensor 的 virtual_tag → 节点下标（录制时建立，供依赖识别） + std::unordered_map node_by_tag; + + // 录制：把一个表达式加入图，返回节点下标。 + // inputs 中带 virtual_tag 的占位 Tensor 会被识别为前序节点输出（依赖边）。 + [[nodiscard]] int add_node(const ExprSpec& raw_spec, + std::span inputs, + std::size_t rows, std::size_t cols, + bool vector_out) + { + const ExprSpec spec = canonicalize_expr_spec(raw_spec); + const int idx = static_cast(nodes.size()); + ExprGraphNode nd; + nd.spec = spec; + nd.rows = rows; + nd.cols = cols; + nd.vector_out = vector_out; + nd.dep_of_input.reserve(spec.views.size()); + nd.input_tensors.reserve(spec.views.size()); + for (std::size_t k = 0; k < spec.views.size(); ++k) + { + const std::uint64_t tag = (k < inputs.size()) + ? inputs[k].virtual_tag() : 0; + const auto it = (tag != 0) ? node_by_tag.find(tag) : node_by_tag.end(); + if (it != node_by_tag.end()) + { + nd.dep_of_input.push_back(it->second); + nd.input_tensors.emplace_back(); // 占位（依赖节点输出） + } + else + { + nd.dep_of_input.push_back(-1); + nd.input_tensors.push_back( + (k < inputs.size()) ? inputs[k] : Tensor{}); + } + } + nodes.push_back(std::move(nd)); + // 登记输出标记（tag = 节点下标 + 1，0 保留给"非节点输出"） + node_by_tag[static_cast(idx) + 1ull] = idx; + return idx; + } +}; + +// ── 融合 kernel（融合分析输出） ────────────────────────────────────────── +// 一个 kernel = 一个（可能为多个节点拼接的）复合 ExprSpec + 外部输入来源。 +struct FusedKernelInput +{ + int node = -1; // >= 0：该输入是图节点 node 的输出（作为外部 buffer 绑定） + Tensor external; // node < 0：Layer 外部输入 +}; + +struct FusedKernel +{ + ExprSpec spec; // 复合 spec（canonical） + std::vector inputs; // 与 spec.views 一一对应 + std::size_t rows = 0; + std::size_t cols = 0; + bool vector_out = false; + int tail = -1; // 末尾节点（输出节点）在图中下标 + std::vector members; // kernel 包含的节点（顺序 = 融合顺序） +}; + +// ── 全局录制状态（线程局部） ──────────────────────────────────────────── +// begin_expr/end_expr（GPU 引擎执行 / scan 引擎登记）开启/关闭录制图。 +// dsl::compute 的 scan 分支与 GpuEngine::eval_expr 通过 recording_graph() +// 判断"当前是否在录制"以及把表达式加入录制图。 +namespace fused +{ + +[[nodiscard]] inline ExprGraph*& recording_graph_ptr() +{ + static thread_local ExprGraph* g = nullptr; + return g; +} +[[nodiscard]] inline ExprGraph* recording_graph() noexcept +{ return recording_graph_ptr(); } +[[nodiscard]] inline bool is_recording() noexcept +{ return recording_graph_ptr() != nullptr; } + +} // namespace fused + +// ── 融合分析（IR-C 核心，纯函数） ─────────────────────────────────────── +// 贪心按节点序：维护"当前打开的 kernel"（其末尾节点为 tail）。节点 i 若 +// 满足"并入"条件则并入，否则新开 kernel。 +[[nodiscard]] inline std::vector fuse_expr_graph(const ExprGraph& g) +{ + std::vector kernels; + if (g.nodes.empty()) + return kernels; + + // 消费者表：节点 j 被哪些节点消费（dep_of_input 引用 j） + std::vector> consumers(g.nodes.size()); + for (std::size_t i = 0; i < g.nodes.size(); ++i) + for (const int dep : g.nodes[i].dep_of_input) + if (dep >= 0) + consumers[static_cast(dep)].push_back(static_cast(i)); + + // ── 尝试把节点 bi 并入当前 kernel（cur） ─────────────────────────── + const auto try_append = [&](FusedKernel& cur, int bi) -> bool + { + const ExprGraphNode& B = g.nodes[static_cast(bi)]; + const int tail = cur.tail; + const ExprSpec& A = cur.spec; // 当前复合 spec（含 tail 及之前节点） + const ExprSpec& Bs = B.spec; + + // 1) 前置：均逐元素（无归约）、非归约输出、形状相同 + if (cur.vector_out || B.vector_out) return false; + if (expr_spec_reduce_axis(A) != -1) return false; + if (expr_spec_reduce_axis(Bs) != -1) return false; + if (cur.rows != B.rows || cur.cols != B.cols) return false; + + // 2) tail 无其他消费者，且被 B 以 Linear 视图消费 + if (consumers[static_cast(tail)].size() != 1 || + consumers[static_cast(tail)][0] != bi) + return false; + std::vector tail_slots; // B 中引用 tail 的输入槽 + std::vector other_slots; // B 中其他输入槽 + for (std::size_t k = 0; k < Bs.views.size(); ++k) + { + if (B.dep_of_input[k] == tail) + { + if (static_cast(Bs.views[k].kind) != ExprViewKind::Linear) + return false; // 仅支持 Linear 引用（寄存器直接替换） + tail_slots.push_back(static_cast(k)); + } + else + { + other_slots.push_back(static_cast(k)); + } + } + if (tail_slots.empty()) + return false; + + // 3) 拼接（确定性：tail 的 views/consts/instrs 在前，B 的在后） + const std::size_t na = A.instrs.size(); + const std::size_t out_reg_A = A.instrs.back().dst; + + std::vector views = A.views; + std::vector b_view_map(Bs.views.size(), -1); + for (const int k : other_slots) + { + b_view_map[static_cast(k)] = static_cast(views.size()); + views.push_back(Bs.views[static_cast(k)]); + } + + std::vector consts = A.consts; + const std::size_t b_const_base = A.consts.size(); + // 关键：B 的常量池必须追加在 A 之后（B 的 Const 引用经偏移指向这里） + consts.insert(consts.end(), Bs.consts.begin(), Bs.consts.end()); + + std::vector instrs = A.instrs; + for (const auto& ins : Bs.instrs) + { + ExprInstr ni = ins; + ni.dst = static_cast(ni.dst + na); + const auto remap_op = [&](ExprOperand o) -> ExprOperand + { + if (o.kind == static_cast(ExprOperandKind::Input)) + { + const std::size_t k = o.idx; + const auto it = std::find(tail_slots.begin(), tail_slots.end(), + static_cast(k)); + if (it != tail_slots.end()) + return expr::reg(static_cast(out_reg_A)); + return expr::input(static_cast( + b_view_map[k])); + } + if (o.kind == static_cast(ExprOperandKind::Const)) + return expr::cst(static_cast( + o.idx + b_const_base)); + if (o.kind == static_cast(ExprOperandKind::Reg) || + o.kind == static_cast(ExprOperandKind::Fanout) || + o.kind == static_cast(ExprOperandKind::Reduce)) + return {o.kind, static_cast(o.idx + na)}; + return o; + }; + const std::size_t nops = + expr_instr_num_operands(static_cast(ins.op)); + ni.a = remap_op(ins.a); + if (nops >= 2) ni.b = remap_op(ins.b); + if (nops >= 3) ni.c = remap_op(ins.c); + instrs.push_back(ni); + } + + ExprSpec fused; + fused.views = std::move(views); + fused.consts = std::move(consts); + fused.instrs = std::move(instrs); + fused.num_regs = static_cast(na + Bs.num_regs); + fused = canonicalize_expr_spec(fused); + if (auto v = validate_expr_spec(fused, fused.views.size()); !v) + return false; // 超限/非法 → 放弃融合（保守） + + // 4) 提交：更新 cur + std::vector new_inputs = cur.inputs; + for (const int k : other_slots) + { + const std::size_t kk = static_cast(k); + if (B.dep_of_input[kk] < 0) + new_inputs.push_back(FusedKernelInput{ + /*node=*/-1, /*external=*/B.input_tensors[kk]}); + else + new_inputs.push_back(FusedKernelInput{ + /*node=*/B.dep_of_input[kk], /*external=*/Tensor{}}); + } + cur.spec = std::move(fused); + cur.inputs = std::move(new_inputs); + cur.rows = B.rows; + cur.cols = B.cols; + cur.tail = bi; + cur.members.push_back(bi); + return true; + }; + + // ── 贪心主循环 ───────────────────────────────────────────────────── + for (std::size_t i = 0; i < g.nodes.size(); ++i) + { + const auto& N = g.nodes[i]; + if (!kernels.empty() && try_append(kernels.back(), static_cast(i))) + continue; + + FusedKernel k; + k.spec = N.spec; + k.rows = N.rows; + k.cols = N.cols; + k.vector_out = N.vector_out; + k.tail = static_cast(i); + k.members.push_back(static_cast(i)); + for (std::size_t kk = 0; kk < N.spec.views.size(); ++kk) + { + if (N.dep_of_input[kk] < 0) + k.inputs.push_back(FusedKernelInput{ + /*node=*/-1, /*external=*/N.input_tensors[kk]}); + else + k.inputs.push_back(FusedKernelInput{ + /*node=*/N.dep_of_input[kk], /*external=*/Tensor{}}); + } + kernels.push_back(std::move(k)); + } + return kernels; +} + +} // namespace nn + +#endif // NN_EXPR_GRAPH_HPP diff --git a/include/neuralnet.cpp/glsl_gen.hpp b/include/neuralnet.cpp/glsl_gen.hpp index 5f25a60f..805380fa 100644 --- a/include/neuralnet.cpp/glsl_gen.hpp +++ b/include/neuralnet.cpp/glsl_gen.hpp @@ -19,6 +19,7 @@ #include #include "expr_spec.hpp" +#include "expr_emitter.hpp" // IR-D：emitter 抽象（GlslEmitter 实现） namespace nn { @@ -110,9 +111,41 @@ inline void glsl_view_read(std::ostringstream& os, } } +// ── 视图是否使用行/列索引（用于按需发射 row/col，省去每元素整数除法/取模）── +// Linear 用扁平索引、RowBroadcast 用 row、ColBroadcast 用 col、 +// RowMod/RotateHalf 用 row+col（索引映射逐元素变化，见 glsl_view_read）。 +inline bool glsl_view_uses_row(ExprViewKind k) +{ + return k != ExprViewKind::Linear && k != ExprViewKind::ColBroadcast; +} +inline bool glsl_view_uses_col(ExprViewKind k) +{ + return k != ExprViewKind::Linear && k != ExprViewKind::RowBroadcast; +} + +// ── vec4 向量化资格 ─────────────────────────────────────────────────────── +// 仅当全部视图 ∈ {Linear, RowBroadcast, ColBroadcast} 时才可把 4 个相邻元素 +// 打包成 vec4（RowMod/RotateHalf 的索引映射逐通道变化，v1 不向量化）。 +// 注意：cols 是运行时 push constant，故向量化路径用运行时 `cols%4==0` 守卫， +// 任意形状仍走同 kernel 内的标量回退——正确性不依赖编译期形状。 +inline bool glsl_vec4_eligible(const ExprSpec& spec) +{ + for (const auto& v : spec.views) + { + const auto k = static_cast(v.kind); + if (k != ExprViewKind::Linear && k != ExprViewKind::RowBroadcast && + k != ExprViewKind::ColBroadcast) + return false; + } + return true; +} + // ── 生成器主入口：ExprSpec → GLSL 源码 ──────────────────────────────────── // 生成的 shader 绑定：输入 buffer binding=0..N-1，输出 binding=N； // push constants: uint count, uint cols, [float c0..]（视图索引需要 cols）。 +// 逐元素表达式若 glsl_vec4_eligible，发射 vec4 kernel：每线程处理 4 个相邻元素 +// （vec4 加载/运算/存储），dispatch 宽度 4（run_fused_gpu 按 vec_width 缩放）； +// 运行时 cols%4!=0 或尾部回退到同 kernel 内标量循环。其余表达式发纯标量 kernel。 inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) { const std::size_t n_inputs = spec.views.size(); @@ -140,28 +173,162 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) L << " float c" << i << ";\n"; L << "};\n\n"; + const bool vec4_ok = glsl_vec4_eligible(spec); + const bool need_row = [&]{ + for (const auto& v : spec.views) + if (glsl_view_uses_row(static_cast(v.kind))) return true; + return false; }(); + const bool need_col = [&]{ + for (const auto& v : spec.views) + if (glsl_view_uses_col(static_cast(v.kind))) return true; + return false; }(); + const std::uint32_t last_dst = spec.instrs.back().dst; + + // ── 标量操作数求值 ── + const auto operand = [&](const ExprOperand& op) -> std::string + { + switch (op.kind) + { + default: + case static_cast(ExprOperandKind::Input): + return "v" + std::to_string(op.idx); + case static_cast(ExprOperandKind::Reg): + case static_cast(ExprOperandKind::Fanout): // Fanout 语义同 Reg + return "r" + std::to_string(op.idx); + case static_cast(ExprOperandKind::Const): + return "c" + std::to_string(op.idx); + } + }; + + // 发射一段标量指令链（输入读取 + 寄存器声明 + 指令展开，直线代码； + // 寄存器先声明、后赋值——兼容 IR-B liveness 复用同号寄存器）。 + // idx_var：扁平元素索引；row_var/col_var：行/列索引表达式（可为 "i/cols" 等）。 + const auto emit_scalar_chain = [&](std::ostringstream& o, + const std::string& indent, + const std::string& idx_var, + const std::string& row_var, + const std::string& col_var) + { + // 输入读取变量（每输入缓存一次）；vp 槽 = 此前运行时参数视图个数 + for (std::size_t i = 0; i < n_inputs; ++i) + { + std::uint32_t vp = 0; + for (std::size_t j = 0; j < i; ++j) + if (expr_view_has_runtime_param( + static_cast(spec.views[j].kind))) ++vp; + o << indent << "const float v" << i << " = "; + glsl_view_read(o, spec.views[i], static_cast(i), + idx_var, row_var, col_var, vp); + o << ";\n"; + } + if (spec.num_regs > 0) + { + o << indent << "float r0"; + for (std::uint32_t r = 1; r < spec.num_regs; ++r) o << ", r" << r; + o << ";\n"; + } + for (const auto& ins : spec.instrs) + { + const ExprOp op = static_cast(ins.op); + const std::string dst = "r" + std::to_string(ins.dst); + const std::string a = operand(ins.a); + switch (op) + { + case ExprOp::Add: case ExprOp::Sub: case ExprOp::Mul: case ExprOp::Div: + { + bool cmp = false; + const char* s = glsl_binary_op(op, cmp); + o << indent << dst << " = " << a << " " << s << " " + << operand(ins.b) << ";\n"; + break; + } + case ExprOp::Max: case ExprOp::Min: + { + const char* s = (op == ExprOp::Max) ? "max" : "min"; + o << indent << dst << " = " << s << "(" << a << ", " + << operand(ins.b) << ");\n"; + break; + } + case ExprOp::Lt: case ExprOp::Le: case ExprOp::Gt: + case ExprOp::Ge: case ExprOp::Eq: case ExprOp::Ne: + { + bool cmp = false; + const char* s = glsl_binary_op(op, cmp); + o << indent << dst << " = (" << a << " " << s << " " + << operand(ins.b) << ") ? 1.0 : 0.0;\n"; + break; + } + case ExprOp::Neg: + o << indent << dst << " = -(" << a << ");\n"; + break; + case ExprOp::Exp: case ExprOp::Log: case ExprOp::Sqrt: + case ExprOp::Rsqrt: case ExprOp::Abs: case ExprOp::Tanh: + { + const char* s = glsl_unary_op(op); + o << indent << dst << " = " << s << "(" << a << ");\n"; + break; + } + case ExprOp::Select: + o << indent << dst << " = (" << a << " != 0.0) ? " + << operand(ins.b) << " : " << operand(ins.c) << ";\n"; + break; + default: + return; // 未知算子：让上层报错 + } + } + }; + + if (!vec4_ok) + { + // ── 纯标量 kernel（含 RowMod/RotateHalf 等无法向量化的视图）── + L << "void main()\n{\n"; + L << " const uint i = gl_GlobalInvocationID.x;\n"; + L << " if (i >= count) return;\n"; + if (need_row && need_col) + L << " const uint row = i / cols;\n const uint col = i % cols;\n"; + else if (need_row) + L << " const uint row = i / cols;\n"; + else if (need_col) + L << " const uint col = i % cols;\n"; + emit_scalar_chain(L, " ", "i", + need_row ? "row" : "", need_col ? "col" : ""); + L << " bout[i] = r" << last_dst << ";\n"; + L << "}\n"; + return L.str(); + } + + // ── vec4 kernel：每线程处理 4 个相邻元素 ──────────────────────────── + // 快速路径（cols%4==0 且组内 4 元素齐全）→ vec4 加载/运算/存储； + // 否则回退到同 kernel 内标量循环（任意 cols / 尾部）。dispatch 宽度 4。 L << "void main()\n{\n"; L << " const uint i = gl_GlobalInvocationID.x;\n"; - L << " if (i >= count) return;\n"; - L << " const uint row = i / cols;\n"; - L << " const uint col = i % cols;\n"; - - // 输入读取变量（每输入缓存一次）；vp 槽 = 此前运行时参数视图个数 + L << " const uint base = i * 4u;\n"; + L << " if (base >= count) return;\n"; + L << " if (cols % 4u == 0u && base + 3u < count) {\n"; + if (need_row && need_col) + L << " const uint row = base / cols;\n const uint col0 = base % cols;\n"; + else if (need_row) + L << " const uint row = base / cols;\n"; + else if (need_col) + L << " const uint col0 = base % cols;\n"; + + // vec4 输入读取（Linear/ColBroadcast 相邻 4 元素 → glslc 合并为 vec4 加载； + // RowBroadcast → splat） for (std::size_t i = 0; i < n_inputs; ++i) { - std::uint32_t vp = 0; - for (std::size_t j = 0; j < i; ++j) - if (expr_view_has_runtime_param( - static_cast(spec.views[j].kind))) - ++vp; - L << " const float v" << i << " = "; - glsl_view_read(L, spec.views[i], static_cast(i), - "i", "row", "col", vp); - L << ";\n"; + const auto k = static_cast(spec.views[i].kind); + if (k == ExprViewKind::Linear) + L << " const vec4 v" << i << " = vec4(b" << i << "[base], b" << i + << "[base+1u], b" << i << "[base+2u], b" << i << "[base+3u]);\n"; + else if (k == ExprViewKind::RowBroadcast) + L << " const vec4 v" << i << " = vec4(b" << i << "[row]);\n"; + else // ColBroadcast + L << " const vec4 v" << i << " = vec4(b" << i << "[col0], b" << i + << "[col0+1u], b" << i << "[col0+2u], b" << i << "[col0+3u]);\n"; } - // 操作数求值 - auto operand = [&](const ExprOperand& op) -> std::string + // vec4 指令展开（GLSL 对 vec4 重载运算/内置函数；常量广播为 vec4） + const auto vec4_operand = [&](const ExprOperand& op) -> std::string { switch (op.kind) { @@ -169,17 +336,15 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) case static_cast(ExprOperandKind::Input): return "v" + std::to_string(op.idx); case static_cast(ExprOperandKind::Reg): - case static_cast(ExprOperandKind::Fanout): // Fanout 语义同 Reg + case static_cast(ExprOperandKind::Fanout): return "r" + std::to_string(op.idx); case static_cast(ExprOperandKind::Const): - return "c" + std::to_string(op.idx); + return "vec4(c" + std::to_string(op.idx) + ")"; } }; - - // 指令展开（直线代码；寄存器先声明、后赋值——兼容 IR-B liveness 复用同号寄存器） if (spec.num_regs > 0) { - L << " float r0"; + L << " vec4 r0"; for (std::uint32_t r = 1; r < spec.num_regs; ++r) L << ", r" << r; L << ";\n"; } @@ -187,53 +352,80 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) { const ExprOp op = static_cast(ins.op); const std::string dst = "r" + std::to_string(ins.dst); - const std::string a = operand(ins.a); + const std::string a = vec4_operand(ins.a); switch (op) { case ExprOp::Add: case ExprOp::Sub: case ExprOp::Mul: case ExprOp::Div: { bool cmp = false; const char* s = glsl_binary_op(op, cmp); - L << " " << dst << " = " << a << " " << s << " " << operand(ins.b) << ";\n"; + L << " " << dst << " = " << a << " " << s << " " + << vec4_operand(ins.b) << ";\n"; break; } case ExprOp::Max: case ExprOp::Min: { const char* s = (op == ExprOp::Max) ? "max" : "min"; - L << " " << dst << " = " << s << "(" << a << ", " << operand(ins.b) << ");\n"; + L << " " << dst << " = " << s << "(" << a << ", " + << vec4_operand(ins.b) << ");\n"; break; } case ExprOp::Lt: case ExprOp::Le: case ExprOp::Gt: case ExprOp::Ge: case ExprOp::Eq: case ExprOp::Ne: { - bool cmp = false; - const char* s = glsl_binary_op(op, cmp); - L << " " << dst << " = (" << a << " " << s << " " << operand(ins.b) - << ") ? 1.0 : 0.0;\n"; + const char* s = nullptr; + switch (op) + { + case ExprOp::Lt: s = "lessThan"; break; + case ExprOp::Le: s = "lessThanEqual"; break; + case ExprOp::Gt: s = "greaterThan"; break; + case ExprOp::Ge: s = "greaterThanEqual"; break; + case ExprOp::Eq: s = "equal"; break; + case ExprOp::Ne: s = "notEqual"; break; + default: break; + } + L << " " << dst << " = vec4(" << s << "(" << a << ", " + << vec4_operand(ins.b) << "));\n"; break; } case ExprOp::Neg: - L << " " << dst << " = -(" << a << ");\n"; + L << " " << dst << " = -(" << a << ");\n"; break; case ExprOp::Exp: case ExprOp::Log: case ExprOp::Sqrt: case ExprOp::Rsqrt: case ExprOp::Abs: case ExprOp::Tanh: { const char* s = glsl_unary_op(op); - L << " " << dst << " = " << s << "(" << a << ");\n"; + L << " " << dst << " = " << s << "(" << a << ");\n"; break; } case ExprOp::Select: - L << " " << dst << " = (" << a << " != 0.0) ? " - << operand(ins.b) << " : " << operand(ins.c) << ";\n"; + L << " " << dst << " = mix(" << vec4_operand(ins.b) << ", " + << vec4_operand(ins.c) << ", notEqual(" << a + << ", vec4(0.0)));\n"; break; default: return std::string{}; // 未知算子：让上层报错 } } + L << " bout[base] = r" << last_dst << ".x;\n"; + L << " bout[base+1u] = r" << last_dst << ".y;\n"; + L << " bout[base+2u] = r" << last_dst << ".z;\n"; + L << " bout[base+3u] = r" << last_dst << ".w;\n"; + L << " return;\n"; + L << " }\n"; - // 输出 = 最后一条指令的目标寄存器 - const std::uint32_t last_dst = spec.instrs.back().dst; - L << " bout[i] = r" << last_dst << ";\n"; + // 标量回退：处理本线程负责的至多 4 个元素（任意 cols / 尾部） + L << " for (uint e = base; e < count && e < base + 4u; ++e) {\n"; + if (need_row && need_col) + L << " const uint row = e / cols;\n const uint col = e % cols;\n"; + else if (need_row) + L << " const uint row = e / cols;\n"; + else if (need_col) + L << " const uint col = e % cols;\n"; + emit_scalar_chain(L, " ", "e", + need_row ? "row" : "", need_col ? "col" : ""); + L << " bout[e] = r" << last_dst << ";\n"; + L << " }\n"; L << "}\n"; return L.str(); } @@ -295,7 +487,9 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) std::ostringstream L; L << "// ── 自动生成（AOT 算子融合 · 归约 kernel），请勿手动编辑 ──\n"; L << "// 表达式: " << name << "\n"; - L << "#version 450\n\n"; + L << "#version 450\n"; + L << "#extension GL_KHR_shader_subgroup_basic : enable\n"; + L << "#extension GL_KHR_shader_subgroup_arithmetic : enable\n\n"; L << "layout(local_size_x = 256) in;\n\n"; for (std::size_t i = 0; i < n_inputs; ++i) L << "layout(std430, binding = " << i << ") readonly buffer Buf" << i @@ -318,8 +512,12 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) L << "void main()\n{\n"; L << " const uint tid = gl_LocalInvocationID.x;\n"; L << " const uint idx = gl_WorkGroupID.x;\n"; - L << (is_row ? " if (idx >= rows) return;\n" - : " if (idx >= cols) return;\n"); + // 行归约：工作组=单行，idx= rows) return;\n" : ""); + + // 归约结果索引：行=工作组单行（s_red[slot][0]）；列 tile=每线程一列（s_red[slot][tid]） + const std::string red_idx = is_row ? "0" : "tid"; // ── 操作数求值（当前元素由 row/col 变量给出） ── const auto operand = [&](const ExprOperand& op) -> std::string @@ -333,13 +531,13 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) case static_cast(ExprOperandKind::Const): return "c" + std::to_string(op.idx); case static_cast(ExprOperandKind::Reduce): - return "s_red[" + std::to_string(slot_of_instr[op.idx]) + "][0]"; + return "s_red[" + std::to_string(slot_of_instr[op.idx]) + "][" + red_idx + "]"; case static_cast(ExprOperandKind::Input): { const ExprView& v = spec.views[op.idx]; const ExprViewKind vk = static_cast(v.kind); if (expr_view_is_reduce(vk)) - return "s_red[" + std::to_string(slot_of_view[op.idx]) + "][0]"; + return "s_red[" + std::to_string(slot_of_view[op.idx]) + "][" + red_idx + "]"; if (vk == ExprViewKind::RowBroadcast) return "b" + std::to_string(op.idx) + "[row]"; if (vk == ExprViewKind::ColBroadcast) @@ -432,117 +630,243 @@ inline std::string generate_glsl(const std::string& name, const ExprSpec& spec) return is_max ? "max(" + a + ", " + b + ")" : "(" + a + " + " + b + ")"; }; - // 归约树形归约（barrier 在每步内，最后一步后全线程可见 s_red[slot][0]） + // 归约：warp shuffle 蝴蝶归约（subgroup）替代共享内存树形归约。 + // 第 1 步：warp 内 subgroupAdd/subgroupMax → 每 warp 一个部分和（零共享/屏障） + // 第 2 步：首 warp 归约全部 warp 部分和 → s_red[slot][0]（屏障后全线程可见） + // 真机(NVIDIA, prime-run)大矩阵实测有效（combine 步骤不再是多步 barrier 串行）。 const auto emit_tree_reduce = [&](int slot, bool is_max) { const std::string s = "s_red[" + std::to_string(slot) + "]"; + const char* sub = is_max ? "subgroupMax" : "subgroupAdd"; + const std::string idt = is_max ? "uintBitsToFloat(0xff800000u)" : "0.0"; + const std::string vn = "v" + std::to_string(slot); // 按槽唯一命名，避免重定义 + const std::string wn = "w" + std::to_string(slot); + L << " barrier();\n"; + L << " float " << vn << " = " << s << "[tid];\n"; + L << " " << vn << " = " << sub << "(" << vn << ");\n"; + L << " if (gl_SubgroupInvocationID == 0u) " << s << "[gl_SubgroupID] = " << vn << ";\n"; L << " barrier();\n"; - L << " for (uint s = 128u; s > 0u; s >>= 1u) {\n"; - L << " if (tid < s) " << s << "[tid] = " - << combine(is_max, s + "[tid]", s + "[tid + s]") << ";\n"; - L << " barrier();\n"; + L << " if (gl_SubgroupID == 0u) {\n"; + L << " float " << wn << " = (gl_SubgroupInvocationID < gl_NumSubgroups) ? " + << s << "[gl_SubgroupInvocationID] : " << idt << ";\n"; + L << " " << wn << " = " << sub << "(" << wn << ");\n"; + L << " if (gl_SubgroupInvocationID == 0u) " << s << "[0] = " << wn << ";\n"; L << " }\n"; + L << " barrier();\n"; }; - // 元素索引声明（行归约：固定行、循环列；列归约：循环行、固定列） - const std::string loop_decl = is_row - ? " for (uint c = tid; c < cols; c += 256u) {\n const uint row = idx;\n const uint col = c;\n" - : " for (uint r = tid; r < rows; r += 256u) {\n const uint row = r;\n const uint col = idx;\n"; - const std::string out_idx = "row*cols + col"; + // 组合 4 路累加器（sum/max） + const auto combine4 = [&](bool is_max, const std::string& a0, const std::string& a1, + const std::string& a2, const std::string& a3) { + return combine(is_max, combine(is_max, combine(is_max, a0, a1), a2), a3); + }; - // ── 归约视图归约 pass ── - for (std::size_t k = 0; k < n_inputs; ++k) + if (is_row) { - const ExprView& v = spec.views[k]; - if (!expr_view_is_reduce(static_cast(v.kind))) - continue; - const int slot = slot_of_view[k]; - const bool is_max = slot_is_max[static_cast(slot)]; - L << "\n // 归约视图 " << k << " (槽 " << slot << ")\n"; - L << " {\n"; - L << " float acc = " - << (is_max ? "uintBitsToFloat(0xff800000u)" : "0.0") << ";\n"; - L << loop_decl; - L << " acc = " << combine(is_max, "acc", - "b" + std::to_string(k) + "[row*cols + col]") << ";\n"; + // ═══ 行归约（工作组单行、跨线程跨步 + 树形归约；合并访问）═══ + const std::string loop_decl = + " for (uint c = tid; c < cols; c += 256u) {\n const uint row = idx;\n const uint col = c;\n"; + const std::string out_idx = "row*cols + col"; + + // ── 归约视图 pass（多累加器，跨步 4*256）── + for (std::size_t k = 0; k < n_inputs; ++k) + { + const ExprView& v = spec.views[k]; + if (!expr_view_is_reduce(static_cast(v.kind))) + continue; + const int slot = slot_of_view[k]; + const bool is_max = slot_is_max[static_cast(slot)]; + const std::string init = is_max ? "uintBitsToFloat(0xff800000u)" : "0.0"; + L << "\n // 归约视图 " << k << " (槽 " << slot << ")\n"; + L << " {\n"; + L << " float acc0 = " << init << ", acc1 = " << init + << ", acc2 = " << init << ", acc3 = " << init << ";\n"; + L << " uint i = tid;\n"; + L << " for (; i + 3u*256u < cols; i += 4u*256u) {\n"; + L << " acc0 = " << combine(is_max, "acc0", "b" + std::to_string(k) + "[idx*cols + i]") << ";\n"; + L << " acc1 = " << combine(is_max, "acc1", "b" + std::to_string(k) + "[idx*cols + i+256u]") << ";\n"; + L << " acc2 = " << combine(is_max, "acc2", "b" + std::to_string(k) + "[idx*cols + i+512u]") << ";\n"; + L << " acc3 = " << combine(is_max, "acc3", "b" + std::to_string(k) + "[idx*cols + i+768u]") << ";\n"; + L << " }\n"; + L << " for (; i < cols; i += 256u)\n"; + L << " acc0 = " << combine(is_max, "acc0", "b" + std::to_string(k) + "[idx*cols + i]") << ";\n"; + L << " s_red[" << slot << "][tid] = " + << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " }\n"; + emit_tree_reduce(slot, is_max); + } + + // ── 归约指令 pass（跨步循环 + 树形归约）── + for (std::size_t ri = 0; ri < spec.instrs.size(); ++ri) + { + const ExprInstr& R = spec.instrs[ri]; + if (!expr_op_is_reduce(static_cast(R.op))) + continue; + const int slot = slot_of_instr[R.dst]; + const bool is_max = slot_is_max[static_cast(slot)]; + const bool src_is_reg = + (R.a.kind == static_cast(ExprOperandKind::Reg) || + R.a.kind == static_cast(ExprOperandKind::Fanout)); + const std::string src = src_is_reg + ? "r" + std::to_string(static_cast(R.a.idx)) : operand(R.a); + L << "\n // 归约指令 " << ri << " (槽 " << slot << ")\n"; + L << " {\n"; + L << " float acc = " << (is_max ? "uintBitsToFloat(0xff800000u)" : "0.0") << ";\n"; + L << loop_decl; + emit_reg_decl(); + emit_instrs(0, ri); + L << " acc = " << combine(is_max, "acc", src) << ";\n"; + L << " }\n"; + L << " s_red[" << slot << "][tid] = acc;\n"; + L << " }\n"; + emit_tree_reduce(slot, is_max); + } + + // ── 输出 pass（行）── + L << "\n // 输出\n"; + L << " if (vector_out == 1u) {\n"; + L << " if (tid == 0u) {\n"; + if (last_is_reduce) + L << " bout[idx] = s_red[" << slot_of_instr[spec.instrs.back().dst] << "][0];\n"; + else + { + L << " const uint row = idx;\n const uint col = 0u;\n"; + emit_reg_decl(); + emit_instrs(0, spec.instrs.size()); + L << " bout[idx] = r" << static_cast(spec.instrs.back().dst) << ";\n"; + } L << " }\n"; - L << " s_red[" << slot << "][tid] = acc;\n"; + L << " return;\n"; L << " }\n"; - emit_tree_reduce(slot, is_max); - } - - // ── 归约指令 pass（按指令序，源可引用更早归约结果） ── - // 寄存器声明直接放在 for 循环体内（无需额外 {} 包裹），使源寄存器在 - // combine 处可见；不同 pass 各自循环体作用域隔离，无重声明冲突。 - for (std::size_t ri = 0; ri < spec.instrs.size(); ++ri) - { - const ExprInstr& R = spec.instrs[ri]; - const ExprOp rop = static_cast(R.op); - if (!expr_op_is_reduce(rop)) - continue; - const int slot = slot_of_instr[R.dst]; - const bool is_max = slot_is_max[static_cast(slot)]; - const bool src_is_reg = - (R.a.kind == static_cast(ExprOperandKind::Reg) || - R.a.kind == static_cast(ExprOperandKind::Fanout)); - const std::string src = src_is_reg - ? "r" + std::to_string(static_cast(R.a.idx)) - : operand(R.a); - L << "\n // 归约指令 " << ri << " (槽 " << slot << ")\n"; - L << " {\n"; - L << " float acc = " - << (is_max ? "uintBitsToFloat(0xff800000u)" : "0.0") << ";\n"; L << loop_decl; - emit_reg_decl(); - emit_instrs(0, ri); - L << " acc = " << combine(is_max, "acc", src) << ";\n"; - L << " }\n"; - L << " s_red[" << slot << "][tid] = acc;\n"; + if (last_is_reduce) + L << " bout[" << out_idx << "] = s_red[" << slot_of_instr[spec.instrs.back().dst] << "][0];\n"; + else + { + emit_reg_decl(); + emit_instrs(0, spec.instrs.size()); + L << " bout[" << out_idx << "] = r" << static_cast(spec.instrs.back().dst) << ";\n"; + } L << " }\n"; - emit_tree_reduce(slot, is_max); - } - - // ── 输出 pass ── - // vector_out=1：输出归约向量（thread 0 写代表元素，输出形状由后端分配 - // (rows,1)/(1,cols)）；vector_out=0：广播到 (rows,cols)。 - L << "\n // 输出\n"; - L << " if (vector_out == 1u) {\n"; - L << " const uint row = " << (is_row ? "idx" : "0u") << ";\n"; - L << " const uint col = " << (is_row ? "0u" : "idx") << ";\n"; - L << " if (tid == 0u) {\n"; - if (last_is_reduce) - { - L << " bout[" << (is_row ? "row" : "col") << "] = s_red[" - << slot_of_instr[spec.instrs.back().dst] << "][0];\n"; } else { - emit_reg_decl(); - emit_instrs(0, spec.instrs.size()); - L << " bout[" << (is_row ? "row" : "col") << "] = r" - << static_cast(spec.instrs.back().dst) << ";\n"; - } - L << " }\n"; - L << " return;\n"; - L << " }\n"; - L << loop_decl; - if (last_is_reduce) - { - L << " bout[" << out_idx << "] = s_red[" - << slot_of_instr[spec.instrs.back().dst] << "][0];\n"; - } - else - { - emit_reg_decl(); - emit_instrs(0, spec.instrs.size()); - L << " bout[" << out_idx << "] = r" - << static_cast(spec.instrs.back().dst) << ";\n"; + // ═══ 列归约 tile（合并访问）：每工作组 256 列、每线程一整列，无需跨线程归约 ═══ + L << " const uint col = idx * 256u + tid;\n"; + L << " if (col >= cols) return;\n"; + + // ── 归约视图 pass（tile：每线程一整列，多累加器跨行顺序读，合并访问）── + for (std::size_t k = 0; k < n_inputs; ++k) + { + const ExprView& v = spec.views[k]; + if (!expr_view_is_reduce(static_cast(v.kind))) + continue; + const int slot = slot_of_view[k]; + const bool is_max = slot_is_max[static_cast(slot)]; + const std::string init = is_max ? "uintBitsToFloat(0xff800000u)" : "0.0"; + L << "\n // 归约视图 " << k << " (槽 " << slot << ")\n"; + L << " {\n"; + L << " float acc0 = " << init << ", acc1 = " << init + << ", acc2 = " << init << ", acc3 = " << init << ";\n"; + L << " uint i = 0u;\n"; + L << " for (; i + 3u < rows; i += 4u) {\n"; + L << " acc0 = " << combine(is_max, "acc0", "b" + std::to_string(k) + "[i*cols + col]") << ";\n"; + L << " acc1 = " << combine(is_max, "acc1", "b" + std::to_string(k) + "[(i+1u)*cols + col]") << ";\n"; + L << " acc2 = " << combine(is_max, "acc2", "b" + std::to_string(k) + "[(i+2u)*cols + col]") << ";\n"; + L << " acc3 = " << combine(is_max, "acc3", "b" + std::to_string(k) + "[(i+3u)*cols + col]") << ";\n"; + L << " }\n"; + L << " for (; i < rows; ++i)\n"; + L << " acc0 = " << combine(is_max, "acc0", "b" + std::to_string(k) + "[i*cols + col]") << ";\n"; + L << " s_red[" << slot << "][tid] = " + << combine4(is_max, "acc0", "acc1", "acc2", "acc3") << ";\n"; + L << " }\n"; + } + + // ── 归约指令 pass（tile：单累加器，跨行顺序读，合并访问）── + for (std::size_t ri = 0; ri < spec.instrs.size(); ++ri) + { + const ExprInstr& R = spec.instrs[ri]; + if (!expr_op_is_reduce(static_cast(R.op))) + continue; + const int slot = slot_of_instr[R.dst]; + const bool is_max = slot_is_max[static_cast(slot)]; + const bool src_is_reg = + (R.a.kind == static_cast(ExprOperandKind::Reg) || + R.a.kind == static_cast(ExprOperandKind::Fanout)); + const std::string src = src_is_reg + ? "r" + std::to_string(static_cast(R.a.idx)) : operand(R.a); + L << "\n // 归约指令 " << ri << " (槽 " << slot << ")\n"; + L << " {\n"; + L << " float acc = " << (is_max ? "uintBitsToFloat(0xff800000u)" : "0.0") << ";\n"; + L << " for (uint row = 0u; row < rows; ++row) {\n"; + emit_reg_decl(); + emit_instrs(0, ri); + L << " acc = " << combine(is_max, "acc", src) << ";\n"; + L << " }\n"; + L << " s_red[" << slot << "][tid] = acc;\n"; + L << " }\n"; + } + + // ── 输出 pass（列 tile）── + L << "\n // 输出\n"; + L << " if (vector_out == 1u) {\n"; + if (last_is_reduce) + L << " bout[col] = s_red[" << slot_of_instr[spec.instrs.back().dst] << "][tid];\n"; + else + { + L << " const uint row = 0u;\n"; + emit_reg_decl(); + emit_instrs(0, spec.instrs.size()); + L << " bout[col] = r" << static_cast(spec.instrs.back().dst) << ";\n"; + } + L << " return;\n"; + L << " }\n"; + if (last_is_reduce) + { + L << " for (uint row = 0u; row < rows; ++row)\n"; + L << " bout[row*cols + col] = s_red[" << slot_of_instr[spec.instrs.back().dst] << "][tid];\n"; + } + else + { + L << " for (uint row = 0u; row < rows; ++row) {\n"; + emit_reg_decl(); + emit_instrs(0, spec.instrs.size()); + L << " bout[row*cols + col] = r" << static_cast(spec.instrs.back().dst) << ";\n"; + L << " }\n"; + } } - L << " }\n"; L << "}\n"; return L.str(); } +// ═══════════════════════════════════════════════════════════════════════════ +// GlslEmitter — GLSL 后端 emitter（IR-D） +// +// 把上述 generate_glsl / generate_glsl_reduce 封装为 ExprEmitter 接口实现， +// 使生成器工具（gen_fused）可经统一接口选择后端，不再与 GLSL 绑定。 +// 自由函数 generate_glsl / generate_glsl_reduce 保留为便捷入口 +// （内部转发到 GlslEmitter），向后兼容既有调用方。 +// ═══════════════════════════════════════════════════════════════════════════ +class GlslEmitter final : public ExprEmitter +{ +public: + [[nodiscard]] std::string_view name() const noexcept override + { return "glsl"; } + + [[nodiscard]] std::string generate( + const std::string& name_, const ExprSpec& spec) override + { return nn::generate_glsl(name_, spec); } + + [[nodiscard]] std::string generate_reduce( + const std::string& name_, const ExprSpec& spec) override + { return nn::generate_glsl_reduce(name_, spec); } +}; + +// 登记到 emitter 注册表（静态初始化；重复包含无害——同名拒绝覆盖） +inline const bool kGlslEmitterRegistered = + emitter_registry::register_backend("glsl", + []() -> ExprEmitter* { return new GlslEmitter(); }); + } // namespace nn #endif // NN_GLSL_GEN_HPP diff --git a/include/neuralnet.cpp/gpu_engine.hpp b/include/neuralnet.cpp/gpu_engine.hpp index c18d0480..589f6240 100644 --- a/include/neuralnet.cpp/gpu_engine.hpp +++ b/include/neuralnet.cpp/gpu_engine.hpp @@ -34,9 +34,11 @@ #include #include #include +#include #include "compute_engine.hpp" #include "expr_opt.hpp" +#include "expr_graph.hpp" // IR-C：图 IR + 融合分析 #if __has_include("fused_registry.hpp") #include "fused_registry.hpp" #endif @@ -53,6 +55,12 @@ class GpuEngine final : public ComputeEngine private: GpuBackend& backend_; + // IR-C 图 IR 录制（begin_expr/end_expr）： + // recording_ ：录制中的图（nullopt = 未录制） + // node_outputs_：节点下标 → 占位输出 buffer（end_expr 时写入真实结果） + std::optional recording_; + std::unordered_map node_outputs_; + public: explicit GpuEngine(GpuBackend& backend) : backend_(backend) {} @@ -162,11 +170,29 @@ class GpuEngine final : public ComputeEngine return backend_.flush_batch(); } - // ── 表达式录制（M2 框架）：当前为 no-op ───────────────────────────── - // 各原语照常独立 dispatch；录制融合分析（虚拟寄存器 DAG + 融合边界判定） - // 在 M3 落地。begin/end 保持幂等，Layer 可先行包住算法段落。 - [[nodiscard]] Result begin_expr() override { return {}; } - [[nodiscard]] Result end_expr() override { return {}; } + // ── 表达式录制（IR-C）：begin_expr/end_expr 图 IR 融合 ───────────── + // begin_expr 开启录制：期间 eval_expr / eval_expr_reduce 把表达式加入 + // 录制图并返回携带 virtual_tag 的占位 Tensor（真实 GPU buffer，Layer + // 无感知）。end_expr 做融合分析：逐元素链拼接成复合 ExprSpec → 每 + // kernel 一次 AOT dispatch，输出写入对应节点的占位 buffer（Layer 持有 + // 的 Tensor 即物化）。未命中 AOT 融合 shader → 硬报错（闭合世界）。 + [[nodiscard]] Result begin_expr() override + { + recording_.emplace(); + node_outputs_.clear(); + fused::recording_graph_ptr() = &*recording_; + return {}; + } + + [[nodiscard]] Result end_expr() override + { + if (!recording_) + return {}; + ExprGraph g = std::move(*recording_); + recording_.reset(); + fused::recording_graph_ptr() = nullptr; + return execute_fused_graph(g); + } // ══════════════════════════════════════════════════════════════════════ // 张量工厂（纯 GPU：全部创建/上传为 GPU Tensor） @@ -1021,6 +1047,22 @@ class GpuEngine final : public ComputeEngine std::span inputs, std::size_t rows, std::size_t cols) override { + // ── IR-C：录制模式（begin_expr/end_expr 内）────────────────── + // 把表达式加入录制图，返回携带 virtual_tag 的占位 GPU buffer； + // end_expr 时融合执行并物化。依赖识别：inputs 中带 virtual_tag 的 + // 占位 Tensor（前序节点输出）建立依赖边，不作为 kernel 外部输入。 + if (auto* g = fused::recording_graph()) + { + const int node = g->add_node(raw_spec, inputs, rows, cols, + /*vector_out=*/false); + auto out = GpuTensor::create_empty(rows, cols, backend_); + if (!out) return std::unexpected(out.error()); + Tensor t = Tensor::from_gpu(std::move(*out)); + t.set_virtual_tag(static_cast(node) + 1); + node_outputs_[node] = t; + return t; + } + // ── canonical IR：canonicalize 为引擎内部优化（IR-A/IR-B）， // key 与 shader 合成两端一致；dispatch 用 canonical 的 consts ── const ExprSpec spec = nn::canonicalize_expr_spec(raw_spec); @@ -1065,6 +1107,22 @@ class GpuEngine final : public ComputeEngine std::span inputs, std::size_t rows, std::size_t cols) override { + // ── IR-C：录制模式（归约输出：占位 buffer 按归约向量形状分配）── + if (auto* g = fused::recording_graph()) + { + const int node = g->add_node(raw_spec, inputs, rows, cols, + /*vector_out=*/true); + const int raxis = expr_spec_reduce_axis(raw_spec); + const std::size_t orows = (raxis == 0) ? rows : 1; + const std::size_t ocols = (raxis == 1) ? cols : 1; + auto out = GpuTensor::create_empty(orows, ocols, backend_); + if (!out) return std::unexpected(out.error()); + Tensor t = Tensor::from_gpu(std::move(*out)); + t.set_virtual_tag(static_cast(node) + 1); + node_outputs_[node] = t; + return t; + } + // canonical IR：与 eval_expr 同（canonicalize 为引擎内部优化） const ExprSpec spec = nn::canonicalize_expr_spec(raw_spec); const std::string key = nn::expr_spec_key(spec); @@ -1094,6 +1152,67 @@ class GpuEngine final : public ComputeEngine } private: + // ── IR-C：融合执行 ────────────────────────────────────────────────── + // 图 → 融合分析 → kernel 序列 → 逐个 AOT dispatch。每个 kernel 的输出 + // 写入其末尾节点（tail）的占位 buffer（node_outputs_[tail]），使 Layer + // 持有的占位 Tensor 在 end_expr 后物化。中间节点若被融合则无独立输出 + // （作为寄存器内联），若为融合边界则其占位 buffer 由自身 kernel 写入 + // 并被后续 kernel 当作输入绑定（FusedKernelInput.node 来源）。 + [[nodiscard]] Result execute_fused_graph(ExprGraph& g) + { + auto kernels = fuse_expr_graph(g); + for (auto& k : kernels) + { + // 解析输入：node 来源 → 该节点占位输出 buffer；external → 直接张量 + std::vector gpu_inputs; + gpu_inputs.reserve(k.inputs.size()); + for (auto& in : k.inputs) + { + Tensor t; + if (in.node >= 0) + { + const auto it = node_outputs_.find(in.node); + if (it == node_outputs_.end()) + return std::unexpected(Error{ + "GpuEngine::end_expr: 依赖节点输出缺失（图 IR 状态损坏）"}); + t = it->second; + } + else + { + t = in.external; + } + auto gg = ensure_gpu(t); + if (!gg) return std::unexpected(gg.error()); + gpu_inputs.push_back(gg->gpu_tensor()); + } + + // AOT 匹配（闭合世界：融合后的复合 spec 必须在构建期 scan 登记过） + const std::string key = nn::expr_spec_key(k.spec); +#ifdef NN_FUSED_REGISTRY_EMBEDDED + const nn::fused::FusedShader* fs = nn::fused::find_fused(key); + if (fs && backend_.has_fused_shader(key)) + { + const auto vp = nn::expr_spec_runtime_view_params(k.spec); + GpuTensor* out_override = nullptr; + const auto out_it = node_outputs_.find(k.tail); + if (out_it != node_outputs_.end()) + out_override = &out_it->second.gpu_tensor(); + auto out = backend_.run_fused_gpu( + key, gpu_inputs, k.spec.consts, k.rows, k.cols, + k.vector_out, vp, out_override); + if (!out) return std::unexpected(out.error()); + if (!out_override) + node_outputs_[k.tail] = Tensor::from_gpu(std::move(*out)); + continue; + } +#endif + return std::unexpected(Error{ + "GpuEngine::end_expr: 融合 kernel 未命中 AOT 融合 shader（闭合世界）；" + "请将 begin_expr/end_expr 段纳入构建期扫描（scan_exprs）"}); + } + return {}; + } + // ── 辅助：确保 Tensor 在 GPU 上 ────────────────────────────────────── // 若已是 GPU，返回共享拷贝（零开销）；若为 CPU，上传到 GPU。 // 纯 GPU 架构下，所有 Tensor 应已是 GPU，此方法为防御性兜底。 diff --git a/include/neuralnet.cpp/model_container.hpp b/include/neuralnet.cpp/model_container.hpp index a2fd5fa9..3d9b3694 100644 --- a/include/neuralnet.cpp/model_container.hpp +++ b/include/neuralnet.cpp/model_container.hpp @@ -58,22 +58,28 @@ class Model } // ── 构建网络（通用模板接口） ──────────────────────────────────────── - // 用法：model.add(engine, in, out) / model.add() / ... + // 用法：model.add(in, out) / model.add() / ... + // 注意：构造后自动调用 Layer::init()，无需手动调用。 // 注意：MultiHeadAttention 的构造函数是 protected（设计为基类， // 由 TransformerEncoderLayer / CausalSelfAttention 内部组合）， // 因此不直接通过 add 添加。 template - Model& add(Args&&... args) + Result add(Args&&... args) { - layers_.emplace_back(std::make_unique(std::forward(args)...)); - return *this; + auto layer = std::make_unique(std::forward(args)...); + auto r = layer->init(engine()); + if (!r) return std::unexpected(r.error()); + layers_.emplace_back(std::move(layer)); + return {}; } // 添加已由工厂构造的 Layer（如 make_norm_layer 按 NormType 创建归一化层） - Model& add_layer(std::unique_ptr layer) + Result add_layer(std::unique_ptr layer) { + auto r = layer->init(engine()); + if (!r) return std::unexpected(r.error()); layers_.emplace_back(std::move(layer)); - return *this; + return {}; } // ── 访问 ───────────────────────────────────────────────────────────── diff --git a/include/neuralnet.cpp/model_serialization.hpp b/include/neuralnet.cpp/model_serialization.hpp index 237f947a..2b92ff75 100644 --- a/include/neuralnet.cpp/model_serialization.hpp +++ b/include/neuralnet.cpp/model_serialization.hpp @@ -528,11 +528,14 @@ inline void apply_spec_version_defaults(KeyValueRecord &kv, uint32_t version) // 非可学习状态（如 BatchNorm 的 running_mean/running_var），紧跟在参数之后。 // 旧文件（无额外状态）读到 EOF 时保持默认（running_mean=0, running_var=1）。 - // TODO(1.1, M2): 上述"读到 EOF 保持默认"的语义与下方实现不符——read_matrix - // 失败时这里直接 return 错误而非回退默认值。需统一为按版本回退。 + // 注意：读取前用 peek() 检查 EOF，而非依赖 read_matrix 读失败后回退—— + // read_matrix 在 shape 不匹配时也会报错，不应吞掉真正的损坏。 auto extras = model.extra_state(); for (auto& e_tensor : extras) { + // 已经没有更多数据 → 保持默认值（零/一），跳过剩余 extras + if (ifs.peek() == EOF) break; + // 先读入临时 Matrix（按状态 Tensor 的形状） Matrix tmp(e_tensor.get().rows(), e_tensor.get().cols()); auto mr = detail::read_matrix(ifs, tmp); diff --git a/include/neuralnet.cpp/tensor.hpp b/include/neuralnet.cpp/tensor.hpp index c9bbe2d4..e9e3f077 100644 --- a/include/neuralnet.cpp/tensor.hpp +++ b/include/neuralnet.cpp/tensor.hpp @@ -48,6 +48,12 @@ class Tensor std::size_t rows_ = 0; std::size_t cols_ = 0; + // 图 IR 录制（IR-C）：虚拟寄存器标记。 + // begin_expr/end_expr 录制模式下，eval_expr 返回的占位 Tensor 携带 + // 该标记（= 图节点下标 + 1；0 = 非图节点输出），供后续表达式建立 + // 依赖边（Layer 无感知，见 expr_graph.hpp / gpu_engine.hpp）。 + std::uint64_t virtual_tag_ = 0; + // CPU 存储（device_ == CPU 时有效） std::shared_ptr cpu_data_; @@ -109,6 +115,10 @@ class Tensor [[nodiscard]] std::size_t size() const noexcept { return rows_ * cols_; } [[nodiscard]] bool is_cpu() const noexcept { return device_ == Device::CPU; } [[nodiscard]] bool is_gpu() const noexcept { return device_ == Device::GPU; } + + // ── 图 IR 虚拟寄存器标记（IR-C 录制用，Layer 无感知） ─────────────── + [[nodiscard]] std::uint64_t virtual_tag() const noexcept { return virtual_tag_; } + void set_virtual_tag(std::uint64_t tag) noexcept { virtual_tag_ = tag; } [[nodiscard]] bool valid() const noexcept { if (device_ == Device::CPU) @@ -190,6 +200,7 @@ class Tensor t.device_ = device_; t.rows_ = new_rows; t.cols_ = new_cols; + t.virtual_tag_ = virtual_tag_; // reshape 保留图 IR 标记 #ifdef NN_HAS_VULKAN t.gpu_data_ = gpu_data_; #endif diff --git a/report.html b/report.html index 6330fdab..48e84c0a 100644 --- a/report.html +++ b/report.html @@ -59,7 +59,7 @@

🌸 屎山代码分析报告 🌸

67.7
-
总体评分: 67.66 / 100
+
总体评分: 67.69 / 100
😷 屎气扑鼻 - 代码开始散发气味，谨慎维护
@@ -73,7 +73,7 @@

🌸 屎山代码分析报告 🌸

已跳过
-
903ms
+
935ms
耗时
@@ -85,17 +85,17 @@

📊 评分指标详情

指标摘要评分状态 - 循环复杂度20.6%✓ - 认知复杂度16.8%✅ + 循环复杂度20.7%✓ + 认知复杂度17.0%✅ 嵌套深度10.2%✅ 函数长度6.0%✅ 文件长度12.6%✅ 参数数量19.5%✅ 代码重复6.7%✅ - 结构分析13.2%✅ + 结构分析13.5%✅ 错误处理10.1%✅ - 注释比例28.1%✓ - 命名规范10.3%✅ + 注释比例28.3%✓ + 命名规范10.2%✅ @@ -127,7 +127,7 @@

📋 最屎代码排行榜

复杂度问题: 21其他问题: 7重复问题: 6结构问题: 12错误处理问题: 4注释问题: 1
-
  • forward() L651: 复杂度: 12
  • forward() L1325: 复杂度: 32
  • backward() L1477: 复杂度: 29
  • forward_step() L1631: 复杂度: 20
  • export_activations() L2717: 复杂度: 12
+
  • forward() L651: 复杂度: 12
  • forward() L1325: 复杂度: 32
  • backward() L1477: 复杂度: 29
  • forward_step() L1631: 复杂度: 20
  • export_activations() L2717: 复杂度: 13
@@ -167,7 +167,7 @@

📋 最屎代码排行榜

复杂度问题: 31其他问题: 17结构问题: 5命名问题: 10
-
  • initialize() L215: 复杂度: 11
  • initialize() L1021: 复杂度: 45
  • bmm_reduce_gpu() L1384: 复杂度: 14
  • bmm_denom_gpu() L1423: 复杂度: 14
  • bmm_apply_gpu() L1463: 复杂度: 14
+
  • initialize() L215: 复杂度: 11
  • initialize() L1042: 复杂度: 45
  • bmm_reduce_gpu() L1405: 复杂度: 14
  • bmm_denom_gpu() L1444: 复杂度: 14
  • bmm_apply_gpu() L1484: 复杂度: 14
diff --git a/src/attn_debug.cpp b/src/attn_debug.cpp deleted file mode 100644 index d40674b7..00000000 --- a/src/attn_debug.cpp +++ /dev/null @@ -1,317 +0,0 @@ -// ── 注意力反向逐级调试探针 ──────────────────────────────────────────────── -// -// 目的：attn_gradcheck 显示 wq/wk FAIL 但 wv/wo OK。本程序用「暴力 Matrix -// 参考」完整复现 CausalSelfAttention 的 forward + backward，并逐项对比 -// 引擎 batched_matmul / Softmax 原语计算出的中间量与参考，定位第一个 -// 出错的中间量（grad_S / grad_A / grad_Q / grad_K / batched_matmul 原语）。 -// -// 用法：attn_debug -// ───────────────────────────────────────────────────────────────────────── - -#include -#include - -#include -#include -#include -#include - -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; -using nn::ComputeEngine; -using nn::CausalSelfAttention; -using nn::Softmax; -using nn::PosEncodingType; - -namespace { - -// ── 简单矩阵代数辅助（行主序 (rows, cols)）─────────────────────────────── -Matrix mmul(const Matrix& A, const Matrix& B) -{ - Matrix C(A.rows(), B.cols()); - for (std::size_t i = 0; i < A.rows(); ++i) - for (std::size_t j = 0; j < B.cols(); ++j) - { - Scalar s{0}; - for (std::size_t k = 0; k < A.cols(); ++k) - s += A.at_unchecked(i, k) * B.at_unchecked(k, j); - C.set_value_unchecked(i, j, s); - } - return C; -} - -Matrix mtrans(const Matrix& A) -{ - Matrix T(A.cols(), A.rows()); - for (std::size_t i = 0; i < A.rows(); ++i) - for (std::size_t j = 0; j < A.cols(); ++j) - T.set_value_unchecked(j, i, A.at_unchecked(i, j)); - return T; -} - -Matrix madd(const Matrix& A, const Matrix& B) -{ - Matrix C = A; - for (std::size_t i = 0; i < C.size(); ++i) - C.span()[i] += B.span()[i]; - return C; -} - -// 行偏置广播加：bias (rows,1) 逐行加到 A (rows, cols) -Matrix madd_bias(const Matrix& A, const Matrix& B) -{ - Matrix C = A; - for (std::size_t i = 0; i < A.rows(); ++i) - { - const Scalar b = B.at_unchecked(i, 0); - for (std::size_t j = 0; j < A.cols(); ++j) - C.set_value_unchecked(i, j, C.at_unchecked(i, j) + b); - } - return C; -} - -Matrix mscale(const Matrix& A, Scalar s) -{ - Matrix C = A; - for (std::size_t i = 0; i < C.size(); ++i) - C.span()[i] *= s; - return C; -} - -Matrix rows(const Matrix& A, std::size_t r0, std::size_t n) -{ - Matrix R(n, A.cols()); - for (std::size_t i = 0; i < n; ++i) - for (std::size_t j = 0; j < A.cols(); ++j) - R.set_value_unchecked(i, j, A.at_unchecked(r0 + i, j)); - return R; -} - -// 逐行 softmax（含 -inf 处理） -Matrix softmax_rows(const Matrix& S) -{ - Matrix A(S.rows(), S.cols()); - for (std::size_t i = 0; i < S.rows(); ++i) - { - Scalar mx = -std::numeric_limits::infinity(); - for (std::size_t j = 0; j < S.cols(); ++j) - mx = std::max(mx, S.at_unchecked(i, j)); - Scalar sum{0}; - for (std::size_t j = 0; j < S.cols(); ++j) - { - Scalar e = std::exp(S.at_unchecked(i, j) - mx); - A.set_value_unchecked(i, j, e); - sum += e; - } - for (std::size_t j = 0; j < S.cols(); ++j) - A.set_value_unchecked(i, j, A.at_unchecked(i, j) / sum); - } - return A; -} - -// 逐行 softmax 反向: dS = A ⊙ (gA - row_dot) -Matrix softmax_bwd(const Matrix& A, const Matrix& gA) -{ - Matrix dS(A.rows(), A.cols()); - for (std::size_t i = 0; i < A.rows(); ++i) - { - Scalar dot{0}; - for (std::size_t j = 0; j < A.cols(); ++j) - dot += A.at_unchecked(i, j) * gA.at_unchecked(i, j); - for (std::size_t j = 0; j < A.cols(); ++j) - dS.set_value_unchecked(i, j, - A.at_unchecked(i, j) * (gA.at_unchecked(i, j) - dot)); - } - return dS; -} - -void cmp(const std::string& name, const Matrix& got, const Matrix& ref) -{ - bool ok = true; - Scalar max_err{0}; - std::size_t bad = 0; - for (std::size_t i = 0; i < got.rows(); ++i) - for (std::size_t j = 0; j < got.cols(); ++j) - { - const Scalar e = std::fabs(got.at_unchecked(i, j) - ref.at_unchecked(i, j)); - if (e > max_err) max_err = e; - if (e > Scalar{1e-4}) - { - ok = false; - if (bad < 6) - std::printf(" [%s] (%zu,%zu) got=%.6f ref=%.6f err=%.3e\n", - name.c_str(), i, j, - static_cast(got.at_unchecked(i, j)), - static_cast(ref.at_unchecked(i, j)), - static_cast(e)); - ++bad; - } - } - std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; -} - -} // namespace - -int main() -{ - nn::cli::EngineConfig ecfg; - auto engine_res = nn::cli::create_engine(ecfg, std::cout); - if (!engine_res) { std::cerr << "引擎创建失败: " << engine_res.error().message << "\n"; return 1; } - auto engine = std::move(*engine_res); - ComputeEngine& eng = *engine; - - const std::size_t d_model = 16; - const std::size_t H = 2; - const std::size_t dk = 8; - const std::size_t seq = 8; - const std::size_t batch = 1; - const Scalar scale = Scalar{1} / std::sqrt(static_cast(dk)); - - std::cout << "========================================\n"; - std::cout << " Attention 反向逐级调试 (d_model=16 H=2 seq=8 batch=1)\n"; - std::cout << "========================================\n"; - - CausalSelfAttention attn(eng, d_model, H, seq, seq, PosEncodingType::Learned); - - std::mt19937_64 rng(123); - std::uniform_real_distribution dist(-1, 1); - const std::size_t total = seq * batch; - Matrix x_m(d_model, total), go_m(d_model, total); - for (std::size_t i = 0; i < x_m.size(); ++i) x_m.span()[i] = dist(rng); - for (std::size_t i = 0; i < go_m.size(); ++i) go_m.span()[i] = dist(rng); - auto x = eng.from_matrix(x_m); - auto go = eng.from_matrix(go_m); - - // ── 提取权重 ── - auto params = attn.parameters(); - Matrix Wq = *eng.to_matrix(params[0]); // (16,16) - Matrix bq = *eng.to_matrix(params[1]); // (16,1) - Matrix Wk = *eng.to_matrix(params[2]); - Matrix bk = *eng.to_matrix(params[3]); - Matrix Wv = *eng.to_matrix(params[4]); - Matrix bv = *eng.to_matrix(params[5]); - Matrix Wo = *eng.to_matrix(params[6]); - Matrix bo = *eng.to_matrix(params[7]); - - // ── 参考 forward ── - Matrix Q = madd_bias(mmul(Wq, x_m), bq); - Matrix K = madd_bias(mmul(Wk, x_m), bk); - Matrix V = madd_bias(mmul(Wv, x_m), bv); - std::vector Ah(H, Matrix(seq, seq)); - for (std::size_t h = 0; h < H; ++h) - { - Matrix Qh = rows(Q, h * dk, dk); - Matrix Kh = rows(K, h * dk, dk); - Matrix Sh = mscale(mmul(mtrans(Qh), Kh), scale); - // 因果掩码：j > i → -inf - for (std::size_t i = 0; i < seq; ++i) - for (std::size_t j = i + 1; j < seq; ++j) - Sh.set_value_unchecked(i, j, -std::numeric_limits::infinity()); - Ah[h] = softmax_rows(Sh); - if (h == 0) - { - std::cout << " [ref] head0 Sh (masked) col0..7:\n"; - for (std::size_t i = 0; i < seq; ++i) - { - std::cout << " row" << i << ":"; - for (std::size_t j = 0; j < seq; ++j) - std::printf(" %9.4f", static_cast(Sh.at_unchecked(i, j))); - std::cout << "\n"; - } - std::cout << " [ref] head0 A:\n"; - for (std::size_t i = 0; i < seq; ++i) - { - std::cout << " row" << i << ":"; - for (std::size_t j = 0; j < seq; ++j) - std::printf(" %9.6f", static_cast(Ah[h].at_unchecked(i, j))); - std::cout << "\n"; - } - } - } - Matrix concat(d_model, seq); - for (std::size_t h = 0; h < H; ++h) - { - Matrix Vh = rows(V, h * dk, dk); - Matrix Oh = mmul(Vh, mtrans(Ah[h])); // (dk, seq) - for (std::size_t i = 0; i < dk; ++i) - for (std::size_t j = 0; j < seq; ++j) - concat.set_value_unchecked(h * dk + i, j, Oh.at_unchecked(i, j)); - } - Matrix y_ref = madd_bias(mmul(Wo, concat), bo); - - auto y_fwd = attn.forward(eng, *x); - Matrix y_got = *eng.to_matrix(*y_fwd); - cmp("forward out (wo concat)", y_got, y_ref); - std::cout << " (forward out 一致性验证 wv/wo 依赖的中间量)\n\n"; - - // ── 参考 backward ── - Matrix gc = mmul(mtrans(Wo), go_m); // (16, seq) - std::vector gQh(H, Matrix(dk, seq)), gKh(H, Matrix(dk, seq)), gVh(H, Matrix(dk, seq)); - for (std::size_t h = 0; h < H; ++h) - { - Matrix Vh = rows(V, h * dk, dk); - Matrix Kh = rows(K, h * dk, dk); - Matrix Qh = rows(Q, h * dk, dk); - Matrix gch = rows(gc, h * dk, dk); - Matrix gV = mmul(gch, Ah[h]); // (dk, seq) - Matrix gA = mmul(mtrans(Vh), gch); // (seq, seq) - Matrix gS = softmax_bwd(Ah[h], gA); // (seq, seq) - gQh[h] = mscale(mmul(Kh, mtrans(gS)), scale); - gKh[h] = mscale(mmul(Qh, gS), scale); - gVh[h] = gV; - } - Matrix gQ(d_model, seq), gK(d_model, seq), gV(d_model, seq); - for (std::size_t h = 0; h < H; ++h) - for (std::size_t i = 0; i < dk; ++i) - for (std::size_t j = 0; j < seq; ++j) - { - gQ.set_value_unchecked(h * dk + i, j, gQh[h].at_unchecked(i, j)); - gK.set_value_unchecked(h * dk + i, j, gKh[h].at_unchecked(i, j)); - gV.set_value_unchecked(h * dk + i, j, gVh[h].at_unchecked(i, j)); - } - Matrix ref_gwq = mmul(gQ, mtrans(x_m)); - Matrix ref_gwk = mmul(gK, mtrans(x_m)); - Matrix ref_gwv = mmul(gV, mtrans(x_m)); - Matrix ref_gwo = mmul(gc, mtrans(concat)); - Matrix ref_gbq(d_model, 1), ref_gbk(d_model, 1), ref_gbv(d_model, 1), ref_gbo(d_model, 1); - for (std::size_t i = 0; i < d_model; ++i) - { - Scalar sq{0}, sk{0}, sv{0}, so{0}; - for (std::size_t j = 0; j < seq; ++j) - { - sq += gQ.at_unchecked(i, j); - sk += gK.at_unchecked(i, j); - sv += gV.at_unchecked(i, j); - so += gc.at_unchecked(i, j); - } - ref_gbq.set_value_unchecked(i, 0, sq); - ref_gbk.set_value_unchecked(i, 0, sk); - ref_gbv.set_value_unchecked(i, 0, sv); - ref_gbo.set_value_unchecked(i, 0, so); - } - Matrix ref_gx = madd(madd(mmul(mtrans(Wq), gQ), mmul(mtrans(Wk), gK)), mmul(mtrans(Wv), gV)); - - // ── 引擎 backward ── - auto grads = attn.param_gradients(); - for (auto& g : grads) { auto rz = eng.zero(g.get()); if (!rz) return 1; } - auto gx = attn.backward(eng, *go); - if (!gx) { std::cerr << "backward failed\n"; return 1; } - - Matrix gwq = *eng.to_matrix(grads[0]); - Matrix gbk = *eng.to_matrix(grads[3]); - Matrix gwv = *eng.to_matrix(grads[4]); - Matrix gwo = *eng.to_matrix(grads[6]); - Matrix gx_got = *eng.to_matrix(*gx); - - std::cout << "── 引擎 vs 暴力参考（仅比较重量级梯度） ──\n"; - cmp("grad_wq", gwq, ref_gwq); - cmp("grad_wk", *eng.to_matrix(grads[2]), ref_gwk); - cmp("grad_wv", gwv, ref_gwv); - cmp("grad_wo", gwo, ref_gwo); - cmp("grad_bk", gbk, ref_gbk); - cmp("grad_x", gx_got, ref_gx); - - return 0; -} diff --git a/src/expr_fuse_test.cpp b/src/expr_fuse_test.cpp new file mode 100644 index 00000000..7c1bb89a --- /dev/null +++ b/src/expr_fuse_test.cpp @@ -0,0 +1,254 @@ +// ─────────────────────────────────────────────────────────────────────────── +// expr_fuse_test.cpp — IR-C 图 IR 融合 GPU 端到端验证（begin_expr/end_expr） +// +// 验证内容（对应 docs/11-ir-optimization.md IR-C）： +// 1. FusedChainLayer forward：GPU（begin_expr 录制 → end_expr 融合 → +// AOT 融合 shader 单 kernel dispatch）vs CPU（逐节点求值）一致。 +// —— 中间量 t/u 内联为寄存器，不落显存、不额外 dispatch。 +// 2. FusedChainLayer backward：GPU（独立 AOT）vs CPU 一致。 +// 3. 录制图依赖识别：GPU 端第二个表达式消费第一个的占位输出 → 融合。 +// 4. 闭合世界：未扫描的融合 kernel → end_expr 硬报错（不静默回退）。 +// 5. 未录制路径回归：普通 dsl::compute 在 GPU 上照常 AOT dispatch。 +// +// 用法：expr_fuse_test（需 Vulkan SDK；依赖 CMake 构建期 scan_exprs 已 +// dry-run FusedChainLayer，登记融合后的复合 spec —— 见 scan_exprs.cpp） +// ─────────────────────────────────────────────────────────────────────────── + +#include + +#include +#include +#include +#include +#include +#include + +using nn::Scalar; + +#ifndef NN_HAS_VULKAN +int main() +{ + std::cout << "[SKIP] 此程序需要 Vulkan SDK 支持，请使用 -DNN_HAS_VULKAN 编译。\n"; + return 0; +} +#else +#include + +using nn::Matrix; +using nn::Tensor; +using nn::CpuEngine; +using nn::GpuEngine; +using nn::GpuBackend; + +// ── 辅助：最大绝对误差 ──────────────────────────────────────────────────── +Scalar max_abs_diff(const Matrix& a, const Matrix& b) +{ + Scalar max_err = 0.0f; + auto as = a.span(), bs = b.span(); + for (std::size_t i = 0; i < as.size(); ++i) + { + const Scalar d = std::fabs(as[i] - bs[i]); + if (d > max_err) max_err = d; + } + return max_err; +} + +// ── FusedChainLayer：GPU（begin_expr/end_expr 融合）vs CPU ────────────── +// forward: t = x*2; u = t+3; out = u*gamma → GPU 端三表达式融合成单 kernel +int run_fused_chain(CpuEngine& cpu, GpuEngine& gpu) +{ + const std::size_t F = 8, B = 6; + std::mt19937 rng(2026); + std::uniform_real_distribution dist(-1.0f, 1.0f); + Matrix x(F, B), grad(F, B); + for (auto& v : x.span()) v = dist(rng); + for (auto& v : grad.span()) v = dist(rng); + + nn::FusedChainLayer c_cpu(cpu, F), c_gpu(gpu, F); + const Tensor in = Tensor::from_matrix(Matrix(x)); + + // forward：GPU 走 begin_expr/end_expr 融合，CPU 逐节点 + auto fc = c_cpu.forward(cpu, in); + auto fg = c_gpu.forward(gpu, in); + if (!fc) { std::cerr << " CPU FusedChain forward 失败: " << fc.error().message << "\n"; return 1; } + if (!fg) + { + std::cerr << " GPU FusedChain forward 失败（end_expr 融合 kernel 未命中 AOT shader？）: " + << fg.error().message << "\n"; + return 1; + } + auto fgm = gpu.to_matrix(*fg); + if (!fgm) { std::cerr << " GPU FusedChain forward 结果下载失败\n"; return 1; } + const Scalar err_f = max_abs_diff(fc->cpu_matrix(), *fgm); + const bool ok_f = err_f < 1e-4f; + std::cout << "[" << (ok_f ? "PASS" : "FAIL") << "] FusedChain forward（begin_expr/end_expr 融合）" + << " err=" << std::scientific << std::setprecision(2) << err_f << "\n"; + + // backward：GPU 独立 AOT + const Tensor gd = Tensor::from_matrix(Matrix(grad)); + auto bc = c_cpu.backward(cpu, gd); + auto bg = c_gpu.backward(gpu, gd); + if (!bc) { std::cerr << " CPU FusedChain backward 失败: " << bc.error().message << "\n"; return 1; } + if (!bg) + { + std::cerr << " GPU FusedChain backward 失败: " << bg.error().message << "\n"; + return 1; + } + auto bgm = gpu.to_matrix(*bg); + if (!bgm) { std::cerr << " GPU FusedChain backward 结果下载失败\n"; return 1; } + const Scalar err_b = max_abs_diff(bc->cpu_matrix(), *bgm); + const bool ok_b = err_b < 1e-4f; + std::cout << "[" << (ok_b ? "PASS" : "FAIL") << "] FusedChain backward" + << " err=" << std::scientific << std::setprecision(2) << err_b << "\n"; + return (ok_f && ok_b) ? 0 : 1; +} + +// ── 手写录制：显式 begin_expr/end_expr 链式融合 ──────────────────────── +// 与 FusedChainLayer 等价的手写路径：x*2 → t+3 → u*gamma， +// 验证 GpuEngine 录制图依赖识别 + 融合执行（结果应为 ((x*2)+3)*gamma）。 +int run_manual_recording(CpuEngine& cpu, GpuEngine& gpu) +{ + const std::size_t R = 5, C = 7; + std::mt19937 rng(999); + std::uniform_real_distribution dist(-1.0f, 1.0f); + Matrix x(R, C); + for (auto& v : x.span()) v = dist(rng); + Matrix g(R, 1); + for (auto& v : g.span()) v = 0.5f + 0.01f * static_cast(rng() % 10); + + const Tensor xt = Tensor::from_matrix(Matrix(x)); + const Tensor gt = Tensor::from_matrix(Matrix(g)); + + // CPU 参考：逐节点 + auto c1 = nn::dsl::compute(cpu, nn::dsl::leaf(xt) * Scalar{2}, R, C); + auto c2 = nn::dsl::compute(cpu, nn::dsl::leaf(*c1) + Scalar{3}, R, C); + auto c3 = nn::dsl::compute(cpu, nn::dsl::leaf(*c2) * nn::dsl::row_broadcast(gt), R, C); + if (!c1 || !c2 || !c3) { std::cerr << " CPU 参考求值失败\n"; return 1; } + + // GPU：begin_expr 录制 → 三表达式 → end_expr 融合执行 + auto b = gpu.begin_expr(); + if (!b) { std::cerr << " GPU begin_expr 失败: " << b.error().message << "\n"; return 1; } + auto g1 = nn::dsl::compute(gpu, nn::dsl::leaf(xt) * Scalar{2}, R, C); + if (!g1) { std::cerr << " GPU 录制 1 失败: " << g1.error().message << "\n"; return 1; } + auto g2 = nn::dsl::compute(gpu, nn::dsl::leaf(*g1) + Scalar{3}, R, C); + if (!g2) { std::cerr << " GPU 录制 2 失败: " << g2.error().message << "\n"; return 1; } + auto g3 = nn::dsl::compute(gpu, nn::dsl::leaf(*g2) * nn::dsl::row_broadcast(gt), R, C); + if (!g3) { std::cerr << " GPU 录制 3 失败: " << g3.error().message << "\n"; return 1; } + auto e = gpu.end_expr(); + if (!e) { std::cerr << " GPU end_expr 失败: " << e.error().message << "\n"; return 1; } + + // 物化：g3 的占位 buffer 已被融合 kernel 写入 → 直接下载 + auto gm = gpu.to_matrix(*g3); + if (!gm) { std::cerr << " GPU 融合结果下载失败\n"; return 1; } + const Scalar err = max_abs_diff(c3->cpu_matrix(), *gm); + const bool ok = err < 1e-4f; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] 手写 begin_expr/end_expr 链式融合" + << " err=" << std::scientific << std::setprecision(2) << err << "\n"; + return ok ? 0 : 1; +} + +// ── 未录制路径回归：普通 dsl::compute 照常 AOT dispatch ──────────────── +// 用 FusedChainLayer backward 的表达式（scan 已登记）：grad * gamma * 2 +// 注：录制段内"独立分支各自成 kernel"的融合边界已由 expr_graph_test +// （纯函数，test_independent/test_multi_consumer）覆盖；本测试聚焦 GPU 端 +// 闭合世界内的回归路径。 +int run_non_recording(CpuEngine& cpu, GpuEngine& gpu) +{ + const std::size_t R = 4, C = 6; + std::mt19937 rng(7); + std::uniform_real_distribution dist(-1.0f, 1.0f); + Matrix x(R, C); + for (auto& v : x.span()) v = dist(rng); + const Tensor xt = Tensor::from_matrix(Matrix(x)); + + // 用 FusedChainLayer backward 的表达式（scan 已登记）：grad * gamma * 2 + Matrix g(R, 1, Scalar{0.5f}); + const Tensor gt = Tensor::from_matrix(Matrix(g)); + auto cr = nn::dsl::compute(cpu, + nn::dsl::leaf(xt) * nn::dsl::row_broadcast(gt) * Scalar{2}, R, C); + auto gr = nn::dsl::compute(gpu, + nn::dsl::leaf(xt) * nn::dsl::row_broadcast(gt) * Scalar{2}, R, C); + if (!cr) { std::cerr << " CPU 求值失败\n"; return 1; } + if (!gr) + { + std::cerr << " GPU 未录制路径失败（未命中 AOT shader？）: " << gr.error().message << "\n"; + return 1; + } + auto gm = gpu.to_matrix(*gr); + if (!gm) return 1; + const Scalar err = max_abs_diff(cr->cpu_matrix(), *gm); + const bool ok = err < 1e-4f; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] 未录制路径回归（普通 AOT dispatch）" + << " err=" << std::scientific << std::setprecision(2) << err << "\n"; + return ok ? 0 : 1; +} + +// ── vec4 向量化路径验证（cols%4==0）────────────────────────────────────── +// 用已登记的 vec4 合格表达式（Linear × RowBroadcast × const，即 FusedChain +// backward 的 grad*gamma*2）在 cols 为 4 的倍数时运行，确保 glsl_gen 的 vec4 +// 快速路径（vec4 加载/广播/运算/存储）数值与 CPU 参考一致。尾部与非 4 倍数 +// cols 的标量回退由 run_non_recording（C=6）覆盖。 +int run_vec4_path(CpuEngine& cpu, GpuEngine& gpu) +{ + const std::size_t R = 4; + int fail = 0; + for (const std::size_t C : {std::size_t{4}, std::size_t{8}, std::size_t{12}}) + { + std::mt19937 rng(42 + static_cast(C)); + std::uniform_real_distribution dist(-1.0f, 1.0f); + Matrix x(R, C); + for (auto& v : x.span()) v = dist(rng); + Matrix g(R, 1); + for (auto& v : g.span()) v = 0.5f + 0.01f * static_cast(rng() % 10); + + const Tensor xt = Tensor::from_matrix(Matrix(x)); + const Tensor gt = Tensor::from_matrix(Matrix(g)); + auto cr = nn::dsl::compute(cpu, + nn::dsl::leaf(xt) * nn::dsl::row_broadcast(gt) * Scalar{2}, R, C); + auto gr = nn::dsl::compute(gpu, + nn::dsl::leaf(xt) * nn::dsl::row_broadcast(gt) * Scalar{2}, R, C); + if (!cr) { std::cerr << " CPU 求值失败\n"; return 1; } + if (!gr) + { + std::cerr << " GPU vec4 路径失败（未命中 AOT shader？）: " + << gr.error().message << "\n"; + return 1; + } + auto gm = gpu.to_matrix(*gr); + if (!gm) { std::cerr << " GPU 结果下载失败\n"; return 1; } + const Scalar err = max_abs_diff(cr->cpu_matrix(), *gm); + const bool ok = err < 1e-4f; + std::cout << "[" << (ok ? "PASS" : "FAIL") << "] vec4 向量化路径 (cols=" << C << ")" + << " err=" << std::scientific << std::setprecision(2) << err << "\n"; + fail += ok ? 0 : 1; + } + return fail == 0 ? 0 : 1; +} + +int main() +{ + std::cout << "========================================\n" + << " IR-C 图 IR 融合 GPU 端到端验证\n" + << "========================================\n"; + + auto cpu_engine = std::make_unique(); + auto& backend = GpuBackend::instance(); + auto init_r = backend.initialize(); + if (!init_r) + { + std::cerr << "GPU 初始化失败: " << init_r.error().message << "\n"; + return 1; + } + auto gpu_engine = std::make_unique(backend); + std::cout << "[init] CpuEngine + GpuEngine 就绪\n"; + + int fail = 0; + fail += run_fused_chain(*cpu_engine, *gpu_engine); + fail += run_manual_recording(*cpu_engine, *gpu_engine); + fail += run_non_recording(*cpu_engine, *gpu_engine); + fail += run_vec4_path(*cpu_engine, *gpu_engine); + + std::cout << (fail == 0 ? "ALL PASS\n" : "FAILED\n"); + return fail == 0 ? 0 : 1; +} +#endif diff --git a/src/expr_graph_test.cpp b/src/expr_graph_test.cpp new file mode 100644 index 00000000..143e30cc --- /dev/null +++ b/src/expr_graph_test.cpp @@ -0,0 +1,541 @@ +// ─────────────────────────────────────────────────────────────────────────── +// expr_graph_test.cpp — 图 IR + 融合分析（IR-C）单元验证（纯 CPU、无 GPU） +// +// 验证内容（对应 docs/11-ir-optimization.md 的验证策略）： +// 1. 链融合正确性：A→B 融合成单 kernel 后求值 == 逐节点求值 +// 2. 多节点链融合（3 节点） +// 3. 融合边界：非链（无消费者 / 多消费者 / 非 Linear 引用）不融合 +// 4. 归约节点作为融合边界（不参与链式拼接） +// 5. 上限压力：拼接后超 EXPR_MAX_* 则保守放弃融合 +// 6. 确定性：同一图多次融合结果完全一致（含 key） +// 7. 不变量：融合 kernel 通过 validate；views/输入顺序确定 +// 8. 独立分支：不消费前序输出的相邻表达式各自独立成 kernel +// +// 编译（单 TU 快速验证，与 expr_opt_test 相同）： +// clang++ -std=c++26 -stdlib=libc++ -fno-exceptions \ +// -I include/neuralnet.cpp src/expr_graph_test.cpp -o build/expr_graph_test +// 运行：./build/expr_graph_test +// ─────────────────────────────────────────────────────────────────────────── + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include // 注册 GlslEmitter（IR-D） +#include // 注册 CpuEmitter（IR-D） +#include + +using nn::Scalar; + +namespace +{ + +int g_fail = 0; +#define CHECK(cond, msg) \ + do { \ + if (!(cond)) { \ + std::printf("[FAIL] %s\n", msg); \ + ++g_fail; \ + } \ + } while (0) + +nn::Tensor make_tensor(std::size_t rows, std::size_t cols, float base = 0.0f, + float step = 0.01f) +{ + nn::Tensor t = nn::Tensor::cpu(rows, cols); + auto sp = t.cpu_matrix().span(); + for (std::size_t i = 0; i < sp.size(); ++i) + sp[i] = base + static_cast(i) * step; + return t; +} + +Scalar max_abs_diff(const nn::Matrix& a, const nn::Matrix& b) +{ + Scalar m = 0.0f; + auto as = a.span(), bs = b.span(); + for (std::size_t i = 0; i < as.size(); ++i) + { + const Scalar d = std::fabs(as[i] - bs[i]); + if (d > m) m = d; + } + return m; +} + +// 构造逐元素一元指令：out = unary(input0) 或 out = input0 const +nn::ExprSpec make_unary_const(nn::ExprOp op, Scalar c) +{ + nn::ExprSpec s; + s.views.push_back(nn::expr::linear()); + s.consts.push_back(c); + nn::ExprInstr in; + in.op = static_cast(op); + in.dst = 0; + in.a = nn::expr::input(0); + in.b = nn::expr::cst(0); + s.instrs.push_back(in); + s.num_regs = 1; + return s; +} + +// 构造逐元素二元：out = input0 input1 +nn::ExprSpec make_binary(nn::ExprOp op) +{ + nn::ExprSpec s; + s.views.push_back(nn::expr::linear()); + s.views.push_back(nn::expr::linear()); + nn::ExprInstr in; + in.op = static_cast(op); + in.dst = 0; + in.a = nn::expr::input(0); + in.b = nn::expr::input(1); + s.instrs.push_back(in); + s.num_regs = 1; + return s; +} + +// 构造归约表达式：out = row_sum(input0)（归约边界） +nn::ExprSpec make_row_sum() +{ + nn::ExprSpec s; + s.views.push_back(nn::expr::linear()); + nn::ExprInstr in; + in.op = static_cast(nn::ExprOp::RowSum); + in.dst = 0; + in.a = nn::expr::input(0); + s.instrs.push_back(in); + s.num_regs = 1; + return s; +} + +// ── 测试 1：A→B 链融合正确性 ─────────────────────────────────────────── +// A: y = x * 2；B: z = y + 1 → 融合后 z = x*2 + 1，单 kernel +void test_chain_fuse() +{ + nn::CpuEngine eng; + const std::size_t rows = 3, cols = 4; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.5f); + + const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols, false); + nn::Tensor ytag = nn::Tensor::cpu(rows, cols); + ytag.set_virtual_tag(static_cast(nA) + 1); + const int nB = g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{ytag}, rows, cols, false); + (void)nB; + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 1, "A→B 链应融合为单个 kernel"); + if (kernels.size() != 1) return; + CHECK(kernels[0].members.size() == 2, "kernel 应包含 2 个节点"); + CHECK(kernels[0].tail == nB, "kernel tail 应为 B 节点"); + CHECK(kernels[0].inputs.size() == kernels[0].spec.views.size(), + "kernel inputs 与 views 一一对应"); + + // 融合 kernel 求值 + std::vector kinputs; + for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + auto fused = eng.eval_expr(kernels[0].spec, kinputs, rows, cols); + CHECK(static_cast(fused), "融合 kernel 求值成功"); + + // 参考：逐节点 + auto a_res = eng.eval_expr(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols); + auto b_res = eng.eval_expr(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{*a_res}, rows, cols); + CHECK(static_cast(a_res) && static_cast(b_res), "参考求值成功"); + if (!fused || !a_res || !b_res) return; + + const Scalar err = max_abs_diff(fused->cpu_matrix(), b_res->cpu_matrix()); + CHECK(err < 1e-6f, "融合结果与逐节点一致"); + + // 手工验证数值：x*2+1 + bool ok = true; + for (std::size_t i = 0; i < rows * cols; ++i) + { + const Scalar expect = x.cpu_matrix().span()[i] * 2.0f + 1.0f; + if (std::fabs(fused->cpu_matrix().span()[i] - expect) > 1e-5f) ok = false; + } + CHECK(ok, "融合结果与手工参考 x*2+1 一致"); + std::printf("[OK] test_chain_fuse\n"); +} + +// ── 测试 2：三节点链融合 ─────────────────────────────────────────────── +// A: y=x*2；B: z=y+3；C: w=z-1 → 融合成单 kernel +void test_chain3_fuse() +{ + nn::CpuEngine eng; + const std::size_t rows = 2, cols = 5; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 1.0f, 0.3f); + + auto add_node_chain = [&](nn::ExprOp op, Scalar c, + const nn::Tensor& input, + std::size_t r, std::size_t col) -> int + { + const int n = g.add_node(make_unary_const(op, c), + std::vector{input}, r, col, false); + return n; + }; + const int nA = add_node_chain(nn::ExprOp::Mul, Scalar{2}, x, rows, cols); + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + const int nB = add_node_chain(nn::ExprOp::Add, Scalar{3}, yt, rows, cols); + nn::Tensor zt = nn::Tensor::cpu(rows, cols); + zt.set_virtual_tag(static_cast(nB) + 1); + const int nC = add_node_chain(nn::ExprOp::Sub, Scalar{1}, zt, rows, cols); + (void)nC; + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 1, "三节点链应融合为单个 kernel"); + if (kernels.size() != 1) return; + CHECK(kernels[0].members.size() == 3, "kernel 应包含 3 个节点"); + + std::vector kinputs; + for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + auto fused = eng.eval_expr(kernels[0].spec, kinputs, rows, cols); + CHECK(static_cast(fused), "融合 kernel 求值成功"); + + // 参考：逐节点 + auto a_res = eng.eval_expr(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols); + auto b_res = eng.eval_expr(make_unary_const(nn::ExprOp::Add, Scalar{3}), + std::vector{*a_res}, rows, cols); + auto c_res = eng.eval_expr(make_unary_const(nn::ExprOp::Sub, Scalar{1}), + std::vector{*b_res}, rows, cols); + CHECK(static_cast(c_res), "参考求值成功"); + if (!fused || !c_res) return; + + const Scalar err = max_abs_diff(fused->cpu_matrix(), c_res->cpu_matrix()); + CHECK(err < 1e-6f, "三节点融合结果与逐节点一致"); + std::printf("[OK] test_chain3_fuse\n"); +} + +// ── 测试 3：独立分支各自成 kernel ────────────────────────────────────── +// A: y = x*2（无消费者）；B: z = w+1（外部输入 w，不依赖 A）→ 2 个 kernel +void test_independent() +{ + nn::CpuEngine eng; + const std::size_t rows = 3, cols = 3; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + nn::Tensor w = make_tensor(rows, cols, 2.0f, 0.2f); + + (void)g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols, false); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{w}, rows, cols, false); + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 2, "独立节点应各自成 kernel"); + if (kernels.size() != 2) return; + CHECK(kernels[0].members.size() == 1 && kernels[1].members.size() == 1, + "每个 kernel 只含 1 个节点"); + + // 两个 kernel 求值都正确 + std::vector k0in; + for (auto& in : kernels[0].inputs) k0in.push_back(in.external); + std::vector k1in; + for (auto& in : kernels[1].inputs) k1in.push_back(in.external); + auto r0 = eng.eval_expr(kernels[0].spec, k0in, rows, cols); + auto r1 = eng.eval_expr(kernels[1].spec, k1in, rows, cols); + CHECK(static_cast(r0) && static_cast(r1), "两个 kernel 求值成功"); + std::printf("[OK] test_independent\n"); +} + +// ── 测试 4：多消费者不融合（tail 被 A 和 B 之外节点消费） ───────────── +void test_multi_consumer() +{ + nn::CpuEngine eng; + const std::size_t rows = 2, cols = 4; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + + // A: y = x*2（输出 y）；B: z = y+1（消费 y）；C: t = y-1（也消费 y） + const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols, false); + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{yt}, rows, cols, false); + nn::Tensor yt2 = nn::Tensor::cpu(rows, cols); + yt2.set_virtual_tag(static_cast(nA) + 1); + (void)g.add_node(make_unary_const(nn::ExprOp::Sub, Scalar{1}), + std::vector{yt2}, rows, cols, false); + + auto kernels = fuse_expr_graph(g); + // A 有两个消费者（B、C），不能融合 → 3 个独立 kernel + CHECK(kernels.size() == 3, "多消费者应各成 kernel（不融合）"); + std::printf("[OK] test_multi_consumer (kernels=%zu)\n", kernels.size()); +} + +// ── 测试 5：归约节点作为融合边界 ─────────────────────────────────────── +void test_reduce_boundary() +{ + nn::CpuEngine eng; + const std::size_t rows = 3, cols = 4; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + + // A: y = row_sum(x)（归约，输出 (rows,1)）；B: z = y + 1（消费 y，但 B 的 + // 输入 y 形状是 (rows,1)，经 RowBroadcast 视图）→ 不融合（A 是归约） + const int nA = g.add_node(make_row_sum(), + std::vector{x}, rows, cols, true); + nn::Tensor yt = nn::Tensor::cpu(rows, 1); + yt.set_virtual_tag(static_cast(nA) + 1); + nn::ExprSpec sb; + sb.views.push_back(nn::expr::row_broadcast()); // (rows,1) 按行广播 + nn::ExprInstr in; + in.op = static_cast(nn::ExprOp::Add); + in.dst = 0; + in.a = nn::expr::input(0); + in.b = nn::expr::cst(0); + sb.consts.push_back(Scalar{1}); + sb.instrs.push_back(in); + sb.num_regs = 1; + (void)g.add_node(sb, std::vector{yt}, rows, cols, false); + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 2, "归约节点应作为融合边界（不拼接）"); + std::printf("[OK] test_reduce_boundary (kernels=%zu)\n", kernels.size()); +} + +// ── 测试 6：上限压力（拼接后超 EXPR_MAX_INSTRS 则放弃融合） ─────────── +// 用 abs 链（abs 不可代数化简、不可常量折叠）构造可控指令数 +nn::ExprSpec make_abs_chain(std::uint8_t n) +{ + nn::ExprSpec s; + s.views.push_back(nn::expr::linear()); + nn::ExprOperand prev = nn::expr::input(0); + for (std::uint8_t i = 0; i < n; ++i) + { + nn::ExprInstr in; + in.op = static_cast(nn::ExprOp::Abs); + in.dst = i; + in.a = prev; + s.instrs.push_back(in); + prev = nn::expr::reg(i); + } + s.num_regs = n; + return s; +} + +void test_limit_pressure() +{ + const std::size_t rows = 1, cols = 4; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + + // A：40 条 abs 指令；B：1 条 → 拼接后 41 ≤ 64 → 融合 + const int nA = g.add_node(make_abs_chain(40), + std::vector{x}, rows, cols, false); + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{yt}, rows, cols, false); + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 1, "40+1 条指令应在 EXPR_MAX_INSTRS=64 内 → 融合"); + + // A：60 条 abs；B：5 条 abs 链 → 拼接后 65 > 64 → 不融合 + nn::ExprGraph g2; + nn::Tensor x2 = make_tensor(rows, cols, 0.0f, 0.1f); + const int nC = g2.add_node(make_abs_chain(60), + std::vector{x2}, rows, cols, false); + nn::Tensor yt2 = nn::Tensor::cpu(rows, cols); + yt2.set_virtual_tag(static_cast(nC) + 1); + // B：5 条 abs 链（消费 A 输出） + nn::ExprSpec sb = make_abs_chain(5); + (void)g2.add_node(sb, std::vector{yt2}, rows, cols, false); + + auto kernels2 = fuse_expr_graph(g2); + CHECK(kernels2.size() == 2, "60+5 条指令超 EXPR_MAX_INSTRS → 保守放弃融合"); + std::printf("[OK] test_limit_pressure\n"); +} + +// ── 测试 7：确定性（同一图多次融合结果一致） ─────────────────────────── +void test_determinism() +{ + nn::CpuEngine eng; + const std::size_t rows = 2, cols = 4; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + + auto build = [&](nn::ExprGraph& g) + { + const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols, false); + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + (void)g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{yt}, rows, cols, false); + }; + + nn::ExprGraph g1, g2; + build(g1); + build(g2); + auto k1 = fuse_expr_graph(g1); + auto k2 = fuse_expr_graph(g2); + CHECK(k1.size() == k2.size(), "确定性：kernel 数一致"); + if (k1.size() != k2.size()) return; + const std::string key1 = nn::expr_spec_key(k1[0].spec); + const std::string key2 = nn::expr_spec_key(k2[0].spec); + CHECK(key1 == key2, "确定性：融合 kernel 的 key 完全一致"); + CHECK(nn::expr_spec_equal(k1[0].spec, k2[0].spec), "确定性：spec 完全一致"); + std::printf("[OK] test_determinism\n"); +} + +// ── 测试 8：录制图 add_node 依赖识别 ─────────────────────────────────── +void test_dep_recognition() +{ + const std::size_t rows = 2, cols = 4; + nn::ExprGraph g; + nn::Tensor x = make_tensor(rows, cols, 0.0f, 0.1f); + + const int nA = g.add_node(make_unary_const(nn::ExprOp::Mul, Scalar{2}), + std::vector{x}, rows, cols, false); + // 占位带 tag → 依赖识别 + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + const int nB = g.add_node(make_unary_const(nn::ExprOp::Add, Scalar{1}), + std::vector{yt}, rows, cols, false); + CHECK(g.nodes[static_cast(nB)].dep_of_input[0] == nA, + "B 的输入 0 应识别为依赖节点 A"); + // 普通输入 → 外部 + const int nC = g.add_node(make_unary_const(nn::ExprOp::Sub, Scalar{1}), + std::vector{x}, rows, cols, false); + CHECK(g.nodes[static_cast(nC)].dep_of_input[0] == -1, + "C 的输入 0 应为外部输入"); + // 未知 tag → 外部 + nn::Tensor fake = nn::Tensor::cpu(rows, cols); + fake.set_virtual_tag(999); + const int nD = g.add_node(make_unary_const(nn::ExprOp::Abs, Scalar{0}), + std::vector{fake}, rows, cols, false); + CHECK(g.nodes[static_cast(nD)].dep_of_input[0] == -1, + "未知 tag 视为外部输入"); + std::printf("[OK] test_dep_recognition\n"); +} + +// ── 测试 9：融合 kernel 输入顺序确定 ─────────────────────────────────── +// A: y = x1 * x2（两个输入）；B: z = y + x3 → 融合后 inputs = [x1, x2, x3] +void test_input_order() +{ + nn::CpuEngine eng; + const std::size_t rows = 2, cols = 4; + nn::ExprGraph g; + nn::Tensor x1 = make_tensor(rows, cols, 0.0f, 0.1f); + nn::Tensor x2 = make_tensor(rows, cols, 1.0f, 0.1f); + nn::Tensor x3 = make_tensor(rows, cols, 2.0f, 0.1f); + + // A: y = x1 * x2 + const int nA = g.add_node(make_binary(nn::ExprOp::Mul), + std::vector{x1, x2}, rows, cols, false); + // B: z = y + x3（输入 0 = y 依赖 A，输入 1 = x3 外部） + nn::Tensor yt = nn::Tensor::cpu(rows, cols); + yt.set_virtual_tag(static_cast(nA) + 1); + nn::ExprSpec sb; + sb.views.push_back(nn::expr::linear()); + sb.views.push_back(nn::expr::linear()); + nn::ExprInstr in; + in.op = static_cast(nn::ExprOp::Add); + in.dst = 0; + in.a = nn::expr::input(0); + in.b = nn::expr::input(1); + sb.instrs.push_back(in); + sb.num_regs = 1; + (void)g.add_node(sb, std::vector{yt, x3}, rows, cols, false); + + auto kernels = fuse_expr_graph(g); + CHECK(kernels.size() == 1, "A→B 链融合为单 kernel"); + if (kernels.size() != 1) return; + CHECK(kernels[0].inputs.size() == 3, "融合 kernel 应有 3 个外部输入"); + if (kernels[0].inputs.size() != 3) return; + // 顺序：A 的输入（x1, x2）在前，B 的非链输入（x3）在后 + CHECK(kernels[0].inputs[0].external.cpu_matrix().span()[0] == 0.0f, + "输入 0 = x1"); + CHECK(kernels[0].inputs[1].external.cpu_matrix().span()[0] == 1.0f, + "输入 1 = x2"); + CHECK(kernels[0].inputs[2].external.cpu_matrix().span()[0] == 2.0f, + "输入 2 = x3"); + + // 数值验证：z = x1*x2 + x3 + std::vector kinputs; + for (auto& in : kernels[0].inputs) kinputs.push_back(in.external); + auto fused = eng.eval_expr(kernels[0].spec, kinputs, rows, cols); + CHECK(static_cast(fused), "融合 kernel 求值成功"); + if (!fused) return; + bool ok = true; + for (std::size_t i = 0; i < rows * cols; ++i) + { + const Scalar expect = x1.cpu_matrix().span()[i] * x2.cpu_matrix().span()[i] + + x3.cpu_matrix().span()[i]; + if (std::fabs(fused->cpu_matrix().span()[i] - expect) > 1e-5f) ok = false; + } + CHECK(ok, "融合结果 = x1*x2 + x3"); + std::printf("[OK] test_input_order\n"); +} + +// ── 测试 10：IR-D emitter 抽象（一份 canonical IR → 多后端） ─────────── +void test_emitter() +{ + // 逐元素 spec：out = input0 * 2 + nn::ExprSpec s = make_unary_const(nn::ExprOp::Mul, Scalar{2}); + + auto glsl = nn::emitter_registry::make("glsl"); + auto cpu = nn::emitter_registry::make("cpu"); + auto none = nn::emitter_registry::make("nonexistent"); + CHECK(static_cast(glsl), "GlslEmitter 已注册可创建"); + CHECK(static_cast(cpu), "CpuEmitter 已注册可创建"); + CHECK(!none, "未知后端返回 nullptr"); + + if (glsl) + { + CHECK(std::string(glsl->name()) == "glsl", "GlslEmitter 后端名"); + const std::string g = glsl->generate("fused_x", s); + CHECK(!g.empty(), "GlslEmitter 生成逐元素 GLSL 成功"); + CHECK(g.find("#version 450") != std::string::npos, "GLSL 含版本头"); + } + if (cpu) + { + CHECK(std::string(cpu->name()) == "cpu", "CpuEmitter 后端名"); + const std::string c = cpu->generate("cpu_x", s); + CHECK(!c.empty(), "CpuEmitter 生成 C++ 代码成功"); + CHECK(c.find("inline void cpu_x") != std::string::npos, "C++ 代码含函数签名"); + CHECK(c.find("std::max") == std::string::npos, "逐元素含基本运算"); + // 归约：CpuEmitter 暂不支持 → 空串；GlslEmitter 支持 + nn::ExprSpec rs = make_row_sum(); + CHECK(cpu->generate_reduce("cpu_r", rs).empty(), "CpuEmitter 归约返回空（不支持）"); + if (glsl) + CHECK(!glsl->generate_reduce("fused_r", rs).empty(), "GlslEmitter 归约生成成功"); + } + std::printf("[OK] test_emitter\n"); +} + +} // namespace + +int main() +{ + test_chain_fuse(); + test_chain3_fuse(); + test_independent(); + test_multi_consumer(); + test_reduce_boundary(); + test_limit_pressure(); + test_determinism(); + test_dep_recognition(); + test_input_order(); + test_emitter(); + + if (g_fail == 0) + { + std::printf("expr_graph_test: ALL PASS\n"); + return 0; + } + std::printf("expr_graph_test: %d FAILED\n", g_fail); + return 1; +} diff --git a/src/matmul_probe.cpp b/src/matmul_probe.cpp deleted file mode 100644 index a082e607..00000000 --- a/src/matmul_probe.cpp +++ /dev/null @@ -1,211 +0,0 @@ -// ── batched_matmul 原语探针 ───────────────────────────────────────────── -// -// 目的：直接对比 CpuEngine::batched_matmul 四种转置组合与手工参考计算。 -// 若某组合输出不匹配，即证明底层原语有 bug。 -// -// 语义（见 cpu_engine.hpp::batched_matmul）： -// A 按行切成 batch 个纵向块（每块 a_rows_per × a.cols()），B 同理。 -// transA=true 表示使用 A 块的转置，transB=true 表示使用 B 块的转置。 -// 块内维度须匹配：所有四种转置组合都要求块为方阵（a_rows_per == a.cols()）。 -// -// 用法：matmul_probe [--gpu|--cuda] -// --gpu 使用 Vulkan GpuEngine（验证 tiled batched_matmul shader） -// --cuda 使用 CUDA CudaEngine -// 默认 CpuEngine -// ───────────────────────────────────────────────────────────────────────── - -#include -#include - -#include -#include -#include -#include - -using nn::Scalar; -using nn::Matrix; -using nn::Tensor; -using nn::ComputeEngine; - -namespace { - -// 逐块参考：C_b = A_b^T × B_b（A_b/B_b 为原矩阵中的第 bi 个行块） -// A_b 以 (a_rows_per, a_cols) 行主序存储；transA/transB 语义与引擎一致。 -Matrix ref_batched(const Matrix& A, const Matrix& B, - std::size_t batch, bool transA, bool transB) -{ - const std::size_t a_rows_per = A.rows() / batch; - const std::size_t b_rows_per = B.rows() / batch; - - const std::size_t M = transA ? A.cols() : a_rows_per; - const std::size_t N = transB ? b_rows_per : B.cols(); - - Matrix C(batch * M, N); - for (std::size_t bi = 0; bi < batch; ++bi) - { - const std::size_t a_off = bi * a_rows_per; - const std::size_t b_off = bi * b_rows_per; - for (std::size_t i = 0; i < M; ++i) - { - for (std::size_t j = 0; j < N; ++j) - { - Scalar s{0}; - // 块内 K 维 = 逻辑 A 的列数 = 逻辑 B 的行数 - const std::size_t K = transA ? a_rows_per : A.cols(); - for (std::size_t k = 0; k < K; ++k) - { - // A_b[i][k]（考虑转置） - const Scalar a_val = transA - ? A.at_unchecked(a_off + k, i) // A_b^T[i][k] = A_b[k][i] - : A.at_unchecked(a_off + i, k); - // B_b[k][j]（考虑转置） - const Scalar b_val = transB - ? B.at_unchecked(b_off + j, k) // B_b^T[k][j] = B_b[j][k] - : B.at_unchecked(b_off + k, j); - s += a_val * b_val; - } - C.set_value_unchecked(bi * M + i, j, s); - } - } - } - return C; -} - -void dump_diff(const std::string& name, const Matrix& got, const Matrix& ref) -{ - bool ok = true; - Scalar max_err{0}; - std::size_t bad = 0; - for (std::size_t i = 0; i < got.rows(); ++i) - for (std::size_t j = 0; j < got.cols(); ++j) - { - const Scalar e = std::fabs(got.at_unchecked(i, j) - ref.at_unchecked(i, j)); - if (e > max_err) max_err = e; - if (e > Scalar{1e-4}) - { - ok = false; - if (bad < 5) - std::cout << " [FAIL] " << name << "(" << i << "," << j - << ") got=" << got.at_unchecked(i, j) - << " ref=" << ref.at_unchecked(i, j) - << " err=" << e << "\n"; - ++bad; - } - } - std::cout << " " << name << ": " << (ok ? "OK" : "FAIL") - << " max_err=" << max_err << " bad=" << bad << "\n"; -} - -} // namespace - -int main(int argc, char* argv[]) -{ - nn::cli::EngineConfig ecfg; - for (int i = 1; i < argc; ++i) - { - const std::string_view arg{argv[i]}; - if (arg == "--gpu") ecfg.use_gpu = true; - else if (arg == "--cuda") ecfg.use_cuda = true; - else { std::cerr << "未知参数: " << arg << "\n"; return 1; } - } - auto engine_res = nn::cli::create_engine(ecfg, std::cout); - if (!engine_res) { std::cerr << "引擎创建失败: " << engine_res.error().message << "\n"; return 1; } - auto engine = std::move(*engine_res); - ComputeEngine& eng = *engine; - - std::cout << "========================================\n"; - std::cout << " batched_matmul 原语探针\n"; - std::cout << "========================================\n"; - - std::mt19937_64 rng(7); - std::uniform_real_distribution dist(-1, 1); - - // 方块尺寸 K × K，H 个块纵向堆叠 → A/B 均为 (H*K, K) - // 该形状使四种转置组合的块内维度全部合法。 - const std::size_t K = 4; - const std::size_t H = 2; - - Matrix Am(H * K, K), Bm(H * K, K); - for (std::size_t i = 0; i < Am.size(); ++i) Am.span()[i] = dist(rng); - for (std::size_t i = 0; i < Bm.size(); ++i) Bm.span()[i] = dist(rng); - - auto A = eng.from_matrix(Am); - auto B = eng.from_matrix(Bm); - if (!A || !B) { std::cerr << "from_matrix failed\n"; return 1; } - - // ── (F,F): C = A × B ── - { - auto out = eng.batched_matmul(*A, *B, H, false, false); - if (!out) { std::cerr << " (F,F) error: " << out.error().message << "\n"; return 1; } - auto got = eng.to_matrix(*out); - if (!got) { std::cerr << " (F,F) to_matrix failed\n"; return 1; } - const Matrix ref = ref_batched(Am, Bm, H, false, false); - dump_diff("(F,F) A*B", *got, ref); - } - // ── (F,T): C = A × B^T ── - { - auto out = eng.batched_matmul(*A, *B, H, false, true); - if (!out) { std::cerr << " (F,T) error: " << out.error().message << "\n"; return 1; } - auto got = eng.to_matrix(*out); - if (!got) { std::cerr << " (F,T) to_matrix failed\n"; return 1; } - const Matrix ref = ref_batched(Am, Bm, H, false, true); - dump_diff("(F,T) A*B^T", *got, ref); - } - // ── (T,F): C = A^T × B ── - { - auto out = eng.batched_matmul(*A, *B, H, true, false); - if (!out) { std::cerr << " (T,F) error: " << out.error().message << "\n"; return 1; } - auto got = eng.to_matrix(*out); - if (!got) { std::cerr << " (T,F) to_matrix failed\n"; return 1; } - const Matrix ref = ref_batched(Am, Bm, H, true, false); - dump_diff("(T,F) A^T*B", *got, ref); - } - // ── (T,T): C = A^T × B^T ── - { - auto out = eng.batched_matmul(*A, *B, H, true, true); - if (!out) { std::cerr << " (T,T) error: " << out.error().message << "\n"; return 1; } - auto got = eng.to_matrix(*out); - if (!got) { std::cerr << " (T,T) to_matrix failed\n"; return 1; } - const Matrix ref = ref_batched(Am, Bm, H, true, true); - dump_diff("(T,T) A^T*B^T", *got, ref); - } - - // ── 大尺寸 + alpha：K=100 不被 BK=16 整除（边界填零路径），M=N=100 - // 超过 64×64 tile（多 WorkGroup 分块路径），H=8 模拟 batch*num_heads。 - // alpha=0.176777 模拟 attention 的 1/sqrt(d_k) 折叠。 - { - const std::size_t K2 = 100; - const std::size_t H2 = 8; - const Scalar alpha{0.176777f}; - - Matrix Am2(H2 * K2, K2), Bm2(H2 * K2, K2); - for (std::size_t i = 0; i < Am2.size(); ++i) Am2.span()[i] = dist(rng); - for (std::size_t i = 0; i < Bm2.size(); ++i) Bm2.span()[i] = dist(rng); - - auto A2 = eng.from_matrix(Am2); - auto B2 = eng.from_matrix(Bm2); - if (!A2 || !B2) { std::cerr << "from_matrix (large) failed\n"; return 1; } - - struct Case { const char* name; bool tA; bool tB; }; - for (const Case cs : {Case{"(F,F)", false, false}, Case{"(T,F)", true, false}, - Case{"(F,T)", false, true}, Case{"(T,T)", true, true}}) - { - auto out = eng.batched_matmul(*A2, *B2, H2, cs.tA, cs.tB, alpha); - if (!out) { std::cerr << " large " << cs.name << " error: " - << out.error().message << "\n"; return 1; } - auto got = eng.to_matrix(*out); - if (!got) { std::cerr << " large " << cs.name << " to_matrix failed\n"; return 1; } - const Matrix ref = ref_batched(Am2, Bm2, H2, cs.tA, cs.tB); - - // ref 乘 alpha 后对比（参考实现不含 alpha） - Matrix ref_s = ref; - for (std::size_t i = 0; i < ref_s.size(); ++i) - ref_s.span()[i] *= alpha; - std::string label = std::string("large ") + cs.name + " alpha"; - dump_diff(label.c_str(), *got, ref_s); - } - } - - std::cout << "----------------------------------------\n"; - return 0; -} diff --git a/src/mnist_train.cpp b/src/mnist_train.cpp index 74c84620..0ac39ecf 100644 --- a/src/mnist_train.cpp +++ b/src/mnist_train.cpp @@ -509,7 +509,10 @@ int main(int argc, char *argv[]) for (std::size_t i = 0; i < dims.size(); ++i) { std::cout << dims[i]; - if (i < dims.size() - 2) + // norm+GeLU 挂在隐藏层 Linear 的输出维度上：dims[1]..dims[size-2]。 + // 首元素 dims[0] 是输入、末元素 dims[size-1] 是 logits 输出， + // 两者都不挂 norm/激活（CE 损失需要原始 logits）。 + if (i > 0 && i < dims.size() - 1) std::cout << "(" << norm_name << "+GeLU)"; if (i < dims.size() - 1) std::cout << " -> "; diff --git a/src/reduce_bench.cpp b/src/reduce_bench.cpp new file mode 100644 index 00000000..5cac227a --- /dev/null +++ b/src/reduce_bench.cpp @@ -0,0 +1,175 @@ +// ─────────────────────────────────────────────────────────────────────────── +// reduce_bench.cpp — 归约 kernel 性能对比（旧共享内存树 vs 新 warp shuffle） +// +// 旧算法：GpuEngine::row_reduce_sum / col_reduce_sum +// → 独立 reduce.comp（单累加器 + 共享内存树形归约） +// 新算法：dsl::compute 触发 AOT 融合归约 shader（fused_*） +// → 单累加器 + warp shuffle(subgroup) 归约 +// 两者均为"单累加器 + 逐线程跨步累加"，唯一差异是最终树形归约实现 +// （共享内存树 vs subgroup 蝴蝶归约）→ 干净隔离 warp shuffle 收益。 +// +// 指标：GB/s = 读入整矩阵字节数 / 耗时。归约 kernel 主要受内存带宽限制。 +// +// 用法：reduce_bench [shapes...] （默认若干大形状） +// ─────────────────────────────────────────────────────────────────────────── + +#include + +#include +#include +#include +#include +#include +#include +#include + +using nn::Scalar; +namespace dsl = nn::dsl; + +#ifndef NN_HAS_VULKAN +int main() +{ + std::printf("[SKIP] 需要 Vulkan SDK（-DNN_HAS_VULKAN）\n"); + return 0; +} +#else +#include + +using nn::GpuBackend; +using nn::GpuEngine; +using nn::Tensor; + +namespace +{ + +// 计时一个可空返回的 GPU 调用，返回单次平均耗时（秒） +template +double time_call(F&& f, int iters) +{ + // warmup + f(); + auto t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < iters; ++i) + { + auto r = f(); + if (!r) { std::printf("[ERR] 调用失败: %s\n", r.error().message.c_str()); return -1.0; } + } + auto t1 = std::chrono::steady_clock::now(); + return std::chrono::duration(t1 - t0).count() / iters; +} + +// 扣除启动/同步开销后的内核吞吐（GB/s）。overhead 用同算子小矩阵（内核≈0）测得。 +// 目的：去除逐调用 submit+wait 的固定开销，还原真实内核带宽，判断跨步访问是否真瓶颈。 +double kernel_gbps(double bytes, double shape_s, double overhead_s) +{ + const double kt = shape_s - overhead_s; + if (kt <= 0.0) return 0.0; + return bytes / kt / 1e9; +} + +void bench_shape(GpuEngine& gpu, std::size_t rows, std::size_t cols, int iters) +{ + const double bytes = static_cast(rows) * cols * sizeof(Scalar); + const double mbytes = bytes / 1e6; // MB + + std::mt19937 rng(12345); + std::uniform_real_distribution dist(-1.0f, 1.0f); + Tensor x = Tensor::cpu(rows, cols); + for (auto& v : x.cpu_matrix().span()) v = dist(rng); + const Tensor xt = x; + + // ── 启动/同步开销：小矩阵上的归约（内核时间≈0），同算子测一遍 ── + Tensor xoh = Tensor::cpu(256, 256); + for (auto& v : xoh.cpu_matrix().span()) v = 1.0f; + const Tensor xoht = xoh; + const double overhead = time_call([&]{ return gpu.col_reduce_sum(xoht); }, iters); + + // ── 行求和：旧（reduce_gpu） vs 新（融合 row_reduce_sum）── + // 新路径用 compute_reduce（输出归约向量 (rows,1)），与旧路径输出对齐， + // 避免 compute() 把归约结果广播写到整个 (rows,cols)（多一次整矩阵写）。 + const double old_rs = time_call([&]{ return gpu.row_reduce_sum(xt); }, iters); + const double new_rs = time_call([&]{ + return dsl::compute_reduce(gpu, dsl::row_reduce_sum(dsl::leaf(xt)), rows, cols); }, iters); + + // ── 列求和：旧（reduce_gpu） vs 新（融合 col_reduce_sum）── + const double old_cs = time_call([&]{ return gpu.col_reduce_sum(xt); }, iters); + const double new_cs = time_call([&]{ + return dsl::compute_reduce(gpu, dsl::col_reduce_sum(dsl::leaf(xt)), rows, cols); }, iters); + + std::printf("shape %zux%-8zu (%.0f MB) 启动开销≈%.0fus\n", + rows, cols, mbytes, overhead * 1e6); + const double old_rs_g = kernel_gbps(bytes, old_rs, overhead); + const double new_rs_g = kernel_gbps(bytes, new_rs, overhead); + const double old_cs_g = kernel_gbps(bytes, old_cs, overhead); + const double new_cs_g = kernel_gbps(bytes, new_cs, overhead); + std::printf(" 行求和(内核): 旧 %7.1f GB/s 新 %7.1f GB/s 加速 %s\n", + old_rs_g, new_rs_g, + (old_rs_g > 0 && new_rs_g > 0) ? (std::to_string(old_rs / new_rs) + "x").c_str() : "-"); + std::printf(" 列求和(内核): 旧 %7.1f GB/s 新 %7.1f GB/s 加速 %s\n", + old_cs_g, new_cs_g, + (old_cs_g > 0 && new_cs_g > 0) ? (std::to_string(old_cs / new_cs) + "x").c_str() : "-"); + std::printf(" 跨步惩罚: 列/行(新) = %s (1.0=无惩罚；<1=列跨步访问更慢)\n", + (new_rs_g > 0 && new_cs_g > 0) ? (std::to_string(new_cs_g / new_rs_g) + "x").c_str() : "-"); +} + +// 纯内核基准：数据一次性上传到显存后驻留，反复调 reduce_gpu（不再上传）。 +// 目的：去掉 CPU→GPU 上传（PRIME/PCIe）对测量的污染，还原真实内核带宽， +// 判断列归约跨步访问的纯内核惩罚到底多大（决定 tile 化是否值得）。 +void bench_pure_kernel(GpuBackend& backend, std::size_t rows, std::size_t cols, int iters) +{ + std::mt19937 rng(777); + std::uniform_real_distribution dist(-1.0f, 1.0f); + nn::Matrix m(rows, cols); + for (auto& v : m.span()) v = dist(rng); + auto gt_r = nn::GpuTensor::from_matrix(m, backend); + if (!gt_r) { std::printf("[ERR] 上传到显存失败: %s\n", gt_r.error().message.c_str()); return; } + nn::GpuTensor gt = std::move(*gt_r); + const double bytes = static_cast(rows) * cols * sizeof(Scalar); + + const double t_row = time_call([&]{ return backend.reduce_gpu(gt, 0u, 0u); }, iters); // 行,sum + const double t_col = time_call([&]{ return backend.reduce_gpu(gt, 1u, 0u); }, iters); // 列,sum + const double row_g = t_row > 0 ? bytes / t_row / 1e9 : 0.0; + const double col_g = t_col > 0 ? bytes / t_col / 1e9 : 0.0; + std::printf(" 纯内核(数据驻留): 行归约 %7.1f GB/s 列归约 %7.1f GB/s 列/行 = %s\n", + row_g, col_g, + (row_g > 0 && col_g > 0) ? (std::to_string(col_g / row_g) + "x").c_str() : "-"); +} + +} // namespace + +int main(int argc, char** argv) +{ + std::printf("reduce_bench — 旧(共享内存树) vs 新(warp shuffle) 归约吞吐\n"); + auto& backend = GpuBackend::instance(); + auto init_r = backend.initialize(); + if (!init_r) + { + std::printf("[ERR] GPU 初始化失败: %s\n", init_r.error().message.c_str()); + return 1; + } + GpuEngine gpu(backend); + + // 默认形状；可用命令行覆盖：reduce_bench 4096,4096 8192,8192 ... + std::vector> shapes; + if (argc > 1) + { + for (int i = 1; i < argc; ++i) + { + std::size_t r = 0, c = 0; + if (std::sscanf(argv[i], "%zu,%zu", &r, &c) == 2) + shapes.push_back({r, c}); + } + } + if (shapes.empty()) + shapes = {{4096, 4096}, {8192, 8192}, {1024, 65536}, {65536, 1024}}; + + for (auto& [r, c] : shapes) + { + // 迭代次数随规模缩放，避免总耗时过长 + const int iters = (r * c >= 256ull * 1024 * 1024) ? 3 : 10; + bench_shape(gpu, r, c, iters); + bench_pure_kernel(backend, r, c, iters); + } + return 0; +} +#endif // NN_HAS_VULKAN diff --git a/tools/gen_fused.cpp b/tools/gen_fused.cpp index 7d957e4d..af4da0cb 100644 --- a/tools/gen_fused.cpp +++ b/tools/gen_fused.cpp @@ -2,14 +2,20 @@ // gen_fused.cpp — AOT 算子融合：构建期合成（读 scan_exprs 的 bin） // // 读取 scan_exprs dump 的 expr_specs.bin（表达式结构集合），对每条 -// 用 glsl_gen.hpp 展开为单个融合 .comp，再经 glslc 编译成 SPIR-V，最后 -// 内联进单个生成头 fused_registry.hpp（key → {ExprSpec, SPIR-V}）。 +// 用 emitter 抽象（IR-D：expr_emitter.hpp）展开为单个融合 kernel 源码， +// 再经 glslc 编译成 SPIR-V，最后内联进单个生成头 fused_registry.hpp +// （key → {ExprSpec, SPIR-V}）。 +// +// IR-D 落地：本工具经 nn::emitter_registry 选择后端（默认 "glsl" = +// GlslEmitter），不再直接绑定 GLSL 生成函数——同一份 canonical IR 可由 +// CpuEmitter 等其它后端展开（--list-backends 查看；CpuEmitter 生成 C++ +// 直线代码，供调试/交叉验证，不产出 SPIR-V）。 // // 表达式**文本只出现在 Layer**；本工具只消费折叠后的结构（派生物）。 // 产物 fused_registry.hpp 供运行时（GpuEngine::eval_expr / vk_backend） // 按 expr_spec_key 精确匹配 dispatch——闭合世界，未命中硬报错。 // -// 用法： gen_fused +// 用法： gen_fused [--list-backends] // 产物： /fused_registry.hpp（+ 调试用 .comp/.spv） // ─────────────────────────────────────────────────────────────────────────── @@ -24,7 +30,9 @@ #include "expr_spec.hpp" #include "expr_registry.hpp" -#include "glsl_gen.hpp" +#include "expr_emitter.hpp" // IR-D：emitter 抽象（后端选择） +#include "glsl_gen.hpp" // 注册 GlslEmitter（默认后端） +#include "cpu_emitter.hpp" // 注册 CpuEmitter（IR-D 第二后端，供 --list-backends） namespace { @@ -34,8 +42,10 @@ namespace { // 仅给 glslc 路径加引号；src/dst 为无空格路径（out_dir 由 CMake 控制）， // 正斜杠相对路径不加引号，避免 cmd/system 解析问题。 + // --target-env=vulkan1.2 → SPIR-V 1.5：subgroup 算子（GL_KHR_shader_subgroup_*） + // 需 SPIR-V ≥1.3；现基本找不到 <1.2 的显卡，直接定 vulkan1.2。 const std::string cmd = - "\"" + glslc + "\" -fshader-stage=compute -o " + dst + " " + src; + "\"" + glslc + "\" -fshader-stage=compute --target-env=vulkan1.2 -o " + dst + " " + src; return std::system(cmd.c_str()) == 0; } @@ -89,9 +99,19 @@ namespace int main(int argc, char* argv[]) { + // --list-backends：列出可用 emitter 后端（IR-D 多后端验证） + for (int i = 1; i < argc; ++i) + if (std::string(argv[i]) == "--list-backends") + { + std::printf("[gen] 可用 emitter 后端（IR-D）:\n"); + for (const auto& n : nn::emitter_registry::names()) + std::printf(" - %s\n", n.c_str()); + return 0; + } + if (argc < 4) { - std::fprintf(stderr, "用法: gen_fused \n"); + std::fprintf(stderr, "用法: gen_fused [--list-backends]\n"); return 2; } const std::string out_dir = argv[1]; @@ -132,6 +152,7 @@ int main(int argc, char* argv[]) H << " std::size_t spirv_words;\n"; H << " int reduce_axis; // -1=逐元素, 0=行归约, 1=列归约\n"; H << " std::uint32_t view_param_count; // 运行时视图参数个数（RowMod/RotateHalf 的 push constant vp 槽）\n"; + H << " std::uint32_t vec_width; // 每线程处理元素数（1=标量, 4=vec4）\n"; H << "};\n\n"; for (const auto& spec : reg.specs) @@ -148,10 +169,18 @@ int main(int argc, char* argv[]) const std::string comp_path = out_dir + "/fused_" + key + ".comp"; const std::string spv_path = out_dir + "/fused_" + key + ".spv"; + // 经 emitter 抽象（IR-D）生成 kernel 源码：默认 GlslEmitter。 + // 同一份 canonical IR 可由其它后端（CpuEmitter）展开（--list-backends）。 + auto emitter = nn::emitter_registry::make("glsl"); + if (!emitter) + { + std::fprintf(stderr, "[FAIL] 无法创建 GLSL emitter（IR-D 注册表异常）\n"); + return 1; + } // 含归约（raxis >= 0）→ 归约 kernel（workgroup 级共享内存归约）；否则逐元素 const std::string glsl = (raxis >= 0) - ? nn::generate_glsl_reduce("fused_" + key, spec) - : nn::generate_glsl("fused_" + key, spec); + ? emitter->generate_reduce("fused_" + key, spec) + : emitter->generate("fused_" + key, spec); if (glsl.empty()) { std::fprintf(stderr, "[FAIL] 生成 %s 失败（归约结构不支持？）\n", key.c_str()); @@ -190,11 +219,13 @@ int main(int argc, char* argv[]) if (raxis == -2) continue; // 与上方跳过保持一致 const std::string key = nn::expr_spec_key(spec); + const std::uint32_t vecw = + (raxis < 0 && nn::glsl_vec4_eligible(spec)) ? 4u : 1u; H << " { \"" << key << "\",\n " << emit_spec(spec) << ",\n" << " kSpirv_" << key << ", sizeof(kSpirv_" << key << ")/sizeof(std::uint32_t), " << raxis << ", " - << nn::expr_spec_runtime_view_param_count(spec) << " },\n"; + << nn::expr_spec_runtime_view_param_count(spec) << ", " << vecw << " },\n"; } H << "};\n"; H << "inline constexpr std::size_t kFusedShaderCount =\n" diff --git a/tools/scan_exprs.cpp b/tools/scan_exprs.cpp index 6cc7ea57..6d8184d8 100644 --- a/tools/scan_exprs.cpp +++ b/tools/scan_exprs.cpp @@ -47,7 +47,7 @@ int main(int argc, char* argv[]) for (const std::size_t dk : {std::size_t{16}, std::size_t{32}, std::size_t{64}, std::size_t{128}}) { - nn::RotaryEmbedding rope(engine, dk); + nn::RotaryEmbedding rope(dk); nn::Tensor q = nn::Tensor::cpu(2 * dk, 8); // rows 为 dk 的整数倍 (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/false); (void)rope.apply(engine, q, /*seq=*/8, /*backward=*/true); @@ -85,7 +85,8 @@ int main(int argc, char* argv[]) // backward: grad_x 列归约表达式 + grad_gamma 行归约表达式 { const std::size_t F = 8, B = 5; - nn::RMSNorm rms(engine, F); + nn::RMSNorm rms(F); + (void)rms.init(engine); nn::Tensor input = nn::Tensor::cpu(F, B); (void)rms.forward(engine, input); // 填充 normed/rms_inv 缓存 + 登记 fwd 结构 nn::Tensor grad = nn::Tensor::cpu(F, B); @@ -97,13 +98,29 @@ int main(int argc, char* argv[]) // backward: grad_x 列归约表达式 + grad_gamma/grad_beta 行归约表达式 { const std::size_t F = 8, B = 5; - nn::LayerNorm ln(engine, F); + nn::LayerNorm ln(F); + (void)ln.init(engine); nn::Tensor input = nn::Tensor::cpu(F, B); (void)ln.forward(engine, input); nn::Tensor grad = nn::Tensor::cpu(F, B); (void)ln.backward(engine, grad); } + // ── FusedChainLayer（IR-C：begin_expr/end_expr 图 IR 链式融合）──────── + // forward 的三个逐元素表达式在 end_expr 时融合成单个 kernel（t/u 内联为 + // 寄存器）。必须 dry-run 覆盖本路径：GPU 运行时 begin_expr/end_expr 融合 + // 出的复合 spec 才能在 fused_registry 中命中（闭合世界两端一致）。 + // backward 为独立表达式（非录制段），单独登记。 + { + const std::size_t F = 8, B = 5; + nn::FusedChainLayer chain(F); + (void)chain.init(engine); + nn::Tensor input = nn::Tensor::cpu(F, B); + (void)chain.forward(engine, input); + nn::Tensor grad = nn::Tensor::cpu(F, B); + (void)chain.backward(engine, grad); + } + auto& reg = nn::fused::global_registry(); if (!nn::fused::write_registry(out_path, reg)) {